patsy_design_matrices_with_r_style_formulas_and_interactions.py

1import pandas as pd
2import numpy as np
3from patsy import dmatrices, dmatrix, DemoData
4
5# Create some demo data
6data = pd.DataFrame({
7    "x1": [1, 2, 3],
8    "x2": [4, 5, 6],
9    "y": [7, 8, 9],
10    "a": ["a1", "a2", "a1"],
11    "b": ["b1", "b1", "b2"]
12})
13
14# Use dmatrices to create design matrices for a linear model (y ~ x1 + x2)
15outcome, predictors = dmatrices("y ~ x1 + x2", data)
16
17# Use dmatrix to create a single design matrix with interactions and categorical data
18# x1:x2 is the interaction, a is a categorical variable automatically expanded
19mat = dmatrix("x1 + x2 + x1:x2 + a", data)
20
21print("Outcome matrix:\n", outcome)
22print("\nPredictors matrix:\n", predictors)
23print("\nComplex design matrix with categorical 'a':\n", mat)