Factor Variables

Categorical regressors and interactions

Overview

This tutorial shows how to include categorical variables as regressors using the i() syntax.

The i() Syntax

Syntax What it does
y ~ x \| region Absorbs region as fixed effect (coefficients not reported)
y ~ x + i(region) Includes region dummies as regressors (coefficients reported)
i(var, ref='value') Specify reference category
treatment:i(var) Interaction for heterogeneous effects

Example Data

import numpy as np
import polars as pl
from leanfe import leanfe

np.random.seed(42)
n_obs = 50_000

regions = ["North", "South", "East", "West"]
region = np.random.choice(regions, n_obs)

# True effects: North=0 (ref), South=1.5, East=2.0, West=-0.5
region_effects = {"North": 0, "South": 1.5, "East": 2.0, "West": -0.5}
region_effect = np.array([region_effects[r] for r in region])

df = pl.DataFrame({
    "firm_id": np.random.randint(1, 200, n_obs),
    "year": np.random.randint(2015, 2025, n_obs),
    "region": region,
    "treatment": np.random.binomial(1, 0.3, n_obs).astype(float),
    "x1": np.random.normal(0, 1, n_obs),
})

firm_fe = np.random.normal(0, 0.5, 201)[df["firm_id"].to_numpy()]
year_fe = np.random.normal(0, 0.3, 11)[df["year"].to_numpy() - 2015]

df = df.with_columns(
    y = (pl.lit(2.5) * pl.col("treatment") 
         + pl.lit(1.5) * pl.col("x1")
         + pl.lit(region_effect)
         + pl.lit(firm_fe) + pl.lit(year_fe)
         + pl.lit(np.random.normal(0, 1, n_obs)))
)

print(f"Data: {df.shape[0]:,} obs")
print(f"Regions: {df['region'].unique().to_list()}")
Data: 50,000 obs
Regions: ['South', 'North', 'East', 'West']

Basic Factor Variable

result = leanfe(
    formula="y ~ treatment + x1 + i(region) | firm_id + year",
    data=df,
    vcov="iid"
)

print("Coefficients:")
for var, coef in result['coefficients'].items():
    print(f"  {var}: {coef:.4f} (SE: {result['std_errors'][var]:.4f})")
Coefficients:
  treatment: 2.4973 (SE: 0.0099)
  x1: 1.4950 (SE: 0.0045)
  region_North: -1.9610 (SE: 0.0128)
  region_South: -0.4613 (SE: 0.0128)
  region_West: -2.4438 (SE: 0.0128)

Custom Reference Category

# Default (North as reference - first alphabetically)
result_north = leanfe(
    data=df,
    formula="y ~ treatment + x1 + i(region, ref='North') | firm_id + year",
    vcov="iid"
)

# South as reference
result_south = leanfe(
    data=df,
    formula="y ~ treatment + x1 + i(region, ref='South') | firm_id + year",
    vcov="iid"
)

print("North as reference:")
for var, coef in result_north['coefficients'].items():
    if 'region' in var:
        print(f"  {var}: {coef:.4f}")

print("\nSouth as reference:")
for var, coef in result_south['coefficients'].items():
    if 'region' in var:
        print(f"  {var}: {coef:.4f}")
North as reference:
  region_East: 2.0151
  region_South: 1.5076
  region_West: -0.4758

South as reference:
  region_East: 0.5434
  region_North: -1.4636
  region_West: -1.9465

Interactions: Heterogeneous Effects

# Generate data with heterogeneous treatment effects
np.random.seed(123)
treatment = np.random.binomial(1, 0.3, n_obs).astype(float)

# True effects: North=2.0, South=3.0, East=1.0, West=2.5
het_effects = {"North": 2.0, "South": 3.0, "East": 1.0, "West": 2.5}
treatment_effect = np.array([het_effects[r] for r in region]) * treatment

df_het = df.with_columns([
    pl.Series("treatment", treatment),
    pl.Series("y", treatment_effect + 1.5 * df["x1"].to_numpy() + firm_fe + year_fe + np.random.normal(0, 1, n_obs))
])

result_het = leanfe(
    formula="y ~ treatment:i(region) + x1 | firm_id + year",
    data=df_het,
    vcov="iid"
)

print("Heterogeneous treatment effects:")
print("(True: North=2.0, South=3.0, East=1.0, West=2.5)")
for var, coef in result_het['coefficients'].items():
    if 'treatment' in var and 'region' in var:
        print(f"  {var}: {coef:.4f}")
Heterogeneous treatment effects:
(True: North=2.0, South=3.0, East=1.0, West=2.5)

Quick Reference

Syntax Use Case
i(var) Categorical variable as regressor
i(var, ref='value') Specify reference category
treatment:i(var) Heterogeneous effects by category

Next Steps