Fix outcome axis for cross-grid stochastic transitions by hmgaudecker · Pull Request #321 · OpenSourceEconomics/pylcm

hmgaudecker · 2026-04-10T15:54:04Z

Summary

For per-target stochastic transitions crossing grid sizes (e.g. 3-state health → 2-state health), _build_outcome_mapping now uses the target regime's grid for the outcome axis instead of the source regime's grid
Fixes shape mismatch (ValueError: Axis must be specified when shapes of a and weights differ) in Q_and_F.py:217 during solve

Root cause

_build_outcome_mapping extracted the state name from the qualified function name (next_health__post65 → health) but looked it up in the source regime's grids. For cross-grid transitions, the source grid has more categories than the target, producing an array with the wrong last dimension.

Test plan

New test test_convert_series_cross_grid_transition — 3-state source → 2-state target, verifies (n_ages, 3, 2) shape
All 824 existing tests pass (pytest -n 7)

🤖 Generated with Claude Code

Shock grid states (Rouwenhorst, Uniform, etc.) are continuous states that should be accepted as float columns in the DataFrame. Previously they were rejected as "unknown columns" because _collect_all_state_names excluded _ShockGrid instances. Split state name collection into required (non-shock states + age) and optional (shock grid states). Shock columns are accepted but not required, since the model draws fresh shock values each period. Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>

The type-assignment loop set discount type for indices 0–7 (low education) but not for indices 8–15 (high education). All high-education agents were incorrectly assigned to the low discount factor type. Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>

fixed_params passed pd.Series raw to functools.partial, causing JAX TypeError during tracing. Runtime params already auto-convert via _maybe_convert_series. Add the same conversion to the fixed_params path using a callback pattern to avoid circular imports between model.py and model_processing.py. Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>

…arams

For per-target stochastic transitions that cross grid sizes (e.g. 3-state health → 2-state health), the outcome axis must use the target regime's grid, not the source's. Without this fix, the converted transition probability array has the wrong last dimension, causing a shape mismatch in jnp.average during Q_and_F evaluation. Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>

read-the-docs-community · 2026-04-10T15:56:33Z

Documentation build overview

📚 pylcm | 🛠️ Build #32254279 | 📁 Comparing 8858f58 against latest (303df59)

🔍 Preview build

Show files changed (32 files in total): 📝 32 modified | ➕ 0 added | ➖ 0 deleted

File	Status
index.html	📝 modified
approximating-continuous-shocks/index.html	📝 modified
benchmarking/index.html	📝 modified
benchmarking-1/index.html	📝 modified
beta-delta/index.html	📝 modified
conventions/index.html	📝 modified
debugging/index.html	📝 modified
defining-models/index.html	📝 modified
dispatchers/index.html	📝 modified
function-representation/index.html	📝 modified
grids/index.html	📝 modified
index-1/index.html	📝 modified
index-2/index.html	📝 modified
index-3/index.html	📝 modified
index-4/index.html	📝 modified
installation/index.html	📝 modified
interpolation/index.html	📝 modified
mahler-yum-2024/index.html	📝 modified
mortality/index.html	📝 modified
pandas-interop/index.html	📝 modified
parameters/index.html	📝 modified
precautionary-savings/index.html	📝 modified
precautionary-savings-health/index.html	📝 modified
regimes/index.html	📝 modified
setup/index.html	📝 modified
shocks/index.html	📝 modified
solving-and-simulating/index.html	📝 modified
stochastic-transitions/index.html	📝 modified
tiny/index.html	📝 modified
tiny-example/index.html	📝 modified
transitions/index.html	📝 modified
write-economics/index.html	📝 modified

Shock states were made optional in the initial commit, but this is wrong: AR(1) shocks depend on the current value for transitions, and observed persistent shocks (e.g., wage residuals) represent real data. Making them optional would silently fill with NaN. Simplify _collect_state_names to return a single set including all states (shocks and non-shocks alike). This is consistent with validate_initial_conditions in simulate(), which already requires them. Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>

…arams

…-mapping

…tion - Remove decorative section-separator comment in test file - Remove unused _PROBS_ARRAY constant - Add type annotations to _make_markov_model - Add Returns section to build_regimes_and_template docstring - Add _validate_param_types after Series conversion in fixed_params callback, matching the runtime params validation path Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>

…-mapping

Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>

