Solve and simulate on reachable discrete states with a composite state (#522)

[timmens]

Closes #522 as one vertical slice: model description to solve and simulate on reachable discrete states. Interface names are placeholders.

What changes

• Model(composite_states={"history": CompositeState(members=(...), start_codes={...})}) replaces the axes of deterministic discrete member states by one generated discrete state whose codes index the member tuples reachable from the start codes.

• src/_lcm/regime_building/reachable_states.py enumerates the tuples reachable at each (period, regime) from the engine's own laws, constraint pools and fixed params. It tries every code of the other discrete variables those read, and skips (and reports) constraints that read continuous variables or free params, which can only enlarge the sets. It follows every target of both phases' static regime graphs. Stochastic member laws, laws reading non-enumerable inputs, and non-discrete members are refused by name.

• src/_lcm/regime_building/composite_states.py lowers the declaration at the top of Model.__init__. A dense build (no jit, no seal) supplies the engine laws. Each carrying regime then holds:

  • the generated grid over the union of the reachable tuples;
  • one decoder function per member, so user functions keep reading member names;
  • one encoder law per carrying target;
  • a generated constraint <name>_is_reachable, which makes each code infeasible at the periods it cannot be reached in, so V reads -inf there and no reachable cell reads it.

  Fixed params of member laws are bound into the encoder and marked consumed. Codes off the set encode to MISSING_CAT_CODE, which reads NaN.

• Refused by name: a phase-specific (Phased) member law or grid, a member held on different grids in different regimes, a member declared by its values (DiscreteGrid(values=...)), a member law that reads a variable declared by its values, and a declared constraint or function that collides with a generated name.

• The enumeration evaluates laws and constraints with each period's numeric age, so sub-annual and irregular AgeGrids build. The generated encoder accepts every numeric or boolean input its member laws read, such as a floating helper, an int64 helper under x64, or the age. Member laws may return int64 codes; the encoder and the enumeration check them against the grid in their own width before narrowing to int32.

• Initial conditions name the members, as code arrays or as a labelled frame. Only subjects whose starting regime carries the members are encoded, and they need every member column; a supplied column is still checked for whole int32 codes and its length. An unreachable start tuple or an unknown label raises InvalidInitialConditionsError. Members reach to_dataframe through additional_targets.

• Certified: only src/lcm/model.py. Model.__init__, Model.simulate and the module surface are re-pinned (check_seals.py --fix, repin_corridors.py --changed-source src/lcm/model.py; check_seals, verify.py and verify.py --self-test exit 0). The slot-read ledger, public API index, model reference and changelog are updated.

Evidence

• tests/test_composite_states.py covers a three-counter toy with a stochastic offer, a cap read from fixed params, and a continuous budget:

  • V agrees with the dense model at every reachable cell (assert_agrees_to_ulp, 4 ULP);
  • the simulated panel equals the dense one;
  • V has one entry per code and per other-state cell;
  • -inf at codes unreachable at their period;
  • an unreachable start and an unknown label are refused;
  • a labelled frame starts subjects at the same codes;
  • V is finite exactly at the codes an independent oracle reaches, at every period;
  • nonzero-label starts decode to their codes;
  • the exact MISSING_CAT_CODE for an in-grid tuple outside the set and for an off-grid tuple;
  • the refusals listed above;
  • a cloudpickle round trip keeps the lowering.

• tests/test_composite_state_boundaries.py:

  • member laws with no other discrete input and edges with no feasible successor;
  • a constraint the enumeration skips keeps the code set on the member grids, and a feasible off-grid law still raises;
  • initial cohorts with and without carrying subjects, optional member columns, and fractional, out-of-range and wrong-length codes;
  • a quarterly AgeGrid with and without a member law reading age, and a member law reading a floating helper, each against dense V (and the dense panel for the helper);
  • a branch reached only in simulation and the phases swapped;
  • two targets of one phase with different member laws, by literal set and by V.

• tests/regime_building/test_reachable_states.py: the per-node sets equal an itertools oracle, the union is their sorted union, a constraint reading a free parameter is skipped with the support unrestricted, and refusals.

• Focused files at 91148ce (tests/test_composite_states.py, tests/test_composite_state_boundaries.py, tests/regime_building/test_reachable_states.py, tests/test_regime_slot_read_ledger.py, tests/test_docs_public_api_coverage.py, tests/ci/test_ci_workloads_manifest.py): pixi run -e tests-cpu pytest <files> -v --precision=64 and --precision=32 on macOS arm64, junit 520 tests, 0 failures at each precision (six int64 tests skip without x64).

• arcidiacono2025 on one A40 (Marvin mlgpu_short, fp32), 24 decision periods, pylcm e5e882ef0 (jobs 28225824, 28231639, 28231640). That build had both a closed and a masked variant; the closed variant is removed here and the masked numbers below are this PR's design.

  rung	V dense	V composite	GPU peak dense / composite	warm solve, checks off	warm solve, checks on	V max rel gap	panel
  r4	6.45 GB	0.23 GB (27.9x)	9.6 / 0.50 GiB	3.5 s / 0.6 s (6.4x)	6.7 s / 5.9 s	2.05e-7	identical
  r5	19.2 GB	0.58 GB (33.2x)	28.4 / 1.2 GiB	8.2 s / 0.8 s (10.7x)	12.2 s / 6.2 s	2.9e-7	identical
  r6 (paper state space)	not run (time limit; extrapolated about 55 GB)	1.67 GB	n/a / 3.4 GiB	1.6 s	8.5 s	n/a	n/a

  r6 holds 31,813 reachable tuples. The r6 job reached its 8-hour limit after the composite run, because the paper's own input preparation took 7.9 hours. At 13 periods the prototype showed dense r6 does not fit an A40 (planner: 34 GB resident vs a 30.4 GB budget), while composite V is 1.07 GB.

Limits and follow-ups

• With runtime_checks=True the composite solve spends about 5 s per solve in validate_discrete_state_laws_all_periods. The encoder returns MISSING_CAT_CODE at infeasible actions, so the check takes its second, feasibility-joint pass every period. The result is correct, but a composite-aware fast path in transition_checks.py (certified) would remove the cost.
• One composite state per model. Members must be carried by every carrying regime on the same grid with the same law in both phases.
• Out of scope:
  • labelled member columns by default;
  • lookup_policy / value_frame in member terms;
  • per-period ragged grids;
  • stochastic members;
  • members declared by their values;
  • a narrower engine factory in place of the inner dense build;
  • final names;
  • the 48-period paper run (input preparation alone would exceed 24 h).

[hmgaudecker]

Thank you. The reachable-set construction and the lowering to one composite axis are
clear.

Two defects on the input boundary of the composite state. Each one stops a valid
model before the solve:

• A sub-annual or irregular AgeGrid sends a Fraction to JAX. Construction fails, even
  when no member law reads age.
• The generated encoder accepts only int32 keyword inputs. A member law that reads a
  deterministic floating helper of the category states fails the runtime type check.

Also, the tests do not detect some defects in the stated behaviour. Please add the cases
that the inline comment on the tests gives.

Please also update the PR description:

• The CI bullet says that four base tests fail and that the fix is not on the base. At
  0d0416e, the cpu, gpu, ty, notebooks and benchmark-pr workflows complete with
  success. Remove or replace this bullet.
• The focused-files bullet names "these two" files. It does not name
  tests/test_composite_state_boundaries.py. Add it.
• The count 446/446 does not name a commit. Give the commit and the log, or remove the
  count.

[hmgaudecker]

Thank you for the fast fixes. The fractional age labels and the int32/int64 encoder
inputs are now correct. One case is still refused, and two tests do not detect the
defects they are about.

The composite encoder refuses a helper that returns a Python float. pylcm allows a
user function to return a Python scalar (UserFunctionResult in
src/lcm/typing.py:175-184). Without a composite state, such a model solves. With a
composite state, the encoder rejects the value before the member laws run:

BeartypeCallHintParamViolation: Method _CompositeEncoder.__call__() parameter kwargs=0.5
violates type hint FloatND | _SignedInts | BoolND

The fix is one annotation. Please see the inline comment.

Two tests do not detect the defect class they target:

• tests/test_composite_state_boundaries.py:401-406: the helper test floors its own
  value (floor(a + 0.5)). An encoder that floors floating inputs passes it.
• tests/test_composite_state_boundaries.py:395-398: the age law reads only age 0. An encoder that replaces the age by 0 or by the
  period passes it.

The collapsed block at the end holds a patch with two new tests and the fix. I
applied it to 91148ce8e and ran it:

• New tests without the fix, fp64: 2 failed, 10 passed. The 2 failures are the
  Python-helper composite cases, with the error above.
• New tests with the fix: 12 passed, fp64 and fp32.
• The three composite test files with the fix: 73 passed at fp64; 67 passed and 6
  skipped at fp32.
• prek on the two changed files: passed (ruff, format, ty).

Patch: two tests and the encoder fix

diff --git a/src/_lcm/regime_building/composite_states.py b/src/_lcm/regime_building/composite_states.py
index 341f2f8c4..93b7dced9 100644
--- a/src/_lcm/regime_building/composite_states.py
+++ b/src/_lcm/regime_building/composite_states.py
@@ -41,6 +41,7 @@ from _lcm.regime_building.reachable_states import (
 )
 from _lcm.typing import (
     DataclassInstance,
+    EconFunctionArg,
     FlatRegimeParams,
     FunctionName,
     ReferenceName,
@@ -58,7 +59,6 @@ from lcm.transition import fixed_transition
 from lcm.typing import (
     BoolND,
     DiscreteState,
-    FloatND,
     IntND,
     ScalarInt,
     UserInitialConditions,
@@ -454,9 +454,9 @@ class _CompositeEncoder:
     laws: Callable[..., tuple[_SignedInts, ...]]
     """The member laws, returning next member codes in member order.
 
-    They read whatever the member laws read: codes, but also floating or boolean
-    helpers, signed integer helpers of any width and the age, so the keyword inputs
-    are not restricted to int32 codes.
+    Inputs use the same value-tree contract as other function forwarders. Helper
+    values, including Python scalars, pass through unchanged. The returned member
+    codes retain their signed-integer contract and are checked before indexing.
     """
 
     law_args: tuple[ReferenceName, ...]
@@ -471,7 +471,7 @@ class _CompositeEncoder:
     shape: tuple[int, ...]
     """Number of codes of every member."""
 
-    def __call__(self, **kwargs: FloatND | _SignedInts | BoolND) -> IntND:
+    def __call__(self, **kwargs: EconFunctionArg) -> IntND:
         values = kwargs | dict(self.bound)
         nexts = self.laws(**{arg: values[arg] for arg in self.law_args})
         return _read_flat(table=self.encoding, shape=self.shape, index=nexts)
diff --git a/tests/test_composite_state_boundaries.py b/tests/test_composite_state_boundaries.py
index f965957de..e46ce0eca 100644
--- a/tests/test_composite_state_boundaries.py
+++ b/tests/test_composite_state_boundaries.py
@@ -9,6 +9,7 @@
 """
 
 from collections.abc import Mapping
+from fractions import Fraction
 
 import jax
 import jax.numpy as jnp
@@ -49,6 +50,7 @@ from lcm.typing import (
     ScalarInt,
     StateName,
     UserFunction,
+    UserInitialConditions,
 )
 from tests.conftest import DECIMAL_PRECISION, assert_agrees_to_ulp
 
@@ -710,3 +712,179 @@ def test_wide_code_outside_the_table_never_wraps_into_it(code: int) -> None:
     table = jnp.arange(4, dtype=jnp.int32)
     got = _read_flat(table=table, shape=(2, 2), index=(jnp.int64(code), jnp.int64(0)))
     assert int(got) == MISSING_CAT_CODE
+
+
+def _python_half(*, a: DiscreteState) -> float:  # noqa: ARG001 - DAG dependency
+    """Return a Python scalar with a declared categorical dependency."""
+    return 0.5
+
+
+def _array_half(*, a: DiscreteState) -> FloatND:
+    """Return the same fractional value as an array in the working precision."""
+    return jnp.zeros_like(a, dtype=float) + 0.5
+
+
+def _next_a_above_quarter(*, half: FloatND) -> DiscreteState:
+    return jnp.asarray(half > 0.25, dtype=jnp.int32)
+
+
+def _next_b_below_zero(*, half: FloatND) -> DiscreteState:
+    return jnp.asarray(half < 0.0, dtype=jnp.int32)
+
+
+def _next_a_above_quarter_py(*, half: float) -> DiscreteState:
+    return jnp.asarray(half > 0.25, dtype=jnp.int32)
+
+
+def _next_b_below_zero_py(*, half: float) -> DiscreteState:
+    return jnp.asarray(half < 0.0, dtype=jnp.int32)
+
+
+def _fractional_helper_model(
+    *, half: UserFunction, composite: bool, enable_jit: bool
+) -> Model:
+    """Build two live periods with both member laws reading the same helper."""
+    return Model(
+        n_periods=3,
+        regimes={
+            "live": Regime(
+                states={"a": DiscreteGrid(_Bit), "b": DiscreteGrid(_Bit)},
+                state_transitions=(
+                    {"a": _next_a_above_quarter_py, "b": _next_b_below_zero_py}
+                    if half is _python_half
+                    else {"a": _next_a_above_quarter, "b": _next_b_below_zero}
+                ),
+                functions={"utility": _identity_utility, "half": half},
+            ),
+            "dead": Regime(functions={"utility": _zero}),
+        },
+        edges={
+            "live": {
+                "live": Periods(values=(0,)),
+                "dead": Periods(values=(1,)),
+            }
+        },
+        regime_id_class=_RegimeId,
+        initial_nodes=InitialNodes(by_period={0: "live"}),
+        execution_config=ExecutionConfig(device_memory_bytes=None),
+        enable_jit=enable_jit,
+        composite_states=(
+            {
+                "history": CompositeState(
+                    members=("a", "b"), start_codes={"a": (0,), "b": (0,)}
+                )
+            }
+            if composite
+            else {}
+        ),
+    )
+
+
+def _assert_fractional_transition(
+    *,
+    model: Model,
+    composite: bool,
+    initial_conditions: UserInitialConditions,
+) -> pd.DataFrame:
+    """Check literal support, value and decoded states for the path 00 to 10."""
+    if composite:
+        by_node = model._composite_states["history"].reachable.by_node
+        assert {node: set(tuples) for node, tuples in by_node.items()} == {
+            (0, "live"): {(0, 0)},
+            (1, "live"): {(1, 0)},
+        }
+    model.validate_transitions(params=_PARAMS)
+    solution = model.solve(params=_PARAMS, log_level="off")
+    values = np.asarray(solution.values[0]["live"])
+    # The initial tuple 00 is first in either representation. Its utilities are
+    # 1 and 2, with discount factor 9/10, so its value is 1 + 9/10 * 2 = 14/5.
+    np.testing.assert_array_almost_equal(
+        values.ravel()[0], 14 / 5, decimal=DECIMAL_PRECISION
+    )
+    panel = model.simulate(
+        params=_PARAMS,
+        solution=solution,
+        initial_conditions=initial_conditions,
+        seed=7,
+        log_level="off",
+    ).to_dataframe(
+        use_labels=False, additional_targets=["a", "b"] if composite else None
+    )
+    live = panel.loc[panel["regime_name"] == "live"].sort_values("period")
+    np.testing.assert_array_equal(live["period"].to_numpy(), [0, 1])
+    np.testing.assert_array_equal(live[["a", "b"]].to_numpy(), [[0, 0], [1, 0]])
+    return live
+
+
+@pytest.mark.parametrize("composite", [False, True], ids=["dense", "composite"])
+@pytest.mark.parametrize("enable_jit", [False, True], ids=["eager", "jit"])
+@pytest.mark.parametrize("half", [_python_half, _array_half], ids=["python", "array"])
+def test_fractional_helper_preserves_value_and_state_path(
+    *, half: UserFunction, composite: bool, enable_jit: bool
+) -> None:
+    """Preserve Python and array helper values through public solve and simulate."""
+    model = _fractional_helper_model(
+        half=half, composite=composite, enable_jit=enable_jit
+    )
+    _assert_fractional_transition(
+        model=model,
+        composite=composite,
+        initial_conditions={
+            "regime_id": jnp.full(1, _RegimeId.live, dtype=jnp.int32),
+            "period": jnp.zeros(1, dtype=jnp.int32),
+            "a": jnp.zeros(1, dtype=jnp.int32),
+            "b": jnp.zeros(1, dtype=jnp.int32),
+        },
+    )
+
+
+def _next_a_from_fractional_age(*, age: FloatND) -> DiscreteState:
+    return jnp.asarray(age > 0.25, dtype=jnp.int32)
+
+
+@pytest.mark.parametrize("composite", [False, True], ids=["dense", "composite"])
+@pytest.mark.parametrize("enable_jit", [False, True], ids=["eager", "jit"])
+def test_nonzero_fractional_age_preserves_value_and_state_path(
+    *, composite: bool, enable_jit: bool
+) -> None:
+    """Use the source age 1/2 at period zero to reach the literal next state 10."""
+    start_age = Fraction(1, 2)
+    model = Model(
+        regimes={
+            "live": Regime(
+                states={"a": DiscreteGrid(_Bit), "b": DiscreteGrid(_Bit)},
+                state_transitions={
+                    "a": _next_a_from_fractional_age,
+                    "b": fixed_transition("b"),
+                },
+                functions={"utility": _identity_utility},
+            ),
+            "dead": Regime(functions={"utility": _zero}),
+        },
+        edges={"live": {"live": start_age, "dead": Fraction(1)}},
+        ages=AgeGrid(start=start_age, inclusive_stop=Fraction(3, 2), step="2Q"),
+        regime_id_class=_RegimeId,
+        initial_nodes=InitialNodes(by_age={start_age: "live"}),
+        execution_config=ExecutionConfig(device_memory_bytes=None),
+        enable_jit=enable_jit,
+        composite_states=(
+            {
+                "history": CompositeState(
+                    members=("a", "b"), start_codes={"a": (0,), "b": (0,)}
+                )
+            }
+            if composite
+            else {}
+        ),
+    )
+    live = _assert_fractional_transition(
+        model=model,
+        composite=composite,
+        initial_conditions={
+            "regime_id": jnp.full(1, _RegimeId.live, dtype=jnp.int32),
+            "age": jnp.asarray([0.5]),
+            "a": jnp.zeros(1, dtype=jnp.int32),
+            "b": jnp.zeros(1, dtype=jnp.int32),
+        },
+    )
+    np.testing.assert_array_equal(live["age"].to_numpy(), [0.5, 1.0])

[hmgaudecker]

Excellent, thank you!