Skip to content

Derive UK person and benunit weights when loading a dataset#479

Draft
vahid-ahmadi wants to merge 1 commit into
mainfrom
fix/uk-dataset-derive-person-benunit-weights
Draft

Derive UK person and benunit weights when loading a dataset#479
vahid-ahmadi wants to merge 1 commit into
mainfrom
fix/uk-dataset-derive-person-benunit-weights

Conversation

@vahid-ahmadi

Copy link
Copy Markdown
Contributor

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.

PolicyEngineUKDataset.load() assumed `person_weight` and `benunit_weight`
were already columns, but published UK microdata carries only
`household_weight` at the household level. Supplying a real dataset file
directly raised `KeyError: 'person_weight'`.

Factor the derivation that create_datasets performed inline into shared
`derive_person_weight` / `derive_benunit_weight` helpers and call them
from load() as well, so the two paths cannot drift. The helpers are
no-ops when the weight columns are already present, and report missing
join columns with a clear message.

Fixes #478

Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

PolicyEngineUKDataset.load() raises KeyError: 'person_weight' for directly-supplied UK datasets

1 participant