diff --git a/HARK/ConsumptionSaving/tests/test_IndShockConsumerType.py b/HARK/ConsumptionSaving/tests/test_IndShockConsumerType.py index f2cd6ad5b..02dbd6a25 100644 --- a/HARK/ConsumptionSaving/tests/test_IndShockConsumerType.py +++ b/HARK/ConsumptionSaving/tests/test_IndShockConsumerType.py @@ -3,6 +3,7 @@ import numpy as np +from HARK.distribution import DiscreteDistributionLabeled from HARK.ConsumptionSaving.ConsIndShockModel import ( ConsIndShockSolverBasic, IndShockConsumerType, @@ -928,6 +929,20 @@ def test_calc_jacobian(self): class test_compare_trans_mats(unittest.TestCase): + + def setUp(self): + + # Grid + self.m_grid = np.linspace(0, 20, 10) + + # Newborn distribution. Will's method assumes everyone starts + # at m=1. + self.newborn_dstn = DiscreteDistributionLabeled( + pmv=np.array([1.0]), + atoms=np.array([[1.0],[1.0]]), + var_names=['PLvl','mNrm'] + ) + def test_compare_will_mateo(self): # No deaths for now in Mateo's method # but Will's requires LivPrb < 1 @@ -942,17 +957,15 @@ def test_compare_will_mateo(self): 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.define_distribution_grid(dist_mGrid=self.m_grid) agent.calc_transition_matrix() tm_will = agent.tran_matrix # Mateo's methods - agent.make_state_grid(mNrmGrid=m_grid) + agent.make_state_grid(mNrmGrid=self.m_grid) agent.full_shock_dstns = agent.IncShkDstn - agent.find_transition_matrices() + agent.find_transition_matrices(newborn_dstn=self.newborn_dstn) tm_mateo = agent.trans_mats[0] # Compare diff --git a/HARK/core.py b/HARK/core.py index 90cf9345b..bfac5bb0f 100644 --- a/HARK/core.py +++ b/HARK/core.py @@ -893,7 +893,7 @@ def make_shock_distributions(self): # Save list of shock distributions self.full_shock_dstns = shock_dstns - def find_transition_matrices(self): + def find_transition_matrices(self, newborn_dstn=None): # Calculate number of periods per cycle, defaults to 1 if all variables are time invariant if len(self.time_vary) > 0: # name = agent.time_vary[0] @@ -916,16 +916,30 @@ def find_transition_matrices(self): these_args = tuple(filter(lambda x: x not in exclude_args, these_args)) # Extract state grid data - grids = [ - self.state_grid.attrs["grids"][x].astype(float) - for x in self.state_grid.attrs["mesh_order"] - ] - # Find values and indices of non-trivial grids - nt_inds, nt_grids = zip(*[[i, x] for i, x in enumerate(grids) if len(x) > 1]) + grids = {} + for x in self.state_grid.attrs["mesh_order"]: + grids[x] = self.state_grid.attrs["grids"][x].astype(float) + + # Find names and values of non-trivial grids + nt_states = [x for x, grid in grids.items() if grid.size > 1] + nt_grids = [grids[x] for x in nt_states] # Number of points in full grid mesh_size = self.state_grid.coords["mesh"].size + # Find newborn distribution + if newborn_dstn is not None: + nb_points = newborn_dstn.dataset[nt_states].to_array().values.T + nb_pmv = newborn_dstn.probability.values + + newborn_mass = mass_to_grid( + points=nb_points, + mass=nb_pmv, + grids=nt_grids, + )[ + np.newaxis, : + ] # Add dimension for broadcasting + # Initialize the list transition matrices trans_mats = [] @@ -950,7 +964,7 @@ def trans_wrapper(shocks_next, solution, state_points): trans_wrapper, solution=self.solution[k], state_points=self.state_grid ) - state_points = state_dstn.dataset.to_array().values[nt_inds, :, :] + state_points = state_dstn.dataset[nt_states].to_array().values pmv = state_dstn.probability.values # Construct transition matrix from the object above @@ -961,6 +975,11 @@ def trans_wrapper(shocks_next, solution, state_points): mass=pmv, grids=nt_grids, ) + + # Add newborns to transition matrix if needed + if newborn_dstn is not None and hasattr(self, "LivPrb"): + tmat = self.LivPrb[k] * tmat + (1.0 - self.LivPrb[k]) * newborn_mass + # Prepend trans_mats.insert(0, tmat)