remove redundant tests and fixtures

Author: juaristi22

policyengine_us_data/tests/test_local_area_calibration/conftest.py

@@ -1,4 +1,7 @@
-"""Shared fixtures for local area calibration tests."""
+"""Shared fixtures for local area calibration tests.
+
+Importantly, this file determines which variables will be included in the sparse matrix and calibrating routine.
+"""
 
 import pytest
 import numpy as np
@@ -16,6 +19,56 @@
     get_calculated_variables,
 )
 
+# Variables to test for state-level value matching
+# Format: (variable_name, rtol)
+#     variable_name as per the targets in policy_data.db
+#     rtol is relative tolerance for comparison
+VARIABLES_TO_TEST = [
+    ("snap", 1e-2),
+    ("health_insurance_premiums_without_medicare_part_b", 1e-2),
+    ("medicaid", 1e-2),
+    ("medicare_part_b_premiums", 1e-2),
+    ("other_medical_expenses", 1e-2),
+    ("over_the_counter_health_expenses", 1e-2),
+    ("salt_deduction", 1e-2),
+    ("spm_unit_capped_work_childcare_expenses", 1e-2),
+    ("spm_unit_capped_housing_subsidy", 1e-2),
+    ("ssi", 1e-2),
+    ("tanf", 1e-2),
+    ("tip_income", 1e-2),
+    ("unemployment_compensation", 1e-2),
+]
+
+# Combined filter config to build matrix with all variables at once
+COMBINED_FILTER_CONFIG = {
+    "stratum_group_ids": [
+        4,  # SNAP targets
+        5,  # Medicaid targets
+        112,  # Unemployment compensation targets
+    ],
+    "variables": [
+        "snap",
+        "health_insurance_premiums_without_medicare_part_b",
+        "medicaid",
+        "medicare_part_b_premiums",
+        "other_medical_expenses",
+        "over_the_counter_health_expenses",
+        "salt_deduction",
+        "spm_unit_capped_work_childcare_expenses",
+        "spm_unit_capped_housing_subsidy",
+        "ssi",
+        "tanf",
+        "tip_income",
+        "unemployment_compensation",
+    ],
+}
+
+# Maximum allowed mismatch rate for state-level value comparison
+MAX_MISMATCH_RATE = 0.02
+
+# Number of samples for cell-level verification tests
+N_VERIFICATION_SAMPLES = 200
+
 
 @pytest.fixture(scope="module")
 def db_uri():
@@ -30,7 +83,7 @@ def dataset_path():
 
 @pytest.fixture(scope="module")
 def test_cds(db_uri):
-    """CDs from NC, HI, MT, AK (manageable size, multiple same-state CDs)."""
+    """CDs from multiple states for comprehensive testing."""
     engine = create_engine(db_uri)
     query = """
     SELECT DISTINCT sc.value as cd_geoid
@@ -43,6 +96,10 @@ def test_cds(db_uri):
         OR sc.value LIKE '150_'
         OR sc.value LIKE '300_'
         OR sc.value = '200' OR sc.value = '201'
+        OR sc.value IN ('101', '102')
+        OR sc.value IN ('601', '602')
+        OR sc.value IN ('3601', '3602')
+        OR sc.value IN ('4801', '4802')
       )
     ORDER BY sc.value
     """
@@ -58,15 +115,15 @@ def sim(dataset_path):
 
 @pytest.fixture(scope="module")
 def matrix_data(db_uri, dataset_path, test_cds, sim):
-    """Build sparse matrix, return (targets_df, X_sparse, household_id_mapping)."""
+    """Build sparse matrix with all configured variables."""
     builder = SparseMatrixBuilder(
         db_uri,
         time_period=2023,
         cds_to_calibrate=test_cds,
         dataset_path=dataset_path,
     )
     targets_df, X_sparse, household_id_mapping = builder.build_matrix(
-        sim, target_filter={"stratum_group_ids": [4], "variables": ["snap"]}
+        sim, target_filter=COMBINED_FILTER_CONFIG
     )
     return targets_df, X_sparse, household_id_mapping
 

