Categorical predictors in Python

C() makes a column categorical, and the formula builds the dummies.

import numpy as np
import pandas as pd
import statsmodels.formula.api as smf

rng = np.random.default_rng(3)
frame = pd.DataFrame({
    "lang": np.repeat(["python", "rust", "css"], 30),
    "practice": rng.integers(1, 60, size=90),
})
offset = frame["lang"].map({"python": 10.0, "rust": 5.0, "css": 0.0})
frame["tpm"] = 80 + offset + 0.3 * frame["practice"] + rng.normal(0, 3, 90)

model = smf.ols(
    "tpm ~ C(lang, Treatment(reference='css')) + practice",
    data=frame).fit()
for name, value in model.params.round(2).items():
    print(f"{name[:44]:<46}{value:>8}")

How it works

One level becomes the reference; the rest get coefficients.
Treatment(reference=...) chooses which level that is.
The names in the output say which contrast each row is.

Keywords and builtins used here

as
for
print

The run, in numbers

Lines: 17
Characters to type: 550
Tokens: 198
Three-star pace: 110 tpm

At the three-star pace of 110 tokens a minute, this run takes about 108 seconds.

Type this snippet

Step 1 of 3 in Model design, step 4 of 19 in Statistics with statsmodels.

← Previous Next →