Age specialization: per-age regime functions/constraints and age-dependent continuous-state grids

DESIGN-age-dependent-grids.md

@@ -0,0 +1,146 @@
+# Design: age-dependent continuous-state grids (shape-invariant)
+
+Branch `feat/age-dependent-grids`, on top of `feat/age-specialized`.
+
+## Motivation
+
+A continuous state whose grid **bounds move with age** — the canonical case is an
+asset state with an age-dependent borrowing floor `a̲(age)` (Guvenen & Smith 2014,
+and any life-cycle model with an age-varying constraint). On a single fixed grid the
+floor must span the loosest bound, so at ages where the bound is tighter the grid
+cells below it are **unreachable and infeasible** (empty feasible set → `-inf`),
+which then poisons the value function by interpolation and propagates backward. An
+age-dependent grid whose floor tracks `a̲(age)` removes those cells entirely.
+
+## Key scoping decision: shape-invariant (fixed `n_points`)
+
+The grid may vary its **bounds/nodes** per age but keeps **the same `n_points`
+(and grid class) at every age**. This is exactly what a moving floor needs, and it
+is the decisive simplifier:
+
+- Every period's V-array keeps the **same shape**, so the rolling-continuation
+  topology, the V-array shapes, and the solver-kernel dedup are **unchanged**.
+- Only two things become period-aware: the **node values** on the current period's
+  state axis, and the **interpolation coordinates** for the continuation (which must
+  use the *next* period's grid).
+
+Fully arbitrary per-period `n_points` would force period-indexed V shapes and a
+period-indexed rolling template — a much larger change — and is **out of scope for
+v1**. (It can be added later; nothing here precludes it.)
+
+## API
+
+Reuse the `AgeSpecialized` marker, now also accepted as a **grid** in
+`Regime.states` (its `build(age)` returns a `ContinuousGrid` instead of a function):
+
+```python
+from lcm import AgeSpecialized, LinSpacedGrid
+
+states = {
+    "assets": AgeSpecialized(
+        build=lambda age: LinSpacedGrid(start=a_bar(age), stop=A_MAX, n_points=40),
+        signature=lambda age: a_bar(age),   # ages with equal floor share a program
+    ),
+    ...,
+}
+```
+
+Contract (validated at `Regime`/`Model` construction):
+- every `build(age)` returns the **same grid class** and the **same `n_points`**;
+- allowed for **continuous** states (Lin/Log/Irreg/Piecewise); process states and
+  discrete states are not age-varying in v1;
+- allowed in non-terminal and terminal regimes (terminal resolves at its single age);
+- **actions** stay fixed-grid in v1 (only states vary) — a continuation-axis is a
+  state, so this covers the borrowing-floor use case.
+
+`signature(age)` is a correctness precondition exactly as for functions: equal
+signature => identical grid.
+
+## Threading (reusing the per-period build loops)
+
+The per-period **function** builders already iterate periods with `age` in scope
+(`processing.py:_build_Q_and_F_per_period`, `_build_next_state_vmapped`,
+`_build_argmax_and_max_Q_over_a_per_period`). Age-dependent grids add per-period
+plumbing at these points:
+
+1. **Grid resolution.** A new resolver (mirroring `age_specialization.resolve_tree`)
+   turns the `states` mapping into a concrete per-period `{name: Grid}` for a given
+   age. Validate the shape-invariance across the horizon once.
+
+2. **Current-period state axis (values).** `create_state_action_space` /
+   `_build_base_state_action_spaces` resolve the age-varying state's grid at the
+   period's age -> period-t node values feed period-t economic functions. Shape is
+   identical across periods, so the V-array topology and rolling template are
+   untouched; only the axis *values* differ. `backward_induction.solve` already
+   calls the kernel per `(regime, period)`, so it passes period-t's axis.
+
+3. **Continuation interpolation (the crux).** In `get_Q_and_F` /
+   `_get_coordinate_finder` (V.py), the interpolation of `V_{t+1}` must use the
+   **target regime's period-(t+1) grid**. Since `_build_Q_and_F_per_period` builds
+   one program per period-group with `age` known, pass the target's
+   period-(t+1) `VInterpolationInfo`. Extend the group signature to fold in the
+   `(period, period+1)` grid identities so periods with different bounds do not
+   false-share a compiled `Q_and_F`.
+
+4. **Simulation.** The per-period argmax kernel and `_lookup_values_from_indices`
+   use the period's resolved grid for the argmax axis and index->value lookup.
+
+5. **Unchanged (shape-invariance):** V-array shapes/shardings, the zero-continuation
+   template and its roll, and `GridSearch` kernel dedup topology.
+
+No change to `coordinates.py` (the coordinate math already runs off
+`(start, stop, n_points)` / `points`).
+
+## Validation / errors
+
+- `n_points` or grid class differing across ages -> `RegimeInitializationError`.
+- age-varying **process** state or **discrete** state -> rejected (v1).
+- `to_dataframe(additional_targets=...)` on an age-varying-grid-dependent target ->
+  rejected, mirroring the AgeSpecialized rule (published functions use one
+  representative age).
+
+## Test plan
+
+- unit: resolver returns the right grid per age; shape-invariance validator;
+  signature dedup.
+- solve: a 2-state life-cycle model with an age-shrinking floor solves with
+  **finite V at every age** where a fixed grid produces `-inf`; compare against the
+  free-assets reformulation (same economics) — value functions must match to
+  interpolation tolerance.
+- simulate: continuous state stays within `[a̲(age), A_MAX]`; policy sensible.
+- regression: an age-invariant `signature` reproduces the plain fixed-grid solve
+  bit-for-bit (dedup collapses to one program).
+
+## Implementation status / remaining plan
+
+**Done (committed):** merge of origin/main into feat/age-specialized (conflicts
+resolved); `_AgeSpecialized` base + rename `AgeSpecialized`→`AgeSpecializedFunction`
++ new `AgeSpecializedGrid`; resolver/validator helpers in `age_specialization.py`
+(`resolve_state_grids`, `state_grids_signature`, `has_age_specialized_grid`,
+`validate_age_specialized_grids`). All 23 age-related tests green.
+
+**Key simplification confirmed:** the current-period state axis is passed to the
+Q_and_F kernel as *runtime* values (shape-invariant ⇒ same trace), so only the
+*continuation* (target) grid is a trace-time constant. Therefore:
+- state axis → resolve per period in the solve loop (runtime values, no recompile);
+- continuation interp → bake period-(t+1) grid into period-t's Q_and_F (per-period
+  kernel, grouping signature extended with target grid identity at t+1).
+
+**Remaining wiring (solve, then simulate, then tests):**
+1. `lcm/regime.py`: accept `AgeSpecializedGrid` in `states` (typing + treat as a
+   continuous state in construction validation; concrete resolution deferred).
+2. `model_processing.py`/`model.py`: call `validate_age_specialized_grids(states,
+   ages)` at Model construction.
+3. `processing.process_regimes`: build representative-resolved regimes (markers →
+   `build(representative_age)`) for the invariant machinery (from_regime, get_grids,
+   create_v_interpolation_info, create_state_action_space) so they are unchanged;
+   build `period_to_regime_to_v_interp` (resolved per period) for the continuation.
+4. `_build_Q_and_F_per_period`: extend the group signature with the target regimes'
+   `state_grids_signature` at period+1; pass the period-(t+1) VInterpolationInfo.
+5. `backward_induction`: per-period state axis (resolve age-varying states' nodes at
+   the period's age; pass as runtime values). SolutionPhase gains an optional
+   per-period grid table.
+6. `simulation`: per-period argmax axis + index→value lookup use the period's grid.
+7. Tests: solve finite-V where a fixed grid gives -inf; equivalence to the
+   free-assets reformulation; age-invariant signature reproduces the fixed-grid
+   solve bit-for-bit; shape-invariance + placement rejections.

docs/myst.yml

@@ -42,6 +42,7 @@ project:
         - file: user_guide/regimes.ipynb
         - file: user_guide/grids.md
         - file: user_guide/transitions.ipynb
+        - file: user_guide/age_specialized.md
         - file: user_guide/continuous_stochastic_processes.md
         - file: user_guide/parameters.md
         - file: user_guide/pandas_interop.md

docs/user_guide/age_specialized.md

@@ -0,0 +1,118 @@
+---
+title: Age-specialized functions
+---
+
+# Age-specialized functions (`AgeSpecializedFunction`)
+
+Some model functions do not just *read* the agent's age — their whole definition changes
+with it. The canonical case is a tax-transfer system pinned to a policy date: as a
+cohort ages through calendar time, the law that applies to it changes, so the function
+computing net income at age 58 is a *different closure* from the one at age 63, not the
+same closure with a different `age` argument.
+
+`AgeSpecializedFunction` marks such a function. At model build, pylcm resolves the
+marker **per period**: each period's age gets its own concrete function, compiled into
+that period's programs. There is no runtime dispatch on age or calendar year.
+
+```python
+from lcm import AgeSpecializedFunction
+
+
+def make_net_income(age: float):
+    """Build the net-income closure for the policy year this age falls in."""
+    policy_env = load_policy_environment(year_for(age))
+
+    def net_income(gross_income):
+        return policy_env.apply(gross_income)
+
+    return net_income
+
+
+def policy_key(age: float):
+    """Identify the closure: ages in the same policy year share a program."""
+    return year_for(age)
+
+
+functions = {
+    "net_income": AgeSpecializedFunction(build=make_net_income, signature=policy_key),
+    ...,
+}
+```
+
+## The two contracts
+
+`AgeSpecializedFunction(build, signature)` places two obligations on the user; pylcm
+cannot verify either automatically.
+
+1. **Invariant call signature.** Every concrete function `build(age)` returns must
+   expose the *same* argument list at every age. Only the constants the closure binds
+   may differ. pylcm's static passes (the parameter template, used-variable validation)
+   read the function at one representative age and rely on this.
+
+1. **`signature` is a correctness precondition, not a performance hint.** Periods whose
+   `signature(age)` values are equal share a single compiled program. An equal signature
+   must therefore imply *identical closure behavior* — same policy date, same price
+   level, same overrides, same everything the closure binds. If two ages collide to one
+   signature but `build` returns different closures, the solve is silently wrong. When
+   in doubt, put more into the signature (a tuple of every varying ingredient), never
+   less.
+
+Deduplication is what keeps this affordable: an age-invariant `signature` collapses to
+one program for the whole horizon (identical to not using the wrapper), and a
+policy-date signature compiles one program per distinct policy year, not per period.
+
+## Where it is allowed
+
+`AgeSpecializedFunction` is accepted in `functions` and `constraints` of
+**non-terminal** regimes. Everything else is rejected at `Regime` construction, loudly:
+
+- **the regime `transition`**, bare or wrapped in `MarkovTransition` — periodizing the
+  regime-transition probability path is not supported;
+- **a regime transition that _reads_ an `AgeSpecializedFunction` function**, directly or
+  through plain helper functions — regime-transition probabilities are built once, so a
+  policy-specialized value flowing into them would reuse one age's closure everywhere;
+- **`MarkovTransition(AgeSpecializedFunction(...))`** — policy-specialized stochastic
+  transitions are not supported;
+- **a direct `state_transitions` value** — see the pattern below instead;
+- **anything in a terminal regime** — the terminal value program is built once and
+  shared across all periods.
+
+## Policy-dependent laws of motion
+
+A state transition whose content depends on the policy year is written as a **plain
+transition reading an age-specialized helper function**:
+
+```python
+functions = {
+    "new_pension_points": AgeSpecializedFunction(build=make_points, signature=policy_key),
+    ...,
+}
+
+state_transitions = {
+    # A plain function — the policy content lives in the helper it reads.
+    "pension_points": lambda pension_points, new_pension_points: (
+        pension_points + new_pension_points
+    ),
+}
+```
+
+The simulate-phase next-state programs are built per period, so the helper resolves to
+the *current period's* closure inside every transition — the law of motion tracks the
+policy year without the transition itself carrying a marker.
+
+## What is resolved when
+
+- **Per period (exact):** the solve `Q_and_F` programs, the simulate decision programs,
+  and the simulate next-state programs. These are what determine model behavior, and
+  they always use each period's own closure.
+- **At one representative age (the regime's first active age):** the parameter template,
+  used-variable validation, and the *published* function sets
+  (`regime.solution.functions`, `regime.simulation.functions`). These consumers only
+  need the (age-invariant) call signature — except one:
+
+`to_dataframe(additional_targets=...)` computes targets from the published simulation
+functions. A target that depends on an age-specialized function would therefore be
+evaluated under the representative age's policy, not each row's period — so pylcm
+**rejects** such targets with `InvalidAdditionalTargetsError`. Quantities you need per
+period from a specialized function should be carried inside the model (for example as a
+state fed by a plain transition reading the helper).

src/_lcm/engine.py

@@ -333,6 +333,17 @@ class SolutionPhase:
     _base_state_action_space: StateActionSpace = dataclasses.field(repr=False)
     """Base state-action space before runtime grid substitution."""
 
+    period_state_axes: (
+        MappingProxyType[int, MappingProxyType[StateOrActionName, object]] | None
+    ) = None
+    """Per-period node arrays for age-varying (`AgeSpecializedGrid`) states.
+
+    `{period: {state_name: nodes}}` — the current period's grid nodes for each
+    age-varying continuous state, used by backward induction to override the
+    (representative) base axis so period `t`'s value function is tabulated on
+    period `t`'s grid. `None` for age-invariant regimes (the base axis is used
+    unchanged)."""
+
     @property
     def state_names(self) -> tuple[StateOrActionName, ...]:
         """Solve-phase state names in canonical (productmap) order."""
@@ -485,8 +496,17 @@ class SimulationPhase:
     argmax_and_max_Q_over_a: MappingProxyType[int, ArgmaxQOverAFunction]
     """Immutable mapping of period to argmax-and-max-Q functions."""
 
-    next_state: NextStateSimulationFunction
-    """Compiled function to compute next-period states."""
+    next_state: MappingProxyType[int, NextStateSimulationFunction]
+    """Immutable mapping of period to next-period-state functions."""
+
+    age_specialized_function_names: frozenset[FunctionName] = frozenset()
+    """Function names that were `AgeSpecializedFunction` in the user regime.
+
+    The published `functions` hold these resolved at the regime's representative
+    age only — the per-period programs (`argmax_and_max_Q_over_a`, `next_state`)
+    carry the true per-age closures. Consumers computing period-specific outputs
+    from `functions` (e.g. `additional_targets`) must reject targets that depend
+    on these names."""
 
     @property
     def state_names(self) -> tuple[StateOrActionName, ...]:

src/_lcm/model_processing.py

@@ -24,6 +24,7 @@
     materialize_granular_transition_params,
 )
 from _lcm.params.sequence_leaf import SequenceLeaf
+from _lcm.regime_building.age_specialization import resolve_node
 from _lcm.regime_building.finalize import FinalizedUserRegime
 from _lcm.regime_building.h_dag import get_dag_targets_consumed_by_H
 from _lcm.regime_building.max_Q_over_a import TASTE_SHOCK_SCALE_PARAM
@@ -164,6 +165,7 @@ def validate_model_inputs(
     regime_id_class: type,
     n_subjects: int | None = None,
     broadcast_variables: Mapping[RegimeName, frozenset[str]] | None = None,
+    ages: AgeGrid | None = None,
 ) -> None:
     """Validate model constructor inputs.
 
@@ -172,6 +174,9 @@ def validate_model_inputs(
     function with their declared types. This function focuses on value and
     cross-field rules.
 
+    `ages` lets the used-variable check resolve `AgeSpecializedFunction` functions at
+    each regime's representative age, so a state read only by a policy-specialized
+    function still counts as used.
     """
     _fail_if_invalid_n_subjects(n_subjects=n_subjects)
 
@@ -216,7 +221,7 @@ def validate_model_inputs(
         )
     error_messages.extend(
         _validate_all_variables_used(
-            user_regimes, broadcast_variables=broadcast_variables
+            user_regimes, broadcast_variables=broadcast_variables, ages=ages
         )
     )
 
@@ -252,6 +257,7 @@ def _validate_all_variables_used(
     user_regimes: Mapping[RegimeName, UserRegime],
     *,
     broadcast_variables: Mapping[RegimeName, frozenset[str]] | None = None,
+    ages: AgeGrid | None = None,
 ) -> list[str]:
     """Validate that all states and actions are used somewhere in each regime.
 
@@ -282,6 +288,21 @@ def _validate_all_variables_used(
         if broadcast_variables is not None:
             variable_names -= broadcast_variables.get(regime_name, frozenset())
         user_functions = dict(user_regime.get_all_functions(phase="solve"))
+        if ages is not None:
+            # Resolve any `AgeSpecializedFunction` marker to its concrete function at a
+            # representative active age so `get_ancestors` sees the real argument
+            # dependencies. The dependency structure is age-invariant, so any active
+            # age serves; the resolved build is memoized and reused by the engine.
+            active_periods = ages.get_periods_where(user_regime.active)
+            if active_periods:
+                representative_age = float(ages.period_to_age(active_periods[0]))
+                user_functions = cast(
+                    "dict[str, Callable[..., object]]",
+                    {
+                        name: resolve_node(func, representative_age)
+                        for name, func in user_functions.items()
+                    },
+                )
 
         targets = [
             "utility",
@@ -434,7 +455,12 @@ def _partial_fixed_params_into_regimes(
                     for period, func in simulation.argmax_and_max_Q_over_a.items()
                 }
             ),
-            next_state=functools.partial(simulation.next_state, **regime_fixed),
+            next_state=MappingProxyType(
+                {
+                    period: functools.partial(func, **regime_fixed)
+                    for period, func in simulation.next_state.items()
+                }
+            ),
             compute_regime_transition_probs=(
                 functools.partial(
                     simulation.compute_regime_transition_probs,