Make UK paper analysis reproducible on Python 3.14

README.md

@@ -45,6 +45,7 @@ print(f"Total UC spending: £{agg.result / 1e9:.1f}bn")
 - `examples/income_distribution_us.py`: Analyse benefit distribution by decile
 - `examples/employment_income_variation_uk.py`: Model employment income phase-outs
 - `examples/policy_change_uk.py`: Analyse policy reform impacts
+- `examples/paper_repro_uk.py`: Reproduce the UK reform analysis used in the JOSS paper draft
 
 ## Installation
 
@@ -136,6 +137,16 @@ echo "Description of change" > changelog.d/my-change.added
 
 On merge, the versioning workflow bumps the version, builds the changelog, and creates a GitHub Release.
 
+## Paper reproduction
+
+Use the pinned interpreter and the UK extra to run the checked-in paper repro:
+
+```bash
+uv run --python 3.14 --extra uk python examples/paper_repro_uk.py
+```
+
+On first run this will create `./data/enhanced_frs_2023_24_year_2026.h5`.
+
 ## Features
 
 - **Multi-country support**: UK and US tax-benefit systems

changelog.d/codex.paper-repro-314.fixed.md

@@ -0,0 +1 @@
+Fixed the UK paper reproduction workflow so the checked-in example runs on Python 3.14 and the associated analysis helpers handle that path cleanly.

examples/paper_repro_uk.py

@@ -0,0 +1,76 @@
+"""Reproduce the UK policy-reform analysis used in the JOSS paper draft.
+
+This script uses the same reform shown in `paper.md`, but adds the missing
+dataset setup so it can run end-to-end from a fresh checkout.
+
+Run:
+    uv run --python 3.14 --extra uk python examples/paper_repro_uk.py
+"""
+
+from datetime import date
+
+from policyengine.core import Parameter, ParameterValue, Policy, Simulation
+from policyengine.tax_benefit_models.uk import (
+    economic_impact_analysis,
+    ensure_datasets,
+    uk_latest,
+)
+
+
+def load_dataset(year: int = 2026):
+    """Load or create the enhanced UK dataset used for population analysis."""
+    datasets = ensure_datasets(
+        datasets=["hf://policyengine/policyengine-uk-data/enhanced_frs_2023_24.h5"],
+        years=[year],
+    )
+    return datasets[f"enhanced_frs_2023_24_{year}"]
+
+
+def create_reform(year: int = 2026) -> Policy:
+    """Set the UK personal allowance to zero for one tax year."""
+    parameter = Parameter(
+        name="gov.hmrc.income_tax.allowances.personal_allowance.amount",
+        tax_benefit_model_version=uk_latest,
+    )
+
+    return Policy(
+        name="Zero personal allowance",
+        parameter_values=[
+            ParameterValue(
+                parameter=parameter,
+                start_date=date(year, 1, 1),
+                end_date=date(year, 12, 31),
+                value=0,
+            )
+        ],
+    )
+
+
+def main():
+    dataset = load_dataset()
+    reform = create_reform()
+
+    baseline = Simulation(
+        dataset=dataset,
+        tax_benefit_model_version=uk_latest,
+    )
+    reformed = Simulation(
+        dataset=dataset,
+        tax_benefit_model_version=uk_latest,
+        policy=reform,
+    )
+
+    analysis = economic_impact_analysis(baseline, reformed)
+
+    first_decile = analysis.decile_impacts.outputs[0]
+
+    print(f"Dataset: {dataset.filepath}")
+    print(f"Households: {len(dataset.data.household):,}")
+    print(f"Baseline Gini: {analysis.baseline_inequality.gini:.4f}")
+    print(f"Reform Gini: {analysis.reform_inequality.gini:.4f}")
+    print(f"Decile 1 mean change: {first_decile.absolute_change:,.2f}")
+    print(f"Programmes analysed: {len(analysis.programme_statistics.outputs)}")
+
+
+if __name__ == "__main__":
+    main()

src/policyengine/outputs/decile_impact.py

