From d18c54c6b7caccd7ad4f6eb3845647ba79c7218f Mon Sep 17 00:00:00 2001 From: Max Ghenis Date: Tue, 3 Nov 2026 07:16:30 -0500 Subject: [PATCH 1/4] Accept decode_enums in Microsimulation.calculate_add and calculate_divide Population.__call__ passes decode_enums=False to calculate_add or calculate_divide for every options=[ADD] or options=[DIVIDE] read in a formula, as it has since 0abe43a60. Simulation's methods take that argument, but Microsimulation's overrides did not, so every such read raised a TypeError in a Microsimulation. Both overrides now take decode_enums, with Simulation's default of False, and pass it to super(). Simulation.calculate_add and calculate_divide accept the flag without using it, so every call that worked before returns the same values. Example tests run formulas that read with ADD and DIVIDE on a small country-template system. Hypothesis properties check that a Microsimulation's unweighted reads equal a Simulation's over generated data and read sequences, and that formula results equal the same arithmetic in NumPy. A strict xfail records that an Enum read with DIVIDE returns labels in a Microsimulation and an EnumArray in a Simulation. Co-Authored-By: Claude Opus 5.5 --- ...-microsim-add-divide-decode-enums.fixed.md | 1 + .../simulations/microsimulation.py | 10 +- .../core/test_microsimulation_option_reads.py | 150 ++++++++ ...t_microsimulation_option_reads_property.py | 110 ++++++ .../fixtures/microsimulation_option_reads.py | 357 ++++++++++++++++++ 5 files changed, 626 insertions(+), 2 deletions(-) create mode 100644 changelog.d/fix-microsim-add-divide-decode-enums.fixed.md create mode 100644 tests/core/test_microsimulation_option_reads.py create mode 100644 tests/core/test_microsimulation_option_reads_property.py create mode 100644 tests/fixtures/microsimulation_option_reads.py diff --git a/changelog.d/fix-microsim-add-divide-decode-enums.fixed.md b/changelog.d/fix-microsim-add-divide-decode-enums.fixed.md new file mode 100644 index 000000000..70fb64daa --- /dev/null +++ b/changelog.d/fix-microsim-add-divide-decode-enums.fixed.md @@ -0,0 +1 @@ +`Microsimulation.calculate_add` and `calculate_divide` now take `decode_enums`, so formulas that read a variable with `options=[ADD]` or `options=[DIVIDE]` compute in a `Microsimulation` instead of raising a `TypeError`. diff --git a/policyengine_core/simulations/microsimulation.py b/policyengine_core/simulations/microsimulation.py index af50a088c..4ee5758a6 100644 --- a/policyengine_core/simulations/microsimulation.py +++ b/policyengine_core/simulations/microsimulation.py @@ -87,8 +87,11 @@ def calculate_add( period: Period = None, map_to: str = None, use_weights: bool = True, + decode_enums: bool = False, ) -> MicroSeries: - values = super().calculate_add(variable_name, period) + # ``Population.__call__`` passes ``decode_enums`` with every ADD or + # DIVIDE read in a formula. + values = super().calculate_add(variable_name, period, decode_enums=decode_enums) if map_to is not None: source_entity = self.tax_benefit_system.get_variable( variable_name, check_existence=True @@ -105,8 +108,11 @@ def calculate_divide( period: Period = None, map_to: str = None, use_weights: bool = True, + decode_enums: bool = False, ) -> MicroSeries: - values = super().calculate_divide(variable_name, period) + values = super().calculate_divide( + variable_name, period, decode_enums=decode_enums + ) if map_to is not None: source_entity = self.tax_benefit_system.get_variable( variable_name, check_existence=True diff --git a/tests/core/test_microsimulation_option_reads.py b/tests/core/test_microsimulation_option_reads.py new file mode 100644 index 000000000..af82faca1 --- /dev/null +++ b/tests/core/test_microsimulation_option_reads.py @@ -0,0 +1,150 @@ +"""Formulas that read with ADD or DIVIDE compute in a Microsimulation. + +``Population.__call__`` passes ``decode_enums=False`` to ``calculate_add`` or +``calculate_divide`` for every ``options=[ADD]`` or ``options=[DIVIDE]`` read +in a formula. ``Simulation.calculate_add`` and ``calculate_divide`` take that +argument, but ``Microsimulation``'s overrides did not, so in a Microsimulation +every such read raised ``TypeError: Microsimulation.calculate_add() got an +unexpected keyword argument 'decode_enums'``. A Simulation computed the same +reads. +""" + +import dataclasses + +import numpy as np +import pytest +from microdf import MicroSeries + +from policyengine_core.model_api import ADD, DIVIDE, MONTH +from tests.fixtures.microsimulation_option_reads import ( + DATASET_YEAR, + READERS, + Inputs, + expected, + run, + simulations, +) + +EXAMPLE = Inputs( + household_sizes=(2, 1), + household_weight=(10.0, 20.0), + monthly_pay=( + tuple(1.0 * month for month in range(1, 13)), + tuple(2.0 * month for month in range(1, 13)), + tuple(5.0 * month for month in range(1, 13)), + ), + yearly_grant=(120.0, 240.0, 600.0), + household_monthly_fee=((5.0,) * 12, (0.0,) * 12), + household_yearly_grant=(12.0, 36.0), +) + +EXPECTED = { + "pay_added": [78.0, 156.0, 390.0], + "take_home_added": [58.5, 117.0, 292.5], + "grant_divided": [10.0, 20.0, 50.0], + "fee_added": [60.0, 0.0], + "pays_fee": [True, False], + "allowance_divided": [2.0, 6.0], + "members_pay_added": [234.0, 390.0], + "members_grant_divided": [30.0, 50.0], + "fee_and_pay_added": [294.0, 390.0], +} + + +def period_for(reader) -> str: + return "2022-06" if reader.definition_period == MONTH else DATASET_YEAR + + +def test_numpy_reference_gives_the_hand_computed_values(): + for reader in READERS: + np.testing.assert_array_equal( + expected(EXAMPLE, reader.__name__), EXPECTED[reader.__name__] + ) + + +@pytest.mark.parametrize("reader", READERS, ids=lambda reader: reader.__name__) +def test_formula_reading_with_an_option_computes_in_a_microsimulation(reader): + simulation, microsimulation = simulations(EXAMPLE) + name, period = reader.__name__, period_for(reader) + + result = microsimulation.calculate(name, period, use_weights=False) + + np.testing.assert_array_equal(result, EXPECTED[name]) + np.testing.assert_array_equal(result, simulation.calculate(name, period)) + + +def test_weighted_result_of_a_formula_reading_with_add_carries_its_weights(): + _, microsimulation = simulations(EXAMPLE) + + result = microsimulation.calculate("pays_fee", DATASET_YEAR) + + assert isinstance(result, MicroSeries) + np.testing.assert_array_equal(result, [True, False]) + np.testing.assert_array_equal(result.weights, [10.0, 20.0]) + + +@pytest.mark.parametrize( + "call", + [ + ("population", "person", "monthly_pay", DATASET_YEAR, ADD), + ("population", "person", "monthly_take_home", DATASET_YEAR, ADD), + ("population", "person", "yearly_grant", "2022-06", DIVIDE), + ("population", "household", "household_monthly_fee", DATASET_YEAR, ADD), + ("population", "household", "household_yearly_allowance", "2022-06", DIVIDE), + ("population", "household", "monthly_pay", DATASET_YEAR, ADD), + ("population", "household", "yearly_grant", "2022-06", DIVIDE), + ], + ids=lambda call: f"{call[1]}-{call[2]}-{call[4]}", +) +def test_population_read_with_an_option_matches_a_simulation(call): + simulation, microsimulation = simulations(EXAMPLE) + + result = run(microsimulation, call) + + assert isinstance(result, list), result + assert result == run(simulation, call) + + +@pytest.mark.parametrize("decode_enums", [False, True]) +@pytest.mark.parametrize( + "method,variable,period", + [ + ("calculate_add", "monthly_pay", DATASET_YEAR), + ("calculate_divide", "yearly_grant", "2022-06"), + ], +) +def test_microsimulation_add_and_divide_take_decode_enums( + method, variable, period, decode_enums +): + _, microsimulation = simulations(EXAMPLE) + calculate = getattr(microsimulation, method) + + unweighted_result = calculate( + variable, period, use_weights=False, decode_enums=decode_enums + ) + weighted_result = calculate(variable, period, decode_enums=decode_enums) + + np.testing.assert_array_equal( + unweighted_result, calculate(variable, period, use_weights=False) + ) + np.testing.assert_array_equal(weighted_result, unweighted_result) + np.testing.assert_array_equal( + weighted_result.weights, microsimulation.get_weights(variable, period) + ) + + +@pytest.mark.xfail( + strict=True, + reason=( + "Simulation.calculate_divide reads the variable with self.calculate, " + "which in a Microsimulation decodes enums by default, so a " + "Microsimulation returns labels where a Simulation returns an " + "EnumArray. No formula in policyengine-us reads an Enum this way." + ), +) +def test_enum_read_with_divide_in_a_microsimulation_matches_a_simulation(): + inputs = dataclasses.replace(EXAMPLE, yearly_tenure=("OWNER", "RENTER", "OWNER")) + simulation, microsimulation = simulations(inputs) + call = ("population", "person", "yearly_tenure", DATASET_YEAR, DIVIDE) + + assert run(microsimulation, call) == run(simulation, call) diff --git a/tests/core/test_microsimulation_option_reads_property.py b/tests/core/test_microsimulation_option_reads_property.py new file mode 100644 index 000000000..cb1186597 --- /dev/null +++ b/tests/core/test_microsimulation_option_reads_property.py @@ -0,0 +1,110 @@ +"""Property: a Microsimulation's unweighted ADD and DIVIDE reads equal a Simulation's. + +For any small dataset of households, people, weights and monthly and yearly +inputs, and any sequence of reads, a Microsimulation returns the same +unweighted values as a Simulation loaded from the same data, or raises the +same kind of error. The reads are formula results (each formula reads with +``options=[ADD]`` or ``options=[DIVIDE]``), the population calls formulas +make, ``calculate_add`` and ``calculate_divide`` with either value of +``decode_enums``, and plain reads at months, years and longer periods. Since +the Simulation ignores weights, weights never change an unweighted value. + +Formula results also equal the same arithmetic done in NumPy, so the two +simulations cannot agree on a wrong value. + +Enum variables are left out: +``test_enum_read_with_divide_in_a_microsimulation_matches_a_simulation`` +records how their reads differ. +""" + +import numpy as np +import pytest + +from policyengine_core.model_api import ADD, DIVIDE, MONTH +from tests.fixtures.microsimulation_option_reads import ( + DATASET_YEAR, + MONTHS, + READERS, + SOURCES, + Inputs, + expected, + run, + simulations, + unweighted, +) + +hypothesis = pytest.importorskip("hypothesis") +st = hypothesis.strategies + +# Whole amounts keep sums exact in float32, so a mismatch is never rounding. +AMOUNTS = st.integers(-10_000, 10_000).map(float) +WEIGHTS = st.floats(min_value=0, max_value=1e6, allow_nan=False) +PERIODS = ["2022", "2022-01", "2022-06", "2022-12", "2023", "month:2022-01:3"] +VARIABLES = sorted(variable.__name__ for variable in SOURCES + READERS) + +REQUESTS = st.one_of( + st.tuples( + st.just("calculate"), + st.sampled_from(VARIABLES), + st.sampled_from(PERIODS), + ), + st.tuples( + st.just("population"), + st.sampled_from(["person", "household"]), + st.sampled_from(VARIABLES), + st.sampled_from(PERIODS), + st.sampled_from([ADD, DIVIDE]), + ), + st.tuples( + st.sampled_from(["add", "divide"]), + st.sampled_from(VARIABLES), + st.sampled_from(PERIODS), + st.booleans(), + ), +) + + +@st.composite +def inputs(draw): + sizes = draw(st.lists(st.integers(1, 3), min_size=1, max_size=4)) + people, households = sum(sizes), len(sizes) + + def values(count): + return tuple(draw(st.lists(AMOUNTS, min_size=count, max_size=count))) + + return Inputs( + household_sizes=tuple(sizes), + household_weight=tuple( + draw(st.lists(WEIGHTS, min_size=households, max_size=households)) + ), + monthly_pay=tuple(values(12) for _ in range(people)), + yearly_grant=values(people), + household_monthly_fee=tuple(values(12) for _ in range(households)), + household_yearly_grant=values(households), + ) + + +@hypothesis.settings(max_examples=200, deadline=None) +@hypothesis.given(inputs(), st.lists(REQUESTS, min_size=1, max_size=6)) +def test_unweighted_microsimulation_reads_equal_simulation_reads(data, requests): + simulation, microsimulation = simulations(data) + for request in requests: + assert run(microsimulation, request) == run(simulation, request), request + + +@hypothesis.settings(max_examples=100, deadline=None) +@hypothesis.given( + inputs(), + st.sampled_from(READERS), + st.sampled_from(MONTHS), +) +def test_formula_reads_with_an_option_equal_numpy_arithmetic(data, reader, month): + period = month if reader.definition_period == MONTH else DATASET_YEAR + reference = expected(data, reader.__name__) + for simulation in simulations(data): + result = simulation.calculate(reader.__name__, period, **unweighted(simulation)) + np.testing.assert_allclose( + np.asarray(result, dtype=float), + np.asarray(reference, dtype=float), + rtol=1e-6, + ) diff --git a/tests/fixtures/microsimulation_option_reads.py b/tests/fixtures/microsimulation_option_reads.py new file mode 100644 index 000000000..4013c1864 --- /dev/null +++ b/tests/fixtures/microsimulation_option_reads.py @@ -0,0 +1,357 @@ +"""A small system whose formulas read variables with ADD and DIVIDE. + +Shared by ``tests/core/test_microsimulation_option_reads.py`` and its +Hypothesis property module, which pytest must be able to skip on its own when +Hypothesis is not installed. +""" + +from dataclasses import dataclass +from functools import lru_cache +from typing import Callable, Optional, Tuple + +import numpy as np +import pandas as pd + +from policyengine_core.country_template import ( + CountryTaxBenefitSystem, + Microsimulation, + Simulation, +) +from policyengine_core.country_template.entities import Household, Person +from policyengine_core.data import Dataset +from policyengine_core.enums import Enum +from policyengine_core.model_api import ADD, DIVIDE, MONTH, YEAR, Variable, add +from policyengine_core.simulations import Microsimulation as CoreMicrosimulation + +DATASET_YEAR = "2022" +MONTHS = [f"{DATASET_YEAR}-{month:02d}" for month in range(1, 13)] + + +# Inputs, and variables with formulas that the readers below read. + + +class monthly_pay(Variable): + value_type = float + entity = Person + definition_period = MONTH + label = "Monthly pay" + + +class monthly_take_home(Variable): + value_type = float + entity = Person + definition_period = MONTH + label = "Monthly take-home pay" + + def formula(person, period, parameters): + return person("monthly_pay", period) * 0.75 + + +class yearly_grant(Variable): + value_type = float + entity = Person + definition_period = YEAR + label = "Yearly grant" + + +class household_monthly_fee(Variable): + value_type = float + entity = Household + definition_period = MONTH + label = "Household monthly fee" + + +class household_yearly_grant(Variable): + value_type = float + entity = Household + definition_period = YEAR + label = "Household yearly grant" + + +class household_yearly_allowance(Variable): + value_type = float + entity = Household + definition_period = YEAR + label = "Household yearly allowance" + + def formula(household, period, parameters): + return household("household_yearly_grant", period) * 2 + + +class Tenure(Enum): + OWNER = "Owner" + RENTER = "Renter" + + +class yearly_tenure(Variable): + value_type = Enum + possible_values = Tenure + default_value = Tenure.RENTER + entity = Person + definition_period = YEAR + label = "Yearly tenure" + + +# Readers: formulas that read with ADD or DIVIDE, as 21 policyengine-us +# formulas do (``spm_unit("snap", period, options=[ADD]) > 0`` in +# mt_liheap_categorical_eligibility, for one). + + +class pay_added(Variable): + value_type = float + entity = Person + definition_period = YEAR + label = "Pay summed over the year" + + def formula(person, period, parameters): + return person("monthly_pay", period, options=[ADD]) + + +class take_home_added(Variable): + value_type = float + entity = Person + definition_period = YEAR + label = "Take-home pay summed over the year" + + def formula(person, period, parameters): + return person("monthly_take_home", period, options=[ADD]) + + +class grant_divided(Variable): + value_type = float + entity = Person + definition_period = MONTH + label = "Monthly share of the yearly grant" + + def formula(person, period, parameters): + return person("yearly_grant", period, options=[DIVIDE]) + + +class fee_added(Variable): + value_type = float + entity = Household + definition_period = YEAR + label = "Household fees summed over the year" + + def formula(household, period, parameters): + return household("household_monthly_fee", period, options=[ADD]) + + +class pays_fee(Variable): + value_type = bool + entity = Household + definition_period = YEAR + label = "Household paid a fee in the year" + + def formula(household, period, parameters): + return household("household_monthly_fee", period, options=[ADD]) > 0 + + +class allowance_divided(Variable): + value_type = float + entity = Household + definition_period = MONTH + label = "Monthly share of the household allowance" + + def formula(household, period, parameters): + return household("household_yearly_allowance", period, options=[DIVIDE]) + + +class members_pay_added(Variable): + value_type = float + entity = Household + definition_period = YEAR + label = "Members' pay summed over the year" + + def formula(household, period, parameters): + return household("monthly_pay", period, options=[ADD]) + + +class members_grant_divided(Variable): + value_type = float + entity = Household + definition_period = MONTH + label = "Monthly share of members' yearly grants" + + def formula(household, period, parameters): + return household("yearly_grant", period, options=[DIVIDE]) + + +class fee_and_pay_added(Variable): + value_type = float + entity = Household + definition_period = YEAR + label = "Household fees and members' pay summed over the year" + + def formula(household, period, parameters): + return add( + household, period, ["household_monthly_fee", "monthly_pay"], options=[ADD] + ) + + +SOURCES = [ + monthly_pay, + monthly_take_home, + yearly_grant, + household_monthly_fee, + household_yearly_grant, + household_yearly_allowance, +] +READERS = [ + pay_added, + take_home_added, + grant_divided, + fee_added, + pays_fee, + allowance_divided, + members_pay_added, + members_grant_divided, + fee_and_pay_added, +] + + +@lru_cache(maxsize=None) +def tax_benefit_system() -> CountryTaxBenefitSystem: + system = CountryTaxBenefitSystem() + for variable in SOURCES + READERS + [yearly_tenure]: + system.add_variable(variable) + return system + + +@dataclass(frozen=True) +class Inputs: + """Households in order, and each input's values. + + ``monthly_pay`` and ``household_monthly_fee`` hold 12 values (January to + December) for each person or household. + """ + + household_sizes: Tuple[int, ...] + household_weight: Tuple[float, ...] + monthly_pay: Tuple[Tuple[float, ...], ...] + yearly_grant: Tuple[float, ...] + household_monthly_fee: Tuple[Tuple[float, ...], ...] + household_yearly_grant: Tuple[float, ...] + yearly_tenure: Optional[Tuple[str, ...]] = None + + @property + def household_of_person(self) -> np.ndarray: + return np.repeat( + np.arange(len(self.household_sizes)), np.array(self.household_sizes) + ) + + +def dataframe(inputs: Inputs) -> pd.DataFrame: + """One row per person. Each member's row repeats its household's values.""" + household = inputs.household_of_person + first_member = np.r_[True, household[1:] != household[:-1]] + pay = np.array(inputs.monthly_pay, dtype=float) + fee = np.array(inputs.household_monthly_fee, dtype=float) + data = { + f"person_id__{DATASET_YEAR}": np.arange(len(household)), + f"person_household_id__{DATASET_YEAR}": household, + f"person_household_role__{DATASET_YEAR}": np.where( + first_member, "parent", "child" + ), + f"household_weight__{DATASET_YEAR}": np.array(inputs.household_weight)[ + household + ], + f"yearly_grant__{DATASET_YEAR}": np.array(inputs.yearly_grant), + f"household_yearly_grant__{DATASET_YEAR}": np.array( + inputs.household_yearly_grant + )[household], + } + for index, month in enumerate(MONTHS): + data[f"monthly_pay__{month}"] = pay[:, index] + data[f"household_monthly_fee__{month}"] = fee[household, index] + if inputs.yearly_tenure is not None: + data[f"yearly_tenure__{DATASET_YEAR}"] = np.array(inputs.yearly_tenure) + return pd.DataFrame(data) + + +def simulations(inputs: Inputs) -> Tuple[Simulation, Microsimulation]: + """A Simulation and a Microsimulation, each loaded from ``inputs``.""" + frame = dataframe(inputs) + return tuple( + simulation_class( + tax_benefit_system=tax_benefit_system(), + dataset=Dataset.from_dataframe(frame, DATASET_YEAR), + ) + for simulation_class in (Simulation, Microsimulation) + ) + + +def unweighted(simulation) -> dict: + """Keyword arguments that make ``simulation`` return unweighted values.""" + if isinstance(simulation, CoreMicrosimulation): + return {"use_weights": False} + return {} + + +def outcome(read: Callable): + """What ``read()`` returns, as a list, or the name of the error it raises.""" + try: + return np.asarray(read()).tolist() + except Exception as error: + return type(error).__name__ + + +def run(simulation, request: tuple): + """Run one read ``request`` on ``simulation`` and return its outcome. + + A request is one of: + + - ``("calculate", variable, period)``; + - ``("population", entity_key, variable, period, option)``, the call a + formula makes; + - ``("add" or "divide", variable, period, decode_enums)``. + + Reads in a Microsimulation are unweighted. + """ + kind, *arguments = request + if kind == "calculate": + variable, period = arguments + return outcome( + lambda: simulation.calculate(variable, period, **unweighted(simulation)) + ) + if kind == "population": + entity_key, variable, period, option = arguments + population = simulation.populations[entity_key] + return outcome(lambda: population(variable, period, options=[option])) + variable, period, decode_enums = arguments + method = {"add": simulation.calculate_add, "divide": simulation.calculate_divide}[ + kind + ] + return outcome( + lambda: method( + variable, period, decode_enums=decode_enums, **unweighted(simulation) + ) + ) + + +def expected(inputs: Inputs, reader: str) -> np.ndarray: + """``reader`` worked out from ``inputs`` with NumPy, without Core. + + Monthly readers return the same value for every month of the year. + """ + household = inputs.household_of_person + households = len(inputs.household_sizes) + pay = np.array(inputs.monthly_pay, dtype=float) + grant = np.array(inputs.yearly_grant, dtype=float) + fee = np.array(inputs.household_monthly_fee, dtype=float) + household_grant = np.array(inputs.household_yearly_grant, dtype=float) + + def members(values): + return np.bincount(household, weights=values, minlength=households) + + return { + "pay_added": pay.sum(axis=1), + "take_home_added": (0.75 * pay).sum(axis=1), + "grant_divided": grant / 12, + "fee_added": fee.sum(axis=1), + "pays_fee": fee.sum(axis=1) > 0, + "allowance_divided": 2 * household_grant / 12, + "members_pay_added": members(pay.sum(axis=1)), + "members_grant_divided": members(grant) / 12, + "fee_and_pay_added": fee.sum(axis=1) + members(pay.sum(axis=1)), + }[reader] From 753bef26d48653225676ef577834948acdd978ac Mon Sep 17 00:00:00 2001 From: Max Ghenis Date: Mon, 9 Nov 2026 14:12:15 -0500 Subject: [PATCH 2/4] Return NumPy arrays from formula reads with ADD or DIVIDE In a Microsimulation, calculate_add and calculate_divide return a MicroSeries, and pandas pairs two Series' values by index label. GroupPopulation.project labels each member's value with its group's position. So a formula that projected such a read and combined it with another Series raised a ValueError, or, when every group had one member and the data listed them out of group order, paired members with another group's value. policyengine-us's wi_wheap_categorically_eligible has that shape (lines 38-44). Population.__call__ now returns the NumPy array a Simulation returns for these reads. They raised a TypeError in every Microsimulation on master, so no value that computed before moves. The fixture now takes each person's household, so data can list members in any order, and adds fee_or_all_paid, a reader with the PE-US formula's shape. The Enum DIVIDE test now records both classes' results instead of failing as a strict xfail. Co-Authored-By: Claude Opus 5.5 --- ...-microsim-add-divide-decode-enums.fixed.md | 2 +- policyengine_core/populations/population.py | 13 +++- .../core/test_microsimulation_option_reads.py | 74 ++++++++++++++----- .../fixtures/microsimulation_option_reads.py | 50 +++++++++---- 4 files changed, 106 insertions(+), 33 deletions(-) diff --git a/changelog.d/fix-microsim-add-divide-decode-enums.fixed.md b/changelog.d/fix-microsim-add-divide-decode-enums.fixed.md index 70fb64daa..238b47b02 100644 --- a/changelog.d/fix-microsim-add-divide-decode-enums.fixed.md +++ b/changelog.d/fix-microsim-add-divide-decode-enums.fixed.md @@ -1 +1 @@ -`Microsimulation.calculate_add` and `calculate_divide` now take `decode_enums`, so formulas that read a variable with `options=[ADD]` or `options=[DIVIDE]` compute in a `Microsimulation` instead of raising a `TypeError`. +`Microsimulation.calculate_add` and `calculate_divide` now take `decode_enums`, so formulas that read a variable with `options=[ADD]` or `options=[DIVIDE]` compute in a `Microsimulation` instead of raising a `TypeError`. Those reads return a NumPy array, as in a `Simulation`. diff --git a/policyengine_core/populations/population.py b/policyengine_core/populations/population.py index b4c99123f..44cad769d 100644 --- a/policyengine_core/populations/population.py +++ b/policyengine_core/populations/population.py @@ -2,6 +2,7 @@ from typing import TYPE_CHECKING, Any, List, Optional, Container import numpy +import pandas as pd from numpy.typing import ArrayLike from policyengine_core import projectors @@ -154,15 +155,23 @@ def __call__( calculate_kwargs["decode_enums"] = False if config.ADD in options: - return self.simulation.calculate_add( + result = self.simulation.calculate_add( variable_name, period, **calculate_kwargs ) elif config.DIVIDE in options: - return self.simulation.calculate_divide( + result = self.simulation.calculate_divide( variable_name, period, **calculate_kwargs ) else: return self.simulation.calculate(variable_name, period, **calculate_kwargs) + # In a Microsimulation these reads return a MicroSeries. pandas pairs + # two Series' values by index label, and ``GroupPopulation.project`` + # labels each member's value with its group's position, so combining a + # projected read with another Series pairs the wrong values or raises. + # Return the NumPy array a Simulation returns. + if isinstance(result, pd.Series): + return result.to_numpy() + return result # Helpers diff --git a/tests/core/test_microsimulation_option_reads.py b/tests/core/test_microsimulation_option_reads.py index af82faca1..81f92277f 100644 --- a/tests/core/test_microsimulation_option_reads.py +++ b/tests/core/test_microsimulation_option_reads.py @@ -7,6 +7,13 @@ every such read raised ``TypeError: Microsimulation.calculate_add() got an unexpected keyword argument 'decode_enums'``. A Simulation computed the same reads. + +In a Microsimulation those two methods return a ``MicroSeries``, and pandas +pairs two Series' values by index label. ``household.project`` labels each +member's value with its household's position, so combining a projected +``MicroSeries`` with another Series raises a ``ValueError`` or pairs members +with another household's value. A formula's read with an option therefore +returns the NumPy array a Simulation returns. """ import dataclasses @@ -15,18 +22,20 @@ import pytest from microdf import MicroSeries +from policyengine_core.enums import EnumArray from policyengine_core.model_api import ADD, DIVIDE, MONTH from tests.fixtures.microsimulation_option_reads import ( DATASET_YEAR, READERS, Inputs, + Tenure, expected, run, simulations, ) EXAMPLE = Inputs( - household_sizes=(2, 1), + household_of_person=(0, 0, 1), household_weight=(10.0, 20.0), monthly_pay=( tuple(1.0 * month for month in range(1, 13)), @@ -48,6 +57,7 @@ "members_pay_added": [234.0, 390.0], "members_grant_divided": [30.0, 50.0], "fee_and_pay_added": [294.0, 390.0], + "fee_or_all_paid": [True, True], } @@ -73,6 +83,31 @@ def test_formula_reading_with_an_option_computes_in_a_microsimulation(reader): np.testing.assert_array_equal(result, simulation.calculate(name, period)) +def test_projected_read_pairs_each_member_with_its_own_household(): + # Two one-person households, with household 1's member listed first. A + # projected MicroSeries would pair each person with the other household's + # fee and return [True, False] without an error. + inputs = Inputs( + household_of_person=(1, 0), + household_weight=(10.0, 20.0), + monthly_pay=((0.0,) * 12, (0.0,) * 12), + yearly_grant=(0.0, 0.0), + household_monthly_fee=((0.0,) * 12, (5.0,) * 12), + household_yearly_grant=(0.0, 0.0), + ) + simulation, microsimulation = simulations(inputs) + + result = microsimulation.calculate( + "fee_or_all_paid", DATASET_YEAR, use_weights=False + ) + + np.testing.assert_array_equal(expected(inputs, "fee_or_all_paid"), [False, True]) + np.testing.assert_array_equal(result, [False, True]) + np.testing.assert_array_equal( + simulation.calculate("fee_or_all_paid", DATASET_YEAR), [False, True] + ) + + def test_weighted_result_of_a_formula_reading_with_add_carries_its_weights(): _, microsimulation = simulations(EXAMPLE) @@ -98,11 +133,13 @@ def test_weighted_result_of_a_formula_reading_with_add_carries_its_weights(): ) def test_population_read_with_an_option_matches_a_simulation(call): simulation, microsimulation = simulations(EXAMPLE) + _, entity_key, variable, period, option = call - result = run(microsimulation, call) + result = microsimulation.populations[entity_key](variable, period, options=[option]) - assert isinstance(result, list), result - assert result == run(simulation, call) + # A NumPy array, as in a Simulation: a MicroSeries pairs values by label. + assert type(result) is np.ndarray + assert result.tolist() == run(simulation, call) @pytest.mark.parametrize("decode_enums", [False, True]) @@ -133,18 +170,21 @@ def test_microsimulation_add_and_divide_take_decode_enums( ) -@pytest.mark.xfail( - strict=True, - reason=( - "Simulation.calculate_divide reads the variable with self.calculate, " - "which in a Microsimulation decodes enums by default, so a " - "Microsimulation returns labels where a Simulation returns an " - "EnumArray. No formula in policyengine-us reads an Enum this way." - ), -) -def test_enum_read_with_divide_in_a_microsimulation_matches_a_simulation(): +def test_enum_read_with_divide_over_a_year_gives_member_names_in_a_microsimulation(): + # This records a difference that remains. Over a year, + # Simulation.calculate_divide reads the variable with self.calculate, + # which in a Microsimulation decodes enums by default. So a + # Microsimulation returns member names where a Simulation returns an + # EnumArray, and comparing the names with a member gives False for + # everyone. No policyengine-us formula reads an Enum with ADD or DIVIDE. inputs = dataclasses.replace(EXAMPLE, yearly_tenure=("OWNER", "RENTER", "OWNER")) - simulation, microsimulation = simulations(inputs) - call = ("population", "person", "yearly_tenure", DATASET_YEAR, DIVIDE) + simulation_read, microsimulation_read = ( + each.populations["person"]("yearly_tenure", DATASET_YEAR, options=[DIVIDE]) + for each in simulations(inputs) + ) - assert run(microsimulation, call) == run(simulation, call) + assert isinstance(simulation_read, EnumArray) + assert simulation_read.decode_to_str().tolist() == ["OWNER", "RENTER", "OWNER"] + assert (simulation_read == Tenure.OWNER).tolist() == [True, False, True] + assert microsimulation_read.tolist() == ["OWNER", "RENTER", "OWNER"] + assert (microsimulation_read == Tenure.OWNER).tolist() == [False, False, False] diff --git a/tests/fixtures/microsimulation_option_reads.py b/tests/fixtures/microsimulation_option_reads.py index 4013c1864..18b2f816e 100644 --- a/tests/fixtures/microsimulation_option_reads.py +++ b/tests/fixtures/microsimulation_option_reads.py @@ -189,6 +189,22 @@ def formula(household, period, parameters): ) +class fee_or_all_paid(Variable): + value_type = bool + entity = Household + definition_period = YEAR + label = "Household paid a fee or every member was paid" + + def formula(household, period, parameters): + # policyengine-us's wi_wheap_categorically_eligible has this shape: it + # projects a group's ADD read to its members and combines the result + # with a person value. + person = household.members + paid = add(person, period, ["monthly_pay"]) > 0 + pays = household("household_monthly_fee", period, options=[ADD]) > 0 + return household.all(paid | household.project(pays)) + + SOURCES = [ monthly_pay, monthly_take_home, @@ -207,6 +223,7 @@ def formula(household, period, parameters): members_pay_added, members_grant_divided, fee_and_pay_added, + fee_or_all_paid, ] @@ -220,13 +237,16 @@ def tax_benefit_system() -> CountryTaxBenefitSystem: @dataclass(frozen=True) class Inputs: - """Households in order, and each input's values. + """People in the order the data lists them, and each input's values. - ``monthly_pay`` and ``household_monthly_fee`` hold 12 values (January to - December) for each person or household. + ``household_of_person`` gives each person's household. Households take the + numbers 0, 1, ... and each has at least one member; the data can list a + household's members anywhere. ``monthly_pay`` and + ``household_monthly_fee`` hold 12 values (January to December) for each + person or household. """ - household_sizes: Tuple[int, ...] + household_of_person: Tuple[int, ...] household_weight: Tuple[float, ...] monthly_pay: Tuple[Tuple[float, ...], ...] yearly_grant: Tuple[float, ...] @@ -234,17 +254,17 @@ class Inputs: household_yearly_grant: Tuple[float, ...] yearly_tenure: Optional[Tuple[str, ...]] = None - @property - def household_of_person(self) -> np.ndarray: - return np.repeat( - np.arange(len(self.household_sizes)), np.array(self.household_sizes) - ) + def __post_init__(self): + # A dataset declares a household only through its members' rows. + if set(self.household_of_person) != set(range(len(self.household_weight))): + raise ValueError("Every household needs a weight and a member.") def dataframe(inputs: Inputs) -> pd.DataFrame: """One row per person. Each member's row repeats its household's values.""" - household = inputs.household_of_person - first_member = np.r_[True, household[1:] != household[:-1]] + household = np.array(inputs.household_of_person) + first_member = np.zeros(len(household), dtype=bool) + first_member[np.unique(household, return_index=True)[1]] = True pay = np.array(inputs.monthly_pay, dtype=float) fee = np.array(inputs.household_monthly_fee, dtype=float) data = { @@ -334,8 +354,8 @@ def expected(inputs: Inputs, reader: str) -> np.ndarray: Monthly readers return the same value for every month of the year. """ - household = inputs.household_of_person - households = len(inputs.household_sizes) + household = np.array(inputs.household_of_person) + households = len(inputs.household_weight) pay = np.array(inputs.monthly_pay, dtype=float) grant = np.array(inputs.yearly_grant, dtype=float) fee = np.array(inputs.household_monthly_fee, dtype=float) @@ -344,6 +364,8 @@ def expected(inputs: Inputs, reader: str) -> np.ndarray: def members(values): return np.bincount(household, weights=values, minlength=households) + paid = pay.sum(axis=1) > 0 + pays = fee.sum(axis=1) > 0 return { "pay_added": pay.sum(axis=1), "take_home_added": (0.75 * pay).sum(axis=1), @@ -354,4 +376,6 @@ def members(values): "members_pay_added": members(pay.sum(axis=1)), "members_grant_divided": members(grant) / 12, "fee_and_pay_added": fee.sum(axis=1) + members(pay.sum(axis=1)), + # ``all`` over each household's members: no member fails the test. + "fee_or_all_paid": members(~(paid | pays[household]) * 1.0) == 0, }[reader] From 141219eb6df049ffedcd8e5f5d3d4c8dfba0d8f4 Mon Sep 17 00:00:00 2001 From: Max Ghenis Date: Mon, 9 Nov 2026 14:12:15 -0500 Subject: [PATCH 3/4] Read options over periods they accept in the property, in any person order Population calls and calculate_add and calculate_divide requests now pair each variable with periods its option accepts, and population calls read only variables the population can read, so they return values instead of mostly raising in both classes. The data lists people in any order. hypothesis.event reports the mix of values, zeros and errors. Co-Authored-By: Claude Opus 5.5 --- ...t_microsimulation_option_reads_property.py | 107 +++++++++++++----- 1 file changed, 79 insertions(+), 28 deletions(-) diff --git a/tests/core/test_microsimulation_option_reads_property.py b/tests/core/test_microsimulation_option_reads_property.py index cb1186597..2e65d5205 100644 --- a/tests/core/test_microsimulation_option_reads_property.py +++ b/tests/core/test_microsimulation_option_reads_property.py @@ -1,26 +1,33 @@ """Property: a Microsimulation's unweighted ADD and DIVIDE reads equal a Simulation's. -For any small dataset of households, people, weights and monthly and yearly -inputs, and any sequence of reads, a Microsimulation returns the same -unweighted values as a Simulation loaded from the same data, or raises the -same kind of error. The reads are formula results (each formula reads with -``options=[ADD]`` or ``options=[DIVIDE]``), the population calls formulas -make, ``calculate_add`` and ``calculate_divide`` with either value of -``decode_enums``, and plain reads at months, years and longer periods. Since -the Simulation ignores weights, weights never change an unweighted value. +For any small dataset of households, people listed in any order, weights and +monthly and yearly inputs, and any sequence of reads, a Microsimulation +returns the same unweighted values as a Simulation loaded from the same data, +or raises the same kind of error. The reads are: + +- formula results: each formula reads with ``options=[ADD]`` or + ``options=[DIVIDE]``, and one projects such a read to members; +- the population calls formulas make; +- ``calculate_add`` and ``calculate_divide``, with either value of + ``decode_enums``; +- plain reads at months, years and longer periods. + +Population calls and the two methods read each variable over periods its +option accepts. Since the Simulation ignores weights, weights never change an +unweighted value. Formula results also equal the same arithmetic done in NumPy, so the two simulations cannot agree on a wrong value. Enum variables are left out: -``test_enum_read_with_divide_in_a_microsimulation_matches_a_simulation`` +``test_enum_read_with_divide_over_a_year_gives_member_names_in_a_microsimulation`` records how their reads differ. """ import numpy as np import pytest -from policyengine_core.model_api import ADD, DIVIDE, MONTH +from policyengine_core.model_api import ADD, DIVIDE, MONTH, YEAR from tests.fixtures.microsimulation_option_reads import ( DATASET_YEAR, MONTHS, @@ -40,40 +47,75 @@ AMOUNTS = st.integers(-10_000, 10_000).map(float) WEIGHTS = st.floats(min_value=0, max_value=1e6, allow_nan=False) PERIODS = ["2022", "2022-01", "2022-06", "2022-12", "2023", "month:2022-01:3"] -VARIABLES = sorted(variable.__name__ for variable in SOURCES + READERS) +VARIABLES = SOURCES + READERS + +# The periods each option reads a variable over, by the variable's definition +# period. Over any other period the read raises in both classes. +OPTION_PERIODS = { + (ADD, MONTH): ["2022", "2022-01", "2022-06", "2023", "month:2022-01:3"], + (ADD, YEAR): ["2022", "2023"], + (DIVIDE, YEAR): ["2022", "2022-01", "2022-06", "2022-12", "2023"], +} + + +def option_reads(variable) -> list: + """Each option and period that reads ``variable`` without an error.""" + return [ + (option, period) + for (option, definition_period), periods in OPTION_PERIODS.items() + if definition_period == variable.definition_period + for period in periods + ] + + +def readable(entity_key: str, variable) -> bool: + # A household's population sums its members' person variables. + return variable.entity.key == entity_key or ( + entity_key == "household" and variable.entity.is_person + ) + +POPULATION_CALLS = [ + ("population", entity_key, variable.__name__, period, option) + for entity_key in ["person", "household"] + for variable in VARIABLES + if readable(entity_key, variable) + for option, period in option_reads(variable) +] +METHOD_CALLS = [ + ({ADD: "add", DIVIDE: "divide"}[option], variable.__name__, period, decode) + for variable in VARIABLES + for option, period in option_reads(variable) + for decode in [False, True] +] REQUESTS = st.one_of( st.tuples( st.just("calculate"), - st.sampled_from(VARIABLES), - st.sampled_from(PERIODS), - ), - st.tuples( - st.just("population"), - st.sampled_from(["person", "household"]), - st.sampled_from(VARIABLES), - st.sampled_from(PERIODS), - st.sampled_from([ADD, DIVIDE]), - ), - st.tuples( - st.sampled_from(["add", "divide"]), - st.sampled_from(VARIABLES), + st.sampled_from(sorted(variable.__name__ for variable in VARIABLES)), st.sampled_from(PERIODS), - st.booleans(), ), + st.sampled_from(POPULATION_CALLS), + st.sampled_from(METHOD_CALLS), ) @st.composite def inputs(draw): sizes = draw(st.lists(st.integers(1, 3), min_size=1, max_size=4)) - people, households = sum(sizes), len(sizes) + households = len(sizes) + # The data can list a household's members anywhere. + household_of_person = draw( + st.permutations( + [household for household, size in enumerate(sizes) for _ in range(size)] + ) + ) + people = len(household_of_person) def values(count): return tuple(draw(st.lists(AMOUNTS, min_size=count, max_size=count))) return Inputs( - household_sizes=tuple(sizes), + household_of_person=tuple(household_of_person), household_weight=tuple( draw(st.lists(WEIGHTS, min_size=households, max_size=households)) ), @@ -84,12 +126,21 @@ def values(count): ) +def outcome_kind(outcome) -> str: + """Whether a read raised, returned only zeros, or returned other values.""" + if isinstance(outcome, str): + return "an error" + return "values" if np.any(np.asarray(outcome, dtype=float)) else "zeros" + + @hypothesis.settings(max_examples=200, deadline=None) @hypothesis.given(inputs(), st.lists(REQUESTS, min_size=1, max_size=6)) def test_unweighted_microsimulation_reads_equal_simulation_reads(data, requests): simulation, microsimulation = simulations(data) for request in requests: - assert run(microsimulation, request) == run(simulation, request), request + result = run(microsimulation, request) + hypothesis.event(f"{request[0]} read gives {outcome_kind(result)}") + assert result == run(simulation, request), request @hypothesis.settings(max_examples=100, deadline=None) From 99a808e6ef4f946ee4dc56026e37721f59b1397f Mon Sep 17 00:00:00 2001 From: Max Ghenis Date: Mon, 9 Nov 2026 14:28:00 -0500 Subject: [PATCH 4/4] Test that option reads return arrays, and check every formula in each run A new property checks that each population call a formula makes with ADD or DIVIDE returns a NumPy array in both classes. The NumPy-arithmetic property now runs for each formula, 25 examples each: with 100 examples drawn across all ten, one seeded run never drew fee_or_all_paid or members_grant_divided. Co-Authored-By: Claude Opus 5.5 --- ...t_microsimulation_option_reads_property.py | 27 ++++++++++++------- 1 file changed, 18 insertions(+), 9 deletions(-) diff --git a/tests/core/test_microsimulation_option_reads_property.py b/tests/core/test_microsimulation_option_reads_property.py index 2e65d5205..13279fc56 100644 --- a/tests/core/test_microsimulation_option_reads_property.py +++ b/tests/core/test_microsimulation_option_reads_property.py @@ -16,8 +16,9 @@ option accepts. Since the Simulation ignores weights, weights never change an unweighted value. -Formula results also equal the same arithmetic done in NumPy, so the two -simulations cannot agree on a wrong value. +The population calls return NumPy arrays in both classes. Formula results +also equal the same arithmetic done in NumPy, so the two simulations cannot +agree on a wrong value. Enum variables are left out: ``test_enum_read_with_divide_over_a_year_gives_member_names_in_a_microsimulation`` @@ -143,13 +144,21 @@ def test_unweighted_microsimulation_reads_equal_simulation_reads(data, requests) assert result == run(simulation, request), request -@hypothesis.settings(max_examples=100, deadline=None) -@hypothesis.given( - inputs(), - st.sampled_from(READERS), - st.sampled_from(MONTHS), -) -def test_formula_reads_with_an_option_equal_numpy_arithmetic(data, reader, month): +@hypothesis.settings(max_examples=50, deadline=None) +@hypothesis.given(inputs(), st.sampled_from(POPULATION_CALLS)) +def test_population_reads_with_an_option_return_numpy_arrays(data, call): + # pandas pairs two Series' values by index label, and a projected Series + # carries each member's group position, so formulas get positional arrays. + _, entity_key, variable, period, option = call + for simulation in simulations(data): + read = simulation.populations[entity_key](variable, period, options=[option]) + assert type(read) is np.ndarray, (type(simulation).__name__, call) + + +@pytest.mark.parametrize("reader", READERS, ids=lambda reader: reader.__name__) +@hypothesis.settings(max_examples=25, deadline=None) +@hypothesis.given(inputs(), st.sampled_from(MONTHS)) +def test_formula_reads_with_an_option_equal_numpy_arithmetic(reader, data, month): period = month if reader.definition_period == MONTH else DATASET_YEAR reference = expected(data, reader.__name__) for simulation in simulations(data):