policyengine_us_data/tests/test_local_area_calibration/test_cross_state.py

@@ -2,17 +2,19 @@
 
 import pytest
 import numpy as np
+from collections import defaultdict
 
 from policyengine_us import Microsimulation
 from policyengine_us_data.datasets.cps.local_area_calibration.calibration_utils import (
     get_calculated_variables,
 )
 
+from .conftest import VARIABLES_TO_TEST, N_VERIFICATION_SAMPLES
+
 
 def test_cross_state_matches_swapped_sim(
     X_sparse,
     targets_df,
-    tracer,
     test_cds,
     dataset_path,
     n_households,
@@ -25,8 +27,10 @@ def test_cross_state_matches_swapped_sim(
     When household moves to different state, X_sparse should contain the
     value calculated from a fresh simulation with state_fips set to
     destination state.
+
+    Uses stratified sampling to ensure all variables in VARIABLES_TO_TEST
+    are covered with approximately equal samples per variable.
     """
-    n_samples = 200
     seed = 42
     rng = np.random.default_rng(seed)
     n_hh = n_households
@@ -48,28 +52,46 @@ def get_state_sim(state):
 
     nonzero_rows, nonzero_cols = X_sparse.nonzero()
 
-    cross_state_indices = []
+    # Group cross-state cells by variable for stratified sampling
+    variable_to_indices = defaultdict(list)
+    variables_to_test = {v[0] for v in VARIABLES_TO_TEST}
+
     for i in range(len(nonzero_rows)):
+        row_idx = nonzero_rows[i]
         col_idx = nonzero_cols[i]
         cd_idx = col_idx // n_hh
         hh_idx = col_idx % n_hh
         cd = test_cds[cd_idx]
         dest_state = int(cd) // 100
         orig_state = int(hh_states[hh_idx])
-        if dest_state != orig_state:
-            cross_state_indices.append(i)
 
-    if not cross_state_indices:
-        pytest.skip("No cross-state non-zero cells found")
+        # Only include cross-state cells
+        if dest_state == orig_state:
+            continue
+
+        # Get variable for this row
+        variable = targets_df.iloc[row_idx]["variable"]
+        if variable in variables_to_test:
+            variable_to_indices[variable].append(i)
+
+    if not variable_to_indices:
+        pytest.skip("No cross-state non-zero cells found for test variables")
 
-    sample_idx = rng.choice(
-        cross_state_indices,
-        min(n_samples, len(cross_state_indices)),
-        replace=False,
+    # Stratified sampling: sample proportionally from each variable
+    samples_per_var = max(
+        1, N_VERIFICATION_SAMPLES // len(variable_to_indices)
     )
+    sample_indices = []
+
+    for variable, indices in variable_to_indices.items():
+        n_to_sample = min(samples_per_var, len(indices))
+        sampled = rng.choice(indices, n_to_sample, replace=False)
+        sample_indices.extend(sampled)
+
     errors = []
+    variables_tested = set()
 
-    for idx in sample_idx:
+    for idx in sample_indices:
         row_idx = nonzero_rows[idx]
         col_idx = nonzero_cols[idx]
         cd_idx = col_idx // n_hh
@@ -83,6 +105,8 @@ def get_state_sim(state):
             state_sim.calculate(variable, map_to="household").values[hh_idx]
         )
 
+        variables_tested.add(variable)
+
         if not np.isclose(actual, expected, atol=0.5):
             errors.append(
                 {
@@ -95,7 +119,13 @@ def get_state_sim(state):
                 }
             )
 
+    # Report which variables were tested
+    missing_vars = variables_to_test - variables_tested
+    if missing_vars:
+        print(f"Warning: No cross-state cells found for: {missing_vars}")
+
     assert not errors, (
-        f"Cross-state verification failed: {len(errors)}/{len(sample_idx)} "
-        f"mismatches. First 5: {errors[:5]}"
+        f"Cross-state verification failed: {len(errors)}/{len(sample_indices)} "
+        f"mismatches across {len(variables_tested)} variables. "
+        f"First 5: {errors[:5]}"
     )