Fix bugs, recalibrate, and improve Transition_Matrix_Example notebook

[llorracc]

Transition_Matrix_Example.ipynb — Change Log

Changes made to Will Du's Transition_Matrix_Example.ipynb notebook,
with diagnosis and rationale for each.

---

1. TM-initialized Monte Carlo via inverse-CDF sampling (cell 55)

Problem: The original notebook initialized the MIT-shock Monte Carlo
simulation using HARK's default initialize_sim(), which starts all agents
with near-zero assets and pLvl=1. This means the MC starts far from
steady state, requiring a long burn-in before the MC and TM paths can be
meaningfully compared.

Fix: After initialize_sim(), we overwrite agent states with draws from
the TM ergodic distribution. Specifically, we compute the CDF of the 1-D
Harmenberg ergodic distribution over normalized market resources, then use
piecewise-linear inverse-CDF sampling to draw continuous mNrm values for
each of the MC agents. Asset positions aNrm are interpolated from
the steady-state asset policy, and pLvl is set to 1 (Harmenberg neutral
measure). This ensures both MC and TM start from the identical steady-state
distribution, so any remaining discrepancy is purely MC sampling noise vs TM
discretization error.

Cells affected: 55

---

2. Newborn transitory-shock suppression (t_age fix) (cell 55)

Problem: After initialize_sim(), all agents have t_age=0, marking
them as "newborns." HARK's get_shocks() method
(ConsIndShockModel.py, lines 2191–2193) forces TranShk=1.0 for all
newborn agents when the parameter NewbornTransShk is False (the
default). Since every agent is a newborn in the first simulated period,
all agents receive TranShk=1.0 — completely suppressing transitory
income variation. This created a systematic +0.7% bias in first-period
aggregate consumption that decayed over ~3 periods.

Diagnosis: Inspecting FinHorizonAgent.history['TranShk'][0,:] revealed
all values were exactly 1.0, despite the IncShkDstn containing distinct
transitory shock values. Traced to:

# HARK/ConsumptionSaving/ConsIndShockModel.py, lines 2191-2193
if not NewbornTransShk:
    TranShkNow[newborn] = 1.0

Fix: After overriding state_now with the TM-sampled initial
conditions, set FinHorizonAgent.t_age = np.ones(N_fin, dtype=int) so
agents are not treated as newborns and draw proper transitory shocks from the
first period onward.

Cells affected: 55

---

3. Permanent income grid range (max_p_fac) (cell 13)

Problem: The default max_p_fac=30 in define_distribution_grid()
creates a permanent income grid extending from ~10⁻¹⁰ to ~7.7 × 10⁹ (20
orders of magnitude). When computing aggregate assets in levels
(grida = p × a), the asset weights at extreme p-grid points reach
7.6 × 10¹³. The ergodic distribution has effectively zero mass at these
extremes (negative values at the last few points — numerical artifacts), but
machine-epsilon perturbations (~10⁻¹⁶) introduced by each transition-matrix
multiplication get amplified by these enormous weights into visible aggregate
drift, manifesting as a declining trend in the TM asset path figure.

Diagnosis: The drift scales with max(|grida|):

max_p_fac	p_max	max(|grida|)	200-step drift
30 (default)	7.7e+9	7.6e+13	−3.1e−04
20	3.9e+6	3.8e+10	−1.4e−06
10	2.0e+3	1.9e+07	+6.4e−10
5	4.5e+1	4.4e+05	−9.2e−11

Fix: Changed call to
define_distribution_grid(num_pointsP=110, max_p_fac=10.0).
This keeps the grid within exp(7.6) ≈ 2000, still covering 99.99%+ of mass.
The 110 grid points (yielding 221 total p-grid points) are concentrated
where they matter rather than wasted on points in the empty tail.

Cells affected: 13

---

4. Consistent plot colors across all figures

Problem: Monte Carlo and Transition Matrix lines swapped colors between
figures — MC was sometimes blue (first plotted) and sometimes orange (second
plotted), making cross-figure comparison confusing.

Fix: Defined three color constants in the imports cell:

