Add Monte Carlo vs Transition Matrix example notebooks and infrastructure

.gitignore

@@ -299,6 +299,18 @@ uv.lock
 .cursorignore
 .cursorindexingignore
 
+# Local simulation notes and scratch work
+sims-about/**
+
+# AI tool config
+.ragignore
+
+# Debug/scratch notebooks (keep locally, don't ship)
+*_debug*.ipynb
+
+# Test scripts in example directories (real tests go in tests/)
+examples/**/test_*.py
+
 # Generated test artifacts
 test.jpg
 test.pdf

HARK/ConsumptionSaving/ConsIndShockModel.py

@@ -1330,6 +1330,56 @@ def initialize_sim(self):
         self.PermShkAggNow = self.PermGroFacAgg  # This never changes during simulation
         self.state_now["PlvlAgg"] = 1.0
         super().initialize_sim()
+        self._initialize_ergodic_ages()
+
+    def _initialize_ergodic_ages(self):
+        """
+        For infinite-horizon models with finite T_age, re-draw agent ages
+        from the ergodic (truncated geometric) distribution and scale pLvl
+        accordingly.
+
+        Without this, all agents start at age 0 and die together every T_age
+        periods, creating a "cohort echo" that biases aggregates (e.g. ~1.25%
+        in E[pLvl] for HAFiscal's T_age=200, LivPrb=0.99375).
+
+        In state_now (after sim_one_period), t_age ranges from 1 to T_age
+        because t_age is incremented at the end of each period.  The ergodic
+        distribution is P(t_age=k) = C * L^(k-1) for k=1,...,T_age, with
+        C = (1-L)/(1-L^T_age).  Agents at t_age=T_age will die immediately
+        in the first sim_one_period call, matching steady-state turnover.
+
+        Skipped for lifecycle models (cycles != 0) where age-0 start is correct.
+        Controlled by attribute `init_ages_ergodic` (default True).
+        """
+        if not getattr(self, "init_ages_ergodic", True):
+            return
+        if self.cycles != 0 or self.T_age is None:
+            return
+
+        L = np.asarray(self.LivPrb[0])
+        if L.ndim > 0:
+            L = float(L[0])
+        else:
+            L = float(L)
+
+        T = self.T_age
+        ages = np.arange(1, T + 1)  # 1 to T_age (state_now convention)
+        if abs(L - 1.0) < 1e-14:
+            probs = np.ones(T) / T
+        else:
+            C = (1.0 - L) / (1.0 - L**T)
+            probs = C * L ** (ages - 1)
+            probs /= probs.sum()
+
+        self.t_age = self.RNG.choice(ages, size=self.AgentCount, p=probs)
+        self.t_cycle = self.t_age % self.T_cycle
+
+        G = np.asarray(self.PermGroFac[0])
+        if G.ndim > 0:
+            G = float(G[0])
+        else:
+            G = float(G)
+        self.state_now["pLvl"] *= G**self.t_age
 
     def sim_birth(self, which_agents):
         """

HARK/ConsumptionSaving/ConsNewKeynesianModel.py

@@ -131,6 +131,7 @@ class NewKeynesianConsumerType(IndShockConsumerType):
     default_ = {
         "params": init_newkeynesian,
         "solver": solve_one_period_ConsIndShock,
+        "model": "ConsIndShock.yaml",
         "track_vars": ["aNrm", "cNrm", "mNrm", "pLvl"],
     }
 

HARK/core.py

@@ -1548,13 +1548,21 @@ def initialize_sim(self):
                         ]
 
                 else:
-                    warn(
-                        "The option for reading shocks was activated but "
-                        + "the model requires state "
-                        + var_name
-                        + ", not contained in "
-                        + "newborn_init_history."
+                    # Only warn for idiosyncratic states (per-agent arrays).
+                    # Aggregate scalars (e.g. PlvlAgg) are not expected in
+                    # newborn_init_history and are set elsewhere.
+                    is_idio = (
+                        isinstance(self.state_now[var_name], np.ndarray)
+                        and len(self.state_now[var_name]) == self.AgentCount
                     )
+                    if is_idio:
+                        warn(
+                            "The option for reading shocks was activated but "
+                            + "the model requires state "
+                            + var_name
+                            + ", not contained in "
+                            + "newborn_init_history."
+                        )
 
         self.clear_history()
 

HARK/distributions/base.py

@@ -176,8 +176,9 @@ class MarkovProcess(Distribution):
     Parameters
     ----------
     transition_matrix : np.array
-        An array of floats representing a probability mass for
-        each state transition.
+        Row-stochastic transition matrix: ``transition_matrix[i, j]`` is the
+        probability of moving to state *j* given the current state is *i*.
+        Each row must sum to 1.
     seed : int
         Seed for random number generator.
 

HARK/simulator.py

@@ -10,7 +10,7 @@
 from sympy.utilities.lambdify import lambdify
 from sympy import symbols, IndexedBase
 from typing import Callable
-from HARK.utilities import NullFunc, make_exponential_grid
+from HARK.utilities import NullFunc, make_exponential_grid, make_grid_exp_mult
 from HARK.distributions import Distribution
 from scipy.sparse import csr_matrix, csc_matrix
 from scipy.sparse.linalg import eigs
@@ -649,7 +649,18 @@ def _build_input_grids(self, grid_specs, arrival_N):
                 is_arrival = True
             except:
                 is_arrival = False
-            if ("min" in spec) and ("max" in spec):
+            if "grid" in spec:
+                new_grid = np.asarray(spec["grid"], dtype=float)
+                is_cont = True
+                grid_orders[var] = 0.0
+            elif "timestonest" in spec:
+                bot = spec["min"]
+                top = spec["max"]
+                N = spec["N"]
+                new_grid = make_grid_exp_mult(bot, top, N, spec["timestonest"])
+                is_cont = True
+                grid_orders[var] = 0.0
+            elif ("min" in spec) and ("max" in spec):
                 Q = spec["order"] if "order" in spec else 1.0
                 bot = spec["min"]
                 top = spec["max"]
@@ -986,11 +997,14 @@ def make_transition_matrices(self, grid_specs, twist=None, norm=None):
         Parameters
         ----------
         grid_specs : dict
-            Dictionary of dictionaries of grid specifications. For now, these have
-            at most a minimum value, a maximum value, a number of nodes, and a poly-
-            nomial order. They are equispaced if a min and max are specified, and
-            polynomially spaced with the specified order > 0 if provided. Otherwise,
-            they are set at 0,..,N if only N is provided.
+            Dictionary of dictionaries of grid specifications. Each entry can be:
+            - {"min", "max", "N"}: linearly spaced grid.
+            - {"min", "max", "N", "order"}: polynomially spaced (x^order).
+            - {"min", "max", "N", "timestonest"}: multi-exponential grid via
+              make_grid_exp_mult, matching HARK's legacy double-exponential
+              grids (e.g. timestonest=2). Requires non-negative bounds.
+            - {"grid": array}: explicit user-supplied grid array.
+            - {"N"}: discrete grid 0,..,N-1.
         twist : dict or None
             Mapping from end-of-period (continuation) variables to successor's
             arrival variables. When this is specified, additional output is created
@@ -1504,10 +1518,14 @@ def make_transition_matrices(
             of all variables of interest. If any arrival variables are omitted,
             they will be given a default trivial grid with one node at 0. This
             should only be done if that arrival variable is closely tied to the
-            Harmenberg normalizing variable; see below. A grid specification must
-            include a number of gridpoints N, and should also include a min and
-            max if the variable is continuous. If the variable is discrete, the
-            grid values are assumed to be 0,..,N.
+            Harmenberg normalizing variable; see below. Each entry can be:
+            - {"min", "max", "N"}: linearly spaced grid.
+            - {"min", "max", "N", "order"}: polynomially spaced (x^order).
+            - {"min", "max", "N", "timestonest"}: multi-exponential grid via
+              make_grid_exp_mult, matching HARK's legacy double-exponential
+              grids (e.g. timestonest=2). Requires non-negative bounds.
+            - {"grid": array}: explicit user-supplied grid array.
+            - {"N"}: discrete grid 0,..,N-1.
         norm : str or None
             Name of the variable for which Harmenberg normalization should be
             applied, if any. This should be a variable that is directly drawn
@@ -3660,15 +3678,20 @@ def aggregate_blobs_onto_polynomial_grid(
     grid of outcome values, based on their origin in the arrival state space. This
     version is for non-continuation variables, returning only the probability array
     mapping from arrival states to the outcome variable.
+
+    When Q > 0, uses the polynomial inverse formula for O(1) index lookup.
+    When Q <= 0, uses binary search (searchsorted) for arbitrary grids.
     """
     bot = grid[0]
     top = grid[-1]
     M = grid.size
     Mm1 = M - 1
     N = pmv.size
