policyengine-simulation-mechanics

What this skill does

# PolicyEngine Simulation Mechanics

This skill covers advanced patterns for working with policyengine.py simulations, including caching, result access, and entity mapping.

## CRITICAL: Environment Setup

**Before writing any code, check the environment.** The policyengine.py package must be installed in the project's `.venv`.

```bash
# Always run from the policyengine.py repo root:
cd /path/to/policyengine.py
uv run python script.py

# Or activate first:
source .venv/bin/activate
python script.py

# NEVER use bare `pip install` — always:
uv pip install -e ".[uk]"   # for UK work
uv pip install -e ".[us]"   # for US work
```

**If `from policyengine.core import Simulation` fails:**
```bash
cd /path/to/policyengine.py
uv pip install -e ".[uk]"
# Then re-run with: uv run python script.py
```

## CRITICAL: Correct Import Paths

**Only these imports exist — do not guess others:**

```python
# Core simulation
from policyengine.core import Simulation

# UK model
from policyengine.tax_benefit_models.uk import (
    uk_latest,          # The model version (pass as tax_benefit_model_version=)
    uk_model,           # The model itself
    PolicyEngineUKDataset,
    UKYearData,
    create_datasets,    # Create & cache datasets from HF source
    load_datasets,      # Load cached datasets from disk
    ensure_datasets,    # Create if missing, load if present (recommended)
)

# US model
from policyengine.tax_benefit_models.us import (
    us_latest,
    PolicyEngineUSDataset,
    ensure_datasets,
)

# Outputs
from policyengine.outputs.aggregate import Aggregate, AggregateType
from policyengine.outputs.change_aggregate import ChangeAggregate, ChangeAggregateType

# Plotting
from policyengine.utils.plotting import COLORS, format_fig
```

**There is NO:**
- `policyengine.core.dataset_registry`
- `policyengine.datasets`
- `policyengine.core.dataset_version.DatasetVersion.list()`

## UK Datasets

### Loading UK datasets

Use `ensure_datasets()` — it returns a `dict[str, PolicyEngineUKDataset]`, building files in `./data/` on first run and loading from disk on subsequent runs.

**WARNING:** `from policyengine.tax_benefit_models.uk import datasets` gives you the Python **submodule**, not a dict. Never index it like a dict.

```python
from policyengine.tax_benefit_models.uk import ensure_datasets

uk = ensure_datasets(
    datasets=[
        "hf://policyengine/policyengine-uk-data/frs_2023_24.h5",
        "hf://policyengine/policyengine-uk-data/enhanced_frs_2023_24.h5",
    ],
    years=[2026],
    data_folder="./data",
)

efrs = uk["enhanced_frs_2023_24_2026"]
frs  = uk["frs_2023_24_2026"]
```

**Dict key format:** `"{stem}_{year}"` e.g. `"enhanced_frs_2023_24_2026"`

**To force regeneration:** delete `./data/` and call `ensure_datasets()` again.

### Loading US datasets

```python
from policyengine.tax_benefit_models.us import ensure_datasets

us = ensure_datasets(
    datasets=["hf://policyengine/policyengine-us-data/enhanced_cps_2024.h5"],
    years=[2026],
    data_folder="./data",
)

ecps = us["enhanced_cps_2024_2026"]
```

**Default US dataset:** `enhanced_cps_2024.h5` (Enhanced CPS), years 2024–2028.

### Inspecting available variables

Always inspect the dataset to find available variable names — never guess:

```python
from policyengine.tax_benefit_models.uk import ensure_datasets

uk = ensure_datasets(years=[2026], data_folder="./data")
d = uk["enhanced_frs_2023_24_2026"]

# Input variables (present in raw data)
print("household:", list(d.data.household.columns))
print("person:   ", list(d.data.person.columns))
print("benunit:  ", list(d.data.benunit.columns))
```

**Input variables** are what's in the raw survey data — demographics, reported incomes, consumption, wealth, flags.

