Anchor consumption grid lower bound to consumption_floor parameter

.pre-commit-config.yaml

@@ -5,7 +5,7 @@ repos:
       - id: check-hooks-apply
       - id: check-useless-excludes
   - repo: https://github.com/tox-dev/pyproject-fmt
-    rev: v2.19.0
+    rev: v2.21.1
     hooks:
       - id: pyproject-fmt
   - repo: https://github.com/lyz-code/yamlfix
@@ -47,7 +47,7 @@ repos:
     hooks:
       - id: yamllint
   - repo: https://github.com/astral-sh/ruff-pre-commit
-    rev: v0.15.6
+    rev: v0.15.12
     hooks:
       - id: ruff-check
         args:

src/aca_model/aca/model.py

@@ -8,42 +8,46 @@
 from typing import Any
 
 from lcm import AgeGrid, DiscreteGrid, Model
+from lcm.typing import UserParams
 
 from aca_model.aca import PolicyVariant
 from aca_model.aca.regimes import build_all_regimes
 from aca_model.baseline.regimes import RegimeId
-from aca_model.config import GRID_CONFIG, MODEL_CONFIG, GridConfig
+from aca_model.config import MODEL_CONFIG, GridConfig
 
 
 def create_model(
     *,
-    policy: PolicyVariant = PolicyVariant.ACA,
-    fixed_params: Mapping[str, Any] | None = None,
-    wage_params: Mapping[str, Any] | None = None,
-    derived_categoricals: Mapping[str, DiscreteGrid | Mapping[str, DiscreteGrid]]
-    | None = None,
-    grid_config: GridConfig = GRID_CONFIG,
+    n_subjects: int,
+    policy: PolicyVariant,
+    fixed_params: UserParams,
+    wage_params: Mapping[str, Any],
+    derived_categoricals: Mapping[str, DiscreteGrid],
+    grid_config: GridConfig,
+    pref_type_grid: DiscreteGrid,
 ) -> Model:
     """Create an ACA policy variant model.
 
     Args:
-        policy: Which ACA policy combination to apply.
-        fixed_params: Parameters to fix at model creation time. These are
-            partialled into compiled functions and removed from the params
-            template. Pass data-derived constants here; only estimation
-            parameters should go through `model.simulate(params=...)`.
+        n_subjects: Forwarded to `lcm.Model(n_subjects=...)`.
+        policy: Which ACA policy combination to apply (e.g.
+            `PolicyVariant.ACA`).
+        fixed_params: Parameters to fix at model creation time. Pass
+            data-derived constants here; only estimation parameters
+            should go through `model.simulate(params=...)`.
         wage_params: Data-derived wage profile dict (`log_ft_wage_mean`,
             `log_ft_wage_std`, `adj_wage_hours_*`) used only at grid-build
             time to size the assets-floor to `-max_annual_labor_income`.
             Not routed to the pylcm Model.
-        derived_categoricals: Extra categorical mappings for derived variables
-            not in the model's state/action grids. Needed when `fixed_params`
-            contains `pd.Series` indexed by DAG function outputs.
-        grid_config: Continuous-grid point counts. Defaults to production
-            values.
+        derived_categoricals: Categorical mappings for `pd.Series`
+            fixed_params index levels that aren't model state/action
+            grids — `target_his`, `his`, `good_health`, `is_married`,
+            `pref_type`.
+        grid_config: Continuous-grid point counts.
+        pref_type_grid: Pref-type `DiscreteGrid`.
 
     Returns:
-        pylcm Model with ACA-specific function overrides.
+        pylcm Model.
 
     """
     ages = AgeGrid(
@@ -56,13 +60,15 @@ def create_model(
         grid_config=grid_config,
         fixed_params=fixed_params,
         wage_params=wage_params,
+        pref_type_grid=pref_type_grid,
     )
 
     return Model(
         regimes=regimes,
         ages=ages,
         regime_id_class=RegimeId,
         description=f"Structural retirement model ({policy.name})",
-        fixed_params=fixed_params or {},
+        fixed_params=fixed_params,
         derived_categoricals=derived_categoricals,
+        n_subjects=n_subjects,
     )

src/aca_model/aca/regimes/__init__.py

@@ -4,25 +4,30 @@
 from collections.abc import Mapping
 from typing import Any
 
-from lcm import Regime
+from lcm import DiscreteGrid, Regime
+from lcm.typing import UserParams
 
 from aca_model.aca.health_insurance import PolicyVariant
 from aca_model.aca.regimes._overrides import apply_aca_overrides
 from aca_model.baseline.regimes import build_all_regimes as baseline_build_all_regimes
 from aca_model.baseline.regimes._common import REGIME_SPECS
-from aca_model.config import GRID_CONFIG, GridConfig
+from aca_model.config import GridConfig
 
 
 def build_all_regimes(
-    policy: PolicyVariant,
-    grid_config: GridConfig = GRID_CONFIG,
     *,
-    fixed_params: Mapping[str, Any] | None = None,
-    wage_params: Mapping[str, Any] | None = None,
+    policy: PolicyVariant,
+    grid_config: GridConfig,
+    fixed_params: UserParams,
+    wage_params: Mapping[str, Any],
+    pref_type_grid: DiscreteGrid,
 ) -> dict[str, Regime]:
     """Build all 19 regimes with ACA policy overrides."""
     regimes = baseline_build_all_regimes(
-        grid_config, fixed_params=fixed_params, wage_params=wage_params
+        grid_config=grid_config,
+        fixed_params=fixed_params,
+        wage_params=wage_params,
+        pref_type_grid=pref_type_grid,
     )
     result = {}
     for name, regime in regimes.items():

src/aca_model/aca/regimes/_overrides.py

@@ -8,11 +8,12 @@
 
 from aca_model.aca import health_insurance as aca_hi
 from aca_model.aca.health_insurance import PolicyVariant
+from aca_model.baseline.regimes._common import RegimeSpec
 
 
 def apply_aca_overrides(
     functions: dict,
-    spec: dict[str, str],
+    spec: RegimeSpec,
     policy: PolicyVariant,
 ) -> None:
     """Override baseline functions with ACA versions in-place.

src/aca_model/agent/assets_and_income.py

@@ -35,41 +35,74 @@ def cash_on_hand(
     return assets + after_tax_income + ssi_benefit - hic_premium
 
 
-def transfers(
-    cash_on_hand: FloatND,
-    consumption_floor: float,
+def consumption_dollars_floor(
+    consumption_equiv_floor: float,
     equivalence_scale: FloatND,
 ) -> FloatND:
-    """Government transfers to enforce consumption floor.
+    """Per-household $-floor on consumption."""
+    return consumption_equiv_floor * equivalence_scale
 
-    tr = max{0, C_min * equivalence_scale - cash_on_hand}
-    """
-    floor = consumption_floor * equivalence_scale
-    return jnp.maximum(0.0, floor - cash_on_hand)
+
+def transfers(
+    cash_on_hand: FloatND,
+    consumption_dollars_floor: FloatND,
+) -> FloatND:
+    """Government transfers to enforce the consumption floor."""
+    return jnp.maximum(0.0, consumption_dollars_floor - cash_on_hand)
 
 
 def next_assets(
     cash_on_hand: FloatND,
     transfers: FloatND,
     pension_assets_adjustment: FloatND,
-    consumption: ContinuousAction,
+    consumption_dollars: ContinuousAction,
     oop_costs: FloatND,
 ) -> ContinuousState:
-    """Compute beginning-of-next-period assets.
+    """Compute beginning-of-next-period assets for non-terminal targets.
 
     OOP health costs are deducted here (not from cash_on_hand) so that the
     consumption choice does not condition on the HCC shock realization.
     """
     return (
-        cash_on_hand + transfers + pension_assets_adjustment - consumption - oop_costs
+        cash_on_hand
+        + transfers
+        + pension_assets_adjustment
+        - consumption_dollars
+        - oop_costs
     )
 
 
-def borrowing_constraint(
-    consumption: ContinuousAction,
+def next_assets_when_dead(
     cash_on_hand: FloatND,
     transfers: FloatND,
-    pension_assets_adjustment: FloatND,
+    consumption_dollars: ContinuousAction,
+    oop_costs: FloatND,
+) -> ContinuousState:
+    """Compute beginning-of-next-period assets for the dead/terminal target.
+
+    No `pension_assets_adjustment` term: with no future, there is no
+    next-period pension wealth to impute against. Avoiding the dependency
+    also keeps the `dead` per-target transition's DAG free of `next_aime`
+    (which would otherwise need to come from a transition `dead` does not
+    have, since `aime` is not a state in the terminal regime).
+    """
+    return cash_on_hand + transfers - consumption_dollars - oop_costs
+
+
+def borrowing_constraint(
+    consumption_dollars: ContinuousAction,
+    cash_on_hand: FloatND,
+    consumption_dollars_floor: FloatND,
 ) -> BoolND:
-    """Consumption cannot exceed available resources (no borrowing)."""
-    return consumption <= cash_on_hand + transfers + pension_assets_adjustment
+    """Consumption cannot exceed post-transfer resources.
+
+    Post-transfer resources are `max(cash_on_hand, consumption_dollars_floor)`:
+    the transfer system tops `cash_on_hand` to the floor when below,
+    otherwise resources are unchanged. The algebraic identity is
+    `cash_on_hand + transfers == max(cash_on_hand, floor)`; the `max`
+    form is preferred because the additive form rounds to `floor + ε`
+    (with `|ε| ~ ULP(|cash_on_hand|)`) at extreme cash, which flips
+    the kink-boundary comparison at large negative values of `assets`.
+    The `max` form returns `floor` exactly.
+    """
+    return consumption_dollars <= jnp.maximum(cash_on_hand, consumption_dollars_floor)