Introduce discrete state choice specific continuous grids

[MaxBlesch]

Overview

• Continuous grids can be supplied as a function instead of an array. Where
  model_config["continuous_states"] took a fixed grid, you can now pass a function of
  the discrete state under the new continuous_grid_functions argument and get a
  different grid per state. If you do decide to supply a function than the former grid
  entry must be None.
• Those grids are state-choice-specific, but choice-dependence is for now confined
  to the budget equation. Grids live on the state-choice space because the solution
  does; choice as an argument, however, is allowed only in budget_constraint.
• income_shock_previous_period → income_shock: the shock is part of this
  period's state. Continuous laws of motion are uniformly defined as end-of-last-period
  → this-period state, which also fixes what states_initial must contain.
• Sparsity is reintroduced when the budget equation does not depend on the current
  period's choice. State-choice space is denser than state space; if budget_constraint
  does not declare choice, the law of motion is evaluated on unique child states
  again and gathered back out.

1. Supplying a function instead of a grid

continuous_grid_functions is a new top-level argument to setup_model /
create_model_dict. It accepts an entry for assets_end_of_period,
assets_begin_of_period, or any additional continuous state, and that entry is a
function of the discrete state-choice evaluated once per state-choice at model build —
in place of the single global array that name would otherwise carry.

The function and the config entry are paired, and the pairing is enforced both ways. A
name left as an array in model_config["continuous_states"] and not given a function
keeps exactly the old behaviour, skipped by the new machinery entirely since a constant
grid is trivially consistent everywhere. A name given a function must be declared
None there instead — no default array at all, the grid comes only from the function —
and declaring None without supplying a function raises, as does supplying both. The
name must appear in continuous_states either way; the functions themselves go in the
top-level continuous_grid_functions argument, since model_config holds data and
never callables.

assets_end_of_period obeys the same rule as every other name. It briefly did not —
its length sizes the fues tuning params (n_constrained_points_to_add,
n_total_wealth_grid), which check_model_config.py computed eagerly, before
continuous_grid_functions is processed and before any state-choice exists to pin a
deferred length against. Those are now deferred too, resolved by
resolve_tuning_params_from_assets_end_of_period once the length is pinned. The upper
envelope, which reads them, is built after the model structure rather than alongside the
other model functions, so it is constructed from resolved sizes rather than closing over
pending ones. The same move dropped a dead compute_upper_envelope from the
alternative-simulation path, which nothing read.

None is what costs the code. Sizes that used to be known at config time —
n_continuous_state_combinations, n_total_wealth_grid — cannot be computed before a
real state-choice exists to evaluate the function against, so check_model_config.py
leaves them unresolved and they are merged back once the state-choice space is built.
The length is pinned by the first state-choice's evaluation and every other
state-choice's grid is validated against it: grids may differ in their values, never in
their size, so solution containers stay rectangular. Save/load needed the mirror image —
load_model_dict rebuilds model_config from the raw user config, unresolving those
sizes, while taking the state-choice space from the pickle, so the lengths are re-pinned
from one state-choice there.

New module: pre_processing/model_structure/continuous_state_grids.py.

2. State-choice-specific, and why choice is budget-only

Grids are indexed by state-choice, not by state, because that is where the solution
itself lives (value_solved / policy_solved / endog_grid_solved). A grid function
may therefore vary across the choices available in one state.

Declaring choice as an argument is a separate question, and the answer narrowed over
the course of the branch. An earlier iteration let any transition function take it; it
is now only budget_constraint, and _check_choice_not_declared raises at model
build for the others. The reason is timing: every other state variable is fixed before
the choice is made, so a choice argument there has no meaning. Wealth is the one
deliberate exception, and it is what makes choice-specific wages of the current choice available
for consumption. choice in
budget_constraint means the choice made in the period whose wealth is being computed.
next_period_continuous_state gets the previous choice the ordinary way, as the state
variable lagged_choice, as this is part of the states.

Because grids are per-state-choice while the continuation value is computed once per
child, something has to decide whose grid feeds the transition into a shared child.
That is the representative parent: any one of the parent state-choices mapping into
that child. Picking arbitrarily is legitimate because
check_continuous_grid_consistency_across_shared_children verifies once at model-build
time that all of them agree on their own grid — checked at build rather than defended
against at solve time.

