NEGM: compose resources as base minus the declared outer cost

The stacked-carry lift is exact only when the inner resources are affine
in the declared outer cost with coefficient exactly -1 — a property no
finite set of runtime probes can certify on an arbitrary user function
(an interaction term can vanish at every probe binding and still bend
the translation elsewhere). So the affineness becomes structural:

- With `NEGM.outer_cost` declared, the regime defines the cost-free base
  `<inner.resources>_before_outer_cost` and `finalize_regimes` composes
  `resources = base - cost` — the single point before validation and
  processing, so every consumer (inner EGM solve, budget constraint,
  parent child-resources read) sees the composed function. A
  user-defined resources function alongside a declared cost, or a
  missing base, is rejected at model build.
- Validation replaces the cost-opaque resources funnel check with the
  direct requirement that the base does not read the outer
  post-decision; the cost-ancestor and no-adjustment checks stand.
- The shift builder derives purely from the declared cost DAG; the two
  runtime consistency probes (reference-perturbed differences, the
  autodiff coefficient check) are deleted as redundant-by-construction,
  which also removes their float32 tolerance surface.
- All NEGM models restructure mechanically (they were already written
  in `base - cost` form); DS-2024 oracle gates unchanged.

Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
Claude-Session: https://claude.ai/code/session_0123hk3uzBgBeZown4eBBQ2C

Author: hmgaudecker

src/_lcm/egm/negm_validation.py

@@ -36,12 +36,15 @@
   the last passive row axis, so a (hard) discrete action or a passive
   continuous state declared after the durable is rejected — the per-durable
   candidate lift would otherwise address the wrong axis,
-- outer-cost contract: the inner resources may depend on the outer
-  post-decision only through the declared `NEGM.outer_cost` function, which
-  itself may read only the durable state, the outer post-decision, and params
-  (with `outer_cost=None`, resources must be independent of the outer
-  post-decision) — otherwise no constant credited-cost translation onto a
-  common cash-on-hand axis exists and the stacked lift would be wrong,
+- outer-cost contract: with a declared `NEGM.outer_cost`, the resources
+  function is composed by `finalize_regimes` as
+  `<resources>_before_outer_cost - <outer_cost>` (affine in the cost by
+  construction; a user-defined resources function is rejected at
+  finalization); the declared cost may read only the durable state, the outer
+  post-decision, and params, and the cost-free base must not read the outer
+  post-decision (with `outer_cost=None`, resources must be independent of the
+  outer post-decision) — otherwise no constant credited-cost translation onto
+  a common cash-on-hand axis exists and the stacked lift would be wrong,
 - the no-adjustment candidate is a unary function of the durable state — it is
   evaluated as `keep(durable)` in the credited-cost lift and the
   child-resources query map.
