Proof of concept of simulating LC distributions

HARK/ConsumptionSaving/tests/test_ConsPortfolioModel.py

@@ -318,19 +318,20 @@ 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()
 
         # Make shock distribution and grid
         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

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