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

policyengine/core/__init__.py

@@ -0,0 +1,22 @@
+from .dataset import Dataset
+from .dataset import map_to_entity as map_to_entity
+from .dataset_version import DatasetVersion as DatasetVersion
+from .dynamic import Dynamic as Dynamic
+from .output import Output as Output
+from .output import OutputCollection as OutputCollection
+from .parameter import Parameter as Parameter
+from .parameter_value import ParameterValue as ParameterValue
+from .policy import Policy as Policy
+from .simulation import Simulation as Simulation
+from .tax_benefit_model import TaxBenefitModel as TaxBenefitModel
+from .tax_benefit_model_version import (
+    TaxBenefitModelVersion as TaxBenefitModelVersion,
+)
+from .variable import Variable as Variable
+
+# Rebuild models to resolve forward references
+Dataset.model_rebuild()
+TaxBenefitModelVersion.model_rebuild()
+Variable.model_rebuild()
+Parameter.model_rebuild()
+ParameterValue.model_rebuild()

policyengine/core/dataset.py

@@ -0,0 +1,260 @@
+from uuid import uuid4
+
+import pandas as pd
+from microdf import MicroDataFrame
+from pydantic import BaseModel, ConfigDict, Field
+
+from .dataset_version import DatasetVersion
+from .tax_benefit_model import TaxBenefitModel
+
+
+class Dataset(BaseModel):
+    """Base class for datasets.
+
+    The data field contains entity-level data as a BaseModel with DataFrame fields.
+
+    Example:
+        class YearData(BaseModel):
+            model_config = ConfigDict(arbitrary_types_allowed=True)
+            person: pd.DataFrame
+            household: pd.DataFrame
+
+        class MyDataset(Dataset):
+            data: YearData | None = None
+    """
+
+    model_config = ConfigDict(arbitrary_types_allowed=True)
+
+    id: str = Field(default_factory=lambda: str(uuid4()))
+    name: str
+    description: str
+    dataset_version: DatasetVersion | None = None
+    filepath: str
+    is_output_dataset: bool = False
+    tax_benefit_model: TaxBenefitModel | None = None
+    year: int
+
+    data: BaseModel | None = None
+
+
+def map_to_entity(
+    entity_data: dict[str, MicroDataFrame],
+    source_entity: str,
+    target_entity: str,
+    person_entity: str = "person",
+    columns: list[str] | None = None,
+) -> MicroDataFrame:
+    """Map data from source entity to target entity using join keys.
+
+    This is a generic entity mapping utility that handles:
+    - Same entity mapping (returns as is)
+    - Person to group entity mapping (aggregates values)
+    - Group to person entity mapping (expands values)
+    - Group to group entity mapping (aggregates through person entity)
+
+    Args:
+        entity_data: Dictionary mapping entity names to their MicroDataFrame data
+        source_entity: The source entity name
+        target_entity: The target entity name
+        person_entity: The name of the person entity (default "person")
+        columns: 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
+    """
+    valid_entities = set(entity_data.keys())
+
+    if source_entity not in valid_entities:
+        raise ValueError(
+            f"Invalid source entity '{source_entity}'. Must be one of {valid_entities}"
+        )
+    if target_entity not in valid_entities:
+        raise ValueError(
+            f"Invalid target entity '{target_entity}'. Must be one of {valid_entities}"
+        )
+
+    # Get source data (convert to plain DataFrame to avoid weighted operations during mapping)
+    source_df = pd.DataFrame(entity_data[source_entity])
+
+    if columns:
+        # Select only requested columns (keep all ID columns for joins)
+        id_cols = {col for col in source_df.columns if col.endswith("_id")}
+        cols_to_keep = list(set(columns) | id_cols)
+        source_df = source_df[cols_to_keep]
+
+    # Determine weight column for target entity
+    target_weight = f"{target_entity}_weight"
+
+    # Same entity - return as is
+    if source_entity == target_entity:
+        return MicroDataFrame(source_df, weights=target_weight)
+
+    # Get target data and key
+    target_df = entity_data[target_entity]
+    target_key = f"{target_entity}_id"
+
+    # Person to group entity: aggregate person-level data to group level
+    if source_entity == person_entity and target_entity != person_entity:
+        # Check for both naming patterns: "entity_id" and "person_entity_id"
+        person_target_key = f"{person_entity}_{target_entity}_id"
+        join_key = (
+            person_target_key
+            if person_target_key in source_df.columns
+            else target_key
+        )
+
+        if join_key in source_df.columns:
+            # Get columns to aggregate (exclude ID and weight columns)
+            id_cols = {col for col in source_df.columns if col.endswith("_id")}
+            weight_cols = {
+                col for col in source_df.columns if col.endswith("_weight")
+            }
+            agg_cols = [
+                c
+                for c in source_df.columns
+                if c not in id_cols and c not in weight_cols
+            ]
+
+            # Group by join key and sum
+            aggregated = source_df.groupby(join_key, as_index=False)[
+                agg_cols
+            ].sum()
+
+            # Rename join key to target key if needed
+            if join_key != target_key:
+                aggregated = aggregated.rename(columns={join_key: target_key})
+
+            # Merge with target, preserving original order
+            target_pd = pd.DataFrame(target_df)[[target_key, target_weight]]
+            target_pd = target_pd.reset_index(drop=False)
+            result = target_pd.merge(aggregated, on=target_key, how="left")
+
+            # Sort back to original order
+            result = (
+                result.sort_values("index")
+                .drop("index", axis=1)
+                .reset_index(drop=True)
+            )
+
+            # Fill NaN with 0 for groups with no members in source entity
+            result[agg_cols] = result[agg_cols].fillna(0)
+
+            return MicroDataFrame(result, weights=target_weight)
+
+    # Group entity to person: expand group-level data to person level
+    if source_entity != person_entity and target_entity == person_entity:
+        source_key = f"{source_entity}_id"
+        # Check for both naming patterns
+        person_source_key = f"{person_entity}_{source_entity}_id"
+
+        target_pd = pd.DataFrame(target_df)
+        join_key = (
+            person_source_key
+            if person_source_key in target_pd.columns
+            else source_key
+        )
+
+        if join_key in target_pd.columns:
+            # Rename source key to match join key if needed
+            if join_key != source_key and source_key in source_df.columns:
+                source_df = source_df.rename(columns={source_key: join_key})
+
+            result = target_pd.merge(source_df, on=join_key, how="left")
+            return MicroDataFrame(result, weights=target_weight)
+
+    # Group to group: go through person table
+    if source_entity != person_entity and target_entity != person_entity:
+        # Get person link table with both entity IDs
+        person_df = pd.DataFrame(entity_data[person_entity])
+        source_key = f"{source_entity}_id"
+
+        # Check for both naming patterns for person-level links
+        person_source_key = f"{person_entity}_{source_entity}_id"
+        person_target_key = f"{person_entity}_{target_entity}_id"
+
+        # Determine which keys exist in person table
+        source_link_key = (
+            person_source_key
+            if person_source_key in person_df.columns
+            else source_key
+        )
+        target_link_key = (
+            person_target_key
+            if person_target_key in person_df.columns
+            else target_key
+        )
+
+        # Link source -> person -> target
+        if (
+            source_link_key in person_df.columns
+            and target_link_key in person_df.columns
+        ):
+            person_link = person_df[
+                [source_link_key, target_link_key]
+            ].drop_duplicates()
+
+            # Rename source key to match link key if needed
+            source_df_copy = source_df.copy()
+            if (
+                source_link_key != source_key
+                and source_key in source_df_copy.columns
+            ):
+                source_df_copy = source_df_copy.rename(
+                    columns={source_key: source_link_key}
+                )
+
+            # Join source data with target key
+            source_with_target = source_df_copy.merge(
+                person_link, on=source_link_key, how="left"
+            )
+
+            # Aggregate to target level
+            id_cols = {
+                col
+                for col in source_with_target.columns
+                if col.endswith("_id")
+            }
+            weight_cols = {
+                col
+                for col in source_with_target.columns
+                if col.endswith("_weight")
+            }
+            agg_cols = [
+                c
+                for c in source_with_target.columns
+                if c not in id_cols and c not in weight_cols
+            ]
+
+            aggregated = source_with_target.groupby(
+                target_link_key, as_index=False
+            )[agg_cols].sum()
+
+            # Rename target link key to target key if needed
+            if target_link_key != target_key:
+                aggregated = aggregated.rename(
+                    columns={target_link_key: target_key}
+                )
+
+            # Merge with target, preserving original order
+            target_pd = pd.DataFrame(target_df)[[target_key, target_weight]]
+            target_pd = target_pd.reset_index(drop=False)
+            result = target_pd.merge(aggregated, on=target_key, how="left")
+
+            # Sort back to original order
+            result = (
+                result.sort_values("index")
+                .drop("index", axis=1)
+                .reset_index(drop=True)
+            )
+
+            # Fill NaN with 0
+            result[agg_cols] = result[agg_cols].fillna(0)
+
+            return MicroDataFrame(result, weights=target_weight)
+
+    raise ValueError(
+        f"Unsupported mapping from {source_entity} to {target_entity}"
+    )