COLOR_MC   = "tab:blue"
COLOR_TM   = "tab:orange"
COLOR_HARM = "tab:green"

Applied color=COLOR_MC, color=COLOR_TM, and color=COLOR_HARM
explicitly to every plt.plot() call across all plotting cells. Also
standardized legend labels (e.g., inconsistent capitalization like
"transition matrix", "transition Matrix", " Transition Matrices" → uniform
"Transition Matrix" or "Transition Matrix (Harmenberg)").

For the grid-comparison figure ([aggregate_assets_grid_comparison]), TM
lines at different grid sizes use an orange gradient (plt.cm.Oranges) to
visually associate them with the TM family while remaining distinguishable.

Cells affected: 3, 19, 23, 25, 29, 31, 39, 42, 61, 63

---

5. Increased MC agent count

Problem: The original AgentCount (in Dict) was lower, producing
noisier MC paths that made it harder to assess whether MC–TM discrepancies
were real or sampling noise.

Fix: Set AgentCount = 200000 in the calibration dictionary.

Cells affected: 5, 47

---

6. Descriptive figure tags and axis labels

Problem: Figures lacked machine-readable tags and some had minimal axis
labels / titles.

Fix: Added bracketed tags (e.g., # [mc_vs_tm_asset_paths]) as the
first line of each plotting cell for cross-referencing. Added or improved
titles, axis labels, and MC summary statistics (mean, std, ±2σ band) to the
asset-path comparison figure.

Cells affected: 19 and other plotting cells

---

7. MSE decomposition diagnostic (cell 21)

Problem: No quantitative comparison of MC vs TM accuracy was provided.

Fix: Added a bias–variance–MSE decomposition for aggregate assets,
treating the MC long-run mean as the reference. This makes explicit the
tradeoff: MC has zero bias but positive variance; TM has zero variance but
nonzero discretization bias.

Cells affected: 21

---

8. Timing instrumentation (cell 13)

Problem: No visibility into computational cost of TM steps.

Fix: Added time.time() calls around define_distribution_grid,
calc_transition_matrix, and calc_ergodic_dist, printing elapsed times.

Cells affected: 13

---

9. Grid size and agent count in figure legends

Problem: Figures did not indicate the TM grid resolution or MC agent
count, making it impossible to judge whether discrepancies were due to
insufficient grid points, insufficient agents, or genuine bugs.

Fix: Legend labels now include the relevant parameters, e.g.:

• "Monte Carlo (200,000 agents)"
• "Transition Matrix (100 m-pts, 2D)"
• "Transition Matrix — Harmenberg (1000 m-points)"
• "TM Harmenberg (500 m-points)" (grid-comparison figure)

Values are computed from the live objects (len(example1.dist_mGrid),
n_agents, etc.) so they stay accurate if parameters change.

Cells affected: 19, 23, 25, 29, 31, 39, 42, 61, 63

---

10. Terminology: "MIT shock" → "anticipated deviation" / "perfect-foresight transition path"

Problem: The notebook described the finite-horizon interest-rate
experiment as a "Perfect foresight MIT shock" and labeled the output plots
"MIT Shock IRF." Both terms are inaccurate for the exercise being performed.
An "MIT shock" in the heterogeneous-agent literature refers to an
unanticipated aggregate disturbance — a zero-probability event that
surprises agents at the moment it occurs (Boppart, Krusell, and Mitman, 2018,
JEDC). An "impulse response function" traces the economy's dynamics from
that moment of surprise onward.

In this notebook, the interest rate deviation at $t = t^*$ is fully
anticipated from $t=0$: the finite-horizon problem is solved by backward
induction over the full $T$-period horizon, so consumption and savings
policies at every date reflect knowledge of the upcoming rate change. This
corresponds to what the SSJ literature calls the response to "s-period-ahead
news" (Auclert, Bardóczy, Rognlie, and Straub, 2021, Econometrica,
Section 2.2) or, equivalently, a nonlinear perfect-foresight transition path
(McKay, 2023, lecture notes).

