Fix formula reads with ADD or DIVIDE in a Microsimulation

changelog.d/fix-microsim-add-divide-decode-enums.fixed.md

@@ -0,0 +1 @@
+`Microsimulation.calculate_add` and `calculate_divide` now take `decode_enums`, so a formula that reads a variable with `options=[ADD]` or `options=[DIVIDE]` gets a writable NumPy array in a `Microsimulation`, as in a `Simulation`, instead of a `TypeError`.

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,27 @@ 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 a NumPy array, as a Simulation does, and copy it: pandas 3
+        # returns a read-only view of a Series' values, and a formula may write
+        # to what it reads. ``numpy.array`` copies the values without calling
+        # ``MicroSeries.to_numpy`` or ``MicroSeries.values``, and microdf warns
+        # on each of those calls that the result carries no weights.
+        if isinstance(result, pd.Series):
+            return numpy.array(result)
+        return result
 
     # Helpers
 

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

tests/core/test_microsimulation_calculate_signatures.py

@@ -0,0 +1,86 @@
+"""Each ``Microsimulation`` calculate method takes the arguments ``Simulation``'s takes.
+
+``Population.__call__`` sends a formula's read to ``calculate``,
+``calculate_add`` or ``calculate_divide`` with one set of keyword arguments.
+``Microsimulation`` overrides those three methods and ``calculate_dataframe``.
+When an override lacks one of the base method's arguments, Python rejects the
+call with a ``TypeError`` in a ``Microsimulation``, and a ``Simulation``
+computes the same call.
+
+Any call that passes its arguments by keyword and binds to a ``Simulation``
+calculate method also binds to the ``Microsimulation`` override. The test
+checks that for every subset of each method's optional arguments.
+"""
+
+import inspect
+from itertools import chain, combinations
+
+import pytest
+
+from policyengine_core.simulations import Microsimulation, Simulation
+
+
+def calculate_overrides() -> list:
+    """The ``Simulation`` calculate methods that ``Microsimulation`` overrides."""
+    return sorted(
+        name
+        for name, member in vars(Microsimulation).items()
+        if name.startswith("calculate")
+        and inspect.isfunction(member)
+        and inspect.isfunction(getattr(Simulation, name, None))
+    )
+
+
+def keyword_calls(name: str) -> list:
+    """Every set of argument names a caller can pass to ``Simulation``'s method."""
+    parameters = inspect.signature(getattr(Simulation, name)).parameters.values()
+    required = [
+        parameter.name
+        for parameter in parameters
+        if parameter.default is parameter.empty
+    ]
+    optional = [
+        parameter.name
+        for parameter in parameters
+        if parameter.default is not parameter.empty
+    ]
+    subsets = chain.from_iterable(
+        combinations(optional, size) for size in range(len(optional) + 1)
+    )
+    return [(*required, *subset) for subset in subsets]
+
+
+CALLS = [
+    (name, arguments)
+    for name in calculate_overrides()
+    for arguments in keyword_calls(name)
+]
+
+
+def test_microsimulation_overrides_the_four_calculate_methods():
+    # An empty list would leave the property below with nothing to check.
+    assert {
+        "calculate",
+        "calculate_add",
+        "calculate_divide",
+        "calculate_dataframe",
+    } <= set(calculate_overrides())
+
+
+@pytest.mark.parametrize(
+    "name,arguments",
+    CALLS,
+    ids=[f"{name}({', '.join(arguments[1:])})" for name, arguments in CALLS],
+)
+def test_keyword_call_to_a_simulation_method_binds_to_the_override(name, arguments):
+    call = dict.fromkeys(arguments)
+    inspect.signature(getattr(Simulation, name)).bind(**call)
+
+    try:
+        inspect.signature(getattr(Microsimulation, name)).bind(**call)
+    except TypeError as error:
+        pytest.fail(
+            f"Microsimulation.{name} rejects a call that Simulation.{name} "
+            f"takes: {error}",
+            pytrace=False,
+        )

tests/core/test_microsimulation_option_reads.py

@@ -0,0 +1,221 @@
+"""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.
+
+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 a NumPy array, as in a Simulation, and one the formula can write to.
+"""
+
+import dataclasses
+import warnings
+
+import numpy as np
+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_of_person=(0, 0, 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],
+    "fee_or_all_paid": [True, True],
+}
+
+
+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_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)
+
+    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)
+    _, entity_key, variable, period, option = call
+
+    result = microsimulation.populations[entity_key](variable, period, options=[option])
+
+    # A NumPy array, as in a Simulation: a MicroSeries pairs values by label.
+    assert type(result) is np.ndarray
+    # Formulas may write to what they read.
+    assert result.flags.writeable
+    assert result.tolist() == run(simulation, call)
+
+
+@pytest.mark.parametrize(
+    "entity_key,variable,period,option",
+    [
+        ("person", "monthly_pay", DATASET_YEAR, ADD),
+        ("person", "yearly_grant", "2022-06", DIVIDE),
+        ("household", "household_monthly_fee", DATASET_YEAR, ADD),
+        ("household", "household_yearly_allowance", "2022-06", DIVIDE),
+    ],
+)
+def test_population_read_with_an_option_gives_no_microseries_warning(
+    entity_key, variable, period, option
+):
+    # microdf warns when a caller takes a MicroSeries' values through
+    # ``to_numpy`` or ``values``, because the result carries no weights. The
+    # formula never asked for weights, so its read gives no such warning.
+    _, microsimulation = simulations(EXAMPLE)
+
+    with warnings.catch_warnings(record=True) as caught:
+        warnings.simplefilter("always")
+        microsimulation.populations[entity_key](variable, period, options=[option])
+
+    assert [
+        str(warning.message)
+        for warning in caught
+        if "MicroSeries" in str(warning.message)
+    ] == []
+
+
+@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)
+    )
+
+
+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_read, microsimulation_read = (
+        each.populations["person"]("yearly_tenure", DATASET_YEAR, options=[DIVIDE])
+        for each in simulations(inputs)
+    )
+
+    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]