Compare with current methods in ConsIndShock

econ-ark · Jul 20, 2023 · d3238a8 · d3238a8
1 parent 0f1f453
commit d3238a8
Show file tree

Hide file tree

Showing 2 changed files with 88 additions and 0 deletions.
diff --git a/HARK/ConsumptionSaving/ConsIndShockModel.py b/HARK/ConsumptionSaving/ConsIndShockModel.py
@@ -15,6 +15,7 @@
 from copy import copy, deepcopy
 
 import numpy as np
+import xarray as xr
 from scipy import sparse as sp
 from scipy.optimize import newton
 
@@ -3441,6 +3442,58 @@ def construct_lognormal_income_process_unemployment(self):
         )
 
         return IncShkDstn, PermShkDstn, TranShkDstn
+
+    def make_state_grid(
+        self,
+        PLvlGrid=None,
+        mNrmGrid=None,
+    ):
+        if PLvlGrid is None:
+            PLvlGrid = np.array([1.0])
+        if mNrmGrid is None:
+            mNrmGrid = np.array([1.0])
+
+        # Create a mesh
+        points = np.meshgrid(PLvlGrid, mNrmGrid, indexing="ij")
+        points = np.stack([x.flatten() for x in points], axis=0)
+
+        mesh = xr.DataArray(
+            points,
+            dims=["var", "mesh"],
+            coords={"var": ["PLvl", "mNrm"]},
+        )
+
+        self.state_grid = xr.Dataset(
+            data_vars={
+                "PLvl": ("mesh", points[0]),
+                "mNrm": ("mesh", points[1]),
+            },
+            coords={"mesh": np.arange(points.shape[1])},
+            attrs={
+                "grids": {
+                    "PLvl": PLvlGrid,
+                    "mNrm": mNrmGrid,
+                },
+                "mesh_order": ["PLvl", "mNrm"],
+            },
+        )
+
+    def state_to_state_trans(self, shocks_next, solution, state, PermGroFac, Rfree):
+
+        # Consumption
+        cNrm = solution.cFunc(state["mNrm"])
+        # Savings
+        aNrm = state["mNrm"] - cNrm
+
+        # Shock realization
+        PermGroShk = shocks_next["PermShk"] * PermGroFac
+        PLvl_next = state["PLvl"] * PermGroShk
+        mNrm_next = aNrm * Rfree / PermGroShk + shocks_next["TranShk"]
+
+        return {
+            "PLvl": PLvl_next,
+            "mNrm": mNrm_next,
+        }
 
 
 class LognormPermIncShk(DiscreteDistribution):

diff --git a/HARK/ConsumptionSaving/tests/test_IndShockConsumerType.py b/HARK/ConsumptionSaving/tests/test_IndShockConsumerType.py
@@ -922,3 +922,38 @@ def test_calc_jacobian(self):
         self.assertAlmostEqual(CJAC_Perm.T[30][29], -0.06120, places=HARK_PRECISION)
         self.assertAlmostEqual(CJAC_Perm.T[30][30], 0.05307, places=HARK_PRECISION)
         self.assertAlmostEqual(CJAC_Perm.T[30][31], 0.04674, places=HARK_PRECISION)
+
+
+# %% Compare with newer transition matrix methods
+
+
+class test_compare_trans_mats(unittest.TestCase):
+    def test_compare_will_mateo(self):
+        # No deaths for now in Mateo's method
+        # but Will's requires LivPrb < 1
+        params = deepcopy(dict_harmenberg)
+        params["LivPrb"] = [0.999]
+
+        # Create and solve agent
+        agent = IndShockConsumerType(**params)
+        agent.cycles = 0
+        agent.solve()
+        # Activate harmenberg
+        agent.neutral_measure = True
+        agent.update_income_process()
+
+        m_grid = np.linspace(0, 20, 10)
+
+        # Will's methods
+        agent.define_distribution_grid(dist_mGrid=m_grid)
+        agent.calc_transition_matrix()
+        tm_will = agent.tran_matrix
+
+        # Mateo's methods
+        agent.make_state_grid(mNrmGrid=m_grid)
+        agent.full_shock_dstns = agent.IncShkDstn
+        agent.find_transition_matrices()
+        tm_mateo = agent.trans_mats[0]
+
+        # Compare
+        self.assertTrue(np.allclose(tm_will, tm_mateo.T, atol=5e-3))