Certainty-equivalent seam for non-EU preferences (Epstein-Zin) by hmgaudecker · Pull Request #395 · OpenSourceEconomics/pylcm

hmgaudecker · 2026-07-02T13:27:54Z

Closes #385.

What this PR does

A regime can declare a nonlinear certainty equivalent over the next-period value distribution, so non-expected-utility recursive preferences (Epstein–Zin, risk-sensitive) are expressible without touching the engine:

Regime(
    functions={"utility": utility, "H": H_epstein_zin},
    certainty_equivalent=PowerMean(),
    ...
)

The solve then aggregates the continuation as

$$\mathrm{CE} = g^{-1}\Bigl(\sum_r p_r , \mathbb{E}_w\bigl[g(V'_r)\bigr]\Bigr)$$

— g applied elementwise before the expectation over stochastic state transitions, per target regime, g⁻¹ once after the regime-probability-weighted sum — and H receives the CE through its E_next_V argument. Solve and simulate share the same decision functions, so simulated choices are CE-consistent; the NaN-diagnostics path reproduces the same math.

User API

Certainty equivalents (lcm.certainty_equivalent, defined engine-side behind a thin re-export façade like lcm.solvers):
- CertaintyEquivalent — ABC the engine dispatches on; certainty_equivalent=None (the default) is the linear expectation and traces a graph byte-identical to the status quo.
- QuasiArithmeticMean(transform=g, inverse=g_inv) — the generic transform pair. Callables take the value array via the reserved argument value; every further signature argument becomes a runtime param under the pseudo-function name certainty_equivalent in the params template.
- PowerMean() — the Epstein–Zin power mean with runtime param risk_aversion (γ). risk_aversion = 1 is the geometric-mean (log) limit exp(E[log V']); risk_aversion = 0 reduces to the linear expectation.
Koopmans aggregators (lcm.temporal_aggregation):
- H_linear — U + β·CE; also the default injected when a non-terminal regime supplies no H.
- H_epstein_zin — ((1-β)·U^ρ + β·CE^ρ)^(1/ρ), parametrized directly by the intertemporal_elasticity_of_substitution ψ (curvature ρ = 1 − 1/ψ computed inside; ψ = 1 is the Cobb–Douglas limit).
Validation at model build: terminal regimes reject a CE (no continuation to aggregate), and DCEGM + CE is rejected with a message naming GridSearch — Euler-inversion EGM assumes expected utility.

Example + tests

lcm_examples/epstein_zin.py: an Epstein–Zin lifecycle model in the spirit of the Atal–Fang–Karlsson–Ziebarth (2025) consumer block — savings, a two-state health Markov chain, health-dependent survival into a terminal bequest regime — parametrized mortality-style (n_periods, grid-size knobs, survival_probs).
docs/examples/epstein_zin.ipynb: executed notebook covering the recursion, the mapping onto pylcm, the pitfalls (positivity/0^{1-γ}=∞, stateless targets, solver restriction), and a rendered figure comparing mean wealth paths for two risk aversions at a common IES.
Solved values and simulated policies are pinned against an independent NumPy backward induction written in the tests (γ = 0.5 and the γ = 1 log limit; rtol=5e-5 / 1e-5); a reduction test confirms risk_aversion = 0 equals the no-CE solve, and a divergence test that a nonlinear CE changes it.
API tests: params-template discovery, name collisions, terminal/DCEGM/Phased rejection, aggregator units (CES form, Cobb–Douglas limit, default-H identity).

Non-goals

Full-distribution CE callables (enabled by the CertaintyEquivalent subclass seam, not shipped), EZ-EGM, CE params from DAG outputs, model-level broadcast of certainty_equivalent.

Coordination

Based on Device-local continuation-V for fixed shards (no all-gather) #391 (feat/type-local-continuation-v); retarget to main once that lands.
feat/dcegm (Add the endogenous-grid solver family: EGM, DC-EGM, and NEGM #390) is a sibling on the same base: shared-file hunks here are deliberately small and local (Q_and_F.py, processing.py, contract.py); whichever merges second resolves them. feat/dcegm's real DCEGM.validate should keep rejecting regimes whose SolverBuildContext.certainty_equivalent is not None.
Also fixes docs/user_guide/tiny_example.ipynb (rotted regime/regime_name column names); every notebook in the docs now executes cleanly from a cold cache.

🤖 Generated with Claude Code

…river) Two tests on the 4-CPU distributed seam: - test_distributed_solve_matches_single_device_per_type: the F6 correctness contract — a sharded solve must equal the single-device solve per type slice. Green now; the co-map fix must preserve it (catches a wrong device-local index). - test_distributed_solve_kernel_does_not_all_gather_continuation_v: the red driver — the backward-induction kernel currently all-gathers the continuation V across the type shard; it must read only its device-local slice. Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>

vmap_1d gains co_mapped_in_axes: an optional per-argument in_axes override so a pytree argument's leading axis can be mapped in lockstep with the mapped variables. The backward-induction co-map uses it to slice each next_regime_to_V_arr leaf to the device-local type, so the continuation-V interpolation reads only its own shard. Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>

Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>

Fixed, distributed states (e.g. a permanent type sharded one block per device) never transition, so a regime's continuation value depends only on its own slice of the next-period V-array. The grid-search solve kernel now co-maps each such state with the matching axis of every next_regime_to_V_arr leaf that carries it: an outer vmap peels the leading axis off both the state grid and the continuation V, so the interpolation reads only the device-local slice and XLA inserts no all-gather of the full V-array onto every device. - max_Q_over_a splits the co-mapped (leading) states from the inner productmap and wraps it in per-state co-map vmaps; the V-interpolator drops those coordinates and Q_and_F omits the sliced next-states. - processing detects the co-mappable states (distributed and identity-transition) and builds per-state, per-leaf in_axes so a target regime that prunes the state keeps its full leaf. - Simulation is unchanged: it keeps the full continuation V (subjects are not type-aligned across devices), so the co-map applies to the solve path only. Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>

…nuation-v

…seam (#385) Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>

Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>

…ectations Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>

Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>

- Add 'age' entry to initial_conditions in epstein_zin.md Run section - Document certainty_equivalent parameter in get_model docstring - Add docstrings to _power_transform and _power_inverse helpers Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>

read-the-docs-community · 2026-07-02T13:31:09Z

Documentation build overview

📚 pylcm | 🛠️ Build #33444304 | 📁 Comparing f301419 against latest (7d8d2aa)

🔍 Preview build

38 files changed · + 1 added · ± 37 modified

+ Added

epstein-zin/index.html

± Modified

github-actions · 2026-07-02T14:44:57Z

Benchmark comparison (main → HEAD)

Comparing f03c1c96 (main) → 1e78e86e (HEAD)

Benchmark	Statistic	before	after	Ratio
aca-baseline	execution time	13.882 s	13.315 s	0.96
	peak GPU mem	588 MB	588 MB	1.00
	compilation time	373.03 s	376.61 s	1.01
	peak CPU mem	6.68 GB	7.10 GB	1.06
aca-baseline-debug	execution time	56.034 s	58.102 s	1.04
	peak GPU mem	587 MB	587 MB	1.00
	compilation time	442.14 s	439.54 s	0.99
	peak CPU mem	7.72 GB	7.79 GB	1.01
Mahler-Yum	execution time	4.664 s	4.579 s	0.98
	peak GPU mem	520 MB	520 MB	1.00
	compilation time	11.53 s	11.30 s	0.98
	peak CPU mem	1.58 GB	1.58 GB	1.00
Precautionary Savings - Solve	execution time	25.0 ms	22.9 ms	0.92
	peak GPU mem	8 MB	8 MB	1.00
	compilation time	1.56 s	1.62 s	1.04
	peak CPU mem	1.16 GB	1.16 GB	1.00
Precautionary Savings - Simulate	execution time	62.7 ms	62.4 ms	1.00
	peak GPU mem	157 MB	157 MB	1.00
	compilation time	3.49 s	3.52 s	1.01
	peak CPU mem	1.32 GB	1.33 GB	1.01
Precautionary Savings - Solve & Simulate	execution time	93.3 ms	95.7 ms	1.03
	peak GPU mem	566 MB	566 MB	1.00
	compilation time	4.77 s	4.74 s	0.99
	peak CPU mem	1.31 GB	1.30 GB	0.99
Precautionary Savings - Solve & Simulate (irreg)	execution time	202.0 ms	201.2 ms	1.00
	peak GPU mem	2.18 GB	2.18 GB	1.00
	compilation time	5.07 s	5.10 s	1.01
	peak CPU mem	1.38 GB	1.37 GB	1.00
IskhakovEtAl2017Simulate	execution time	198.4 ms	188.4 ms	0.95
	compilation time	4.18 s	4.17 s	1.00
	peak CPU mem	1.30 GB	1.29 GB	1.00
IskhakovEtAl2017Solve	execution time	46.9 ms	44.8 ms	0.96
	compilation time	0.68 s	0.69 s	1.00
	peak CPU mem	1.15 GB	1.15 GB	1.00
IskhakovEtAl2017SimulateGpuPeakMem	peak GPU mem	281 MB	281 MB	1.00
IskhakovEtAl2017SolveGpuPeakMem	peak GPU mem	67 MB	67 MB	1.00

Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>

…model Rename `TransformedExpectation` to `QuasiArithmeticMean` and `PowerCertaintyEquivalent` to `PowerMean`; the `certainty_equivalent` field, params pseudo-function, and `risk_aversion` parameter are unchanged. Move the engine-side implementation (`CE_VALUE_ARG`, the power transform pair, and `resolve_certainty_equivalent`) into `_lcm/certainty_equivalent.py` so the public module is a thin, deep-module namespace and the solver seam in `Q_and_F.py` only imports the resolver. Collapse the toy `tests/test_models/epstein_zin_health.py` and the example into a single parametrized `lcm_examples.epstein_zin` (`EZRegimeId`, `get_model` with a required `certainty_equivalent` and grid-size knobs whose defaults reproduce the toy numerically). The numpy pinning references pass unchanged. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>

Replace `docs/examples/epstein_zin.md` with `docs/examples/epstein_zin.ipynb`: the same recursion, mapping, and pitfalls as markdown cells, plus code cells that solve and simulate a 20-period model for two risk-aversion values and render a plotly figure of mean wealth by age (grey vs accent, direct labels). Update the toctree in `myst.yml` and the link in `examples/index.md`. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>

Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>

review-notebook-app · 2026-07-02T17:07:07Z

Check out this pull request on

See visual diffs & provide feedback on Jupyter Notebooks.

Powered by ReviewNB

…he example - lcm.aggregators exposes the two standard Koopmans aggregators; the default H is now the public H_linear, and H_epstein_zin is parametrized by the intertemporal elasticity of substitution (curvature rho = 1 - 1/psi computed inside, psi = 1 as the Cobb-Douglas limit). - PowerMean handles risk_aversion = 1 as the geometric-mean (log) limit exp(E[log V']) instead of rejecting it; the numpy reference and pinning tests cover it. - Example notebook: merged the redundant positivity pitfalls, added an expected-utility baseline to the wealth figure, log_level='off'. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>

… trace Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>

Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>

The example model gains three parameters, all defaulting to the previous behavior: income (above the consumption floor, saving becomes possible), health_cost (an out-of-pocket expense while in bad health - uninsurable expense risk in the spirit of the Atal et al. medical spending), and bequest_scale (prices the bequest in consumption-equivalent units; H_epstein_zin is a weighted power mean, so the alive value sits at the scale of per-period consumption and an unscaled sqrt(wealth) bequest makes death the good branch of the certainty equivalent). next_wealth clips to the wealth grid so none of the knobs can push states off-grid. The docs page is rewritten around the fixed model: - Atal, Fang, Karlsson & Ziebarth (2025) is now characterized correctly: their baseline is time-separable CARA expected utility; their robustness specification has exactly the CES-aggregator-around-a- certainty-equivalent structure used here, with a CARA certainty equivalent (expressible as a QuasiArithmeticMean). - The pylcm-mapping section reflects the shipped H_linear/H_epstein_zin and explains that per-period utility must live in consumption units because H is a power mean. - A new pitfall documents the bequest-scaling trap. - The figure simulates 1,000 subjects with saving, health costs, and a gentler mortality hazard, so survivor counts stay meaningful at every plotted age, and the closing text explains the economics: both types dissave early out of impatience, the risk-averse type holds and rebuilds a larger precautionary buffer, and the ranking flips when the bad-health spell becomes too expensive to self-insure. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>

…IES/RA sweep The intro now describes the original Atal, Fang, Karlsson & Ziebarth (2025) model accurately: an annual Yaari life-cycle savings problem over ages 25-94, a seven-category health Markov chain driving expenditure and mortality, and the guaranteed-renewable GLTHI vs short-term insurance contracts. Their baseline is time-separable expected utility (CARA gamma=4e-4, CRRA sigma=4 robustness, delta=0.966); the headline is that GLTHI reaches ~96% of first-best welfare, robust to disentangling risk aversion and the IES (sec. VI.D.1, within 0.7%). A 'what this keeps and drops' paragraph is explicit that the equilibrium contract/premium layer is out of scope, so the paper's headline welfare gap is not something this consumer-block example reproduces. A new closing section makes the disentangling concrete: a 2D sweep of the welfare cost of the uninsurable health-expense risk over a log-2 grid of risk aversion and the IES, with expected utility marked as the exact anti-diagonal (IES = 1/gamma). The cost is a risk premium — it roughly triples down the risk-aversion axis and barely moves along the IES axis — so an EU model, confined to the anti-diagonal, confounds the two. This is exactly the degree of freedom the certainty_equivalent seam adds and the lever the paper's robustness section pulls. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>

The example model enters at age 25 rather than 60, matching the annual horizon of Atal et al. instead of a retirement-only slice. Only the entry age changes and periods stay annual, so the discount factor needs no recompounding. The docs page runs the full 25-to-85 lifecycle: 2,000 subjects under a Gompertz-like mortality hazard (low when young, rising with age), which keeps a well-populated surviving cohort through midlife so both figures read cleanly. The risk-averse agent holds a persistently larger precautionary buffer, and the welfare cost of the health-expense risk roughly doubles down the risk-aversion axis while barely moving along the IES axis. Prose and quoted numbers track the new lifecycle. The notebook executes in about ten seconds on CPU, within the Read-the-Docs build budget. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>

The intro markdown cell had lost its newlines and collapsed into a single line, so its headings and paragraphs ran together on the rendered page. Rebuild it with one array element per line. Also switch the recursion's display equation from a ```{math}``` directive — which this project's MyST does not process inside notebook cells, leaking raw LaTeX — to the $$ form used by every other notebook, so it typesets. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>

mj023

Looks good! I am wondering if it might be good to increase the scope of the h function, such that it also contains the whole aggregation of next_V_at_stochastic_states. Then it would be even more flexible, but we could still offer the defaults for normal expectations and Epstein-Zin preferences. Maybe it's easier to keep it seperate for people who only want state dependent beta etc. though.

mj023 · 2026-07-03T17:32:27Z

+def power_inverse(value: FloatND, risk_aversion: FloatND) -> FloatND:
+    """Apply `g^(-1)(v) = v^(1 / (1 - risk_aversion))`; `exp(v)` in the log case."""
+    # The unselected power branch must not divide by zero at `risk_aversion = 1`.
+    safe_risk_aversion = jnp.where(risk_aversion == 1.0, 0.0, risk_aversion)


Not sure if it matters if we divide by zero, if we then select the other value anyways.

mj023 · 2026-07-03T17:45:28Z

+    """
+    rho = 1.0 - 1.0 / intertemporal_elasticity_of_substitution
+    # The unselected CES branch must not divide by zero at `ψ = 1`.
+    safe_rho = jnp.where(rho == 0.0, 1.0, rho)


Same as above.

The two safe_* guards read as divide-by-zero protection, prompting the question of why they matter when jnp.where selects the finite branch anyway. The forward pass is NaN-free without them; they exist solely to keep the reverse-mode gradient finite at the risk_aversion=1 / psi=1 limits. Reword the comments to say so. Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>

…lent

Integrate the 28 commits feat/dcegm advanced since the branch point: the Epstein-Zin certainty-equivalent seam (#395), device-local continuation-V (#391), persistent-compilation-cache fix (#397), per-subject terminal rows (#396), and the CI/tooling bumps (ty prek hook, action/pixi pins, main merges). Conflict resolution: - src/lcm/__init__.py: keep both the NB-EGM case_piece exports and the new certainty_equivalent exports. - src/_lcm/regime_building/Q_and_F.py: take feat/dcegm's version — its #395 refactor relocated the continuation-operator logic into _lcm/certainty_equivalent.py, superseding nb-egm's inline copy. The MappingLeaf-payload concern nb-egm's deleted unit test guarded is covered by the Q-bundle threading contract and the solve-level nbegm_mappingleaf_threshold agreement tests. - tests/regime_building/test_continuation_operator.py: accept the deletion (replaced by tests/test_certainty_equivalent.py). Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com> Claude-Session: https://claude.ai/code/session_016PCdJtoqhhjBWhGAo7AXuz

hmgaudecker and others added 16 commits June 21, 2026 19:09

Type the HLO test helper's return as non-optional

f9a4248

Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>

Merge remote-tracking branch 'origin/main' into feat/type-local-conti…

dd94481

…nuation-v

Add design spec and implementation plan for the certainty-equivalent …

bcad4ed

…seam (#385) Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>

Add the certainty-equivalent user API (transform-pair seam)

72e4c63

Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>

Add Regime.certainty_equivalent with terminal- and DCEGM-rejection

7933ddc

Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>

Surface certainty-equivalent params in the regime params template

943427b

Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>

Inline the certainty-equivalent template block to keep the diff local

f841b05

Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>

Apply the certainty-equivalent transform around both continuation exp…

937bbd1

…ectations Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>

Pin Epstein-Zin solve and simulation against a numpy reference

0b8138b

Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>

Scope the template pseudo-function helper to the certainty equivalent

3a6aa75

Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>

Add an Epstein-Zin lifecycle example with docs

6f8dd73

Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>

Type the test model-builder kwargs so ty accepts the Regime unpacking

76a10bf

Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>

hmgaudecker and others added 4 commits July 2, 2026 18:39

Remove agent planning documents from the repo

99d1eff

Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>

Define the CE classes engine-side behind a thin public facade

b557155

Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>

hmgaudecker and others added 4 commits July 2, 2026 19:49

Label the example figure with both preference parameters, drop the EU…

d3f0f0d

… trace Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>

Fix the rotted regime column names in the tiny-example notebook

81165cd

Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>

Rename lcm.aggregators to lcm.temporal_aggregation

07d8b6f

Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>

hmgaudecker requested a review from mj023 July 2, 2026 20:37

hmgaudecker mentioned this pull request Jul 2, 2026

Emit one terminal-regime row per subject in to_dataframe #396

Merged

hmgaudecker linked an issue Jul 3, 2026 that may be closed by this pull request

ENH: Expose/generalize the aggregator so the certainty-equivalent sees the value distribution (non-EU preferences) #385

Closed

hmgaudecker and others added 3 commits July 3, 2026 08:56

hmgaudecker force-pushed the feat/certainty-equivalent branch from bd99589 to c0dab41 Compare July 3, 2026 09:49

mj023 approved these changes Jul 3, 2026

View reviewed changes

Base automatically changed from feat/type-local-continuation-v to main July 6, 2026 05:07

Merge remote-tracking branch 'origin/main' into feat/certainty-equiva…

1e78e86

…lent

hmgaudecker merged commit d1e84bd into main Jul 6, 2026
10 of 11 checks passed

hmgaudecker deleted the feat/certainty-equivalent branch July 6, 2026 05:13

hmgaudecker mentioned this pull request Aug 3, 2026

BUG Aggregate the continuation lottery stably under a certainty equivalent #413

Closed

Navigation Menu

Certainty-equivalent seam for non-EU preferences (Epstein-Zin) - #395

Certainty-equivalent seam for non-EU preferences (Epstein-Zin)#395
hmgaudecker merged 30 commits into
mainfrom
feat/certainty-equivalent

hmgaudecker commented Jul 2, 2026 •

edited

Loading

Uh oh!

read-the-docs-community Bot commented Jul 2, 2026 •

edited

Loading

Uh oh!

github-actions Bot commented Jul 2, 2026 •

edited

Loading

Uh oh!

review-notebook-app Bot commented Jul 2, 2026

Uh oh!

mj023 left a comment

Uh oh!

mj023 Jul 3, 2026

Uh oh!

mj023 Jul 3, 2026

Uh oh!

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

2 participants

Conversation

hmgaudecker commented Jul 2, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

What this PR does

User API

Example + tests

Non-goals

Coordination

Uh oh!

read-the-docs-community Bot commented Jul 2, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Documentation build overview

Uh oh!

github-actions Bot commented Jul 2, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Benchmark comparison (main → HEAD)

Uh oh!

review-notebook-app Bot commented Jul 2, 2026

Uh oh!

mj023 left a comment

Choose a reason for hiding this comment

Uh oh!

mj023 Jul 3, 2026

Choose a reason for hiding this comment

Uh oh!

mj023 Jul 3, 2026

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

2 participants

hmgaudecker commented Jul 2, 2026 •

edited

Loading

read-the-docs-community Bot commented Jul 2, 2026 •

edited

Loading

github-actions Bot commented Jul 2, 2026 •

edited

Loading