regime_template: exempt next_<state> names from fixed_param extraction (#342)

Co-authored-by: Claude Opus 4.7 (1M context) <noreply@anthropic.com>

Author: hmgaudecker

.ai-instructions

@@ -294,6 +294,150 @@ initial_conditions = {
 - `model.n_periods` - Number of periods in the model (derived from `ages`)
 - `model.regime_names_to_ids` - Immutable mapping from regime names to integer indices
 
+## Testing
+
+### Test-Driven Development — always
+
+**Always write the test first, watch it fail, then implement.** No exceptions for new
+behavior or bug fixes. Tests are not an afterthought, they are the spec.
+
+The cycle:
+
+1. **Red.** Write a failing test that asserts the desired behavior in user-facing terms.
+   Run it. Confirm it fails for the *right* reason (the missing behavior — not a typo,
+   not an import error).
+1. **Green.** Write the smallest amount of code that makes the test pass.
+1. **Refactor.** Clean up while keeping the test green.
+
+Apply per case:
+
+- **New feature** → red-green-refactor.
+- **Bug fix** → reproduce as a failing test before writing the fix. The test then
+  prevents regression.
+- **Refactor (no behavior change)** → existing tests are the spec. Keep them green
+  before, during, and after. No new test needed if behavior is unchanged; if you find a
+  behavior gap, fill it with a new test *before* refactoring.
+
+### Test docstrings — describe behavior, not history
+
+Test docstrings state what *should* be true, in user-facing terms. Pretend the reader
+has never seen the PR. They should not need to.
+
+```python
+# Good — behavior, in plain language
+def test_simulate_with_chained_transitions_yields_expected_next_wealth():
+    """`next_wealth_t = wealth_t - c_t + 0.1 * next_aime_t` holds in simulation."""
+
+
+# Bad — rehearses the prior bug or implementation history
+def test_solve_resolves_chain_via_dags():
+    """Before the fix, `_resolve_fixed_params` raised
+    `InvalidParamsError: Missing required parameter: ...` because
+    `create_regime_params_template` classified ..."""
+```
+
+Rule of thumb: **would the docstring still make sense in 9 months without the PR
+context?** If not, rewrite it.
+
+### Concrete-value assertions
+
+Assert *what* the result is, not just that it didn't crash.
+
+```python
+# Good — analytical value with explicit tolerance
+np.testing.assert_allclose(curr["wealth"], expected_next_wealth, atol=1e-6)
+
+# Bad — passes whether the math is right or not
+assert not jnp.any(jnp.isnan(V_arr))
+assert df["wealth"].notna().all()
+```
+
+`not isnan` and `no exception raised` belong in CI smoke tests, not in the unit tests
+for the feature itself.
+
+### Mechanics
+
+- Use plain pytest functions, never test classes (`class TestFoo`)
+- Use `@pytest.mark.parametrize` for test variations
+
+## Docstring Style
+
+Docstrings and inline comments describe the code's *current* state in user-facing terms.
+The 9-month-without-PR-context reader is the audience: a docstring that survives that
+test stays useful; one that rehearses the diff or the prior implementation rots
+immediately.
+
+This applies to **all** docstrings and comments — source and tests. For tests
+specifically, see also "Test docstrings — describe behavior, not history" above.
+
+### Describe state, not history
+
+State what is true now. Don't reference prior designs, removed code, or what was
+changed. Words like "earlier", "previously", "now", "formerly", "the old", "before the
+fix" are red flags.
+
+```python
+# Good — forward-looking constraint
+class _DiagnosticRow:
+    """Metadata captured during the backward-induction loop.
+
+    Holds only Python-scalar metadata — no device-array references —
+    so every (regime, period) row stays at a few bytes regardless of
+    grid size.
+    """
+
+
+# Bad — rehearses prior design
+class _DiagnosticRow:
+    """Metadata captured during the backward-induction loop.
+
+    Holds only Python-scalar metadata. The earlier design captured
+    state_action_space and a closure directly on each row, which
+    pinned every period's V template in device memory until the
+    post-loop flush.
+    """
+```
+
+### No PR numbers, no model-specific magic numbers
+
+PR references (`#334 removed the host stalls`, `the bug was fixed in #42`) rot as the
+codebase evolves and provide no useful signal to a reader who isn't already in context.
+Magic numbers tied to a specific model size or hardware
+(`~2 MB at production grid sizes`, `fits on a 16 GB device`) imply a fixed scale that's
+only true on whichever model/box the comment was written against. State the qualitative
+dependency instead.
+
+```python
+# Good — qualitative dependency
+# Frees per-period intermediate buffers (V_arr-shaped, so
+# model-dependent) so they don't stack up across the loop.
+
+# Bad — PR reference + magic number
+# Frees per-period intermediate buffers (~2 MB each at production
+# grid sizes) so we don't re-introduce the host stalls that #334
+# removed.
+```
+
+### Bulleted lists for enumerated cases
+
+When describing a fixed set of cases (log levels, regime kinds, parameter types,
+dispatch strategies), use one bullet per case rather than running prose. Bullets scan;
+prose hides cases.
+
+```python
+# Good — scannable
+# Gate falls out of the public log level:
+# - `"off"` ⇒ nothing (skips even the NaN fail-fast)
+# - `"warning"` / `"progress"` ⇒ NaN/Inf only
+# - `"debug"` ⇒ adds the min/max/mean trio
+
+
+# Bad — buried in prose
+# Gate falls out of the public log level: `"off"` ⇒ nothing,
+# `"warning"` / `"progress"` ⇒ NaN/Inf only, `"debug"` ⇒ adds the
+# min/max/mean trio. `"off"` skips even the NaN fail-fast.
+```
+
 ## Development Notes
 
 ### JAX Integration
@@ -401,11 +545,6 @@ Code structure should be self-evident from function names and ordering.
   display math, and `[text](url)` for links. Never use rST-style ``` `` code `` ```,
   `:math:`, `:func:`, or `` `link <url>`_ ``.
 
-### Testing Style
-
-- Use plain pytest functions, never test classes (`class TestFoo`)
-- Use `@pytest.mark.parametrize` for test variations
-
 ### Plotting
 
 - Always use **plotly** for visualizations, never matplotlib. Use `plotly.graph_objects`

AGENTS.md

@@ -17,20 +17,18 @@
 )
 
 
-def create_regime_params_template(
-    regime: Regime,
-) -> RegimeParamsTemplate:
+def create_regime_params_template(regime: Regime) -> RegimeParamsTemplate:
     """Create parameter template from a regime specification.
 
-    Discover parameters from function signatures via `dags.tree`. Parameters are
-    function arguments that are not states, actions, other regime functions, or
-    special variables (period, age, E_next_V).
+    Discover parameters from function signatures via `dags.tree`. Parameters
+    are function arguments that are not states, actions, regime functions,
+    `next_<state>` outputs, or special variables (`period`, `age`, `E_next_V`).
 
     For `SolveSimulateFunctionPair` entries, the template contains the **union**
     of both variants' parameters so the user can provide a single flat params
     dict that satisfies both phases.
 
-    Grids with runtime-supplied values (IrregSpacedGrid without points,
+    Grids with runtime-supplied values (`IrregSpacedGrid` without points,
     `_ShockGrid` without full shock_params) add entries to the template under
     pseudo-function keys matching the state or action name.
 
@@ -41,8 +39,15 @@ def create_regime_params_template(
         The regime parameter template with type annotations as values.
 
     """
-    H_variables = {*regime.functions, "period", "age", "E_next_V"}
-    variables = H_variables | set(regime.actions) | set(regime.states)
+    variables = {
+        *set(regime.states),
+        *set(regime.actions),
+        *regime.functions,
+        *(f"next_{name}" for name in regime.states),
+        "period",
+        "age",
+        "E_next_V",
+    }
 
     function_params: dict[FunctionName, dict[str, str]] = {}
 

pixi.lock

@@ -6,7 +6,7 @@
 import jax
 import pandas as pd
 from dags import concatenate_functions, with_signature
-from dags.tree import qname_from_tree_path, tree_path_from_qname
+from dags.tree import qname_from_tree_path
 from jax import Array
 
 from lcm.grids import Grid
@@ -69,6 +69,17 @@ def get_next_state_function_for_simulation(
 ) -> NextStateSimulationFunction:
     """Get function that computes the next states during the simulation.
 
+    Builds one DAG per target regime using unqualified `next_<state>` keys, mirroring
+    the per-target structure of {func}`get_next_state_function_for_solution`. This
+    lets a transition function or auxiliary regime function consume another
+    transition's `next_<state>` output via plain name resolution within the same
+    target's DAG. The combined function returns a nested mapping keyed by target
+    regime name, with each inner dict using unqualified `next_<state>` keys.
+
+    Stochastic-transition wrappers expose `key_<target>__next_<state>` and
+    `weight_<target>__next_<state>` as external arguments so callers can pass a
+    distinct random key and pre-computed weight per target.
+
     Args:
         transitions: Nested mapping of target regime names to transition functions.
         functions: Immutable mapping of auxiliary functions of a regime.
@@ -78,26 +89,31 @@ def get_next_state_function_for_simulation(
 
     Returns:
         Function that computes the next states. Depends on states and actions of the
-        current period, and the regime parameters ("params"). If target is "simulate",
-        the function also depends on the dictionary of random keys ("keys"), which
-        corresponds to the names of stochastic next functions.
+        current period, and the regime parameters ("params"). The function also
+        depends on the dictionary of random keys ("keys") for stochastic transitions.
+        Returns `{target_regime_name: {next_<state>: array}}`.
 
     """
-    flat_transitions = flatten_regime_namespace(transitions)
-
-    # For the simulation target, we need to extend the functions dictionary with
-    # stochastic next states functions and their weights.
-    extended_transitions = _extend_transitions_for_simulation(
-        all_grids=all_grids,
-        flat_transitions=flat_transitions,
-        variable_info=variable_info,
-        stochastic_transition_names=stochastic_transition_names,
-    )
-    functions_to_concatenate = extended_transitions | dict(functions)
+    per_target_funcs: dict[RegimeName, Callable[..., dict[str, Array]]] = {}
+    for target, target_transitions in transitions.items():
+        extended = _extend_target_transitions_for_simulation(
+            target=target,
+            target_transitions=target_transitions,
+            all_grids=all_grids,
+            variable_info=variable_info,
+            stochastic_transition_names=stochastic_transition_names,
+        )
+        per_target_funcs[target] = concatenate_functions(
+            functions=dict(extended) | dict(functions),
+            targets=list(extended.keys()),
+            return_type="dict",
+            enforce_signature=False,
+            set_annotations=True,
+        )
 
     return concatenate_functions(
-        functions=functions_to_concatenate,
-        targets=list(flat_transitions.keys()),
+        functions=per_target_funcs,
+        targets=list(per_target_funcs.keys()),
         return_type="dict",
         enforce_signature=False,
         set_annotations=True,
@@ -137,64 +153,59 @@ def get_next_stochastic_weights_function(
     )
 
 
-def _extend_transitions_for_simulation(
+def _extend_target_transitions_for_simulation(
     *,
+    target: RegimeName,
+    target_transitions: MappingProxyType[TransitionFunctionName, Callable[..., Array]],
     all_grids: MappingProxyType[RegimeName, MappingProxyType[StateOrActionName, Grid]],
-    flat_transitions: FunctionsMapping,
     variable_info: pd.DataFrame,
     stochastic_transition_names: frozenset[TransitionFunctionName],
 ) -> dict[TransitionFunctionName, Callable[..., Array]]:
-    """Extend the functions dictionary for the simulation target.
+    """Replace stochastic transitions for one target with realisation wrappers.
+
+    Deterministic transitions are passed through unchanged. Stochastic transitions
+    are replaced by wrappers that draw a realisation from a precomputed weight
+    vector and a random key. The wrapper's external argument names use
+    target-qualified form (`key_<target>__<next_state>`,
+    `weight_<target>__<next_state>`) so multi-target callers can supply distinct
+    random keys per target. The dict key keeps the unqualified `next_<state>` so
+    other transitions or regime functions in the same target's DAG can resolve
+    it by name.
 
     Args:
+        target: Target regime name.
+        target_transitions: Mapping of unqualified `next_<state>` transition names
+            to functions, restricted to one target regime.
         all_grids: Immutable mapping of regime names to Grid spec objects.
-        flat_transitions: Flattened mapping of transition names to functions.
         variable_info: Variable info of the current regime.
         stochastic_transition_names: Frozenset of stochastic transition function names.
 
     Returns:
-        Extended functions dictionary.
+        Extended transitions dictionary keyed by unqualified `next_<state>` names.
 
     """
     shock_names: set[ShockName] = set(variable_info.query("is_shock").index.to_list())
     flat_grids = flatten_regime_namespace(all_grids)
-    discrete_stochastic_targets = [
-        func_name
-        for func_name in flat_transitions
-        if tree_path_from_qname(func_name)[-1] in stochastic_transition_names
-        and tree_path_from_qname(func_name)[-1].removeprefix("next_") not in shock_names
-    ]
-    continuous_stochastic_targets = [
-        func_name
-        for func_name in flat_transitions
-        if tree_path_from_qname(func_name)[-1] in stochastic_transition_names
-        and tree_path_from_qname(func_name)[-1].removeprefix("next_") in shock_names
-    ]
-    # Handle stochastic next states functions
-    # ----------------------------------------------------------------------------------
-    # We generate stochastic next states functions that simulate the next state given
-    # a random key (think of a seed) and the weights corresponding to the labels of the
-    # stochastic variable. The weights are computed using the stochastic weight
-    # functions, which we add the to functions dict. `dags.concatenate_functions` then
-    # generates a function that computes the weights and simulates the next state in
-    # one go.
-    # ----------------------------------------------------------------------------------
-    discrete_stochastic_next = {
-        name: _create_discrete_stochastic_next_func(
-            name=name, labels=flat_grids[name.replace("next_", "")].to_jax()
-        )
-        for name in discrete_stochastic_targets
-    }
-    continuous_stochastic_next = {
-        name: _create_continuous_stochastic_next_func(name=name, flat_grids=flat_grids)
-        for name in continuous_stochastic_targets
-    }
-
-    # Overwrite regime transitions with generated stochastic next states functions
-    # ----------------------------------------------------------------------------------
-    return (
-        dict(flat_transitions) | discrete_stochastic_next | continuous_stochastic_next
+    extended: dict[TransitionFunctionName, Callable[..., Array]] = dict(
+        target_transitions
     )
+    for next_state_name in target_transitions:
+        if next_state_name not in stochastic_transition_names:
+            continue
+        qname = qname_from_tree_path((target, next_state_name))
+        raw_state_name = next_state_name.removeprefix("next_")
+        if raw_state_name in shock_names:
+            extended[next_state_name] = _create_continuous_stochastic_next_func(
+                name=qname, flat_grids=flat_grids
+            )
+        else:
+            extended[next_state_name] = _create_discrete_stochastic_next_func(
+                name=qname,
+                labels=flat_grids[
+                    qname_from_tree_path((target, raw_state_name))
+                ].to_jax(),
+            )
+    return extended
 
 
 def _create_discrete_stochastic_next_func(