Let H consume DAG function outputs#324
Open
hmgaudecker wants to merge 8 commits intoimprove/parallel-aot-compilationfrom
Open
Let H consume DAG function outputs#324hmgaudecker wants to merge 8 commits intoimprove/parallel-aot-compilationfrom
hmgaudecker wants to merge 8 commits intoimprove/parallel-aot-compilationfrom
Conversation
pylcm's default `_default_H(utility, E_next_V, discount_factor)` and any user custom H used to only see values from `states_actions_params` (states + actions + flat regime params). This meant a regime function named like an H argument (e.g. a `discount_factor` DAG fn that indexes a `pref_type` state) was silently invisible to H — `H_kwargs` never picked it up, and the solve failed with `TypeError: _default_H() missing 1 required positional argument`. Extend both the solve-phase `Q_and_F` and the diagnostic `compute_intermediates` paths so that any `_H_accepted_params` name that is also in `regime.functions` (other than `utility`, `feasibility`, or `H`) is computed as a DAG output at call time and merged into `H_kwargs`. The two pools (user params / DAG outputs) are disjoint by construction (`create_regime_params_template` excludes function names from every non-H function's params dict), so no precedence rule is needed. Factor the build-time compilation into `_get_h_dag_fn`, which returns `None` when H does not need DAG outputs — the common case — keeping the existing scalar-H path fully unchanged. Extend `tests/solution/test_custom_aggregator.py` to cover the new pattern: reuse its `_make_model` factory with a `with_pref_type` switch that adds a three-category `pref_type` state and a `discount_factor(pref_type, discount_factor_by_type)` DAG function. Assert the default H then produces value functions monotone in the per-type discount factor. Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>
Previously the lcm example's Mahler-Yum model only used the mean discount factor at solve time; `create_inputs` returned a `discount_factor_type` array but no consumer plumbed it into the model. This change moves the heterogeneity inside the regime, matching the pattern in `/home/hmg/econ/aca-dev/aca-model` and the reference test `tests/solution/test_custom_aggregator.py::test_dag_output_feeds_default_h_monotone_in_discount_factor`. - New `DiscountType` (`small`/`large`) categorical. - `discount_type` added to `ALIVE_REGIME.states` and `DEAD_REGIME.states` as a fixed state (`state_transitions = None`). - New DAG function `discount_factor(discount_type, discount_factor_by_type)` registered on `ALIVE_REGIME.functions`; pylcm's default Bellman aggregator picks up the scalar via the DAG-output path added in `2945765 "Let H consume DAG function outputs"`. - `utility` accepts an unused `discount_type` arg to satisfy pylcm's state-usage check (same pattern as the reference test). - `create_inputs` now takes `beta: dict[str, float]` and returns `(params, initial_states)` (dropped the third tuple element). Computes `discount_factor_by_type = [mean - std, mean + std]` once and exposes it under `params["discount_factor"]`. Callers (`bench_mahler_yum`, regression fixture generator / test, docs example) pass `**START_PARAMS` without filtering and consume a single params dict. Also dedup a `test_nan_diagnostics_end_to_end` block duplicated by an earlier cascade merge. Note: `tests/data/regression_tests/f64/mahler_yum_simulation.pkl` needs regeneration on a GPU host — the simulation output now contains a `discount_type` column with both small/large types mixed. Follow-up commit. Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>
A prior cascade merge with -X ours kept two copies of `_build_nan_model` + `test_nan_diagnostics_end_to_end`. Ruff F811 on CI's pre-commit.ci caught it. Keep the first, drop the second. Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>
The simulation output now includes a `discount_type` column and reflects one single solve+simulate with both discount types handled via the `discount_type` state (not two separate solves concatenated). Regenerated on local GPU with `pixi run --environment tests-cuda13 python tests/data/regression_tests/generate_benchmark_data.py --precision 64`. `test_regression_mahler_yum` passes on the new fixture. Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>
Ty narrowed the dict's value type from the three initial entries, then rejected `functions["discount_factor"] = discount_factor_from_type`. Explicit `dict[str, Callable]` annotation fixes the inference; no behavioral change. Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>
CLAUDE.md / AGENTS.md: "`func` for callable abbreviations — use `func`, `func_name`, `func_params` (never `fn`)." Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>
Benchmark comparison (main → HEAD)Comparing
|
A state reachable only via `H → discount_factor(state, ...)` was wrongly flagged as unused, forcing `utility` to declare a fake `discount_type` arg with `# noqa: ARG001`. Extend the state-usage walk to include H-DAG target names as reachability targets so states consumed only via H's DAG dependencies count as used. Factor `get_h_dag_target_names` + `get_h_accepted_params` into `regime_building/h_dag.py` so Q_and_F (runtime) and the validator share the same target-set definition. Drop the fake `discount_type` arg in the Mahler-Yum example and the `utility_with_pref_type` workaround in the reference test. Terminal regimes still need to declare unused states in their utility (they have no H), so `dead_utility` retains its workaround. Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>
hmgaudecker
added a commit
to mj023/Replication_MY2024
that referenced
this pull request
Apr 18, 2026
pylcm now walks H's DAG dependencies during state-usage validation (OpenSourceEconomics/pylcm#324), so `utility` no longer needs to accept `discount_type` as an unused kwarg to satisfy the check. Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>
Add five regression tests exercising H consuming a continuous state, continuous action, discrete action, discrete state, and all kinds simultaneously. Rewrite _get_h_dag_func's docstring so it states the full contract positively — H may name any argument supported by regime functions — instead of implying states/actions are excluded. Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
1 participant
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Summary
Let pylcm's default Bellman aggregator (and any user-supplied H) consume DAG function outputs as H inputs, not only user-supplied scalars (through
params/fixed_params). This unlocks the "per-type discount factor" pattern used in aca-dev — the discount factor is a DAG function indexed by a state (e.g.pref_type,discount_type), and_default_H(utility, E_next_V, discount_factor)resolves it automatically.Before this PR, a
discount_factorregime function was silently invisible to H:H_kwargsonly included scalar values from user params, so the solve failed withTypeError: _default_H() missing 1 required positional argument.Core change (
src/lcm/regime_building/Q_and_F.py):get_Q_and_Fandget_compute_intermediatesso that every_H_accepted_paramsname that is also inregime.functions(other thanutility,feasibility, orH) is computed as a DAG output at call time and merged intoH_kwargs._get_h_dag_fn, which returnsNonewhen H does not need DAG outputs — the common case — so the scalar-H path stays fully unchanged.create_regime_params_templateexcludes function names from every non-H function's params dict, so no precedence rule is needed.Reference test (
tests/solution/test_custom_aggregator.py):_make_modelvia a newwith_pref_typeswitch that adds a three-categorypref_typestate and wiresdiscount_factor(pref_type, discount_factor_by_type)as a DAG function.test_dag_output_feeds_default_h_monotone_in_discount_factorasserts V is monotone in the per-type β, confirming the DAG-output path works end-to-end with the default H.Apply to the Mahler & Yum example:
DiscountTypecategorical (small/large).discount_typeas a fixed state on both alive and dead regimes (state_transitions["discount_type"] = None).discount_factor(discount_type, discount_factor_by_type)onALIVE_REGIME.functions.create_inputsnow returns(params, initial_states)— dropped the third tuple element; exposes the two-valued β array viaparams["discount_factor"]["discount_factor_by_type"]. Callers (benchmark, regression generator/test, docs example) simplify to a single dict and a single simulate call.utilityaccepts an unuseddiscount_typearg to satisfy pylcm's state-usage check (same pattern as the reference test).tests/data/regression_tests/f64/mahler_yum_simulation.pklto reflect the new simulation output (single solve;discount_typecolumn present; both types in one stream).Also ships a minor cleanup: drops a duplicated
_build_nan_model+test_nan_diagnostics_end_to_endblock left over from an earlier cascade merge.Test plan
tests/solution/test_custom_aggregator.py::test_dag_output_feeds_default_h_monotone_in_discount_factor— new coverage for the DAG-output → H path.tests/test_regression_test.py::test_regression_mahler_yum— passes on the regenerated fixture (GPU + f64).tyclean,prekclean.🤖 Generated with Claude Code