Claude
Skills
Sign in
Back

policyengine-code-style

Included with Lifetime
$97 forever

PolicyEngine code writing style guide - formula optimization, direct returns, eliminating unnecessary variables

Writing & Docs

What this skill does


# PolicyEngine Code Writing Style Guide

Essential patterns for writing clean, efficient PolicyEngine formulas.

## Core Principles

1. **Eliminate unnecessary intermediate variables**
2. **Use direct parameter/variable access**
3. **Return directly when possible**
4. **Combine boolean logic**
5. **Use correct period access** (period vs period.this_year)
6. **NO hardcoded values** - use parameters or constants

---

## Pattern 1: Direct Parameter Access

### ❌ Bad - Unnecessary intermediate variable

```python
def formula(spm_unit, period, parameters):
    countable = spm_unit("tn_tanf_countable_resources", period)
    p = parameters(period).gov.states.tn.dhs.tanf.resource_limit
    resource_limit = p.amount  # ❌ Unnecessary
    return countable <= resource_limit
```

### ✅ Good - Direct access

```python
def formula(spm_unit, period, parameters):
    countable = spm_unit("tn_tanf_countable_resources", period)
    p = parameters(period).gov.states.tn.dhs.tanf.resource_limit
    return countable <= p.amount
```

---

## Pattern 2: Direct Return

### ❌ Bad - Unnecessary result variable

```python
def formula(spm_unit, period, parameters):
    assets = spm_unit("spm_unit_assets", period.this_year)
    p = parameters(period).gov.states.tn.dhs.tanf.resource_limit
    vehicle_exemption = p.vehicle_exemption  # ❌ Unnecessary
    countable = max_(assets - vehicle_exemption, 0)  # ❌ Unnecessary
    return countable
```

### ✅ Good - Direct return

```python
def formula(spm_unit, period, parameters):
    assets = spm_unit("spm_unit_assets", period.this_year)
    p = parameters(period).gov.states.tn.dhs.tanf.resource_limit
    return max_(assets - p.vehicle_exemption, 0)
```

---

## Pattern 3: Combined Boolean Logic

### ❌ Bad - Too many intermediate booleans

```python
def formula(spm_unit, period, parameters):
    person = spm_unit.members
    age = person("age", period.this_year)
    is_disabled = person("is_disabled", period.this_year)

    caretaker_is_60_or_older = spm_unit.any(age >= 60)  # ❌ Unnecessary
    caretaker_is_disabled = spm_unit.any(is_disabled)   # ❌ Unnecessary
    eligible = caretaker_is_60_or_older | caretaker_is_disabled  # ❌ Unnecessary

    return eligible
```

### ✅ Good - Combined logic

```python
def formula(spm_unit, period, parameters):
    person = spm_unit.members
    age = person("age", period.this_year)
    is_disabled = person("is_disabled", period.this_year)

    return spm_unit.any((age >= 60) | is_disabled)
```

---

## Pattern 4: Period Access - period vs period.this_year

```python
# ❌ Bad
age = person("age", period)  # Gives age/12
monthly_income = person("employment_income", period.this_year) / MONTHS_IN_YEAR  # Redundant

# ✅ Good
age = person("age", period.this_year)  # Gets actual age
monthly_income = person("employment_income", period)  # Auto-converts to monthly
```

