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

policyengine/database/simulation_table.py

@@ -1,50 +0,0 @@
-from datetime import datetime
-from uuid import uuid4
-
-from sqlmodel import Field, SQLModel
-
-from policyengine.models import Simulation
-from policyengine.utils.compress import compress_data, decompress_data
-
-from .link import TableLink
-
-
-class SimulationTable(SQLModel, table=True):
-    __tablename__ = "simulations"
-
-    id: str = Field(default_factory=lambda: str(uuid4()), primary_key=True)
-    created_at: datetime = Field(default_factory=datetime.now)
-    updated_at: datetime = Field(default_factory=datetime.now)
-
-    policy_id: str | None = Field(
-        default=None, foreign_key="policies.id", ondelete="SET NULL"
-    )
-    dynamic_id: str | None = Field(
-        default=None, foreign_key="dynamics.id", ondelete="SET NULL"
-    )
-    dataset_id: str = Field(foreign_key="datasets.id", ondelete="CASCADE")
-    model_id: str = Field(foreign_key="models.id", ondelete="CASCADE")
-    model_version_id: str | None = Field(
-        default=None, foreign_key="model_versions.id", ondelete="SET NULL"
-    )
-
-    result: bytes | None = Field(default=None)
-
-
-simulation_table_link = TableLink(
-    model_cls=Simulation,
-    table_cls=SimulationTable,
-    model_to_table_custom_transforms=dict(
-        policy_id=lambda s: s.policy.id if s.policy else None,
-        dynamic_id=lambda s: s.dynamic.id if s.dynamic else None,
-        dataset_id=lambda s: s.dataset.id,
-        model_id=lambda s: s.model.id,
-        model_version_id=lambda s: s.model_version.id
-        if s.model_version
-        else None,
-        result=lambda s: compress_data(s.result) if s.result else None,
-    ),
-    table_to_model_custom_transforms=dict(
-        result=lambda b: decompress_data(b) if b else None,
-    ),
-)

policyengine/database/user_table.py

@@ -1,28 +0,0 @@
-import uuid
-from datetime import datetime
-
-from sqlmodel import Field, SQLModel
-
-from policyengine.models.user import User
-
-from .link import TableLink
-
-
-class UserTable(SQLModel, table=True, extend_existing=True):
-    __tablename__ = "users"
-
-    id: str = Field(
-        primary_key=True, default_factory=lambda: str(uuid.uuid4())
-    )
-    username: str = Field(nullable=False, unique=True)
-    first_name: str | None = Field(default=None)
-    last_name: str | None = Field(default=None)
-    email: str | None = Field(default=None)
-    created_at: datetime = Field(default_factory=datetime.utcnow)
-    updated_at: datetime = Field(default_factory=datetime.utcnow)
-
-
-user_table_link = TableLink(
-    model_cls=User,
-    table_cls=UserTable,
-)

policyengine/database/versioned_dataset_table.py

@@ -1,28 +0,0 @@
-from uuid import uuid4
-
-from sqlmodel import Field, SQLModel
-
-from policyengine.models import VersionedDataset
-
-from .link import TableLink
-
-
-class VersionedDatasetTable(SQLModel, table=True):
-    __tablename__ = "versioned_datasets"
-
-    id: str = Field(default_factory=lambda: str(uuid4()), primary_key=True)
-    name: str = Field(nullable=False)
-    description: str = Field(nullable=False)
-    model_id: str | None = Field(
-        default=None, foreign_key="models.id", ondelete="SET NULL"
-    )
-
-
-versioned_dataset_table_link = TableLink(
-    model_cls=VersionedDataset,
-    table_cls=VersionedDatasetTable,
-    model_to_table_custom_transforms=dict(
-        model_id=lambda vd: vd.model.id if vd.model else None,
-    ),
-    table_to_model_custom_transforms={},
-)

policyengine/models/__init__.py

@@ -1,30 +0,0 @@
-from .aggregate import Aggregate as Aggregate
-from .aggregate import AggregateType as AggregateType
-from .baseline_parameter_value import (
-    BaselineParameterValue as BaselineParameterValue,
-)
-from .baseline_variable import BaselineVariable as BaselineVariable
-from .dataset import Dataset as Dataset
-from .dynamic import Dynamic as Dynamic
-from .model import Model as Model
-from .model_version import ModelVersion as ModelVersion
-from .parameter import Parameter as Parameter
-from .parameter_value import ParameterValue as ParameterValue
-from .policy import Policy as Policy
-from .policyengine_uk import (
-    policyengine_uk_latest_version as policyengine_uk_latest_version,
-)
-from .policyengine_uk import (
-    policyengine_uk_model as policyengine_uk_model,
-)
-from .policyengine_us import (
-    policyengine_us_latest_version as policyengine_us_latest_version,
-)
-from .policyengine_us import (
-    policyengine_us_model as policyengine_us_model,
-)
-from .report import Report as Report
-from .report_element import ReportElement as ReportElement
-from .simulation import Simulation as Simulation
-from .user import User as User
-from .versioned_dataset import VersionedDataset as VersionedDataset

