diff --git a/changelog.d/microsimulation-map-to.fixed.md b/changelog.d/microsimulation-map-to.fixed.md new file mode 100644 index 00000000..851148d7 --- /dev/null +++ b/changelog.d/microsimulation-map-to.fixed.md @@ -0,0 +1 @@ +Fix `Microsimulation.calculate_add` and `calculate_divide` to honor `map_to` for both values and weights. diff --git a/policyengine_core/simulations/microsimulation.py b/policyengine_core/simulations/microsimulation.py index 135be6df..af50a088 100644 --- a/policyengine_core/simulations/microsimulation.py +++ b/policyengine_core/simulations/microsimulation.py @@ -88,10 +88,15 @@ def calculate_add( map_to: str = None, use_weights: bool = True, ) -> MicroSeries: - values = super().calculate_add(variable_name, period, map_to) + values = super().calculate_add(variable_name, period) + if map_to is not None: + source_entity = self.tax_benefit_system.get_variable( + variable_name, check_existence=True + ).entity.key + values = self.map_result(np.array(values), source_entity, map_to) if not use_weights: return values - weights = self.get_weights(variable_name, period) + weights = self.get_weights(variable_name, period, map_to) return MicroSeries(np.array(values), weights=weights) def calculate_divide( @@ -101,10 +106,15 @@ def calculate_divide( map_to: str = None, use_weights: bool = True, ) -> MicroSeries: - values = super().calculate_divide(variable_name, period, map_to) + values = super().calculate_divide(variable_name, period) + if map_to is not None: + source_entity = self.tax_benefit_system.get_variable( + variable_name, check_existence=True + ).entity.key + values = self.map_result(np.array(values), source_entity, map_to) if not use_weights: return values - weights = self.get_weights(variable_name, period) + weights = self.get_weights(variable_name, period, map_to) return MicroSeries(np.array(values), weights=weights) def calculate_dataframe( diff --git a/tests/core/test_microsimulation_weights.py b/tests/core/test_microsimulation_weights.py index 8d5facba..f794f76c 100644 --- a/tests/core/test_microsimulation_weights.py +++ b/tests/core/test_microsimulation_weights.py @@ -16,6 +16,7 @@ def _weighted_dataset(include_person_weight: bool = True) -> Dataset: "person_household_role__2022": ["parent", "child", "parent"], "household_weight__2022": [10.0, 10.0, 20.0], "salary__2022-01": [100.0, 200.0, 300.0], + "housing_tax__2022": [120.0, 120.0, 240.0], } if include_person_weight: data["person_weight__2022"] = [1.0, 2.0, 3.0] @@ -65,6 +66,38 @@ def test__given_dataframe_mapped_to_household__then_weights_are_household_weight np.testing.assert_array_equal(dataframe.weights, np.array([10.0, 20.0])) +def test__given_calculate_add_mapped_to_household__then_values_and_weights_match_households(): + # Given + simulation = Microsimulation(dataset=_weighted_dataset()) + + # When + result = simulation.calculate_add( + "salary", + "2022-01", + map_to="household", + ) + + # Then + np.testing.assert_array_equal(result, np.array([300.0, 300.0])) + np.testing.assert_array_equal(result.weights, np.array([10.0, 20.0])) + + +def test__given_calculate_divide_mapped_to_person__then_values_and_weights_match_persons(): + # Given + simulation = Microsimulation(dataset=_weighted_dataset()) + + # When + result = simulation.calculate_divide( + "housing_tax", + "2022-01", + map_to="person", + ) + + # Then + np.testing.assert_array_equal(result, np.array([10.0, 10.0, 20.0])) + np.testing.assert_array_equal(result.weights, np.array([10.0, 10.0, 20.0])) + + Family = build_entity( key="family", plural="families",