238 lines
7.4 KiB
Python
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)
|