Fix: Replaced cell 32 with a thorough description of the experiment,
citing Auclert et al. (2021), Boppart et al. (2018), McKay (2023), and
Harmenberg (2021). Updated section headers and plot titles to use
"Perfect-Foresight Transition" and "Anticipated Rate Deviation" instead of
"MIT Shock IRF."

Cells affected: 32, 60, 61, 62, 63

---

11. Terminal condition bug: cFunc_terminal_ vs solution_terminal (cell 49)

Problem: The notebook set
FinHorizonAgent.cFunc_terminal_ = deepcopy(ss.solution[0].cFunc) to give the
finite-horizon agent the steady-state consumption function as its terminal
condition. However, cFunc_terminal_ is a parameter attribute that does not
propagate to solution_terminal, which is the actual object HARK's backward-
induction solver uses. solution_terminal retained its default "consume
everything" behavior (c(m) = m), meaning the last period's agent consumed
all market resources.

Diagnosis: Comparing FinHorizonAgent.cFunc_terminal_(m) vs
FinHorizonAgent.solution_terminal.cFunc(m) showed that the former matched
the SS consumption function while the latter was the identity function.
The solver reads solution_terminal, not cFunc_terminal_.

Fix: Replaced cFunc_terminal_ assignment with:

agent.solution_terminal = deepcopy(ss.solution[0])

This sets the full terminal solution (cFunc, vFunc, vPfunc, hNrm, MPCmin,
etc.) correctly. Now encapsulated in the create_finite_horizon_agent()
helper.

Impact: Small improvement near the terminal period, but did not resolve the
main asset-level discrepancy (see Fix 12).

Cells affected: 3 (helper function), 49

---

12. Critical bug: neutral-measure IncShkDstn used to solve finite-horizon agent

Problem: The finite-horizon agent's consumption problem was solved using
the Harmenberg neutral-measure income shock distribution instead of the
original (true) income shock distribution. The Harmenberg neutral measure
reweights shock probabilities to efficiently track the distribution via
transition matrices, but the agent's optimization problem (Euler equation)
must use the true income distribution.

The bug arose from ordering: ss.neutral_measure = True and
ss.update_income_process() were called before the finite-horizon agent was
created, overwriting ss.IncShkDstn[0] with the neutral-measure version.
Then FinHorizonAgent.IncShkDstn = T_cycle * [ss.IncShkDstn[0]] passed the
wrong distribution to the solver.

Diagnosis: The neutral-measure distribution has the same shock values but
different probability weights (pmv). This caused the Euler equation to
produce systematically wrong consumption policies, leading to aggregate
assets ~3% below the true steady state.

Fix: Save the original IncShkDstn before enabling the neutral measure:

orig_IncShkDstn = deepcopy(ss.IncShkDstn[0])
ss.neutral_measure = True
ss.update_income_process()

Then use orig_IncShkDstn for the finite-horizon solve (now encapsulated
in create_finite_horizon_agent()). After solving, the neutral measure is
enabled on FinHorizonAgent for transition matrix computation (this is
correct — TMs should use the neutral measure).

Cells affected: 36, 49 (and the create_finite_horizon_agent helper in cell 3)

---

13. Recalibrated to cstwMPC baseline (Carroll, Slacalek, Tokuoka & White 2017)

Problem: The previous calibration appeared to be a mix of annual-frequency
and quarterly-frequency values. In particular, DiscFac = 0.975 is a
typical annual discount factor (implying ~2.5 % annual discounting), but
when used at a quarterly frequency it implies ~10 % annual discounting—far
too impatient. Similarly, UnempPrb = 0.0 omitted unemployment risk
entirely, and no single published calibration matched the parameter set.

Fix: Adopted the quarterly infinite-horizon "β-Point" calibration from the
cstwMPC paper (Carroll et al., Quantitative Economics 2017, 8(3), 977–1020;
DOI: 10.3982/QE694). Key parameter changes:

