policyengine-vectorization

What this skill does

# PolicyEngine Vectorization Patterns

Critical patterns for vectorized operations in PolicyEngine. Scalar logic with arrays will crash the microsimulation.

## The Golden Rule

**PolicyEngine processes multiple households simultaneously using NumPy arrays. NEVER use if-elif-else with entity data.**

---

## 1. Critical: What Will Crash

### ❌ NEVER: if-elif-else with Arrays

```python
# THIS WILL CRASH - household data is an array
def formula(household, period, parameters):
    income = household("income", period)
    if income > 1000:  # ❌ CRASH: "truth value of array is ambiguous"
        return 500
    else:
        return 100
```

### ✅ ALWAYS: Vectorized Operations

```python
# CORRECT - works with arrays
def formula(household, period, parameters):
    income = household("income", period)
    return where(income > 1000, 500, 100)  # ✅ Vectorized
```

---

## 2. Common Vectorization Patterns

### Pattern 1: Simple Conditions → `where()`

```python
# Instead of if-else
❌ if age >= 65:
    amount = senior_amount
else:
    amount = regular_amount

✅ amount = where(age >= 65, senior_amount, regular_amount)
```

### Pattern 2: Multiple Conditions → `select()`

```python
# Instead of if-elif-else
❌ if age < 18:
    benefit = child_amount
elif age >= 65:
    benefit = senior_amount
else:
    benefit = adult_amount

✅ benefit = select(
    [age < 18, age >= 65],
    [child_amount, senior_amount],
    default=adult_amount
)
```

### Pattern 2b: Enum Dispatch → List ALL Conditions

When dispatching on an Enum variable, list all values explicitly in `select()`. Set the `default` to match the Enum's `default_value` so unexpected values get consistent behavior.

```python
# ❌ BAD — hides one option in default, unclear which is the fallback
provider_type = person("ri_ccap_provider_type", period)
types = provider_type.possible_values
return select(
    [
        provider_type == types.LICENSED_CENTER,
        provider_type == types.LICENSED_FAMILY,
    ],
    [center_rate, family_rate],
    default=exempt_rate,  # Why exempt as fallback?
)

# ✅ GOOD — all options listed, default matches Enum default_value
provider_type = person("ri_ccap_provider_type", period)
types = provider_type.possible_values
return select(
    [
        provider_type == types.LICENSED_CENTER,
        provider_type == types.LICENSED_FAMILY,
        provider_type == types.LICENSE_EXEMPT,
    ],
    [center_rate, family_rate, exempt_rate],
    default=center_rate,  # Matches default_value = LICENSED_CENTER
)
```

### Pattern 3: Boolean Operations

```python
# Combining conditions
eligible = (age >= 18) & (income < threshold)  # Use & not 'and'
eligible = (is_disabled | is_elderly)          # Use | not 'or'
eligible = ~is_excluded                        # Use ~ not 'not'
```

### Pattern 4: Clipping Values

```python
# Instead of if for bounds checking
❌ if amount < 0:
    amount = 0
elif amount > maximum:
    amount = maximum

✅ amount = clip(amount, 0, maximum)
# Or: amount = max_(0, min_(amount, maximum))
```

### Pattern 5: Flooring Subtraction Results (CRITICAL)

When subtracting values and wanting to floor at zero, you must wrap the **entire subtraction** in `max_()`:

```python
# Common scenario: income after deductions/losses
❌ WRONG - Creates phantom negative values:
income = max_(income, 0) - capital_loss  # If capital_loss > income, result is negative!

✅ CORRECT - Properly floors at zero:
income = max_(income - capital_loss, 0)  # Entire subtraction floored

# Real example from MT income tax bug:
❌ WRONG - Tax on phantom negative income:
def formula(tax_unit, period, parameters):
    income = tax_unit("adjusted_gross_income", period)
    capital_gains = tax_unit("capital_gains", period)

    # BUG: If capital_gains is negative (loss), this creates negative income
    # But max_() only floors income, not the result
    regular_income = max_(income, 0) - capital_gains
    return calculate_tax(regular_income)  # Tax on negative number!

✅ CORRECT - No phantom income:
def formula(tax_unit, period, parameters):
    income = tax_unit("adjusted_gross_income", period)
    capital_gains = tax_unit("capital_gains", period)

    # Properly floors the entire result
    regular_income = max_(income - capital_gains, 0)
    return calculate_tax(regular_income)  # Never negative
```

