Merge pull request #855 from sbenthall/i838

Add namespace for control variables analogous to for state variables.
econ-ark · Oct 29, 2020 · 866f78e · 866f78e
2 parents 6d1d604 + 94fe400
commit 866f78e
Show file tree

Hide file tree

Showing 13 changed files with 34 additions and 27 deletions.
diff --git a/Documentation/CHANGELOG.md b/Documentation/CHANGELOG.md
@@ -17,6 +17,7 @@ Release Data: TBD
 * Namespace variables for the Market class [#765](https://github.com/econ-ark/HARK/pull/765)
 * We now have a Numba based implementation of PerfForesightConsumerType model available as PerfForesightConsumerTypeFast [#774](https://github.com/econ-ark/HARK/pull/774)
 * Namespace for exogenous shocks [#803](https://github.com/econ-ark/HARK/pull/803)
+* Namespace for controls [#855](https://github.com/econ-ark/HARK/pull/855)
 * State and poststate attributes replaced with state_now and state_prev namespaces [#836](https://github.com/econ-ark/HARK/pull/836)
 
 #### Minor Changes 

diff --git a/HARK/ConsumptionSaving/ConsAggShockModel.py b/HARK/ConsumptionSaving/ConsAggShockModel.py
@@ -382,7 +382,8 @@ def getControls(self):
             MPCnow[these] = self.solution[t].cFunc.derivativeX(
                 self.state_now["mNrmNow"][these], MaggNow[these]
             )  # Marginal propensity to consume
-        self.cNrmNow = cNrmNow
+
+        self.controls["cNrmNow"] = cNrmNow
         self.MPCnow = MPCnow
         return None
 
@@ -625,7 +626,7 @@ def getControls(self):
                     .cFunc[i]
                     .derivativeX(self.state_now["mNrmNow"][those], MaggNow[those])
                 )
-        self.cNrmNow = cNrmNow
+        self.controls["cNrmNow"] = cNrmNow
         self.MPCnow = MPCnow
         return None
 
@@ -1041,13 +1042,13 @@ def getControls(self):
         cNow[employed] = self.solution[0].cFunc[emp_idx](
             self.state_now["mNow"][employed], Mnow[employed]
         )
-        self.cNow = cNow
+        self.controls["cNow"] = cNow
 
     def getPostStates(self):
         """
         Gets each agent's retained assets after consumption and stores MrkvNow as MrkvPrev.
         """
-        self.state_now['aNow'] = self.state_now["mNow"] - self.cNow
+        self.state_now['aNow'] = self.state_now["mNow"] - self.controls["cNow"]
         self.MrkvPrev = self.MrkvNow
 
 

diff --git a/HARK/ConsumptionSaving/ConsGenIncProcessModel.py b/HARK/ConsumptionSaving/ConsGenIncProcessModel.py
@@ -1372,7 +1372,7 @@ def getControls(self):
             MPCnow[these] = self.solution[t].cFunc.derivativeX(
                 self.state_now["mLvlNow"][these], self.state_now["pLvlNow"][these]
             )
-        self.cLvlNow = cLvlNow
+        self.controls["cLvlNow"] = cLvlNow
         self.MPCnow = MPCnow
 
     def getPostStates(self):
@@ -1388,7 +1388,7 @@ def getPostStates(self):
         -------
         None
         """
-        self.state_now['aLvlNow'] = self.state_now["mLvlNow"] - self.cLvlNow
+        self.state_now['aLvlNow'] = self.state_now["mLvlNow"] - self.controls["cLvlNow"]
         # moves now to prev
         AgentType.getPostStates(self)
 

diff --git a/HARK/ConsumptionSaving/ConsIndShockModel.py b/HARK/ConsumptionSaving/ConsIndShockModel.py
@@ -1817,7 +1817,9 @@ def getControls(self):
             cNrmNow[these], MPCnow[these] = self.solution[t].cFunc.eval_with_derivative(
                 self.state_now['mNrmNow'][these]
             )
-        self.cNrmNow = cNrmNow
+        self.controls['cNrmNow'] = cNrmNow
+
+        # MPCnow is not really a control
         self.MPCnow = MPCnow
         return None
 
@@ -1834,7 +1836,7 @@ def getPostStates(self):
         None
         """
         # should this be "Now", or "Prev"?!?
-        self.state_now['aNrmNow'] = self.state_now['mNrmNow'] - self.cNrmNow
+        self.state_now['aNrmNow'] = self.state_now['mNrmNow'] - self.controls['cNrmNow']
         # Useful in some cases to precalculate asset level
         self.state_now['aLvlNow'] = self.state_now['aNrmNow'] * self.state_now['pLvlNow']
 

diff --git a/HARK/ConsumptionSaving/ConsLaborModel.py b/HARK/ConsumptionSaving/ConsLaborModel.py
@@ -516,9 +516,9 @@ def getControls(self):
             LbrNow[these] = self.solution[t].LbrFunc(
                 self.state_now['bNrmNow'][these], self.shocks["TranShkNow"][these]
             )  # Assign labor supply
-        self.cNrmNow = cNrmNow
+        self.controls['cNrmNow'] = cNrmNow
         self.MPCnow = MPCnow
-        self.LbrNow = LbrNow
+        self.controls['LbrNow'] = LbrNow
 
     def getPostStates(self):
         """
@@ -538,9 +538,9 @@ def getPostStates(self):
             these = t == self.t_cycle
             mNrmNow[these] = (
                 self.state_now['bNrmNow'][these]
-                + self.LbrNow[these] * self.shocks["TranShkNow"][these]
+                + self.controls['LbrNow'][these] * self.shocks["TranShkNow"][these]
             )  # mNrm = bNrm + yNrm
-            aNrmNow[these] = mNrmNow[these] - self.cNrmNow[these]  # aNrm = mNrm - cNrm
+            aNrmNow[these] = mNrmNow[these] - self.controls['cNrmNow'][these]  # aNrm = mNrm - cNrm
         self.state_now['mNrmNow'] = mNrmNow
         self.state_now['aNrmNow'] = aNrmNow
 

diff --git a/HARK/ConsumptionSaving/ConsMarkovModel.py b/HARK/ConsumptionSaving/ConsMarkovModel.py
@@ -1185,7 +1185,7 @@ def getControls(self):
                 cNrmNow[these], MPCnow[these] = (
                     self.solution[t].cFunc[j].eval_with_derivative(self.state_now['mNrmNow'][these])
                 )
-        self.cNrmNow = cNrmNow
+        self.controls["cNrmNow"] = cNrmNow
         self.MPCnow = MPCnow
 
     def calcBoundingValues(self):

diff --git a/HARK/ConsumptionSaving/ConsMedModel.py b/HARK/ConsumptionSaving/ConsMedModel.py
@@ -852,7 +852,7 @@ def getShocks(self):
                 MedShkNow[these] = self.MedShkDstn[t].drawDiscrete(N)
                 MedPriceNow[these] = self.MedPrice[t]
         self.shocks["MedShkNow"] = MedShkNow
-        self.MedPriceNow = MedPriceNow
+        self.shocks["MedPriceNow"] = MedPriceNow
 
     def getControls(self):
         """
@@ -876,8 +876,8 @@ def getControls(self):
                 self.state_now['pLvlNow'][these],
                 self.shocks["MedShkNow"][these],
             )
-        self.cLvlNow = cLvlNow
-        self.MedNow = MedNow
+        self.controls['cLvlNow'] = cLvlNow
+        self.controls['MedNow'] = MedNow
         return None
 
     def getPostStates(self):
@@ -892,7 +892,7 @@ def getPostStates(self):
         -------
         None
         """
-        self.state_now['aLvlNow'] = self.state_now['mLvlNow'] - self.cLvlNow - self.MedPriceNow * self.MedNow
+        self.state_now['aLvlNow'] = self.state_now['mLvlNow'] - self.controls['cLvlNow'] - self.shocks["MedPriceNow"] * self.controls['MedNow']
 
         # moves now to prev
         AgentType.getPostStates(self)

diff --git a/HARK/ConsumptionSaving/ConsPortfolioModel.py b/HARK/ConsumptionSaving/ConsPortfolioModel.py
@@ -406,7 +406,7 @@ def getRfree(self):
             return factor.  Will be used by getStates() to calculate mNrmNow, where it
             will be mislabeled as "Rfree".
         """
-        Rport = self.ShareNow * self.shocks['RiskyNow'] + (1.0 - self.ShareNow) * self.Rfree
+        Rport = self.controls["ShareNow"] * self.shocks['RiskyNow'] + (1.0 - self.controls["ShareNow"]) * self.Rfree
         self.RportNow = Rport
         return Rport
 
@@ -425,7 +425,7 @@ def initializeSim(self):
         """
         # these need to be set because "post states",
         # but are a control variable and shock, respectively
-        self.ShareNow = np.zeros(self.AgentCount)
+        self.controls["ShareNow"] = np.zeros(self.AgentCount)
         self.shocks['AdjustNow'] = np.zeros(self.AgentCount, dtype=bool)
         IndShockConsumerType.initializeSim(self)
 
@@ -506,8 +506,8 @@ def getControls(self):
             )
 
         # Store controls as attributes of self
-        self.cNrmNow = cNrmNow
-        self.ShareNow = ShareNow
+        self.controls["cNrmNow"] = cNrmNow
+        self.controls["ShareNow"] = ShareNow
 
 
 # Define a non-object-oriented one period solver

diff --git a/HARK/ConsumptionSaving/ConsPrefShockModel.py b/HARK/ConsumptionSaving/ConsPrefShockModel.py
@@ -207,7 +207,7 @@ def getControls(self):
             cNrmNow[these] = self.solution[t].cFunc(
                 self.state_now['mNrmNow'][these], self.shocks["PrefShkNow"][these]
             )
-        self.cNrmNow = cNrmNow
+        self.controls['cNrmNow'] = cNrmNow
         return None
 
     def calcBoundingValues(self):

diff --git a/HARK/ConsumptionSaving/ConsRepAgentModel.py b/HARK/ConsumptionSaving/ConsRepAgentModel.py
@@ -392,7 +392,7 @@ def getControls(self):
         """
         t = self.t_cycle[0]
         i = self.MrkvNow[0]
-        self.cNrmNow = self.solution[t].cFunc[i](self.mNrmNow)
+        self.controls["cNrmNow"] = self.solution[t].cFunc[i](self.mNrmNow)
 
 
 # Define the default dictionary for a representative agent type

diff --git a/HARK/ConsumptionSaving/TractableBufferStockModel.py b/HARK/ConsumptionSaving/TractableBufferStockModel.py
@@ -681,7 +681,7 @@ def getControls(self):
         cLvlNow = np.zeros(self.AgentCount)
         cLvlNow[employed] = self.solution[0].cFunc(self.state_now['mLvlNow'][employed])
         cLvlNow[unemployed] = self.solution[0].cFunc_U(self.state_now['mLvlNow'][unemployed])
-        self.cLvlNow = cLvlNow
+        self.controls["cLvlNow"] = cLvlNow
 
     def getPostStates(self):
         """
@@ -695,7 +695,7 @@ def getPostStates(self):
         -------
         None
         """
-        self.state_now['aLvlNow'] = self.state_now['mLvlNow'] - self.cLvlNow
+        self.state_now['aLvlNow'] = self.state_now['mLvlNow'] - self.controls["cLvlNow"]
         return None
 
 

diff --git a/HARK/ConsumptionSaving/tests/test_ConsPortfolioModel.py b/HARK/ConsumptionSaving/tests/test_ConsPortfolioModel.py
@@ -51,11 +51,11 @@ def test_simOnePeriod(self):
         self.pcct.simOnePeriod()
 
         self.assertAlmostEqual(
-            self.pcct.ShareNow[0],
+            self.pcct.controls["ShareNow"][0],
             0.8627164488246847
         )
         self.assertAlmostEqual(
-            self.pcct.cNrmNow[0],
+            self.pcct.controls["cNrmNow"][0],
             1.67874799
         )
 

diff --git a/HARK/core.py b/HARK/core.py
@@ -233,6 +233,7 @@ def __init__(
         self.track_vars = []  # NOQA
         self.state_now = {sv : None for sv in self.state_vars}
         self.state_prev = self.state_now.copy()
+        self.controls = {}
         self.shocks = {}
         self.read_shocks = False  # NOQA
         self.shock_history = {}
@@ -783,6 +784,8 @@ def simulate(self, sim_periods=None):
                         ]
                     elif var_name in self.shocks:
                         self.history[var_name][self.t_sim, :] = self.shocks[var_name]
+                    elif var_name in self.controls:
+                        self.history[var_name][self.t_sim, :] = self.controls[var_name]
                     else:
                         self.history[var_name][self.t_sim, :] = getattr(self, var_name)
                 self.t_sim += 1