Parameter	Old	New (cstwMPC)
CRRA	2	1.01 (near-log; 1.01 for numerical stability)
Rfree	1.04^¼ ≈ 1.010	1.01 / LivPrb ≈ 1.016 (mortality-adjusted)
DiscFac	0.975	0.9867 (estimated to match K/Y and Lorenz curve)
PermShkStd	0.06	(0.01×4/11)^½ ≈ 0.0603 (same, now with formula)
TranShkStd	0.2	(0.01×4)^½ = 0.2 (same, now with formula)
PermShkCount	5	7
TranShkCount	5	7
UnempPrb	0.00	0.07
IncUnemp	0.0	0.15 (15 % of permanent income)

Unchanged: LivPrb = 0.99375, PermGroFac = 1.0, BoroCnstArt = 0.0.

Income-process estimates originate from Sabelhaus & Song (2010) as used in
the cstwMPC paper.

Also added: A markdown calibration table immediately above the parameter
dictionary (cell 4) with a citation to the cstwMPC paper.

Cells affected: 4, 5

---

14. Documentation improvements

Problem: Several aspects of the notebook were unclear to new readers:

• Cell 32's experiment description stated Δr = -0.05 and t = 10, but the
  code uses dx = -0.01 and shock_t = 100 — a factual contradiction.
• ~15 HARK-specific attributes and methods were used without indicating which
  source files define them.
• The class name NewKeynesianConsumerType would lead a reader to infer this
  is a HANK/NK model, when in fact no NK features are present.
• The column-stochastic matrix convention (critical for correct forward
  propagation: T @ pi, not pi @ T) was mentioned only once in passing.
• Two agent objects (example1 and ss) were created from the same Dict
  without explanation.
• The Bellman equation was referenced but never written down.

Fixes (six changes):

1. Fixed stale parameter values in cell 32. Replaced hardcoded values
   with symbolic Δr and t*, with a parenthetical giving the actual code
   values.

2. Added HARK API source mapping (cell 1, introduction). Maps key methods
   and attributes to their defining source files.

3. Clarified NewKeynesianConsumerType naming (cell 6). Added a note
   explaining the class is used solely for its TM infrastructure.

4. Added column-stochastic convention note (cell 58). Explicitly states
   that HARK matrices are column-stochastic.

5. Explained example1 / ss distinction (cell 33). Notes that ss is
   a fresh instance because the distribution-grid methods mutate internal
   state.

6. Added the one-period Bellman equation (cell 6). States the normalized
   recursive problem with all variables defined and cites Carroll (2006) for
   the EGM solution method.

Cells affected: 1, 6, 32, 33, 58

---

15. Lognormal newborn permanent income (pLogInitStd = 0.4)

Problem: With the default pLogInitStd = 0.0, all newborn MC agents enter
at exactly pLvl = 1.0. Combined with the 7-point discrete approximation to
permanent shocks (PermShkCount = 7), this creates a "multiplicative lattice"
of preferred pLvl values near 1.0. Each period ~0.625% of agents die
(Blanchard–Yaari) and are replaced at the degenerate point, replenishing the
lattice indefinitely. The TM does not show this artifact because it propagates
a continuous probability vector via grid interpolation.

Root cause: Degenerate newborn pLvl distribution + discrete permanent
shock approximation → lattice of products of shock values that never fully
smooths out in the MC histogram.

Fix (two parts):

1. MC side (cell 5): Set pLogInitStd = 0.4 and pLogInitMean = −0.08
   (Jensen correction so E[pLvl] = 1.0), with pLvlInitCount = 25.
   The value 0.4 follows the cstwMPC life-cycle model
   (Carroll, Slacalek, Tokuoka, and White (2017)),
   which calibrates it to match the total variance of income among young
   households in the SCF 2004 data.

2. TM side (cell 13): HARK's calc_transition_matrix hardcodes
   NewBornDist at pLvl = 1.0. We post-process the transition matrix
   via the correct_newborn_dist() helper: subtract
   (1 − LivPrb) × old_NBD and add (1 − LivPrb) × new_NBD
   (lognormal-distributed across the p-grid). This preserves column sums
   and ensures MC and TM solve the same economic model.

   For the Harmenberg (1D) case, no correction is needed: the neutral measure
   integrates out the permanent income dimension, and the Jensen correction
   ensures E[pLvl] = 1.0 for newborns.

