Merge pull request #633 from sbenthall/i632-markov

Fixes for ConsMarkovModel, example, tests

Author: mnwhite

HARK/ConsumptionSaving/ConsMarkovModel.py

@@ -897,6 +897,7 @@ def getShocks(self):
                 if N > 0:
                     IncomeDstnNow    = self.IncomeDstn[t-1][j] # set current income distribution
                     PermGroFacNow    = self.PermGroFac[t-1][j] # and permanent growth factor
+
                     Indices          = np.arange(IncomeDstnNow[0].size) # just a list of integers
                     # Get random draws of income shocks from the discrete distribution
                     EventDraws       = DiscreteDistribution(

HARK/ConsumptionSaving/tests/test_ConsMarkovModel.py

@@ -1,6 +1,7 @@
 import numpy as np
 from HARK.ConsumptionSaving.ConsIndShockModel import init_idiosyncratic_shocks
 from HARK.ConsumptionSaving.ConsMarkovModel import MarkovConsumerType
+from HARK.distribution import DiscreteDistribution
 from copy import copy
 import unittest
 
@@ -17,7 +18,7 @@ def setUp(self):
         p_unemploy_good = p_reemploy * urate_good / (1 - urate_good)
         p_unemploy_bad = p_reemploy * urate_bad / (1 - urate_bad)
         boom_prob = 1.0 / recession_length
-        self.MrkvArray = np.array(
+        MrkvArray = np.array(
             [
                 [
                     (1 - p_unemploy_good) * (1 - bust_prob),
@@ -47,15 +48,32 @@ def setUp(self):
         )
 
         init_serial_unemployment = copy(init_idiosyncratic_shocks)
-        init_serial_unemployment["MrkvArray"] = [self.MrkvArray]       
-        self.model = MarkovConsumerType(**init_serial_unemployment) 
-    
+        init_serial_unemployment["MrkvArray"] = [MrkvArray]       
+        init_serial_unemployment["UnempPrb"] = 0  # to make income distribution when employed
+        init_serial_unemployment["global_markov"] = False
+        self.model = MarkovConsumerType(**init_serial_unemployment)
+        self.model.cycles = 0
+        self.model.vFuncBool = False  # for easy toggling here
+
+        # Replace the default (lognormal) income distribution with a custom one
+        employed_income_dist = DiscreteDistribution(np.ones(1), [np.ones(1), np.ones(1)])  # Definitely get income
+        unemployed_income_dist = DiscreteDistribution(np.ones(1), [np.ones(1), np.zeros(1)]) # Definitely don't
+        self.model.IncomeDstn = [
+            [
+                employed_income_dist,
+                unemployed_income_dist,
+                employed_income_dist,
+                unemployed_income_dist,
+            ]
+        ]
+
     def test_checkMarkovInputs(self):
         # check Rfree
         self.assertRaises(ValueError, self.model.checkMarkovInputs)
         # fix Rfree
         self.model.Rfree = np.array(4 * [self.model.Rfree])
         # check MrkvArray, first mess up the setup
+        self.MrkvArray = self.model.MrkvArray
         self.model.MrkvArray = np.random.rand(3, 3)
         self.assertRaises(ValueError, self.model.checkMarkovInputs)
         # then fix it back
@@ -68,3 +86,7 @@ def test_checkMarkovInputs(self):
         self.assertRaises(ValueError, self.model.checkMarkovInputs)
         # fix PermGroFac
         self.model.PermGroFac = [np.array(4 * self.model.PermGroFac)]
+
+    def test_solve(self):
+        self.test_checkMarkovInputs()
+        self.model.solve()

examples/ConsumptionSaving/example_ConsMarkovModel.ipynb

@@ -12,6 +12,7 @@
     "import numpy as np\n",
     "from HARK.ConsumptionSaving.ConsIndShockModel import init_idiosyncratic_shocks\n",
     "from HARK.ConsumptionSaving.ConsMarkovModel import MarkovConsumerType\n",
+    "from HARK.distribution import DiscreteDistribution\n",
     "mystr = lambda number: \"{:.4f}\".format(number)\n",
     "do_simulation = True"
    ]
@@ -65,7 +66,11 @@
   {
    "cell_type": "code",
    "execution_count": 3,
-   "metadata": {},
+   "metadata": {
+    "jupyter": {
+     "source_hidden": true
+    }
+   },
    "outputs": [],
    "source": [
     "# Make a consumer with serially correlated unemployment, subject to boom and bust cycles\n",
@@ -85,8 +90,8 @@
    "outputs": [],
    "source": [
     "# Replace the default (lognormal) income distribution with a custom one\n",
-    "employed_income_dist = [np.ones(1), np.ones(1), np.ones(1)]  # Definitely get income\n",
-    "unemployed_income_dist = [np.ones(1), np.ones(1), np.zeros(1)]  # Definitely don't\n",
+    "employed_income_dist = DiscreteDistribution(np.ones(1), [np.ones(1), np.ones(1)])  # Definitely get income\n",
+    "unemployed_income_dist = DiscreteDistribution(np.ones(1), [np.ones(1), np.zeros(1)]) # Definitely don't\n",
     "SerialUnemploymentExample.IncomeDstn = [\n",
     "    [\n",
     "        employed_income_dist,\n",
@@ -120,7 +125,7 @@
      "name": "stdout",
      "output_type": "stream",
      "text": [
-      "Solving a Markov consumer with serially correlated unemployment took 0.2911 seconds.\n",
+      "Solving a Markov consumer with serially correlated unemployment took 0.2833 seconds.\n",
       "Consumption functions for each discrete state:\n"
      ]
     },
@@ -190,16 +195,16 @@
    "outputs": [],
    "source": [
     "StateCount = ImmunityT + 1  # Total number of Markov states\n",
-    "IncomeDstnReg = [\n",
+    "IncomeDstnReg = DiscreteDistribution(\n",
     "    np.array([1 - UnempPrb, UnempPrb]),\n",
-    "    np.array([1.0, 1.0]),\n",
-    "    np.array([1.0 / (1.0 - UnempPrb), 0.0]),\n",
-    "]  # Ordinary income distribution\n",
-    "IncomeDstnImm = [\n",
-    "    np.array([1.0]),\n",
+    "    [np.array([1.0, 1.0]),\n",
+    "     np.array([1.0 / (1.0 - UnempPrb), 0.0])]\n",
+    ")  # Ordinary income distribution\n",
+    "IncomeDstnImm = DiscreteDistribution(\n",
     "    np.array([1.0]),\n",
-    "    np.array([1.0]),\n",
-    "]  # Income distribution when unemployed\n",
+    "    [np.array([1.0]),\n",
+    "     np.array([1.0])]\n",
+    ")\n",
     "IncomeDstn = [IncomeDstnReg] + ImmunityT * [\n",
     "    IncomeDstnImm\n",
     "]  # Income distribution for each Markov state, in a list"
@@ -260,7 +265,7 @@
      "name": "stdout",
      "output_type": "stream",
      "text": [
-      "Solving an \"unemployment immunity\" consumer took 0.3454 seconds.\n",
+      "Solving an \"unemployment immunity\" consumer took 0.3062 seconds.\n",
       "Consumption functions for each discrete state:\n"
      ]
     },
@@ -320,11 +325,11 @@
    "metadata": {},
    "outputs": [],
    "source": [
-    "IncomeDstnReg = [\n",
+    "IncomeDstnReg = DiscreteDistribution(\n",
     "    np.array([1 - UnempPrb, UnempPrb]),\n",
-    "    np.array([1.0, 1.0]),\n",
-    "    np.array([1.0, 0.0]),\n",
-    "]\n",
+    "    [np.array([1.0, 1.0]),\n",
+    "     np.array([1.0, 0.0])]\n",
+    ")\n",
     "IncomeDstn = StateCount * [\n",
     "    IncomeDstnReg\n",
     "]  # Same simple income distribution in each state"
@@ -375,7 +380,7 @@
      "name": "stdout",
      "output_type": "stream",
      "text": [
-      "Solving a serially correlated growth consumer took 0.2647 seconds.\n",
+      "Solving a serially correlated growth consumer took 0.2594 seconds.\n",
       "Consumption functions for each discrete state:\n"
      ]
     },
@@ -431,7 +436,7 @@
      "name": "stdout",
      "output_type": "stream",
      "text": [
-      "Solving a serially correlated interest consumer took 0.2409 seconds.\n",
+      "Solving a serially correlated interest consumer took 0.2642 seconds.\n",
       "Consumption functions for each discrete state:\n"
      ]
     },

examples/ConsumptionSaving/example_ConsMarkovModel.py

@@ -5,6 +5,7 @@
 import numpy as np
 from HARK.ConsumptionSaving.ConsIndShockModel import init_idiosyncratic_shocks
 from HARK.ConsumptionSaving.ConsMarkovModel import MarkovConsumerType
+from HARK.distribution import DiscreteDistribution
 mystr = lambda number: "{:.4f}".format(number)
 do_simulation = True
 
@@ -60,8 +61,8 @@
 
 # %%
 # Replace the default (lognormal) income distribution with a custom one
-employed_income_dist = [np.ones(1), np.ones(1), np.ones(1)]  # Definitely get income
-unemployed_income_dist = [np.ones(1), np.ones(1), np.zeros(1)]  # Definitely don't
+employed_income_dist = DiscreteDistribution(np.ones(1), [np.ones(1), np.ones(1)])  # Definitely get income
+unemployed_income_dist = DiscreteDistribution(np.ones(1), [np.ones(1), np.zeros(1)]) # Definitely don't
 SerialUnemploymentExample.IncomeDstn = [
     [
         employed_income_dist,
@@ -114,16 +115,16 @@
 
 # %%
 StateCount = ImmunityT + 1  # Total number of Markov states
-IncomeDstnReg = [
+IncomeDstnReg = DiscreteDistribution(
     np.array([1 - UnempPrb, UnempPrb]),
-    np.array([1.0, 1.0]),
-    np.array([1.0 / (1.0 - UnempPrb), 0.0]),
-]  # Ordinary income distribution
-IncomeDstnImm = [
+    [np.array([1.0, 1.0]),
+     np.array([1.0 / (1.0 - UnempPrb), 0.0])]
+)  # Ordinary income distribution
+IncomeDstnImm = DiscreteDistribution(
     np.array([1.0]),
-    np.array([1.0]),
-    np.array([1.0]),
-]  # Income distribution when unemployed
+    [np.array([1.0]),
+     np.array([1.0])]
+)
 IncomeDstn = [IncomeDstnReg] + ImmunityT * [
     IncomeDstnImm
 ]  # Income distribution for each Markov state, in a list
@@ -185,11 +186,11 @@
 )  # Probability of getting the same permanent income growth rate next period
 
 # %%
-IncomeDstnReg = [
+IncomeDstnReg = DiscreteDistribution(
     np.array([1 - UnempPrb, UnempPrb]),
-    np.array([1.0, 1.0]),
-    np.array([1.0, 0.0]),
-]
+    [np.array([1.0, 1.0]),
+     np.array([1.0, 0.0])]
+)
 IncomeDstn = StateCount * [
     IncomeDstnReg
 ]  # Same simple income distribution in each state