@@ -97,32 +97,48 @@ def run(self):
 
 
 def calculate_decile_impacts(
-    dataset: Dataset,
-    tax_benefit_model_version: TaxBenefitModelVersion,
+    dataset: Dataset | None = None,
+    tax_benefit_model_version: TaxBenefitModelVersion | None = None,
     baseline_policy: Policy | None = None,
     reform_policy: Policy | None = None,
     dynamic: Dynamic | None = None,
     income_variable: str = "equiv_hbai_household_net_income",
     entity: str | None = None,
     quantiles: int = 10,
+    baseline_simulation: Simulation | None = None,
+    reform_simulation: Simulation | None = None,
 ) -> OutputCollection[DecileImpact]:
     """Calculate decile-by-decile impact of a reform.
 
     Returns:
         OutputCollection containing list of DecileImpact objects and DataFrame
     """
-    baseline_simulation = Simulation(
-        dataset=dataset,
-        tax_benefit_model_version=tax_benefit_model_version,
-        policy=baseline_policy,
-        dynamic=dynamic,
-    )
-    reform_simulation = Simulation(
-        dataset=dataset,
-        tax_benefit_model_version=tax_benefit_model_version,
-        policy=reform_policy,
-        dynamic=dynamic,
-    )
+    if (baseline_simulation is None) != (reform_simulation is None):
+        raise ValueError(
+            "baseline_simulation and reform_simulation must be provided together"
+        )
+
+    if baseline_simulation is None:
+        if dataset is None or tax_benefit_model_version is None:
+            raise ValueError(
+                "dataset and tax_benefit_model_version are required when simulations are not provided"
+            )
+
+        baseline_simulation = Simulation(
+            dataset=dataset,
+            tax_benefit_model_version=tax_benefit_model_version,
+            policy=baseline_policy,
+            dynamic=dynamic,
+        )
+        reform_simulation = Simulation(
+            dataset=dataset,
+            tax_benefit_model_version=tax_benefit_model_version,
+            policy=reform_policy,
+            dynamic=dynamic,
+        )
+
+    baseline_simulation.ensure()
+    reform_simulation.ensure()
 
     results = []
     for decile in range(1, quantiles + 1):

src/policyengine/tax_benefit_models/uk/__init__.py

@@ -31,6 +31,7 @@
     UKYearData.model_rebuild()
     PolicyEngineUKDataset.model_rebuild()
     PolicyEngineUKLatest.model_rebuild()
