Anchor consumption grid lower bound to consumption_floor parameter#8
Open
hmgaudecker wants to merge 55 commits intomainfrom
Open
Anchor consumption grid lower bound to consumption_floor parameter#8hmgaudecker wants to merge 55 commits intomainfrom
hmgaudecker wants to merge 55 commits intomainfrom
Conversation
Consumption is now declared as `IrregSpacedGrid(n_points=N)` (no fixed points). Callers inject log-spaced gridpoints from `consumption_floor` to $300k via `aca_model.consumption_grid. inject_consumption_points(params=..., model=...)` before solving. This means the lowest consumption choice equals the per-iteration floor, removing a degree of freedom from the grid and eliminating the previous mismatch where c < floor was a legal grid choice. Requires pylcm support for runtime-supplied points on continuous action grids (PR OpenSourceEconomics/pylcm#338). aca-model CI now installs pylcm from the matching `feature/runtime-action-grids` branch. Other changes: - `consumption_grid.py`: new module with `compute_consumption_points` and `inject_consumption_points` helpers. - `benchmark.get_benchmark_params(*, model=None)`: when `model` is given, returns params with consumption points injected. - `benchmark.get_benchmark_initial_conditions`: switch from `.start` / `.stop` to `to_jax().min()` / `.max()` so it works on both `LinSpacedGrid` and `PiecewiseLinSpacedGrid` (the AIME grid is now piecewise; this was a pre-existing bug surfacing as `AttributeError`). Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>
`utility_scale_factor` was registered as a regime function returning a (n_pref_types,) array, then re-indexed by `pref_type` inside `bequest` and `utility`. pylcm broadcasts function outputs to per-cell scalars before consumption, so that `[pref_type]` indexing produced silent NaN in the dead regime's V — surfaced as the all-NaN failure on the ASV benchmark. Mirror the `discount_factor` pattern: take the state as input, return a per-cell scalar. Drop the `[pref_type]` indexing on `utility_scale_factor` from `utility` and `bequest` (those still index the params-Series `consumption_weight` and `coefficient_rra`, which is the supported pattern — only DAG function outputs are pre-broadcast). The matching pylcm validator (PR #338) now raises a clear `RegimeInitializationError` when a function output is consumed via state-indexing in a downstream consumer; this aca-model change is the fix that lets the dead regime construct under that validator. Tests in `test_preferences.py` and `test_model_components.py` updated to pass scalar `utility_scale_factor` and supply the new `pref_type` arg to `utility_scale_factor`. Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>
Reduce n_assets_batch_size from 2 to 1 in MODEL_CONFIG so the assets state axis is streamed one slice at a time, lowering peak GPU memory during solve on the V100-PCIE-16GB. Benchmark grid config is unchanged. Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>
The OOM at production grid sizes came from pylcm's deferred diagnostics flush in solve_brute (`_emit_deferred_diagnostics` materialising a fused per-period reduction graph at end-of-solve), not from per-period peak. Halving the assets batch did not address that; reverting so the production loop runs at its previous throughput. Workaround for the diagnostics OOM lives in aca-estimation's simulate tasks (log_level="off"). Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>
5 tasks
PR #339's per-period `block_until_ready` made the OOM surface inside the loop instead of at the post-loop diagnostic flush, but the 7.26 GiB allocation request was the same — it isn't the diagnostic accumulator, it's a real per-period `max_Q_over_a` working set at production grid sizes (`n_consumption=70`, `n_assets=24`, `n_aime=12`, plus the per-target next-V gather across reachable regimes). Cutting the assets-axis chunk back to 1 reduces the per-kernel peak. Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>
Production solve allocates a per-period Q intermediate of shape `(non-assets-states × actions)` per assets-batch slot. With `n_assets_batch_size=1` we already chunk that axis to the minimum; the remaining outer-state product (aime × wage_res × hcc × pref_type × health × ...) times the action grid still pushes past the V100 16 GB once `pref_type` is split off into its own partition lift, which removes a free factor that previously thinned the kernel. Add a sibling `n_aime_batch_size` knob (default 1, 0 in `BENCHMARK_GRID_CONFIG`) and thread it through both AIME grid types in `_build_aime_grid`. AIME has 12 prod gridpoints in the LinSpaced fallback and 32 in the PiecewiseLinSpaced production path, so a unit batch shrinks the live Q intermediate by roughly that factor — enough headroom to land back inside V100 memory. Pairs with the pylcm-side fix that stops `_DiagnosticRow` pinning per-period V templates in device memory (lazy-solve-diagnostics branch). The diagnostic leak masked the underlying batching gap; once it's gone, the Q intermediate is the next thing to size for the device. Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>
The grid floor already tracks the per-iteration `consumption_floor` parameter; the ceiling was a hardcoded 300k constant. Surface it as a fixed param via a marker function (`consumption_grid_upper_bound`) so callers can declare the bracket per model creation, and read it back at inject time from each regime's `resolved_fixed_params`. The marker function's output is intentionally unused — its only job is to put `max_consumption` in the regime params template so pylcm's fixed-param machinery captures it. dags.tree pruning drops the call at solve / simulate. Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>
The runtime-upper-bound change requires every caller to supply `max_consumption` via `fixed_params`; estimation tasks (e.g. `task_simulate_aca`) hit a `KeyError` mid-pipeline because they construct the model from data-derived `fixed_params` that have no reason to mention a grid bracket. Centralise the default in both `baseline.model.create_model` and `aca.model.create_model` so existing callers keep working with the prior 300k bracket and only opt-in callers need to override.
Lets callers opt in to pylcm's simulate-AOT path (`Model(n_subjects=...)`) without bypassing the aca-model factories.
The aca-model factories now require `n_subjects` as a kw-only int with no default — there's never a good reason for an aca-model caller to leave it unspecified, and silently letting it default to `None` (= no AOT, lazy-compile path) was exactly how the simulate-AOT benefit went unused on the prod estimation loop. Forcing each caller to make a deliberate choice catches that. Tests pass `n_subjects=1` for bare `get_params_template()` / shock-grid-inspection paths that never simulate.
… to Model The marker-function-via-DAG pattern didn't survive pylcm's pruning: `consumption_grid_upper_bound`'s output is unused, so dags.tree drops it before its `max_consumption` parameter reaches the params template, and `broadcast_to_template` has nowhere to put the value. Result: `resolved_fixed_params["max_consumption"]` was always missing, `inject_consumption_points` raised KeyError. Sidestep pylcm's params machinery for this knob: - Drop the `consumption_grid_upper_bound` marker function and the `_with_max_consumption_default` helper. - Add `max_consumption: float` (kw-only, required, no default) to all three factories: `baseline.create_model`, `aca.create_model`, `create_benchmark_model`. - Each factory attaches the value directly to the returned `Model` instance (`model.max_consumption = ...`). - `inject_consumption_points` reads `model.max_consumption` directly. No defaults — every caller passes the bracket explicitly.
Adds `MAX_CONSUMPTION = 300_000.0` to `baseline/regimes/_common.py` next to the other grid bounds (assets `stop=500_000.0`, AIME `stop=8_000.0`). The two `create_model` factories and `create_benchmark_model` no longer take `max_consumption` as a kwarg; each factory reads the constant directly and attaches it onto `model.max_consumption`. `inject_consumption_points` is unchanged — it still reads `model.max_consumption` (the legitimate consumer that combines it with the per-iteration `consumption_floor`). Routed via the Model attribute rather than `fixed_params` because pylcm validates fixed_params keys against the regime DAG and rejects entries no function consumes (`InvalidParamsError: Unknown keys: ['max_consumption']`). Also pins the pylcm CI ref to 6c610d1 — the squash-merge of pylcm #341 (int32 lock-in) into feat/simulate-aot-n-subjects — to make this build deterministic against pylcm drift. Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>
With consumption now declared as `IrregSpacedGrid(n_points=N)` and points filled at runtime from `geomspace(consumption_floor, max_consumption, N)`, the grid clusters densely just above `consumption_floor`. At the lowest-asset / highest-OOP-shock corner, those near-floor consumption choices push `next_assets = cash_on_hand - OOP - consumption` slightly below the assets grid's old lower bound (`0` for the bare model, `-max_annual_labor_income` when wage_params are available). Out-of-bounds interpolation of next-period V then injects NaN, which propagates back through E[V] and eventually fails `validate_V`. Symptom on the production solve: `Value function at age 93 in regime 'retiree_oamc_forced_forcedout': 7317 of 207360 values are NaN`, with the `[NOTE]` showing E[V] NaN concentrated at the lowest assets indices and the highest hcc_transitory shock. Subtract `MAX_CONSUMPTION` from the assets floor to give a worst-case single-period drain margin. With 24 linspace points spanning the wider range, the per-point density change is negligible; the dead state and the bare-model fallback get the margin too. The asymmetry fix is the cheapest one — no change to the consumption grid type, no change to per-iteration parameters, no new constraints. Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>
This reverts commit 714fee0.
Two new DAG functions in canwork & ss != "forced" regimes: - target_his(his, labor_supply, is_medicaid_eligible): HIS class of the surviving target regime. Mirrors the cross-HIS branches inside _make_transition_canwork (tied → nongroup when stopping work, Medicaid override → nongroup). - imputed_pension_wealth_next_period(next_aime, target_his, period, ...): computes pw_next_imputed = benefit_imputed(next_pia, next_period, target_his) · epdv_constant_pension[next_period] using bare-name parameters into 1-period-shifted views of the imputation arrays (`*_next_period`). Inlining is required because pylcm's AST shape inference doesn't trace nested calls into pensions.benefit. next_assets continues to consume pension_assets_adjustment, which now sees a real imputed_pension_wealth_next_period via the DAG (previously fixed to 0.0 in aca-estimation). The chained dependency next_aime → imputed_pension_wealth_next_period → pension_assets_adjustment is unblocked by pylcm exempting next_<state> names from fixed_param extraction (PR pylcm#342). Also drops pension_assets_adjustment from borrowing_constraint: a negative correction at a cross-HIS transition can leave no feasible action and inject `-inf` into V via `argmax_and_max(initial=-inf, where=F_arr)`, which then cancels with `0 * -inf = NaN`. The correction is a post-decision shift on next-period assets and must not gate the current consumption choice.
…iod key The frozen benchmark_params.pkl was generated when aca-estimation's _assemble_params.py still wrote the placeholder `fp["imputed_pension_wealth_next_period"] = 0.0` into fixed_params. Now that the regime registers `imputed_pension_wealth_next_period` as a DAG function (pension imputation correction in 4ae4446), pylcm's `_resolve_fixed_params` rejects the stale key with `InvalidParamsError: Unknown keys: ['imputed_pension_wealth_next_period']`. Drop the key on load so the snapshot stays valid. Regenerating `benchmark_params.pkl` end-to-end would also remove it; the filter is a no-op for a fresh snapshot.
The frozen `benchmark_params.pkl` predates aca-data's `_shift_one_period_forward` change, so the 1-period-shifted views the pension correction consumes are missing. Synthesise them on load with the same transformation aca-data applies. Regenerating the snapshot end-to-end would also produce the keys; this filter is a no-op for a fresh snapshot.
target_his is a DAG function returning an HealthInsuranceState int, used to index 2D imputation arrays inside imputed_pension_wealth_next_period. pylcm needs the categorical mapping declared so array_from_series can reshape (age, target_his)-indexed Series correctly. Mirrors the existing 'his' entry — same enum class.
The shifted imputation arrays (`imp_*_next_period`) are consumed by `imputed_pension_wealth_next_period(target_his, period, ...)`. pylcm's `_validate_and_reorder_levels` matches Series MultiIndex level names against the function's parameter names, so the level needs to be `target_his`, not `his`.
…sion chain) state_transitions["assets"] becomes a per-target dict. The dead target gets a simpler `next_assets_terminal` (cash + transfers - consumption - oop) without the `pension_assets_adjustment` chain, because: 1. There is no future for a dead agent — the imputation correction is meaningless. 2. `pension_assets_adjustment` consumes `imputed_pension_wealth_next_period` which consumes `next_aime`. The dead per-target transitions don't include `next_aime` (dead has no aime state), so dags can't resolve it and pylcm leaks `next_aime` into the kernel signature with no value to pass. Non-dead targets keep `assets_and_income.next_assets` (full version with the pension correction).
The pension imputation correction's `imputed_pension_wealth_next_period` indexes shifted arrays via `arr[period, target_his]`, where `target_his` is a DAG output (computed by `health_insurance.target_his` on nongroup/tied/retiree regimes), not a state. pylcm reads the level name `target_his` off the function body via AST inference and rejects matching `pd.Series` fixed_params unless `target_his` is declared as a derived categorical. Production `task_simulate_baseline` calls `create_model(...)` directly, which previously only forwarded the user's `derived_categoricals` arg. The benchmark module was masking this by injecting target_his via `_DERIVED_CATEGORICALS`. Move the declaration to `create_model` itself so the correction works in production without per-caller setup. Tighten the param annotation: pylcm's `Model.derived_categoricals` is a flat `Mapping[str, DiscreteGrid]`, never the nested form.
Same fix as baseline.model.create_model e1a3eb2: ACA variant model creation also takes its own path through `Model(...)`, so the production `task_simulate_aca_*` flows hit the same "Unrecognised indexing parameter 'target_his'" error after the pension correction landed. Move the derived-categorical declaration into the function itself rather than relying on per-caller setup. Tighten the param annotation to match pylcm's flat `Mapping[str, DiscreteGrid]`.
…ation Asserts that `validate_initial_conditions` admits a subject placed at `assets = -1_000_000` in `retiree_nomc_inelig_canwork` under the benchmark model. Encodes the economic story: with the consumption floor / transfer system, any past assets level is representable — `c = c_floor` is always feasible because `transfers` tops up cash-on-hand to the floor. The test passes today on benchmark params; it doesn't reproduce the gpu-01 failure (production-side, separate setup loaded by `aca-estimation`'s `assemble_fixed_params`). Kept as a permanent regression guard so a future change that re-introduces a constraint shape that rejects extreme negatives is caught immediately at benchmark scale. Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>
…ellation
The expression `cash_on_hand + transfers` suffers float32 catastrophic
cancellation when `|cash_on_hand|` is much larger than `consumption_floor`.
For a subject at $-1{,}000{,}000$ in starting assets:
cash_on_hand ≈ -1e6 (dominated by assets)
transfers = max(0, c_floor - cash_on_hand) ≈ c_floor + 1e6
cash_on_hand + transfers ≈ c_floor ± 0.1 (fp32 error at 1e6 magnitude)
The lowest grid `c` is exactly `c_floor`. With unfavorable rounding,
`c_floor <= c_floor - 0.1` is False — every action gets rejected and
`validate_initial_conditions` raises. This is exactly the failure
gpu-01 hit on `task_simulate_aca_*`: the per-constraint diagnostic
showed `borrowing_constraint = False` (rejects every action by itself)
while `positive_leisure = True`.
The algebraic identity `cash_on_hand + transfers == max(cash_on_hand,
floor)` (where `floor = c_floor * equivalence_scale`) holds exactly
because `transfers` is defined as `max(0, floor - cash_on_hand)`.
Substituting in:
cash_on_hand + max(0, floor - cash_on_hand)
= max(cash_on_hand, cash_on_hand + floor - cash_on_hand)
= max(cash_on_hand, floor)
The `max` form has no cancellation: it returns `floor` exactly when
`cash_on_hand << floor`, and `cash_on_hand` exactly otherwise. Switch
the constraint to take `consumption_floor` and `equivalence_scale`
directly and compute `floor = consumption_floor * equivalence_scale`
in-line.
Add a precision-specific unit test asserting `c = c_floor` is admitted
at `cash_on_hand = -$1M` in fp32. The pre-existing benchmark-based
regression guard (`test_extreme_negative_assets_subject_passes_
validation`) didn't catch the bug because benchmark params land on the
favorable side of the rounding; the new test exercises the exact
cancellation case.
Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>
The previous pin (6c610d1, "Lock integer dtype to int32 end-to-end")
predates pylcm #342, so the test_initial_conditions_extreme_assets
test (and any other test that solves a benchmark regime carrying the
pension-imputation correction) raised:
InvalidParamsError: Missing required parameter:
'retiree_nomc_inelig_canwork__imputed_pension_wealth_next_period__next_aime'
#342's `regime_template` change exempts `next_<state>` references
inside transition signatures from `fixed_param` extraction, which the
correction's `imputed_pension_wealth_next_period(next_aime, ...)`
signature relies on. The new pin tracks `feat/simulate-aot-n-subjects`,
which carries #342, #339, #340 (n_subjects API used by
`create_benchmark_model`), and the per-constraint validation
diagnostic.
Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>
Production failure root cause: `consumption_floor` is a Python fp64 float (≈ 1597.0921419521899); `consumption` arrives from the model's fp32 grid (`jnp.geomspace(consumption_floor, ...)`), quantized to 1597.0921630859375 — one fp32 ulp above the input. Without an explicit dtype cast on the floor, `consumption_floor * equivalence_scale` keeps its fp64 type, the comparison promotes to fp64, and the lowest grid point evaluates as 1597.0921630859375 > 1597.0921419521899 → False. Constraint rejects every action. Cast `consumption_floor` to `consumption.dtype` before the multiply so both sides of the `max` use the same precision. Constraint then admits c=c_floor by exact equality in fp32. Diagnosed via the per-constraint admissibility table (pylcm 838473e/ e4cae2a): production showed `borrowing_constraint=False` at modest asset levels (e.g. -$42k), where neither cash_on_hand magnitude nor NaN propagation could explain the rejection. Local repro pinned the ulp mismatch. Add `test_borrowing_constraint_admits_c_floor_with_python_float_floor` as a regression guard at the precise production scenario. Drop the debug script; it served its purpose. Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>
`jnp.asarray(consumption_floor, dtype=consumption.dtype)` quantized the Python-float `consumption_floor` to the action grid's dtype to match the fp32-quantized consumption grid, so the `c == c_floor` boundary compared as exact equality. The pylcm canonical-float boundary cast (#345) routes every continuous-grid `to_jax()` through `canonical_float_dtype()`. Under `jax_enable_x64=True` (set in `aca_model/__init__.py`) that's `fp64`, so the action grid no longer quantizes the floor and Python-float / grid-value cannot disagree on dtype in the first place. Drop the regression test pinned to the cast workaround; the `max(cash_on_hand, floor)` cancellation guard and the full validate- initial-conditions integration test stay in place. Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>
hmgaudecker
force-pushed
the
feature/runtime-consumption-points
branch
from
8cb2936 to
c895bd9
Compare
`from tests.helpers.social_security import …` collided with the sibling `tests/__init__.py` packages in aca-data and aca-estimation when pytest collected from the aca-dev workspace root — whichever `tests` package got imported first shadowed the others. Use a relative import so each test module resolves its own helpers package unambiguously. Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>
Reverts the relative-import attempt and instead removes the empty tests/__init__.py (which was colliding with aca-data and aca-estimation's identically named stubs across the aca-dev workspace). A new tests/conftest.py prepends the tests directory to sys.path so `from helpers.social_security import ...` resolves unambiguously. Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>
Cleanup driven by pylcm's canonical-float boundary cast (#345). With every input pinned to fp64 under `jax_enable_x64=True` (which `aca_model/__init__.py` sets at import), aca-side precision workarounds no longer have a hook. Source: - `borrowing_constraint`: switch from `consumption <= max(cash_on_hand, floor)` to `consumption <= cash_on_hand + transfers`. The two are algebraically identical (`cash_on_hand + transfers == max(cash_on_hand, floor)`); the `max` form was justified by float32 catastrophic cancellation at extreme negative cash_on_hand, which cannot occur under fp64. The constraint now consumes `transfers` directly instead of recomputing `consumption_floor * equivalence_scale` — `transfers` is already a DAG node, so the resolved interface is shorter. Defaults dropped (callers must pass everything explicitly): - `aca_model.benchmark.create_benchmark_model`: `pref_type_grid`. - `aca_model.benchmark.get_benchmark_params`: `model`. - `aca_model.benchmark.get_benchmark_initial_conditions`: `n_subjects`, `seed`. - `aca_model.baseline.model.create_model`: `fixed_params`, `wage_params`, `derived_categoricals`, `grid_config`, `pref_type_grid`. - `aca_model.aca.model.create_model`: `policy`, `fixed_params`, `wage_params`, `derived_categoricals`, `grid_config`. - `aca_model.baseline.regimes.build_all_regimes`: same five. - `aca_model.aca.regimes.build_all_regimes`: same four. - `aca_model.baseline.regimes._common.build_grids`: same four. - Drop `GRID_CONFIG` import where it was only used as a default value. Tests: - New `tests/helpers/model.py` exposes `make_baseline_model` and `make_aca_model` factories that wrap `create_model` with `None` for every optional input. Tests that don't need fixed params reach the factories through the helper rather than spelling out six `None`s each. Production code stays default-free. - New `test_benchmark_simulate_obeys_borrowing_constraint`: pins the invariant `consumption <= cash_on_hand + transfers` on every alive row of the benchmark simulation. Catches a regression that drops the constraint from a regime, replaces transfers with something looser, or lets an action grid skip the floor. - `test_initial_conditions_extreme_assets`: drop the fp32-specific cancellation regression test (the runtime no longer reaches that path); replace with a pair of unit tests for the new `borrowing_constraint(consumption, cash_on_hand, transfers)` signature.
…e cash
The `consumption <= cash_on_hand + transfers` form (algebraically
identical to `consumption <= max(cash_on_hand, floor)`) rounds short by
sub-ULP at extreme `|cash_on_hand|` ~ 1e6 — for HRS-bottom-coded
subjects at `assets=-$1{,}000{,}000$`, the additive RHS comes in at
`floor - 5.7e-11` (fp64), flipping the kink-boundary `<=` for the
lowest consumption gridpoint. Production task_simulate_aca_no_mandate
on HPC fails at validate_initial_conditions for those subjects.
The `max(cash_on_hand, floor)` form has no cancellation and returns
`floor` exactly when `cash_on_hand < floor`. This is a general
floating-point precision concern at extreme operands, not an
fp32-specific workaround. Docstring updated accordingly.
Reverts the signature back to
`(consumption, cash_on_hand, consumption_floor, equivalence_scale)`.
Tests:
- `test_borrowing_constraint_admits_floor_at_million_dollar_negative_cash`:
unit-level reproducer of the production failure — passes only with
the `max` form.
- The two new `_at_floor` / `_above_post_transfer_resources` unit tests
switch back to the new signature.
- `test_benchmark_simulate_obeys_borrowing_constraint`: post-hoc check
uses `max(cash_on_hand, floor)` rather than `cash_on_hand +
transfers` (the additive form has the same sub-ULP issue and would
spuriously trip on the same rows).
`jnp.geomspace(consumption_floor, max_consumption, num=n)` returns `consumption_floor * r^0 == consumption_floor` mathematically, but some XLA backends drift the first point by sub-ULP. CUDA at n=70 produces `consumption_floor + 2.27e-13`. The borrowing_constraint compares `consumption[0]` against `max(cash_on_hand, consumption_floor)` and any positive drift above `consumption_floor` flips the kink- boundary `<=` for subjects with very negative cash — explaining the HPC-only `task_simulate` failures (~250 subjects) that didn't reproduce on CPU. Pin the first gridpoint back to `consumption_floor` after geomspace. The same drift exists at the upper end (`pts[-1] != max_consumption` exactly) but doesn't flip any constraint comparison, so it's left alone. `tests/test_consumption_grid.py` parametrises the invariant over `n_points = 5, 16, 64, 70, 100` so a future XLA / JAX upgrade that introduces drift at any of these counts surfaces here rather than at `validate_initial_conditions` on HPC.
Sweeps in the dtype-barrier polish, simulate AOT-during-solve, and the persistence/benchmark fixes from feat/canonical-float-dtype. Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>
hmgaudecker
force-pushed
the
feature/runtime-consumption-points
branch
2 times, most recently
from
9dd1e2f to
5a15184
Compare
Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>
hmgaudecker
force-pushed
the
feature/runtime-consumption-points
branch
from
5a15184 to
9e39a06
Compare
…cstrings Bare `consumption` is the raw-$ household action that gets divided by `equivalence_scale` to compute the per-equivalent input to utility. Renaming it `consumption_unequiv` (and the local `equiv_cons` → `consumption_equiv`) makes the equiv/unequiv distinction explicit at every reference. `consumption_weight` (the CES weight α) and `coefficient_rra` carry no equiv/unequiv meaning and stay as-is. Renames applied (case-sensitive substring + word-boundary token): - consumption (action) → consumption_unequiv - consumption_floor → consumption_unequiv_floor - max_consumption / MAX_CONSUMPTION → max_consumption_unequiv / MAX_CONSUMPTION_UNEQUIV - n_consumption_gridpoints → n_consumption_unequiv_gridpoints - inject_consumption_points → inject_consumption_unequiv_points - _compute_consumption_points → _compute_consumption_unequiv_points - average_consumption → average_consumption_unequiv - equiv_cons → consumption_equiv - module file consumption_grid.py → consumption_unequiv_grid.py - test file test_consumption_grid.py → test_consumption_unequiv_grid.py Also drop pref_type_grid docstring references to `DispatchStrategy`, `PARTITION_SCAN`, `PARTITION_VMAP`, and `FUSED_VMAP` from baseline/model.py, baseline/regimes/__init__.py, and baseline/regimes/_common.py — that pylcm enum was never merged. `benchmark.py` and `_benchmark_data/` are out of scope for this commit; their `_STALE_FIXED_KEYS` and `_rename_his_level` workarounds plus the DispatchStrategy docstrings in benchmark.py will be cleaned up when the benchmark snapshot is regenerated and benchmarks become a true special case of the production model. The benchmark.py imports get the minimum updates needed for it to load. `tests/test_initial_conditions_extreme_assets.py::test_extreme_negative_assets_subject_passes_validation` is skipped because it uses the benchmark snapshot, which still carries `average_consumption` (renamed to `average_consumption_unequiv`). Unskip after the snapshot is regenerated. Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>
The frozen `benchmark_params.pkl` is now produced by `aca-dev/scripts/regen_benchmark_params.py` against the current aca-data + aca-estimation + aca-model code. With a fresh snapshot: - `consumption_unequiv_floor` and `average_consumption_unequiv` keys match the post-rename model — no more old-key carryover. - `imp_*_next_period` shifted views come from aca-data's `_shift_one_period_forward` directly — no synthesis needed at load. - `imputed_pension_wealth_next_period` is no longer a fixed param (regimes resolve it via DAG) — no filter needed. - pref-type-indexed Series are pre-truncated to BenchmarkPrefType's two rows — no truncation needed at load. Drop from `benchmark.py`: - `_STALE_FIXED_KEYS` filter - `_SHIFTED_IMPUTATION_KEYS` + `_add_shifted_imputation_arrays` synthesis - `_shift_one_period_forward` + `_rename_his_level` helpers - `_truncate_pref_type_indexed` - The `DispatchStrategy` / `PARTITION_SCAN` / `PARTITION_VMAP` docstring references — that pylcm enum was never merged. Drop the unused `pandas` import. Unskip `test_extreme_negative_assets_subject_passes_validation` — the benchmark snapshot it loads now matches the model's expected keys. Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>
Major refactor splitting the consumption floor into a per-equivalent estimation parameter (`consumption_equiv_floor`, ~$1.6k/yr) and a household-$ DAG-derived value (`consumption_unequiv_floor = consumption_equiv_floor * equivalence_scale`). Same for consumption: `consumption_unequiv` is the action variable, `consumption_equiv = consumption_unequiv / equivalence_scale` is a new DAG function fed into `utility`. `transfers` and `borrowing_constraint` now consume `consumption_unequiv_floor` (DAG output) directly — no inline multiplication. `utility` consumes `consumption_equiv` directly — no inline division. `MAX_CONSUMPTION_UNEQUIV` is no longer attached as a dynamic Model attribute; `inject_consumption_unequiv_points` imports it as a module constant. The runtime grid's first two points are pinned exactly to the singles' floor and the married floor (`* 2 ** exponent`), with geomspace from there to `MAX_CONSUMPTION_UNEQUIV` — so both household-floor levels land on a feasible action regardless of sub-ULP drift in `jnp.geomspace`. Pref-type-indexed Series renamed to plural form (`consumption_weights`, `coefficients_rra`); the per-cell scalars (after `[pref_type]` indexing) take the singular form. Drops the `alpha` / `gamma` aliases. `scaled_bequest_weight` continues to consume the per-cell scalars (pylcm broadcasts pref-type-indexed Series to per-cell scalars before consumption). Also tighten `_build_per_target_*` annotations to `dict[RegimeName, ...]`, drop stale `consumption_unequiv` prose in docstrings/comments where bare "consumption" was meant, and update all callers + tests for the new signatures. The benchmark snapshot is regenerated against the new keys/structure. Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>
Pylcm's typing aliases (Period, IntND) are now strict — Python ints no longer satisfy them. Wrap kwarg literals in test_pensions, test_pension_integration, test_social_security, test_ss_benefit_integration, and test_model_components. Hoist top-level integer constants (PERIOD, his/old_his/new_his) to jnp.int32. Convert _RATIO_NP to pd.Series at the compute_di_dropout_scale call site. aca-model: drop the leftover MAX_CONSUMPTION_UNEQUIV import + the test_max_consumption_unequiv_attached_from_canonical_constant test (the dynamic Model attribute is gone). Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>
- `next_assets_terminal` → `next_assets_when_dead` (definition + register site in `_build_per_target_next_assets`). - `consumption_unequiv_floor`, `transfers`, `consumption_equiv`, `utility`: drop stale formula / cross-reference paragraphs. - `borrowing_constraint` + tests: HRS phrasing → "large negative values of `assets`"; drop `pension_assets_adjustment` paragraph. - `inject_consumption_unequiv_points`: walk the canonical `consumption_unequiv` action on each non-terminal regime, raise if it is missing or not a runtime-points IrregSpacedGrid. - `_compute_consumption_unequiv_points` docstring: drop the jax-arithmetic paragraph. - `baseline/model.py`: replace the verbose `target_his` block with one line. - Pre-existing PLC0415 in-function imports moved to module-level. Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>
…e_reference_hours → reference_hours `utility_scale_factor` is a multiplicative constant — shifting the calibration target age rescales every utility uniformly, so the distinction between `scale_reference_age` and `reference_age` is unobservable. Drop the redundant parameter and the `fixed_cost_of_work_age_trend * age_offset` term it gated. Rename `scale_reference_hours` → `reference_hours` for consistency with `reference_age`. Regenerated benchmark snapshot and updated test_preferences.py regression values to reflect the new `average_leisure = 4500` (was 4000 under the old `age_offset=10` shift). Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>
- `preferences.utility` → `preferences.u_working_life` (CRRA + leisure aggregator; `leisure` stays a DAG input). - `utility.retired` → `preferences.u_retired` (forcedout — computes `leisure_retired` inline). - `utility.dead` → `preferences.u_dead` (terminal bequest wrapper around `preferences.bequest`). - Delete `agent/utility.py`; update `select_utility` dispatch. The `u_*` prefix disambiguates regime-utility entry points from auxiliary scalars (`utility_scale_factor`, `discount_factor`, `scaled_bequest_weight`). Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>
Add `consumption_weight(consumption_weights, pref_type)` and `coefficient_rra(coefficients_rra, pref_type)` mirroring `discount_factor`. Register both on every regime in `build_common_functions` and on the dead regime. Drop `pref_type`, `consumption_weights`, `coefficients_rra` from `u_working_life`, `u_retired`, `u_dead`, `bequest`, `utility_scale_factor`: each now takes the scalar `consumption_weight` / `coefficient_rra` as a DAG-resolved input. `scaled_bequest_weight` already consumed scalars — no change. aca-data / aca-estimation param shapes are unchanged: pref-type- indexed Series for `consumption_weights` / `coefficients_rra` still flow through fixed_params; the new DAG functions resolve them per-cell at runtime. Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>
…goricals
`create_model` (baseline + ACA) now requires `fixed_params`,
`wage_params`, `derived_categoricals`, and `pref_type_grid`. No `None`
defaults, no `or {}` shortcuts. Same for `build_all_regimes` and
`build_grids`; the AIME piecewise-fallback and the static assets
floor are gone.
Add `aca_model.baseline.derived_categoricals.BASE_DERIVED_CATEGORICALS`
constant — `target_his` lives there now and both `create_model`
factories merge it into the caller's `derived_categoricals` so callers
no longer maintain a per-file `base_derived = {"target_his": ...}`
block.
Expose `get_hcc_persistent_shock(grid_config)` /
`get_hcc_persistent_grid_points(grid_config)` so callers (aca-data
`task_predicted_hcc_insurer`, aca-estimation `_assemble_params`) can
derive the shock without a bare model.
Benchmark snapshot now bundles `wage_params` alongside `fixed_params`
and `params`. `get_benchmark_params` returns a 3-tuple. Test helpers
in `tests/helpers/model.py` thread the snapshot through with the
benchmark `derived_categoricals` set.
Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>
…gimeSpec TypedDict `precompute_targets(spec: Mapping[str, str]) -> dict[str, int]` becomes `precompute_target_regimes(spec: RegimeSpec) -> MappingProxyType[str, int]`. The four-axis `RegimeSpec` `TypedDict` makes the valid keys (`his`, `mc`, `ss`, `canwork`) explicit and propagates through every spec-consuming helper. Helper renames make the dispatch verb read as "per target *regime*", matching the new function name: - `_build_per_target_next_assets` → `_build_per_target_regime_next_assets` - `_build_per_target_health` → `_build_per_target_regime_health` - `_build_per_target_claimed_ss` → `_build_per_target_regime_claimed_ss` - `_build_per_target_lagged_labor_supply` → `_build_per_target_regime_lagged_labor_supply` Every `targets` local that's the result of the precompute call is now `target_regimes`. Spec parameters on the HIS-specific builders (`_retiree`, `_tied`, `_nongroup`) and ACA overrides also adopt the new `RegimeSpec` type. Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>
`ca66ba9` was an interim commit on the pre-squash `distributed` lineage; the squash-force-push to that branch (#346) made it unreachable. Bump to pylcm main HEAD, which carries the canonical-float-dtype work this branch depends on. Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>
Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>
… factory magic - consumption_unequiv → consumption_dollars (state, action, params, module, test module). The dollar / dollars suffix replaces the unequiv / equiv pair on the in-$ side; `consumption_equiv` (the utility-equivalized form) keeps its name. - consumption_dollar_grid.py → consumption_dollars_grid.py. - `_DERIVED_CATEGORICALS` dicts now include `target_his` alongside `his`, `good_health`, `is_married`, `pref_type`. The previous `BASE_DERIVED_CATEGORICALS` constant + per-factory merge is gone — callers pass the full dict. - `RegimeSpec` TypedDict fields tightened to `Literal[...]` for the four axes. - `fixed_params` annotations switched to `lcm.typing.UserParams` (the pylcm-side alias) across `create_model`, `build_all_regimes`, `build_grids`. - `u_working_life` → `u_can_work`; `u_retired` → `u_cannot_work`. - `_build_per_target_regime_next_assets` → `_build_per_target_regime_assets`; ordering of dispatch site and definition aligned (assets → health → claimed_ss → lagged_labor_supply). `lagged_supply_transition` → `lagged_labor_supply_transition`. - `MAX_CONSUMPTION_DOLLARS` docstring: drop the in-line rationale, add a TODO referencing pylcm#348. - `aca.create_model` return doc → "pylcm Model" (no "ACA-specific" qualifier). - CI: pylcm pin uses `@main` (PyPI release lags so the git pin stays). 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
consumptionbecomes anIrregSpacedGrid(n_points=N)action grid; the actual gridpoints are supplied per estimation iteration via the newaca_model.consumption_grid.inject_consumption_points(params=..., model=...)helper, log-spaced from the per-iterationconsumption_floorparameter to $300k.c < floorwas a legal grid choice. The agent can now never consume below the transfer floor.Depends on
feature/runtime-action-gridsuntil that lands.Notes
get_benchmark_initial_conditionsnow usesto_jax().min()/.max()instead of.start/.stopso it works on the piecewise AIME grid (pre-existing AttributeError once AIME became piecewise).Test plan
pixi run -e tests-cpu tests aca-model/tests/(199 passed)🤖 Generated with Claude Code