@@ -360,29 +363,27 @@ def _fail_if_outer_cost_contract_violated(
     cash-on-hand axis by crediting a constant per (durable, outer-node) cell.
     That constant exists exactly when the inner resources depend on the outer
     post-decision through a single additive cost term that itself varies only
-    with the durable margin. The structural half of that contract is enforced
-    here, fail-closed:
-
-    - `NEGM.outer_cost` declared ⇒ it must be a regime function; its DAG
-      ancestors may contain no state or action other than the durable state and
-      the outer post-decision (params are fine); and the inner resources — with
-      the cost function made opaque — must not reach the outer post-decision
-      through any other channel,
+    with the durable margin. The contract is enforced fail-closed:
+
+    - `NEGM.outer_cost` declared ⇒ the resources function is composed by
+      `finalize_regimes` as `<resources>_before_outer_cost - <outer_cost>`, so
+      its affine use of the cost holds by construction (a user-defined
+      resources function is rejected at finalization). Here: the cost must be
+      a regime function whose DAG ancestors contain no state or action other
+      than the durable state and the outer post-decision (params are fine),
+      and the cost-free base must not read the outer post-decision — its only
+      outer-margin channel is the subtracted cost itself,
     - `NEGM.outer_cost=None` ⇒ the inner resources must be independent of the
       outer post-decision altogether (every candidate already shares the
       keeper's axis; the shift is zero).
-
-    The additive *use* of the cost inside resources (coefficient exactly `-1`,
-    no nonlinear wrapping) is not a graph property; the shift builder probes it
-    at runtime and raises `InvalidParamsError` on a mismatch.
     """
     inner = solver.inner
     durable_state = solver.outer_post_decision.removeprefix("next_")
-    resources_func = functions.get(inner.resources)
-    if resources_func is None:
-        return
 
     if solver.outer_cost is None:
+        resources_func = functions.get(inner.resources)
+        if resources_func is None:
+            return
         resources_ancestors = _dag_ancestors(
             functions=functions, target_func=resources_func
         )
@@ -425,21 +426,22 @@ def _fail_if_outer_cost_contract_violated(
         )
         raise ModelInitializationError(msg)
 
-    opaque_functions = _without(functions=functions, names={solver.outer_cost})
-    resources_ancestors = _dag_ancestors(
-        functions=opaque_functions, target_func=resources_func
-    )
-    if solver.outer_post_decision in resources_ancestors:
-        msg = (
-            f"In regime '{regime_name}', the inner resources "
-            f"'{inner.resources}' reads the outer post-decision "
-            f"'{solver.outer_post_decision}' outside the declared outer cost "
-            f"'{solver.outer_cost}'. The NEGM lift requires all outer-margin "
-            "dependence of resources to flow through `NEGM.outer_cost` — route "
-            f"it through '{solver.outer_cost}', or use `GridSearch` for this "
-            "regime."
-        )
-        raise ModelInitializationError(msg)
+    base_func = functions.get(f"{inner.resources}_before_outer_cost")
+    if base_func is not None:
+        base_ancestors = _dag_ancestors(functions=functions, target_func=base_func)
+        if solver.outer_post_decision in base_ancestors:
+            msg = (
+                f"In regime '{regime_name}', the cost-free resources base "
+                f"'{inner.resources}_before_outer_cost' reads the outer "
+                f"post-decision '{solver.outer_post_decision}'. It must not "
+                "read the outer post-decision: pylcm composes "
+                f"`{inner.resources} = {inner.resources}_before_outer_cost - "
+                f"{solver.outer_cost}`, so the base's only outer-margin "
+                "channel is the subtracted declared cost itself. Route the "
+                f"dependence through '{solver.outer_cost}', or use "
+                "`GridSearch` for this regime."
+            )
+            raise ModelInitializationError(msg)
 
 
 def _fail_if_no_adjustment_candidate_not_unary(

src/_lcm/regime_building/finalize.py

@@ -12,14 +12,21 @@
 
 from collections.abc import Mapping
 from types import MappingProxyType
+from typing import cast
+
+from dags import get_annotations, with_signature
+from dags.annotations import ensure_annotations_are_strings
 
 from _lcm.grids import DiscreteGrid
 from _lcm.typing import FunctionName, RegimeName
 from _lcm.user_regime_validation import _validate_completeness
 from _lcm.utils.error_messages import format_messages
 from lcm.exceptions import ModelInitializationError, RegimeInitializationError
+from lcm.phased import Phased
 from lcm.regime import Regime as UserRegime
+from lcm.solvers import NEGM
 from lcm.temporal_aggregation import H_linear
+from lcm.typing import FloatND, UserFunction
 
 # A user `Regime` after model-build finalization. Runtime-equivalent to
 # `lcm.regime.Regime`; internal signatures use this alias to mark values
@@ -67,6 +74,11 @@ def finalize_regimes(
         # Terminal regimes don't need H since Q = U directly (no E_next_V).
         if user_regime.transition is not None and "H" not in functions:
             functions["H"] = H_linear
+        _compose_negm_resources(
+            regime_name=regime_name,
+            user_regime=user_regime,
+            functions=functions,
+        )
         finalized = user_regime.replace(
             functions=functions, derived_categoricals=merged
         )
@@ -79,6 +91,90 @@ def finalize_regimes(
     return MappingProxyType(result)
 
 
+def _compose_negm_resources(
+    *,
+    regime_name: RegimeName,
+    user_regime: UserRegime,
+    functions: dict[FunctionName, UserFunction | Phased | None],
+) -> None:
+    """Inject `resources = base - outer_cost` for an NEGM regime with a cost.
+
+    The NEGM stacked-carry lift is exact only when the inner resources are
+    affine in the declared outer cost with coefficient exactly `-1`. That is
+    not a property finitely many probes can certify on an arbitrary user
+    function, so pylcm performs the subtraction itself: the user supplies the
+    cost-free base `<resources>_before_outer_cost`, and the resources function
+    every consumer sees — the inner EGM solve, the budget constraint, the
+    parent's child-resources read — is composed here, at the single point
+    before validation and processing. Mutates `functions` in place; regimes
+    whose solver is not `NEGM`, or whose `outer_cost` is `None`, are untouched
+    (they define the resources function directly).
+
+    Raises:
+        ModelInitializationError: If the regime defines the resources function
+            itself, or the cost-free base or the declared cost is missing.
+
+    """
+    solver = user_regime.solver
+    if not isinstance(solver, NEGM) or solver.outer_cost is None:
+        return
+    resources_name = solver.inner.resources
+    base_name = f"{resources_name}_before_outer_cost"
+    cost_name = solver.outer_cost
+
+    if resources_name in functions:
+        msg = (
+            f"Regime '{regime_name}' defines the resources function "
+            f"'{resources_name}' alongside a declared `NEGM.outer_cost` "
+            f"('{cost_name}'). With a declared outer cost the resources "
+            f"function is composed by pylcm as `{base_name} - {cost_name}`, "
+            "so its affine use of the cost holds by construction — define "
+            f"the cost-free base '{base_name}' instead of '{resources_name}'."
+        )
+        raise ModelInitializationError(msg)
+    if base_name not in functions:
+        msg = (
+            f"Regime '{regime_name}' declares `NEGM.outer_cost` "
+            f"('{cost_name}') but no cost-free resources base "
+            f"'{base_name}'. pylcm composes the resources function as "
+            f"`{base_name} - {cost_name}` — declare the base function."
+        )
+        raise ModelInitializationError(msg)
+    if cost_name not in functions:
+        msg = (
+            f"NEGM.outer_cost '{cost_name}' is not a declared function of "
+            f"regime '{regime_name}'. The credited outer cost must be a "
+            "regime function reading only the durable state, the outer "
+            "post-decision, and params."
+        )
+        raise ModelInitializationError(msg)
+
+    base_annotation = _return_annotation(functions[base_name])
+    cost_annotation = _return_annotation(functions[cost_name])
+
+    @with_signature(
+        args={base_name: base_annotation, cost_name: cost_annotation},
+        return_annotation=base_annotation,
+    )
+    def composed_resources(**kwargs: FloatND) -> FloatND:
+        return kwargs[base_name] - kwargs[cost_name]
+
+    composed_resources.__name__ = resources_name
+    functions[resources_name] = cast("UserFunction", composed_resources)
+
+
+def _return_annotation(func: UserFunction | Phased | None) -> str:
+    """Return a function's stringified return annotation, defaulting to `FloatND`.
+
+    The composed resources function copies its producers' annotations so the
+    DAG's annotation-consistency check stays satisfied.
+    """
+    if not callable(func):
+        return "FloatND"
+    annotations = ensure_annotations_are_strings(get_annotations(func))
+    return annotations.get("return", "FloatND")
+
+
 def _merge_derived_categoricals(
     *,
     regime_name: RegimeName,