3. income_shock, and the timing of continuous laws of motion

The shock is realized in the period whose wealth the budget equation computes, so it is
part of that period's state, not the previous one's. The old name
income_shock_previous_period described a habit of the shipped example models — income
earned against lagged_choice, so a retiring agent still collects last period's wage —
rather than anything the framework requires. The two non-state arguments of
budget_constraint are deliberately asymmetric in timing and the names now say so:
asset_end_of_previous_period really is carried over from t-1, income_shock is
realized in t.

Generalized: every continuous law of motion maps end-of-last-period quantities to a
this-period state. The consequence for simulation is that states_initial takes
assets_end_of_previous_period, not assets_begin_of_period — what agents carry
into period 0, not what they have to spend in it. Period 0 then runs the budget
equation like every other period and stops being a special case, and the simulated panel
becomes self-consistent row by row: applying your budget constraint to the previous
period's savings and this period's own income_shock column reproduces this period's
assets_begin_of_period.

Both old names now raise a pointed ValueError rather than being silently
reinterpreted or surfacing as a bare KeyError from the signature filter. Templates,
toy models and tests are updated. Simulation also computes beginning-of-period wealth
for every choice before the choice is made and gathers the realized one afterwards —
the counterpart to a choice-dependent budget equation — and
alternative_sim_functions recomputes choice-dependence from the alternative budget
equation rather than inheriting it.

4. Sparsity reintroduced when the budget is choice-independent

State-choice space is denser than state space, and evaluating the law of motion on it
does redundant work: unless the budget equation depends on the current choice, every
state-choice sharing a child state computes a bit-identical transition.

calc_law_of_motion is now the single entry point every caller uses, dispatching on one
model-static flag, budget_depends_on_choice:

Budget declares choice	Law of motion evaluated once per
yes	child state-choice
no (the common case)	unique child state, then gathered out

These are not two implementations of the transition math — the coarse branch calls the
fine one with deduplicated rows. Since the branch is known at setup, only that branch's
dedup arrays are bundled and threaded through the backward-induction scan; the other is
never built.

5. What fell out of the single entry point

Having one entry point forced two identities apart that had been conflated, and the
separation fixed a latent bug. A transition into a child now distinguishes:

Role	Supplies
the child, non-proxy, at its real state	the arguments to the transition functions
a representative parent state-choice	the continuous grid values fed into them
the proxy slot	the stored value/policy that is read back

Pulling the proxy apart from the child is behavioural, not plumbing. With a value-reuse
proxy — a sparsity condition redirecting every dead state to one last-period death slot,
say — the law of motion now runs on the child's real state, with its real age and
experience, while only the value/policy lookup follows the proxy to where that value was
stored. Previously the proxy identity leaked into the transition. The user-facing
corollary is now in the sparsity docs: a living parent re-evaluates the periodic
utility and marginal_utility at the proxied child's own identity during EGM, so those
must branch on the terminal state and agree with utility_functions_final_period, or
every parent with a dead child has a wrong Euler equation. (That branch fires for every
period except the last two, which bypass the periodic functions.)

