Lift discrete fixed states as partition dimensions

[hmgaudecker]

Summary

state_transitions[name] = None on a DiscreteGrid state has always meant "this value never changes." Today pylcm implements that by vmap'ing an auto-generated identity transition alongside every other state, paying memory linearly in the fixed state's cardinality. Because the Bellman backward induction never couples values across such a dimension, each fixed state partitions the state space into independent sub-problems: pylcm can compile one reduced sub-model and run it once per point in the product of all partition grids.

This PR adds that machinery end-to-end. User-facing API is unchanged — None transitions still declare a fixed state; internally pylcm lifts them out of variable_info / state-action space / Q_and_F, iterates the partition product in Model.solve / Model.simulate, and restores partition axes at the API boundary.

What changed

• src/lcm/regime_building/partitions.py (new). Shared helpers: detect_model_partitions, lift_partitions_from_regime, iterate_partition_points, inject_partition_scalars, stack_partition_V_arrays, group_subjects_by_partition, slice_V_at_partition_point, slice_initial_conditions. All handle the empty-grid case transparently (one iteration, no stacking, no grouping).
• src/lcm/regime_building/processing.py. process_regimes detects model-level partitions up-front and passes reduced regimes to every downstream builder. InternalRegime gains a partitions field; downstream machinery never sees partition states as vmap'd axes.
• src/lcm/model.py. Model._partition_grid aggregates partitions across regimes (validates category agreement). Model.solve wraps backward induction in a partition loop, stacks sub-V-arrays along trailing axes. Model.simulate groups subjects by their partition-value tuple, slices V-arrays per group, and concatenates PeriodRegimeSimulationData with globally-correct subject_ids.
• src/lcm/simulation/{simulate,result}.py. PeriodRegimeSimulationData gains an explicit subject_ids field; simulate() accepts it as a kwarg so per-partition dispatch groups stay aligned after concatenation. to_dataframe emits partition names as "state" columns with correct per-regime categorical mapping.
• src/lcm/simulation/initial_conditions.py. Partition names are required in initial_conditions — they supply each subject's fixed value.

Qualification rules for partition lifting

A discrete state qualifies as a partition iff:

1. The grid is a DiscreteGrid.
2. Every non-terminal regime where the state appears declares None.
3. The state appears in at least one non-terminal regime (target-only states populated by per-target transitions at the boundary are excluded).
4. Categories agree across regimes.

Terminal regimes (which must have empty state_transitions by validation) count the absence of an entry as the equivalent of None.

Continuous fixed states (None on a continuous grid) fall through to an identity transition — partitioning a continuous value is meaningless.

Mahler-Yum impact

Four fixed states (productivity, education, discount_type, health_type) auto-lift on this branch, giving a 32-point partition product.

• GPU peak memory (cuda13, 100 subjects): 501 MB. Comparable to or slightly below the pre-heterogeneity baseline (~400 MB on older main, 522 MB on the h-consumes-dag-outputs branch). The win is larger for Replication_MY2024's 10K-subject / 37-period runs where memory pressure is real.
• Execution time dominated by compile once + run N times: only one JIT compile across all partition points (shape invariant across points).

Test plan

• pixi run --environment tests-cpu tests -n 7 — 856 pass, 5 skipped (843 existing + 13 new partition tests).
• pixi run ty — clean.
• prek run --all-files — clean.
• Mahler-Yum GPU benchmark runs end-to-end on cuda13.
• test_dag_output_feeds_default_h_monotone_in_discount_factor updated: skip the last working-life period where V = U (deterministic transition to dead, V_dead = 0) is pref_type-independent regardless of discount factor.

Not in scope

• Removing _IdentityTransition entirely — continuous fixed states still need it. Narrower cleanup is a follow-up.
• Memory-optimal stacking via jax.lax.scan into a pre-allocated output buffer (current implementation holds all sub-V-arrays during stack). Peak could drop another ~30 %.

🤖 Generated with Claude Code

[read-the-docs-community]

Documentation build overview

📚 pylcm | 🛠️ Build #32322723 | 📁 Comparing 8c14948 against latest (719b533)

  🔍 Preview build  

32 files changed · ± 32 modified

± Modified

• index.html
• approximating-continuous-shocks/index.html
• benchmarking/index.html
• benchmarking-1/index.html
• beta-delta/index.html
• conventions/index.html
• debugging/index.html
• defining-models/index.html
• dispatchers/index.html
• function-representation/index.html
• and 22 more...

[hmgaudecker]

Code review

User asked for ALL issues, not just those ≥80. Scoring is from a follow-up Haiku pass against the rubric in the code-review skill (0 = false positive, 100 = certain & frequent).

Confidence 80+

1. Missing kw-only separator on detect_model_partitions (CLAUDE.md/MEMORY.md: "all function parameters should be keyword-only (*,) unless JAX tracing requires positional args") — score 100.
   https://github.com/OpenSourceEconomics/pylcm/blob/85e4d06568c5d10a4cd8d9fae58f44717ff2c2bf/src/lcm/regime_building/partitions.py#L36-L40

