Add tests that a matrix valued distribution should pass
@@ -123,6 +123,61 @@ class DiscreteDistributionTests(unittest.TestCase):
self.assertAlmostEqual(ce9[3], 9.518015322143837)
+class MatrixDiscreteDistributionTests(unittest.TestCase):
+ """
+ Tests matrix-valued discrete distribution.
+ """
+
+ def setUp(self):
+
+ self.draw_1 = np.array([
+ [
+ [1.0,2.0,3.0],
+ [4.0,5.0,6.0]
+ ],
+ [
+ [7.0,8.0,9.0],
+ [10.0,11.0,12.0]
+ ],
+ ])
+
+ self.draw_2 = -1*self.draw_1
+
+ X = np.stack([self.draw_1, self.draw_2], axis = -1)
+ pmf = np.array([0.5,0.5])
+
+ self.mat_distr = DiscreteDistribution(pmf, X, seed = 0)
+
+ def test_draw(self):
+ """
+ Check that the draws are the matrices we
+ want them to be
+ """
+
+ draw = self.mat_distr.draw(1)
+ self.assertTrue(
+ np.allclose(draw[...,0], self.draw_2)
+ )
+
+
+ def test_expected(self):
+
+ # Expectation without transformation
+ exp = calc_expectation(self.mat_distr)
+
+ # Check the expectation is of the shape we want
+ self.assertTrue(exp.shape[0] == self.draw_1.shape[0])
+ self.assertTrue(exp.shape[1] == self.draw_1.shape[1])
+
+ # Check that its value is what we expect
+ self.assertTrue(
+ np.allclose(exp, 0.0)
+ )
+
+ # Expectation of the sum
+ exp = calc_expectation(self.mat_distr, func = np.sum)
+ self.assertTrue(float(exp) == 0.0)
+
class DistributionClassTests(unittest.TestCase):
"""
Tests for distribution.py sampling distributions
