[WIP] Transition matrices (general, life-cycle)

HARK/ConsumptionSaving/ConsPortfolioModel.py

@@ -28,7 +28,13 @@
     MargValueFuncCRRA,
     ValueFuncCRRA,
 )
+from HARK.distribution import (
+    combine_indep_dstns,
+    DiscreteDistribution,
+    DiscreteDistributionLabeled,
+)
 from HARK.metric import MetricObject
+import xarray as xr
 
 
 # Define a class to represent the single period solution of the portfolio choice problem
@@ -145,6 +151,20 @@
         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):
     """
@@ -316,6 +336,114 @@
         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
+
+    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)
+
+        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)),
+            },
+            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):
+        # Consumption
+        cNrm = solution.cFunc(state["mNrm"], state["Share"], state["Adjust"])
+        # Savings
+        aNrm = state["mNrm"] - cNrm
+        # Share
+        Share_next = solution.ShareFunc(state["mNrm"], state["Share"], state["Adjust"])
+        # Shock transition
+        state_next = post_state_transition(
+            shocks_next,
+            state["PLvl"],
+            aNrm,
+            Share_next,
+            PermGroFac,
+            Rfree,
+        )
+
+        return state_next
+
+
+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 * 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):
     def __init__(self, verbose=False, quiet=False, **kwds):

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,70 @@ 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))
+
+
+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)
+        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))

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:
     """
@@ -850,6 +852,121 @@
             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))
+
+        # 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])
+
+        # Number of points in full grid
+        mesh_size = self.state_grid.coords["mesh"].size
+
+        # 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, 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((mesh_size, mesh_size))
+            for i in range(mesh_size):
+                tmat[i, :] = mass_to_grid(
+                    points=state_points[:, i, :].T,
+                    mass=pmv,
+                    grids=nt_grids,
+                )
+            # Prepend
+            trans_mats.insert(0, tmat)
+
+        # Save matrices
+        self.trans_mats = trans_mats
+
 
 def solve_agent(agent, verbose):
     """