Changes names and start adding tests.

HARK/interpolation.py

@@ -3367,7 +3367,7 @@ def _derY(self,x,y):
         dfdy = y_alpha*dfda + y_beta*dfdb
         return dfdy
 
-def discreteEnvelope(Vs, sigma):
+def discreteLogSumProb(Vs, sigma):
     '''
     Returns the final optimal value and policies given the choice specific value
     functions Vs. Policies are degenerate if sigma == 0.0.
@@ -3385,22 +3385,68 @@ def discreteEnvelope(Vs, sigma):
         A numpy.array that holds the discrete choice probabilities
     '''
 
+    return discreteLogSum(Vs, sigma), discreteProb(Vs, sigma)
+
+def discreteProb(Vs, sigma):
+    '''
+    Returns the policies given the choice specific value functions Vs. Policies
+    are degenerate if sigma == 0.0.
+
+    Parameters
+    ----------
+    Vs : [numpy.array]
+        A numpy.array that holds choice specific values at common grid points.
+    sigma : float
+        A number that controls the variance of the taste shocks
+    Returns
+    -------
+    P : [numpy.array]
+        A numpy.array that holds the discrete choice probabilities
+    '''
+
     # Assumes that NaNs have been replaced by -numpy.inf or similar
     if sigma == 0.0:
         # We could construct a linear index here and use unravel_index.
         Pflat = np.argmax(Vs, axis=0)
         P = np.zeros(Vs.shape)
         for i in range(Vs.shape[0]):
             P[i][Pflat==i] = 1
+        return P
+
+    P = np.divide(np.exp((Vs-Vs[0])/sigma), np.sum(np.exp((Vs-Vs[0])/sigma), axis=0))
+    return P
+
+
+def discreteLogSum(Vs, sigma):
+    '''
+    Returns the optimal value given the choice specific value functions Vs.
+
+    Parameters
+    ----------
+    Vs : [numpy.array]
+        A numpy.array that holds choice specific values at common grid points.
+    sigma : float
+        A number that controls the variance of the taste shocks
+    Returns
+    -------
+    V : [numpy.array]
+        A numpy.array that holds the integrated value function.
+    '''
+
+    # Assumes that NaNs have been replaced by -numpy.inf or similar
+    if sigma == 0.0:
+        # We could construct a linear index here and use unravel_index.
         V = np.amax(Vs, axis=0)
-        return V, P
+        return V
+
+    # else we have a taste shock
     maxV = Vs.max()
-    P = np.divide(np.exp((Vs-Vs[0])/sigma), np.sum(np.exp((Vs-Vs[0])/sigma), axis=0))
 
+    # calculate maxV+sigma*log(sum_i=1^J exp((V[i]-maxV))/sigma)
     sumexp = np.sum(np.exp((Vs-maxV)/sigma), axis=0)
     V = np.log(sumexp)
     V = maxV + sigma*V
-    return V, P
+    return V
 
 def main():
     print("Sorry, HARK.interpolation doesn't actually do much on its own.")

HARK/tests/test_discrete.py

@@ -0,0 +1,37 @@
+"""
+This file implements unit tests to check discrete choice functions
+"""
+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
+import unittest
+import numpy as np
+
+
+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.CRRA_vals = np.linspace(1.,10.,10)
+
+    def test_noShock2D(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.discreteLogSum(self.Vs2D, sigma)
+        self.assertTrue((V, P), (Vref, Pref))
+
+
+
+#        self.assertTrue((np.array([3., 2., 4.]), array([[0., 0., 1.], [0., 1., 0.], [1., 0., 0.]]))