more 3d.

econ-ark · Jan 17, 2019 · 6271c50 · 6271c50
1 parent 3f73ba4
commit 6271c50
Showing 1 changed file with 4 additions and 2 deletions.
diff --git a/HARK/tests/test_discrete.py b/HARK/tests/test_discrete.py
@@ -19,7 +19,9 @@ def setUp(self):
         self.Pref2D = np.array([[0.0, 0.0, 0.0], [1.0, 1.0, 1.0]])
 
         self.Vs3D = np.array([[0.0, 1.0, 4.0], [1.0, 2.0, 0.0], [3.0, 0.0, 2.0]])
-        self.Vref3DFirst = 3.0+np.log((np.exp(0.0-3.0)+np.exp(1.0-3.0)+np.exp(3.0-3.0)))
+        maxV = self.Vs3D.max()
+        self.Vref3D = maxV + np.log(sum(np.exp(self.Vs3D-maxV),axis=0))
+        self.Pref3D = np.log(sum(np.exp(self.Vs3D-maxV),axis=0))
 
     def test_noShock2DBoth(self):
         # Test the value functions and policies of the 2D case
@@ -32,4 +34,4 @@ def test_noShock2DIndividual(self):
         sigma = 0.0
         V = interpolation.discreteLogSum(self.Vs3D, sigma)
         P = interpolation.discreteProb(self.Vs3D, sigma)
-        self.assertTrue(V[0], self.Vref3DFirst)
+        self.assertTrue(V, self.Vref3D)