Cells affected: 3 (helper), 5, 13

---

16. Distribution plots: probability density, not probability mass

Problem: All four cross-sectional distribution plots (normalized market
resources, permanent income, market resources in levels, liquid assets) plotted
probability mass per bin — histogram_counts / total_counts — rather than
probability density. With the double-exponential grids (timestonest=2),
bin widths vary by an order of magnitude. Narrower bins near the lower bound
(or near pLvl = 1.0 for the p-grid) naturally contain less mass, creating
spurious dips in the plotted distribution even when the true density is smooth.

This was the root cause of the "hole near pLvl = 1.0" that appeared in the
permanent income distribution.

Fix: For TM, divide probability mass at each grid point by the midpoint
bin width to obtain proper probability density. For MC, use
np.histogram(..., density=True) with 200 uniform bins. Change y-axis label
from "Probability" to "Probability Density."

Additional issue in liquid assets plot: aPol_Grid has ~22 entries at
exactly 0.0 (the borrowing constraint — for low-m agents, optimal savings
is a = 0). This creates zero-width bins, producing division-by-zero when
converting to density. Fix: skip to the first positive aPol_Grid entry
and report the mass point at the constraint separately as a printed
diagnostic.

Cells affected: 23 (mNrm), 25 (pLvl), 29 (mLvl), 31 (aLvl)

---

17. Code quality and structural improvements

Problem: The notebook accumulated various code quality issues during
iterative development: scattered imports, duplicated code blocks, Python
loops where vectorized numpy operations apply, variable shadowing bugs,
inconsistent plot formatting, and magic numbers.

Fixes:

1. Consolidated imports. All imports (init_newkeynesian,
   jump_to_grid_2D, LognormalDist) moved to the single imports cell
   (cell 3), eliminating scattered from ... import statements in cells 5
   and 13.

2. Extracted helper functions. Two reusable functions defined in cell 3:

   • correct_newborn_dist(agent, param_dict) — patches the TM newborn
     distribution for lognormal pLvl (was ~20 inline lines in cell 13).
   • create_finite_horizon_agent(ss, Dict, T_cycle, shock_t, dx, orig_IncShkDstn) — creates, configures, and returns a finite-horizon
     agent (was ~30 duplicated lines in cells 49 and 68).
   • jump_to_grid_fast(m_vals, probs, dist_mGrid) — moved from cell 28 to
     the utilities cell for earlier availability.

3. Vectorized loops.

   • gridc/grida computation (cell 14): replaced for j in range(len(p))
     loop with np.outer(c_2d, p_grid_2d).
   • mLvl_vals/aLvl_vals (cell 27): replaced double-nested append loops
     with np.outer(...).ravel().

4. Fixed variable shadowing.

   • Cell 27's inner loop used p as the loop variable, overwriting
     p = example1.dist_pGrid from cell 13. Eliminated by vectorization.
   • Cell 37 assigned c = ss.cPol_Grid and a = ss.aPol_Grid, overwriting
     c and asset from cell 13. Renamed to c_harm / a_harm; 2D-grid
     variables renamed to c_2d / asset_2d.

5. Named constants. Introduced BURNIN = 500, N_MC_BINS = 200,
   N_P_DISC = 50, MAX_P_FAC = 10.0 in cell 3, replacing bare literals
   throughout.

6. Unified variable names. Removed num_pts (duplicate of n_m_grid).
   Defined n_m_grid, n_p_grid, and n_agents once in cell 13 after the
   TM computation, rather than inside plotting cells. Renamed N to
   N_fin in cell 55 for clarity.

7. Standardized plot formatting. Set plt.rcParams defaults for figure
   size, label size, title size, legend size, and line width in cell 3.
   Removed per-cell linewidth=, fontsize=, and prop={"size": ...}
   overrides. Added plt.tight_layout() and explicit figsize to all
   plotting cells.

8. Standardized density plot ranges. All four density plots now use
   the same percentile-based range (0.5–99.5) for both MC and TM.

