Use PolicyEngine formulas for oracle targets

Author: MaxGhenis

src/microplex_us/__init__.py

@@ -167,6 +167,7 @@
     "infer_policyengine_us_variable_bindings",
     "load_policyengine_us_entity_tables",
     "materialize_policyengine_us_variables",
+    "policyengine_us_formula_variables_for_targets",
     "policyengine_us_variables_to_materialize",
     "project_frame_to_time_period_arrays",
     "write_policyengine_us_time_period_dataset",
@@ -356,6 +357,7 @@ def __getattr__(name: str) -> Any:
     "infer_policyengine_us_variable_bindings",
     "load_policyengine_us_entity_tables",
     "materialize_policyengine_us_variables",
+    "policyengine_us_formula_variables_for_targets",
     "policyengine_us_variables_to_materialize",
     "project_frame_to_time_period_arrays",
     "write_policyengine_us_time_period_dataset",

src/microplex_us/pipelines/summarize_policyengine_oracle_target_drilldown.py

@@ -65,7 +65,7 @@ def summarize_us_policyengine_oracle_target_drilldown(
         _supported_targets,
         _constraints,
         _feasibility_filter_summary,
-        _materialized_variables,
+        calibration_materialized_variables,
         _materialization_failures,
     ) = pipeline._resolve_policyengine_calibration_targets(
         tables,
@@ -100,6 +100,8 @@ def summarize_us_policyengine_oracle_target_drilldown(
         str(variable)
         for variable in manifest.get("calibration", {}).get("materialized_variables", ())
     }
+    materialized_variables.update(str(variable) for variable in calibration_materialized_variables)
+    materialized_variables.update(str(variable) for variable in report.materialized_variables)
     ledger_by_name = {
         str(entry["target_name"]): dict(entry)
         for entry in target_ledger

src/microplex_us/pipelines/us.py

@@ -72,6 +72,7 @@
     infer_policyengine_us_variable_bindings,
     load_us_pipeline_checkpoint,
     materialize_policyengine_us_variables_safely,
+    policyengine_us_formula_variables_for_targets,
     policyengine_us_variables_to_materialize,
     resolve_policyengine_excluded_export_variables,
     save_us_pipeline_checkpoint,
@@ -3831,9 +3832,15 @@ def _resolve_policyengine_calibration_targets(
             period=target_period,
             for_calibration=True,
         ).targets
+        force_materialize_variables = policyengine_us_formula_variables_for_targets(
+            canonical_targets,
+            simulation_cls=self.config.policyengine_simulation_cls,
+            direct_override_variables=self.config.policyengine_direct_override_variables,
+        )
         missing_variables = policyengine_us_variables_to_materialize(
             canonical_targets,
             bindings,
+            force_materialize_variables=force_materialize_variables,
         )
         materialization_failures: dict[str, str] = {}
         materialized_variables: set[str] = set()
@@ -3844,9 +3851,20 @@ def _resolve_policyengine_calibration_targets(
                 period=target_period,
                 dataset_year=self.config.policyengine_dataset_year or target_period,
                 simulation_cls=self.config.policyengine_simulation_cls,
+                direct_override_variables=self.config.policyengine_direct_override_variables,
                 batch_size=self.config.policyengine_materialize_batch_size,
             )
             tables = materialization_result.tables
+            unmaterialized_forced_variables = (
+                force_materialize_variables
+                & missing_variables
+                - set(materialization_result.bindings)
+            )
+            bindings = {
+                variable: binding
+                for variable, binding in bindings.items()
+                if variable not in unmaterialized_forced_variables
+            }
             bindings = {
                 **bindings,
                 **materialization_result.bindings,

src/microplex_us/policyengine/__init__.py

@@ -39,6 +39,7 @@
     infer_policyengine_us_variable_bindings,
     load_policyengine_us_entity_tables,
     materialize_policyengine_us_variables,
+    policyengine_us_formula_variables_for_targets,
     policyengine_us_variables_to_materialize,
     project_frame_to_time_period_arrays,
     write_policyengine_us_time_period_dataset,
@@ -79,6 +80,7 @@
     "infer_policyengine_us_variable_bindings",
     "load_policyengine_us_entity_tables",
     "materialize_policyengine_us_variables",
+    "policyengine_us_formula_variables_for_targets",
     "policyengine_us_variables_to_materialize",
     "project_frame_to_time_period_arrays",
     "write_policyengine_us_time_period_dataset",

src/microplex_us/policyengine/comparison.py

@@ -34,6 +34,8 @@
     infer_policyengine_us_variable_bindings,
     load_policyengine_us_entity_tables,
     materialize_policyengine_us_variables_safely,
+    policyengine_us_formula_variables_for_targets,
+    policyengine_us_variables_to_materialize,
 )
 
 POLICYENGINE_US_BENCHMARK_GROUP_FIELDS = (
@@ -363,20 +365,35 @@ def evaluate_policyengine_us_target_set(
     target_list = _normalize_target_list(targets)
     working_tables = tables
     bindings = infer_policyengine_us_variable_bindings(working_tables)
+    force_materialize_variables = policyengine_us_formula_variables_for_targets(
+        target_list,
+        simulation_cls=simulation_cls,
+        direct_override_variables=direct_override_variables,
+    )
+    variables_to_materialize = policyengine_us_variables_to_materialize(
+        target_list,
+        bindings,
+        force_materialize_variables=force_materialize_variables,
+    )
     materialization_result = materialize_policyengine_us_variables_safely(
         working_tables,
-        variables=tuple(
-            feature
-            for target in target_list
-            for feature in target.required_features
-            if feature not in bindings
-        ),
+        variables=tuple(sorted(variables_to_materialize)),
         period=period,
         dataset_year=dataset_year,
         simulation_cls=simulation_cls,
         direct_override_variables=direct_override_variables,
     )
     working_tables = materialization_result.tables
+    unmaterialized_forced_variables = (
+        force_materialize_variables
+        & variables_to_materialize
+        - set(materialization_result.bindings)
+    )
+    bindings = {
+        variable: binding
+        for variable, binding in bindings.items()
+        if variable not in unmaterialized_forced_variables
+    }
     bindings = {
         **bindings,
         **materialization_result.bindings,

src/microplex_us/policyengine/us.py

@@ -286,7 +286,7 @@ class PolicyEngineUSVariableMaterializationResult:
     "other_medical_expenses",
     "over_the_counter_health_expenses",
     "self_employment_income_before_lsr",
-    "social_security_retirement",
+    "social_security_retirement_reported",
     "social_security_disability",
     "social_security_survivors",
     "social_security_dependents",
@@ -327,6 +327,7 @@ class PolicyEngineUSVariableMaterializationResult:
 
 POLICYENGINE_US_EXPORT_COLUMN_ALIASES: dict[str, str] = {
     "race": "cps_race",
+    "social_security_retirement": "social_security_retirement_reported",
 }
 
 POLICYENGINE_US_EXPORT_DEFAULTS: dict[str, Any] = {
@@ -1866,18 +1867,70 @@ def compile_supported_policyengine_us_household_linear_constraints(
     return supported_targets, unsupported_targets, tuple(constraints)
 
 
+def _policyengine_us_target_required_variables(targets: list[TargetSpec]) -> set[str]:
+    return {
+        feature
+        for target in targets
+        for feature in target.required_features
+    }
+
+
+def policyengine_us_formula_variables_for_targets(
+    targets: list[TargetSpec],
+    *,
+    simulation_cls: Any | None = None,
+    tax_benefit_system: Any | None = None,
+    direct_override_variables: tuple[str, ...] = (),
+) -> set[str]:
+    """Return target features that should be recalculated by PolicyEngine."""
+    required_variables = _policyengine_us_target_required_variables(targets)
+    if not required_variables:
+        return set()
+    if tax_benefit_system is None:
+        tax_benefit_system = _resolve_policyengine_us_tax_benefit_system(
+            simulation_cls
+        )
+    variables = getattr(tax_benefit_system, "variables", {})
+    direct_overrides = set(direct_override_variables)
+    formula_variables: set[str] = set()
+    for variable in required_variables:
+        if variable in direct_overrides:
+            continue
+        variable_metadata = variables.get(variable)
+        if variable_metadata is None:
+            continue
+        if _policyengine_us_variable_is_calculated(variable_metadata):
+            formula_variables.add(variable)
+    return formula_variables
+
+
+def _policyengine_us_variable_is_calculated(variable_metadata: Any) -> bool:
+    if getattr(variable_metadata, "formulas", {}):
+        return True
+    if getattr(variable_metadata, "adds", ()) or getattr(variable_metadata, "subtracts", ()):
+        return True
+    is_input_variable = getattr(variable_metadata, "is_input_variable", None)
+    if callable(is_input_variable):
+        try:
+            return not bool(is_input_variable())
+        except TypeError:
+            return False
+    return False
+
+
 def policyengine_us_variables_to_materialize(
     targets: list[TargetSpec],
     bindings: dict[str, PolicyEngineUSVariableBinding],
+    *,
+    force_materialize_variables: set[str] | tuple[str, ...] | None = None,
 ) -> set[str]:
     """Compute the missing features required to score the given targets."""
-    requested_variables = {
-        feature
-        for target in targets
-        for feature in target.required_features
-    }
+    requested_variables = _policyengine_us_target_required_variables(targets)
+    force_variables = set(force_materialize_variables or ())
     return {
-        variable for variable in requested_variables if variable not in bindings
+        variable
+        for variable in requested_variables
+        if variable not in bindings or variable in force_variables
     }
 
 

tests/pipelines/test_artifacts.py

@@ -176,19 +176,9 @@ def _create_policyengine_targets_db(path: Path) -> None:
             t.value,
             t.period,
             t.active,
-            CASE
-                WHEN t.variable = 'snap' THEN 'state'
-                ELSE 'district'
-            END AS geo_level,
-            CASE
-                WHEN t.variable = 'snap' THEN '06'
-                ELSE '0601'
-            END AS geographic_id,
-            CASE
-                WHEN t.variable = 'snap' THEN 'snap'
-                WHEN t.variable = 'household_count' THEN 'snap'
-                ELSE NULL
-            END AS domain_variable
+            'state' AS geo_level,
+            '06' AS geographic_id,
+            'household_count' AS domain_variable
         FROM targets AS t;
         """
     )
@@ -216,7 +206,6 @@ def _create_policyengine_targets_db(path: Path) -> None:
         """,
         [
             (1, "household_count", 2024, 1, 0, 3.0, 1, None, "test", "count"),
-            (2, "snap", 2024, 1, 0, 250.0, 1, None, "test", "snap"),
         ],
     )
     conn.commit()
@@ -604,12 +593,11 @@ def test_writes_policyengine_harness_when_baseline_and_targets_are_provided(
             TargetSet(
                 [
                     TargetSpec(
-                        name="snap_total",
+                        name="household_count",
                         entity=EntityType.HOUSEHOLD,
-                        value=250.0,
+                        value=3.0,
                         period=2024,
-                        measure="snap",
-                        aggregation="sum",
+                        aggregation="count",
                     ),
                 ]
             )
@@ -622,9 +610,9 @@ def test_writes_policyengine_harness_when_baseline_and_targets_are_provided(
             policyengine_baseline_dataset=baseline_dataset,
             policyengine_harness_slices=(
                 PolicyEngineUSHarnessSlice(
-                    name="snap",
-                    description="SNAP parity",
-                    query=TargetQuery(period=2024, names=("snap_total",)),
+                    name="household_count",
+                    description="Household count parity",
+                    query=TargetQuery(period=2024, names=("household_count",)),
                 ),
             ),
             policyengine_harness_metadata={"baseline_dataset": baseline_dataset.name},
@@ -838,7 +826,7 @@ def test_writes_policyengine_harness_from_build_config_defaults(self, tmp_path):
                 policyengine_dataset_year=2024,
                 policyengine_targets_db=str(targets_db),
                 policyengine_baseline_dataset=str(baseline_dataset),
-                policyengine_target_variables=("snap", "household_count"),
+                policyengine_target_variables=("household_count",),
             ),
             seed_data=pd.DataFrame({"income": [10.0], "hh_weight": [1.0]}),
             synthetic_data=pd.DataFrame({"income": [10.0, 20.0], "weight": [1.0, 1.0]}),
@@ -921,10 +909,7 @@ def test_writes_policyengine_harness_from_build_config_defaults(self, tmp_path):
         assert harness_payload["metadata"]["targets_db"] == "policyengine_targets.db"
         assert harness_payload["metadata"]["harness_suite"] == "policyengine_us_all_targets"
         assert harness_payload["metadata"]["harness_slice_names"] == ["all_targets"]
-        assert harness_payload["metadata"]["target_variables"] == [
-            "snap",
-            "household_count",
-        ]
+        assert harness_payload["metadata"]["target_variables"] == ["household_count"]
         assert harness_payload["metadata"]["policyengine_us_runtime_version"] is not None
         assert [slice_payload["name"] for slice_payload in harness_payload["slices"]] == [
             "all_targets",