-    scale = 1.0 / (top - bot)
-    order = 1.0 / Q
     diffs = grid[1:] - grid[:-1]
+    use_poly = Q > 0.0
+    if use_poly:
+        scale = 1.0 / (top - bot)
+        order = 1.0 / Q
 
     probs = np.zeros((J, M))
 
@@ -3677,7 +3700,14 @@ def aggregate_blobs_onto_polynomial_grid(
         jj = origins[n]
         p = pmv[n]
         if (x > bot) and (x < top):
-            ii = int(np.floor(((x - bot) * scale) ** order * Mm1))
+            if use_poly:
+                ii = int(np.floor(((x - bot) * scale) ** order * Mm1))
+            else:
+                ii = np.searchsorted(grid, x) - 1
+                if ii < 0:
+                    ii = 0
+                if ii >= Mm1:
+                    ii = Mm1 - 1
             temp = (x - grid[ii]) / diffs[ii]
             probs[jj, ii] += (1.0 - temp) * p
             probs[jj, ii + 1] += temp * p
@@ -3698,15 +3728,20 @@ def aggregate_blobs_onto_polynomial_grid_alt(
     version is for continuation variables, returning the probability array mapping
     from arrival states to the outcome variable, the index in the outcome variable grid
     for each blob, and the alpha weighting between gridpoints.
+
+    When Q > 0, uses the polynomial inverse formula for O(1) index lookup.
+    When Q <= 0, uses binary search (searchsorted) for arbitrary grids.
     """
     bot = grid[0]
     top = grid[-1]
     M = grid.size
     Mm1 = M - 1
     N = pmv.size
-    scale = 1.0 / (top - bot)
-    order = 1.0 / Q
     diffs = grid[1:] - grid[:-1]
+    use_poly = Q > 0.0
+    if use_poly:
+        scale = 1.0 / (top - bot)
+        order = 1.0 / Q
 
     probs = np.zeros((J, M))
     idx = np.empty(N, dtype=np.dtype(np.int32))
@@ -3717,7 +3752,14 @@ def aggregate_blobs_onto_polynomial_grid_alt(
         jj = origins[n]
         p = pmv[n]
         if (x > bot) and (x < top):
-            ii = int(np.floor(((x - bot) * scale) ** order * Mm1))
+            if use_poly:
+                ii = int(np.floor(((x - bot) * scale) ** order * Mm1))
+            else:
+                ii = np.searchsorted(grid, x) - 1
+                if ii < 0:
+                    ii = 0
+                if ii >= Mm1:
+                    ii = Mm1 - 1
             temp = (x - grid[ii]) / diffs[ii]
             probs[jj, ii] += (1.0 - temp) * p
             probs[jj, ii + 1] += temp * p

HARK/utilities.py

@@ -850,6 +850,84 @@ def gen_tran_matrix_2D(
     return TranMatrix
 
 
+@numba.njit(parallel=True)
+def gen_tran_matrix_1D_markov(
+    dist_mGrid,
+    aPol_Grid,
+    MrkvArray,
+    Rfree_arr,
+    PermGroFac_arr,
+    LivPrb_arr,
+    shk_prbs,
+    perm_shks,
+    tran_shks,
+    NewBornDist,
+):  # pragma: nocover
+    """
+    Computes the block-structured transition matrix for a MarkovConsumerType
+    using the permanent-income-neutral measure (1D m-grid per Markov state).
+
+    The state space is (m, j) flattened into a vector of length M*J, where
+    indices [j*M : (j+1)*M] correspond to Markov state j.  The matrix is
+    column-stochastic: column src represents "starting in state src" and
+    rows represent "arriving in state dst".
+
+    Parameters
+    ----------
+    dist_mGrid : np.array, shape (M,)
+        Grid over normalized market resources (same for all Markov states).
+    aPol_Grid : np.array, shape (J, M)
+        End-of-period asset policy for each Markov state j evaluated on dist_mGrid.
+    MrkvArray : np.array, shape (J, J)
+        Row-stochastic Markov transition matrix.  MrkvArray[j, jp] = P(jp | j).
+    Rfree_arr : np.array, shape (J,)
+        Risk-free interest factor for each Markov state.
+    PermGroFac_arr : np.array, shape (J,)
+        Permanent income growth factor for each Markov state.
+    LivPrb_arr : np.array, shape (J,)
+        Survival probability for each Markov state.
+    shk_prbs : np.array
+        Shock probabilities (neutral-measure weights).
+    perm_shks : np.array
+        Permanent shock values (neutral-measure adjusted).
+    tran_shks : np.array
+        Transitory shock values.
+    NewBornDist : np.array, shape (M*J,)
+        Distribution of newborns across the full (m, j) state space.
+
+    Returns
+    -------
+    TranMatrix : np.array, shape (M*J, M*J)
+        Column-stochastic transition matrix.
+    """
+    J = MrkvArray.shape[0]
+    M = len(dist_mGrid)
+    N = M * J
+    TranMatrix = np.zeros((N, N))
+
+    for src in numba.prange(N):
+        j = src // M
+        i = src % M
+        LivPrb_j = LivPrb_arr[j]
+
+        for jp in range(J):
+            markov_prob = MrkvArray[j, jp]
+            if markov_prob < 1e-15:
+                continue
+
+            bNext_i = Rfree_arr[jp] * aPol_Grid[j, i]
+            mNext_shks = bNext_i / (perm_shks * PermGroFac_arr[jp]) + tran_shks
+            lottery_1d = jump_to_grid_1D(mNext_shks, shk_prbs, dist_mGrid)
+
+            TranMatrix[jp * M : (jp + 1) * M, src] += (
+                markov_prob * LivPrb_j * lottery_1d
+            )
+
+        TranMatrix[:, src] += (1.0 - LivPrb_j) * NewBornDist
+
+    return TranMatrix
+
+
 # ==============================================================================
 # ============== Some basic plotting tools  ====================================
 # ==============================================================================

docs/example_notebooks/Include_list.txt

@@ -42,4 +42,7 @@ examples/SequenceSpaceJacobians/KS-HARK-presentation.ipynb
 examples/SequenceSpaceJacobians/SSJ_explanation.ipynb
 examples/SequenceSpaceJacobians/SSJ-tutorial.ipynb
 examples/SequenceSpaceJacobians/SSJ-advanced-examples.ipynb
-examples/SequenceSpaceJacobians/HANKFiscal_example.ipynb
+examples/SequenceSpaceJacobians/HANKFiscal_example.ipynb
+examples/MonteCarlovsTransitionMatrix/PE_MarkovConsumerType.ipynb
+examples/MonteCarlovsTransitionMatrix/GE_KrusellSmith.ipynb
+examples/MonteCarlovsTransitionMatrix/Validation_and_SSJ.ipynb

docs/guides/index.rst

@@ -8,6 +8,7 @@ Guides
    quick_start
    installation
    simulation
+   transition_matrix_methods
    krusell_smith
    migration_case_study