diff --git a/HARK/ConsumptionSaving/ConsIndShockModel.py b/HARK/ConsumptionSaving/ConsIndShockModel.py index 10cec6c99..5a2183560 100644 --- 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 index 92923bd1d..f2cd6ad5b 100644 --- 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))