Register policy and solver helpers only where consumed

src/aca_model/aca/regimes/_overrides.py

@@ -8,6 +8,7 @@
 
 from aca_model.aca import health_insurance as aca_hi
 from aca_model.aca.health_insurance import PolicyVariant
+from aca_model.baseline import health_insurance as baseline_hi
 from aca_model.baseline.regimes._common import RegimeSpec
 from aca_model.environment import taxes
 
@@ -55,6 +56,7 @@ def apply_aca_overrides(
     )
 
     if has_medicaid_expansion:
+        functions["aca_magi"] = baseline_hi.aca_magi
         functions["is_medicaid_eligible"] = aca_hi.is_medicaid_eligible
 
     if has_subsidies and spec["his"] == "nongroup" and spec["mc"] == "nomc":
@@ -64,8 +66,9 @@ def apply_aca_overrides(
         # dict, not a DAG function — no entry needed here.
         functions["plan_premium"] = aca_hi.community_rated_premium
         # The community-rated premium replaces the risk-rated one, so the
-        # risk-rated premium's zero-profit intercept has no consumer left.
+        # risk-rated inputs have no consumer left.
         functions.pop("private_premium_intercept", None)
+        functions.pop("predicted_hcc_insurer", None)
         functions["hic_premium_subsidy"] = aca_hi.premium_subsidy
         # Split at the bracket, which takes finitely many values, so the
         # out-of-pocket bill can be averaged per bracket.

src/aca_model/baseline/regimes/_common.py

@@ -845,12 +845,10 @@ def build_model_functions(*, solver: SolverName = "brute_force") -> dict:
     # assets adjustment, both of which scale by the marginal income tax rate.
     functions["marginal_tax_rate"] = taxes.marginal_rate
 
-    # HIC premium
-    functions["predicted_hcc_insurer"] = health_insurance.hcc_insurer_predicted
-
     # Transfers
     functions["consumption_dollars_floor"] = assets_and_income.consumption_dollars_floor
-    functions["transfers"] = assets_and_income.transfers
+    if solver == "brute_force":
+        functions["transfers"] = assets_and_income.transfers
     functions["consumption_equiv"] = preferences.consumption_equiv
 
     return functions
@@ -999,10 +997,8 @@ def build_common_functions(spec: RegimeSpec, *, subtract_bill: bool = False) ->
         if spec["mc"] == "oamc"
         else health_insurance.is_disabled_from_health
     )
-    # MAGI for the ACA Medicaid expansion track; pruned in the baseline DAG.
-    functions["aca_magi"] = health_insurance.aca_magi
-
     # Swapped per policy variant by the ACA overlay, hence regime-level
+    functions["predicted_hcc_insurer"] = health_insurance.hcc_insurer_predicted
     functions["is_medicaid_eligible"] = health_insurance.is_medicaid_eligible
     functions["after_tax_income"] = taxes.after_tax_income
     if subtract_bill:

tests/test_aca_plan.py

@@ -83,6 +83,21 @@ def test_aca_non_group_premium_is_community_rated(
     assert _aca_premium(policy, spousal_income) == pytest.approx(expected)
 
 
+@pytest.mark.parametrize("policy", list(PolicyVariant))
+def test_aca_non_group_risk_prediction_is_registered_only_without_reform(
+    policy: PolicyVariant,
+) -> None:
+    """The community-rated premium does not register a risk prediction node."""
+    functions = {**build_model_functions(), **_build_functions(REGIME_SPECS[REGIME])}
+    apply_aca_overrides(functions, REGIME_SPECS[REGIME], policy)
+    expected = (
+        health_insurance.hcc_insurer_predicted
+        if policy == PolicyVariant.ACA_ONLY_MEDICAID_EXPANSION
+        else None
+    )
+    assert functions.get("predicted_hcc_insurer") is expected
+
+
 RISK_RATED_PREMIUM_PARAMS = frozenset(
     {
         "premium_markup",

tests/test_model_creation.py

@@ -13,11 +13,15 @@
 from aca_model.aca import health_insurance as aca_hi
 from aca_model.aca.health_insurance import PolicyVariant
 from aca_model.aca.regimes import build_all_regimes as _build_aca_regimes
+from aca_model.agent import assets_and_income
 from aca_model.agent.preferences import BenchmarkPrefType
+from aca_model.baseline import health_insurance as baseline_hi
 from aca_model.baseline.regimes import REGIME_SPECS, RegimeId
 from aca_model.baseline.regimes import build_regime as _build_regime
 from aca_model.baseline.regimes._common import (
+    SolverName,
     build_grids,
+    build_model_functions,
     build_model_state_transitions,
     build_model_states,
 )
@@ -290,6 +294,36 @@ def test_aca_model_creates_successfully() -> None:
     assert model.n_periods == 46
 
 
+@pytest.mark.parametrize("name", REGIME_SPECS)
+def test_baseline_regime_omits_expansion_magi(name: str) -> None:
+    """Baseline eligibility does not register the ACA expansion income node."""
+    assert "aca_magi" not in build_regime(name).functions
+
+
+@pytest.mark.parametrize("policy", list(PolicyVariant))
+def test_aca_regime_registers_magi_only_for_expansion(policy: PolicyVariant) -> None:
+    """The expansion eligibility consumer owns its MAGI income node."""
+    has_expansion = policy in (
+        PolicyVariant.ACA,
+        PolicyVariant.ACA_NO_MANDATE,
+        PolicyVariant.ACA_ONLY_MEDICAID_EXPANSION,
+    )
+    expected = baseline_hi.aca_magi if has_expansion else None
+    regimes = build_aca_regimes(policy)
+    assert {
+        name: regimes[name].functions.get("aca_magi") for name in REGIME_SPECS
+    } == dict.fromkeys(REGIME_SPECS, expected)
+
+
+@pytest.mark.parametrize("solver", ["brute_force", "dcegm", "nbegm"])
+def test_model_registers_transfers_only_for_brute_asset_laws(
+    solver: SolverName,
+) -> None:
+    """Savings-form asset laws already include transfers through resources."""
+    expected = assets_and_income.transfers if solver == "brute_force" else None
+    assert build_model_functions(solver=solver).get("transfers") is expected
+
+
 def test_aca_nongroup_inelig_has_real_functions() -> None:
     """Nongroup+nomc regimes get real ACA functions under ACA policy."""
     regimes = build_aca_regimes(PolicyVariant.ACA)