Add calcLogSum, calcChoiceProbs and combined calcLogSumChoiceProbs to HARK.interpolation.

HARK/interpolation.py

@@ -3367,6 +3367,87 @@ def _derY(self,x,y):
         dfdy = y_alpha*dfda + y_beta*dfdb
         return dfdy
 
+def calcLogSumChoiceProbs(Vals, sigma):
+    '''
+    Returns the final optimal value and choice probabilities given the choice
+    specific value functions `Vals`. Probabilities are degenerate if sigma == 0.0.
+    Parameters
+    ----------
+    Vals : [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.
+    P : [numpy.array]
+        A numpy.array that holds the discrete choice probabilities
+    '''
+
+    return calcLogSum(Vals, sigma), calcChoiceProbs(Vals, sigma)
+
+def calcChoiceProbs(Vals, sigma):
+    '''
+    Returns the choice probabilities given the choice specific value functions
+    `Vals`. Probabilities are degenerate if sigma == 0.0.
+
+    Parameters
+    ----------
+    Vals : [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
+    -------
+    Probs : [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(Vals, axis=0)
+        Probs = np.zeros(Vals.shape)
+        for i in range(Vals.shape[0]):
+            Probs[i][Pflat==i] = 1
+        return Probs
+
+    Probs = np.divide(np.exp((Vals-Vals[0])/sigma), np.sum(np.exp((Vals-Vals[0])/sigma), axis=0))
+    return Probs
+
+
+def calcLogSum(Vals, sigma):
+    '''
+    Returns the optimal value given the choice specific value functions Vals.
+
+    Parameters
+    ----------
+    Vals : [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(Vals, axis=0)
+        return V
+
+    # else we have a taste shock
+    maxV = Vals.max()
+
+    # calculate maxV+sigma*log(sum_i=1^J exp((V[i]-maxV))/sigma)
+    sumexp = np.sum(np.exp((Vals-maxV)/sigma), axis=0)
+    V = np.log(sumexp)
+    V = maxV + sigma*V
+    return V
+
 def main():
     print("Sorry, HARK.interpolation doesn't actually do much on its own.")
     print("To see some examples of its interpolation methods in action, look at any")

HARK/tests/test_discrete.py

@@ -0,0 +1,70 @@
+"""
+This file implements unit tests to check discrete choice functions
+"""
+from __future__ import print_function, division
+from __future__ import absolute_import
+
+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.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.Vs3D = np.array([[0.0, 1.0, 4.0], [1.0, 2.0, 0.0], [3.0, 0.0, 2.0]])
+        self.Vref3D = np.array([[3.0, 2.0, 4.0]])
+        self.Pref3D = np.array([[0, 0, 1], [0, 1, 0], [1, 0, 0]])
+        # maxV = self.Vs3D.max()
+        # self.Vref3D = maxV + np.log(np.sum(np.exp(self.Vs3D-maxV),axis=0))
+        # self.Pref3D = np.log(np.sum(np.exp(self.Vs3D-maxV),axis=0))
+
+    def test_noShock2DBothEqualValue(self):
+        # Test the value functions and policies of the 2D case
+        sigma = 0.0
+        V, P = interpolation.calcLogSumChoiceProbs(self.Vs2D, sigma)
+        self.assertTrue((V == self.Vref2D).all)
+        self.assertTrue((P == self.Pref2D).all)
+
+    def test_noShock2DBoth(self):
+        # Test the value functions and policies of the 2D case
+        sigma = 0.0
+        V, P = interpolation.calcLogSumChoiceProbs(self.Vs2D, sigma)
+        self.assertTrue((V == self.Vref2D).all)
+        self.assertTrue((P == self.Pref2D).all)
+
+    def test_noShock2DIndividual(self):
+        # Test the value functions and policies of the 2D case
+        sigma = 0.0
+        V = interpolation.calcLogSum(self.Vs2D, sigma)
+        P = interpolation.calcChoiceProbs(self.Vs2D, sigma)
+        self.assertTrue((V == self.Vref2D).all())
+        self.assertTrue((P == self.Pref2D).all())
+
+    def test_noShock3DBothEqualValue(self):
+        # Test the value functions and policies of the 3D case
+        sigma = 0.0
+        V, P = interpolation.calcLogSumChoiceProbs(self.Vs3D, sigma)
+        self.assertTrue((V == self.Vref3D).all)
+        self.assertTrue((P == self.Pref3D).all)
+
+    def test_noShock3DBoth(self):
+        # Test the value functions and policies of the 3D case
+        sigma = 0.0
+        V, P = interpolation.calcLogSumChoiceProbs(self.Vs3D, sigma)
+        self.assertTrue((V == self.Vref3D).all)
+        self.assertTrue((P == self.Pref3D).all)
+
+    def test_noShock3DIndividual(self):
+        # Test the value functions and policies of the 3D case
+        sigma = 0.0
+        V = interpolation.calcLogSum(self.Vs3D, sigma)
+        P = interpolation.calcChoiceProbs(self.Vs3D, sigma)
+        self.assertTrue((V == self.Vref3D).all())
+        self.assertTrue((P == self.Pref3D).all())