policyengine 3.0.0__py3-none-any.whl → 3.1.1__py3-none-any.whl

policyengine/tax_benefit_models/us/analysis.py

@@ -0,0 +1,99 @@
+"""General utility functions for US policy reform analysis."""
+
+import pandas as pd
+from pydantic import BaseModel
+
+from policyengine.core import OutputCollection, Simulation
+from policyengine.outputs.decile_impact import (
+    DecileImpact,
+    calculate_decile_impacts,
+)
+
+from .outputs import ProgramStatistics
+
+
+class PolicyReformAnalysis(BaseModel):
+    """Complete policy reform analysis result."""
+
+    decile_impacts: OutputCollection[DecileImpact]
+    program_statistics: OutputCollection[ProgramStatistics]
+
+
+def general_policy_reform_analysis(
+    baseline_simulation: Simulation,
+    reform_simulation: Simulation,
+) -> PolicyReformAnalysis:
+    """Perform comprehensive analysis of a policy reform.
+
+    Returns:
+        PolicyReformAnalysis containing decile impacts and program statistics
+    """
+    # Decile impact (using household_net_income for US)
+    decile_impacts = calculate_decile_impacts(
+        baseline_simulation=baseline_simulation,
+        reform_simulation=reform_simulation,
+        income_variable="household_net_income",
+    )
+
+    # Major programs to analyse
+    programs = {
+        # Federal taxes
+        "income_tax": {"entity": "tax_unit", "is_tax": True},
+        "payroll_tax": {"entity": "person", "is_tax": True},
+        # State and local taxes
+        "state_income_tax": {"entity": "tax_unit", "is_tax": True},
+        # Benefits
+        "snap": {"entity": "spm_unit", "is_tax": False},
+        "tanf": {"entity": "spm_unit", "is_tax": False},
+        "ssi": {"entity": "person", "is_tax": False},
+        "social_security": {"entity": "person", "is_tax": False},
+        "medicare": {"entity": "person", "is_tax": False},
+        "medicaid": {"entity": "person", "is_tax": False},
+        "eitc": {"entity": "tax_unit", "is_tax": False},
+        "ctc": {"entity": "tax_unit", "is_tax": False},
+    }
+
+    program_statistics = []
+
+    for program_name, program_info in programs.items():
+        entity = program_info["entity"]
+        is_tax = program_info["is_tax"]
+
+        stats = ProgramStatistics(
+            baseline_simulation=baseline_simulation,
+            reform_simulation=reform_simulation,
+            program_name=program_name,
+            entity=entity,
+            is_tax=is_tax,
+        )
+        stats.run()
+        program_statistics.append(stats)
+
+    # Create DataFrame
+    program_df = pd.DataFrame(
+        [
+            {
+                "baseline_simulation_id": p.baseline_simulation.id,
+                "reform_simulation_id": p.reform_simulation.id,
+                "program_name": p.program_name,
+                "entity": p.entity,
+                "is_tax": p.is_tax,
+                "baseline_total": p.baseline_total,
+                "reform_total": p.reform_total,
+                "change": p.change,
+                "baseline_count": p.baseline_count,
+                "reform_count": p.reform_count,
+                "winners": p.winners,
+                "losers": p.losers,
+            }
+            for p in program_statistics
+        ]
+    )
+
+    program_collection = OutputCollection(
+        outputs=program_statistics, dataframe=program_df
+    )
+
+    return PolicyReformAnalysis(
+        decile_impacts=decile_impacts, program_statistics=program_collection
+    )

policyengine/tax_benefit_models/us/datasets.py

