Merge pull request #398 from PolicyEngine/treasury

A SQLite Database for Calibration Targets

Author: baogorek

.github/workflows/reusable_test.yaml

@@ -58,6 +58,10 @@ jobs:
         if: inputs.full_suite
         run: make download
 
+      - name: Create and load calibration targets database 
+        if: inputs.full_suite
+        run: make database
+
       - name: Build datasets
         if: inputs.full_suite
         run: make data
@@ -90,4 +94,4 @@ jobs:
         with:
           branch: gh-pages
           folder: docs/_build/html
-          clean: true
+          clean: true

Makefile

@@ -22,12 +22,6 @@ changelog:
 download:
 	python policyengine_us_data/storage/download_private_prerequisites.py
 
-targets:
-	python policyengine_us_data/storage/calibration_targets/pull_hardcoded_targets.py
-	python policyengine_us_data/storage/calibration_targets/pull_age_targets.py
-	python policyengine_us_data/storage/calibration_targets/pull_soi_targets.py
-	python policyengine_us_data/storage/calibration_targets/pull_snap_targets.py
-
 upload:
 	python policyengine_us_data/storage/upload_completed_datasets.py
 
@@ -61,10 +55,12 @@ documentation-dev:
 
 database:
 	python policyengine_us_data/db/create_database_tables.py
-	python policyengine_us_data/db/load_age_targets.py
-
-clean-database:
-	rm *.db
+	python policyengine_us_data/db/create_initial_strata.py
+	python policyengine_us_data/db/etl_age.py
+	python policyengine_us_data/db/etl_medicaid.py
+	python policyengine_us_data/db/etl_snap.py
+	python policyengine_us_data/db/etl_irs_soi.py
+	python policyengine_us_data/db/validate_database.py
 
 data:
 	python policyengine_us_data/utils/uprating.py
@@ -80,6 +76,7 @@ data:
 
 clean:
 	rm -f policyengine_us_data/storage/*.h5
+	rm -f policyengine_us_data/storage/*.db
 	git clean -fX -- '*.csv'
 	rm -rf policyengine_us_data/docs/_build
 

changelog_entry.yaml

@@ -0,0 +1,4 @@
+- bump: minor
+  changes:
+    added:
+    - add SQLite database for calibration targets 

policyengine_us_data/db/create_database_tables.py

@@ -1,16 +1,20 @@
 import logging
 import hashlib
 from typing import List, Optional
+from enum import Enum
 
 from sqlalchemy import event, UniqueConstraint
 from sqlalchemy.orm.attributes import get_history
-
 from sqlmodel import (
     Field,
     Relationship,
     SQLModel,
     create_engine,
 )
+from policyengine_us.system import system
+
+from policyengine_us_data.storage import STORAGE_FOLDER
+
 
 logging.basicConfig(
     level=logging.INFO,
@@ -20,6 +24,12 @@
 logger = logging.getLogger(__name__)
 
 
+# An Enum type to ensure the variable exists in policyengine-us
+USVariable = Enum(
+    "USVariable", {name: name for name in system.variables.keys()}, type=str
+)
+
+
 class Stratum(SQLModel, table=True):
     """Represents a unique population subgroup (stratum)."""
 
@@ -79,7 +89,7 @@ class StratumConstraint(SQLModel, table=True):
     __tablename__ = "stratum_constraints"
 
     stratum_id: int = Field(foreign_key="strata.stratum_id", primary_key=True)
-    constraint_variable: str = Field(
+    constraint_variable: USVariable = Field(
         primary_key=True,
         description="The variable the constraint applies to (e.g., 'age').",
     )
@@ -112,7 +122,7 @@ class Target(SQLModel, table=True):
     )
 
     target_id: Optional[int] = Field(default=None, primary_key=True)
-    variable: str = Field(
+    variable: USVariable = Field(
         description="A variable defined in policyengine-us (e.g., 'income_tax')."
     )
     period: int = Field(
@@ -171,12 +181,11 @@ def calculate_definition_hash(mapper, connection, target: Stratum):
     fingerprint_text = "\n".join(constraint_strings)
     h = hashlib.sha256(fingerprint_text.encode("utf-8"))
     target.definition_hash = h.hexdigest()
-    logger.info(
-        f"Set definition_hash for Stratum to '{target.definition_hash}'"
-    )
 
 
-def create_database(db_uri="sqlite:///policy_data.db"):
+def create_database(
+    db_uri: str = f"sqlite:///{STORAGE_FOLDER / 'policy_data.db'}",
+):
     """
     Creates a SQLite database and all the defined tables.
 

policyengine_us_data/db/create_initial_strata.py

@@ -0,0 +1,75 @@
+from typing import Dict
+
+import pandas as pd
+from sqlmodel import Session, create_engine
+
+from policyengine_us_data.storage import STORAGE_FOLDER
+
+
+from policyengine_us.variables.household.demographic.geographic.ucgid.ucgid_enum import (
+    UCGID,
+)
+from policyengine_us_data.db.create_database_tables import (
+    Stratum,
+    StratumConstraint,
+)
+
+
+def main():
+    # Get the implied hierarchy by the UCGID enum --------
+    rows = []
+    for node in UCGID:
+        codes = node.get_hierarchical_codes()
+        rows.append(
+            {
+                "name": node.name,
+                "code": codes[0],
+                "parent": codes[1] if len(codes) > 1 else None,
+            }
+        )
+
+    hierarchy_df = (
+        pd.DataFrame(rows)
+        .sort_values(["parent", "code"], na_position="first")
+        .reset_index(drop=True)
+    )
+
+    DATABASE_URL = f"sqlite:///{STORAGE_FOLDER / 'policy_data.db'}"
+    engine = create_engine(DATABASE_URL)
+
+    # map the ucgid_str 'code' to auto-generated 'stratum_id'
+    code_to_stratum_id: Dict[str, int] = {}
+
+    with Session(engine) as session:
+        for _, row in hierarchy_df.iterrows():
+            parent_code = row["parent"]
+
+            parent_id = (
+                code_to_stratum_id.get(parent_code) if parent_code else None
+            )
+
+            new_stratum = Stratum(
+                parent_stratum_id=parent_id,
+                notes=f'{row["name"]} (ucgid {row["code"]})',
+                stratum_group_id=1,
+            )
+
+            new_stratum.constraints_rel = [
+                StratumConstraint(
+                    constraint_variable="ucgid_str",
+                    operation="in",
+                    value=row["code"],
+                )
+            ]
+
+            session.add(new_stratum)
+
+            session.flush()
+
+            code_to_stratum_id[row["code"]] = new_stratum.stratum_id
+
+        session.commit()
+
+
+if __name__ == "__main__":
+    main()