update test_same_test so it compares to original values instead of fresh calculations

Author: juaristi22

policyengine_us_data/tests/test_local_area_calibration/test_same_state.py

@@ -1,31 +1,26 @@
-"""Test same-state values match fresh simulations."""
+"""Test same-state values match original simulation values."""
 
 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_same_state_matches_original(
+    sim,
     X_sparse,
     targets_df,
     test_cds,
-    dataset_path,
     n_households,
     household_ids,
     household_states,
 ):
     """
-    Same-state non-zero cells must match fresh same-state simulation.
+    Same-state non-zero cells must match ORIGINAL simulation values.
 
     When household stays in same state, X_sparse should contain the value
-    calculated from a fresh simulation with state_fips set to that state.
+    from the original simulation (ground truth from H5 dataset).
 
     Uses stratified sampling to ensure all variables in VARIABLES_TO_TEST
     are covered with approximately equal samples per variable.
@@ -36,19 +31,6 @@ def test_same_state_matches_original(
     hh_ids = household_ids
     hh_states = household_states
 
-    state_sims = {}
-
-    def get_state_sim(state):
-        if state not in state_sims:
-            s = Microsimulation(dataset=dataset_path)
-            s.set_input(
-                "state_fips", 2023, np.full(n_hh, state, dtype=np.int32)
-            )
-            for var in get_calculated_variables(s):
-                s.delete_arrays(var)
-            state_sims[state] = s
-        return state_sims[state]
-
     nonzero_rows, nonzero_cols = X_sparse.nonzero()
 
     # Group same-state cells by variable for stratified sampling
@@ -68,7 +50,6 @@ def get_state_sim(state):
         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)
@@ -87,6 +68,16 @@ def get_state_sim(state):
         sampled = rng.choice(indices, n_to_sample, replace=False)
         sample_indices.extend(sampled)
 
+    # Cache original values per variable to avoid repeated calculations
+    original_values_cache = {}
+
+    def get_original_values(variable):
+        if variable not in original_values_cache:
+            original_values_cache[variable] = sim.calculate(
+                variable, map_to="household"
+            ).values
+        return original_values_cache[variable]
+
     errors = []
     variables_tested = set()
 
@@ -95,28 +86,32 @@ def get_state_sim(state):
         col_idx = nonzero_cols[idx]
         cd_idx = col_idx // n_hh
         hh_idx = col_idx % n_hh
-        cd = test_cds[cd_idx]
-        dest_state = int(cd) // 100
         variable = targets_df.iloc[row_idx]["variable"]
         actual = float(X_sparse[row_idx, col_idx])
-        state_sim = get_state_sim(dest_state)
-        expected = float(
-            state_sim.calculate(variable, map_to="household").values[hh_idx]
-        )
+
+        # Compare to ORIGINAL simulation values (ground truth)
+        original_values = get_original_values(variable)
+        expected = float(original_values[hh_idx])
 
         variables_tested.add(variable)
 
         if not np.isclose(actual, expected, atol=0.5):
             errors.append(
                 {
                     "hh_id": hh_ids[hh_idx],
+                    "hh_idx": hh_idx,
                     "variable": variable,
                     "actual": actual,
                     "expected": expected,
+                    "diff": actual - expected,
+                    "rel_diff": (
+                        (actual - expected) / expected
+                        if expected != 0
+                        else np.inf
+                    ),
                 }
             )
 
-    # Report which variables were tested
     missing_vars = variables_to_test - variables_tested
     if missing_vars:
         print(f"Warning: No same-state cells found for: {missing_vars}")