A 3D test.

HARK/tests/test_discrete.py

@@ -4,11 +4,6 @@
 from __future__ import print_function, division
 from __future__ import absolute_import
 
-from builtins import str
-from builtins import zip
-from builtins import range
-from builtins import object
-
 from HARK import interpolation
 
 # Bring in modules we need
@@ -19,19 +14,22 @@
 class testsForDiscreteChoice(unittest.TestCase):
 
     def setUp(self):
-        self.Vs2D    = np.stack((np.zeros(3), np.ones(3)))
-#        self.Vs3D    = np.array([[0.0, 1.0, 4.0], [1.0, 2.0, 0.0], [3.0, 0.0, 2.0]])
+        self.Vs2D = np.stack((np.zeros(3), np.ones(3)))
+        self.Vref2D = np.array([1.0, 1.0, 1.0])
+        self.Pref2D = np.array([[0.0, 0.0, 0.0], [1.0, 1.0, 1.0]])
 
-        self.CRRA_vals = np.linspace(1.,10.,10)
+        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)))
 
-    def test_noShock2D(self):
+    def test_noShock2DBoth(self):
         # Test the value functions and policies of the 2D case
         sigma = 0.0
-        Vref = np.array([1.0, 1.0, 1.0])
-        Pref = np.array([[0.0, 0.0, 0.0], [1.0, 1.0, 1.0]])
         V, P = interpolation.discreteLogSumProb(self.Vs2D, sigma)
-        self.assertTrue((V, P), (Vref, Pref))
-
-
+        self.assertTrue((V, P), (self.Vref2D, self.Pref2D))
 
-#        self.assertTrue((np.array([3., 2., 4.]), array([[0., 0., 1.], [0., 1., 0.], [1., 0., 0.]]))
+    def test_noShock2DIndividual(self):
+        # Test the value functions and policies of the 2D case
+        sigma = 0.0
+        V = interpolation.discreteLogSum(self.Vs3D, sigma)
+        P = interpolation.discreteProb(self.Vs3D, sigma)
+        self.assertTrue(V[0], self.Vref3DFirst)