Files
2026-07-13 13:22:34 +08:00

238 lines
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)