Repository: econ-ark/HARK

Issue:
refactor MrkvNow into model state variable in ConsAggShockModel
In ConsAggShockModel, the model variable MrkvNow is currently stored as an object attribute.

Is it a state? A shock?

Either way, it should be changed into a namespaced variable for consistency with the new style.

Discussion:

Patch:
diff --git a/Documentation/CHANGELOG.md b/Documentation/CHANGELOG.md
index 3b83b45d..0cee9541 100644
--- a/Documentation/CHANGELOG.md
+++ b/Documentation/CHANGELOG.md
@@ -27,6 +27,8 @@ interpolation for problems with CRRA utility. See [#888](https://github.com/econ
 * Fix the return fields of `dcegm/calcCrossPoints`[#909](https://github.com/econ-ark/HARK/pull/909).
 * Corrects location of constructor documentation to class string for Sphinx rendering [#908](https://github.com/econ-ark/HARK/pull/908)
 * Adds a module for producing life-cycle profiles of income shock variances from [Sabelhaus and Song (2010)](https://www.sciencedirect.com/science/article/abs/pii/S0304393210000358). See [#921](https://github.com/econ-ark/HARK/pull/921).
+* Moves state MrkvNow to shocks['Mrkv'] in AggShockMarkov and KrusellSmith models [#935](https://github.com/econ-ark/HARK/pull/935)
+
 #### Minor Changes
 
 * Move AgentType constructor parameters docs to class docstring so it is rendered by Sphinx.
diff --git a/HARK/ConsumptionSaving/ConsAggShockModel.py b/HARK/ConsumptionSaving/ConsAggShockModel.py
index 545bc548..0865fe45 100644
--- a/HARK/ConsumptionSaving/ConsAggShockModel.py
+++ b/HARK/ConsumptionSaving/ConsAggShockModel.py
@@ -493,6 +493,9 @@ class AggShockMarkovConsumerType(AggShockConsumerType):
         params.update(kwds)
         kwds = params
         AggShockConsumerType.__init__(self, **kwds)
+
+        self.shocks['Mrkv'] = None
+
         self.addToTimeInv("MrkvArray")
         self.solveOnePeriod = solveConsAggMarkov
 
@@ -566,7 +569,7 @@ class AggShockMarkovConsumerType(AggShockConsumerType):
             N = np.sum(these)
             if N > 0:
                 IncShkDstnNow = self.IncShkDstn[t - 1][
-                    self.MrkvNow
+                    self.shocks['Mrkv']
                 ]  # set current income distribution
                 PermGroFacNow = self.PermGroFac[t - 1]  # and permanent growth factor
 
@@ -582,7 +585,7 @@ class AggShockMarkovConsumerType(AggShockConsumerType):
         if N > 0:
             these = newborn
             IncShkDstnNow = self.IncShkDstn[0][
-                self.MrkvNow
+                self.shocks['Mrkv']
             ]  # set current income distribution
             PermGroFacNow = self.PermGroFac[0]  # and permanent growth factor
 
@@ -642,7 +645,7 @@ class AggShockMarkovConsumerType(AggShockConsumerType):
         return None
 
     def getMrkvNow(self):  # This function exists to be overwritten in StickyE model
-        return self.MrkvNow * np.ones(self.AgentCount, dtype=int)
+        return self.shocks['Mrkv'] * np.ones(self.AgentCount, dtype=int)
 
 
 init_KS_agents = {
@@ -699,7 +702,9 @@ class KrusellSmithType(AgentType):
             "EmpNow" : None
         }
 
-        self.shock_vars = {}
+        self.shock_vars = {
+            "Mrkv" : None
+        }
 
         self.solveOnePeriod = solveKrusellSmith
         self.update()
@@ -733,7 +738,7 @@ class KrusellSmithType(AgentType):
         )  # Need to be able to track as many periods as economy runs
         self.kInit = Economy.KSS  # Initialize simulation assets to steady state
         self.MrkvInit = Economy.sow_init[
-            "MrkvNow"
+            "Mrkv"
         ]  # Starting Markov state for the macroeconomy
         self.Mgrid = (
             Economy.MSS * self.MgridBase
@@ -972,8 +977,7 @@ class KrusellSmithType(AgentType):
         self.simulate(1)
 
     def initializeSim(self):
-        self.MrkvNow = self.MrkvInit
-        self.MrkvPrev = self.MrkvInit
+        self.shocks['Mrkv'] = self.MrkvInit
         AgentType.initializeSim(self)
         self.state_now["EmpNow"] = self.state_now["EmpNow"].astype(bool)
         self.makeEmpIdxArrays()
@@ -990,10 +994,10 @@ class KrusellSmithType(AgentType):
         if N == 0:
             return
 
-        if self.MrkvNow == 0:
+        if self.shocks['Mrkv'] == 0:
             unemp_N = int(np.round(self.UrateB * N))
             emp_N = self.AgentCount - unemp_N
-        elif self.MrkvNow == 1:
+        elif self.shocks['Mrkv'] == 1:
             unemp_N = int(np.round(self.UrateG * N))
             emp_N = self.AgentCount - unemp_N
         else:
@@ -1001,6 +1005,7 @@ class KrusellSmithType(AgentType):
         EmpNew = np.concatenate(
             [np.zeros(unemp_N, dtype=bool), np.ones(emp_N, dtype=bool)]
         )
+
         self.state_now["EmpNow"][which] = self.RNG.permutation(EmpNew)
         self.state_now["aNow"][which] = self.kInit
 
@@ -1012,12 +1017,20 @@ class KrusellSmithType(AgentType):
         employed = self.state_prev["EmpNow"].copy().astype(bool)
         unemployed = np.logical_not(employed)
 
+        # derive from past employment rate rather than store previous value
+        mrkv_prev = int((unemployed.sum() / float(self.AgentCount)) != self.UrateB)
+
         # Transition some agents between unemployment and employment
-        emp_permute = self.emp_permute[self.MrkvPrev][self.MrkvNow]
-        unemp_permute = self.unemp_permute[self.MrkvPrev][self.MrkvNow]
+        emp_permute = self.emp_permute[mrkv_prev][self.shocks['Mrkv']]
+        unemp_permute = self.unemp_permute[mrkv_prev][self.shocks['Mrkv']]
         # TODO: replace poststate_vars functionality with shocks here
         EmpNow = self.state_now["EmpNow"]
 
+        if employed.sum() != emp_permute.size:
+            import pdb; pdb.set_trace()
+
+        # It's really this permutation that is the shock...
+        # This apparatus is trying to 'exact match' the 'internal' Markov process.
         EmpNow[employed] = self.RNG.permutation(emp_permute)
         EmpNow[unemployed] = self.RNG.permutation(unemp_permute)
 
@@ -1035,10 +1048,10 @@ class KrusellSmithType(AgentType):
         unemployed = np.logical_not(employed)
 
         # Get the discrete index for (un)employed agents
-        if self.MrkvNow == 0:  # Bad macroeconomic conditions
+        if self.shocks['Mrkv'] == 0:  # Bad macroeconomic conditions
             unemp_idx = 0
             emp_idx = 1
-        elif self.MrkvNow == 1:  # Good macroeconomic conditions
+        elif self.shocks['Mrkv'] == 1:  # Good macroeconomic conditions
             unemp_idx = 2
             emp_idx = 3
         else:
@@ -1057,10 +1070,9 @@ class KrusellSmithType(AgentType):
 
     def getPostStates(self):
         """
-        Gets each agent's retained assets after consumption and stores MrkvNow as MrkvPrev.
+        Gets each agent's retained assets after consumption.
         """
         self.state_now['aNow'] = self.state_now["mNow"] - self.controls["cNow"]
-        self.MrkvPrev = self.MrkvNow
 
 
 ###############################################################################
@@ -2395,7 +2407,7 @@ class CobbDouglasMarkovEconomy(CobbDouglasEconomy):
             "PermShkAggNow",
             "TranShkAggNow",
             "KtoLnow",
-            "MrkvNow",  # This one is new
+            "Mrkv",  # This one is new
         ],
         **kwds
     ):
@@ -2412,7 +2424,7 @@ class CobbDouglasMarkovEconomy(CobbDouglasEconomy):
             **params
         )
 
-        self.sow_init["MrkvNow"] = params["MrkvNow_init"]
+        self.sow_init["Mrkv"] = params["MrkvNow_init"]
 
     def update(self):
         """
@@ -2584,7 +2596,7 @@ class CobbDouglasMarkovEconomy(CobbDouglasEconomy):
         MrkvNow_hist = np.zeros(self.act_T_orig, dtype=int)
         loops = 0
         go = True
-        MrkvNow = self.sow_init["MrkvNow"]
+        MrkvNow = self.sow_init["Mrkv"]
         t = 0
         StateCount = self.MrkvArray.shape[0]
 
@@ -2829,9 +2841,9 @@ class KrusellSmithEconomy(Market):
             self,
             agents=agents,
             tolerance=tolerance,
-            sow_vars=["Mnow", "Aprev", "MrkvNow", "Rnow", "Wnow"],
+            sow_vars=["Mnow", "Aprev", "Mrkv", "Rnow", "Wnow"],
             reap_vars=["aNow", "EmpNow"],
-            track_vars=["MrkvNow", "Aprev", "Mnow", "Urate"],
+            track_vars=["Mrkv", "Aprev", "Mnow", "Urate"],
             dyn_vars=["AFunc"],
             **params
         )
@@ -2870,7 +2882,7 @@ class KrusellSmithEconomy(Market):
         self.sow_init["Wnow"] = self.WSS
         self.PermShkAggNow_init = 1.0
         self.TranShkAggNow_init = 1.0
-        self.sow_init["MrkvNow"] = 0
+        self.sow_init["Mrkv"] = 0
         self.makeMrkvArray()
 
     def reset(self):
@@ -2939,21 +2951,14 @@ class KrusellSmithEconomy(Market):
         MrkvNow_hist.  This variable is binary (0 bad, 1 good) in the KS model.
         """
         # Initialize the Markov history and set up transitions
-        MrkvNow_hist = np.zeros(self.act_T, dtype=int)
+        self.MrkvNow_hist = np.zeros(self.act_T, dtype=int)
         MrkvNow = self.MrkvNow_init
-        t = 0
-
-        # Add histories until each state has been visited at least state_T_min times
-        draws = Uniform(seed=0).draw(N=self.act_T)
 
         markov_process = MarkovProcess(self.MrkvArray, seed= 0)
         for s in range(self.act_T):  # Add act_T_orig more periods
-            MrkvNow_hist[t] = MrkvNow
+            self.MrkvNow_hist[s] = MrkvNow
             MrkvNow = markov_process.draw(MrkvNow)
-            t += 1
 
-        # Store the result as attribute of self
-        self.MrkvNow_hist = MrkvNow_hist
 
     def millRule(self, aNow, EmpNow):
         """
diff --git a/HARK/core.py b/HARK/core.py
index e1be6f20..3fe5a04f 100644
--- a/HARK/core.py
+++ b/HARK/core.py
@@ -533,8 +533,12 @@ class AgentType(Model):
         # state_{t-1}
         for var in self.state_now:
             self.state_prev[var] = self.state_now[var]
-            # note: this is not type checked for aggregate variables.
-            self.state_now[var] = np.empty(self.AgentCount)
+
+            if isinstance(self.state_now[var], np.ndarray):
+                self.state_now[var] = np.empty(self.AgentCount)
+            else:
+                # Probably an aggregate variable. It may be getting set by the Market.
+                pass
 
         if self.read_shocks:  # If shock histories have been pre-specified, use those
             self.readShocks()
@@ -1251,6 +1255,8 @@ class Market(Model):
             for this_type in self.agents:
                 if sow_var in this_type.state_now:
                     this_type.state_now[sow_var] = self.sow_state[sow_var]
+                if sow_var in this_type.shocks:
+                    this_type.shocks[sow_var] = self.sow_state[sow_var]
                 else:
                     setattr(this_type, sow_var, self.sow_state[sow_var])
 


Test patch:
diff --git a/HARK/ConsumptionSaving/tests/test_ConsAggShockModel.py b/HARK/ConsumptionSaving/tests/test_ConsAggShockModel.py
index a24bfa4f..3fe9f76a 100644
--- a/HARK/ConsumptionSaving/tests/test_ConsAggShockModel.py
+++ b/HARK/ConsumptionSaving/tests/test_ConsAggShockModel.py
@@ -263,6 +263,6 @@ class KrusellSmithEconomyTestCase(KrusellSmithTestCase):
 
         self.assertAlmostEqual(self.economy.history["Aprev"][4], 11.009107526443584)
 
-        self.assertAlmostEqual(self.economy.history["MrkvNow"][40], 1)
+        self.assertAlmostEqual(self.economy.history["Mrkv"][40], 1)
 
         self.assertAlmostEqual(self.economy.history["Urate"][12], 0.040000000000000036)