Generalized Linear Models

!sudo /bin/bash -c "(source /venv/bin/activate; pip install --quiet jupyterlab-vim)"
!jupyter labextension enable

!sudo /bin/bash -c "(source /venv/bin/activate; pip install --quiet graphviz)"

!sudo /bin/bash -c "(source /venv/bin/activate; pip install --quiet dataframe_image)"

%load_ext autoreload
%autoreload 2


import arviz as az
import pandas as pd
import xarray as xr
import pymc as pm
import numpy as np
import seaborn as sns
import matplotlib.pyplot as plt

dir_name = "./L07_data"

!ls $dir_name

import msml610.tutorials.msml610_utils as ut

ut.config_notebook()

np.random.seed(1)

# Number of samples.
N = 100

# Parameters.
alpha_real = 2.5
beta_real = 0.9
sigma_eps_real = 0.5

# Generate data.
x = np.random.normal(10, 1, N)
y_real = alpha_real + beta_real * x

# Add noise.
eps_real = np.random.normal(0, sigma_eps_real, size=N)
y = y_real + eps_real

plt.scatter(x, y_real)
plt.scatter(x, y)

df = pd.DataFrame({"X": x, "Y": y})
sns.regplot(
    x="X",
    y="Y",
    data=df,
    scatter_kws={"color": "blue"},
    line_kws={"color": "red"},
)

with pm.Model() as model_g:
    alpha = pm.Normal("alpha", mu=0, sigma=10)
    beta = pm.Normal("beta", mu=0, sigma=1)
    sigma = pm.HalfCauchy("sigma", 5)
    #
    mu = pm.Deterministic("mu", alpha + beta * x)
    y_pred = pm.Normal("y_pred", mu=mu, sigma=sigma, observed=y)
    idata_g = pm.sample(2000, tune=1000)

pm.model_to_graphviz(model_g)

az.plot_trace(idata_g, var_names=["alpha", "beta", "sigma"])

az.summary(idata_g, var_names="alpha beta sigma".split(), kind="stats")

bikes = pd.read_csv(dir_name + "/bikes.csv")
bikes.plot(x="temperature", y="rented", figsize=(12, 3), kind="scatter")

bikes.head()

with pm.Model() as model_lb:
    alpha = pm.Normal("alpha", mu=0, sigma=100)
    beta = pm.Normal("beta", mu=0, sigma=10)
    sigma = pm.HalfCauchy("sigma", 10)
    #
    mu = pm.Deterministic("mu", alpha + beta * bikes.temperature)
    #
    y_pred = pm.Normal("y_pred", mu=mu, sigma=sigma, observed=bikes.rented)
    idata_lb = pm.sample()

# Plot all the vars, excluding mu.
az.plot_posterior(idata_lb, var_names=["~mu"])

# Sample from the posterior.
posterior = az.extract(idata_lb, num_samples=50)

# Create a data array of 50 equally-spaced points from min to max temperature along the `plot_id` axis.
x_plot = xr.DataArray(
    np.linspace(bikes.temperature.min(), bikes.temperature.max(), 50),
    dims="plot_id",
)

# Compute the expected value of the model for the points.
mean_line = posterior["alpha"].mean() + posterior["beta"].mean() * x_plot


# Compute 50 lines using the posterior.
lines = posterior["alpha"] + posterior["beta"] * x_plot

idata_lb_pp = pm.sample_posterior_predictive(idata_lb, model=model_lb)
mean_line = idata_lb.posterior["mu"].mean(("chain", "draw"))


def plot_data_and_model(bikes, idata):
    # Generate a vector with the temperatures and a bit of jitter.
    temperatures = np.random.normal(bikes.temperature.values, 0.01)

    # Sort in increasing order.
    idx = np.argsort(temperatures)

    # Sample the temperature intervals.
    x = np.linspace(temperatures.min(), temperatures.max(), 15)

    # Compute the quantiles, flattening over chain and draw.
    y_pred_q = idata.posterior_predictive["y_pred"].quantile(
        [0.03, 0.97, 0.25, 0.75], dim=["chain", "draw"]
    )
    # obj_to_str(y_pred_q, only_schema=False)

    from scipy.interpolate import PchipInterpolator

    y_hat_bounds = [
        PchipInterpolator(temperatures[idx], y_pred_q[i][idx])(x)
        for i in range(4)
    ]
    # print(y_hat_bounds)

    # Plot the data set.
    plt.plot(bikes.temperature, bikes.rented, "C2.", zorder=-3)
    # Plot the mean.
    plt.plot(bikes.temperature[idx], mean_line[idx], c="C0")

    # Plot the quantiles.
    lb, ub = y_hat_bounds[0], y_hat_bounds[1]
    plt.fill_between(x, lb, ub, color="C1", alpha=0.2)
    lb, ub = y_hat_bounds[2], y_hat_bounds[3]
    plt.fill_between(x, lb, ub, color="C1", alpha=0.2)