**Why this matters:**
- If `capital_gains = -3000` (loss), then `income - capital_gains = income + 3000`
- The wrong pattern `max_(income, 0) - capital_gains` allows the subtraction to make the result negative
- This creates "phantom income" where none exists, leading to incorrect tax calculations

**Rule:** When the formula is `A - B` and you want the result floored at zero, use `max_(A - B, 0)`, NOT `max_(A, 0) - B`.

---

## 3. When if-else IS Acceptable

### ✅ OK: Parameter-Only Conditions

```python
# OK - parameters are scalars, not arrays
def formula(entity, period, parameters):
    p = parameters(period).gov.program

    # This is fine - p.enabled is a scalar boolean
    if p.enabled:
        base = p.base_amount
    else:
        base = 0

    # But must vectorize when using entity data
    income = entity("income", period)
    return where(income < p.threshold, base, 0)
```

### ✅ OK: Control Flow (Not Data)

```python
# OK - controlling which calculation to use
def formula(entity, period, parameters):
    year = period.start.year

    if year >= 2024:
        # Use new formula (still vectorized)
        return entity("new_calculation", period)
    else:
        # Use old formula (still vectorized)
        return entity("old_calculation", period)
```

---

## 4. Common Vectorization Mistakes

### Mistake 1: Scalar Comparison with Array

```python
❌ WRONG:
if household("income", period) > 1000:
    # Error: truth value of array is ambiguous

✅ CORRECT:
income = household("income", period)
high_income = income > 1000  # Boolean array
benefit = where(high_income, low_benefit, high_benefit)
```

### Mistake 2: Using Python's and/or/not

```python
❌ WRONG:
eligible = is_elderly or is_disabled  # Python's 'or'

✅ CORRECT:
eligible = is_elderly | is_disabled   # NumPy's '|'
```

### Mistake 3: Nested if Statements

```python
❌ WRONG:
if eligible:
    if income < threshold:
        return full_benefit
    else:
        return partial_benefit
else:
    return 0

✅ CORRECT:
return where(
    eligible,
    where(income < threshold, full_benefit, partial_benefit),
    0
)
```

---

## 5. CRITICAL: Avoiding Divide-by-Zero Warnings

### The Problem with `where()` for Division

`where()` evaluates **BOTH branches** before selecting. This causes divide-by-zero warnings even when the zero case wouldn't be selected:

```python
# ❌ WRONG - causes divide-by-zero warning
proportion = where(
    total_income > 0,
    person_income / total_income,  # Still evaluated when total_income = 0!
    0,
)
```

### ✅ CORRECT: Use `np.divide` with `where` Parameter

```python
# ✅ CORRECT - only divides where mask is True
# The `out` parameter IS the default value - positions where mask=False keep this value
mask = total_income > 0
proportion = np.divide(
    person_income,
    total_income,
    out=np.zeros_like(person_income),  # Default to 0 where mask is False
    where=mask,
)
```

**How `out` works as the default:**
- `out=np.zeros_like(...)` → default is 0
- `out=np.ones_like(...)` → default is 1
- Positions where `where=False` keep their `out` value unchanged

### ✅ CORRECT: Alternative Mask Pattern

```python
# ✅ CORRECT - traditional mask assignment
proportion = np.zeros_like(total_income)
mask = total_income > 0
proportion[mask] = person_income[mask] / total_income[mask]
```

### Common Use Cases

**Proportional allocation (e.g., splitting deductions between spouses):**
```python
# Allocate proportionally by income
unit_income = tax_unit.sum(person_income)
mask = unit_income > 0
share = np.divide(
    person_income,
    unit_income,
    out=np.zeros_like(person_income),
    where=mask,
)
# Default share when unit has no incom