Two cleanups come with it. The final period no longer runs a duplicated transition
path — it goes through calc_law_of_motion like everything else, with its own dedup
arrays built in the new batches/last_two_periods.py. And the batch dedup logic,
duplicated identically across both batch-size algorithms (algo_batch_size.py's
largest_block and single_segment.py's period_max), is one implementation in
batches/child_state_dedup.py that also exposes the state-level dedup section 4 needs.

6. Two additions riding along

dcegm.get_solve_from_small_models (new interfaces/combined_solve.py, exported
from dcegm/__init__.py) — for models whose pooled solve does not fit in GPU memory.
Split the model by hand into sub-models whose discrete state-choice spaces tile the big
one exactly, typically one per invariant type (education × sex) where the problem is
block-diagonal and each block is an exact standalone sub-problem. You get back a
solve(params) -> model_solved for the big model that solves each block and scatters it
into the big container, so the full solution is only ever stored in one piece, never
solved in one. Sequential by default, one block's device arrays live at a time;
parallel=True dispatches blocks round-robin across jax.devices(). The tiling is
verified to be an exact partition rather than assumed.

interpolation/interpnd_regular.py — N-dimensional regular-grid interpolation, for
reading stored solutions back on their own (possibly state-specific) regular grids under
Druedahl-Jorgensen.

Docs

Two new pages: background/choice_dependent_sparsity.rst (why sparsity_condition takes
no choice, and the recipe for each of the two distinct things people mean when they ask
for it) and development/internals/law_of_motion.rst (evaluation granularity, parent vs.
child vs. proxy, writing a choice-dependent budget equation). Updated:
sparsity_conditions.rst with the proxied-terminal-state utility requirement, and
practitioner_guide.rst with the initial-assets contract and choice-dependent budgets.
Both new pages carry an AI-drafted warning banner — worth deciding whether that stays on
merge, as is whether internals/state_specific_continuous_grids_plan.md, a 237-line
working plan, should ship at all.

Tests

Area	Tests
Config plumbing (continuous_grid_functions, None-declared grids), batch dedup consistency	test_state_specific_grids_config.py
Representative-parent grid selection; the two granularity branches bit-for-bit equivalent	test_state_specific_grids_law_of_motion.py
End-to-end via solve() / simulate() / choice_values_for_states	test_state_specific_grids_end_to_end.py
Absolute correctness vs. closed form and a hand-solved NumPy reference — incl. a second state-specific grid, and a choice-dependent budget over 4 periods with stochastic transitions	test_state_specific_grids_reference.py
Proxy/transition decoupling: closed-form Euler oracle under certain death; proxied vs. explicitly-solved death agree with an experience-dependent bequest	tests/sparse_death/test_death_proxy_euler.py, tests/sparse_death/test_sparse_death_period_grid.py
Initial-assets contract — period-0 wealth, shock attribution, old key rejected	test_initial_assets_contract.py
Sub-model tiling solver	test_combined_solve.py
N-D regular-grid interpolation	test_interpnd_regular_own_grids.py, test_interpolation.py
State-specific grids survive save/load	test_load_save.py
Final period / general regression	test_two_period_continuous_experience.py, test_law_of_motion.py

Breaking changes

1. budget_constraint: income_shock_previous_period → income_shock (old name raises).
2. states_initial: assets_begin_of_period → assets_end_of_previous_period (old key raises).
3. Any transition function other than budget_constraint declaring choice raises at model build.

[codecov]

Codecov Report

❌ Patch coverage is 95.52670% with 31 lines in your changes missing coverage. Please review.
✅ Project coverage is 93.86%. Comparing base (6914875) to head (334eea0).

Files with missing lines	Patch %	Lines
src/dcegm/interfaces/combined_solve.py	81.57%	21 Missing ⚠️
src/dcegm/egm/interpolate_marginal_utility.py	78.57%	3 Missing ⚠️
src/dcegm/interfaces/sol_interface.py	33.33%	2 Missing ⚠️
src/dcegm/interpolation/interp_interfaces.py	85.71%	2 Missing ⚠️
src/dcegm/interpolation/simulation_interp.py	92.30%	1 Missing ⚠️
...ocessing/model_structure/continuous_state_grids.py	98.83%	1 Missing ⚠️
src/dcegm/solve_single_period.py	92.30%	1 Missing ⚠️

Additional details and impacted files

@@            Coverage Diff             @@
##             main     #212      +/-   ##
==========================================
+ Coverage   93.07%   93.86%   +0.79%     
==========================================
  Files          71       74       +3     
  Lines        2598     3128     +530     
==========================================
+ Hits         2418     2936     +518     
- Misses        180      192      +12     

☔ View full report in Codecov by Harness.
📢 Have feedback on the report? Share it here.

🚀 New features to boost your workflow:

• ❄️ Test Analytics: Detect flaky tests, report on failures, and find test suite problems.

[review-notebook-app]

Check out this pull request on 

See visual diffs & provide feedback on Jupyter Notebooks.

---

Powered by ReviewNB