print(x_plot.shape)

print(mean_line.shape)

# - Plot the data
# zorder is to plot behind the line.
plt.plot(bikes.temperature, bikes.rented, "C2.", zorder=-3)

# - Plot the 50 models from the posterior.
# lines.T.values are the 50 lines.
lines_ = plt.plot(x_plot, lines.T.values, c="C1", alpha=0.2, label="lines")
# Remove the label for all the lines but the first one.
plt.setp(lines_[1:], label="_")

# Plot the mean line.
# plt.plot(x_plot, mean_line, c="C0", label="mean line");

plt.xlabel("temp")
plt.ylabel("rented bikes")

plt.legend()

plot_data_and_model(bikes, idata_lb_pp)

np.random.seed(42)
with pm.Model() as model_neg:
    alpha = pm.Normal("alpha", mu=0, sigma=100)
    beta = pm.Normal("beta", mu=0, sigma=10)
    # We use exp to have all positive numbers.
    mu = pm.Deterministic("mu", pm.math.exp(alpha + beta * bikes.temperature))
    # NegativeBinomial has an extra param alpha to control the variance.
    sigma = pm.HalfNormal("sigma", 10)
    y_pred = pm.NegativeBinomial(
        "y_pred", mu=mu, alpha=sigma, observed=bikes.rented
    )
    #
    idata_neg = pm.sample()
    idata_neg.extend(pm.sample_posterior_predictive(idata_neg))

pm.model_to_graphviz(model_neg)

az.plot_trace(idata_neg, var_names=["~mu"])

plot_data_and_model(bikes, idata_neg)

az.plot_ppc(idata_lb_pp, num_pp_samples=200, alpha=0.1, mean=False)

az.plot_ppc(idata_neg, num_pp_samples=200, alpha=0.1, mean=False)

ans = pd.read_csv(dir_name + "/anscombe_3.csv")
display(ans.head())

ans.plot("x", "y", kind="scatter")

import scipy

beta_c, alpha_c, *_ = scipy.stats.linregress(ans.x, ans.y)

_, ax = plt.subplots()
ax.plot(ans.x, (alpha_c + beta_c * ans.x), "C0:", label="non-robust")
ax.plot(ans.x, ans.y, "C0o")

ut.save_ax(ax, "Lesson07_Non_robust_regression1.png")

with pm.Model() as model_t:
    # Alpha is normal centered around the mean of the y data.
    alpha = pm.Normal("alpha", mu=ans.y.mean(), sigma=1)
    # Beta is a standard (0, 1).
    beta = pm.Normal("beta", mu=0, sigma=1)
    sigma = pm.HalfNormal("sigma", 5)
    # The exponential puts too much weight close to 0, so we shift by 1.
    nu_ = pm.Exponential("nu_", 1 / 29)
    nu = pm.Deterministic("nu", nu_ + 1)

    # Model.
    mu = pm.Deterministic("mu", alpha + beta * ans.x)
    y_pred = pm.StudentT("y_pred", mu=mu, sigma=sigma, nu=nu, observed=ans.y)
    idata_t = pm.sample(2000, tune=2000)

    idata_t.extend(pm.sample_posterior_predictive(idata_t))

# dot = pm.model_to_graphviz(model_t)
ut.save_dot(model_t, "Lesson07_Robust_regression_model")

var_names = "alpha beta sigma nu".split()
az.plot_trace(idata_t, var_names=var_names)
az.summary(idata_t, var_names=var_names, round_to=2, kind="stats")

_, ax = plt.subplots()

# Non-robust.
ax.plot(ans.x, (alpha_c + beta_c * ans.x), "C0:", label="non-robust")
ax.plot(ans.x, ans.y, "C0o")

# Robust.
alpha_m = idata_t.posterior["alpha"].mean(("chain", "draw"))
beta_m = idata_t.posterior["beta"].mean(("chain", "draw"))

