项目文件夹

文件
2026-07-13 13:22:34 +08:00

238 行
7.4 KiB
Python

from typing import Any, NamedTuple
import numpy as np
import pandas as pd
import statsmodels.api as sm
from scipy.linalg import toeplitz
from statsmodels.tsa.arima.model import ARIMA
from statsmodels.tsa.arima_process import arma_generate_sample
from mlflow.models import ModelSignature
from mlflow.types.schema import Schema, TensorSpec
class ModelWithResults(NamedTuple):
model: Any
alg: Any
inference_dataframe: Any
"""
Fixtures for a number of models available in statsmodels
https://www.statsmodels.org/dev/api.html
"""
def ols_model(**kwargs):
# Ordinary Least Squares (OLS)
np.random.seed(9876789)
nsamples = 100
x = np.linspace(0, 10, 100)
X = np.column_stack((x, x**2))
beta = np.array([1, 0.1, 10])
e = np.random.normal(size=nsamples)
X = sm.add_constant(X)
y = np.dot(X, beta) + e
ols = sm.OLS(y, X)
model = ols.fit(**kwargs)
return ModelWithResults(model=model, alg=ols, inference_dataframe=X)
def ols_model_signature():
return ModelSignature(
inputs=Schema([TensorSpec(np.dtype("float64"), (-1, 3))]),
outputs=Schema([TensorSpec(np.dtype("float64"), (-1,))]),
)
def failing_logit_model():
X = pd.DataFrame(
{
"x0": np.array([2.0, 3.0, 1.0, 2.0, 20.0, 30.0, 10.0, 20.0]),
"x1": np.array([2.0, 3.0, 1.0, 2.0, 20.0, 30.0, 10.0, 20.0]),
},
columns=["x0", "x1"],
)
y = np.array([0, 0, 0, 0, 1, 1, 1, 1])
# building the model and fitting the data
log_reg = sm.Logit(y, X)
model = log_reg.fit()
return ModelWithResults(model=model, alg=log_reg, inference_dataframe=X)
def get_dataset(name):
dataset_module = getattr(sm.datasets, name)
data = dataset_module.load()
data.exog = np.asarray(data.exog)
data.endog = np.asarray(data.endog)
return data
def gls_model():
# Generalized Least Squares (GLS)
data = get_dataset("longley")
data.exog = sm.add_constant(data.exog)
ols_resid = sm.OLS(data.endog, data.exog).fit().resid
res_fit = sm.OLS(ols_resid[1:], ols_resid[:-1]).fit()
rho = res_fit.params
order = toeplitz(np.arange(16))
sigma = rho**order
gls = sm.GLS(data.endog, data.exog, sigma=sigma)
model = gls.fit()
return ModelWithResults(model=model, alg=gls, inference_dataframe=data.exog)
def glsar_model():
# Generalized Least Squares with AR covariance structure
X = range(1, 8)
X = sm.add_constant(X)
Y = [1, 3, 4, 5, 8, 10, 9]
glsar = sm.GLSAR(Y, X, rho=2)
model = glsar.fit()
return ModelWithResults(model=model, alg=glsar, inference_dataframe=X)
def wls_model():
# Weighted Least Squares
Y = [1, 3, 4, 5, 2, 3, 4]
X = range(1, 8)
X = sm.add_constant(X)
wls = sm.WLS(Y, X, weights=list(range(1, 8)))
model = wls.fit()
return ModelWithResults(model=model, alg=wls, inference_dataframe=X)
def recursivels_model():
# Recursive Least Squares
dta = sm.datasets.copper.load_pandas().data
dta.index = pd.date_range("1951-01-01", "1975-01-01", freq="AS")
endog = dta.WORLDCONSUMPTION
# To the regressors in the dataset, we add a column of ones for an intercept
exog = sm.add_constant(dta[["COPPERPRICE", "INCOMEINDEX", "ALUMPRICE", "INVENTORYINDEX"]])
rls = sm.RecursiveLS(endog, exog)
model = rls.fit()
inference_dataframe = pd.DataFrame([["1951-01-01", "1975-01-01"]], columns=["start", "end"])
return ModelWithResults(model=model, alg=rls, inference_dataframe=inference_dataframe)
def rolling_ols_model():
# Rolling Ordinary Least Squares (Rolling OLS)
from statsmodels.regression.rolling import RollingOLS
data = get_dataset("longley")
exog = sm.add_constant(data.exog, prepend=False)
rolling_ols = RollingOLS(data.endog, exog)
model = rolling_ols.fit(reset=50)
return ModelWithResults(model=model, alg=rolling_ols, inference_dataframe=exog)
def rolling_wls_model():
# Rolling Weighted Least Squares (Rolling WLS)
from statsmodels.regression.rolling import RollingWLS
data = get_dataset("longley")
exog = sm.add_constant(data.exog, prepend=False)
rolling_wls = RollingWLS(data.endog, exog)
model = rolling_wls.fit(reset=50)
return ModelWithResults(model=model, alg=rolling_wls, inference_dataframe=exog)
def gee_model():
# Example taken from
# https://www.statsmodels.org/devel/examples/notebooks/generated/gee_nested_simulation.html
np.random.seed(9876789)
p = 5
groups_var = 1
level1_var = 2
level2_var = 3
resid_var = 4
n_groups = 100
group_size = 20
level1_size = 10
level2_size = 5
n = n_groups * group_size * level1_size * level2_size
xmat = np.random.normal(size=(n, p))
# Construct labels showing which group each observation belongs to at each level.
groups_ix = np.kron(np.arange(n // group_size), np.ones(group_size)).astype(int)
level1_ix = np.kron(np.arange(n // level1_size), np.ones(level1_size)).astype(int)
level2_ix = np.kron(np.arange(n // level2_size), np.ones(level2_size)).astype(int)
# Simulate the random effects.
groups_re = np.sqrt(groups_var) * np.random.normal(size=n // group_size)
level1_re = np.sqrt(level1_var) * np.random.normal(size=n // level1_size)
level2_re = np.sqrt(level2_var) * np.random.normal(size=n // level2_size)
# Simulate the response variable
y = groups_re[groups_ix] + level1_re[level1_ix] + level2_re[level2_ix]
y += np.sqrt(resid_var) * np.random.normal(size=n)
# Put everything into a dataframe.
df = pd.DataFrame(xmat, columns=[f"x{j}" for j in range(p)])
df["y"] = y + xmat[:, 0] - xmat[:, 3]
df["groups_ix"] = groups_ix
df["level1_ix"] = level1_ix
df["level2_ix"] = level2_ix
# Fit the model
cs = sm.cov_struct.Nested()
dep_fml = "0 + level1_ix + level2_ix"
gee = sm.GEE.from_formula(
"y ~ x0 + x1 + x2 + x3 + x4", cov_struct=cs, dep_data=dep_fml, groups="groups_ix", data=df
)
model = gee.fit()
return ModelWithResults(model=model, alg=gee, inference_dataframe=df)
def glm_model():
# Generalized Linear Model (GLM)
data = get_dataset("scotland")
data.exog = sm.add_constant(data.exog)
glm = sm.GLM(data.endog, data.exog, family=sm.families.Gamma())
model = glm.fit()
return ModelWithResults(model=model, alg=glm, inference_dataframe=data.exog)
def glmgam_model():
# Generalized Additive Model (GAM)
from statsmodels.gam.tests.test_penalized import df_autos
x_spline = df_autos[["weight", "hp"]]
bs = sm.gam.BSplines(x_spline, df=[12, 10], degree=[3, 3])
alpha = np.array([21833888.8, 6460.38479])
gam_bs = sm.GLMGam.from_formula(
"city_mpg ~ fuel + drive", data=df_autos, smoother=bs, alpha=alpha
)
model = gam_bs.fit()
return ModelWithResults(model=model, alg=gam_bs, inference_dataframe=df_autos)
def arma_model():
# Autoregressive Moving Average (ARMA)
np.random.seed(12345)
arparams = np.array([1, -0.75, 0.25])
maparams = np.array([1, 0.65, 0.35])
nobs = 250
y = arma_generate_sample(arparams, maparams, nobs)
dates = pd.date_range("1980-1-1", freq="M", periods=nobs)
y = pd.Series(y, index=dates)
arima = ARIMA(y, order=(2, 0, 2), trend="n")
model = arima.fit()
inference_dataframe = pd.DataFrame([["1999-06-30", "2001-05-31"]], columns=["start", "end"])
return ModelWithResults(model=model, alg=arima, inference_dataframe=inference_dataframe)