Fix post-reform calibration test compatibility

Author: MaxGhenis

policyengine_us_data/calibration/unified_matrix_builder.py

@@ -19,6 +19,7 @@
 from scipy import sparse
 from sqlalchemy import create_engine, text
 
+from policyengine_us_data.db.create_database_tables import refresh_sql_views
 from policyengine_us_data.storage import STORAGE_FOLDER
 from policyengine_us_data.utils.census import STATE_NAME_TO_FIPS
 from policyengine_us_data.calibration.calibration_utils import (
@@ -928,6 +929,8 @@ def __init__(
     ):
         self.db_uri = db_uri
         self.engine = create_engine(db_uri)
+        # Existing SQLite checkpoints may carry an older target_overview view.
+        refresh_sql_views(self.engine)
         self.time_period = time_period
         self.dataset_path = dataset_path
         self._entity_rel_cache = None

policyengine_us_data/calibration/validate_staging.py

@@ -41,6 +41,7 @@
 from policyengine_us_data.calibration.sanity_checks import (
     run_sanity_checks,
 )
+from policyengine_us_data.db.create_database_tables import refresh_sql_views
 
 logger = logging.getLogger(__name__)
 
@@ -537,6 +538,7 @@ def _validate_single_area(
     from sqlalchemy import create_engine as _create_engine
 
     engine = _create_engine(f"sqlite:///{db_path}")
+    refresh_sql_views(engine)
 
     logger.info("Loading sim from %s", h5_path)
     try:
@@ -1015,6 +1017,7 @@ def main(argv=None):
     from policyengine_us import Microsimulation
 
     engine = create_engine(f"sqlite:///{args.db_path}")
+    refresh_sql_views(engine)
 
     all_targets = _query_all_active_targets(engine, args.period)
     logger.info("Loaded %d active targets from DB", len(all_targets))

policyengine_us_data/db/create_database_tables.py

@@ -353,6 +353,16 @@ def validate_parent_child_constraints(mapper, connection, target: Stratum):
 """
 
 
+def refresh_sql_views(engine) -> None:
+    """Recreate derived SQL views so existing DBs pick up schema changes."""
+    with engine.connect() as conn:
+        conn.execute(text("DROP VIEW IF EXISTS target_overview"))
+        conn.execute(text("DROP VIEW IF EXISTS stratum_domain"))
+        conn.execute(text(STRATUM_DOMAIN_VIEW))
+        conn.execute(text(TARGET_OVERVIEW_VIEW))
+        conn.commit()
+
+
 def create_validation_triggers(engine) -> None:
     """Create SQL triggers that validate fields against field_valid_values.
 
@@ -506,11 +516,8 @@ def create_database(
     # Create validation triggers
     create_validation_triggers(engine)
 
-    # Create SQL views
-    with engine.connect() as conn:
-        conn.execute(text(STRATUM_DOMAIN_VIEW))
-        conn.execute(text(TARGET_OVERVIEW_VIEW))
-        conn.commit()
+    # Recreate SQL views so existing DB files do not keep stale definitions.
+    refresh_sql_views(engine)
 
     logger.info(f"Database and tables created successfully at {db_uri}")
     return engine

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

@@ -211,7 +211,7 @@ def test_reform_targets_preserved(self):
 
     def test_inactive_targets_are_excluded(self):
         b = self._make_builder(time_period=2024)
-        df = b._query_targets({"stratum_ids": [1], "variables": ["aca_ptc"]})
+        df = b._query_targets({"target_ids": [1, 18, 19]})
         baseline_rows = df[(df["variable"] == "aca_ptc") & (df["reform_id"] == 0)]
         self.assertEqual(len(baseline_rows), 1)
         self.assertEqual(int(baseline_rows.iloc[0]["period"]), 2022)
@@ -526,6 +526,7 @@ def test_return_structure_no_takeup(self, mock_msim_cls, mock_gcv):
             sim=None,
             target_vars={"snap"},
             constraint_vars={"income"},
+            reform_vars=set(),
             geography=geo,
             rerandomize_takeup=False,
         )
@@ -560,6 +561,7 @@ def test_fresh_sim_per_state(self, mock_msim_cls, mock_gcv):
             sim=None,
             target_vars={"snap"},
             constraint_vars=set(),
+            reform_vars=set(),
             geography=geo,
             rerandomize_takeup=False,
         )
@@ -582,6 +584,7 @@ def test_state_fips_set_correctly(self, mock_msim_cls, mock_gcv):
             sim=None,
             target_vars={"snap"},
             constraint_vars=set(),
+            reform_vars=set(),
             geography=geo,
             rerandomize_takeup=False,
         )
@@ -617,6 +620,7 @@ def test_takeup_vars_forced_true(self, mock_msim_cls, mock_gcv):
             sim=None,
             target_vars={"snap"},
             constraint_vars=set(),
+            reform_vars=set(),
             geography=geo,
             rerandomize_takeup=True,
         )
@@ -664,6 +668,7 @@ def test_count_vars_skipped(self, mock_msim_cls, mock_gcv):
             sim=None,
             target_vars={"snap", "snap_count"},
             constraint_vars=set(),
+            reform_vars=set(),
             geography=geo,
             rerandomize_takeup=False,
         )
@@ -914,6 +919,7 @@ def test_workers_gt1_creates_pool(self, mock_msim_cls, mock_gcv, mock_pool_cls):
                 sim=None,
                 target_vars={"snap"},
                 constraint_vars=set(),
+                reform_vars=set(),
                 geography=geo,
                 rerandomize_takeup=False,
                 workers=2,
@@ -939,6 +945,7 @@ def test_workers_1_skips_pool(self, mock_msim_cls, mock_gcv):
                 sim=None,
                 target_vars={"snap"},
                 constraint_vars=set(),
+                reform_vars=set(),
                 geography=geo,
                 rerandomize_takeup=False,
                 workers=1,

policyengine_us_data/tests/test_schema_views_and_lookups.py

@@ -11,13 +11,15 @@
 import tempfile
 import unittest
 
+from sqlalchemy import text
 from sqlmodel import Session
 
 from policyengine_us_data.db.create_database_tables import (
     Stratum,
     StratumConstraint,
     Target,
     create_database,
+    refresh_sql_views,
 )
 from policyengine_us_data.utils.db import get_geographic_strata
 from policyengine_us_data.calibration.calibration_utils import (
@@ -399,6 +401,24 @@ def test_reform_id_passthrough(self):
         self.assertEqual(len(matches), 1)
         self.assertEqual(matches[0][reform_idx], 1)
 
+    def test_refresh_sql_views_updates_existing_target_overview(self):
+        """Refreshing views updates stale target_overview definitions."""
+        with self.engine.connect() as conn:
+            conn.execute(text("DROP VIEW IF EXISTS target_overview"))
+            conn.execute(
+                text(
+                    "CREATE VIEW target_overview AS "
+                    "SELECT target_id, stratum_id, variable, value, period, active "
+                    "FROM targets"
+                )
+            )
+            conn.commit()
+
+        refresh_sql_views(self.engine)
+
+        cols = self._overview_columns()
+        self.assertIn("reform_id", cols)
+
     # ----------------------------------------------------------------
     # get_geographic_strata()
     # ----------------------------------------------------------------