+    ProgrammeStatistics.model_rebuild()
 
     __all__ = [
         "UKYearData",

src/policyengine/tax_benefit_models/uk/analysis.py

@@ -202,33 +202,32 @@ def economic_impact_analysis(
 
     # Decile impact
     decile_impacts = calculate_decile_impacts(
-        dataset=baseline_simulation.dataset,
-        tax_benefit_model_version=baseline_simulation.tax_benefit_model_version,
-        baseline_policy=baseline_simulation.policy,
-        reform_policy=reform_simulation.policy,
-        dynamic=baseline_simulation.dynamic,
+        baseline_simulation=baseline_simulation,
+        reform_simulation=reform_simulation,
     )
 
     # Major programmes to analyse
     programmes = {
         # Tax
-        "income_tax": {"entity": "person", "is_tax": True},
-        "national_insurance": {"entity": "person", "is_tax": True},
-        "vat": {"entity": "household", "is_tax": True},
-        "council_tax": {"entity": "household", "is_tax": True},
+        "income_tax": {"is_tax": True},
+        "national_insurance": {"is_tax": True},
+        "vat": {"is_tax": True},
+        "council_tax": {"is_tax": True},
         # Benefits
-        "universal_credit": {"entity": "person", "is_tax": False},
-        "child_benefit": {"entity": "person", "is_tax": False},
-        "pension_credit": {"entity": "person", "is_tax": False},
-        "income_support": {"entity": "person", "is_tax": False},
-        "working_tax_credit": {"entity": "person", "is_tax": False},
-        "child_tax_credit": {"entity": "person", "is_tax": False},
+        "universal_credit": {"is_tax": False},
+        "child_benefit": {"is_tax": False},
+        "pension_credit": {"is_tax": False},
+        "income_support": {"is_tax": False},
+        "working_tax_credit": {"is_tax": False},
+        "child_tax_credit": {"is_tax": False},
     }
 
     programme_statistics = []
 
     for programme_name, programme_info in programmes.items():
-        entity = programme_info["entity"]
+        entity = baseline_simulation.tax_benefit_model_version.get_variable(
+            programme_name
+        ).entity
         is_tax = programme_info["is_tax"]
 
         stats = ProgrammeStatistics(

src/policyengine/tax_benefit_models/uk/model.py

@@ -108,6 +108,8 @@ class PolicyEngineUKLatest(TaxBenefitModelVersion):
             # Income and benefits
             "universal_credit",
             "child_benefit",
+            "pension_credit",
+            "income_support",
             "working_tax_credit",
             "child_tax_credit",
         ],

src/policyengine/tax_benefit_models/uk/outputs.py

@@ -1,27 +1,22 @@
 """UK-specific output templates."""
 
-from typing import TYPE_CHECKING
-
 from pydantic import ConfigDict
 
-from policyengine.core import Output
+from policyengine.core import Output, Simulation
 from policyengine.outputs.aggregate import Aggregate, AggregateType
 from policyengine.outputs.change_aggregate import (
     ChangeAggregate,
     ChangeAggregateType,
 )
 
-if TYPE_CHECKING:
-    from policyengine.core.simulation import Simulation
-
 
 class ProgrammeStatistics(Output):
     """Single programme's statistics from a policy reform - represents one database row."""
 
     model_config = ConfigDict(arbitrary_types_allowed=True)
 
-    baseline_simulation: "Simulation"
-    reform_simulation: "Simulation"
+    baseline_simulation: Simulation
+    reform_simulation: Simulation
     programme_name: str
     entity: str
     is_tax: bool = False

src/policyengine/tax_benefit_models/us/__init__.py

@@ -31,6 +31,7 @@
     USYearData.model_rebuild()
     PolicyEngineUSDataset.model_rebuild()
     PolicyEngineUSLatest.model_rebuild()
+    ProgramStatistics.model_rebuild()
 
     __all__ = [
         "USYearData",

src/policyengine/tax_benefit_models/us/analysis.py

@@ -218,11 +218,8 @@ def economic_impact_analysis(
 
     # Decile impact (using household_net_income for US)
     decile_impacts = calculate_decile_impacts(
-        dataset=baseline_simulation.dataset,
-        tax_benefit_model_version=baseline_simulation.tax_benefit_model_version,
-        baseline_policy=baseline_simulation.policy,
-        reform_policy=reform_simulation.policy,
-        dynamic=baseline_simulation.dynamic,
+        baseline_simulation=baseline_simulation,
+        reform_simulation=reform_simulation,
         income_variable="household_net_income",
     )
 

src/policyengine/tax_benefit_models/us/outputs.py

@@ -1,27 +1,22 @@
 """US-specific output templates."""
 
-from typing import TYPE_CHECKING
-
 from pydantic import ConfigDict
 
-from policyengine.core import Output
+from policyengine.core import Output, Simulation
 from policyengine.outputs.aggregate import Aggregate, AggregateType
 from policyengine.outputs.change_aggregate import (
     ChangeAggregate,
     ChangeAggregateType,
 )
 
-if TYPE_CHECKING:
-    from policyengine.core.simulation import Simulation
-
 
 class ProgramStatistics(Output):
     """Single program's statistics from a policy reform - represents one database row."""
 
     model_config = ConfigDict(arbitrary_types_allowed=True)
 
-    baseline_simulation: "Simulation"
-    reform_simulation: "Simulation"
+    baseline_simulation: Simulation
+    reform_simulation: Simulation
     program_name: str
     entity: str
     is_tax: bool = False