diff --git a/changelog.d/v4-dict-reforms.added.md b/changelog.d/v4-dict-reforms.added.md new file mode 100644 index 00000000..02405cdc --- /dev/null +++ b/changelog.d/v4-dict-reforms.added.md @@ -0,0 +1 @@ +``Simulation(policy={...})`` and ``Simulation(dynamic={...})`` now accept the same flat ``{"param.path": value}`` / ``{"param.path": {date: value}}`` dict that ``pe.{uk,us}.calculate_household(reform=...)`` accepts. Dicts are compiled to full ``Policy`` / ``Dynamic`` objects on construction using the ``tax_benefit_model_version`` for parameter-path validation and ``dataset.year`` for scalar effective-date defaulting. Removes the last place where population microsim required building ``Parameter`` / ``ParameterValue`` by hand. diff --git a/src/policyengine/core/simulation.py b/src/policyengine/core/simulation.py index 5002b141..9fddd240 100644 --- a/src/policyengine/core/simulation.py +++ b/src/policyengine/core/simulation.py @@ -1,9 +1,9 @@ import logging from datetime import datetime -from typing import Optional +from typing import Any, Optional, Union from uuid import uuid4 -from pydantic import BaseModel, Field +from pydantic import BaseModel, Field, model_validator from .cache import LRUCache from .dataset import Dataset @@ -22,8 +22,21 @@ class Simulation(BaseModel): created_at: datetime = Field(default_factory=datetime.now) updated_at: datetime = Field(default_factory=datetime.now) - policy: Optional[Policy] = None - dynamic: Optional[Dynamic] = None + policy: Optional[Union[Policy, dict[str, Any]]] = Field( + default=None, + description=( + "Reform policy. Pass a ``Policy`` directly, or a flat " + "``{'param.path': value}`` / ``{'param.path': {date: value}}`` " + "dict and it will be compiled against " + "``tax_benefit_model_version`` at run time." + ), + ) + dynamic: Optional[Union[Dynamic, dict[str, Any]]] = Field( + default=None, + description=( + "Behavioural-response overlay. Same dict shape as ``policy``." + ), + ) dataset: Dataset = None scoping_strategy: Optional[ScopingStrategy] = Field( @@ -44,6 +57,47 @@ class Simulation(BaseModel): output_dataset: Optional[Dataset] = None + @model_validator(mode="after") + def _compile_dict_reforms(self) -> "Simulation": + """Coerce dict ``policy`` / ``dynamic`` inputs into proper objects. + + We can't do this in a ``field_validator`` because compiling a + reform requires the ``tax_benefit_model_version`` (for parameter + path validation) and the ``dataset.year`` (for the scalar + effective-date default). By the time ``model_validator(mode="after")`` + fires, both are already on ``self``. + """ + from policyengine.tax_benefit_models.common.reform import ( + compile_reform_to_dynamic, + compile_reform_to_policy, + ) + + if isinstance(self.policy, dict): + if self.tax_benefit_model_version is None: + raise ValueError( + "Cannot compile a dict policy without " + "tax_benefit_model_version; pass model_version or a Policy." + ) + year = getattr(self.dataset, "year", None) + self.policy = compile_reform_to_policy( + self.policy, + year=year, + model_version=self.tax_benefit_model_version, + ) + if isinstance(self.dynamic, dict): + if self.tax_benefit_model_version is None: + raise ValueError( + "Cannot compile a dict dynamic without " + "tax_benefit_model_version; pass model_version or a Dynamic." + ) + year = getattr(self.dataset, "year", None) + self.dynamic = compile_reform_to_dynamic( + self.dynamic, + year=year, + model_version=self.tax_benefit_model_version, + ) + return self + def run(self): self.tax_benefit_model_version.run(self) diff --git a/src/policyengine/tax_benefit_models/common/__init__.py b/src/policyengine/tax_benefit_models/common/__init__.py index 6f6efa25..654f350d 100644 --- a/src/policyengine/tax_benefit_models/common/__init__.py +++ b/src/policyengine/tax_benefit_models/common/__init__.py @@ -10,5 +10,7 @@ MicrosimulationModelVersion as MicrosimulationModelVersion, ) from .reform import compile_reform as compile_reform +from .reform import compile_reform_to_dynamic as compile_reform_to_dynamic +from .reform import compile_reform_to_policy as compile_reform_to_policy from .result import EntityResult as EntityResult from .result import HouseholdResult as HouseholdResult diff --git a/src/policyengine/tax_benefit_models/common/reform.py b/src/policyengine/tax_benefit_models/common/reform.py index 0bb83182..0289c730 100644 --- a/src/policyengine/tax_benefit_models/common/reform.py +++ b/src/policyengine/tax_benefit_models/common/reform.py @@ -28,11 +28,14 @@ from __future__ import annotations +import datetime from collections.abc import Mapping from difflib import get_close_matches from typing import TYPE_CHECKING, Any, Optional if TYPE_CHECKING: + from policyengine.core.dynamic import Dynamic + from policyengine.core.policy import Policy from policyengine.core.tax_benefit_model_version import TaxBenefitModelVersion @@ -80,3 +83,95 @@ def compile_reform( else: compiled[parameter_path] = {default_date: spec} return compiled + + +def _reform_dict_to_parameter_values( + reform: Mapping[str, Any], + *, + year: Optional[int], + model_version: "TaxBenefitModelVersion", +) -> list: + """Compile a flat reform dict into a list of ``ParameterValue`` objects. + + Uses :func:`compile_reform` for path validation and effective-date + defaulting, then materialises each ``{path: {date: value}}`` pair + as an open-ended ``ParameterValue`` bound to a + ``Parameter(name=path, tax_benefit_model_version=model_version)``. + """ + from policyengine.core.parameter import Parameter + from policyengine.core.parameter_value import ParameterValue + + compiled = compile_reform(reform, year=year, model_version=model_version) + if compiled is None: + return [] + + parameter_values: list[ParameterValue] = [] + for path, date_to_value in compiled.items(): + for effective_date, value in date_to_value.items(): + data_type = type(value) if isinstance(value, (int, float, bool)) else float + parameter_values.append( + ParameterValue( + parameter=Parameter( + name=path, + tax_benefit_model_version=model_version, + data_type=data_type, + ), + start_date=datetime.datetime.strptime( + effective_date, "%Y-%m-%d" + ), + end_date=None, + value=value, + ) + ) + return parameter_values + + +def compile_reform_to_policy( + reform: Optional[Mapping[str, Any]], + *, + year: Optional[int], + model_version: "TaxBenefitModelVersion", + name: Optional[str] = None, +) -> "Optional[Policy]": + """Compile a flat reform dict into a fully-assembled ``Policy``. + + Accepts the same ``{param.path: value}`` / + ``{param.path: {date: value}}`` shape as + :func:`compile_reform`, but returns a ready-to-use ``Policy`` with + :class:`~policyengine.core.parameter_value.ParameterValue` objects + instead of a raw dict. This lets ``Simulation(policy={"..."}: ...)`` + work without the caller building ``Parameter`` / ``ParameterValue`` + by hand. + """ + from policyengine.core.policy import Policy + + parameter_values = _reform_dict_to_parameter_values( + reform or {}, year=year, model_version=model_version + ) + if not parameter_values: + return None + return Policy(name=name or "Reform", parameter_values=parameter_values) + + +def compile_reform_to_dynamic( + reform: Optional[Mapping[str, Any]], + *, + year: Optional[int], + model_version: "TaxBenefitModelVersion", + name: Optional[str] = None, +) -> "Optional[Dynamic]": + """Compile a flat reform dict into a ready-to-use ``Dynamic``. + + See :func:`compile_reform_to_policy` — this is the ``Dynamic`` + counterpart for behavioural responses. + """ + from policyengine.core.dynamic import Dynamic + + parameter_values = _reform_dict_to_parameter_values( + reform or {}, year=year, model_version=model_version + ) + if not parameter_values: + return None + return Dynamic( + name=name or "Dynamic response", parameter_values=parameter_values + ) diff --git a/tests/test_dict_reforms_on_simulation.py b/tests/test_dict_reforms_on_simulation.py new file mode 100644 index 00000000..86460597 --- /dev/null +++ b/tests/test_dict_reforms_on_simulation.py @@ -0,0 +1,123 @@ +"""``Simulation(policy={...})`` and ``Simulation(dynamic={...})``. + +These tests pin the v4 contract: the same flat reform dict shape that +``pe.{uk,us}.calculate_household(reform=...)`` accepts is also accepted +by ``Simulation(policy=...)`` / ``Simulation(dynamic=...)``, and is +compiled into the full ``Policy`` / ``Dynamic`` object on construction. +We exercise only the coercion path — no country microsim is run — so +the tests are fast and don't need HF credentials. +""" + +from __future__ import annotations + +import pytest + +pytest.importorskip("policyengine_us") + +import policyengine as pe +from policyengine.core import Dynamic, Policy, Simulation +from tests.fixtures.filtering_fixtures import us_test_dataset # noqa: F401 + + +@pytest.fixture +def tiny_dataset(us_test_dataset): + """In-memory US dataset pinned to 2026. Simulation is never .run() in these tests.""" + us_test_dataset.year = 2026 + return us_test_dataset + + +class TestDictPolicyCoercion: + def test__dict_policy__then_compiled_to_policy_with_parameter_values(self, tiny_dataset): + sim = Simulation( + dataset=tiny_dataset, + tax_benefit_model_version=pe.us.model, + policy={"gov.irs.credits.ctc.amount.base[0].amount": 3_000}, + ) + assert isinstance(sim.policy, Policy) + assert len(sim.policy.parameter_values) == 1 + + pv = sim.policy.parameter_values[0] + assert pv.parameter.name == "gov.irs.credits.ctc.amount.base[0].amount" + assert pv.value == 3_000 + # Scalar reforms default the effective date to {year}-01-01. + assert pv.start_date.year == 2026 + assert pv.start_date.month == 1 + + def test__dict_policy_with_effective_date__then_start_date_matches(self, tiny_dataset): + sim = Simulation( + dataset=tiny_dataset, + tax_benefit_model_version=pe.us.model, + policy={ + "gov.irs.credits.ctc.amount.base[0].amount": { + "2026-07-01": 2_500, + "2027-01-01": 3_000, + }, + }, + ) + assert isinstance(sim.policy, Policy) + assert len(sim.policy.parameter_values) == 2 + starts = sorted(pv.start_date for pv in sim.policy.parameter_values) + assert [d.strftime("%Y-%m-%d") for d in starts] == [ + "2026-07-01", + "2027-01-01", + ] + + def test__unknown_parameter_path__raises_with_suggestion(self, tiny_dataset): + with pytest.raises(ValueError) as exc: + Simulation( + dataset=tiny_dataset, + tax_benefit_model_version=pe.us.model, + policy={ + # plausible typo of the real path + "gov.irs.credits.ctc.amount.base[0].amont": 3_000, + }, + ) + assert "not defined" in str(exc.value) + assert "did you mean" in str(exc.value) + + def test__existing_policy_object_passes_through_unchanged(self, tiny_dataset): + import datetime + + from policyengine.core import Parameter, ParameterValue + + existing = Policy( + name="Existing", + parameter_values=[ + ParameterValue( + parameter=Parameter( + name="gov.irs.credits.ctc.amount.base[0].amount", + tax_benefit_model_version=pe.us.model, + data_type=float, + ), + start_date=datetime.datetime(2026, 1, 1), + end_date=None, + value=2_750, + ) + ], + ) + sim = Simulation( + dataset=tiny_dataset, + tax_benefit_model_version=pe.us.model, + policy=existing, + ) + assert sim.policy is existing + + def test__dict_without_model_version__raises(self, tiny_dataset): + with pytest.raises(ValueError) as exc: + Simulation( + dataset=tiny_dataset, + policy={"gov.irs.credits.ctc.amount.base[0].amount": 3_000}, + ) + assert "tax_benefit_model_version" in str(exc.value) + + +class TestDictDynamicCoercion: + def test__dict_dynamic__then_compiled_to_dynamic(self, tiny_dataset): + sim = Simulation( + dataset=tiny_dataset, + tax_benefit_model_version=pe.us.model, + dynamic={"gov.irs.credits.ctc.amount.base[0].amount": 2_800}, + ) + assert isinstance(sim.dynamic, Dynamic) + assert len(sim.dynamic.parameter_values) == 1 + assert sim.dynamic.parameter_values[0].value == 2_800