如何在Python的patsy中实现自由度为2的自然样条？

背景 #

R语言的ns()函数支持直接设置自由度为2的自然样条，示例代码：

x <- 0:10
y <- x**2
lm(y ~ ns(x, df=2))

但Python中patsy库的cr()函数要求自由度≥3，无法直接实现相同效果。

解决方案 #

自由度为2的自然样条本质上等价于线性拟合（自然样条在df=2时，约束边界二阶导数为0，最终拟合结果就是一条直线），因此有两种实现方式：

方法1：直接用线性项替代（最简方案）

直接使用x的线性项即可完全匹配R中ns(x, df=2)的拟合结果：

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

# 构造数据集
df = pd.DataFrame({
    "x": np.arange(11),
    "y": np.arange(11)**2
})

# 公式直接使用线性项
formula = "y ~ x"
reg = smf.ols(formula, data=df).fit()
print(reg.summary())

方法2：自定义patsy转换函数（保持样条写法）

如果需要在公式中保留类似R的样条函数调用形式，可以自定义一个ns2转换函数，手动构造自由度为2的自然样条基：

import numpy as np
import pandas as pd
import statsmodels.formula.api as smf
from patsy import user_transform
from scipy.interpolate import BSpline

@user_transform
def ns2(x):
    # 构造自然样条基：df=2对应一次多项式，满足边界二阶导数为0的约束
    x_min, x_max = x.min(), x.max()
    # 自然样条节点：边界节点重复3次以满足边界条件
    knots = np.array([x_min, x_min, x_min, x_max, x_max, x_max])
    # 生成自由度为2的样条基
    basis = BSpline(knots, np.eye(2), k=1)(x)
    return basis

# 在公式中使用自定义的ns2函数
formula = "y ~ ns2(x)"
reg = smf.ols(formula, data=df).fit()
print(reg.summary())

说明 #

两种方法的拟合结果完全一致，都等价于R中ns(x, df=2)的输出。如果仅需实现功能，推荐方法1；如果需要统一代码风格（保持样条函数调用形式），可以使用方法2。

内容的提问来源于stack exchange，提问作者EuginePickett