IndShockConsumerType now uses dictionary for shocks

econ-ark · Aug 18, 2020 · 2947470 · 2947470
1 parent fe8ebdb
commit 2947470
Show file tree

Hide file tree

Showing 3 changed files with 17 additions and 19 deletions.
diff --git a/HARK/ConsumptionSaving/ConsIndShockModel.py b/HARK/ConsumptionSaving/ConsIndShockModel.py
@@ -1625,8 +1625,8 @@ def getShocks(self):
         None
         '''
         PermGroFac = np.array(self.PermGroFac)
-        self.PermShkNow = PermGroFac[self.t_cycle-1] # cycle time has already been advanced
-        self.TranShkNow = np.ones(self.AgentCount)
+        self.shocks['PermShkNow'] = PermGroFac[self.t_cycle-1] # cycle time has already been advanced
+        self.shocks['TranShkNow'] = np.ones(self.AgentCount)
 
     def getRfree(self):
         '''
@@ -1662,11 +1662,11 @@ def getStates(self):
         RfreeNow = self.getRfree()
 
         # Calculate new states: normalized market resources and permanent income level
-        self.pLvlNow = pLvlPrev*self.PermShkNow # Updated permanent income level
+        self.pLvlNow = pLvlPrev*self.shocks['PermShkNow'] # Updated permanent income level
         self.PlvlAggNow = self.PlvlAggNow*self.PermShkAggNow # Updated aggregate permanent productivity level
-        ReffNow      = RfreeNow/self.PermShkNow # "Effective" interest factor on normalized assets
+        ReffNow      = RfreeNow/self.shocks['PermShkNow'] # "Effective" interest factor on normalized assets
         self.bNrmNow = ReffNow*aNrmPrev         # Bank balances before labor income
-        self.mNrmNow = self.bNrmNow + self.TranShkNow # Market resources after income
+        self.mNrmNow = self.bNrmNow + self.shocks['TranShkNow'] # Market resources after income
         return None
 
     def getControls(self):
@@ -2034,8 +2034,8 @@ def getShocks(self):
         -------
         None
         '''
-        PermShkNow = np.zeros(self.AgentCount) # Initialize shock arrays
-        TranShkNow = np.zeros(self.AgentCount)
+        self.shocks['PermShkNow'] = np.zeros(self.AgentCount) # Initialize shock arrays
+        self.shocks['TranShkNow'] = np.zeros(self.AgentCount)
         newborn = self.t_age == 0
         for t in range(self.T_cycle):
             these = t == self.t_cycle
@@ -2045,8 +2045,8 @@ def getShocks(self):
                 PermGroFacNow    = self.PermGroFac[t-1] # and permanent growth factor
                 # Get random draws of income shocks from the discrete distribution
                 IncShks = IncomeDstnNow.drawDiscrete(N)
-                PermShkNow[these] = IncShks[0,:]*PermGroFacNow # permanent "shock" includes expected growth
-                TranShkNow[these] = IncShks[1,:]
+                self.shocks['PermShkNow'][these] = IncShks[0,:]*PermGroFacNow # permanent "shock" includes expected growth
+                self.shocks['TranShkNow'][these] = IncShks[1,:]
 
         # That procedure used the *last* period in the sequence for newborns, but that's not right
         # Redraw shocks for newborns, using the *first* period in the sequence.  Approximation.
@@ -2058,17 +2058,14 @@ def getShocks(self):
 
             # Get random draws of income shocks from the discrete distribution
             EventDraws       = IncomeDstnNow.draw_events(N)
-            PermShkNow[these] = IncomeDstnNow.X[0][EventDraws]*PermGroFacNow # permanent "shock" includes expected growth
-            TranShkNow[these] = IncomeDstnNow.X[1][EventDraws]
+            self.shocks['PermShkNow'][these] = IncomeDstnNow.X[0][EventDraws]*PermGroFacNow # permanent "shock" includes expected growth
+            self.shocks['TranShkNow'][these] = IncomeDstnNow.X[1][EventDraws]
 #        PermShkNow[newborn] = 1.0
-        TranShkNow[newborn] = 1.0
+        self.shocks['TranShkNow'][newborn] = 1.0
 
         # Store the shocks in self
         self.EmpNow = np.ones(self.AgentCount,dtype=bool)
-        self.EmpNow[TranShkNow == self.IncUnemp] = False
-        self.PermShkNow = PermShkNow
-        self.TranShkNow = TranShkNow
-
+        self.EmpNow[self.shocks['TranShkNow'] == self.IncUnemp] = False
 
     def calcBoundingValues(self):
         '''

diff --git a/HARK/ConsumptionSaving/tests/test_IndShockConsumerType.py b/HARK/ConsumptionSaving/tests/test_IndShockConsumerType.py
@@ -22,11 +22,11 @@ def test_getShocks(self):
 
         self.agent.getShocks()
 
-        self.assertEqual(self.agent.PermShkNow[0],
+        self.assertEqual(self.agent.shocks['PermShkNow'][0],
                          1.0427376294215103)
-        self.assertEqual(self.agent.PermShkNow[1],
+        self.assertEqual(self.agent.shocks['PermShkNow'][1],
                          0.9278094171517413)
-        self.assertEqual(self.agent.TranShkNow[0],
+        self.assertEqual(self.agent.shocks['TranShkNow'][0],
                          0.881761797501595)
 
     def test_ConsIndShockSolverBasic(self):

diff --git a/HARK/core.py b/HARK/core.py
@@ -209,6 +209,7 @@ def __init__(self, solution_terminal=None, cycles=1, pseudo_terminal=True,
         self.tolerance          = tolerance # NOQA
         self.seed               = seed # NOQA
         self.track_vars         = [] # NOQA
+        self.shocks             = {}
         self.poststate_vars     = [] # NOQA
         self.read_shocks        = False # NOQA
         self.history            = {}