policyengine/models/aggregate.py

@@ -1,92 +0,0 @@
-from enum import Enum
-from typing import TYPE_CHECKING, Literal
-
-import pandas as pd
-from microdf import MicroDataFrame
-from pydantic import BaseModel
-
-if TYPE_CHECKING:
-    from policyengine.models import Simulation
-
-
-class AggregateType(str, Enum):
-    SUM = "sum"
-    MEAN = "mean"
-    COUNT = "count"
-
-
-class Aggregate(BaseModel):
-    simulation: "Simulation"
-    entity: str
-    variable_name: str
-    year: int | None = None
-    filter_variable_name: str | None = None
-    filter_variable_value: str | None = None
-    filter_variable_leq: float | None = None
-    filter_variable_geq: float | None = None
-    aggregate_function: Literal[
-        AggregateType.SUM, AggregateType.MEAN, AggregateType.COUNT
-    ]
-
-    value: float | None = None
-
-    @staticmethod
-    def run(aggregates: list["Aggregate"]) -> list["Aggregate"]:
-        # Assumes that all aggregates are from the same simulation
-        results = []
-
-        tables = aggregates[0].simulation.result
-        for table in tables:
-            tables[table] = pd.DataFrame(tables[table])
-            weight_col = f"{table}_weight"
-            if weight_col in tables[table].columns:
-                tables[table] = MicroDataFrame(
-                    tables[table], weights=weight_col
-                )
-
-        for agg in aggregates:
-            if agg.entity not in tables:
-                raise ValueError(
-                    f"Entity {agg.entity} not found in simulation results"
-                )
-            table = tables[agg.entity]
-
-            if agg.variable_name not in table.columns:
-                raise ValueError(
-                    f"Variable {agg.variable_name} not found in entity {agg.entity}"
-                )
-
-            df = table
-
-            if agg.year is None:
-                agg.year = aggregates[0].simulation.dataset.year
-
-            if agg.filter_variable_name is not None:
-                if agg.filter_variable_name not in df.columns:
-                    raise ValueError(
-                        f"Filter variable {agg.filter_variable_name} not found in entity {agg.entity}"
-                    )
-                if agg.filter_variable_value is not None:
-                    df = df[
-                        df[agg.filter_variable_name]
-                        == agg.filter_variable_value
-                    ]
-                if agg.filter_variable_leq is not None:
-                    df = df[
-                        df[agg.filter_variable_name] <= agg.filter_variable_leq
-                    ]
-                if agg.filter_variable_geq is not None:
-                    df = df[
-                        df[agg.filter_variable_name] >= agg.filter_variable_geq
-                    ]
-
-            if agg.aggregate_function == AggregateType.SUM:
-                agg.value = float(df[agg.variable_name].sum())
-            elif agg.aggregate_function == AggregateType.MEAN:
-                agg.value = float(df[agg.variable_name].mean())
-            elif agg.aggregate_function == AggregateType.COUNT:
-                agg.value = float((df[agg.variable_name] > 0).sum())
-
-            results.append(agg)
-
-        return results

policyengine/models/baseline_parameter_value.py

@@ -1,14 +0,0 @@
-from datetime import datetime
-
-from pydantic import BaseModel
-
-from .model_version import ModelVersion
-from .parameter import Parameter
-
-
-class BaselineParameterValue(BaseModel):
-    parameter: Parameter
-    model_version: ModelVersion
-    value: float | int | str | bool | list | None = None
-    start_date: datetime
-    end_date: datetime | None = None

policyengine/models/baseline_variable.py

@@ -1,12 +0,0 @@
-from pydantic import BaseModel
-
-from .model_version import ModelVersion
-
-
-class BaselineVariable(BaseModel):
-    id: str
-    model_version: ModelVersion
-    entity: str
-    label: str | None = None
-    description: str | None = None
-    data_type: type | None = None

policyengine/models/dataset.py