**Computed variables** (`household_net_income`, `income_tax`, `universal_credit`, etc.) are **not** in the raw dataset — they are calculated by the simulation. To see what's available after running:

```python
from policyengine.core import Simulation
from policyengine.tax_benefit_models.uk import uk_latest

sim = Simulation(dataset=d, tax_benefit_model_version=uk_latest)
sim.run()

print("household (post-sim):", list(sim.output_dataset.data.household.columns))
print("person (post-sim):   ", list(sim.output_dataset.data.person.columns))
```

The computed variables available are defined by `uk_latest.entity_variables` — inspect this to see the full list without running a simulation:

```python
from policyengine.tax_benefit_models.uk import uk_latest
print(uk_latest.entity_variables)  # dict: entity → [variable names]
```

## For Analysts: Core Concepts

When running simulations with policyengine.py (the microsimulation package, not the API client), you work with three key components:

1. **`Simulation.ensure()`** - Smart caching to avoid redundant computation
2. **`simulation.output_dataset.data`** - Accessing calculated results
3. **`map_to_entity()`** - Converting data between entity levels (person ↔ household)

**Note:** This is for microsimulation with policyengine.py, not the policyengine Python API client (which uses `Simulation(situation=...)`).

## Simulation Lifecycle

### The Four Methods

```python
from policyengine.core import Simulation
from policyengine.tax_benefit_models.uk import uk_latest

simulation = Simulation(
    dataset=dataset,
    tax_benefit_model_version=uk_latest,
)

# Method 1: Always run (no caching)
simulation.run()

# Method 2: Run only if needed (recommended)
simulation.ensure()

# Method 3: Save results to disk
simulation.save()

# Method 4: Load results from disk
simulation.load()
```

### When to Use Each

- **`run()`**: Use when you need fresh results or parameters changed
- **`ensure()`**: Use for iterative development (checks cache → disk → run)
- **`save()`**: Use to persist large simulation results
- **`load()`**: Use to resume from previous session

### How ensure() Works

```python
def ensure(self):
    # 1. Check in-memory LRU cache (100 simulations)
    cached = _cache.get(self.id)
    if cached:
        self.output_dataset = cached.output_dataset
        return

    # 2. Try loading from disk
    try:
        self.tax_benefit_model_version.load(self)
    except Exception:
        # 3. Only run if both cache and disk fail
        self.run()
        self.save()

    # 4. Add to cache for next ensure() call
    _cache.add(self.id, self)
```

**Performance impact:**
- First call: Full simulation runtime (seconds to minutes)
- Same session: Instant (in-memory cache)
- New session: Fast (disk load, no recomputation)

### Example: Reusing Baseline Across Reforms

```python
# Run baseline once
baseline = Simulation(dataset=dataset, tax_benefit_model_version=uk_latest)
baseline.ensure()  # First call: runs simulation
baseline.save()    # Persist to disk

# Test multiple reforms
for reform in [reform1, reform2, reform3]:
    baseline.ensure()  # Instant from cache!

    reform_sim = Simulation(
        dataset=dataset,
        tax_benefit_model_version=uk_latest,
        policy=reform
    )
    reform_sim.run()  # Only reform needs to run

    # Compare results...
```

## Accessing Results: output_dataset.data

After running a simulation, all calculated variables are in `simulation.output_dataset.data`.

### Structure (UK Example)

```python
simulation.run()

# Access output container
output = simulation.output_dataset.data

# Entity-level MicroDataFrames
output.person      # Person-level results
output.benunit     # Benefit unit results
output.household   # Household-level results
```

### US Entity Structure

```python
# US has more entities
output.person
output.tax_unit       # Federal tax filing unit
output.spm_unit       # Supplemental Poverty Measure unit
output.family         # Census family definition
output.marital_unit   # Married couple or single
output.household
```

### Available Variables

Each dataframe contains **input variables** + **calculated variables**:

```python
# Person-level (UK)
print(output.person.columns)
# ['person_i