- Income/flows → Use `period` (want monthly from annual)
- Age/assets/counts/booleans → Use `period.this_year` (don't divide by 12)

See **policyengine-period-patterns** skill for the full decision tree and more examples.

---

## Pattern 5: No Hardcoded Values

### ❌ Bad - Hardcoded numbers

```python
def formula(spm_unit, period, parameters):
    size = spm_unit.nb_persons()
    capped_size = min_(size, 10)  # ❌ Hardcoded

    age = person("age", period.this_year)
    income = person("income", period) / 12  # ❌ Use MONTHS_IN_YEAR

    # ❌ Hardcoded thresholds
    if age >= 18 and age <= 65 and income < 2000:
        return True
```

### ✅ Good - Parameterized

```python
def formula(spm_unit, period, parameters):
    p = parameters(period).gov.program
    capped_size = min_(spm_unit.nb_persons(), p.max_unit_size)  # ✅

    age = person("age", period.this_year)
    monthly_income = person("income", period)  # ✅ Auto-converts (no manual /12)

    age_eligible = (age >= p.age_min) & (age <= p.age_max)  # ✅
    income_eligible = monthly_income < p.income_threshold  # ✅

    return age_eligible & income_eligible
```

### Use Framework Constants, Not Custom Parameters

PolicyEngine provides constants for universal conversion factors. Don't create parameters for these:

```python
# ❌ BAD — created a weeks_per_month.yaml parameter:
weekly_subsidy * p.weeks_per_month

# ✅ GOOD — use framework constants:
weekly_subsidy * (WEEKS_IN_YEAR / MONTHS_IN_YEAR)
```

Available constants: `MONTHS_IN_YEAR` (12), `WEEKS_IN_YEAR` (52). Derive others from these.

**CRITICAL: `WEEKS_IN_YEAR` is the integer 52, not 52.1429.** Never use `WEEKS_IN_YEAR * 7` for days per year — that gives 364, not 365. Use the literal `365` for days-per-year calculations.

**Regulatory conversion factors take precedence.** When a regulation cites a specific factor (e.g., "multiply by 4.3"), use that exact value even if a framework constant is close (`WEEKS_IN_YEAR / MONTHS_IN_YEAR` ≈ 4.333). The discrepancy may affect benefit amounts at boundary conditions.

---

## Pattern 6: Streamline Variable Access

### ❌ Bad - Redundant steps

```python
def formula(spm_unit, period, parameters):
    unit_size = spm_unit.nb_persons()  # ❌ Unnecessary
    max_size = 10  # ❌ Hardcoded
    capped_size = min_(unit_size, max_size)

    p = parameters(period).gov.states.tn.dhs.tanf.benefit
    spa = p.standard_payment_amount[capped_size]  # ❌ Unnecessary
    dgpa = p.differential_grant_payment_amount[capped_size]  # ❌ Unnecessary

    eligible = spm_unit("eligible_for_dgpa", period)
    return where(eligible, dgpa, spa)
```

### ✅ Good - Streamlined

```python
def formula(spm_unit, period, parameters):
    p = parameters(period).gov.states.tn.dhs.tanf.benefit
    capped_size = min_(spm_unit.nb_persons(), p.max_unit_size)
    eligible = spm_unit("eligible_for_dgpa", period)

    return where(
        eligible,
        p.differential_grant_payment_amount[capped_size],
        p.standard_payment_amount[capped_size]
    )
```

---

## When to Keep Intermediate Variables

### ✅ Keep when value is used multiple times

```python
def formula(tax_unit, period, parameters):
    p = parameters(period).gov.irs.credits
    filing_status = tax_unit("filing_status", period)

    # ✅ Used multiple times - keep as variable
    threshold = p.phase_out.start[filing_status]

    income = tax_unit("adjusted_gross_income", period)
    excess = max_(0, income - threshold)
    reduction = (excess / p.phase_out.width) * threshold

    return max_(0, threshold - reduction)
```

### ✅ Keep when calculation is complex

```python
def formula(spm_unit, period, parameters):
    p = parameters(period).gov.program
    gross_earned = spm_unit("gross_earned_income", period)

    # ✅ Complex multi-step calculation - break it down
    work_expense_deduction = min_(gross_earned * p.work_expense_rate, p.work_expense_max)
    after_work_expense = gross_earned - work_expense_deduction

    earned_disregard = after_work_expense * p.earned_disregard_rate
    countable_earned = after_work_expense - earned_disregard

    dependent_care = spm_unit("dependent_care_expenses", period)

    return max_(0, countable_earned - dependent_care)
```

### ✅ Break out complex expressions inside function calls

**Don't inline complex calculations inside `where()`, `max_()`, or other function calls - give them descriptive names.**

```python
# ❌ BAD - Complex expression inlined in where()
return where(
    above_trigger,
    reduced_payment,
    max_(maximum_benefit - countable_income, 0),  # Hard to read
)

# ✅ GOOD - Break out into named variable
standard_payment = max_(maximum_benefit - countable_income, 0)
return where(
    above_trigger,
    reduced_payment,
    standard_payment,  # Clear what this represents
)
```

**Another example:**

```python
# ❌ BAD - Multiple complex inlined expressions
return where(
    income > add(spm_unit, period, ["earned", "unearned"]) * p.rate,
    max_(benefit - (income * p.reduction_rate), 0),
    benefit,
)

# ✅ GOOD - Named variables explain the logic
gross_income = add(spm_unit, period, ["earned", "unearned"])
income_threshold 

Related in Writing & Docs