The counted-once check raises or misreports when a head's baseline value is NaN #2773

The counted-once check (policyengine_uk/utils/aggregate_impact/guard.py) perturbs each component from its baseline value, holds other heads at their baseline values, and measures each head's change from its baseline. #2759 covers responses that turn infinite or NaN from a finite baseline. When a head is already NaN at baseline, the check either stops with policyengine-core's ValueError or fails components that the head ignores, and its reset check reports a changed baseline.

Minimal case

One household on policyengine-core 3.33.12. tax_b is an input of 5. tax_a is sqrt(-1), so it is NaN at baseline and changes with no other head. gov_tax adds the heads.

import numpy as np
from policyengine_core.country_template import CountryTaxBenefitSystem
from policyengine_core.country_template.entities import Household
from policyengine_core.periods import YEAR
from policyengine_core.simulations import Simulation
from policyengine_core.variables import Variable
from policyengine_uk.utils.aggregate_impact.guard import counted_once
from policyengine_uk.utils.aggregate_impact.partition import parse_partition


def variable(name, formula=None, default=0.0):
    attrs = {"label": name, "value_type": float, "entity": Household,
             "definition_period": YEAR, "default_value": default}
    if formula is not None:
        attrs["formula"] = lambda household, period: formula(
            lambda v: household(v, period)
        )
    return type(name, (Variable,), attrs)


def simulation(formulas):
    system = CountryTaxBenefitSystem()
    for name, formula in formulas.items():
        system.update_variable(variable(name, formula, 5.0 if formula is None else 0.0))
    system.update_variable(variable("gov_tax", lambda f: sum(f(v) for v in formulas)))
    system.update_variable(variable("household_weight"))
    return Simulation(tax_benefit_system=system, situation={
        "persons": {"p": {}},
        "households": {"h": {"parents": ["p"], "household_weight": {"2026": 1.0}}},
    })


def partition(heads):
    return parse_partition({
        "schema_version": 1,
        "totals": {"gov_tax": {}},
        "tax_heads": [{"id": v, "label": v, "variable": v, "totals": ["gov_tax"],
                       "entity": "household", "reads": []} for v in heads],
    })


formulas = {"tax_b": None, "tax_a": lambda f: np.sqrt(-1 + 0 * f("tax_b"))}

What happens

These outputs come from #2759's branch at df997704d.

1. Perturbing a component whose baseline is NaN raises. counted_once(simulation(formulas), partition(["tax_a", "tax_b"]), 2026, 1e-6) adds 1 to tax_a through _perturb and _set_guard_value, and Core's Simulation.set_input refuses the NaN:

ValueError: Unable to set value for variable "tax_a", as the input contains NaN values.

2. Holding a head whose baseline is NaN raises. Add tax_c = tax_b, which moves with tax_b, and check only tax_b (variables_to_check=["tax_b"]). The check measures tax_c again with every other head held at its baseline value, and _perturb raises the same ValueError when it holds tax_a at NaN.

3. A head whose baseline is NaN fails a component it ignores. Without tax_c, checking only tax_b runs to the end. tax_a's change is NaN − NaN, so the check fails tax_b, although tax_a doesn't change with it:

gov_tax: tax_a moves, or is not finite, when tax_b does in fewer than 10 households
gov_tax: the model's gov_tax moves by other than its heads, or is not finite, when tax_b does in fewer than 10 households

4. The reset check treats NaN as differing from NaN. The same run returns reset_reproduces_baseline: False, though tax_a recomputes to NaN, as it started. The check compares each recomputed head with np.array_equal(heads_again[v], heads0[v]), and np.array_equal(np.array([np.nan]), np.array([np.nan])) returns False.

Cases 1, 2 and 4 behave the same on #2758 (a122f2355). There, case 3 reports only gov_tax: tax_a moves when tax_b does in fewer than 10 households.