@@ -1,18 +0,0 @@
-from typing import Any
-from uuid import uuid4
-
-from pydantic import BaseModel, Field
-
-from .model import Model
-from .versioned_dataset import VersionedDataset
-
-
-class Dataset(BaseModel):
-    id: str = Field(default_factory=lambda: str(uuid4()))
-    name: str
-    description: str | None = None
-    version: str | None = None
-    versioned_dataset: VersionedDataset | None = None
-    year: int | None = None
-    data: Any | None = None
-    model: Model | None = None

policyengine/models/dynamic.py

@@ -1,15 +0,0 @@
-from collections.abc import Callable
-from datetime import datetime
-from uuid import uuid4
-
-from pydantic import BaseModel, Field
-
-
-class Dynamic(BaseModel):
-    id: str = Field(default_factory=lambda: str(uuid4()))
-    name: str
-    description: str | None = None
-    parameter_values: list[str] = []
-    simulation_modifier: Callable | None = None
-    created_at: datetime = Field(default_factory=datetime.now)
-    updated_at: datetime = Field(default_factory=datetime.now)

policyengine/models/model.py

@@ -1,124 +0,0 @@
-from collections.abc import Callable
-from datetime import datetime
-from typing import TYPE_CHECKING
-
-from pydantic import BaseModel
-
-if TYPE_CHECKING:
-    from .baseline_parameter_value import BaselineParameterValue
-    from .baseline_variable import BaselineVariable
-    from .parameter import Parameter
-
-
-class Model(BaseModel):
-    id: str
-    name: str
-    description: str | None = None
-    simulation_function: Callable
-
-    def create_seed_objects(self, model_version):
-        from policyengine_core.parameters import Parameter as CoreParameter
-
-        from .baseline_parameter_value import BaselineParameterValue
-        from .baseline_variable import BaselineVariable
-        from .parameter import Parameter
-
-        if self.id == "policyengine_uk":
-            from policyengine_uk.tax_benefit_system import system
-        elif self.id == "policyengine_us":
-            from policyengine_us.system import system
-        else:
-            raise ValueError("Unsupported model.")
-
-        parameters = []
-        baseline_parameter_values = []
-        baseline_variables = []
-        seen_parameter_ids = set()
-
-        for parameter in system.parameters.get_descendants():
-            # Skip if we've already processed this parameter ID
-            if parameter.name in seen_parameter_ids:
-                continue
-            seen_parameter_ids.add(parameter.name)
-            param = Parameter(
-                id=parameter.name,
-                description=parameter.description,
-                data_type=None,
-                model=self,
-            )
-            parameters.append(param)
-            if isinstance(parameter, CoreParameter):
-                values = parameter.values_list[::-1]
-                param.data_type = type(values[-1].value)
-                for i in range(len(values)):
-                    value_at_instant = values[i]
-                    instant_str = safe_parse_instant_str(
-                        value_at_instant.instant_str
-                    )
-                    if i + 1 < len(values):
-                        next_instant_str = safe_parse_instant_str(
-                            values[i + 1].instant_str
-                        )
-                    else:
-                        next_instant_str = None
-                    baseline_param_value = BaselineParameterValue(
-                        parameter=param,
-                        model_version=model_version,
-                        value=value_at_instant.value,
-                        start_date=instant_str,
-                        end_date=next_instant_str,
-                    )
-                    baseline_parameter_values.append(baseline_param_value)
-
-        for variable in system.variables.values():
-            baseline_variable = BaselineVariable(
-                id=variable.name,
-                model_version=model_version,
-                entity=variable.entity.key,
-                label=variable.label,
-                description=variable.documentation,
-                data_type=variable.value_type,
-            )
-            baseline_variables.append(baseline_variable)
-
-        return SeedObjects(
-            parameters=parameters,
-            baseline_parameter_values=baseline_parameter_values,
-            baseline_variables=baseline_variables,
-        )
-
-
-def safe_parse_instant_str(instant_str: str) -> datetime:
-    if instant_str == "0000-01-01":
-        return datetime(1, 1, 1)
-    else:
-        try:
-            return datetime.strptime(instant_str, "%Y-%m-%d")
-        except ValueError:
-            # Handle invalid dates like 2021-06-31
-            # Try to parse year and month, then use last valid day
-            parts = instant_str.split("-")
-            if len(parts) == 3:
-                year = int(parts[0])
-                month = int(parts[1])
-                day = int(parts[2])
-
-                # Find the last valid day of the month
-                import calendar
-
-                last_day = calendar.monthrange(year, month)[1]
-                if day > last_day:
-                    print(
-                        f"Warning: Invalid date {instant_str}, using {year}-{month:02d}-{last_day:02d}"
-                    )
-                    return datetime(year, month, last_day)
-
-            # If we can't parse it at all, print and raise
-            print(f"Error: Cannot parse date {instant_str}")
-            raise
-
-
-class SeedObjects(BaseModel):
-    parameters: list["Parameter"]
-    baseline_parameter_values: list["BaselineParameterValue"]
-    baseline_variables: list["BaselineVariable"]

