From 6c87cf51a54e3afd7a7462d64ad83951d00184fa Mon Sep 17 00:00:00 2001 From: Mateo VG Date: Wed, 14 Jun 2023 20:38:09 -0400 Subject: [PATCH 01/28] Shock distribution maker --- HARK/core.py | 45 +++++++++++++++++++++++++++++++++++++++++++++ 1 file changed, 45 insertions(+) diff --git a/HARK/core.py b/HARK/core.py index 3149ebf0d..8f5407ce7 100644 --- a/HARK/core.py +++ b/HARK/core.py @@ -1054,6 +1054,51 @@ def one_period_solver(**kwds): return one_period_solver +# ======================================================================== +# Transition matrix methods +# ======================================================================== + +def make_shock_distributions(agent, dstn_engine): + + # Calculate number of periods per cycle, defaults to 1 if all variables are time invariant + if len(agent.time_vary) > 0: + # name = agent.time_vary[0] + # T = len(eval('agent.' + name)) + T = len(agent.__dict__[agent.time_vary[0]]) + else: + T = 1 + + dstn_dict = {parameter: agent.__dict__[parameter] for parameter in agent.time_inv} + dstn_dict.update({parameter: None for parameter in agent.time_vary}) + + # Initialize the list of shock distributions for this cycle, + # then iterate on periods + shock_dstns = [] + + cycles_range = [0] + list(range(T - 1, 0, -1)) + for k in range(T - 1, -1, -1) if agent.cycles == 1 else cycles_range: + + if hasattr(agent.shock_dstn_engine, "dstn_args"): + these_args = agent.shock_dstn_engine.dstn_args + else: + these_args = get_arg_names(agent.shock_dstn_engine) + + # Update time-varying single period inputs + for name in agent.time_vary: + if name in these_args: + dstn_dict[name] = agent.__dict__[name][k] + + # Make a temporary dictionary for this period + temp_dict = {name: dstn_dict[name] for name in these_args} + + # Construct this period's shock distribution one period + # Add it to the solution, and move to the next period + dstn_t = agent.shock_dstn_engine(**temp_dict) + shock_dstns.insert(0, dstn_t) + + # Return the list of distributions + return shock_dstns + # ======================================================================== # ======================================================================== From 23f46bd463916a821b2c2a953a244131e45c775b Mon Sep 17 00:00:00 2001 From: Mateo VG Date: Wed, 14 Jun 2023 20:39:45 -0400 Subject: [PATCH 02/28] Make args work --- HARK/core.py | 18 ++++++++++-------- 1 file changed, 10 insertions(+), 8 deletions(-) diff --git a/HARK/core.py b/HARK/core.py index 8f5407ce7..36c95d5b5 100644 --- a/HARK/core.py +++ b/HARK/core.py @@ -1058,8 +1058,8 @@ def one_period_solver(**kwds): # Transition matrix methods # ======================================================================== -def make_shock_distributions(agent, dstn_engine): - + +def make_shock_distributions(agent): # Calculate number of periods per cycle, defaults to 1 if all variables are time invariant if len(agent.time_vary) > 0: # name = agent.time_vary[0] @@ -1071,18 +1071,19 @@ def make_shock_distributions(agent, dstn_engine): dstn_dict = {parameter: agent.__dict__[parameter] for parameter in agent.time_inv} dstn_dict.update({parameter: None for parameter in agent.time_vary}) + if hasattr(agent.shock_dstn_engine, "dstn_args"): + these_args = agent.shock_dstn_engine.dstn_args + else: + these_args = get_arg_names(agent.shock_dstn_engine) + + these_args = tuple(filter(lambda x: x != "self", these_args)) + # Initialize the list of shock distributions for this cycle, # then iterate on periods shock_dstns = [] cycles_range = [0] + list(range(T - 1, 0, -1)) for k in range(T - 1, -1, -1) if agent.cycles == 1 else cycles_range: - - if hasattr(agent.shock_dstn_engine, "dstn_args"): - these_args = agent.shock_dstn_engine.dstn_args - else: - these_args = get_arg_names(agent.shock_dstn_engine) - # Update time-varying single period inputs for name in agent.time_vary: if name in these_args: @@ -1099,6 +1100,7 @@ def make_shock_distributions(agent, dstn_engine): # Return the list of distributions return shock_dstns + # ======================================================================== # ======================================================================== From a3a74a63f7c4032a058e71af33698eae2ef5dacb Mon Sep 17 00:00:00 2001 From: Mateo VG Date: Wed, 14 Jun 2023 20:41:03 -0400 Subject: [PATCH 03/28] Create transition matrix method sketch --- HARK/core.py | 139 ++++++++++++++++++++++++++++++++++----------------- 1 file changed, 92 insertions(+), 47 deletions(-) diff --git a/HARK/core.py b/HARK/core.py index 36c95d5b5..25c427bfa 100644 --- a/HARK/core.py +++ b/HARK/core.py @@ -850,6 +850,98 @@ def clear_history(self): self.history[var_name] = np.empty((self.T_sim, self.AgentCount)) self.history[var_name].fill(np.nan) + def make_shock_distributions(self): + # 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] + # T = len(eval('agent.' + name)) + T = len(self.__dict__[self.time_vary[0]]) + else: + T = 1 + + dstn_dict = {parameter: self.__dict__[parameter] for parameter in self.time_inv} + dstn_dict.update({parameter: None for parameter in self.time_vary}) + + if hasattr(self.shock_dstn_engine, "dstn_args"): + these_args = self.shock_dstn_engine.dstn_args + else: + these_args = get_arg_names(self.shock_dstn_engine) + + these_args = tuple(filter(lambda x: x != "self", these_args)) + + # Initialize the list of shock distributions for this cycle, + # then iterate on periods + shock_dstns = [] + + cycles_range = [0] + list(range(T - 1, 0, -1)) + for k in range(T - 1, -1, -1) if self.cycles == 1 else cycles_range: + # Update time-varying single period inputs + for name in self.time_vary: + if name in these_args: + dstn_dict[name] = self.__dict__[name][k] + + # Make a temporary dictionary for this period + temp_dict = {name: dstn_dict[name] for name in these_args} + + # Construct this period's shock distribution one period + # Add it to the solution, and move to the next period + dstn_t = self.shock_dstn_engine(**temp_dict) + shock_dstns.insert(0, dstn_t) + + # Save list of shock distributions + self.full_shock_dstns = shock_dstns + + def find_transition_matrices(self): + # 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] + # T = len(eval('agent.' + name)) + T = len(self.__dict__[self.time_vary[0]]) + else: + T = 1 + + trans_dict = { + parameter: self.__dict__[parameter] for parameter in self.time_inv + } + trans_dict.update({parameter: None for parameter in self.time_vary}) + + if hasattr(self.state_to_state_trans, "trans_args"): + these_args = self.state_to_state_trans.trans_args + else: + these_args = get_arg_names(self.state_to_state_trans) + + exclude_args = ["self", "shocks_next", "state", "solution"] + these_args = tuple(filter(lambda x: x not in exclude_args, these_args)) + + # Initialize the list transition matrices + trans_mats = [] + + cycles_range = [0] + list(range(T - 1, 0, -1)) + for k in range(T - 1, -1, -1) if self.cycles == 1 else cycles_range: + # Update time-varying single period inputs + for name in self.time_vary: + if name in these_args: + trans_dict[name] = self.__dict__[name][k] + + # Make a temporary dictionary for this period + temp_dict = {name: trans_dict[name] for name in these_args} + + shock_dstn = self.full_shock_dstns[k] + + def trans_wrapper(shocks_next, solution, state_points): + return self.state_to_state_trans( + shocks_next, solution, state_points, **temp_dict + ) + + state_dstn = shock_dstn.dist_of_func( + trans_wrapper, self.solution[k], self.state_grid + ) + + # TODO: construct transition matrix from the object above + + # Save matrices + self.trans_mats = trans_mats + def solve_agent(agent, verbose): """ @@ -1054,53 +1146,6 @@ def one_period_solver(**kwds): return one_period_solver -# ======================================================================== -# Transition matrix methods -# ======================================================================== - - -def make_shock_distributions(agent): - # Calculate number of periods per cycle, defaults to 1 if all variables are time invariant - if len(agent.time_vary) > 0: - # name = agent.time_vary[0] - # T = len(eval('agent.' + name)) - T = len(agent.__dict__[agent.time_vary[0]]) - else: - T = 1 - - dstn_dict = {parameter: agent.__dict__[parameter] for parameter in agent.time_inv} - dstn_dict.update({parameter: None for parameter in agent.time_vary}) - - if hasattr(agent.shock_dstn_engine, "dstn_args"): - these_args = agent.shock_dstn_engine.dstn_args - else: - these_args = get_arg_names(agent.shock_dstn_engine) - - these_args = tuple(filter(lambda x: x != "self", these_args)) - - # Initialize the list of shock distributions for this cycle, - # then iterate on periods - shock_dstns = [] - - cycles_range = [0] + list(range(T - 1, 0, -1)) - for k in range(T - 1, -1, -1) if agent.cycles == 1 else cycles_range: - # Update time-varying single period inputs - for name in agent.time_vary: - if name in these_args: - dstn_dict[name] = agent.__dict__[name][k] - - # Make a temporary dictionary for this period - temp_dict = {name: dstn_dict[name] for name in these_args} - - # Construct this period's shock distribution one period - # Add it to the solution, and move to the next period - dstn_t = agent.shock_dstn_engine(**temp_dict) - shock_dstns.insert(0, dstn_t) - - # Return the list of distributions - return shock_dstns - - # ======================================================================== # ======================================================================== From 0430b26d782b5b93915c138b224f2c229f801180 Mon Sep 17 00:00:00 2001 From: Mateo VG Date: Wed, 14 Jun 2023 20:42:33 -0400 Subject: [PATCH 04/28] Construct transition matrix from conditional distribution --- HARK/core.py | 19 +++++++++++++++++-- 1 file changed, 17 insertions(+), 2 deletions(-) diff --git a/HARK/core.py b/HARK/core.py index 25c427bfa..16fcd03aa 100644 --- a/HARK/core.py +++ b/HARK/core.py @@ -26,6 +26,8 @@ from HARK.parallel import multi_thread_commands, multi_thread_commands_fake from HARK.utilities import NullFunc, get_arg_names +from HARK.mat_methods import mass_to_grid + class Model: """ @@ -913,6 +915,12 @@ def find_transition_matrices(self): exclude_args = ["self", "shocks_next", "state", "solution"] these_args = tuple(filter(lambda x: x not in exclude_args, these_args)) + # Extract state grid data + meshpoints = self.state_grid["points"] + grids = [ + self.state_grid["grids"][x].astype(float) for x in self.state_grid["order"] + ] + # Initialize the list transition matrices trans_mats = [] @@ -934,10 +942,17 @@ def trans_wrapper(shocks_next, solution, state_points): ) state_dstn = shock_dstn.dist_of_func( - trans_wrapper, self.solution[k], self.state_grid + trans_wrapper, self.solution[k], meshpoints ) - # TODO: construct transition matrix from the object above + # Construct transition matrix from the object above + tmat = np.zeros((meshpoints.shape[1], meshpoints.shape[1])) + for i in range(meshpoints.shape[1]): + tmat[i, :] = mass_to_grid( + points=state_dstn.atoms[:, i, :].T, mass=state_dstn.pmv, grids=grids + ) + # Prepend + trans_mats.insert(0, tmat) # Save matrices self.trans_mats = trans_mats From a73f6dc6dab438fb5b45b41198cea830d7be9f51 Mon Sep 17 00:00:00 2001 From: Mateo VG Date: Wed, 14 Jun 2023 20:43:46 -0400 Subject: [PATCH 05/28] Support masking states --- HARK/core.py | 6 +++++- 1 file changed, 5 insertions(+), 1 deletion(-) diff --git a/HARK/core.py b/HARK/core.py index 16fcd03aa..e9733d8ba 100644 --- a/HARK/core.py +++ b/HARK/core.py @@ -920,6 +920,8 @@ def find_transition_matrices(self): grids = [ self.state_grid["grids"][x].astype(float) for x in self.state_grid["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]) # Initialize the list transition matrices trans_mats = [] @@ -949,7 +951,9 @@ def trans_wrapper(shocks_next, solution, state_points): tmat = np.zeros((meshpoints.shape[1], meshpoints.shape[1])) for i in range(meshpoints.shape[1]): tmat[i, :] = mass_to_grid( - points=state_dstn.atoms[:, i, :].T, mass=state_dstn.pmv, grids=grids + points=state_dstn.atoms[nt_inds, i, :].T, + mass=state_dstn.pmv, + grids=nt_grids, ) # Prepend trans_mats.insert(0, tmat) From b8e6fdaa1ef7400e5aa7ee6ba23dea4e3d8f1e2b Mon Sep 17 00:00:00 2001 From: Mateo VG Date: Thu, 15 Jun 2023 10:17:41 -0400 Subject: [PATCH 06/28] Shock engine for consportfolio --- HARK/ConsumptionSaving/ConsPortfolioModel.py | 31 ++++++++++++++++++++ 1 file changed, 31 insertions(+) diff --git a/HARK/ConsumptionSaving/ConsPortfolioModel.py b/HARK/ConsumptionSaving/ConsPortfolioModel.py index 731f603de..40e52332e 100644 --- a/HARK/ConsumptionSaving/ConsPortfolioModel.py +++ b/HARK/ConsumptionSaving/ConsPortfolioModel.py @@ -28,6 +28,11 @@ MargValueFuncCRRA, ValueFuncCRRA, ) +from HARK.distribution import ( + combine_indep_dstns, + DiscreteDistribution, + DiscreteDistributionLabeled, +) from HARK.metric import MetricObject @@ -316,6 +321,32 @@ def get_controls(self): self.controls["cNrm"] = cNrmNow self.controls["Share"] = ShareNow + def shock_dstn_engine(self, ShockDstn, AdjustPrb): + """ + Creates a joint labeled distribution of all the shocks in the model + """ + + # ShockDstn is created by RiskyAssetConsumerType and it contains, in order: + # PermShk, TranShk, and Risky + + # Create a distribution for the Adjust shock. + # TODO: self.AdjustDstn already exists, but it is a FrozenDist type of object + # that does not work with combine_indep_dstns. This should be fixed. + if AdjustPrb < 1.0: + AdjustDstn = DiscreteDistribution( + np.array([1.0 - AdjustPrb, AdjustPrb]), [False, True] + ) + else: + AdjustDstn = DiscreteDistribution(np.array([1.0]), [True]) + + LabeledShkDstn = DiscreteDistributionLabeled.from_unlabeled( + dist=combine_indep_dstns(ShockDstn, AdjustDstn), + name="Full shock distribution", + var_names=["PermShk", "TranShk", "Risky", "Adjust"], + ) + + return LabeledShkDstn + class SequentialPortfolioConsumerType(PortfolioConsumerType): def __init__(self, verbose=False, quiet=False, **kwds): From 9310d2518fdd78bf8a4d612e2f1b7503c29576d0 Mon Sep 17 00:00:00 2001 From: Mateo VG Date: Thu, 15 Jun 2023 10:25:41 -0400 Subject: [PATCH 07/28] State grid for consportfolio --- HARK/ConsumptionSaving/ConsPortfolioModel.py | 32 ++++++++++++++++++++ 1 file changed, 32 insertions(+) diff --git a/HARK/ConsumptionSaving/ConsPortfolioModel.py b/HARK/ConsumptionSaving/ConsPortfolioModel.py index 40e52332e..e5eb29eb0 100644 --- a/HARK/ConsumptionSaving/ConsPortfolioModel.py +++ b/HARK/ConsumptionSaving/ConsPortfolioModel.py @@ -346,6 +346,38 @@ def shock_dstn_engine(self, ShockDstn, AdjustPrb): ) return LabeledShkDstn + + def make_state_grid( + self, + PLvlGrid=None, + mNrmGrid=None, + ShareGrid=None, + AdjustGrid=None, + ): + if PLvlGrid is None: + PLvlGrid = np.array([1.0]) + if mNrmGrid is None: + mNrmGrid = np.array([1.0]) + if ShareGrid is None: + ShareGrid = np.array([1.0]) + if AdjustGrid is None: + AdjustGrid = np.array([True]) + + # Create a mesh + points = np.meshgrid(PLvlGrid, mNrmGrid, ShareGrid, AdjustGrid, indexing="ij") + points = np.stack([x.flatten() for x in points], axis=0) + + # Store a dictionary with individual grids, mesh points and order + self.state_grid = { + "grids": { + "PLvl": PLvlGrid, + "mNrm": mNrmGrid, + "Share": ShareGrid, + "Adjust": AdjustGrid, + }, + "points": points, + "order": ["PLvl", "mNrm", "Share", "Adjust"], + } class SequentialPortfolioConsumerType(PortfolioConsumerType): From 88b71cdaaac458c777737d7501497eb42f56e30f Mon Sep 17 00:00:00 2001 From: Mateo VG Date: Thu, 15 Jun 2023 10:57:36 -0400 Subject: [PATCH 08/28] Transition equation for ConsPortfolio --- HARK/ConsumptionSaving/ConsPortfolioModel.py | 44 ++++++++++++++++++++ 1 file changed, 44 insertions(+) diff --git a/HARK/ConsumptionSaving/ConsPortfolioModel.py b/HARK/ConsumptionSaving/ConsPortfolioModel.py index e5eb29eb0..85ebe04a2 100644 --- a/HARK/ConsumptionSaving/ConsPortfolioModel.py +++ b/HARK/ConsumptionSaving/ConsPortfolioModel.py @@ -379,6 +379,50 @@ def make_state_grid( "order": ["PLvl", "mNrm", "Share", "Adjust"], } + def state_to_state_trans(self, shocks_next, solution, state, PermGroFac, Rfree): + # TODO: + # Would be good to have transitions receive + # state as a labeled xarray and return a labeled xarray + # Also define dist_of_func for labeled distributions that returns a labeled distribution + + # Unpack next period's states + PLvl, mNrm, Share, Adjust = state[0], state[1], state[2], state[3] + Adjust = Adjust.astype(bool) + + # Consumption + cNrm = np.empty_like(mNrm) + cNrm[Adjust] = solution.cFuncAdj(mNrm[Adjust]) + cNrm[~Adjust] = solution.cFuncFxd(mNrm[~Adjust], Share[~Adjust]) + # Savings + aNrm = mNrm - cNrm + # Share + Share_next = np.empty_like(Share) + Share_next[Adjust] = solution.ShareFuncAdj(mNrm[Adjust]) + Share_next[~Adjust] = solution.ShareFuncFxd(mNrm[~Adjust], Share[~Adjust]) + + PLvl_next, mNrm_next, Share_next, Adjust_next = post_state_transition( + shocks_next, + PLvl, + aNrm, + Share_next, + PermGroFac, + Rfree, + ) + + return np.stack([PLvl_next, mNrm_next, Share_next, Adjust_next], axis=0) + +def post_state_transition(shocks_next, PLvl, aNrm, Share_next, PermGroFac, Rfree): + + PermGroShk = shocks_next['PermShk'] * PermGroFac + PLvl_next = PLvl * PermGroShk + Rport = Rfree + Share_next * (shocks_next['Risky'] - Rfree) + mNrm_next = aNrm * Rport / PermGroShk + shocks_next['TranShk'] + + # Augment dimensions if needed + Share_next = Share_next * np.ones_like(PLvl_next) + Adjust_next = shocks_next['Adjust'] * np.ones_like(PLvl_next, dtype=bool) + + return PLvl_next, mNrm_next, Share_next, Adjust_next class SequentialPortfolioConsumerType(PortfolioConsumerType): def __init__(self, verbose=False, quiet=False, **kwds): From ff42bdf22b3cc73774ce31a41afa01b3a7d48a5e Mon Sep 17 00:00:00 2001 From: Mateo VG Date: Thu, 15 Jun 2023 10:58:14 -0400 Subject: [PATCH 09/28] Test for frictionless and calvo agent --- .../tests/test_ConsPortfolioModel.py | 43 +++++++++++++++++++ 1 file changed, 43 insertions(+) diff --git a/HARK/ConsumptionSaving/tests/test_ConsPortfolioModel.py b/HARK/ConsumptionSaving/tests/test_ConsPortfolioModel.py index 96b047a60..a2fd49f1d 100644 --- a/HARK/ConsumptionSaving/tests/test_ConsPortfolioModel.py +++ b/HARK/ConsumptionSaving/tests/test_ConsPortfolioModel.py @@ -5,6 +5,7 @@ import HARK.ConsumptionSaving.ConsPortfolioModel as cpm from HARK import make_one_period_oo_solver from HARK.tests import HARK_PRECISION +from copy import copy class PortfolioConsumerTypeTestCase(unittest.TestCase): @@ -304,3 +305,45 @@ def test_draws(self): # Adjust self.assertTrue(np.all(Adjust_draws[t_age == 1] == 0)) self.assertTrue(np.all(Adjust_draws[t_age == 2] == 1)) + + + +class test_transition_mat(unittest.TestCase): + + def setUp(self): + pass + + def test_adjust(self): + + # Create agent + agent = cpm.PortfolioConsumerType(**cpm.init_portfolio) + agent.make_shock_distributions() + agent.make_state_grid( + PLvlGrid=None, + mNrmGrid=np.linspace(0, 30, 50), + ShareGrid=None, + AdjustGrid=None, + ) + agent.solve() + agent.find_transition_matrices() + self.assertTrue(agent.trans_mats[0].size == np.power(50,2)) + + def test_calvo(self): + + # Create agent that has some chance of not being able to + # adjust + params = copy(cpm.init_portfolio) + params["AdjustPrb"] = 0.5 + + agent = cpm.PortfolioConsumerType(**params) + agent.make_shock_distributions() + # Share and adjust become states, so we need grids for them + agent.make_state_grid( + PLvlGrid=None, + mNrmGrid=np.linspace(0, 30, 50), + ShareGrid=np.linspace(0, 1, 10), + AdjustGrid=np.array([False, True]), + ) + agent.solve() + agent.find_transition_matrices() + self.assertTrue(agent.trans_mats[0].size == np.power(50*10*2,2)) \ No newline at end of file From c4c5064bea6665d096886fe049fc5ec3922c27fc Mon Sep 17 00:00:00 2001 From: Mateo VG Date: Thu, 15 Jun 2023 10:59:05 -0400 Subject: [PATCH 10/28] Black --- HARK/ConsumptionSaving/ConsPortfolioModel.py | 15 ++++++++------- .../tests/test_ConsPortfolioModel.py | 10 +++------- 2 files changed, 11 insertions(+), 14 deletions(-) diff --git a/HARK/ConsumptionSaving/ConsPortfolioModel.py b/HARK/ConsumptionSaving/ConsPortfolioModel.py index 85ebe04a2..ec3a3f337 100644 --- a/HARK/ConsumptionSaving/ConsPortfolioModel.py +++ b/HARK/ConsumptionSaving/ConsPortfolioModel.py @@ -346,7 +346,7 @@ def shock_dstn_engine(self, ShockDstn, AdjustPrb): ) return LabeledShkDstn - + def make_state_grid( self, PLvlGrid=None, @@ -399,7 +399,7 @@ def state_to_state_trans(self, shocks_next, solution, state, PermGroFac, Rfree): Share_next = np.empty_like(Share) Share_next[Adjust] = solution.ShareFuncAdj(mNrm[Adjust]) Share_next[~Adjust] = solution.ShareFuncFxd(mNrm[~Adjust], Share[~Adjust]) - + PLvl_next, mNrm_next, Share_next, Adjust_next = post_state_transition( shocks_next, PLvl, @@ -411,19 +411,20 @@ def state_to_state_trans(self, shocks_next, solution, state, PermGroFac, Rfree): return np.stack([PLvl_next, mNrm_next, Share_next, Adjust_next], axis=0) -def post_state_transition(shocks_next, PLvl, aNrm, Share_next, PermGroFac, Rfree): - PermGroShk = shocks_next['PermShk'] * PermGroFac +def post_state_transition(shocks_next, PLvl, aNrm, Share_next, PermGroFac, Rfree): + PermGroShk = shocks_next["PermShk"] * PermGroFac PLvl_next = PLvl * PermGroShk - Rport = Rfree + Share_next * (shocks_next['Risky'] - Rfree) - mNrm_next = aNrm * Rport / PermGroShk + shocks_next['TranShk'] + Rport = Rfree + Share_next * (shocks_next["Risky"] - Rfree) + mNrm_next = aNrm * Rport / PermGroShk + shocks_next["TranShk"] # Augment dimensions if needed Share_next = Share_next * np.ones_like(PLvl_next) - Adjust_next = shocks_next['Adjust'] * np.ones_like(PLvl_next, dtype=bool) + Adjust_next = shocks_next["Adjust"] * np.ones_like(PLvl_next, dtype=bool) return PLvl_next, mNrm_next, Share_next, Adjust_next + class SequentialPortfolioConsumerType(PortfolioConsumerType): def __init__(self, verbose=False, quiet=False, **kwds): params = init_portfolio.copy() diff --git a/HARK/ConsumptionSaving/tests/test_ConsPortfolioModel.py b/HARK/ConsumptionSaving/tests/test_ConsPortfolioModel.py index a2fd49f1d..1cc46ad92 100644 --- a/HARK/ConsumptionSaving/tests/test_ConsPortfolioModel.py +++ b/HARK/ConsumptionSaving/tests/test_ConsPortfolioModel.py @@ -307,14 +307,11 @@ def test_draws(self): self.assertTrue(np.all(Adjust_draws[t_age == 2] == 1)) - class test_transition_mat(unittest.TestCase): - def setUp(self): pass def test_adjust(self): - # Create agent agent = cpm.PortfolioConsumerType(**cpm.init_portfolio) agent.make_shock_distributions() @@ -326,15 +323,14 @@ def test_adjust(self): ) agent.solve() agent.find_transition_matrices() - self.assertTrue(agent.trans_mats[0].size == np.power(50,2)) + self.assertTrue(agent.trans_mats[0].size == np.power(50, 2)) def test_calvo(self): - # Create agent that has some chance of not being able to # adjust params = copy(cpm.init_portfolio) params["AdjustPrb"] = 0.5 - + agent = cpm.PortfolioConsumerType(**params) agent.make_shock_distributions() # Share and adjust become states, so we need grids for them @@ -346,4 +342,4 @@ def test_calvo(self): ) agent.solve() agent.find_transition_matrices() - self.assertTrue(agent.trans_mats[0].size == np.power(50*10*2,2)) \ No newline at end of file + self.assertTrue(agent.trans_mats[0].size == np.power(50 * 10 * 2, 2)) From 11d70597dfe5cc82c96bea643ecc1a3743a9a1be Mon Sep 17 00:00:00 2001 From: MateoVG Date: Wed, 5 Jul 2023 16:53:20 -0400 Subject: [PATCH 11/28] Add a life-cycle example --- .../tests/test_ConsPortfolioModel.py | 29 +++++++++++++++++++ 1 file changed, 29 insertions(+) diff --git a/HARK/ConsumptionSaving/tests/test_ConsPortfolioModel.py b/HARK/ConsumptionSaving/tests/test_ConsPortfolioModel.py index 1cc46ad92..2864ea23b 100644 --- a/HARK/ConsumptionSaving/tests/test_ConsPortfolioModel.py +++ b/HARK/ConsumptionSaving/tests/test_ConsPortfolioModel.py @@ -307,10 +307,39 @@ def test_draws(self): self.assertTrue(np.all(Adjust_draws[t_age == 2] == 1)) +from HARK.ConsumptionSaving.ConsIndShockModel import init_lifecycle +from HARK.ConsumptionSaving.ConsRiskyAssetModel import risky_asset_parms + + class test_transition_mat(unittest.TestCase): def setUp(self): pass + def test_LC(self): + # Create an lc agent + lc_pars = copy(init_lifecycle) + lc_pars.update(risky_asset_parms) + agent = cpm.PortfolioConsumerType(**lc_pars) + agent.solve() + + # Make shock distribution and grid + agent.make_shock_distributions() + agent.make_state_grid( + PLvlGrid=None, + # Low number of points, else RAM reqs are high + mNrmGrid=np.linspace(0, 10, 5), + ShareGrid=None, + AdjustGrid=None, + ) + # Solve + agent.solve() + # Check that it is indeed an LC model + assert len(agent.solution) > 10 + + # Get transition matrices + agent.find_transition_matrices() + assert len(agent.solution) - 1 == len(agent.trans_mats) + def test_adjust(self): # Create agent agent = cpm.PortfolioConsumerType(**cpm.init_portfolio) From eed2d012bf63f040d3decb44864894a90a8e4dbf Mon Sep 17 00:00:00 2001 From: MateoVG Date: Tue, 18 Jul 2023 17:00:03 -0400 Subject: [PATCH 12/28] Start working xrrays in --- HARK/ConsumptionSaving/ConsPortfolioModel.py | 54 ++++++++++++-------- HARK/core.py | 5 +- 2 files changed, 34 insertions(+), 25 deletions(-) diff --git a/HARK/ConsumptionSaving/ConsPortfolioModel.py b/HARK/ConsumptionSaving/ConsPortfolioModel.py index ec3a3f337..1a53bb76f 100644 --- a/HARK/ConsumptionSaving/ConsPortfolioModel.py +++ b/HARK/ConsumptionSaving/ConsPortfolioModel.py @@ -34,6 +34,7 @@ DiscreteDistributionLabeled, ) from HARK.metric import MetricObject +import xarray as xr # Define a class to represent the single period solution of the portfolio choice problem @@ -367,17 +368,30 @@ def make_state_grid( points = np.meshgrid(PLvlGrid, mNrmGrid, ShareGrid, AdjustGrid, indexing="ij") points = np.stack([x.flatten() for x in points], axis=0) - # Store a dictionary with individual grids, mesh points and order - self.state_grid = { - "grids": { - "PLvl": PLvlGrid, - "mNrm": mNrmGrid, - "Share": ShareGrid, - "Adjust": AdjustGrid, + mesh = xr.DataArray( + points, + dims=['var', 'mesh'], + coords = {'var': ["PLvl", "mNrm", "Share", "Adjust"]} + ) + + self.state_grid = xr.Dataset( + data_vars={ + "PLvl": ('mesh', points[0]), + "mNrm": ('mesh', points[1]), + 'Share': ('mesh', points[2]), + 'Adjust': ('mesh', points[3].astype(bool)) }, - "points": points, - "order": ["PLvl", "mNrm", "Share", "Adjust"], - } + coords={'mesh': np.arange(points.shape[1])}, + attrs= { + "grids": { + "PLvl": PLvlGrid, + "mNrm": mNrmGrid, + "Share": ShareGrid, + "Adjust": AdjustGrid, + }, + "mesh_order": ["PLvl", "mNrm", "Share", "Adjust"], + } + ) def state_to_state_trans(self, shocks_next, solution, state, PermGroFac, Rfree): # TODO: @@ -385,24 +399,20 @@ def state_to_state_trans(self, shocks_next, solution, state, PermGroFac, Rfree): # state as a labeled xarray and return a labeled xarray # Also define dist_of_func for labeled distributions that returns a labeled distribution - # Unpack next period's states - PLvl, mNrm, Share, Adjust = state[0], state[1], state[2], state[3] - Adjust = Adjust.astype(bool) - # Consumption - cNrm = np.empty_like(mNrm) - cNrm[Adjust] = solution.cFuncAdj(mNrm[Adjust]) - cNrm[~Adjust] = solution.cFuncFxd(mNrm[~Adjust], Share[~Adjust]) + cNrm = np.empty_like(state['mNrm']) + cNrm[state['Adjust']] = solution.cFuncAdj(state['mNrm'][state['Adjust']]) + cNrm[~state['Adjust']] = solution.cFuncFxd(state['mNrm'][~state['Adjust']], state['Share'][~state['Adjust']]) # Savings - aNrm = mNrm - cNrm + aNrm = state['mNrm'] - cNrm # Share - Share_next = np.empty_like(Share) - Share_next[Adjust] = solution.ShareFuncAdj(mNrm[Adjust]) - Share_next[~Adjust] = solution.ShareFuncFxd(mNrm[~Adjust], Share[~Adjust]) + Share_next = np.empty_like(state['Share']) + Share_next[state['Adjust']] = solution.ShareFuncAdj(state['mNrm'][state['Adjust']]) + Share_next[~state['Adjust']] = solution.ShareFuncFxd(state['mNrm'][~state['Adjust']], state['Share'][~state['Adjust']]) PLvl_next, mNrm_next, Share_next, Adjust_next = post_state_transition( shocks_next, - PLvl, + state['PLvl'], aNrm, Share_next, PermGroFac, diff --git a/HARK/core.py b/HARK/core.py index e9733d8ba..64ad0888c 100644 --- a/HARK/core.py +++ b/HARK/core.py @@ -916,9 +916,8 @@ def find_transition_matrices(self): these_args = tuple(filter(lambda x: x not in exclude_args, these_args)) # Extract state grid data - meshpoints = self.state_grid["points"] grids = [ - self.state_grid["grids"][x].astype(float) for x in self.state_grid["order"] + 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]) @@ -944,7 +943,7 @@ def trans_wrapper(shocks_next, solution, state_points): ) state_dstn = shock_dstn.dist_of_func( - trans_wrapper, self.solution[k], meshpoints + trans_wrapper, self.solution[k], self.state_grid ) # Construct transition matrix from the object above From c7a816f856ea7208d9231f4aaab5af364ebbacd4 Mon Sep 17 00:00:00 2001 From: Mateo VG Date: Tue, 18 Jul 2023 23:41:27 -0400 Subject: [PATCH 13/28] Use xarray --- HARK/ConsumptionSaving/ConsPortfolioModel.py | 57 ++++++++++++-------- HARK/core.py | 19 ++++--- 2 files changed, 47 insertions(+), 29 deletions(-) diff --git a/HARK/ConsumptionSaving/ConsPortfolioModel.py b/HARK/ConsumptionSaving/ConsPortfolioModel.py index 1a53bb76f..ccd8a6099 100644 --- a/HARK/ConsumptionSaving/ConsPortfolioModel.py +++ b/HARK/ConsumptionSaving/ConsPortfolioModel.py @@ -370,19 +370,19 @@ def make_state_grid( mesh = xr.DataArray( points, - dims=['var', 'mesh'], - coords = {'var': ["PLvl", "mNrm", "Share", "Adjust"]} + dims=["var", "mesh"], + coords={"var": ["PLvl", "mNrm", "Share", "Adjust"]}, ) self.state_grid = xr.Dataset( data_vars={ - "PLvl": ('mesh', points[0]), - "mNrm": ('mesh', points[1]), - 'Share': ('mesh', points[2]), - 'Adjust': ('mesh', points[3].astype(bool)) + "PLvl": ("mesh", points[0]), + "mNrm": ("mesh", points[1]), + "Share": ("mesh", points[2]), + "Adjust": ("mesh", points[3].astype(bool)), }, - coords={'mesh': np.arange(points.shape[1])}, - attrs= { + coords={"mesh": np.arange(points.shape[1])}, + attrs={ "grids": { "PLvl": PLvlGrid, "mNrm": mNrmGrid, @@ -390,7 +390,7 @@ def make_state_grid( "Adjust": AdjustGrid, }, "mesh_order": ["PLvl", "mNrm", "Share", "Adjust"], - } + }, ) def state_to_state_trans(self, shocks_next, solution, state, PermGroFac, Rfree): @@ -400,26 +400,32 @@ def state_to_state_trans(self, shocks_next, solution, state, PermGroFac, Rfree): # Also define dist_of_func for labeled distributions that returns a labeled distribution # Consumption - cNrm = np.empty_like(state['mNrm']) - cNrm[state['Adjust']] = solution.cFuncAdj(state['mNrm'][state['Adjust']]) - cNrm[~state['Adjust']] = solution.cFuncFxd(state['mNrm'][~state['Adjust']], state['Share'][~state['Adjust']]) + cNrm = xr.full_like(state["mNrm"], np.nan) + cNrm[state["Adjust"]] = solution.cFuncAdj(state["mNrm"][state["Adjust"]]) + cNrm[~state["Adjust"]] = solution.cFuncFxd( + state["mNrm"][~state["Adjust"]], state["Share"][~state["Adjust"]] + ) # Savings - aNrm = state['mNrm'] - cNrm + aNrm = state["mNrm"] - cNrm # Share - Share_next = np.empty_like(state['Share']) - Share_next[state['Adjust']] = solution.ShareFuncAdj(state['mNrm'][state['Adjust']]) - Share_next[~state['Adjust']] = solution.ShareFuncFxd(state['mNrm'][~state['Adjust']], state['Share'][~state['Adjust']]) + Share_next = xr.full_like(state["Share"], np.nan) + Share_next[state["Adjust"]] = solution.ShareFuncAdj( + state["mNrm"][state["Adjust"]] + ) + Share_next[~state["Adjust"]] = solution.ShareFuncFxd( + state["mNrm"][~state["Adjust"]], state["Share"][~state["Adjust"]] + ) - PLvl_next, mNrm_next, Share_next, Adjust_next = post_state_transition( + state_next = post_state_transition( shocks_next, - state['PLvl'], + state["PLvl"], aNrm, Share_next, PermGroFac, Rfree, ) - return np.stack([PLvl_next, mNrm_next, Share_next, Adjust_next], axis=0) + return state_next def post_state_transition(shocks_next, PLvl, aNrm, Share_next, PermGroFac, Rfree): @@ -429,10 +435,15 @@ def post_state_transition(shocks_next, PLvl, aNrm, Share_next, PermGroFac, Rfree mNrm_next = aNrm * Rport / PermGroShk + shocks_next["TranShk"] # Augment dimensions if needed - Share_next = Share_next * np.ones_like(PLvl_next) - Adjust_next = shocks_next["Adjust"] * np.ones_like(PLvl_next, dtype=bool) - - return PLvl_next, mNrm_next, Share_next, Adjust_next + Share_next = Share_next * xr.ones_like(PLvl_next) + Adjust_next = shocks_next["Adjust"] * xr.ones_like(PLvl_next, dtype=bool) + + return { + "PLvl": PLvl_next, + "mNrm": mNrm_next, + "Share": Share_next, + "Adjust": Adjust_next, + } class SequentialPortfolioConsumerType(PortfolioConsumerType): diff --git a/HARK/core.py b/HARK/core.py index 64ad0888c..90cf9345b 100644 --- a/HARK/core.py +++ b/HARK/core.py @@ -917,11 +917,15 @@ def find_transition_matrices(self): # Extract state grid data grids = [ - self.state_grid.attrs["grids"][x].astype(float) for x in self.state_grid.attrs["mesh_order"] + 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]) + # Number of points in full grid + mesh_size = self.state_grid.coords["mesh"].size + # Initialize the list transition matrices trans_mats = [] @@ -943,15 +947,18 @@ def trans_wrapper(shocks_next, solution, state_points): ) state_dstn = shock_dstn.dist_of_func( - trans_wrapper, self.solution[k], self.state_grid + trans_wrapper, solution=self.solution[k], state_points=self.state_grid ) + state_points = state_dstn.dataset.to_array().values[nt_inds, :, :] + pmv = state_dstn.probability.values + # Construct transition matrix from the object above - tmat = np.zeros((meshpoints.shape[1], meshpoints.shape[1])) - for i in range(meshpoints.shape[1]): + tmat = np.zeros((mesh_size, mesh_size)) + for i in range(mesh_size): tmat[i, :] = mass_to_grid( - points=state_dstn.atoms[nt_inds, i, :].T, - mass=state_dstn.pmv, + points=state_points[:, i, :].T, + mass=pmv, grids=nt_grids, ) # Prepend From d3539d395d93a3bccb1f3d407e1f83f67117dc08 Mon Sep 17 00:00:00 2001 From: Mateo VG Date: Tue, 18 Jul 2023 23:50:57 -0400 Subject: [PATCH 14/28] Create conditional policyfuncs --- HARK/ConsumptionSaving/ConsPortfolioModel.py | 35 ++++++++++---------- 1 file changed, 17 insertions(+), 18 deletions(-) diff --git a/HARK/ConsumptionSaving/ConsPortfolioModel.py b/HARK/ConsumptionSaving/ConsPortfolioModel.py index ccd8a6099..b19802dcd 100644 --- a/HARK/ConsumptionSaving/ConsPortfolioModel.py +++ b/HARK/ConsumptionSaving/ConsPortfolioModel.py @@ -151,6 +151,20 @@ def __init__( self.EndOfPrddvds_fxd = EndOfPrddvds_fxd self.AdjPrb = AdjPrb + def cFunc(self, mNrm, Share, Adjust): + cNrm = xr.full_like(mNrm, np.nan) + cNrm[Adjust] = self.cFuncAdj(mNrm[Adjust]) + no_adj = ~Adjust + cNrm[no_adj] = self.cFuncFxd(mNrm[no_adj], Share[no_adj]) + return cNrm + + def ShareFunc(self, mNrm, Share, Adjust): + ShareNext = xr.full_like(mNrm, np.nan) + ShareNext[Adjust] = self.ShareFuncAdj(mNrm[Adjust]) + no_adj = ~Adjust + ShareNext[no_adj] = self.ShareFuncFxd(mNrm[no_adj], Share[no_adj]) + return ShareNext + class PortfolioConsumerType(RiskyAssetConsumerType): """ @@ -394,28 +408,13 @@ def make_state_grid( ) def state_to_state_trans(self, shocks_next, solution, state, PermGroFac, Rfree): - # TODO: - # Would be good to have transitions receive - # state as a labeled xarray and return a labeled xarray - # Also define dist_of_func for labeled distributions that returns a labeled distribution - # Consumption - cNrm = xr.full_like(state["mNrm"], np.nan) - cNrm[state["Adjust"]] = solution.cFuncAdj(state["mNrm"][state["Adjust"]]) - cNrm[~state["Adjust"]] = solution.cFuncFxd( - state["mNrm"][~state["Adjust"]], state["Share"][~state["Adjust"]] - ) + cNrm = solution.cFunc(state["mNrm"], state["Share"], state["Adjust"]) # Savings aNrm = state["mNrm"] - cNrm # Share - Share_next = xr.full_like(state["Share"], np.nan) - Share_next[state["Adjust"]] = solution.ShareFuncAdj( - state["mNrm"][state["Adjust"]] - ) - Share_next[~state["Adjust"]] = solution.ShareFuncFxd( - state["mNrm"][~state["Adjust"]], state["Share"][~state["Adjust"]] - ) - + Share_next = solution.ShareFunc(state["mNrm"], state["Share"], state["Adjust"]) + # Shock transition state_next = post_state_transition( shocks_next, state["PLvl"], From d3238a8dbb86f330d21c59f46b78b1a8c2901d38 Mon Sep 17 00:00:00 2001 From: MateoVG Date: Thu, 20 Jul 2023 09:31:16 -0400 Subject: [PATCH 15/28] Compare with current methods in ConsIndShock --- HARK/ConsumptionSaving/ConsIndShockModel.py | 53 +++++++++++++++++++ .../tests/test_IndShockConsumerType.py | 35 ++++++++++++ 2 files changed, 88 insertions(+) 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)) From 6b502aae9701a327e71afef5391982d9f4e03d2c Mon Sep 17 00:00:00 2001 From: Mateo VG Date: Thu, 20 Jul 2023 21:25:21 -0400 Subject: [PATCH 16/28] Implement deaths (but results don't match) Check that the newborn distn is the same --- .../tests/test_IndShockConsumerType.py | 23 +++++++++--- HARK/core.py | 35 ++++++++++++++----- 2 files changed, 45 insertions(+), 13 deletions(-) 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) From f155fe69018edc408b567a429ebc550a1c7605f5 Mon Sep 17 00:00:00 2001 From: MateoVG Date: Fri, 21 Jul 2023 10:02:17 -0400 Subject: [PATCH 17/28] Fix treatment of newborns in infinite horizon models --- .../ConsumptionSaving/tests/test_IndShockConsumerType.py | 9 +++++---- HARK/core.py | 8 ++++++-- 2 files changed, 11 insertions(+), 6 deletions(-) diff --git a/HARK/ConsumptionSaving/tests/test_IndShockConsumerType.py b/HARK/ConsumptionSaving/tests/test_IndShockConsumerType.py index 02dbd6a25..d8105c159 100644 --- a/HARK/ConsumptionSaving/tests/test_IndShockConsumerType.py +++ b/HARK/ConsumptionSaving/tests/test_IndShockConsumerType.py @@ -944,10 +944,10 @@ def setUp(self): ) def test_compare_will_mateo(self): - # No deaths for now in Mateo's method - # but Will's requires LivPrb < 1 + + # Deaths params = deepcopy(dict_harmenberg) - params["LivPrb"] = [0.999] + params["LivPrb"] = [0.9] # Create and solve agent agent = IndShockConsumerType(**params) @@ -969,4 +969,5 @@ def test_compare_will_mateo(self): tm_mateo = agent.trans_mats[0] # Compare - self.assertTrue(np.allclose(tm_will, tm_mateo.T, atol=5e-3)) + self.assertTrue(np.allclose(tm_will, tm_mateo.T, atol=1e-10)) + diff --git a/HARK/core.py b/HARK/core.py index bfac5bb0f..8900b720d 100644 --- a/HARK/core.py +++ b/HARK/core.py @@ -922,7 +922,7 @@ def find_transition_matrices(self, newborn_dstn=None): # 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] + nt_grids = tuple(grids[x] for x in nt_states) # Number of points in full grid mesh_size = self.state_grid.coords["mesh"].size @@ -978,7 +978,11 @@ def trans_wrapper(shocks_next, solution, state_points): # 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 + if T==1: + tmat = self.LivPrb[k] * tmat + (1.0 - self.LivPrb[k]) * newborn_mass + else: + # TODO: Life-cycle treatment of death to be defined + pass # Prepend trans_mats.insert(0, tmat) From 516f4a01317ddbd604386324648d17c9674c5ce0 Mon Sep 17 00:00:00 2001 From: MateoVG Date: Tue, 25 Jul 2023 15:43:23 -0400 Subject: [PATCH 18/28] Add transition matrix class --- .../tests/test_ConsPortfolioModel.py | 8 ++-- .../tests/test_IndShockConsumerType.py | 10 ++--- HARK/core.py | 29 ++++++------- HARK/mat_methods.py | 43 +++++++++++++++++++ 4 files changed, 64 insertions(+), 26 deletions(-) diff --git a/HARK/ConsumptionSaving/tests/test_ConsPortfolioModel.py b/HARK/ConsumptionSaving/tests/test_ConsPortfolioModel.py index 2864ea23b..b04845567 100644 --- a/HARK/ConsumptionSaving/tests/test_ConsPortfolioModel.py +++ b/HARK/ConsumptionSaving/tests/test_ConsPortfolioModel.py @@ -338,7 +338,7 @@ def test_LC(self): # Get transition matrices agent.find_transition_matrices() - assert len(agent.solution) - 1 == len(agent.trans_mats) + assert len(agent.solution) - 1 == len(agent.trans_mat.living_transitions) def test_adjust(self): # Create agent @@ -352,7 +352,7 @@ def test_adjust(self): ) agent.solve() agent.find_transition_matrices() - self.assertTrue(agent.trans_mats[0].size == np.power(50, 2)) + self.assertTrue(agent.trans_mat.living_transitions[0].size == np.power(50, 2)) def test_calvo(self): # Create agent that has some chance of not being able to @@ -371,4 +371,6 @@ def test_calvo(self): ) agent.solve() agent.find_transition_matrices() - self.assertTrue(agent.trans_mats[0].size == np.power(50 * 10 * 2, 2)) + self.assertTrue( + agent.trans_mat.living_transitions[0].size == np.power(50 * 10 * 2, 2) + ) diff --git a/HARK/ConsumptionSaving/tests/test_IndShockConsumerType.py b/HARK/ConsumptionSaving/tests/test_IndShockConsumerType.py index d8105c159..bc1c75029 100644 --- a/HARK/ConsumptionSaving/tests/test_IndShockConsumerType.py +++ b/HARK/ConsumptionSaving/tests/test_IndShockConsumerType.py @@ -929,9 +929,7 @@ def test_calc_jacobian(self): class test_compare_trans_mats(unittest.TestCase): - def setUp(self): - # Grid self.m_grid = np.linspace(0, 20, 10) @@ -939,12 +937,11 @@ def setUp(self): # at m=1. self.newborn_dstn = DiscreteDistributionLabeled( pmv=np.array([1.0]), - atoms=np.array([[1.0],[1.0]]), - var_names=['PLvl','mNrm'] + atoms=np.array([[1.0], [1.0]]), + var_names=["PLvl", "mNrm"], ) def test_compare_will_mateo(self): - # Deaths params = deepcopy(dict_harmenberg) params["LivPrb"] = [0.9] @@ -966,8 +963,7 @@ def test_compare_will_mateo(self): agent.make_state_grid(mNrmGrid=self.m_grid) agent.full_shock_dstns = agent.IncShkDstn agent.find_transition_matrices(newborn_dstn=self.newborn_dstn) - tm_mateo = agent.trans_mats[0] + tm_mateo = agent.trans_mat.get_full_tmat() # Compare self.assertTrue(np.allclose(tm_will, tm_mateo.T, atol=1e-10)) - diff --git a/HARK/core.py b/HARK/core.py index 8900b720d..cbf904763 100644 --- a/HARK/core.py +++ b/HARK/core.py @@ -26,7 +26,7 @@ from HARK.parallel import multi_thread_commands, multi_thread_commands_fake from HARK.utilities import NullFunc, get_arg_names -from HARK.mat_methods import mass_to_grid +from HARK.mat_methods import mass_to_grid, transition_mat class Model: @@ -936,12 +936,12 @@ def find_transition_matrices(self, newborn_dstn=None): points=nb_points, mass=nb_pmv, grids=nt_grids, - )[ - np.newaxis, : - ] # Add dimension for broadcasting + ) + else: + newborn_mass = None - # Initialize the list transition matrices - trans_mats = [] + # Initialize the list of matrices conditional on survival + surv_mats = [] cycles_range = [0] + list(range(T - 1, 0, -1)) for k in range(T - 1, -1, -1) if self.cycles == 1 else cycles_range: @@ -976,19 +976,16 @@ def trans_wrapper(shocks_next, solution, state_points): grids=nt_grids, ) - # Add newborns to transition matrix if needed - if newborn_dstn is not None and hasattr(self, "LivPrb"): - if T==1: - tmat = self.LivPrb[k] * tmat + (1.0 - self.LivPrb[k]) * newborn_mass - else: - # TODO: Life-cycle treatment of death to be defined - pass - # Prepend - trans_mats.insert(0, tmat) + surv_mats.insert(0, tmat) # Save matrices - self.trans_mats = trans_mats + self.trans_mat = transition_mat( + living_transitions=surv_mats, + surv_probs=self.LivPrb, + newborn_dstn=newborn_mass, + life_cycle= T > 1, + ) def solve_agent(agent, verbose): diff --git a/HARK/mat_methods.py b/HARK/mat_methods.py index abe2cb4fd..da65a0b58 100644 --- a/HARK/mat_methods.py +++ b/HARK/mat_methods.py @@ -254,3 +254,46 @@ def mass_to_grid( distr = sum_weights(weights, dims, add_inds) return distr + + +class transition_mat: + def __init__( + self, + living_transitions: list, + surv_probs: list, + newborn_dstn: np.ndarray, + life_cycle: bool, + ) -> None: + self.living_transitions = living_transitions + self.surv_probs = surv_probs + self.newborn_dstn = newborn_dstn + self.life_cycle = life_cycle + + if self.life_cycle: + assert len(self.living_transitions) == len( + self.surv_probs + ), "living_transitions must be a list of length len(surv_probs) + 1 if life_cycle is True" + else: + assert ( + len(self.living_transitions) == 1 + ), "living_transitions must be a list of length 1 if life_cycle is False" + assert ( + len(self.surv_probs) == 1 + ), "surv_probs must be a list of length 1 if life_cycle is False" + + self.T = len(self.living_transitions) + 1 + + self.grid_len = self.living_transitions[0].shape[0] + + def get_full_tmat(self): + if self.life_cycle: + # Life cycle + pass + else: + # Infinite horizon + full_mat = ( + self.surv_probs[0] * self.living_transitions[0] + + (1 - self.surv_probs[0]) * self.newborn_dstn[np.newaxis, :] + ) + + return full_mat From cb583c8f657f11afa5345b36e0ce0ed4a40c3270 Mon Sep 17 00:00:00 2001 From: MateoVG Date: Tue, 25 Jul 2023 16:50:17 -0400 Subject: [PATCH 19/28] Add a bruteforce LC transition matrix --- .../tests/test_ConsPortfolioModel.py | 27 +++++++++++++++---- HARK/mat_methods.py | 22 ++++++++++++++- 2 files changed, 43 insertions(+), 6 deletions(-) diff --git a/HARK/ConsumptionSaving/tests/test_ConsPortfolioModel.py b/HARK/ConsumptionSaving/tests/test_ConsPortfolioModel.py index b04845567..5660da805 100644 --- a/HARK/ConsumptionSaving/tests/test_ConsPortfolioModel.py +++ b/HARK/ConsumptionSaving/tests/test_ConsPortfolioModel.py @@ -6,6 +6,7 @@ from HARK import make_one_period_oo_solver from HARK.tests import HARK_PRECISION from copy import copy +from HARK.distribution import DiscreteDistributionLabeled class PortfolioConsumerTypeTestCase(unittest.TestCase): @@ -316,6 +317,9 @@ def setUp(self): pass def test_LC(self): + # Low number of points, else RAM reqs are high + npoints = 3 + # Create an lc agent lc_pars = copy(init_lifecycle) lc_pars.update(risky_asset_parms) @@ -326,8 +330,7 @@ def test_LC(self): agent.make_shock_distributions() agent.make_state_grid( PLvlGrid=None, - # Low number of points, else RAM reqs are high - mNrmGrid=np.linspace(0, 10, 5), + mNrmGrid=np.linspace(0, 10, npoints), ShareGrid=None, AdjustGrid=None, ) @@ -336,23 +339,37 @@ def test_LC(self): # Check that it is indeed an LC model assert len(agent.solution) > 10 + # Define some default newborn distribution + newborn_dstn = DiscreteDistributionLabeled( + pmv=np.array([1.0]), + atoms=np.array([[1.0], [1.0]]), + var_names=["PLvl", "mNrm"], + ) + # Get transition matrices - agent.find_transition_matrices() + agent.find_transition_matrices(newborn_dstn=newborn_dstn) assert len(agent.solution) - 1 == len(agent.trans_mat.living_transitions) + full_mat = agent.trans_mat.get_full_tmat() + # Rows of valid transition matrix sum to 1.0 + self.assertTrue(np.allclose(np.sum(full_mat, 1), 1.0)) + def test_adjust(self): # Create agent + npoints = 5 agent = cpm.PortfolioConsumerType(**cpm.init_portfolio) agent.make_shock_distributions() agent.make_state_grid( PLvlGrid=None, - mNrmGrid=np.linspace(0, 30, 50), + mNrmGrid=np.linspace(0, 10, npoints), ShareGrid=None, AdjustGrid=None, ) agent.solve() agent.find_transition_matrices() - self.assertTrue(agent.trans_mat.living_transitions[0].size == np.power(50, 2)) + self.assertTrue( + agent.trans_mat.living_transitions[0].size == np.power(npoints, 2) + ) def test_calvo(self): # Create agent that has some chance of not being able to diff --git a/HARK/mat_methods.py b/HARK/mat_methods.py index da65a0b58..b23cdfdce 100644 --- a/HARK/mat_methods.py +++ b/HARK/mat_methods.py @@ -288,7 +288,27 @@ def __init__( def get_full_tmat(self): if self.life_cycle: # Life cycle - pass + dim = self.T * self.grid_len + full_mat = np.zeros((dim, dim)) + for k in range(self.T - 1): + row_init = k * self.grid_len + row_end = row_init + self.grid_len + # Living-to-newborn + full_mat[row_init:row_end, : self.grid_len] += ( + 1 - self.surv_probs[k] + ) * self.newborn_dstn[np.newaxis, :] + # Living-to-age+1 + col_init = row_init + self.grid_len + col_end = col_init + self.grid_len + full_mat[row_init:row_end, col_init:col_end] += ( + self.surv_probs[k] * self.living_transitions[k] + ) + + # In at the end of the last age, everyone turns into a newborn + full_mat[ + (self.T - 1) * self.grid_len :, : self.grid_len + ] += self.newborn_dstn[np.newaxis, :] + else: # Infinite horizon full_mat = ( From 4678d6002813ba72cccfe5bc18ecda3c161f17bb Mon Sep 17 00:00:00 2001 From: Mateo VG Date: Wed, 26 Jul 2023 12:31:53 -0400 Subject: [PATCH 20/28] Proof of concept of simulating LC distributions --- .../tests/test_ConsPortfolioModel.py | 19 ++++++++++-- HARK/mat_methods.py | 30 +++++++++++++++++++ 2 files changed, 47 insertions(+), 2 deletions(-) diff --git a/HARK/ConsumptionSaving/tests/test_ConsPortfolioModel.py b/HARK/ConsumptionSaving/tests/test_ConsPortfolioModel.py index 5660da805..025e005af 100644 --- a/HARK/ConsumptionSaving/tests/test_ConsPortfolioModel.py +++ b/HARK/ConsumptionSaving/tests/test_ConsPortfolioModel.py @@ -318,11 +318,12 @@ def setUp(self): def test_LC(self): # Low number of points, else RAM reqs are high - npoints = 3 + npoints = 50 # Create an lc agent lc_pars = copy(init_lifecycle) lc_pars.update(risky_asset_parms) + lc_pars["DiscFac"] = 0.9 agent = cpm.PortfolioConsumerType(**lc_pars) agent.solve() @@ -330,7 +331,7 @@ def test_LC(self): agent.make_shock_distributions() agent.make_state_grid( PLvlGrid=None, - mNrmGrid=np.linspace(0, 10, npoints), + mNrmGrid=np.linspace(0, 20, npoints), ShareGrid=None, AdjustGrid=None, ) @@ -350,10 +351,24 @@ def test_LC(self): agent.find_transition_matrices(newborn_dstn=newborn_dstn) assert len(agent.solution) - 1 == len(agent.trans_mat.living_transitions) + # Check the bruteforce representation that treats age as a state. full_mat = agent.trans_mat.get_full_tmat() # Rows of valid transition matrix sum to 1.0 self.assertTrue(np.allclose(np.sum(full_mat, 1), 1.0)) + # Check iterating distributions forward + + # Set an initial distribution where everyone starts at the youngest age, + # in the first gridpoint. + dstn = np.zeros((npoints, len(agent.trans_mat.living_transitions))) + dstn[0, 0] = 1.0 + + for _ in range(1000): + dstn = agent.trans_mat.iterate_dstn_forward(dstn) + + # Rows of valid transition matrix sum to 1.0 + self.assertTrue(np.allclose(np.sum(full_mat, 1), 1.0)) + def test_adjust(self): # Create agent npoints = 5 diff --git a/HARK/mat_methods.py b/HARK/mat_methods.py index b23cdfdce..559fefcc2 100644 --- a/HARK/mat_methods.py +++ b/HARK/mat_methods.py @@ -285,6 +285,36 @@ def __init__( self.grid_len = self.living_transitions[0].shape[0] + def iterate_dstn_forward(self, dstn_init: np.ndarray) -> np.ndarray: + # Initialize final distribution + dstn_final = np.zeros_like(dstn_init) + + if self.life_cycle: + for k in range(self.T - 2): + # Living-to-age+1 + dstn_final[:, k + 1] += self.surv_probs[k] * np.dot( + dstn_init[:, k], self.living_transitions[k] + ) + # Living-to-newborn + dstn_final[:, 0] += ( + (1 - self.surv_probs[k]) + * np.sum(dstn_init[:, k]) + * self.newborn_dstn + ) + + # In at the end of the last age, everyone turns into a newborn + dstn_final[:, 0] += np.sum(dstn_init[:, -1]) * self.newborn_dstn + + else: + # Living-to-age+1 + dstn_final += self.surv_probs[0] * np.dot( + self.living_transitions[0], dstn_init + ) + # Living-to-newborn + dstn_final += (1 - self.surv_probs[0]) * self.newborn_dstn + + return dstn_final + def get_full_tmat(self): if self.life_cycle: # Life cycle From d4e37de96911e16cf03f5ace5f9f4c9fe62a9a87 Mon Sep 17 00:00:00 2001 From: Mateo VG Date: Wed, 26 Jul 2023 13:12:33 -0400 Subject: [PATCH 21/28] Fix and test infinite horizon simulator --- .../tests/test_ConsPortfolioModel.py | 27 +++++++++++-------- HARK/mat_methods.py | 4 +-- 2 files changed, 18 insertions(+), 13 deletions(-) diff --git a/HARK/ConsumptionSaving/tests/test_ConsPortfolioModel.py b/HARK/ConsumptionSaving/tests/test_ConsPortfolioModel.py index 025e005af..1222f70e2 100644 --- a/HARK/ConsumptionSaving/tests/test_ConsPortfolioModel.py +++ b/HARK/ConsumptionSaving/tests/test_ConsPortfolioModel.py @@ -314,7 +314,12 @@ def test_draws(self): class test_transition_mat(unittest.TestCase): def setUp(self): - pass + # Define some default newborn distribution over all states + self.newborn_dstn = DiscreteDistributionLabeled( + pmv=np.array([1.0]), + atoms=np.array([[1.0], [1.0], [0.5], [1.0]]), + var_names=["PLvl", "mNrm", "Share", "Adjust"], + ) def test_LC(self): # Low number of points, else RAM reqs are high @@ -340,15 +345,8 @@ def test_LC(self): # Check that it is indeed an LC model assert len(agent.solution) > 10 - # Define some default newborn distribution - newborn_dstn = DiscreteDistributionLabeled( - pmv=np.array([1.0]), - atoms=np.array([[1.0], [1.0]]), - var_names=["PLvl", "mNrm"], - ) - # Get transition matrices - agent.find_transition_matrices(newborn_dstn=newborn_dstn) + agent.find_transition_matrices(newborn_dstn=self.newborn_dstn) assert len(agent.solution) - 1 == len(agent.trans_mat.living_transitions) # Check the bruteforce representation that treats age as a state. @@ -381,7 +379,7 @@ def test_adjust(self): AdjustGrid=None, ) agent.solve() - agent.find_transition_matrices() + agent.find_transition_matrices(newborn_dstn=self) self.assertTrue( agent.trans_mat.living_transitions[0].size == np.power(npoints, 2) ) @@ -402,7 +400,14 @@ def test_calvo(self): AdjustGrid=np.array([False, True]), ) agent.solve() - agent.find_transition_matrices() + agent.find_transition_matrices(newborn_dstn=self.newborn_dstn) self.assertTrue( agent.trans_mat.living_transitions[0].size == np.power(50 * 10 * 2, 2) ) + + # Check that we can simulate it + dstn = np.zeros((len(agent.state_grid.coords["mesh"]), 1)) + dstn[0, 0] = 1.0 + + for _ in range(1000): + dstn = agent.trans_mat.iterate_dstn_forward(dstn) diff --git a/HARK/mat_methods.py b/HARK/mat_methods.py index 559fefcc2..dda41c8c3 100644 --- a/HARK/mat_methods.py +++ b/HARK/mat_methods.py @@ -308,10 +308,10 @@ def iterate_dstn_forward(self, dstn_init: np.ndarray) -> np.ndarray: else: # Living-to-age+1 dstn_final += self.surv_probs[0] * np.dot( - self.living_transitions[0], dstn_init + self.living_transitions[0].T, dstn_init ) # Living-to-newborn - dstn_final += (1 - self.surv_probs[0]) * self.newborn_dstn + dstn_final[:, 0] += (1 - self.surv_probs[0]) * self.newborn_dstn return dstn_final From fea5c2ccc8acc1f12a5c7ab944358ecf617bc22f Mon Sep 17 00:00:00 2001 From: Mateo VG Date: Wed, 26 Jul 2023 17:50:54 -0400 Subject: [PATCH 22/28] Add tool to find the steady state --- .../tests/test_ConsPortfolioModel.py | 8 +- HARK/mat_methods.py | 78 ++++++++++++------- 2 files changed, 50 insertions(+), 36 deletions(-) diff --git a/HARK/ConsumptionSaving/tests/test_ConsPortfolioModel.py b/HARK/ConsumptionSaving/tests/test_ConsPortfolioModel.py index 1222f70e2..bd475297f 100644 --- a/HARK/ConsumptionSaving/tests/test_ConsPortfolioModel.py +++ b/HARK/ConsumptionSaving/tests/test_ConsPortfolioModel.py @@ -360,12 +360,8 @@ def test_LC(self): # in the first gridpoint. dstn = np.zeros((npoints, len(agent.trans_mat.living_transitions))) dstn[0, 0] = 1.0 - - for _ in range(1000): - dstn = agent.trans_mat.iterate_dstn_forward(dstn) - - # Rows of valid transition matrix sum to 1.0 - self.assertTrue(np.allclose(np.sum(full_mat, 1), 1.0)) + # Find steady_state + ss_dstn = agent.trans_mat.find_steady_state_dstn(dstn_init=dstn, max_iter=1e4) def test_adjust(self): # Create agent diff --git a/HARK/mat_methods.py b/HARK/mat_methods.py index dda41c8c3..da9ce4a9f 100644 --- a/HARK/mat_methods.py +++ b/HARK/mat_methods.py @@ -285,36 +285,6 @@ def __init__( self.grid_len = self.living_transitions[0].shape[0] - def iterate_dstn_forward(self, dstn_init: np.ndarray) -> np.ndarray: - # Initialize final distribution - dstn_final = np.zeros_like(dstn_init) - - if self.life_cycle: - for k in range(self.T - 2): - # Living-to-age+1 - dstn_final[:, k + 1] += self.surv_probs[k] * np.dot( - dstn_init[:, k], self.living_transitions[k] - ) - # Living-to-newborn - dstn_final[:, 0] += ( - (1 - self.surv_probs[k]) - * np.sum(dstn_init[:, k]) - * self.newborn_dstn - ) - - # In at the end of the last age, everyone turns into a newborn - dstn_final[:, 0] += np.sum(dstn_init[:, -1]) * self.newborn_dstn - - else: - # Living-to-age+1 - dstn_final += self.surv_probs[0] * np.dot( - self.living_transitions[0].T, dstn_init - ) - # Living-to-newborn - dstn_final[:, 0] += (1 - self.surv_probs[0]) * self.newborn_dstn - - return dstn_final - def get_full_tmat(self): if self.life_cycle: # Life cycle @@ -347,3 +317,51 @@ def get_full_tmat(self): ) return full_mat + + def iterate_dstn_forward(self, dstn_init: np.ndarray) -> np.ndarray: + # Initialize final distribution + dstn_final = np.zeros_like(dstn_init) + + if self.life_cycle: + for k in range(self.T - 2): + # Living-to-age+1 + dstn_final[:, k + 1] += self.surv_probs[k] * np.dot( + dstn_init[:, k], self.living_transitions[k] + ) + # Living-to-newborn + dstn_final[:, 0] += ( + (1 - self.surv_probs[k]) + * np.sum(dstn_init[:, k]) + * self.newborn_dstn + ) + + # In at the end of the last age, everyone turns into a newborn + dstn_final[:, 0] += np.sum(dstn_init[:, -1]) * self.newborn_dstn + + else: + # Living-to-age+1 + dstn_final += self.surv_probs[0] * np.dot( + self.living_transitions[0].T, dstn_init + ) + # Living-to-newborn + dstn_final[:, 0] += (1 - self.surv_probs[0]) * self.newborn_dstn + + return dstn_final + + def find_steady_state_dstn( + self, dstn_init, tol=1e-10, max_iter=1000, check_every=10, normalize_every=20 + ): + # Initialize + dstn = dstn_init + err = tol + 1 + i = 0 + while err > tol and i < max_iter: + dstn_new = self.iterate_dstn_forward(dstn) + if i % normalize_every == 0: + dstn_new /= np.sum(dstn_new) + if i % check_every == 0: + err = np.max(np.abs(dstn_new - dstn)) + dstn = dstn_new + i += 1 + + return dstn From 126e8c543e309cee378c315a8cb6988ebf974b7c Mon Sep 17 00:00:00 2001 From: MateoVG Date: Tue, 1 Aug 2023 16:36:56 -0400 Subject: [PATCH 23/28] Use tool for evaluating outcomes of interest as functions --- HARK/core.py | 38 +++++++++++++++++++++++++++++++++++++- 1 file changed, 37 insertions(+), 1 deletion(-) diff --git a/HARK/core.py b/HARK/core.py index cbf904763..2f43a2a7a 100644 --- a/HARK/core.py +++ b/HARK/core.py @@ -16,6 +16,7 @@ import numpy as np import pandas as pd from xarray import DataArray +from inspect import signature from HARK.distribution import ( Distribution, @@ -893,6 +894,41 @@ def make_shock_distributions(self): # Save list of shock distributions self.full_shock_dstns = shock_dstns + def eval_outcomes_on_mesh(self, outcomes): + # TODO: stuff that depends on time varying parameters + + T = len(self.solution) + M = len(self.state_grid.coords["mesh"]) + + outcome_mesh = {} + + for out_name, out_fn in outcomes.items(): + # Initialize + outcome_mesh[out_name] = np.zeros((M, T)) + + # Parse args + args = list(signature(out_fn).parameters) + if len(args) > 0 and args[0] == "solution": + uses_sol = True + state_args = args[1:] + else: + uses_sol = False + state_args = args + + if uses_sol: + for t in range(T): + outcome_mesh[out_name][:, t] = out_fn( + *([self.solution[t]] + [self.state_grid[x] for x in state_args]) + ) + else: + # TODO: this could change with time-varying parameters + for t in range(T): + outcome_mesh[out_name][:, t] = out_fn( + *[self.state_grid[x] for x in state_args] + ) + + return outcome_mesh + 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: @@ -984,7 +1020,7 @@ def trans_wrapper(shocks_next, solution, state_points): living_transitions=surv_mats, surv_probs=self.LivPrb, newborn_dstn=newborn_mass, - life_cycle= T > 1, + life_cycle=T > 1, ) From 10210f14c1aef314c2eec93ebae0c870f4f0ed4b Mon Sep 17 00:00:00 2001 From: MateoVG Date: Tue, 1 Aug 2023 16:37:19 -0400 Subject: [PATCH 24/28] Compare with Will in more detail --- .../tests/test_IndShockConsumerType.py | 73 +++++++++++++------ HARK/mat_methods.py | 13 +++- 2 files changed, 62 insertions(+), 24 deletions(-) diff --git a/HARK/ConsumptionSaving/tests/test_IndShockConsumerType.py b/HARK/ConsumptionSaving/tests/test_IndShockConsumerType.py index bc1c75029..b73b3a7d3 100644 --- a/HARK/ConsumptionSaving/tests/test_IndShockConsumerType.py +++ b/HARK/ConsumptionSaving/tests/test_IndShockConsumerType.py @@ -942,28 +942,57 @@ def setUp(self): ) def test_compare_will_mateo(self): - # Deaths - params = deepcopy(dict_harmenberg) - params["LivPrb"] = [0.9] - - # Create and solve agent - agent = IndShockConsumerType(**params) - agent.cycles = 0 - agent.solve() - # Activate harmenberg - agent.neutral_measure = True - agent.update_income_process() - - # Will's methods - agent.define_distribution_grid(dist_mGrid=self.m_grid) - agent.calc_transition_matrix() - tm_will = agent.tran_matrix + # Create and solve agents + will = IndShockConsumerType(**dict_harmenberg) + mateo = IndShockConsumerType(**dict_harmenberg) + for agent in [will, mateo]: + agent.cycles = 0 + agent.solve() + # Activate harmenberg + agent.neutral_measure = True + agent.update_income_process() + + # %% 1. Transition matrices + + # Define grids (Will) + will.define_distribution_grid() + m_grid = will.dist_mGrid + # Define grids and shock distributions (Mateo) + mateo.make_state_grid(mNrmGrid=m_grid) + mateo.full_shock_dstns = mateo.IncShkDstn + + # Calculate transition matrix (Will) + will.calc_transition_matrix() + tm_will = will.tran_matrix + # Calculate transition matrix (Mateo) + mateo.find_transition_matrices(newborn_dstn=self.newborn_dstn) + tm_mateo = mateo.trans_mat.get_full_tmat() + # Compare + self.assertTrue(np.allclose(tm_will, tm_mateo.T, atol=1e-10)) - # Mateo's methods - agent.make_state_grid(mNrmGrid=self.m_grid) - agent.full_shock_dstns = agent.IncShkDstn - agent.find_transition_matrices(newborn_dstn=self.newborn_dstn) - tm_mateo = agent.trans_mat.get_full_tmat() + # %% 2. Ergodic distributions + # Steady state distribution (Will) + will.calc_ergodic_dist() + will_ss = will.vec_erg_dstn + # Steady state distribution (Mateo) + mateo_ss = mateo.trans_mat.find_steady_state_dstn(tol=1e-14, max_iter=1e4) # Compare - self.assertTrue(np.allclose(tm_will, tm_mateo.T, atol=1e-10)) + self.assertTrue(np.allclose(will_ss, mateo_ss, atol=1e-14)) + + # Outcome grids (Will) + a_points_will = will.aPol_Grid + c_points_will = will.cPol_Grid + # Outcome grids (Mateo) + out_fns = { + "c": lambda solution, mNrm: solution.cFunc(mNrm), + "a": lambda solution, mNrm: mNrm - solution.cFunc(mNrm), + } + points_mateo = mateo.eval_outcomes_on_mesh(out_fns) + # Compare + self.assertTrue( + np.allclose(a_points_will, points_mateo["a"].flatten(), atol=1e-14) + ) + self.assertTrue( + np.allclose(c_points_will, points_mateo["c"].flatten(), atol=1e-14) + ) diff --git a/HARK/mat_methods.py b/HARK/mat_methods.py index da9ce4a9f..c5b715af2 100644 --- a/HARK/mat_methods.py +++ b/HARK/mat_methods.py @@ -349,8 +349,17 @@ def iterate_dstn_forward(self, dstn_init: np.ndarray) -> np.ndarray: return dstn_final def find_steady_state_dstn( - self, dstn_init, tol=1e-10, max_iter=1000, check_every=10, normalize_every=20 + self, dstn_init=None, tol=1e-10, max_iter=1000, check_every=10, normalize_every=20 ): + + if dstn_init is None: + # Create an initial distribution that concentrates + # on the first gridpoint of the first age + dstn_init = dstn = np.zeros( + (len(self.newborn_dstn), len(self.living_transitions)) + ) + dstn[0, 0] = 1.0 + # Initialize dstn = dstn_init err = tol + 1 @@ -364,4 +373,4 @@ def find_steady_state_dstn( dstn = dstn_new i += 1 - return dstn + return dstn \ No newline at end of file From 6aa5f345c87a6897a78f7c127d0095c28d7460a2 Mon Sep 17 00:00:00 2001 From: Mateo VG Date: Tue, 1 Aug 2023 20:10:19 -0400 Subject: [PATCH 25/28] Typo --- HARK/ConsumptionSaving/tests/test_ConsPortfolioModel.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/HARK/ConsumptionSaving/tests/test_ConsPortfolioModel.py b/HARK/ConsumptionSaving/tests/test_ConsPortfolioModel.py index bd475297f..10dc99706 100644 --- a/HARK/ConsumptionSaving/tests/test_ConsPortfolioModel.py +++ b/HARK/ConsumptionSaving/tests/test_ConsPortfolioModel.py @@ -375,7 +375,7 @@ def test_adjust(self): AdjustGrid=None, ) agent.solve() - agent.find_transition_matrices(newborn_dstn=self) + agent.find_transition_matrices(newborn_dstn=self.newborn_dstn) self.assertTrue( agent.trans_mat.living_transitions[0].size == np.power(npoints, 2) ) From 9555210e6cb67671853c80c428cacae1a87cda5f Mon Sep 17 00:00:00 2001 From: MateoVG Date: Wed, 23 Aug 2023 14:56:57 -0400 Subject: [PATCH 26/28] Implement pre- and post-multiplication by transition mats --- HARK/mat_methods.py | 107 ++++++++++++++++++++++++++++++++- HARK/tests/test_mat_methods.py | 70 ++++++++++++++++++++- 2 files changed, 173 insertions(+), 4 deletions(-) diff --git a/HARK/mat_methods.py b/HARK/mat_methods.py index c5b715af2..6e39d89b8 100644 --- a/HARK/mat_methods.py +++ b/HARK/mat_methods.py @@ -318,6 +318,103 @@ def get_full_tmat(self): return full_mat + def post_multiply(self, mat): + # Check dimension compatibility + n_rows, n_cols = mat.shape + if self.life_cycle: + ncols_fullmat = self.T * self.grid_len + else: + ncols_fullmat = self.grid_len + + if n_rows != ncols_fullmat: + raise Exception( + "Matrix has {} rows, but should have {}".format(n_rows, ncols_fullmat) + ) + + if self.life_cycle: + full_mat_dim = self.T * self.grid_len + prod = np.zeros((full_mat_dim, n_cols)) + + for k in range(self.T): + row_init = k * self.grid_len + row_end = row_init + self.grid_len + if k < self.T - 1: + sp = self.surv_probs[k] + else: + sp = 0.0 + + for j in range(n_cols): + # From the newborn dstn + prod[row_init:row_end, j] += (1 - sp) * np.dot( + self.newborn_dstn[np.newaxis, :], mat[: self.grid_len, j] + ) + if k < self.T - 1: + # From the living dstn + prod[row_init:row_end, j] += sp * np.dot( + self.living_transitions[k], + mat[row_end : (row_end + self.grid_len), j], + ) + + return prod + + else: + # Infinite horizon + prod = self.surv_probs[0] * np.dot(self.living_transitions[0], mat) + prod += (1 - self.surv_probs[0]) * np.dot( + self.newborn_dstn[np.newaxis, :], mat + ) + return prod + + def pre_multiply(self, mat): + + # Check dimension compatibility + n_rows, n_cols = mat.shape + if self.life_cycle: + nrows_fullmat = self.T * self.grid_len + else: + nrows_fullmat = self.grid_len + + if n_cols != nrows_fullmat: + raise Exception( + "Matrix has {} cols, but should have {}".format(n_cols, nrows_fullmat) + ) + + if self.life_cycle: + full_mat_dim = self.T * self.grid_len + prod = np.zeros((n_rows, nrows_fullmat)) + + for k in range(self.T): + col_init = k * self.grid_len + col_end = col_init + self.grid_len + if k < self.T - 1: + sp = self.surv_probs[k] + else: + sp = 0.0 + + # Newborns contribute to first block of + # cols + nb_mat = np.tile(self.newborn_dstn, (self.grid_len, 1)) + prod[:, : self.grid_len] += (1 - sp) * np.dot( + mat[:, col_init:col_end], nb_mat + ) + # Living contribute to other columns + if k < self.T - 1: + prod[:, col_end : (col_end + self.grid_len)] += sp * np.dot( + mat[:, col_init:col_end], self.living_transitions[k].T + ) + + return prod + + else: + # Infinite horizon + prod = np.dot( + mat, + self.surv_probs[0] * self.living_transitions[0] + + (1 - self.surv_probs[0]) * self.newborn_dstn[np.newaxis, :], + ) + + return prod + def iterate_dstn_forward(self, dstn_init: np.ndarray) -> np.ndarray: # Initialize final distribution dstn_final = np.zeros_like(dstn_init) @@ -349,9 +446,13 @@ def iterate_dstn_forward(self, dstn_init: np.ndarray) -> np.ndarray: return dstn_final def find_steady_state_dstn( - self, dstn_init=None, tol=1e-10, max_iter=1000, check_every=10, normalize_every=20 + self, + dstn_init=None, + tol=1e-10, + max_iter=1000, + check_every=10, + normalize_every=20, ): - if dstn_init is None: # Create an initial distribution that concentrates # on the first gridpoint of the first age @@ -373,4 +474,4 @@ def find_steady_state_dstn( dstn = dstn_new i += 1 - return dstn \ No newline at end of file + return dstn diff --git a/HARK/tests/test_mat_methods.py b/HARK/tests/test_mat_methods.py index 7e5b84374..71c257af3 100644 --- a/HARK/tests/test_mat_methods.py +++ b/HARK/tests/test_mat_methods.py @@ -1,7 +1,7 @@ import unittest import numpy as np from HARK.utilities import jump_to_grid_1D, jump_to_grid_2D -from HARK.mat_methods import mass_to_grid +from HARK.mat_methods import mass_to_grid, transition_mat # Compare general mass_to_grid with jump_to_grid_1D @@ -85,3 +85,71 @@ def test_simple_3d(self): self.assertTrue(grid_mass[1, 0, 1] == 1 / 8) self.assertTrue(grid_mass[1, 1, 0] == 1 / 8) self.assertTrue(grid_mass[1, 1, 1] == (1 / 8 + 1.0)) + + +class TestTransMatMultiplication(unittest.TestCase): + + def setUp(self): + + # Create dummy 2-gridpoint problems + + # A newborn "distribution" + self.newborn_dstn = np.array([0.7, 0.3]) + + # Infinite horizon transition + inf_h_trans = np.array([[0.9, 0.1], [0.1, 0.9]]) + self.inf_horizon_mat = transition_mat( + living_transitions=[inf_h_trans], + surv_probs=[0.9], + newborn_dstn=self.newborn_dstn, + life_cycle=False, + ) + + # Life cycle transition + lc_trans = [ + np.array( + [ + [x, 1-x], + [1-x, x], + ] + ) + for x in [0.1, 0.2, 0.3, 0.4, 0.5] + ] + self.life_cycle_mat = transition_mat( + living_transitions=lc_trans, + surv_probs=[0.95, 0.93, 0.92, 0.91, 0.90], + newborn_dstn=self.newborn_dstn, + life_cycle=True, + ) + + def test_post_multiply(self): + + # Infinite horizon + nrows = self.inf_horizon_mat.grid_len + mat = np.random.rand(nrows, 3) + res = self.inf_horizon_mat.post_multiply(mat) + res2 = np.dot(self.inf_horizon_mat.get_full_tmat(), mat) + self.assertTrue(np.allclose(res, res2)) + + # Life cycle + nrows = self.life_cycle_mat.grid_len * self.life_cycle_mat.T + mat = np.random.rand(nrows, 3) + res = self.life_cycle_mat.post_multiply(mat) + res2 = np.dot(self.life_cycle_mat.get_full_tmat(), mat) + self.assertTrue(np.allclose(res, res2)) + + def test_pre_multiply(self): + + # Infinite horizon + ncols = self.inf_horizon_mat.grid_len + mat = np.random.rand(3, ncols) + res = self.inf_horizon_mat.pre_multiply(mat) + res2 = np.dot(mat,self.inf_horizon_mat.get_full_tmat()) + self.assertTrue(np.allclose(res, res2)) + + # Life cycle + ncols = self.life_cycle_mat.grid_len * self.life_cycle_mat.T + mat = np.random.rand(3, ncols) + res = self.life_cycle_mat.pre_multiply(mat) + res2 = np.dot(mat, self.life_cycle_mat.get_full_tmat()) + self.assertTrue(np.allclose(res, res2)) From d96fe710ee7230150e6ade0ef4703875458af4e0 Mon Sep 17 00:00:00 2001 From: MateoVG Date: Wed, 23 Aug 2023 15:05:41 -0400 Subject: [PATCH 27/28] Fix pre-multiplication bug caught with 3x3 test --- HARK/mat_methods.py | 3 +-- HARK/tests/test_mat_methods.py | 17 +++++++---------- 2 files changed, 8 insertions(+), 12 deletions(-) diff --git a/HARK/mat_methods.py b/HARK/mat_methods.py index 6e39d89b8..c699cc70a 100644 --- a/HARK/mat_methods.py +++ b/HARK/mat_methods.py @@ -366,7 +366,6 @@ def post_multiply(self, mat): return prod def pre_multiply(self, mat): - # Check dimension compatibility n_rows, n_cols = mat.shape if self.life_cycle: @@ -400,7 +399,7 @@ def pre_multiply(self, mat): # Living contribute to other columns if k < self.T - 1: prod[:, col_end : (col_end + self.grid_len)] += sp * np.dot( - mat[:, col_init:col_end], self.living_transitions[k].T + mat[:, col_init:col_end], self.living_transitions[k] ) return prod diff --git a/HARK/tests/test_mat_methods.py b/HARK/tests/test_mat_methods.py index 71c257af3..69b7a12f2 100644 --- a/HARK/tests/test_mat_methods.py +++ b/HARK/tests/test_mat_methods.py @@ -88,16 +88,14 @@ def test_simple_3d(self): class TestTransMatMultiplication(unittest.TestCase): - def setUp(self): - # Create dummy 2-gridpoint problems - + # A newborn "distribution" - self.newborn_dstn = np.array([0.7, 0.3]) + self.newborn_dstn = np.array([0.5, 0.3, 0.2]) # Infinite horizon transition - inf_h_trans = np.array([[0.9, 0.1], [0.1, 0.9]]) + inf_h_trans = np.array([[0.9, 0.1, 0.0], [0.1, 0.8, 0.1], [0.0, 0.1, 0.9]]) self.inf_horizon_mat = transition_mat( living_transitions=[inf_h_trans], surv_probs=[0.9], @@ -109,8 +107,9 @@ def setUp(self): lc_trans = [ np.array( [ - [x, 1-x], - [1-x, x], + [x, 1 - x, 0], + [x / 2, 1 - x / 2, 0], + [0, 1 - x, x], ] ) for x in [0.1, 0.2, 0.3, 0.4, 0.5] @@ -123,7 +122,6 @@ def setUp(self): ) def test_post_multiply(self): - # Infinite horizon nrows = self.inf_horizon_mat.grid_len mat = np.random.rand(nrows, 3) @@ -139,12 +137,11 @@ def test_post_multiply(self): self.assertTrue(np.allclose(res, res2)) def test_pre_multiply(self): - # Infinite horizon ncols = self.inf_horizon_mat.grid_len mat = np.random.rand(3, ncols) res = self.inf_horizon_mat.pre_multiply(mat) - res2 = np.dot(mat,self.inf_horizon_mat.get_full_tmat()) + res2 = np.dot(mat, self.inf_horizon_mat.get_full_tmat()) self.assertTrue(np.allclose(res, res2)) # Life cycle From 497a7f289895dc83f5878e76cc7729b61439505c Mon Sep 17 00:00:00 2001 From: MateoVG Date: Wed, 23 Aug 2023 15:09:39 -0400 Subject: [PATCH 28/28] Update distribution iterator --- HARK/mat_methods.py | 27 +-------------------------- 1 file changed, 1 insertion(+), 26 deletions(-) diff --git a/HARK/mat_methods.py b/HARK/mat_methods.py index c699cc70a..4e7983e85 100644 --- a/HARK/mat_methods.py +++ b/HARK/mat_methods.py @@ -415,32 +415,7 @@ def pre_multiply(self, mat): return prod def iterate_dstn_forward(self, dstn_init: np.ndarray) -> np.ndarray: - # Initialize final distribution - dstn_final = np.zeros_like(dstn_init) - - if self.life_cycle: - for k in range(self.T - 2): - # Living-to-age+1 - dstn_final[:, k + 1] += self.surv_probs[k] * np.dot( - dstn_init[:, k], self.living_transitions[k] - ) - # Living-to-newborn - dstn_final[:, 0] += ( - (1 - self.surv_probs[k]) - * np.sum(dstn_init[:, k]) - * self.newborn_dstn - ) - - # In at the end of the last age, everyone turns into a newborn - dstn_final[:, 0] += np.sum(dstn_init[:, -1]) * self.newborn_dstn - - else: - # Living-to-age+1 - dstn_final += self.surv_probs[0] * np.dot( - self.living_transitions[0].T, dstn_init - ) - # Living-to-newborn - dstn_final[:, 0] += (1 - self.surv_probs[0]) * self.newborn_dstn + dstn_final = self.pre_multiply(dstn_init.T).T return dstn_final