policyengine/core/dataset_version.py

@@ -0,0 +1,16 @@
+from typing import TYPE_CHECKING
+from uuid import uuid4
+
+from pydantic import BaseModel, Field
+
+from .tax_benefit_model import TaxBenefitModel
+
+if TYPE_CHECKING:
+    from .dataset import Dataset
+
+
+class DatasetVersion(BaseModel):
+    id: str = Field(default_factory=lambda: str(uuid4()))
+    dataset: "Dataset"
+    description: str
+    tax_benefit_model: TaxBenefitModel = None

policyengine/core/dynamic.py

@@ -0,0 +1,43 @@
+from collections.abc import Callable
+from datetime import datetime
+from uuid import uuid4
+
+from pydantic import BaseModel, Field
+
+from .parameter_value import ParameterValue
+
+
+class Dynamic(BaseModel):
+    id: str = Field(default_factory=lambda: str(uuid4()))
+    name: str
+    description: str | None = None
+    parameter_values: list[ParameterValue] = []
+    simulation_modifier: Callable | None = None
+    created_at: datetime = Field(default_factory=datetime.now)
+    updated_at: datetime = Field(default_factory=datetime.now)
+
+    def __add__(self, other: "Dynamic") -> "Dynamic":
+        """Combine two dynamics by appending parameter values and chaining simulation modifiers."""
+        if not isinstance(other, Dynamic):
+            return NotImplemented
+
+        # Combine simulation modifiers
+        combined_modifier = None
+        if self.simulation_modifier is not None and other.simulation_modifier is not None:
+
+            def combined_modifier(sim):
+                sim = self.simulation_modifier(sim)
+                sim = other.simulation_modifier(sim)
+                return sim
+
+        elif self.simulation_modifier is not None:
+            combined_modifier = self.simulation_modifier
+        elif other.simulation_modifier is not None:
+            combined_modifier = other.simulation_modifier
+
+        return Dynamic(
+            name=f"{self.name} + {other.name}",
+            description=f"Combined dynamic: {self.name} and {other.name}",
+            parameter_values=self.parameter_values + other.parameter_values,
+            simulation_modifier=combined_modifier,
+        )