x_plot = xr.DataArray(np.linspace(ans.x.min(), ans.x.max(), 50), dims="plot_id")
ax.plot(x_plot, alpha_m + beta_m * x_plot, c="C0", label="robust")
az.plot_hdi(ans.x, az.hdi(idata_t.posterior["mu"])["mu"].T, ax=ax)
ax.set_xlabel("x")
ax.set_ylabel("y", rotation=0)
ax.legend(loc=2)

ut.save_ax(ax, "Lesson07_Non_robust_regression2")

# ?pm.sample_posterior_predictive
# ?az.plot_ppc

# Posterior predictive check.
ppc = pm.sample_posterior_predictive(
    idata_t,
    # samples=200,
    model=model_t,
    random_seed=2,
)
az.plot_ppc(idata_t, mean=True, num_pp_samples=100)
plt.xlim(0, 20)

iris = pd.read_csv(dir_name + "/iris.csv")
iris.head()

ut.save_df(iris.head(), "Lesson07_Logistic_regression_df.png")

# Filter the dataframe keeping only 2 values for species.
df = iris.query("species == ('setosa', 'versicolor')")
df.head()

# Get the predicted variable.
y_0 = pd.Categorical(df["species"]).codes
y_0

# Get the sepal length as feature.
x_n = "sepal_length"
x_0 = df[x_n].values
x_c = x_0 - x_0.mean()

with pm.Model() as model_lrs:
    # Linear part.
    alpha = pm.Normal("alpha", mu=0, sigma=1)
    beta = pm.Normal("beta", mu=0, sigma=5)
    mu = alpha + x_c * beta
    # Sigmoid.
    theta = pm.Deterministic("theta", pm.math.sigmoid(mu))
    # Model.
    yl = pm.Bernoulli("yl", p=theta, observed=y_0)
    # Intercept?
    bd = pm.Deterministic("bd", -alpha / beta)
    #
    idata_lrs = pm.sample(random_seed=123)

ut.save_dot(model_lrs, "Lesson07_Logistic_regression_model.png")

var_names = ["~bd", "~theta"]
az.summary(idata_lrs, var_names=var_names, round_to=2, kind="stats")

ax = az.plot_trace(idata_lrs, var_names=var_names)

ut.save_fig(ax, "Lesson07_Logistic_regression_result.png")

posterior = idata_lrs.posterior
theta = posterior["theta"].mean(("chain", "draw"))
idx = np.argsort(x_c)
# print(posterior)
# print(idx)

# Plot the model.
_, ax = plt.subplots()
ax.plot(x_c[idx], theta[idx], color="C0", lw=2)

# Plot the decision.
ax.vlines(posterior["bd"].mean(("chain", "draw")), 0, 1, color="C2", zorder=0)
bd_hdi = az.hdi(posterior["bd"])
ax.fill_betweenx(
    [0, 1], bd_hdi["bd"][0], bd_hdi["bd"][1], color="C2", alpha=0.6, lw=0
)

# Plot the data.
ax.scatter(x_c, np.random.normal(y_0, 0.02), marker=".")

az.plot_hdi(x_c, posterior["theta"], color="C0", ax=ax)

ut.save_ax(ax, "Lesson07_Logistic_regression_result2.png")

# Load data.
data = pd.read_csv(dir_name + "/babies.csv")
data.columns = ["month", "length"]
data.plot.scatter("month", "length")
display(data.head())

ax = data.plot.scatter("month", "length")
ut.save_ax(ax, "Lesson07_Variable_variance_data.png")

with pm.Model() as model_vv:
    # Create a shared variable so that the data can change after the model is created.
    x_shared = pm.Data("x_shared", data.month.values.astype(float))
    # Linear model for the mean is a function of sqrt(x).
    alpha = pm.Normal("alpha", sigma=10)
    beta = pm.Normal("beta", sigma=10)
    mu = pm.Deterministic("mu", alpha + beta * x_shared**0.5)
    # Linear model for the std dev.
    gamma = pm.HalfNormal("gamma", sigma=10)
    delta = pm.HalfNormal("delta", sigma=10)
    sigma = pm.Deterministic("sigma", gamma + delta * x_shared)
    # Fit.
    y_pred = pm.Normal("y_pred", mu=mu, sigma=sigma, observed=data.length)
    #
    idata_vv = pm.sample(random_seed=123)

pm.model_to_graphviz(model_vv)

save_dot(model_vv, "Lesson07_Variable_variance_model.png")

