Implement deaths (but results don't match)

Check that the newborn distn is the same

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

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)