policyengine/core/output.py

@@ -0,0 +1,26 @@
+from typing import TypeVar
+
+import pandas as pd
+from pydantic import BaseModel, ConfigDict
+
+T = TypeVar("T", bound="Output")
+
+
+class Output(BaseModel):
+    """Base class for all output templates."""
+
+    def run(self):
+        """Calculate and populate the output fields.
+
+        Must be implemented by subclasses.
+        """
+        raise NotImplementedError("Subclasses must implement run()")
+
+
+class OutputCollection[T: "Output"](BaseModel):
+    """Container for a collection of outputs with their DataFrame representation."""
+
+    model_config = ConfigDict(arbitrary_types_allowed=True)
+
+    outputs: list[T]
+    dataframe: pd.DataFrame

policyengine/{models → core}/parameter.py

@@ -2,11 +2,13 @@ from uuid import uuid4
 
 from pydantic import BaseModel, Field
 
-from .model import Model
+from .tax_benefit_model_version import TaxBenefitModelVersion
 
 
 class Parameter(BaseModel):
     id: str = Field(default_factory=lambda: str(uuid4()))
+    name: str
     description: str | None = None
     data_type: type | None = None
-    model: Model | None = None
+    tax_benefit_model_version: TaxBenefitModelVersion
+    unit: str | None = None

policyengine/{models → core}/parameter_value.py

@@ -8,7 +8,7 @@ from .parameter import Parameter
 
 class ParameterValue(BaseModel):
     id: str = Field(default_factory=lambda: str(uuid4()))
-    parameter: Parameter
+    parameter: Parameter | None = None
     value: float | int | str | bool | list | None = None
     start_date: datetime
     end_date: datetime | None = None

policyengine/core/policy.py

@@ -0,0 +1,43 @@
+from collections.abc import Callable
+from datetime import datetime
+from uuid import uuid4
+
+from pydantic import BaseModel, Field
+
+from .parameter_value import ParameterValue
+
+
+class Policy(BaseModel):
+    id: str = Field(default_factory=lambda: str(uuid4()))
+    name: str
+    description: str | None = None
+    parameter_values: list[ParameterValue] = []
+    simulation_modifier: Callable | None = None
+    created_at: datetime = Field(default_factory=datetime.now)
+    updated_at: datetime = Field(default_factory=datetime.now)
+
+    def __add__(self, other: "Policy") -> "Policy":
+        """Combine two policies by appending parameter values and chaining simulation modifiers."""
+        if not isinstance(other, Policy):
+            return NotImplemented
+
+        # Combine simulation modifiers
+        combined_modifier = None
+        if self.simulation_modifier is not None and other.simulation_modifier is not None:
+
+            def combined_modifier(sim):
+                sim = self.simulation_modifier(sim)
+                sim = other.simulation_modifier(sim)
+                return sim
+
+        elif self.simulation_modifier is not None:
+            combined_modifier = self.simulation_modifier
+        elif other.simulation_modifier is not None:
+            combined_modifier = other.simulation_modifier
+
+        return Policy(
+            name=f"{self.name} + {other.name}",
+            description=f"Combined policy: {self.name} and {other.name}",
+            parameter_values=self.parameter_values + other.parameter_values,
+            simulation_modifier=combined_modifier,
+        )