policyengine/models/model_version.py

@@ -1,14 +0,0 @@
-from datetime import datetime
-from uuid import uuid4
-
-from pydantic import BaseModel, Field
-
-from .model import Model
-
-
-class ModelVersion(BaseModel):
-    id: str = Field(default_factory=lambda: str(uuid4()))
-    model: Model
-    version: str
-    description: str | None = None
-    created_at: datetime = Field(default_factory=datetime.now)

policyengine/models/policy.py

@@ -1,17 +0,0 @@
-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)

policyengine/models/policyengine_uk.py

@@ -1,114 +0,0 @@
-import importlib.metadata
-
-import pandas as pd
-
-from ..models import Dataset, Dynamic, Model, ModelVersion, Policy
-
-
-def run_policyengine_uk(
-    dataset: "Dataset",
-    policy: "Policy | None" = None,
-    dynamic: "Dynamic | None" = None,
-) -> dict[str, "pd.DataFrame"]:
-    data: dict[str, pd.DataFrame] = dataset.data
-
-    from policyengine_uk import Microsimulation
-    from policyengine_uk.data import UKSingleYearDataset
-
-    pe_input_data = UKSingleYearDataset(
-        person=data["person"],
-        benunit=data["benunit"],
-        household=data["household"],
-        fiscal_year=dataset.year,
-    )
-
-    sim = Microsimulation(dataset=pe_input_data)
-    sim.default_calculation_period = dataset.year
-
-    def simulation_modifier(sim: Microsimulation):
-        if policy is not None and len(policy.parameter_values) > 0:
-            for parameter_value in policy.parameter_values:
-                sim.tax_benefit_system.parameters.get_child(
-                    parameter_value.parameter.id
-                ).update(
-                    value=parameter_value.value,
-                    start=parameter_value.start_date.strftime("%Y-%m-%d"),
-                    stop=parameter_value.end_date.strftime("%Y-%m-%d")
-                    if parameter_value.end_date
-                    else None,
-                )
-
-        if dynamic is not None and len(dynamic.parameter_values) > 0:
-            for parameter_value in dynamic.parameter_values:
-                sim.tax_benefit_system.parameters.get_child(
-                    parameter_value.parameter.id
-                ).update(
-                    value=parameter_value.value,
-                    start=parameter_value.start_date.strftime("%Y-%m-%d"),
-                    stop=parameter_value.end_date.strftime("%Y-%m-%d")
-                    if parameter_value.end_date
-                    else None,
-                )
-
-        if dynamic is not None and dynamic.simulation_modifier is not None:
-            dynamic.simulation_modifier(sim)
-        if policy is not None and policy.simulation_modifier is not None:
-            policy.simulation_modifier(sim)
-
-    simulation_modifier(sim)
-
-    # Skip reforms for now
-
-    output_data = {}
-
-    variable_blacklist = [  # TEMPORARY: we need to fix policyengine-uk to make these only take a long time with non-default parameters set to true.
-        "is_uc_entitled_baseline",
-        "income_elasticity_lsr",
-        "child_benefit_opts_out",
-        "housing_benefit_baseline_entitlement",
-        "baseline_ctc_entitlement",
-        "pre_budget_change_household_tax",
-        "pre_budget_change_household_net_income",
-        "is_on_cliff",
-        "marginal_tax_rate_on_capital_gains",
-        "relative_capital_gains_mtr_change",
-        "pre_budget_change_ons_equivalised_income_decile",
-        "substitution_elasticity",
-        "marginal_tax_rate",
-        "cliff_evaluated",
-        "cliff_gap",
-        "substitution_elasticity_lsr",
-        "relative_wage_change",
-        "relative_income_change",
-        "pre_budget_change_household_benefits",
-    ]
-
-    for entity in ["person", "benunit", "household"]:
-        output_data[entity] = pd.DataFrame()
-        for variable in sim.tax_benefit_system.variables.values():
-            correct_entity = variable.entity.key == entity
-            if variable.name in variable_blacklist:
-                continue
-            if variable.definition_period != "year":
-                continue
-            if correct_entity:
-                output_data[entity][variable.name] = sim.calculate(
-                    variable.name
-                )
-
-    return output_data
-
-
-policyengine_uk_model = Model(
-    id="policyengine_uk",
-    name="PolicyEngine UK",
-    description="PolicyEngine's open-source tax-benefit microsimulation model.",
-    simulation_function=run_policyengine_uk,
-)
-
-# Get policyengine-uk version
-
-policyengine_uk_latest_version = ModelVersion(
-    model=policyengine_uk_model,
-    version=importlib.metadata.distribution("policyengine_uk").version,
-)

