[vahid-ahmadi]

Fixes #478

The failure

Constructing a PolicyEngineUKDataset against any real published UK dataset file raises KeyError: 'person_weight' at construction time:

import policyengine as pe
from policyengine.tax_benefit_models.uk.model import PolicyEngineUKDataset

P = ".../policyengine_uk_data/storage/enhanced_frs_2024_25.h5"
ds = PolicyEngineUKDataset(name="x", description="x", filepath=P, year=2024)
sim = pe.Simulation(id="t", dataset=ds, tax_benefit_model_version=pe.uk.model,
                    extra_variables={"person": ["capital_gains"]})
sim.ensure()   # -> KeyError: 'person_weight'

Why

PolicyEngineUKDataset.load() assumed person_weight and benunit_weight were already columns in the file. Published UK microdata carries only household_weight, at the household level — person and benunit weights are derived from it. Verified absent in enhanced_frs_2024_25.h5, enhanced_frs_2023_24.h5, and the certified populace_uk_2023.h5.

create_datasets already did that derivation inline (merging household_weight onto the person frame via person_household_id, and onto benunits via the person frame) and saved the rewritten file. So a manifest-resolved dataset worked while a directly-supplied filepath did not.

The change

• Factored the derivation out of create_datasets into shared derive_person_weight / derive_benunit_weight helpers, and call them from load() too, so the two paths cannot drift.
• Derivation happens only when the column is absent — a dataset carrying its own person-level weights is returned untouched.
• Missing join columns now raise a clear, explanatory error rather than a bare KeyError.

The manifest path is unchanged. create_datasets performs exactly the same merges, in the same order, producing the same columns and values — verified below against the pre-change code.

User-facing benefit

Callers can run policyengine.py on their own UK dataset — for example a newer Enhanced FRS release — without waiting for it to appear in the certified bundle manifest.

Verification

Reproduced the KeyError on main, then on this branch, against enhanced_frs_2024_25.h5:

person_weight equal:  True
benunit_weight equal: True
person cols equal: True   benunit cols equal: True
sample weights: [2106.553  365.01886  383.83633  2106.553]   nan count: 0

(equal compares against the exact merge code create_datasets used before this change — the value-equivalence check, not just "it no longer raises".)

End-to-end sim.ensure() from the report now succeeds on both enhanced_frs_2024_25.h5 and enhanced_frs_2023_24.h5.

Tests

New tests/test_uk_dataset_weights.py (9 tests) covers derivation when columns are missing, no-op when present (both helper-level and through load()), clear errors for missing join columns, and equivalence with the original create_datasets merge.

Full suite: 747 passed.

[vahid-ahmadi]

Still needed: on main (485dc23) PolicyEngineUKDataset.load() on a file with only household_weight still raises KeyError: 'person_weight'. Merged main in (784d7d6). Main has since moved the weight merge into _year_data_with_weights and reads files through _read_year_data, so the derivation now sits in _read_year_data (also covering the data-year tables) and _year_data_with_weights uses the shared helpers. The repro now loads; the weights tests, year-file data-year tests and dataset tests pass locally (58 passed), and ruff passes.

🤖 Generated with Claude Code