2. Missing kw-only separator on iterate_partition_points (same rule) — score 100.
   https://github.com/OpenSourceEconomics/pylcm/blob/85e4d06568c5d10a4cd8d9fae58f44717ff2c2bf/src/lcm/regime_building/partitions.py#L139-L143

3. Missing kw-only separator on _build_partition_grid (same rule) — score 100.
   https://github.com/OpenSourceEconomics/pylcm/blob/85e4d06568c5d10a4cd8d9fae58f44717ff2c2bf/src/lcm/model.py#L578-L581

Confidence 75 (highly likely real)

4. Same RNG seed passed to every partition group in Model.simulate — each sub-simulate calls jax.random.key(seed=seed) with the same integer. Stochastic state transitions apply identical draws across groups. Silent correctness bug when seed is not None and there are ≥2 partition groups with stochastic transitions. Fix: split key per group (jax.random.split(jax.random.key(seed), n_groups)) or hash (seed, partition_point) before passing.
   https://github.com/OpenSourceEconomics/pylcm/blob/85e4d06568c5d10a4cd8d9fae58f44717ff2c2bf/src/lcm/model.py#L381-L394

5. Test docstring non-imperative mood — _utility starts with "Utility scales…" (descriptive). AGENTS.md: "Imperative mood for docstring summary lines."
   https://github.com/OpenSourceEconomics/pylcm/blob/85e4d06568c5d10a4cd8d9fae58f44717ff2c2bf/tests/regime_building/test_partitions.py#L53-L62

Confidence 50 (real but edge-case or stylistic)

6. Shadowed partition_names in process_regimes — outer model-level frozenset is re-bound inside the per-regime loop. Functionally correct today (outer binding isn't read after the loop) but a refactoring trap.
   https://github.com/OpenSourceEconomics/pylcm/blob/85e4d06568c5d10a4cd8d9fae58f44717ff2c2bf/src/lcm/regime_building/processing.py#L100-L152

7. _merge_sub_simulation_results fast-path shape mismatch for single-point partitions — when len(sub_results) == 1 AND partition_grid is non-empty, the returned period_to_regime_to_V_arr has the partition axis appended but raw_results V_arrays do not. Mismatch visible if downstream jointly reads both.
   https://github.com/OpenSourceEconomics/pylcm/blob/85e4d06568c5d10a4cd8d9fae58f44717ff2c2bf/src/lcm/model.py#L497-L520

8. Merge does not assert matching state_names across slices — state_names = tuple(slices[0].states) silently drops keys present in later slices but missing from the first.
   https://github.com/OpenSourceEconomics/pylcm/blob/85e4d06568c5d10a4cd8d9fae58f44717ff2c2bf/src/lcm/model.py#L540-L563

Confidence 25 (plausible, unverified)

9. Module layout: detect_model_partitions / lift_partitions_from_regime precede the primary-API helpers named in the module docstring. .ai-instructions/AGENTS.md: "Write 'deep' modules: important public function(s) at the top, private helpers below."
   https://github.com/OpenSourceEconomics/pylcm/blob/85e4d06568c5d10a4cd8d9fae58f44717ff2c2bf/src/lcm/regime_building/partitions.py#L36-L165

10. subject_id semantic change — previously jnp.arange(len(result.in_regime)) (dense local row index), now result.subject_ids (global ID). Design is cleaner but a silent break for downstream consumers assuming 0..n-1 density. No in-repo consumer relies on density.
    https://github.com/OpenSourceEconomics/pylcm/blob/85e4d06568c5d10a4cd8d9fae58f44717ff2c2bf/src/lcm/simulation/result.py#L526-L539

11. Partition-state reachability gap in _validate_all_variables_used — a partition state consumed only as an index into a params-dict/Series (not as a direct function arg) may not be reachable through the get_ancestors walk. No test covers this.
    https://github.com/OpenSourceEconomics/pylcm/blob/85e4d06568c5d10a4cd8d9fae58f44717ff2c2bf/src/lcm/model_processing.py#L195-L260

Confidence 0 (false positive / pre-existing, documented for completeness)

12. Axis-order invariant concern (stacked V has continuous axes non-last): Model.simulate always calls slice_V_at_partition_point before handing V to the interpolator, so the invariant is preserved at the consumption point. (Score 0, but worth a comment at slice_V_at_partition_point documenting the invariant.)
13. PR Support heterogeneous state sets in initial conditions #315 NaN→int32 cast: not reintroduced.
14. PR Support heterogeneous state sets in initial conditions #315 wrong exception type in _raise_feasibility_type_error: not touched.
15. PR Skip unreachable target regimes in continuation value loop #316 io_callback per scalar: pre-existing, Q_and_F not modified.
16. PR Auto-convert pd.Series in fixed_params #308 helper-inlining readability: partition logic is in dedicated partitions.py; only a thin wrapper stays in model.py.

🤖 Generated with Claude Code

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

[hmgaudecker]

Closing to reopen as a clean single-commit stack; replaced by the new PR in the same slot. No content change.

[hmgaudecker]

Reopened as #329 (clean single-commit stack).