Let H consume DAG function outputs

benchmarks/bench_mahler_yum.py

@@ -20,22 +20,13 @@ def _build(self):
             create_inputs,
         )
 
-        start_params_without_beta = {
-            k: v for k, v in START_PARAMS.items() if k != "beta"
-        }
-
         self.model = MAHLER_YUM_MODEL
-        common_params, initial_states, _discount_factor_type = create_inputs(
+        common_params, initial_states = create_inputs(
             seed=0,
             n_simulation_subjects=_N_SUBJECTS,
-            **start_params_without_beta,
+            **START_PARAMS,
         )
-        self.model_params = {
-            "alive": {
-                "discount_factor": START_PARAMS["beta"]["mean"],
-                **common_params,
-            },
-        }
+        self.model_params = {"alive": common_params}
         self.initial_conditions = {
             **initial_states,
             "regime": jnp.full(

docs/examples/mahler_yum_2024.md

@@ -27,35 +27,26 @@ from lcm_examples.mahler_yum_2024 import (
     create_inputs,
 )
 
-# Build inputs (params, initial states, discount factor types)
-start_params_without_beta = {k: v for k, v in START_PARAMS.items() if k != "beta"}
-common_params, initial_states, discount_factor_types = create_inputs(
+n_subjects = 1_000
+
+# Build inputs: per-subject initial states include `discount_type`
+# (small/large), and `params["discount_factor"]["discount_factor_by_type"]`
+# carries the two-element beta array that the `discount_factor` DAG
+# function indexes with the state.
+common_params, initial_states = create_inputs(
     seed=7235,
-    n_simulation_subjects=1_000,
-    **start_params_without_beta,
+    n_simulation_subjects=n_subjects,
+    **START_PARAMS,
 )
 
-beta_mean = START_PARAMS["beta"]["mean"]
-beta_std = START_PARAMS["beta"]["std"]
-
-# Select initial states with high discount factor type
-selected_ids_high = jnp.flatnonzero(discount_factor_types)
-initial_states_high = {
-    state: values[selected_ids_high] for state, values in initial_states.items()
-}
-
-# Solve and simulate for high discount factor type
+# One solve, one simulate — both discount types are handled inside the
+# regime via the `discount_type` state.
 result = MAHLER_YUM_MODEL.simulate(
-    params={
-        "alive": {
-            "discount_factor": beta_mean + beta_std,
-            **common_params,
-        },
-    },
+    params={"alive": common_params},
     initial_conditions={
-        **initial_states_high,
+        **initial_states,
         "regime": jnp.full(
-            selected_ids_high.shape[0],
+            n_subjects,
             MAHLER_YUM_MODEL.regime_names_to_ids["alive"],
         ),
     },

src/lcm/model_processing.py

@@ -24,6 +24,10 @@
 )
 from lcm.params.sequence_leaf import SequenceLeaf
 from lcm.regime import Regime
+from lcm.regime_building.h_dag import (
+    get_h_accepted_params,
+    get_h_dag_target_names,
+)
 from lcm.regime_building.processing import (
     InternalRegime,
     process_regimes,
@@ -205,6 +209,7 @@ def _validate_all_variables_used(regimes: Mapping[str, Regime]) -> list[str]:
     Each state or action must appear in at least one of:
     - The concurrent valuation (utility or constraints)
     - A transition function
+    - An H-DAG target function (a regime function whose output H consumes)
 
     Args:
         regimes: Mapping of regime names to regimes to validate.
@@ -219,6 +224,11 @@ def _validate_all_variables_used(regimes: Mapping[str, Regime]) -> list[str]:
         variable_names = set(regime.states) | set(regime.actions)
         user_functions = dict(regime.get_all_functions(phase="solve"))
 
+        h_accepted_params = get_h_accepted_params(user_functions)
+        h_dag_target_names = get_h_dag_target_names(
+            functions=user_functions, h_accepted_params=h_accepted_params
+        )
+
         targets = [
             "utility",
             *list(regime.constraints),
@@ -228,6 +238,7 @@ def _validate_all_variables_used(regimes: Mapping[str, Regime]) -> list[str]:
                 if name.startswith("next_")
                 and not getattr(user_functions[name], "_is_auto_identity", False)
             ),
+            *h_dag_target_names,
         ]
         reachable = get_ancestors(
             user_functions, targets=targets, include_targets=False

src/lcm/regime_building/Q_and_F.py

@@ -6,6 +6,7 @@
 from dags import concatenate_functions, with_signature
 from jax import Array
 
+from lcm.regime_building.h_dag import get_h_dag_target_names
 from lcm.regime_building.next_state import (
     get_next_state_function_for_solution,
     get_next_stochastic_weights_function,
@@ -135,6 +136,13 @@ def get_Q_and_F(
         exclude=frozenset(),
     )
 
+    # Resolve H arguments that are regime-function outputs (e.g. a
+    # `discount_factor` DAG function that indexes a per-type Series by a
+    # state). `None` when H only needs state/action/user-param values.
+    _h_dag_func = _get_h_dag_func(
+        functions=functions, h_accepted_params=_H_accepted_params
+    )
+
     @with_signature(
         args=arg_names_of_Q_and_F, return_annotation="tuple[FloatND, BoolND]"
     )
@@ -203,6 +211,8 @@ def Q_and_F(
         H_kwargs = {
             k: v for k, v in states_actions_params.items() if k in _H_accepted_params
         }
+        if _h_dag_func is not None:
+            H_kwargs |= _h_dag_func(**states_actions_params)
         Q_arr = _H_func(utility=U_arr, E_next_V=E_next_V, **H_kwargs)
 
         # Handle cases when there is only one state.
@@ -301,6 +311,9 @@ def get_compute_intermediates(
     _H_accepted_params = frozenset(
         get_union_of_args([_H_func]) - {"utility", "E_next_V"}
     )
+    _h_dag_func = _get_h_dag_func(
+        functions=functions, h_accepted_params=_H_accepted_params
+    )
 
     arg_names_of_compute_intermediates = _get_arg_names_of_Q_and_F(
         [
@@ -357,6 +370,8 @@ def compute_intermediates(
         H_kwargs = {
             k: v for k, v in states_actions_params.items() if k in _H_accepted_params
         }
+        if _h_dag_func is not None:
+            H_kwargs |= _h_dag_func(**states_actions_params)
         Q_arr = _H_func(utility=U_arr, E_next_V=E_next_V, **H_kwargs)
 
         return U_arr, F_arr, E_next_V, Q_arr, active_regime_probs
@@ -536,6 +551,54 @@ def _outer(**kwargs: Float1D) -> FloatND:
     )
 
 
+def _get_h_dag_func(
+    *,
+    functions: FunctionsMapping,
+    h_accepted_params: frozenset[str],
+) -> Callable[..., dict[str, Any]] | None:
+    """Compile a DAG that resolves H arguments computed by regime functions.
+
+    `H` may name any argument supported by regime functions: states,
+    actions, flat params, or outputs of other user-provided functions.
+    Names in H's signature are resolved at runtime from, in order:
+
+    1. `states_actions_params` (states, actions, and flat params — the
+       same scalar pool every regime function draws from), and
+    2. DAG-output functions, compiled here.
+
+    This helper handles only (2): for every name in H's signature that
+    is also a user-provided function, compile a DAG target so its
+    output can be merged into `H_kwargs` alongside the values supplied
+    by (1). If no such names exist, return `None`.
+
+    Args:
+        functions: Regime functions (user and generated).
+        h_accepted_params: Names H accepts beyond `utility` / `E_next_V`.
+
+    Returns:
+        A callable mapping `states_actions_params` kwargs to a dict of
+        the resolved DAG outputs, or `None` if H needs no DAG outputs.
+
+    """
+    dag_targets = tuple(
+        sorted(
+            get_h_dag_target_names(
+                functions=functions, h_accepted_params=h_accepted_params
+            )
+        )
+    )
+
+    if not dag_targets:
+        return None
+
+    return concatenate_functions(
+        functions={k: v for k, v in functions.items() if k != "H"},
+        targets=list(dag_targets),
+        return_type="dict",
+        enforce_signature=False,
+    )
+
+
 def _get_U_and_F(
     *,
     functions: FunctionsMapping,

src/lcm/regime_building/h_dag.py

@@ -0,0 +1,72 @@
+"""H's DAG-target bookkeeping, shared between runtime and validation.
+
+The default Bellman aggregator `H(utility, E_next_V, discount_factor)` —
+and any user-supplied H — may declare parameters that are not
+states/actions/user-params but are outputs of regime functions
+registered under the same name (e.g. a `discount_factor` DAG function
+that indexes a per-type Series by a `discount_type` state).
+
+This module exposes:
+
+- `get_h_accepted_params`: H's signature minus `utility` / `E_next_V`.
+- `get_h_dag_target_names`: those H parameters that are *also* regime
+  functions. Q_and_F compiles these into a runtime DAG;
+  `_validate_all_variables_used` uses them as reachability targets so
+  states consumed only via H's DAG dependencies count as "used".
+"""
+
+from collections.abc import Callable, Mapping
+from typing import Any
+
+from lcm.utils.functools import get_union_of_args
+
+
+def get_h_accepted_params(
+    functions: Mapping[str, Callable[..., Any]],
+) -> frozenset[str]:
+    """H's signature parameters, minus `utility` and `E_next_V`.
+
+    Empty when the regime has no `H` (terminal regimes).
+
+    Args:
+        functions: Mapping of regime function names to callables (user
+            and generated).
+
+    Returns:
+        Frozenset of parameter names H accepts beyond `utility` / `E_next_V`.
+
+    """
+    h_func = functions.get("H")
+    if h_func is None:
+        return frozenset()
+    return frozenset(get_union_of_args([h_func]) - {"utility", "E_next_V"})
+
+
+def get_h_dag_target_names(
+    *,
+    functions: Mapping[str, Callable[..., Any]],
+    h_accepted_params: frozenset[str],
+) -> frozenset[str]:
+    """Names of regime functions whose outputs H consumes via the DAG.
+
+    These are H's signature parameters that are also regime functions,
+    minus `H`, `utility`, `feasibility` (H cannot consume its own
+    output; `utility` is wired directly from `U_and_F`; `feasibility`
+    is never a legitimate H input).
+
+    Args:
+        functions: Mapping of regime function names to callables (user
+            and generated).
+        h_accepted_params: Names H accepts beyond `utility` / `E_next_V`
+            (typically the output of `get_h_accepted_params`).
+
+    Returns:
+        Frozenset of regime function names whose outputs are routed
+        into H at runtime.
+
+    """
+    return frozenset(h_accepted_params) & set(functions) - {
+        "H",
+        "utility",
+        "feasibility",
+    }

src/lcm_examples/mahler_yum_2024/__init__.py

@@ -9,6 +9,7 @@
     DEAD_REGIME,
     MAHLER_YUM_MODEL,
     START_PARAMS,
+    DiscountType,
     Education,
     Effort,
     Health,
@@ -26,6 +27,7 @@
     "DEAD_REGIME",
     "MAHLER_YUM_MODEL",
     "START_PARAMS",
+    "DiscountType",
     "Education",
     "Effort",
     "Health",