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

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 @@
         )
 
         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):

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,8 @@
 import HARK.ConsumptionSaving.ConsPortfolioModel as cpm
 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):
@@ -304,3 +306,104 @@ 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):
+        # 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
+        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()
+
+        # Make shock distribution and grid
+        agent.make_shock_distributions()
+        agent.make_state_grid(
+            PLvlGrid=None,
+            mNrmGrid=np.linspace(0, 20, npoints),
+            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(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.
+        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
+        # Find steady_state
+        ss_dstn = agent.trans_mat.find_steady_state_dstn(dstn_init=dstn, max_iter=1e4)
+
+    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, 10, npoints),
+            ShareGrid=None,
+            AdjustGrid=None,
+        )
+        agent.solve()
+        agent.find_transition_matrices(newborn_dstn=self.newborn_dstn)
+        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
+        # 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(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)