9. Removed "Known HARK Issues and Workarounds" section. The three
   HARK workarounds (t_age suppression, cFunc_terminal_ vs
   solution_terminal, neutral-measure IncShkDstn leakage) are documented
   with brief inline comments at their respective workaround sites (cells
   36, 49/helper, 55). A standalone section was unnecessary.

10. Unhid cell 5. Removed "source_hidden": true from the calibration
    dictionary cell metadata so the parameters are visible by default.

Cells affected: 3, 5, 13, 14, 18, 19, 23, 25, 27, 28, 29, 31, 37, 39,
42, 49, 51, 55, 61, 63, 68, 71

Made with Cursor

[llorracc]

⚠️ Not ready to merge — needs review — may be updated.

[Copilot]

Pull request overview

This PR fixes how MarkovConsumerType.get_shocks() handles newborns’ permanent income growth in Markov models, and adds regression tests to prevent reintroducing the issue.

Changes:

• Update MarkovConsumerType.get_shocks() so newborns’ PermShk includes deterministic PermGroFac (instead of suppressing growth).
• Add unit tests validating newborn pLvl grows by PermGroFac and that newborns do not receive an idiosyncratic permanent shock dispersion.

Reviewed changes

Copilot reviewed 2 out of 3 changed files in this pull request and generated 3 comments.

File	Description
HARK/ConsumptionSaving/ConsMarkovModel.py	Adjust newborn permanent shock assignment to include deterministic permanent income growth.
tests/ConsumptionSaving/test_ConsMarkovModel.py	Add regression tests for newborn permanent shock behavior under MarkovConsumerType simulation.

---

💡 Add Copilot custom instructions for smarter, more guided reviews. Learn how to get started.

[Copilot]

MarkovConsumerType has a NewbornTransShk parameter (see init_indshk_markov), but get_shocks() still forces TranShkNow[newborn] = 1.0 unconditionally. This makes the parameter ineffective and prevents enabling transitory shock variation for newborns. Consider mirroring IndShockConsumerType.get_shocks() behavior: only overwrite newborn TranShk when NewbornTransShk is False; otherwise leave the drawn transitory shocks in place (and potentially redraw newborn shocks from period-0 distributions for correctness).

Suggested change

	for j in range(self.MrkvArray[0].shape[0]):
	these_nb = np.logical_and(newborn, j == MrkvNow)
	if np.any(these_nb):
	PermShkNow[these_nb] = self.PermGroFac[0][j]
	TranShkNow[newborn] = 1.0
	NewbornTransShk = getattr(self, "NewbornTransShk", False)
	for j in range(self.MrkvArray[0].shape[0]):
	these_nb = np.logical_and(newborn, j == MrkvNow)
	if np.any(these_nb):
	# Deterministic permanent growth for newborns in Markov state j
	PermShkNow[these_nb] = self.PermGroFac[0][j]
	if not NewbornTransShk:
	# Default behavior: no idiosyncratic transitory shock in birth period
	TranShkNow[these_nb] = 1.0
	else:
	# When enabled, give newborns transitory shocks drawn from
	# the period-0 income distribution for their Markov state.
	IncShkDstn_init = self.IncShkDstn[0][j]
	N_nb = np.sum(these_nb)
	EventDraws_nb = IncShkDstn_init.draw_events(N_nb)
	TranShkNow[these_nb] = IncShkDstn_init.atoms[1][EventDraws_nb]

[Copilot]

test_newborn_pLvl_grows_deterministic asserts on np.mean(pLvl_t0) even though the docstring says newborn pLvl should be exactly pLvl_0 * G. Using the mean could mask cross-sectional dispersion or a subset of agents not receiving PermGroFac. Consider asserting pLvl_t0 is elementwise equal to G (similar to the second test) to make the regression check stricter.

[Copilot]

The PR title/description focuses on Transition_Matrix_Example.ipynb, but the diff shown here changes core library code (ConsMarkovModel.py) and adds unit tests. If the notebook changes are intended, they don’t appear in this PR diff; either include the notebook file changes or update the PR metadata to reflect that this PR is primarily fixing Markov newborn shock behavior and adding regression tests.