@@ -0,0 +1,307 @@
+import warnings
+from pathlib import Path
+
+import pandas as pd
+from microdf import MicroDataFrame
+from pydantic import BaseModel, ConfigDict
+
+from policyengine.core import Dataset, map_to_entity
+
+
+class USYearData(BaseModel):
+    """Entity-level data for a single year."""
+
+    model_config = ConfigDict(arbitrary_types_allowed=True)
+
+    person: MicroDataFrame
+    marital_unit: MicroDataFrame
+    family: MicroDataFrame
+    spm_unit: MicroDataFrame
+    tax_unit: MicroDataFrame
+    household: MicroDataFrame
+
+    def map_to_entity(
+        self, source_entity: str, target_entity: str, columns: list[str] = None
+    ) -> MicroDataFrame:
+        """Map data from source entity to target entity using join keys.
+
+        Args:
+            source_entity (str): The source entity name.
+            target_entity (str): The target entity name.
+            columns (list[str], optional): List of column names to map. If None, maps all columns.
+
+        Returns:
+            MicroDataFrame: The mapped data at the target entity level.
+
+        Raises:
+            ValueError: If source or target entity is invalid.
+        """
+        entity_data = {
+            "person": self.person,
+            "marital_unit": self.marital_unit,
+            "family": self.family,
+            "spm_unit": self.spm_unit,
+            "tax_unit": self.tax_unit,
+            "household": self.household,
+        }
+        return map_to_entity(
+            entity_data=entity_data,
+            source_entity=source_entity,
+            target_entity=target_entity,
+            person_entity="person",
+            columns=columns,
+        )
+
+
+class PolicyEngineUSDataset(Dataset):
+    """US dataset with multi-year entity-level data."""
+
+    data: USYearData | None = None
+
+    def model_post_init(self, __context) -> None:
+        """Called after Pydantic initialization."""
+        # Make sure we are synchronised between in-memory and storage, at least on initialisation
+        if self.data is not None:
+            self.save()
+        elif self.filepath and not self.data:
+            try:
+                self.load()
+            except FileNotFoundError:
+                # File doesn't exist yet, that's OK
+                pass
+
+    def save(self) -> None:
+        """Save dataset to HDF5 file."""
+        filepath = Path(self.filepath)
+        if not filepath.parent.exists():
+            filepath.parent.mkdir(parents=True, exist_ok=True)
+        with warnings.catch_warnings():
+            warnings.filterwarnings(
+                "ignore",
+                category=pd.errors.PerformanceWarning,
+                message=".*PyTables will pickle object types.*",
+            )
+            with pd.HDFStore(filepath, mode="w") as store:
+                store["person"] = pd.DataFrame(self.data.person)
+                store["marital_unit"] = pd.DataFrame(self.data.marital_unit)
+                store["family"] = pd.DataFrame(self.data.family)
+                store["spm_unit"] = pd.DataFrame(self.data.spm_unit)
+                store["tax_unit"] = pd.DataFrame(self.data.tax_unit)
+                store["household"] = pd.DataFrame(self.data.household)
+
+    def load(self) -> None:
+        """Load dataset from HDF5 file into this instance."""
+        filepath = self.filepath
+        with pd.HDFStore(filepath, mode="r") as store:
+            self.data = USYearData(
+                person=MicroDataFrame(
+                    store["person"], weights="person_weight"
+                ),
+                marital_unit=MicroDataFrame(
+                    store["marital_unit"], weights="marital_unit_weight"
+                ),
+                family=MicroDataFrame(
+                    store["family"], weights="family_weight"
+                ),
+                spm_unit=MicroDataFrame(
+                    store["spm_unit"], weights="spm_unit_weight"
+                ),
+                tax_unit=MicroDataFrame(
+                    store["tax_unit"], weights="tax_unit_weight"
+                ),
+                household=MicroDataFrame(
+                    store["household"], weights="household_weight"
+                ),
+            )
+
+    def __repr__(self) -> str:
+        if self.data is None:
+            return f"<PolicyEngineUSDataset id={self.id} year={self.year} filepath={self.filepath} (not loaded)>"
+        else:
+            n_people = len(self.data.person)
+            n_marital_units = len(self.data.marital_unit)
+            n_families = len(self.data.family)
+            n_spm_units = len(self.data.spm_unit)
+            n_tax_units = len(self.data.tax_unit)
+            n_households = len(self.data.household)
+            return f"<PolicyEngineUSDataset id={self.id} year={self.year} filepath={self.filepath} people={n_people} marital_units={n_marital_units} families={n_families} spm_units={n_spm_units} tax_units={n_tax_units} households={n_households}>"
+
+
+def create_datasets(
+    datasets: list[str] = [
+        "hf://policyengine/policyengine-us-data/enhanced_cps_2024.h5",
+    ],
+    years: list[int] = [2024, 2025, 2026, 2027, 2028],
+) -> None:
+    """Create PolicyEngineUSDataset instances from HuggingFace dataset paths.
+
+    Args:
+        datasets: List of HuggingFace dataset paths (e.g., "hf://policyengine/policyengine-us-data/cps_2024.h5")
+        years: List of years to extract data for
+    """
+    from policyengine_us import Microsimulation
+
+    for dataset in datasets:
+        sim = Microsimulation(dataset=dataset)
+
+        for year in years:
+            # Get all input variables from the simulation
+            # We'll calculate each input variable for the specified year
+            entity_data = {
+                "person": {},
+                "household": {},
+                "marital_unit": {},
+                "family": {},
+                "spm_unit": {},
+                "tax_unit": {},
+            }
+
+            # First, get ID columns which are structural (not input variables)
+            # These define entity membership and relationships
+            # For person-level links to group entities, use person_X_id naming
+            id_variables = {
+                "person": [
+                    "person_id",
+                    "person_household_id",
+                    "person_marital_unit_id",
+                    "person_family_id",
+                    "person_spm_unit_id",
+                    "person_tax_unit_id",
+                ],
+                "household": ["household_id"],
+                "marital_unit": ["marital_unit_id"],
+                "family": ["family_id"],
+                "spm_unit": ["spm_unit_id"],
+                "tax_unit": ["tax_unit_id"],
+            }
+
+            for entity_key, var_names in id_variables.items():
+                for id_var in var_names:
+                    if id_var in sim.tax_benefit_system.variables:
+                        values = sim.calculate(id_var, period=year).values
+                        entity_data[entity_key][id_var] = values
+
+            # Get input variables and calculate them for this year
+            for variable_name in sim.input_variables:
+                variable = sim.tax_benefit_system.variables[variable_name]
+                entity_key = variable.entity.key
+
+                # Calculate the variable for the given year
+                values = sim.calculate(variable_name, period=year).values
+
+                # Store in the appropriate entity dictionary
+                entity_data[entity_key][variable_name] = values
+
+            # Build entity DataFrames
+            person_df = pd.DataFrame(entity_data["person"])
+            household_df = pd.DataFrame(entity_data["household"])
+            marital_unit_df = pd.DataFrame(entity_data["marital_unit"])
+            family_df = pd.DataFrame(entity_data["family"])
+            spm_unit_df = pd.DataFrame(entity_data["spm_unit"])
+            tax_unit_df = pd.DataFrame(entity_data["tax_unit"])
+
+            # Add weight columns - household weights are primary, map to all entities
+            # Person weights = household weights (mapped via person_household_id)
+            if "household_weight" in household_df.columns:
+                # Only add person_weight if it doesn't already exist
+                if "person_weight" not in person_df.columns:
+                    person_df = person_df.merge(
+                        household_df[["household_id", "household_weight"]],
+                        left_on="person_household_id",
+                        right_on="household_id",
+                        how="left",
+                    )
+                    person_df = person_df.rename(
+                        columns={"household_weight": "person_weight"}
+                    )
+                    person_df = person_df.drop(
+                        columns=["household_id"], errors="ignore"
+                    )
+
+                # Map household weights to other group entities via person table
+                for entity_name, entity_df, person_id_col, entity_id_col in [
+                    (
+                        "marital_unit",
+                        marital_unit_df,
+                        "person_marital_unit_id",
+                        "marital_unit_id",
+                    ),
+                    ("family", family_df, "person_family_id", "family_id"),
+                    (
+                        "spm_unit",
+                        spm_unit_df,
+                        "person_spm_unit_id",
+                        "spm_unit_id",
+                    ),
+                    (
+                        "tax_unit",
+                        tax_unit_df,
+                        "person_tax_unit_id",
+                        "tax_unit_id",
+                    ),
+                ]:
+                    # Only add entity weight if it doesn't already exist
+                    if f"{entity_name}_weight" not in entity_df.columns:
+                        # Get household_id for each entity from person table
+                        entity_household_map = person_df[
+                            [person_id_col, "person_household_id"]
+                        ].drop_duplicates()
+                        entity_df = entity_df.merge(
+                            entity_household_map,
+                            left_on=entity_id_col,
+                            right_on=person_id_col,
+                            how="left",
+                        )
+                        entity_df = entity_df.merge(
+                            household_df[["household_id", "household_weight"]],
+                            left_on="person_household_id",
+                            right_on="household_id",
+                            how="left",
+                        )
+                        entity_df = entity_df.rename(
+                            columns={
+                                "household_weight": f"{entity_name}_weight"
+                            }
+                        )
+                        entity_df = entity_df.drop(
+                            columns=[
+                                "household_id",
+                                "person_household_id",
+                                person_id_col,
+                            ],
+                            errors="ignore",
+                        )
+
+                    # Update the entity_data
+                    if entity_name == "marital_unit":
+                        marital_unit_df = entity_df
+                    elif entity_name == "family":
+                        family_df = entity_df
+                    elif entity_name == "spm_unit":
+                        spm_unit_df = entity_df
+                    elif entity_name == "tax_unit":
+                        tax_unit_df = entity_df
+
+            us_dataset = PolicyEngineUSDataset(
+                name=f"{dataset}-year-{year}",
+                description=f"US Dataset for year {year} based on {dataset}",
+                filepath=f"./data/{Path(dataset).stem}_year_{year}.h5",
+                year=year,
+                data=USYearData(
+                    person=MicroDataFrame(person_df, weights="person_weight"),
+                    household=MicroDataFrame(
+                        household_df, weights="household_weight"
+                    ),
+                    marital_unit=MicroDataFrame(
+                        marital_unit_df, weights="marital_unit_weight"
+                    ),
+                    family=MicroDataFrame(family_df, weights="family_weight"),
+                    spm_unit=MicroDataFrame(
+                        spm_unit_df, weights="spm_unit_weight"
+                    ),
+                    tax_unit=MicroDataFrame(
+                        tax_unit_df, weights="tax_unit_weight"
+                    ),
+                ),
+            )
+            us_dataset.save()