Add BEA state wage calibration targets (#1034)

* Add BEA state wage calibration targets

* Bump policyengine-us to 1.696.0

* Prioritize BEA wage controls in legacy reweighting

* Format BEA wage loss weights

* Relax Modal status seam timeout

* Bump policyengine-us to 1.697.0

* Retry Modal PR secret sync

Author: MaxGhenis

.github/scripts/sync_modal_secrets.py

@@ -1,10 +1,28 @@
 import os
 import subprocess
+import time
+
+
+def create_secret_with_retry(args: list[str]) -> None:
+    max_attempts = 5
+    for attempt in range(1, max_attempts + 1):
+        try:
+            subprocess.run(args, check=True)
+            return
+        except subprocess.CalledProcessError:
+            if attempt == max_attempts:
+                raise
+            delay = min(2**attempt, 10)
+            print(
+                "Modal secret creation failed; retrying "
+                f"in {delay}s ({attempt}/{max_attempts})"
+            )
+            time.sleep(delay)
 
 
 def main() -> None:
     env_name = os.environ["MODAL_ENVIRONMENT"]
-    subprocess.run(
+    create_secret_with_retry(
         [
             "modal",
             "secret",
@@ -15,9 +33,8 @@ def main() -> None:
             "huggingface-token",
             f"HUGGING_FACE_TOKEN={os.environ['HUGGING_FACE_TOKEN']}",
         ],
-        check=True,
     )
-    subprocess.run(
+    create_secret_with_retry(
         [
             "modal",
             "secret",
@@ -31,7 +48,6 @@ def main() -> None:
                 f"{os.environ['GOOGLE_APPLICATION_CREDENTIALS']}"
             ),
         ],
-        check=True,
     )
 
 

AGENTS.md

@@ -29,8 +29,10 @@ Manually sourced national or local-file calibration targets must be registered
 in every active target path before merging:
 
 1. `policyengine_us_data/utils/loss.py` for the ECPS loss matrix.
-2. `policyengine_us_data/db/etl_national_targets.py` for `policy_data.db` and
-   local H5 validation inputs.
+2. The appropriate `policyengine_us_data/db/etl_*.py` loader for
+   `policy_data.db` and local H5 validation inputs. National targets usually
+   belong in `etl_national_targets.py`; state or local targets should use a
+   state/local ETL module and must still be added to this DB path.
 3. `policyengine_us_data/calibration/target_config.yaml` when the default
    calibration uses an `include:` list; otherwise the target can exist in
    `policy_data.db` but still be omitted from calibration.

Makefile

@@ -84,6 +84,7 @@ database:
 	python -m policyengine_us_data.db.create_database_tables
 	python -m policyengine_us_data.db.create_initial_strata --year $(YEAR)
 	python -m policyengine_us_data.db.etl_national_targets --year $(YEAR)
+	python -m policyengine_us_data.db.etl_bea_state_wages --year $(YEAR)
 	python -m policyengine_us_data.db.etl_age --year $(YEAR)
 	python -m policyengine_us_data.db.etl_medicaid --year $(YEAR)
 	python -m policyengine_us_data.db.etl_snap --year $(YEAR)

changelog.d/1033.added

@@ -0,0 +1 @@
+Add BEA regional state wage targets to calibration.

docs/engineering/skills/calibration_targets.md

@@ -8,8 +8,10 @@ calibration targets.
 New targets must be registered in both active target systems:
 
 - `policyengine_us_data/utils/loss.py` for the ECPS `build_loss_matrix()` path.
-- `policyengine_us_data/db/etl_national_targets.py` for `policy_data.db`, local
-  H5 outputs, and validation inputs.
+- The appropriate `policyengine_us_data/db/etl_*.py` loader for
+  `policy_data.db`, local H5 outputs, and validation inputs. National targets
+  usually belong in `etl_national_targets.py`; state and local targets should
+  use a state/local ETL module and must still be present in this DB path.
 
 If the default calibration path uses `policyengine_us_data/calibration/target_config.yaml`
 with an `include:` list, also add the matching include rule there. A target can
@@ -19,10 +21,10 @@ from `target_config.yaml`.
 ## Tests
 
 Every target change should add or update tests that prove the target is wired
-through every active path. For manually sourced national targets, cover:
+through every active path. For manually sourced targets, cover:
 
 - the ECPS loss matrix registration in `tests/unit/calibration/test_loss_targets.py`;
-- the DB ETL row in `tests/unit/test_etl_national_targets.py`;
+- the DB ETL row in the matching `tests/unit/test_etl_*.py` file;
 - the default calibration include rule in
   `tests/unit/calibration/test_target_config.py`;
 - any publication guard in `tests/unit/test_upload_completed_datasets.py` when

policyengine_us_data/calibration/target_config.yaml

@@ -58,6 +58,11 @@ include:
     geo_level: state
   - variable: adjusted_gross_income
     geo_level: state
+  # BEA regional wage totals. New macro targets like this must be wired into
+  # both legacy utils/loss.py and the target DB ETL, or one publication path
+  # will silently miss the constraint.
+  - variable: employment_income_before_lsr
+    geo_level: state
   - variable: rent
     geo_level: state
   - variable: spm_unit_count

policyengine_us_data/datasets/cps/enhanced_cps.py

@@ -8,6 +8,7 @@
     ABSOLUTE_ERROR_SCALE_TARGETS,
     build_loss_matrix,
     get_target_error_normalisation,
+    get_target_loss_weights,
     HardConcrete,
     print_reweighting_diagnostics,
     set_seeds,
@@ -506,7 +507,9 @@ def reweight(
     normalisation_factor = np.where(
         is_national, nation_normalisation_factor, state_normalisation_factor
     )
+    target_loss_weights = get_target_loss_weights(target_names)
     normalisation_factor = torch.tensor(normalisation_factor, dtype=torch.float32)
+    target_loss_weights = torch.tensor(target_loss_weights, dtype=torch.float32)
     targets_array = torch.tensor(targets_array, dtype=torch.float32)
     numerator_shift = torch.tensor(numerator_shift_np, dtype=torch.float32)
     error_denominator = torch.tensor(error_denominator_np, dtype=torch.float32)
@@ -525,6 +528,7 @@ def loss(weights):
             (estimate - targets_array + numerator_shift) / error_denominator
         ) ** 2
         rel_error_normalized = inv_mean_normalisation * rel_error * normalisation_factor
+        rel_error_normalized = rel_error_normalized * target_loss_weights
         if torch.isnan(rel_error_normalized).any():
             raise ValueError("Relative error contains NaNs")
         return rel_error_normalized.mean()

policyengine_us_data/db/create_field_valid_values.py

@@ -69,6 +69,7 @@ def populate_field_valid_values(session: Session) -> None:
     # Static values for source field
     source_values = [
         ("source", "Census ACS S0101", "survey"),
+        ("source", "BEA Regional SAINC4", "administrative"),
         ("source", "IRS SOI", "administrative"),
         ("source", "IRS EITC Central", "administrative"),
         ("source", "CMS Marketplace", "administrative"),

policyengine_us_data/db/etl_bea_state_wages.py

@@ -0,0 +1,121 @@
+"""ETL for BEA regional state wage calibration targets."""
+
+import logging
+
+import pandas as pd
+from sqlmodel import Session, create_engine
+
+from policyengine_us_data.db.etl_national_targets import (
+    BEA_NIPA_WAGES_AND_SALARIES_2024,
+    _register_target_variable,
+    _upsert_baseline_target,
+)
+from policyengine_us_data.storage import STORAGE_FOLDER
+from policyengine_us_data.utils.bea_regional import (
+    BEA_STATE_WAGES_SOURCE,
+    BEA_STATE_WAGES_SOURCE_URL,
+    get_bea_state_wage_targets,
+)
+from policyengine_us_data.utils.db import etl_argparser, get_geographic_strata
+
+logger = logging.getLogger(__name__)
+
+TARGET_VARIABLE = "employment_income_before_lsr"
+
+
+def extract_bea_state_wage_targets(year: int) -> tuple[pd.DataFrame, int]:
+    """Extract BEA state wage targets scaled to the national NIPA total."""
+    return get_bea_state_wage_targets(
+        year,
+        national_total=BEA_NIPA_WAGES_AND_SALARIES_2024,
+    )
+
+
+def load_bea_state_wage_targets(
+    targets: pd.DataFrame,
+    *,
+    target_year: int,
+    source_year: int,
+) -> int:
+    """Load BEA state wage targets into state geographic strata."""
+    if targets.empty:
+        return 0
+
+    database_url = f"sqlite:///{STORAGE_FOLDER / 'calibration' / 'policy_data.db'}"
+    engine = create_engine(database_url)
+    loaded = 0
+
+    with Session(engine) as session:
+        _register_target_variable(session, TARGET_VARIABLE)
+        geo_strata = get_geographic_strata(session)
+        state_strata = geo_strata.get("state", {})
+
+        for row in targets.itertuples(index=False):
+            state_fips = int(row.state_fips)
+            stratum_id = state_strata.get(state_fips)
+            if stratum_id is None:
+                logger.warning(
+                    "No geographic stratum found for state %s (FIPS %s), skipping",
+                    row.state_code,
+                    state_fips,
+                )
+                continue
+
+            _upsert_baseline_target(
+                session,
+                stratum_id=stratum_id,
+                variable=TARGET_VARIABLE,
+                period=target_year,
+                value=float(row.employment_income_before_lsr),
+                source=BEA_STATE_WAGES_SOURCE,
+                notes=(
+                    "BEA SAINC4 line 50 wages and salaries by state, adjusted "
+                    "to a residence basis by allocating line 42's residence "
+                    "adjustment to wages in proportion to place-of-work "
+                    "net-compensation components, then scaled to the national "
+                    "BEA NIPA Table 2.1 wages and salaries target. "
+                    f"Source year: {source_year}; state: {row.state_code}; "
+                    f"raw residence-adjusted state wages: "
+                    f"${row.wages_and_salaries:,.0f}; "
+                    f"national scaling factor: {row.scale_factor:.8f}; "
+                    f"Source: {BEA_STATE_WAGES_SOURCE_URL}"
+                ),
+            )
+            loaded += 1
+
+        session.commit()
+
+    logger.info("Loaded %s BEA state wage targets", loaded)
+    return loaded
+
+
+def main():
+    logging.basicConfig(
+        level=logging.INFO,
+        format="%(asctime)s - %(name)s - %(levelname)s - %(message)s",
+    )
+    _, year = etl_argparser(
+        "ETL for BEA regional state wage calibration targets",
+        allow_year=True,
+    )
+
+    targets, source_year = extract_bea_state_wage_targets(year)
+    loaded = load_bea_state_wage_targets(
+        targets,
+        target_year=year,
+        source_year=source_year,
+    )
+
+    logger.info(
+        "BEA State Wage Targets Summary:\n"
+        "  Source year: %s\n"
+        "  States loaded: %s\n"
+        "  Target total: $%.1fT",
+        source_year,
+        loaded,
+        targets["employment_income_before_lsr"].sum() / 1e12,
+    )
+
+
+if __name__ == "__main__":
+    main()

policyengine_us_data/utils/__init__.py

@@ -13,6 +13,7 @@
     "HardConcrete",
     "build_loss_matrix",
     "get_target_error_normalisation",
+    "get_target_loss_weights",
     "print_reweighting_diagnostics",
     "set_seeds",
 ]