policyengine/models/policyengine_us.py

@@ -1,115 +0,0 @@
-import importlib.metadata
-
-import pandas as pd
-
-from ..models import Dataset, Dynamic, Model, ModelVersion, Policy
-
-
-def run_policyengine_us(
-    dataset: "Dataset",
-    policy: "Policy | None" = None,
-    dynamic: "Dynamic | None" = None,
-) -> dict[str, "pd.DataFrame"]:
-    data: dict[str, pd.DataFrame] = dataset.data
-
-    person_df = pd.DataFrame()
-
-    for table_name, table in data.items():
-        if table_name == "person":
-            for col in table.columns:
-                person_df[f"{col}__{dataset.year}"] = table[col].values
-        else:
-            foreign_key = data["person"][f"person_{table_name}_id"]
-            primary_key = data[table_name][f"{table_name}_id"]
-
-            projected = table.set_index(primary_key).loc[foreign_key]
-
-            for col in projected.columns:
-                person_df[f"{col}__{dataset.year}"] = projected[col].values
-
-    from policyengine_us import Microsimulation
-
-    sim = Microsimulation(dataset=person_df)
-    sim.default_calculation_period = dataset.year
-
-    def simulation_modifier(sim: Microsimulation):
-        if policy is not None and len(policy.parameter_values) > 0:
-            for parameter_value in policy.parameter_values:
-                sim.tax_benefit_system.parameters.get_child(
-                    parameter_value.parameter.id
-                ).update(
-                    parameter_value.value,
-                    start=parameter_value.start_date.strftime("%Y-%m-%d"),
-                    stop=parameter_value.end_date.strftime("%Y-%m-%d")
-                    if parameter_value.end_date
-                    else None,
-                )
-
-        if dynamic is not None and len(dynamic.parameter_values) > 0:
-            for parameter_value in dynamic.parameter_values:
-                sim.tax_benefit_system.parameters.get_child(
-                    parameter_value.parameter.id
-                ).update(
-                    parameter_value.value,
-                    start=parameter_value.start_date.strftime("%Y-%m-%d"),
-                    stop=parameter_value.end_date.strftime("%Y-%m-%d")
-                    if parameter_value.end_date
-                    else None,
-                )
-
-        if dynamic is not None and dynamic.simulation_modifier is not None:
-            dynamic.simulation_modifier(sim)
-        if policy is not None and policy.simulation_modifier is not None:
-            policy.simulation_modifier(sim)
-
-    simulation_modifier(sim)
-
-    # Skip reforms for now
-
-    output_data = {}
-
-    variable_whitelist = [
-        "household_net_income",
-    ]
-
-    for variable in variable_whitelist:
-        sim.calculate(variable)
-
-    for entity in [
-        "person",
-        "marital_unit",
-        "family",
-        "tax_unit",
-        "spm_unit",
-        "household",
-    ]:
-        output_data[entity] = pd.DataFrame()
-        for variable in sim.tax_benefit_system.variables.values():
-            correct_entity = variable.entity.key == entity
-            if str(dataset.year) not in list(
-                map(str, sim.get_holder(variable.name).get_known_periods())
-            ):
-                continue
-            if variable.definition_period != "year":
-                continue
-            if not correct_entity:
-                continue
-            output_data[entity][variable.name] = sim.calculate(variable.name)
-
-    return output_data
-
-
-policyengine_us_model = Model(
-    id="policyengine_us",
-    name="PolicyEngine US",
-    description="PolicyEngine's open-source tax-benefit microsimulation model.",
-    simulation_function=run_policyengine_us,
-)
-
-# Get policyengine-uk version
-
-
-policyengine_us_latest_version = ModelVersion(
-    model=policyengine_us_model,
-    version=importlib.metadata.distribution("policyengine_us").version,
-)