…-mapping

hmgaudecker · 2026-04-10T17:19:07Z

Code review

Found 3 issues (only changes relative to PR #308):

_Rid naming inconsistency. New test defines class _Rid (lowercase d) but all other test-local regime ID classes in the file use _RId (capital D) — see lines 527, 1435, 1525, 1607, 1698. CLAUDE.md says "Consistent naming across a file. When multiple functions in the same file use the same concept, use the same parameter name everywhere." (Score: 75)

pylcm/tests/test_pandas_utils.py

Lines 447 to 449 in ee49b73

    
           ages = model.ages.exact_values 
        
           sr = pd.Series([1.0, 2.0, 3.0, 4.0], index=pd.Index(ages, name="age")) 
        
           # Should not raise despite heterogeneous health grids

_build_outcome_mapping docstring not updated. The docstring lists two cases (state transitions and regime transitions) but the PR adds a third case (per-target transitions using target regime's grid). The inline comment is helpful but the function-level docstring should also describe this behavior. (Score: 50)

pylcm/src/lcm/pandas_utils.py

Lines 656 to 668 in ee49b73

    
               Returns: 
        
                   `_LevelMapping` for the outcome (last) axis. 
        
               """ 
        
               if func_name == "next_regime": 
        
                   regime_ids = dict(model.regime_names_to_ids) 
        
                   return _LevelMapping( 
        
                       name="next_regime", 
        
                       size=len(regime_ids), 
        
                       get_code_from_label=regime_ids.__getitem__, 
        
                       valid_labels=tuple(regime_ids), 
        
                   )

ty failure: arr.shape on int | float | Array union. Test accesses .shape on a value typed as int | float | Array. Needs type narrowing (e.g., assert isinstance(arr, Array)) or # ty: ignore[unresolved-attribute]. (Score: 75)

https://github.com/OpenSourceEconomics/pylcm/blob/ee49b734bc11c239b2f311a4749f6f79a97650f4/tests/test_pandas_utils.py#L1768-L1770

🤖 Generated with Claude Code

_{- If this code review was useful, please react with 👍. Otherwise, react with 👎.}

- Update _build_outcome_mapping docstring to describe per-target transition case (target regime's grid for outcome axis) - Add ty: ignore for arr.shape on union type in cross-grid test Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>

…alls The test compared two stochastic simulations without fixing the random seed, causing different MarkovTransition draws on macOS and Windows. Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>

…-mapping

github-actions · 2026-04-10T20:03:17Z

Benchmark comparison (main → HEAD)

Comparing 303df59b (main) → 8858f583 (HEAD)

Benchmark	Statistic	before	after	Ratio
Mahler-Yum	execution time	3.64±0.01s	3.63±0.02s	1.00
	peak GPU mem	262M	262M	1.00
	compilation time	1.96m	1.98m	1.01
	peak CPU mem	2.24G	2.24G	1.00
Mortality	execution time	251±5ms	248±8ms	0.99
	peak GPU mem	542M	542M	1.00
	compilation time	10.7s	10.6s	1.00
	peak CPU mem	1.24G	1.25G	1.01
Precautionary Savings - Solve	execution time	32.9±2ms	34.3±3ms	1.04
	peak GPU mem	8.44M	8.44M	1.00
	compilation time	5.09s	5.08s	1.00
	peak CPU mem	1.07G	1.07G	1.00
Precautionary Savings - Simulate	execution time	142±2ms	146±3ms	1.03
	peak GPU mem	138M	138M	1.00
	compilation time	7.06s	7.15s	1.01
	peak CPU mem	1.2G	1.2G	1.00
Precautionary Savings - Solve & Simulate	execution time	153±0.6ms	151±0.6ms	0.99
	peak GPU mem	565M	565M	1.00
	compilation time	11.3s	11.5s	1.01
	peak CPU mem	1.21G	1.21G	1.00
Precautionary Savings - Solve & Simulate (irreg)	execution time	291±4ms	297±2ms	1.02
	peak GPU mem	2.18G	2.18G	1.00
	compilation time	12.1s	12.2s	1.01
	peak CPU mem	1.27G	1.27G	1.00

…arams

…-mapping

- Change convert_series_in_params, array_from_series, and internal helpers to take regimes + ages + regime_names_to_ids instead of model: Model. This breaks the circular import between pandas_utils and model. - Change initial_conditions_from_dataframe to take regimes + regime_names_to_ids instead of model: Model (public API change). - Move _validate_param_types and _check_leaf to model_processing.py. - Remove convert_fixed_params callback from build_regimes_and_template; call convert_series_in_params and _validate_param_types directly. - Remove TYPE_CHECKING import of Model from pandas_utils.py. Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>

…-mapping # Conflicts: # src/lcm/pandas_utils.py

…sion The parameter is needed when fixed_params contain pd.Series indexed by derived categorical outputs (DAG function results like is_married). Without it, convert_series_in_params cannot resolve those indices. Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>

…-mapping

- Add missing `regimes` param to `initial_conditions_from_dataframe` docstring - Fix "Mapping" -> "Immutable mapping" for `regime_names_to_ids` in `build_regimes_and_template` docstring - Add test exercising fixed_params with pd.Series indexed by a derived categorical (DAG function output not in model grids) Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>

…-mapping

derived_categoricals describes outcome spaces of DAG function outputs — it belongs on Regime where the functions live. Model-level entries are broadcast to all regimes at init time (convenience sugar). Raises on conflict if a regime already has a different grid for the same key. - Add derived_categoricals field to Regime (Mapping[str, DiscreteGrid]) - Remove from solve() and simulate() signatures - Simplify pandas_utils chain: remove _resolve_categorical_entry, read derived_categoricals from regime directly - Add broadcasting tests (merge, match, conflict, coexistence) - Update docs and AGENTS.md Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>

…-mapping

- derived_categoricals accepts both categorical classes and DiscreteGrid; raw classes are normalized to DiscreteGrid in Regime.__post_init__ - Inline _maybe_convert_series and _maybe_convert_dataframe into solve() and simulate() — they were trivial one-liner wrappers - Remove max_compilation_workers and enable_jit from solve() calls (leaked from downstream PRs, not on this branch's solve_brute.solve) - Update docs and error messages to show raw class syntax Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>

…-mapping

- Revert derived_categoricals type to Mapping[str, DiscreteGrid] (raw categorical class acceptance was premature) - Remove InvalidValueFunctionError try/except blocks that leaked from improve/lazy-diagnostics via stash pop - Remove unnecessary dict() conversion in _build_outcome_mapping - Restore DiscreteGrid(...) in docs and tests Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>

…teless - Add Model._process_params(): broadcast, convert Series, validate — replaces duplicated 6-line blocks in solve() and simulate() - Extract _apply_fixed_params() from build_regimes_and_template so the parent function has no variable reassignment Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>

…riable - _build_regimes_and_template_with_fixed_params calls process_regimes and create_params_template internally, so build_regimes_and_template has no variable reassignment - Rename invalid -> invalid_regimes for symmetry with valid_regimes Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>

Consistent parameter order: ages, regimes, regime_names_to_ids across all def sites, call sites, and docstrings in model.py, model_processing.py, and pandas_utils.py. Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>

Both _resolve_categoricals and _build_discrete_grid_lookup now iterate regime.states and regime.actions symmetrically. Delete the redundant _build_discrete_action_lookup helper. Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>

…rams `any(v for v in fixed_internal.values())` is False for params like {"a": 0} or {"a": 0.0}, silently skipping partialling. Remove the check entirely — _partial_fixed_params_into_regimes handles empty dicts correctly. Also rename internal_regimes/params_template to raw_* for clarity. Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>

Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>

- Fix enable_jit docstring: "jit" -> "JIT-compile", singular -> plural - Fix array_from_series regime_name arg: mention derived categorical lookup - Move import pytest to top-level in test_static_params.py - Include derived_categoricals in _resolve_categoricals else-branch (regime_name=None); add test Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>

…-mapping # Conflicts: # tests/test_pandas_utils.py

…e-mapping # Conflicts: # src/lcm/pandas_utils.py # tests/test_pandas_utils.py

Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>

mj023

Looks correct, nothing from my side.

hmgaudecker and others added 8 commits April 8, 2026 19:28

Merge branch 'fix-dataframe-shock-columns' into fix-series-in-fixed-p…

2365d5c

…arams

Merge branch 'main' into fix-dataframe-shock-columns

6d7fb96

Merge branch 'fix-dataframe-shock-columns' into fix-series-in-fixed-p…

44d092e

…arams

hmgaudecker and others added 7 commits April 10, 2026 18:08

Merge branch 'fix-dataframe-shock-columns' into fix-series-in-fixed-p…

16d9e65

…arams

Merge branch 'fix-series-in-fixed-params' into fix/cross-grid-outcome…

44b1926

…-mapping

Merge branch 'fix-series-in-fixed-params' into fix/cross-grid-outcome…

0d150d2

…-mapping

Fix ty: use to_numpy() instead of .values for type-safe array comparison

eb581a8

Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>

Merge branch 'fix-series-in-fixed-params' into fix/cross-grid-outcome…

ee49b73

…-mapping

hmgaudecker and others added 3 commits April 10, 2026 19:33

Merge branch 'fix-series-in-fixed-params' into fix/cross-grid-outcome…

f06be2f

…-mapping

hmgaudecker mentioned this pull request Apr 10, 2026

Skip unreachable target regimes in continuation value loop #316

Merged

3 tasks

hmgaudecker changed the title ~~Fix cross-grid outcome axis in Series→array conversion~~ Fix outcome axis for cross-grid stochastic transitions Apr 10, 2026

hmgaudecker and others added 7 commits April 12, 2026 16:54

Merge remote-tracking branch 'origin/main' into fix-series-in-fixed-p…

0638b55

…arams

Merge branch 'fix-series-in-fixed-params' into fix/cross-grid-outcome…

42022aa

…-mapping

Merge branch 'fix-series-in-fixed-params' into fix/cross-grid-outcome…

d4a9586

…-mapping # Conflicts: # src/lcm/pandas_utils.py

Merge branch 'fix-series-in-fixed-params' into fix/cross-grid-outcome…

80a659b

…-mapping

hmgaudecker and others added 14 commits April 13, 2026 07:52

Merge branch 'fix-series-in-fixed-params' into fix/cross-grid-outcome…

e81e85e

…-mapping

Merge branch 'fix-series-in-fixed-params' into fix/cross-grid-outcome…

1eb2ba8

…-mapping

Merge branch 'fix-series-in-fixed-params' into fix/cross-grid-outcome…

839754e

…-mapping

Rename _remove_fixed_from_template -> _remove_fixed_params_from_template

4f763cc

Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>

Merge branch 'fix-series-in-fixed-params' into fix/cross-grid-outcome…

19fc695

…-mapping # Conflicts: # tests/test_pandas_utils.py

Base automatically changed from fix-series-in-fixed-params to main April 14, 2026 10:27

hmgaudecker and others added 2 commits April 14, 2026 12:30

Merge remote-tracking branch 'origin/main' into fix/cross-grid-outcom…

a9266ab

…e-mapping # Conflicts: # src/lcm/pandas_utils.py # tests/test_pandas_utils.py

Fix indentation of Returns block in _build_outcome_mapping docstring

8858f58

Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>

hmgaudecker requested a review from mj023 April 14, 2026 10:50

mj023 approved these changes Apr 14, 2026

View reviewed changes

hmgaudecker merged commit 49cb0e5 into main Apr 15, 2026
10 checks passed

hmgaudecker deleted the fix/cross-grid-outcome-mapping branch April 15, 2026 04:40

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

Fix outcome axis for cross-grid stochastic transitions#321

Fix outcome axis for cross-grid stochastic transitions#321
hmgaudecker merged 41 commits into
mainfrom
fix/cross-grid-outcome-mapping

hmgaudecker commented Apr 10, 2026

Uh oh!

read-the-docs-community Bot commented Apr 10, 2026 •

edited

Loading

Uh oh!

hmgaudecker commented Apr 10, 2026

Uh oh!

github-actions Bot commented Apr 10, 2026 •

edited

Loading

Uh oh!

mj023 left a comment

Uh oh!

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

2 participants

Conversation

hmgaudecker commented Apr 10, 2026

Summary

Root cause

Test plan

Uh oh!

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

Uh oh!

Documentation build overview

Uh oh!

hmgaudecker commented Apr 10, 2026

Code review

Uh oh!

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

Uh oh!

Benchmark comparison (main → HEAD)

Uh oh!

mj023 left a comment

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

2 participants

read-the-docs-community Bot commented Apr 10, 2026 •

edited

Loading

github-actions Bot commented Apr 10, 2026 •

edited

Loading