policyengine/{models → core}/simulation.py

@@ -1,14 +1,12 @@
 from datetime import datetime
-from typing import Any
 from uuid import uuid4
 
 from pydantic import BaseModel, Field
 
 from .dataset import Dataset
 from .dynamic import Dynamic
-from .model import Model
-from .model_version import ModelVersion
 from .policy import Policy
+from .tax_benefit_model_version import TaxBenefitModelVersion
 
 
 class Simulation(BaseModel):
@@ -18,17 +16,15 @@ class Simulation(BaseModel):
 
     policy: Policy | None = None
     dynamic: Dynamic | None = None
-    dataset: Dataset
+    dataset: Dataset = None
 
-    model: Model
-    model_version: ModelVersion
-    result: Any | None = None
+    tax_benefit_model_version: TaxBenefitModelVersion = None
+    output_dataset: Dataset | None = None
+
+    variables: dict[str, list[str]] | None = Field(
+        default=None,
+        description="Optional dictionary mapping entity names to lists of variable names to calculate. If None, uses model defaults.",
+    )
 
     def run(self):
-        self.result = self.model.simulation_function(
-            dataset=self.dataset,
-            policy=self.policy,
-            dynamic=self.dynamic,
-        )
-        self.updated_at = datetime.now()
-        return self.result
+        self.tax_benefit_model_version.run(self)

policyengine/core/tax_benefit_model.py

@@ -0,0 +1,11 @@
+from typing import TYPE_CHECKING
+
+from pydantic import BaseModel
+
+if TYPE_CHECKING:
+    pass
+
+
+class TaxBenefitModel(BaseModel):
+    id: str
+    description: str | None = None

policyengine/core/tax_benefit_model_version.py

@@ -0,0 +1,34 @@
+from datetime import datetime
+from typing import TYPE_CHECKING
+from uuid import uuid4
+
+from pydantic import BaseModel, Field
+
+from .tax_benefit_model import TaxBenefitModel
+
+if TYPE_CHECKING:
+    from .parameter import Parameter
+    from .parameter_value import ParameterValue
+    from .simulation import Simulation
+    from .variable import Variable
+
+
+class TaxBenefitModelVersion(BaseModel):
+    id: str = Field(default_factory=lambda: str(uuid4()))
+    model: TaxBenefitModel
+    version: str
+    description: str | None = None
+    created_at: datetime | None = Field(default_factory=datetime.utcnow)
+
+    variables: list["Variable"] = Field(default_factory=list)
+    parameters: list["Parameter"] = Field(default_factory=list)
+    parameter_values: list["ParameterValue"] = Field(default_factory=list)
+
+    def run(self, simulation: "Simulation") -> "Simulation":
+        raise NotImplementedError(
+            "The TaxBenefitModel class must define a method to execute simulations."
+        )
+
+    def __repr__(self) -> str:
+        # Give the id and version, and the number of variables, parameters, parameter values
+        return f"<TaxBenefitModelVersion id={self.id} variables={len(self.variables)} parameters={len(self.parameters)} parameter_values={len(self.parameter_values)}>"

policyengine/core/variable.py

@@ -0,0 +1,15 @@
+from typing import Any
+
+from pydantic import BaseModel
+
+from .tax_benefit_model_version import TaxBenefitModelVersion
+
+
+class Variable(BaseModel):
+    id: str
+    name: str
+    tax_benefit_model_version: TaxBenefitModelVersion
+    entity: str
+    description: str | None = None
+    data_type: type = None
+    possible_values: list[Any] | None = None

policyengine/outputs/__init__.py

@@ -0,0 +1,21 @@
+from policyengine.core import Output, OutputCollection
+from policyengine.outputs.aggregate import Aggregate, AggregateType
+from policyengine.outputs.change_aggregate import (
+    ChangeAggregate,
+    ChangeAggregateType,
+)
+from policyengine.outputs.decile_impact import (
+    DecileImpact,
+    calculate_decile_impacts,
+)
+
+__all__ = [
+    "Output",
+    "OutputCollection",
+    "Aggregate",
+    "AggregateType",
+    "ChangeAggregate",
+    "ChangeAggregateType",
+    "DecileImpact",
+    "calculate_decile_impacts",
+]