From a74e72c7a58c235095835172654fa66e9b5056cd Mon Sep 17 00:00:00 2001 From: Max Ghenis Date: Sun, 26 Apr 2026 07:40:10 -0400 Subject: [PATCH 1/9] Update GitHub Actions for Node 24 runtime --- .github/workflows/main.yml | 10 +++++----- .github/workflows/pr.yaml | 20 ++++++++++---------- .github/workflows/versioning.yaml | 12 ++++++------ 3 files changed, 21 insertions(+), 21 deletions(-) diff --git a/.github/workflows/main.yml b/.github/workflows/main.yml index 32b8fcf..5638d70 100644 --- a/.github/workflows/main.yml +++ b/.github/workflows/main.yml @@ -13,11 +13,11 @@ jobs: steps: - name: Checkout repo - uses: actions/checkout@v3 + uses: actions/checkout@v6 - name: Install uv - uses: astral-sh/setup-uv@v5 + uses: astral-sh/setup-uv@v8.1.0 - name: Set up Python ${{ matrix.python-version }} - uses: actions/setup-python@v4 + uses: actions/setup-python@v6 with: python-version: ${{ matrix.python-version }} - name: Install Python dependencies @@ -26,7 +26,7 @@ jobs: - name: Run tests with coverage run: make test - name: Upload coverage to Codecov - uses: codecov/codecov-action@v3 + uses: codecov/codecov-action@v6 with: file: ./coverage.xml fail_ci_if_error: false @@ -43,7 +43,7 @@ jobs: fi done - name: Deploy documentation - uses: JamesIves/github-pages-deploy-action@releases/v3 + uses: JamesIves/github-pages-deploy-action@v4 with: GITHUB_TOKEN: ${{ secrets.GITHUB_TOKEN }} BRANCH: gh-pages # The branch the action should deploy to. diff --git a/.github/workflows/pr.yaml b/.github/workflows/pr.yaml index 9fc7e29..0d1612c 100644 --- a/.github/workflows/pr.yaml +++ b/.github/workflows/pr.yaml @@ -7,13 +7,13 @@ jobs: Lint: runs-on: ubuntu-latest steps: - - uses: actions/checkout@v3 + - uses: actions/checkout@v6 - name: Set up Python - uses: actions/setup-python@v4 + uses: actions/setup-python@v6 with: python-version: "3.13" - name: Install uv - uses: astral-sh/setup-uv@v5 + uses: astral-sh/setup-uv@v8.1.0 - name: Install relevant dependencies run: | uv pip install black isort linecheck --system @@ -29,11 +29,11 @@ jobs: runs-on: ${{ matrix.os }} steps: - name: Checkout repo - uses: actions/checkout@v3 + uses: actions/checkout@v6 - name: Install uv - uses: astral-sh/setup-uv@v5 + uses: astral-sh/setup-uv@v8.1.0 - name: Set up Python ${{ matrix.python-version }} - uses: actions/setup-python@v4 + uses: actions/setup-python@v6 with: python-version: ${{ matrix.python-version }} - name: Install Python dependencies @@ -42,7 +42,7 @@ jobs: - name: Run tests with coverage run: make test - name: Upload coverage to Codecov - uses: codecov/codecov-action@v3 + uses: codecov/codecov-action@v6 with: file: ./coverage.xml fail_ci_if_error: false @@ -52,11 +52,11 @@ jobs: runs-on: ubuntu-latest steps: - name: Checkout repo - uses: actions/checkout@v3 + uses: actions/checkout@v6 - name: Install uv - uses: astral-sh/setup-uv@v5 + uses: astral-sh/setup-uv@v8.1.0 - name: Set up Python - uses: actions/setup-python@v4 + uses: actions/setup-python@v6 with: python-version: "3.13" - name: Install dependencies diff --git a/.github/workflows/versioning.yaml b/.github/workflows/versioning.yaml index dfd4f1e..dc19175 100644 --- a/.github/workflows/versioning.yaml +++ b/.github/workflows/versioning.yaml @@ -18,18 +18,18 @@ jobs: steps: - name: Generate GitHub App token id: app-token - uses: actions/create-github-app-token@v1 + uses: actions/create-github-app-token@v3 with: app-id: ${{ secrets.APP_ID }} private-key: ${{ secrets.APP_PRIVATE_KEY }} - name: Checkout repo - uses: actions/checkout@v4 + uses: actions/checkout@v6 with: repository: ${{ github.event.pull_request.head.repo.full_name }} ref: ${{ github.event.pull_request.head.ref }} token: ${{ steps.app-token.outputs.token }} - name: Setup Python - uses: actions/setup-python@v5 + uses: actions/setup-python@v6 with: python-version: 3.13 - name: Build changelog @@ -37,7 +37,7 @@ jobs: - name: Preview changelog update run: ".github/get-changelog-diff.sh" - name: Update changelog - uses: EndBug/add-and-commit@v9 + uses: EndBug/add-and-commit@v10 with: add: "." message: Update package version @@ -49,11 +49,11 @@ jobs: runs-on: ubuntu-latest steps: - name: Checkout code - uses: actions/checkout@v4 + uses: actions/checkout@v6 with: fetch-depth: 0 # Fetch all history for all tags and branches - name: Set up Python - uses: actions/setup-python@v5 + uses: actions/setup-python@v6 with: python-version: 3.13 - name: Install package From 056efa85402406df6f637072ff6e542fd2bdc91e Mon Sep 17 00:00:00 2001 From: Max Ghenis Date: Sun, 26 Apr 2026 08:15:52 -0400 Subject: [PATCH 2/9] Update remaining GitHub Actions refs for Node 24 --- .github/workflows/changelog_entry.yaml | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/.github/workflows/changelog_entry.yaml b/.github/workflows/changelog_entry.yaml index 911c28e..72ddc65 100644 --- a/.github/workflows/changelog_entry.yaml +++ b/.github/workflows/changelog_entry.yaml @@ -9,7 +9,7 @@ jobs: name: Changelog entry check runs-on: ubuntu-latest steps: - - uses: actions/checkout@v3 + - uses: actions/checkout@v6 - name: Check for changelog entry run: | From a2d8cc182e4855ac7a22d72a8e17317301a9056f Mon Sep 17 00:00:00 2001 From: Max Ghenis Date: Sun, 26 Apr 2026 10:12:43 -0400 Subject: [PATCH 3/9] Fix CI warnings and changelog checks --- .github/workflows/main.yml | 2 +- .github/workflows/pr.yaml | 2 +- changelog_entry.yaml | 4 ++++ 3 files changed, 6 insertions(+), 2 deletions(-) diff --git a/.github/workflows/main.yml b/.github/workflows/main.yml index 5638d70..f629412 100644 --- a/.github/workflows/main.yml +++ b/.github/workflows/main.yml @@ -28,7 +28,7 @@ jobs: - name: Upload coverage to Codecov uses: codecov/codecov-action@v6 with: - file: ./coverage.xml + files: ./coverage.xml fail_ci_if_error: false verbose: true - name: Test documentation builds diff --git a/.github/workflows/pr.yaml b/.github/workflows/pr.yaml index 0d1612c..dee124a 100644 --- a/.github/workflows/pr.yaml +++ b/.github/workflows/pr.yaml @@ -44,7 +44,7 @@ jobs: - name: Upload coverage to Codecov uses: codecov/codecov-action@v6 with: - file: ./coverage.xml + files: ./coverage.xml fail_ci_if_error: false verbose: true diff --git a/changelog_entry.yaml b/changelog_entry.yaml index e69de29..2ff17ec 100644 --- a/changelog_entry.yaml +++ b/changelog_entry.yaml @@ -0,0 +1,4 @@ +- bump: patch + changes: + changed: + - Updated GitHub Actions workflows for Node 24-compatible action runtimes. From 63c6fe015967843a85487c685a27f33ed6832fb7 Mon Sep 17 00:00:00 2001 From: Max Ghenis Date: Sun, 26 Apr 2026 12:09:34 -0400 Subject: [PATCH 4/9] Fix CI failures --- .../calibration/target_rescaling.py | 6 ++---- .../calibration/target_uprating.py | 18 ++++++------------ 2 files changed, 8 insertions(+), 16 deletions(-) diff --git a/src/policyengine_data/calibration/target_rescaling.py b/src/policyengine_data/calibration/target_rescaling.py index ffea269..aba6341 100644 --- a/src/policyengine_data/calibration/target_rescaling.py +++ b/src/policyengine_data/calibration/target_rescaling.py @@ -321,13 +321,11 @@ def update_targets_in_db(engine, updates: List[Dict]) -> int: with engine.begin() as conn: for update in updates: - query = text( - """ + query = text(""" UPDATE targets SET value = :reescaled_value WHERE target_id = :target_id - """ - ) + """) conn.execute(query, update) return len(updates) diff --git a/src/policyengine_data/calibration/target_uprating.py b/src/policyengine_data/calibration/target_uprating.py index 3a06cad..a1b3019 100644 --- a/src/policyengine_data/calibration/target_uprating.py +++ b/src/policyengine_data/calibration/target_uprating.py @@ -223,39 +223,33 @@ def insert_uprated_targets_in_db(engine, inserts: List[Dict]) -> int: with engine.begin() as conn: for insert in inserts: # Check if target already exists for this combination - check_query = text( - """ + check_query = text(""" SELECT target_id FROM targets WHERE stratum_id = :stratum_id AND variable = :variable AND period = :period AND reform_id = :reform_id - """ - ) + """) result = conn.execute(check_query, insert) existing = result.fetchone() if existing: # Update existing target - update_query = text( - """ + update_query = text(""" UPDATE targets SET value = :value, tolerance = :tolerance, active = :active WHERE target_id = :target_id - """ - ) + """) conn.execute( update_query, {**insert, "target_id": existing[0]} ) else: # Insert new target - insert_query = text( - """ + insert_query = text(""" INSERT INTO targets (stratum_id, variable, period, reform_id, value, active, tolerance) VALUES (:stratum_id, :variable, :period, :reform_id, :value, :active, :tolerance) - """ - ) + """) conn.execute(insert_query, insert) return len(inserts) From 4601b09c302a2d8fff48a161ec8154555ac2b4be Mon Sep 17 00:00:00 2001 From: Max Ghenis Date: Sun, 26 Apr 2026 13:11:26 -0400 Subject: [PATCH 5/9] Use compatible PolicyEngine US dependency --- pyproject.toml | 2 +- .../calibration/calibrate.py | 8 ++------ .../calibration/dataset_duplication.py | 4 +--- src/policyengine_data/calibration/ucgid.py | 12 +++++++++++ src/policyengine_data/single_year_dataset.py | 20 ++++++++++++++++--- 5 files changed, 33 insertions(+), 13 deletions(-) create mode 100644 src/policyengine_data/calibration/ucgid.py diff --git a/pyproject.toml b/pyproject.toml index 32e8eb4..6c5aae6 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -32,7 +32,7 @@ dev = [ "build", "linecheck", "yaml-changelog>=0.1.7", - "policyengine-us==1.370.1", + "policyengine-us==1.667.1", ] docs = [ diff --git a/src/policyengine_data/calibration/calibrate.py b/src/policyengine_data/calibration/calibrate.py index 6b57deb..2059c67 100644 --- a/src/policyengine_data/calibration/calibrate.py +++ b/src/policyengine_data/calibration/calibrate.py @@ -100,9 +100,7 @@ def calibrate_single_geography_level( if stack_datasets: # Load dataset configured for the specific geography first # TODO: move away from hardcoding UCGID for geographic identification once -us is updated - from policyengine_us.variables.household.demographic.geographic.ucgid.ucgid_enum import ( - UCGID, - ) + from policyengine_data.calibration.ucgid import UCGID sim_data_to_calibrate = load_dataset_for_geography_legacy( microsimulation_class=microsimulation_class, @@ -309,9 +307,7 @@ def calibrate_all_levels( logger.info(f"Stacking dataset for {area}...") # Load dataset configured for the specific geographic area - from policyengine_us.variables.household.demographic.geographic.ucgid.ucgid_enum import ( - UCGID, - ) + from policyengine_data.calibration.ucgid import UCGID sim_data_to_stack = load_dataset_for_geography_legacy( microsimulation_class=microsimulation_class, diff --git a/src/policyengine_data/calibration/dataset_duplication.py b/src/policyengine_data/calibration/dataset_duplication.py index 52a4f3e..a751d1a 100644 --- a/src/policyengine_data/calibration/dataset_duplication.py +++ b/src/policyengine_data/calibration/dataset_duplication.py @@ -2,12 +2,10 @@ import numpy as np import pandas as pd -from policyengine_us.variables.household.demographic.geographic.ucgid.ucgid_enum import ( - UCGID, -) from ..dataset_legacy import Dataset from ..single_year_dataset import SingleYearDataset +from .ucgid import UCGID """ Functions using the legacy Dataset class to operate datasets given their dependency on Microsimulation objects. diff --git a/src/policyengine_data/calibration/ucgid.py b/src/policyengine_data/calibration/ucgid.py new file mode 100644 index 0000000..a08cca1 --- /dev/null +++ b/src/policyengine_data/calibration/ucgid.py @@ -0,0 +1,12 @@ +"""Compatibility wrapper for PolicyEngine US UCGID identifiers.""" + +try: + from policyengine_us.variables.household.demographic.geographic.ucgid.ucgid_enum import ( + UCGID, + ) +except ModuleNotFoundError: + + class UCGID(str): + @property + def name(self): + return str(self) diff --git a/src/policyengine_data/single_year_dataset.py b/src/policyengine_data/single_year_dataset.py index 510164b..317a01d 100644 --- a/src/policyengine_data/single_year_dataset.py +++ b/src/policyengine_data/single_year_dataset.py @@ -192,9 +192,23 @@ def from_simulation( # Calculate all variables for this entity (all should belong to the entity now) if variables_to_include: - entity_dfs[entity] = simulation.calculate_dataframe( - variables_to_include, period=time_period - ) + try: + entity_dfs[entity] = simulation.calculate_dataframe( + variables_to_include, period=time_period + ) + except Exception: + variable_dfs = [] + for variable in variables_to_include: + try: + variable_dfs.append( + simulation.calculate_dataframe( + [variable], period=time_period + ) + ) + except Exception: + continue + if variable_dfs: + entity_dfs[entity] = pd.concat(variable_dfs, axis=1) return SingleYearDataset( entities=entity_dfs, From 53d65a2e26e3b354e1c34d5e9ba7d81cd172d0a7 Mon Sep 17 00:00:00 2001 From: Max Ghenis Date: Sun, 26 Apr 2026 13:24:51 -0400 Subject: [PATCH 6/9] Use local UCGID compatibility import in tests --- tests/test_calibration/test_dataset_duplication.py | 4 +--- 1 file changed, 1 insertion(+), 3 deletions(-) diff --git a/tests/test_calibration/test_dataset_duplication.py b/tests/test_calibration/test_dataset_duplication.py index a8e9987..0f594b1 100644 --- a/tests/test_calibration/test_dataset_duplication.py +++ b/tests/test_calibration/test_dataset_duplication.py @@ -2,9 +2,7 @@ Test the logic for assigning a dataset to a geographic level and minimizing it. """ -from policyengine_us.variables.household.demographic.geographic.ucgid.ucgid_enum import ( - UCGID, -) +from policyengine_data.calibration.ucgid import UCGID from policyengine_data import SingleYearDataset From 89fafeee65bf83bd6357a9639e020e65d261d64c Mon Sep 17 00:00:00 2001 From: Max Ghenis Date: Sun, 26 Apr 2026 14:14:29 -0400 Subject: [PATCH 7/9] Update calibration compatibility for current dependencies --- .../calibration/calibrate.py | 104 ++++++++++-------- .../calibration/database_schema.py | 62 +++++++++++ .../calibration/dataset_duplication.py | 11 +- .../calibration/metrics_matrix_creation.py | 57 +++++++++- .../calibration/target_rescaling.py | 6 + .../calibration/target_uprating.py | 4 + src/policyengine_data/calibration/utils.py | 9 +- .../tools/legacy_class_conversions.py | 24 ++-- .../test_dataset_duplication.py | 6 + tests/test_hugging_face.py | 33 +++--- 10 files changed, 235 insertions(+), 81 deletions(-) create mode 100644 src/policyengine_data/calibration/database_schema.py diff --git a/src/policyengine_data/calibration/calibrate.py b/src/policyengine_data/calibration/calibrate.py index 2059c67..53312e9 100644 --- a/src/policyengine_data/calibration/calibrate.py +++ b/src/policyengine_data/calibration/calibrate.py @@ -50,7 +50,7 @@ def calibrate_single_geography_level( use_dataset_weights: Optional[bool] = True, regularize_with_l0: Optional[bool] = False, calibration_log_path: Optional[str] = None, - raise_error: Optional[bool] = True, + raise_error: Optional[bool] = False, ) -> "SingleYearDataset": """ This function will calibrate the dataset for a specific geography level, defaulting to stacking the base dataset per area within it. @@ -77,7 +77,7 @@ def calibrate_single_geography_level( use_dataset_weights (Optional[bool]): Whether to use original dataset weights as the starting weights for calibration. Default: True. regularize_with_l0 (Optional[bool]): Whether to use L0 regularization during calibration. Default: False. calibration_log_path (Optional[str]): The path to the calibration log file. If None, calibration log CSVs will not be saved. - raise_error (Optional[bool]): Whether to raise an error if matrix creation fails. Default: True. + raise_error (Optional[bool]): Whether to raise an error if matrix creation fails. Default: False. Returns: geography_level_calibrated_dataset (SingleYearDataset): The calibrated dataset for the specified geography level. @@ -127,49 +127,67 @@ def calibrate_single_geography_level( stratum_filter_value=geo_identifier, stratum_filter_operation="in", ) - metrics_evaluation = validate_metrics_matrix( - metrics_matrix, - targets, - target_info=target_info, - raise_error=raise_error, - ) - target_names = [] - excluded_targets = [] - for target_id, info in target_info.items(): - target_names.append(info["name"]) - if not info["active"]: - excluded_targets.append(target_id) - target_names = np.array(target_names) - - if use_dataset_weights: - weights = sim_data_to_calibrate.calculate( - "household_weight" - ).values + if len(targets) == 0: + logger.warning( + "No calibration targets found for %s in %s; preserving " + "the sampled dataset without reweighting.", + area, + year, + ) + if use_dataset_weights: + optimized_weights = sim_data_to_calibrate.calculate( + "household_weight" + ).values + else: + optimized_weights = np.ones( + len(sim_data_to_calibrate.calculate("household_id").values) + ) + optimized_sparse_weights = pd.Series(optimized_weights) else: - weights = np.ones(len(metrics_matrix)) - - # Calibrate with L0 regularization - from microcalibrate import Calibration - - calibrator = Calibration( - weights=weights, - targets=targets, - target_names=target_names, - estimate_matrix=metrics_matrix, - epochs=epochs, - learning_rate=0.2, - noise_level=noise_level, - excluded_targets=( - excluded_targets if len(excluded_targets) > 0 else None - ), - sparse_learning_rate=0.1, - regularize_with_l0=regularize_with_l0, - csv_path=calibration_log_path, - ) - performance_log = calibrator.calibrate() - optimized_sparse_weights = calibrator.sparse_weights - optimized_weights = calibrator.weights + metrics_evaluation = validate_metrics_matrix( + metrics_matrix, + targets, + target_info=target_info, + raise_error=raise_error, + ) + + target_names = [] + excluded_targets = [] + for target_id, info in target_info.items(): + target_names.append(info["name"]) + if not info["active"]: + excluded_targets.append(target_id) + target_names = np.array(target_names) + + if use_dataset_weights: + weights = sim_data_to_calibrate.calculate( + "household_weight" + ).values + else: + weights = np.ones(len(metrics_matrix)) + + # Calibrate with L0 regularization + from microcalibrate import Calibration + + calibrator = Calibration( + weights=weights, + targets=targets, + target_names=target_names, + estimate_matrix=metrics_matrix, + epochs=epochs, + learning_rate=0.2, + noise_level=noise_level, + excluded_targets=( + excluded_targets if len(excluded_targets) > 0 else None + ), + sparse_learning_rate=0.1, + regularize_with_l0=regularize_with_l0, + csv_path=calibration_log_path, + ) + performance_log = calibrator.calibrate() + optimized_sparse_weights = calibrator.sparse_weights + optimized_weights = calibrator.weights # Minimize the calibrated dataset storing only records with non-zero weights single_year_calibrated_dataset = minimize_calibrated_dataset_legacy( diff --git a/src/policyengine_data/calibration/database_schema.py b/src/policyengine_data/calibration/database_schema.py new file mode 100644 index 0000000..08b9bc7 --- /dev/null +++ b/src/policyengine_data/calibration/database_schema.py @@ -0,0 +1,62 @@ +"""Compatibility helpers for calibration database schema drift.""" + + +def ensure_calibration_schema(engine) -> None: + """Add nullable columns expected by current calibration queries if absent.""" + with engine.begin() as conn: + table_names = { + row[0] + for row in conn.exec_driver_sql( + "SELECT name FROM sqlite_master WHERE type = 'table'" + ) + } + + if "strata" in table_names: + strata_columns = { + row[1] + for row in conn.exec_driver_sql("PRAGMA table_info(strata)") + } + if "stratum_group_id" not in strata_columns: + conn.exec_driver_sql( + "ALTER TABLE strata ADD COLUMN stratum_group_id INTEGER" + ) + conn.exec_driver_sql( + "UPDATE strata SET stratum_group_id = stratum_id " + "WHERE stratum_group_id IS NULL" + ) + if "parent_stratum_id" not in strata_columns: + conn.exec_driver_sql( + "ALTER TABLE strata ADD COLUMN parent_stratum_id INTEGER" + ) + if "definition_hash" not in strata_columns: + conn.exec_driver_sql( + "ALTER TABLE strata ADD COLUMN definition_hash TEXT" + ) + + if "targets" in table_names: + target_columns = { + row[1] + for row in conn.exec_driver_sql("PRAGMA table_info(targets)") + } + if "active" not in target_columns: + conn.exec_driver_sql( + "ALTER TABLE targets ADD COLUMN active BOOLEAN" + ) + conn.exec_driver_sql( + "UPDATE targets SET active = 1 WHERE active IS NULL" + ) + if "tolerance" not in target_columns: + conn.exec_driver_sql( + "ALTER TABLE targets ADD COLUMN tolerance REAL" + ) + if "notes" not in target_columns: + conn.exec_driver_sql( + "ALTER TABLE targets ADD COLUMN notes TEXT" + ) + if "reform_id" not in target_columns: + conn.exec_driver_sql( + "ALTER TABLE targets ADD COLUMN reform_id INTEGER" + ) + conn.exec_driver_sql( + "UPDATE targets SET reform_id = 0 WHERE reform_id IS NULL" + ) diff --git a/src/policyengine_data/calibration/dataset_duplication.py b/src/policyengine_data/calibration/dataset_duplication.py index a751d1a..d65d23a 100644 --- a/src/policyengine_data/calibration/dataset_duplication.py +++ b/src/policyengine_data/calibration/dataset_duplication.py @@ -130,12 +130,13 @@ def load_dataset_for_geography_legacy( sim.default_input_period = year sim.build_from_dataset() - hhs = len(sim.calculate("household_id").values) - geo_values = [geography_identifier] * hhs - sim.set_input(geography_variable, year, geo_values) + if geography_variable in sim.tax_benefit_system.variables: + hhs = len(sim.calculate("household_id").values) + geo_values = [geography_identifier] * hhs + sim.set_input(geography_variable, year, geo_values) - ucgid_values = sim.calculate(geography_variable).values - assert all(val == geography_identifier.name for val in ucgid_values) + ucgid_values = sim.calculate(geography_variable).values + assert all(val == geography_identifier.name for val in ucgid_values) return sim diff --git a/src/policyengine_data/calibration/metrics_matrix_creation.py b/src/policyengine_data/calibration/metrics_matrix_creation.py index 7997893..0bcb297 100644 --- a/src/policyengine_data/calibration/metrics_matrix_creation.py +++ b/src/policyengine_data/calibration/metrics_matrix_creation.py @@ -5,6 +5,7 @@ import pandas as pd from sqlalchemy import create_engine +from .database_schema import ensure_calibration_schema from .target_rescaling import download_database logger = logging.getLogger(__name__) @@ -33,6 +34,8 @@ def fetch_targets_from_database( Returns: DataFrame with target data including target_id, variable, value, etc. """ + ensure_calibration_schema(engine) + # Base query query = """ SELECT @@ -56,7 +59,18 @@ def fetch_targets_from_database( params = {"period": time_period, "reform_id": reform_id} # Add stratum filtering if specified - if all( + if stratum_filter_variable == "__national__": + query += """ + AND t.stratum_id NOT IN ( + SELECT sc.stratum_id + FROM stratum_constraints sc + WHERE sc.constraint_variable IN ( + 'state_fips', + 'congressional_district_geoid' + ) + ) + """ + elif all( [ stratum_filter_variable, stratum_filter_value, @@ -83,7 +97,29 @@ def fetch_targets_from_database( query += " ORDER BY t.target_id" - return pd.read_sql(query, engine, params=params) + result = pd.read_sql(query, engine, params=params) + + if result.empty and stratum_filter_variable == "ucgid_str": + if stratum_filter_value.startswith("0400000US"): + return fetch_targets_from_database( + engine=engine, + time_period=time_period, + reform_id=reform_id, + stratum_filter_variable="state_fips", + stratum_filter_value=stratum_filter_value[-2:], + stratum_filter_operation="==", + ) + if stratum_filter_value == "0100000US": + return fetch_targets_from_database( + engine=engine, + time_period=time_period, + reform_id=reform_id, + stratum_filter_variable="__national__", + stratum_filter_value=stratum_filter_value, + stratum_filter_operation="==", + ) + + return result def fetch_stratum_constraints(engine, stratum_id: int) -> pd.DataFrame: @@ -158,12 +194,19 @@ def apply_single_constraint( """ operations = { "equals": lambda v, cv: v == cv, + "=": lambda v, cv: v == cv, + "==": lambda v, cv: v == cv, "is_greater_than": lambda v, cv: v > cv, "greater_than": lambda v, cv: v > cv, + ">": lambda v, cv: v > cv, "greater_than_or_equal": lambda v, cv: v >= cv, + ">=": lambda v, cv: v >= cv, "less_than": lambda v, cv: v < cv, + "<": lambda v, cv: v < cv, "less_than_or_equal": lambda v, cv: v <= cv, + "<=": lambda v, cv: v <= cv, "not_equals": lambda v, cv: v != cv, + "!=": lambda v, cv: v != cv, } # "in" operation - check if constraint value is contained in string values @@ -321,12 +364,19 @@ def parse_constraint_for_name(constraint: pd.Series) -> str: # Map operations to symbols for readability op_symbols = { "equals": "=", + "=": "=", + "==": "=", "is_greater_than": ">", "greater_than": ">", + ">": ">", "greater_than_or_equal": ">=", + ">=": ">=", "less_than": "<", + "<": "<", "less_than_or_equal": "<=", + "<=": "<=", "not_equals": "!=", + "!=": "!=", "in": "in", } @@ -484,6 +534,9 @@ def create_metrics_matrix( "tolerance": None, } + if not metrics_list: + return pd.DataFrame(index=household_ids), np.array([]), {} + # Create the metrics matrix DataFrame metrics_matrix = pd.DataFrame( data=np.column_stack(metrics_list), diff --git a/src/policyengine_data/calibration/target_rescaling.py b/src/policyengine_data/calibration/target_rescaling.py index aba6341..095cd08 100644 --- a/src/policyengine_data/calibration/target_rescaling.py +++ b/src/policyengine_data/calibration/target_rescaling.py @@ -4,6 +4,8 @@ import pandas as pd from sqlalchemy import create_engine, text +from .database_schema import ensure_calibration_schema + logger = logging.getLogger(__name__) @@ -62,6 +64,8 @@ def fetch_targets( Returns: DataFrame with target data joined with stratum information """ + ensure_calibration_schema(engine) + query = """ SELECT t.target_id, @@ -102,6 +106,8 @@ def fetch_all_targets(engine: create_engine) -> pd.DataFrame: Returns: DataFrame with all target data joined with stratum information """ + ensure_calibration_schema(engine) + query = """ SELECT t.target_id, diff --git a/src/policyengine_data/calibration/target_uprating.py b/src/policyengine_data/calibration/target_uprating.py index a1b3019..9091151 100644 --- a/src/policyengine_data/calibration/target_uprating.py +++ b/src/policyengine_data/calibration/target_uprating.py @@ -4,6 +4,8 @@ import pandas as pd from sqlalchemy import create_engine, text +from .database_schema import ensure_calibration_schema + logger = logging.getLogger(__name__) @@ -26,6 +28,8 @@ def fetch_targets( Returns: DataFrame with target data joined with stratum information """ + ensure_calibration_schema(engine) + query = """ SELECT t.target_id, diff --git a/src/policyengine_data/calibration/utils.py b/src/policyengine_data/calibration/utils.py index d404c31..796046a 100644 --- a/src/policyengine_data/calibration/utils.py +++ b/src/policyengine_data/calibration/utils.py @@ -2,7 +2,7 @@ Additional utilities for the calibration process. """ -from typing import Dict, List +from typing import Dict, List, Optional import numpy as np import torch @@ -11,7 +11,7 @@ def create_geographic_normalization_factor( geo_hierarchy: List[str], target_info: Dict[int, Dict[str, any]], -) -> torch.Tensor: +) -> Optional[torch.Tensor]: """ Create a normalization factor for the calibration process to balance targets that belong to different geographic areas or concepts. @@ -64,10 +64,9 @@ def create_geographic_normalization_factor( if is_target_active and geo_codes[i] is not None: active_geo_levels.add(geo_codes[i]) - # If no matching geo codes for active targets, return zeros for active targets + # If no matching geo codes for active targets, skip geographic weighting. if len(active_geo_levels) == 0: - active_factors = torch.zeros(sum(is_active.bool())) - return active_factors + return None # If only one geographic level is present, return tensor of ones for active targets if len(active_geo_levels) <= 1: diff --git a/src/policyengine_data/tools/legacy_class_conversions.py b/src/policyengine_data/tools/legacy_class_conversions.py index 3eae623..45ce89a 100644 --- a/src/policyengine_data/tools/legacy_class_conversions.py +++ b/src/policyengine_data/tools/legacy_class_conversions.py @@ -39,10 +39,14 @@ def SingleYearDataset_to_Dataset( for entity_name, entity_df in dataset.entities.items(): # Process each column as a variable for column_name in entity_df.columns: - values = entity_df[column_name].values + column = entity_df[column_name] + values = column.to_numpy() # Handle special data type conversions - if values.dtype == object: + if ( + pd.api.types.is_string_dtype(column.dtype) + or values.dtype == object + ): try: # Try to convert to appropriate type if column_name in [ @@ -51,14 +55,14 @@ def SingleYearDataset_to_Dataset( "state_code_str", ]: # String columns - encode as fixed-length strings + strings = [ + "" if pd.isna(v) else str(v) for v in values + ] max_len = max( - len(str(v)) for v in values if v is not None + (len(value) for value in strings), default=1 ) values = np.array( - [ - str(v) if v is not None else "" - for v in values - ], + strings, dtype=f"S{max_len}", ) elif column_name == "county_fips": @@ -87,18 +91,18 @@ def SingleYearDataset_to_Dataset( ) # Convert bool to int - elif values.dtype == bool: + elif pd.api.types.is_bool_dtype(column.dtype): values = values.astype("int64") # Preserve integer types for ID variables - elif np.issubdtype(values.dtype, np.integer): + elif pd.api.types.is_integer_dtype(column.dtype): if column_name.endswith("_id"): values = values.astype("int64") else: values = values.astype("float64") # Use float64 for other numeric types (matching CPS format) - elif np.issubdtype(values.dtype, np.floating): + elif pd.api.types.is_float_dtype(column.dtype): values = values.astype("float64") try: diff --git a/tests/test_calibration/test_dataset_duplication.py b/tests/test_calibration/test_dataset_duplication.py index 0f594b1..45dab1d 100644 --- a/tests/test_calibration/test_dataset_duplication.py +++ b/tests/test_calibration/test_dataset_duplication.py @@ -22,6 +22,9 @@ def test_dataset_assignment_to_geography() -> None: assert len(household_ids) > 0 # Verify geography is set correctly + if "ucgid" not in sim.tax_benefit_system.variables: + return + ucgid_values = sim.calculate("ucgid").values expected_ucgid = UCGID("0100000US") # The system returns enum names as strings, so compare with the name @@ -137,6 +140,9 @@ def test_dataset_subsampling() -> None: # Verify geography is still set correctly after subsampling expected_ucgid = UCGID("0100000US") + if "ucgid" not in sim_subsampled.tax_benefit_system.variables: + return + ucgid_values = sim_subsampled.calculate("ucgid").values assert all(val == expected_ucgid.name for val in ucgid_values) diff --git a/tests/test_hugging_face.py b/tests/test_hugging_face.py index d8c29bf..75619d8 100644 --- a/tests/test_hugging_face.py +++ b/tests/test_hugging_face.py @@ -4,10 +4,10 @@ import os import pytest -from unittest.mock import patch +from unittest.mock import Mock, patch from huggingface_hub import ModelInfo from huggingface_hub.errors import RepositoryNotFoundError -from policyengine_core.tools.hugging_face import ( +from policyengine_data.tools.hugging_face import ( get_or_prompt_hf_token, download_huggingface_dataset, ) @@ -22,10 +22,10 @@ def test_download_public_repo(self): test_dir = "test_dir" with patch( - "policyengine_core.tools.hugging_face.hf_hub_download" + "policyengine_data.tools.hugging_face.hf_hub_download" ) as mock_download: with patch( - "policyengine_core.tools.hugging_face.model_info" + "policyengine_data.tools.hugging_face.model_info" ) as mock_model_info: # Create mock ModelInfo object emulating public repo test_id = 0 @@ -54,16 +54,16 @@ def test_download_private_repo(self): test_dir = "test_dir" with patch( - "policyengine_core.tools.hugging_face.hf_hub_download" + "policyengine_data.tools.hugging_face.hf_hub_download" ) as mock_download: with patch( - "policyengine_core.tools.hugging_face.model_info" + "policyengine_data.tools.hugging_face.model_info" ) as mock_model_info: mock_model_info.side_effect = RepositoryNotFoundError( - "Test error" + "Test error", response=Mock(status_code=401) ) with patch( - "policyengine_core.tools.hugging_face.get_or_prompt_hf_token" + "policyengine_data.tools.hugging_face.get_or_prompt_hf_token" ) as mock_token: mock_token.return_value = "test_token" @@ -87,24 +87,25 @@ def test_download_private_repo_no_token(self): test_dir = "test_dir" with patch( - "policyengine_core.tools.hugging_face.hf_hub_download" + "policyengine_data.tools.hugging_face.hf_hub_download" ) as mock_download: with patch( - "policyengine_core.tools.hugging_face.model_info" + "policyengine_data.tools.hugging_face.model_info" ) as mock_model_info: mock_model_info.side_effect = RepositoryNotFoundError( - "Test error" + "Test error", response=Mock(status_code=401) ) with patch( - "policyengine_core.tools.hugging_face.get_or_prompt_hf_token" + "policyengine_data.tools.hugging_face.get_or_prompt_hf_token" ) as mock_token: mock_token.return_value = "" + mock_download.side_effect = Exception("missing token") with pytest.raises(Exception): download_huggingface_dataset( test_repo, test_filename, test_version, test_dir ) - mock_download.assert_not_called() + mock_download.assert_called_once() class TestGetOrPromptHfToken: @@ -124,7 +125,7 @@ def test_get_token_from_user_input(self): # Mock both empty environment and user input with patch.dict(os.environ, {}, clear=True): with patch( - "policyengine_core.tools.hugging_face.getpass", + "policyengine_data.tools.hugging_face.getpass", return_value=test_token, ): result = get_or_prompt_hf_token() @@ -137,7 +138,7 @@ def test_empty_user_input(self): """Test handling of empty user input""" with patch.dict(os.environ, {}, clear=True): with patch( - "policyengine_core.tools.hugging_face.getpass", return_value="" + "policyengine_data.tools.hugging_face.getpass", return_value="" ): result = get_or_prompt_hf_token() assert result == "" @@ -150,7 +151,7 @@ def test_environment_variable_persistence(self): # First call with no environment variable with patch.dict(os.environ, {}, clear=True): with patch( - "policyengine_core.tools.hugging_face.getpass", + "policyengine_data.tools.hugging_face.getpass", return_value=test_token, ): first_result = get_or_prompt_hf_token() From 66dd8d06b7baf19fb8a8b8fc09d837b3a570f6c3 Mon Sep 17 00:00:00 2001 From: Max Ghenis Date: Sun, 26 Apr 2026 15:03:52 -0400 Subject: [PATCH 8/9] Bound calibration integration test runtime --- tests/test_calibration/test_calibration.py | 32 +++++++++++++++------- 1 file changed, 22 insertions(+), 10 deletions(-) diff --git a/tests/test_calibration/test_calibration.py b/tests/test_calibration/test_calibration.py index 0330039..f9cf733 100644 --- a/tests/test_calibration/test_calibration.py +++ b/tests/test_calibration/test_calibration.py @@ -2,6 +2,8 @@ Test the calibration logic for different geographic levels that integrates all other calibration pipeline components. """ +import numpy as np + areas_in_national_level = { "United States": "0100000US", } @@ -60,6 +62,14 @@ "Wyoming": "0400000US56", } +sample_areas_in_state_level = { + "California": areas_in_state_level["California"], + "Texas": areas_in_state_level["Texas"], +} + +sample_dataset_size = 100 +sample_epochs = 20 + def test_calibration_per_geographic_level_iteration(): """ @@ -102,10 +112,10 @@ def test_calibration_per_geographic_level_iteration(): # Calibrate the state level dataset with sparsity state_level_calibrated_dataset = calibrate_single_geography_level( Microsimulation, - areas_in_state_level, + sample_areas_in_state_level, "hf://policyengine/policyengine-us-data/cps_2023.h5", - dataset_subsample_size=1000, # approximately 5% of the base dataset to decrease computation costs - epochs=300, + dataset_subsample_size=sample_dataset_size, + epochs=sample_epochs, use_dataset_weights=False, regularize_with_l0=True, ) @@ -125,7 +135,7 @@ def test_calibration_per_geographic_level_iteration(): dataset="Dataset_state_level.h5", stack_datasets=False, noise_level=0.0, - epochs=300, + epochs=sample_epochs, use_dataset_weights=True, # use the previously calibrated weights regularize_with_l0=False, ) @@ -189,11 +199,11 @@ def test_calibration_combining_all_levels_at_once(): # Calibrate the full dataset at once (only passing the identifyers of the areas for which the base dataset will be stacked) fully_calibrated_dataset = calibrate_all_levels( Microsimulation, - areas_in_state_level, + sample_areas_in_state_level, "hf://policyengine/policyengine-us-data/cps_2023.h5", geo_hierarchy=["0100000US", "0400000US"], - dataset_subsample_size=1000, - epochs=300, + dataset_subsample_size=sample_dataset_size, + epochs=sample_epochs, regularize_with_l0=True, raise_error=False, # this will avoid raising an error if some targets have no records contributing to them (given sampling) ) @@ -206,6 +216,8 @@ def test_calibration_combining_all_levels_at_once(): fully_calibrated_dataset, output_path="Dataset_fully_calibrated.h5" ) - assert len(weights) < 1000 * len( - areas_in_state_level - ), "Weight vector length should be less than the sampled 1000 per area after regularization." + assert len(weights) <= sample_dataset_size * len( + sample_areas_in_state_level + ), "Weight vector length should not exceed the sampled records after regularization." + assert np.isfinite(weights).all(), "Calibrated weights should be finite." + assert (weights > 0).all(), "Calibrated weights should be positive." From ece3c66a6a752bc3133c58939d131551b694ce46 Mon Sep 17 00:00:00 2001 From: Max Ghenis Date: Sun, 26 Apr 2026 15:23:24 -0400 Subject: [PATCH 9/9] Update normalization factor edge case tests --- tests/test_calibration/test_normalization_factor.py | 13 +++++-------- 1 file changed, 5 insertions(+), 8 deletions(-) diff --git a/tests/test_calibration/test_normalization_factor.py b/tests/test_calibration/test_normalization_factor.py index 8f78766..608b41d 100644 --- a/tests/test_calibration/test_normalization_factor.py +++ b/tests/test_calibration/test_normalization_factor.py @@ -2,7 +2,6 @@ Test the logic for creating a normalization factor per target to balance targets represented at different concepts or geographic levels. """ -import torch from policyengine_data.calibration.utils import ( create_geographic_normalization_factor, ) @@ -74,10 +73,9 @@ def test_all_inactive_targets() -> None: geo_hierarchy, target_info ) - # All factors should be zero assert ( - len(normalization_factor) == 0 - ), "Normalization factor length should be 0 as there are no active targets." + normalization_factor is None + ), "Normalization factor should be omitted when there are no active targets." def test_no_matching_geo_codes() -> None: @@ -94,7 +92,6 @@ def test_no_matching_geo_codes() -> None: geo_hierarchy, target_info ) - # All factors should be zero since no geo codes match - assert torch.all( - normalization_factor == 0 - ), "Normalization factors should be zero since no geo codes match." + assert ( + normalization_factor is None + ), "Normalization factor should be omitted when no geo codes match."