# Plot the data.
plt.plot(data.month, data.length, "C0.", alpha=0.1)

# Compute the posterior mean and sigma.
posterior = az.extract(idata_vv)
mu_m = posterior["mu"].mean("sample").values
sigma_m = posterior["sigma"].mean("sample").values

# Plot 1 and 2 std dev of the model.
plt.plot(data.month, mu_m, c="k")
plt.fill_between(
    data.month, mu_m + 1 * sigma_m, mu_m - 1 * sigma_m, alpha=0.6, color="C1"
)
plt.fill_between(
    data.month, mu_m + 2 * sigma_m, mu_m - 2 * sigma_m, alpha=0.4, color="C1"
)

ut.save_plt("Lesson07_Variable_variance_result.png")

np.random.seed(314)

N = 100
# N = 1000
alpha_real = 2.5
beta_real = [0.9, 1.5]
eps_stddev_real = 0.5
eps_real = np.random.normal(0, eps_stddev_real, size=N)

# Independent variables.
X = np.array(
    [
        np.random.normal(i, j, N)
        for i, j in zip(
            # mean of gaussian.
            [10, 2],
            # std dev.
            [1, 1.5],
        )
    ]
).T
X_mean = X.mean(axis=0, keepdims=True)
X_centered = X - X_mean

# Create samples.
y = alpha_real + np.dot(X, beta_real) + eps_real

def scatter_plot(x, y):
    # plt.figure(figsize=(10, 10))
    # Plot y vs [x_1, x_2].
    for idx, x_i in enumerate(x.T):
        plt.subplot(2, 2, idx + 1)
        plt.scatter(x_i, y)
        plt.xlabel(f"x_{idx + 1}")
        plt.ylabel("y", rotation=0)
    # Plot x_2 vs x_1.
    plt.subplot(2, 2, idx + 2)
    plt.scatter(x[:, 0], x[:, 1])
    plt.xlabel(f"x_{idx}")
    plt.ylabel(f"x_{idx + 1}", rotation=0)


scatter_plot(X_centered, y)
ut.save_plt("Lesson07_Multiple_linear_regression3.png")

with pm.Model() as model_mlr:
    alpha_tmp = pm.Normal("alpha_tmp", mu=0, sigma=10)
    # Beta is a vector.
    beta = pm.Normal("beta", mu=0, sigma=1, shape=2)
    eps = pm.HalfCauchy("eps", 5)
    # mu.
    mu = alpha_tmp + pm.math.dot(X_centered, beta)
    # Extract alpha.
    alpha = pm.Deterministic("alpha", alpha_tmp - pm.math.dot(X_mean, beta))

    # Model.
    y_pred = pm.Normal("y_pred", mu=mu, sigma=eps, observed=y)

    idata_mlr = pm.sample(2000)

ut.save_dot(model_mlr, "Lesson07_Multiple_linear_regression_model.png")
pm.model_to_graphviz(model_mlr)

var_names = ["alpha", "beta", "eps"]
az.plot_trace(idata_mlr, var_names=var_names)
ut.save_plt("Lesson07_Multiple_linear_regression_results1.png")

df = az.summary(idata_mlr, var_names=var_names, round_to=2, kind="stats")
ut.save_df(df, "Lesson07_Multiple_linear_regression_results2.png")
df

with pm.Model() as model_mlb:
    alpha = pm.Normal("alpha", mu=0, sigma=1)
    beta0 = pm.Normal("beta0", mu=0, sigma=10)
    beta1 = pm.Normal("beta1", mu=0, sigma=10)
    sigma = pm.HalfNormal("sigma", 10)
    mu = pm.Deterministic(
        "mu", pm.math.exp(alpha + beta0 * bikes.temperature + beta1 * bikes.hour)
    )
    _ = pm.NegativeBinomial("y_pred", mu=mu, alpha=sigma, observed=bikes.rented)
    #
    idata_mlb = pm.sample()

pm.model_to_graphviz(model_mlb)

ut.save_dot(
    model_mlb, "Lesson07_Multiple_linear_regression_model_RentedBikes_model.png"
)

var_names = ["alpha", "beta0", "beta1", "sigma"]
az.plot_trace(idata_mlb, var_names=var_names)
ut.save_plt(
    "Lesson07_Multiple_linear_regression_model_RentedBikes_model_trace.png"
)

df = az.summary(idata_mlb, var_names=var_names, round_to=2, kind="stats")
df