Interactions and transforms in Python

The formula language covers products, powers and functions inline.

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

rng = np.random.default_rng(4)
frame = pd.DataFrame({
    "practice": rng.uniform(1, 60, size=120),
    "editor": rng.integers(0, 2, size=120),
})
frame["tpm"] = (70 + 0.5 * frame["practice"]
                + 6 * frame["editor"]
                + 0.1 * frame["practice"] * frame["editor"]
                + rng.normal(0, 3, size=120))

model = smf.ols("tpm ~ practice * editor", data=frame).fit()
print(sorted(model.params.index.tolist()))
print(model.params.round(3).to_dict())

curved = smf.ols("tpm ~ practice + I(practice ** 2)", data=frame).fit()
print(round(curved.rsquared, 4), round(model.rsquared, 4))

How it works

1. a:b is the interaction alone, a*b is both plus the interaction.
2. I(x ** 2) protects arithmetic from the formula parser.
3. np.log(x) works directly inside a formula.

Keywords and builtins used here

• as
• print
• round
• sorted

The run, in numbers

Lines

20

Characters to type

633

Tokens

212

Three-star pace

110 tpm

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

Type this snippet

Step 2 of 3 in Model design, step 5 of 19 in Statistics with statsmodels.

← Previous Next →