Fix unified calibration test drift

Author: MaxGhenis

policyengine_us_data/calibration/unified_matrix_builder.py

@@ -931,6 +931,7 @@ def __init__(
         self.time_period = time_period
         self.dataset_path = dataset_path
         self._entity_rel_cache = None
+        self._target_overview_columns = None
 
     # ---------------------------------------------------------------
     # Entity relationships
@@ -959,8 +960,8 @@ def _build_state_values(
         sim,
         target_vars: set,
         constraint_vars: set,
-        reform_vars: set,
-        geography,
+        reform_vars: set = None,
+        geography=None,
         rerandomize_takeup: bool = True,
         workers: int = 1,
     ) -> dict:
@@ -997,6 +998,9 @@ def _build_state_values(
             TAKEUP_AFFECTED_TARGETS,
         )
 
+        if geography is None:
+            raise ValueError("geography is required")
+
         unique_states = sorted(set(int(s) for s in geography.state_fips))
         n_hh = geography.n_records
 
@@ -1022,7 +1026,7 @@ def _build_state_values(
         # Convert sets to sorted lists for deterministic iteration
         target_vars_list = sorted(target_vars)
         constraint_vars_list = sorted(constraint_vars)
-        reform_vars_list = sorted(reform_vars)
+        reform_vars_list = sorted(reform_vars or set())
 
         state_values = {}
 
@@ -1518,63 +1522,103 @@ def _get_stratum_constraints(self, stratum_id: int) -> List[dict]:
             )
         return df.to_dict("records")
 
+    def _get_target_overview_columns(self) -> set:
+        if self._target_overview_columns is None:
+            with self.engine.connect() as conn:
+                rows = conn.execute(
+                    text("PRAGMA table_info(target_overview)")
+                ).fetchall()
+            self._target_overview_columns = {row[1] for row in rows}
+        return self._target_overview_columns
+
     def _query_targets(self, target_filter: dict) -> pd.DataFrame:
         """Query targets via target_overview view with
         best-period selection."""
-        or_conditions = []
+        and_conditions = []
 
         if "domain_variables" in target_filter:
             dvs = target_filter["domain_variables"]
             ph = ",".join(f"'{dv}'" for dv in dvs)
-            or_conditions.append(f"tv.domain_variable IN ({ph})")
+            and_conditions.append(f"tv.domain_variable IN ({ph})")
 
         if "variables" in target_filter:
             vs = ",".join(f"'{v}'" for v in target_filter["variables"])
-            or_conditions.append(f"tv.variable IN ({vs})")
+            and_conditions.append(f"tv.variable IN ({vs})")
 
         if "target_ids" in target_filter:
             ids = ",".join(map(str, target_filter["target_ids"]))
-            or_conditions.append(f"tv.target_id IN ({ids})")
+            and_conditions.append(f"tv.target_id IN ({ids})")
 
         if "stratum_ids" in target_filter:
             ids = ",".join(map(str, target_filter["stratum_ids"]))
-            or_conditions.append(f"tv.stratum_id IN ({ids})")
+            and_conditions.append(f"tv.stratum_id IN ({ids})")
 
-        if not or_conditions:
+        if not and_conditions:
             where_clause = "1=1"
         else:
-            where_clause = " OR ".join(f"({c})" for c in or_conditions)
-
-        query = f"""
-        WITH filtered_targets AS (
-            SELECT tv.target_id, tv.stratum_id, tv.variable, tv.reform_id,
-                   tv.value, tv.period, tv.geo_level,
-                   tv.geographic_id, tv.domain_variable
-            FROM target_overview tv
-            WHERE tv.active = 1
-              AND ({where_clause})
-        ),
-        best_periods AS (
-            SELECT stratum_id, variable, reform_id,
-                CASE
-                    WHEN MAX(CASE WHEN period <= :time_period
-                             THEN period END) IS NOT NULL
-                    THEN MAX(CASE WHEN period <= :time_period
-                             THEN period END)
-                    ELSE MIN(period)
-                END as best_period
-            FROM filtered_targets
-            GROUP BY stratum_id, variable, reform_id
-        )
-        SELECT ft.*
-        FROM filtered_targets ft
-        JOIN best_periods bp
-            ON ft.stratum_id = bp.stratum_id
-            AND ft.variable = bp.variable
-            AND ft.reform_id = bp.reform_id
-            AND ft.period = bp.best_period
-        ORDER BY ft.target_id
-        """
+            where_clause = " AND ".join(f"({c})" for c in and_conditions)
+
+        if "reform_id" in self._get_target_overview_columns():
+            query = f"""
+            WITH filtered_targets AS (
+                SELECT tv.target_id, tv.stratum_id, tv.variable, tv.reform_id,
+                       tv.value, tv.period, tv.geo_level,
+                       tv.geographic_id, tv.domain_variable
+                FROM target_overview tv
+                WHERE tv.active = 1
+                  AND ({where_clause})
+            ),
+            best_periods AS (
+                SELECT stratum_id, variable, reform_id,
+                    CASE
+                        WHEN MAX(CASE WHEN period <= :time_period
+                                 THEN period END) IS NOT NULL
+                        THEN MAX(CASE WHEN period <= :time_period
+                                 THEN period END)
+                        ELSE MIN(period)
+                    END as best_period
+                FROM filtered_targets
+                GROUP BY stratum_id, variable, reform_id
+            )
+            SELECT ft.*
+            FROM filtered_targets ft
+            JOIN best_periods bp
+                ON ft.stratum_id = bp.stratum_id
+                AND ft.variable = bp.variable
+                AND ft.reform_id = bp.reform_id
+                AND ft.period = bp.best_period
+            ORDER BY ft.target_id
+            """
+        else:
+            query = f"""
+            WITH filtered_targets AS (
+                SELECT tv.target_id, tv.stratum_id, tv.variable,
+                       0 AS reform_id, tv.value, tv.period, tv.geo_level,
+                       tv.geographic_id, tv.domain_variable
+                FROM target_overview tv
+                WHERE tv.active = 1
+                  AND ({where_clause})
+            ),
+            best_periods AS (
+                SELECT stratum_id, variable,
+                    CASE
+                        WHEN MAX(CASE WHEN period <= :time_period
+                                 THEN period END) IS NOT NULL
+                        THEN MAX(CASE WHEN period <= :time_period
+                                 THEN period END)
+                        ELSE MIN(period)
+                    END as best_period
+                FROM filtered_targets
+                GROUP BY stratum_id, variable
+            )
+            SELECT ft.*
+            FROM filtered_targets ft
+            JOIN best_periods bp
+                ON ft.stratum_id = bp.stratum_id
+                AND ft.variable = bp.variable
+                AND ft.period = bp.best_period
+            ORDER BY ft.target_id
+            """
 
         with self.engine.connect() as conn:
             return pd.read_sql(

policyengine_us_data/db/create_database_tables.py

@@ -508,6 +508,8 @@ def create_database(
 
     # Create SQL views
     with engine.connect() as conn:
+        conn.execute(text("DROP VIEW IF EXISTS stratum_domain"))
+        conn.execute(text("DROP VIEW IF EXISTS target_overview"))
         conn.execute(text(STRATUM_DOMAIN_VIEW))
         conn.execute(text(TARGET_OVERVIEW_VIEW))
         conn.commit()

policyengine_us_data/tests/test_calibration/test_unified_calibration.py

@@ -459,7 +459,7 @@ def test_county_var_uses_county_values(self):
         person_hh_idx = np.array([0, 1, 2, 3])
 
         builder = UnifiedMatrixBuilder.__new__(UnifiedMatrixBuilder)
-        hh_vars, _ = builder._assemble_clone_values(
+        hh_vars, _, _ = builder._assemble_clone_values(
             state_values,
             clone_states,
             person_hh_idx,
@@ -499,7 +499,7 @@ def test_non_county_var_uses_state_values(self):
         person_hh_idx = np.array([0, 1, 2, 3])
 
         builder = UnifiedMatrixBuilder.__new__(UnifiedMatrixBuilder)
-        hh_vars, _ = builder._assemble_clone_values(
+        hh_vars, _, _ = builder._assemble_clone_values(
             state_values,
             clone_states,
             person_hh_idx,

policyengine_us_data/tests/test_calibration/test_unified_matrix_builder.py

@@ -61,6 +61,59 @@ def _create_test_db(db_path):
     return db_uri, engine
 
 
+def _create_legacy_target_overview(engine):
+    legacy_view = """\
+CREATE VIEW target_overview AS
+SELECT
+    t.target_id,
+    t.stratum_id,
+    t.variable,
+    t.value,
+    t.period,
+    t.active,
+    CASE
+        WHEN MAX(CASE
+            WHEN sc.constraint_variable = 'congressional_district_geoid'
+                THEN 1
+            WHEN sc.constraint_variable = 'ucgid_str'
+                AND length(sc.value) = 13 THEN 1
+            ELSE 0 END) = 1 THEN 'district'
+        WHEN MAX(CASE
+            WHEN sc.constraint_variable = 'state_fips' THEN 1
+            WHEN sc.constraint_variable = 'ucgid_str'
+                AND length(sc.value) = 11 THEN 1
+            ELSE 0 END) = 1 THEN 'state'
+        ELSE 'national'
+    END AS geo_level,
+    COALESCE(
+        MAX(CASE
+            WHEN sc.constraint_variable = 'congressional_district_geoid'
+            THEN sc.value END),
+        MAX(CASE
+            WHEN sc.constraint_variable = 'state_fips'
+            THEN sc.value END),
+        MAX(CASE
+            WHEN sc.constraint_variable = 'ucgid_str'
+            THEN sc.value END),
+        'US'
+    ) AS geographic_id,
+    GROUP_CONCAT(DISTINCT CASE
+        WHEN sc.constraint_variable NOT IN (
+            'state_fips', 'congressional_district_geoid',
+            'tax_unit_is_filer', 'ucgid_str'
+        ) THEN sc.constraint_variable
+    END) AS domain_variable
+FROM targets t
+LEFT JOIN stratum_constraints sc ON t.stratum_id = sc.stratum_id
+GROUP BY t.target_id, t.stratum_id, t.variable,
+         t.value, t.period, t.active;
+"""
+    with engine.connect() as conn:
+        conn.execute(text("DROP VIEW target_overview"))
+        conn.execute(text(legacy_view))
+        conn.commit()
+
+
 def _insert_aca_ptc_data(engine):
     with engine.connect() as conn:
         strata = [1, 2, 3, 4, 5, 6, 7, 8, 9]
@@ -217,6 +270,20 @@ def test_inactive_targets_are_excluded(self):
         self.assertEqual(int(baseline_rows.iloc[0]["period"]), 2022)
         self.assertEqual(float(baseline_rows.iloc[0]["value"]), 10000.0)
 
+    def test_legacy_target_overview_without_reform_id(self):
+        _create_legacy_target_overview(self.engine)
+        try:
+            b = self._make_builder()
+            df = b._query_targets({"domain_variables": ["aca_ptc"]})
+            self.assertGreater(len(df), 0)
+            self.assertIn("reform_id", df.columns)
+            self.assertTrue((df["reform_id"] == 0).all())
+        finally:
+            with self.engine.connect() as conn:
+                conn.execute(text("DROP VIEW target_overview"))
+                conn.execute(text(TARGET_OVERVIEW_VIEW))
+                conn.commit()
+
     def test_target_name_adds_expenditure_suffix_for_reforms(self):
         name = UnifiedMatrixBuilder._make_target_name(
             "salt_deduction",