SimulateRegressionModel.py

# %% [markdown]
"""
# Simulated Regression and State Space Modeling

This program is inspired by a semiconductor fabrication analytics project. The
client was forecasting a yield-critical measurement from related sensor signals
with linear regression, but my exploratory analysis showed temporal dynamics
that argued for a state space approach. The script recreates those patterns with
synthetic data, compares the modeling options, and saves all figures to the
`artifacts/` directory.
"""

# %%
from pathlib import Path

import matplotlib.pyplot as plt
import numpy as np
import pandas as pd
from matplotlib import animation
from scipy import stats
from sklearn.metrics import mean_absolute_error
import statsmodels.api as sm

try:  # Jupyter-friendly display if available
    from IPython import get_ipython  # type: ignore
except ImportError:  # pragma: no cover - fallback when IPython missing
    get_ipython = lambda: None  # type: ignore

try:
    from IPython.display import display as notebook_display  # type: ignore
except ImportError:  # pragma: no cover - fallback when IPython missing
    notebook_display = None

plt.style.use("seaborn-v0_8")

ARTIFACTS = Path("artifacts")
ARTIFACTS.mkdir(exist_ok=True)


def savefig(fig: plt.Figure, filename: str, *, show: bool = True) -> None:
    """Persist figure and optionally display the result before closing."""

    fig.tight_layout()
    fig.savefig(ARTIFACTS / filename, dpi=150)

    if show and notebook_display is not None and get_ipython():
        notebook_display(fig)

    plt.close(fig)


def _series_to_array(series: pd.Series, length: int) -> np.ndarray:
    """Align a series with integer indices to a dense NumPy array."""
    values = np.full(length, np.nan)
    values[np.asarray(series.index)] = series.to_numpy()
    return values


def create_regression_animation(
    x: np.ndarray,
    y: np.ndarray,
    intercept: np.ndarray,
    slope: np.ndarray,
    y_pred: np.ndarray,
    filename: str,
    *,
    start: int = 2,
    interval: int = 200,
) -> None:
    """Render an ArtistAnimation that follows one-step predictions over time."""

    x = np.asarray(x)
    y = np.asarray(y)
    intercept = np.asarray(intercept)
    slope = np.asarray(slope)
    y_pred = np.asarray(y_pred)

    fig, ax = plt.subplots(figsize=(5, 4))
    xx = np.array([float(np.nanmin(x)), float(np.nanmax(x))])
    y_min = float(np.nanmin(y))
    y_max = float(np.nanmax(y))

    frames = []
    for t in range(start, len(y)):
        if np.isnan(intercept[t]) or np.isnan(slope[t]) or np.isnan(y_pred[t]):
            continue

        line = slope[t] * xx + intercept[t]
        frame = [
            ax.scatter(x[:t], y[:t], color="gray", s=20),
            ax.plot(xx, line, color="crimson")[0],
            ax.scatter(x[t], y[t], color="navy", s=40),
            ax.scatter(x[t], y_pred[t], color="crimson", s=40),
        ]
        frames.append(frame)

    if not frames:
        plt.close(fig)
        return

    ax.set_ylim(np.floor(y_min) - 2, np.ceil(y_max) + 2)
    ax.set_xlabel("x")
    ax.set_ylabel("y")

    ani = animation.ArtistAnimation(fig, frames, interval=interval)
    ani.save(ARTIFACTS / filename, writer="pillow")
    plt.close(fig)


# %% [markdown]
"""
## Dataset 1: Local Level with Static Regression Coefficient

We simulate a latent local level process, a slowly varying regressor, and noisy
measurements.  The true relationship uses a fixed slope `beta = 0.8`.
"""

# %%
N1 = 50
A_SD = 1.8
B_SD = 1.0
Y_SD = 1.2
A0 = 10.0
BETA = 0.8

rng = np.random.default_rng(4)

state = A0 + rng.normal(0.0, A_SD, N1).cumsum()
base = rng.normal(0.0, B_SD, N1).cumsum()
trend = np.repeat((0.1, -0.05, 0.1), (12, 20, 18))
x1 = base + trend
y1 = state + BETA * x1 + rng.normal(0.0, Y_SD, N1)
idx1 = np.arange(N1)

# %%
fig, axes = plt.subplots(2, 1, sharex=True, figsize=(7, 5))
axes[0].scatter(idx1, y1, s=20)
axes[0].set_ylabel("y")
axes[1].scatter(idx1, x1, s=20)
axes[1].set_ylabel("x")
axes[1].set_xlabel("time")
savefig(fig, "dataset1_timeseries.png")

fig, ax = plt.subplots(figsize=(5, 4))
sc = ax.scatter(x1, y1, c=idx1, s=20, cmap="viridis")
cb = fig.colorbar(sc, ax=ax, label="time")
ax.set_xlabel("x")
ax.set_ylabel("y")
savefig(fig, "dataset1_scatter.png")

fig = sm.graphics.tsa.plot_acf(y1, lags=20)
fig.set_size_inches(5, 4)
savefig(fig, "dataset1_acf.png")

# %% [markdown]
"""
### Online Linear Regression Baseline

A sliding estimator refits ordinary least squares at every step using all data
up to time `t-1` and generates a one-step-ahead prediction.
"""

# %%
def online_linear_regression(x: np.ndarray, y: np.ndarray, start: int = 2) -> pd.DataFrame:
    records = []
    for t in range(start, len(x)):
        slope, intercept, _, _, _ = stats.linregress(x[:t], y[:t])
        records.append((t, slope, intercept, slope * x[t] + intercept))
    df = pd.DataFrame(records, columns=["t", "slope", "intercept", "y_pred"])
    df.set_index("t", inplace=True)
    return df


ols_df1 = online_linear_regression(x1, y1)
mae_ols1 = mean_absolute_error(y1[2:], ols_df1["y_pred"])
print(f"Dataset1 OLS one-step MAE = {mae_ols1:.3f}")

# %%
T_SNAPSHOT = 48
xx = np.array([x1.min(), x1.max()])
yy = ols_df1.loc[T_SNAPSHOT, "slope"] * xx + ols_df1.loc[T_SNAPSHOT, "intercept"]

fig, ax = plt.subplots(figsize=(4.5, 4))
ax.scatter(x1[:T_SNAPSHOT], y1[:T_SNAPSHOT], color="gray", s=20)
ax.plot(xx, yy, c="crimson")
ax.scatter(x1[T_SNAPSHOT], y1[T_SNAPSHOT], c="navy", label="truth")
ax.scatter(x1[T_SNAPSHOT], ols_df1.loc[T_SNAPSHOT, "y_pred"], c="crimson", label="pred")
ax.legend()
ax.set_xlabel("x")
ax.set_ylabel("y")
ax.set_title(f"OLS snapshot t={T_SNAPSHOT}")
savefig(fig, "dataset1_ols_snapshot.png")

# %%
fig, ax = plt.subplots(figsize=(7, 3))
ax.scatter(idx1, y1, s=20, alpha=0.8, label="observed")
ax.scatter(ols_df1.index, ols_df1["y_pred"], s=20, alpha=0.8, label="pred")
ax.legend()
ax.set_xlabel("time")
ax.set_ylabel("y")
savefig(fig, "dataset1_ols_vs_obs.png")

ols_intercept_arr1 = _series_to_array(ols_df1["intercept"], N1)
ols_slope_arr1 = _series_to_array(ols_df1["slope"], N1)
ols_pred_arr1 = _series_to_array(ols_df1["y_pred"], N1)
create_regression_animation(
    x1,
    y1,
    ols_intercept_arr1,
    ols_slope_arr1,
    ols_pred_arr1,
    "dataset1_ols_animation.gif",
    start=int(ols_df1.index.min()),
)

# %% [markdown]
"""
### State Space Model: Constant Regression Coefficient

We fit an Unobserved Components model with a local level and a fixed regression
coefficient on `x`.  The Kalman filter provides filtered states and predictive
intervals.
"""

# %%
ss_mod1_const = sm.tsa.UnobservedComponents(y1, "llevel", exog=x1)
ss_res1_const = ss_mod1_const.fit(disp=False)
print(ss_res1_const.summary())

pred_frame1_const = ss_res1_const.get_prediction().summary_frame()
pred_frame1_const.rename(columns={
    "mean": "y_pred",
    "mean_se": "pred_se",
    "mean_ci_lower": "lwr",
    "mean_ci_upper": "upr",
}, inplace=True)
pred_frame1_const["intercept"] = ss_res1_const.predicted_state[0, :N1]
const_slope = ss_res1_const.params[-1]

# %%
fig, ax = plt.subplots(figsize=(7, 3))
ax.scatter(idx1, y1, s=20, alpha=0.8, label="observed")
ax.scatter(idx1[2:], pred_frame1_const.loc[2:, "y_pred"], s=20, label="pred")
ax.fill_between(idx1[2:], pred_frame1_const.loc[2:, "lwr"], pred_frame1_const.loc[2:, "upr"],
                color="orange", alpha=0.2, label="80% CI")
ax.legend()
ax.set_xlabel("time")
ax.set_ylabel("y")
savefig(fig, "dataset1_ss_const_vs_obs.png")

const_intercept_arr1 = pred_frame1_const["intercept"].to_numpy()
const_slope_arr1 = np.full(N1, const_slope)
const_pred_arr1 = pred_frame1_const["y_pred"].to_numpy()
create_regression_animation(
    x1,
    y1,
    const_intercept_arr1,
    const_slope_arr1,
    const_pred_arr1,
    "dataset1_ss_const_animation.gif",
    start=2,
)

mae_ss1_const = mean_absolute_error(y1[2:], pred_frame1_const.loc[2:, "y_pred"])
print(f"Dataset1 state space (const beta) MAE = {mae_ss1_const:.3f}")

# %%
fig, ax = plt.subplots(figsize=(4.5, 4))
ax.scatter(x1[:T_SNAPSHOT], y1[:T_SNAPSHOT], color="gray", s=20)
ax.plot(xx, const_slope * xx + pred_frame1_const.loc[T_SNAPSHOT, "intercept"], color="crimson")
ax.scatter(x1[T_SNAPSHOT], y1[T_SNAPSHOT], color="navy", s=40, label="truth")
ax.scatter(x1[T_SNAPSHOT], pred_frame1_const.loc[T_SNAPSHOT, "y_pred"], color="crimson", s=40, label="pred")
ax.legend()
ax.set_xlabel("x")
ax.set_ylabel("y")
ax.set_title(f"State space snapshot t={T_SNAPSHOT}")
savefig(fig, "dataset1_ss_const_snapshot.png")

# %% [markdown]
"""
### State Space Model: Time-Varying Regression Coefficient

Setting `mle_regression=False` lets the slope evolve as part of the state
vector.  This better accommodates gradual structural change.
"""

# %%
ss_mod1_tv = sm.tsa.UnobservedComponents(y1, "llevel", exog=x1, mle_regression=False)
ss_res1_tv = ss_mod1_tv.fit(disp=False)
print(ss_res1_tv.summary())

pred_frame1_tv = ss_res1_tv.get_prediction().summary_frame()
pred_frame1_tv.rename(columns={
    "mean": "y_pred",
    "mean_se": "pred_se",
    "mean_ci_lower": "lwr",
    "mean_ci_upper": "upr",
}, inplace=True)
pred_frame1_tv["intercept"] = ss_res1_tv.predicted_state[0, :N1]
pred_frame1_tv["slope"] = ss_res1_tv.predicted_state[1, :N1]

mae_ss1_tv = mean_absolute_error(y1[2:], pred_frame1_tv.loc[2:, "y_pred"])
print(f"Dataset1 state space (time-varying beta) MAE = {mae_ss1_tv:.3f}")

# %%
fig, ax = plt.subplots(figsize=(7, 3))
ax.scatter(idx1, y1, s=20, alpha=0.8, label="observed")
ax.scatter(idx1[2:], pred_frame1_tv.loc[2:, "y_pred"], s=20, label="pred")
ax.fill_between(idx1[2:], pred_frame1_tv.loc[2:, "lwr"], pred_frame1_tv.loc[2:, "upr"],
                color="orange", alpha=0.2)
ax.legend()
ax.set_xlabel("time")
ax.set_ylabel("y")
savefig(fig, "dataset1_ss_tv_vs_obs.png")

tv_intercept_arr1 = pred_frame1_tv["intercept"].to_numpy()
tv_slope_arr1 = pred_frame1_tv["slope"].to_numpy()
tv_pred_arr1 = pred_frame1_tv["y_pred"].to_numpy()
create_regression_animation(
    x1,
    y1,
    tv_intercept_arr1,
    tv_slope_arr1,
    tv_pred_arr1,
    "dataset1_ss_tv_animation.gif",
    start=2,
)

fig, ax = plt.subplots(figsize=(7, 3))
ax.plot(idx1, pred_frame1_tv["slope"], label="estimated slope")
ax.axhline(BETA, color="black", linestyle="--", label="true beta")
ax.legend()
ax.set_xlabel("time")
ax.set_ylabel("slope")
savefig(fig, "dataset1_ss_tv_slope.png")

# %% [markdown]
"""
## Dataset 2: Level Shifts and Change Points

We extend the simulation with abrupt level shifts that are known through dummy
variables.  This stresses models that assume smooth latent behavior.
"""

# %%
N2 = 100
rng = np.random.default_rng(1)

state2 = A0 + rng.normal(0.0, 0.6, N2).cumsum()
base2 = rng.normal(0.0, 1.0, N2).cumsum()
trend2 = np.repeat((0.1, -0.05, 0.1), (24, 40, 36))
x2 = base2 + trend2
y2 = state2 + 0.7 * x2 + rng.normal(0.0, 0.9, N2)

period = np.zeros(N2, dtype=int)
period[30:] += 1
period[70:] += 1
period_dummies = pd.get_dummies(period, prefix="period", drop_first=True, dtype=float)

level_shift = period_dummies.to_numpy().dot(np.array([8.0, -18.0]))
y2 += level_shift
idx2 = np.arange(N2)

# %%
fig, axes = plt.subplots(2, 1, sharex=True, figsize=(7, 5))
axes[0].scatter(idx2, y2, s=20)
axes[0].set_ylabel("y")
axes[1].scatter(idx2, x2, s=20)
axes[1].set_ylabel("x")
axes[1].set_xlabel("time")
savefig(fig, "dataset2_timeseries.png")

fig, ax = plt.subplots(figsize=(5, 4))
sc = ax.scatter(x2, y2, c=idx2, s=20, cmap="viridis")
fig.colorbar(sc, ax=ax, label="time")
ax.set_xlabel("x")
ax.set_ylabel("y")
savefig(fig, "dataset2_scatter.png")

# %% [markdown]
"""
### Online Linear Regression Baseline
"""

# %%
ols_df2 = online_linear_regression(x2, y2)
mae_ols2 = mean_absolute_error(y2[2:], ols_df2["y_pred"])
print(f"Dataset2 OLS one-step MAE = {mae_ols2:.3f}")

fig, ax = plt.subplots(figsize=(7, 3))
ax.scatter(idx2, y2, s=20, alpha=0.8, label="observed")
ax.scatter(ols_df2.index, ols_df2["y_pred"], s=20, alpha=0.8, label="pred")
ax.legend()
ax.set_xlabel("time")
ax.set_ylabel("y")
savefig(fig, "dataset2_ols_vs_obs.png")

ols_intercept_arr2 = _series_to_array(ols_df2["intercept"], N2)
ols_slope_arr2 = _series_to_array(ols_df2["slope"], N2)
ols_pred_arr2 = _series_to_array(ols_df2["y_pred"], N2)
create_regression_animation(
    x2,
    y2,
    ols_intercept_arr2,
    ols_slope_arr2,
    ols_pred_arr2,
    "dataset2_ols_animation.gif",
    start=int(ols_df2.index.min()),
)

# %% [markdown]
"""
### State Space: Constant Regression Coefficient
"""

# %%
ss_mod2_const = sm.tsa.UnobservedComponents(y2, "llevel", exog=x2)
ss_res2_const = ss_mod2_const.fit(disp=False)
print(ss_res2_const.summary())

pred_frame2_const = ss_res2_const.get_prediction().summary_frame()
pred_frame2_const.rename(columns={
    "mean": "y_pred",
    "mean_se": "pred_se",
    "mean_ci_lower": "lwr",
    "mean_ci_upper": "upr",
}, inplace=True)
pred_frame2_const["intercept"] = ss_res2_const.predicted_state[0, :N2]

mae_ss2_const = mean_absolute_error(y2[2:], pred_frame2_const.loc[2:, "y_pred"])
print(f"Dataset2 state space (const beta) MAE = {mae_ss2_const:.3f}")

fig, ax = plt.subplots(figsize=(7, 3))
ax.scatter(idx2, y2, s=20, alpha=0.8, label="observed")
ax.scatter(idx2[2:], pred_frame2_const.loc[2:, "y_pred"], s=20, label="pred")
ax.fill_between(idx2[2:], pred_frame2_const.loc[2:, "lwr"], pred_frame2_const.loc[2:, "upr"],
                color="orange", alpha=0.2)
ax.legend()
ax.set_xlabel("time")
ax.set_ylabel("y")
savefig(fig, "dataset2_ss_const_vs_obs.png")

const_intercept_arr2 = pred_frame2_const["intercept"].to_numpy()
const_slope_arr2 = np.full(N2, ss_res2_const.params[-1])
const_pred_arr2 = pred_frame2_const["y_pred"].to_numpy()
create_regression_animation(
    x2,
    y2,
    const_intercept_arr2,
    const_slope_arr2,
    const_pred_arr2,
    "dataset2_ss_const_animation.gif",
    start=2,
)

# %% [markdown]
"""
### State Space: Time-Varying Regression Coefficient

The time-varying specification often inflates the process variance to absorb
the jumps, leading to wide predictive intervals.
"""

# %%
ss_mod2_tv = sm.tsa.UnobservedComponents(y2, "llevel", exog=x2, mle_regression=False)
ss_res2_tv = ss_mod2_tv.fit(disp=False)
print(ss_res2_tv.summary())

pred_frame2_tv = ss_res2_tv.get_prediction().summary_frame()
pred_frame2_tv.rename(columns={
    "mean": "y_pred",
    "mean_se": "pred_se",
    "mean_ci_lower": "lwr",
    "mean_ci_upper": "upr",
}, inplace=True)
pred_frame2_tv["intercept"] = ss_res2_tv.predicted_state[0, :N2]
pred_frame2_tv["slope"] = ss_res2_tv.predicted_state[1, :N2]

mae_ss2_tv = mean_absolute_error(y2[2:], pred_frame2_tv.loc[2:, "y_pred"])
print(f"Dataset2 state space (time-varying beta) MAE = {mae_ss2_tv:.3f}")

fig, ax = plt.subplots(figsize=(7, 3))
ax.scatter(idx2, y2, s=20, alpha=0.8, label="observed")
ax.scatter(idx2[2:], pred_frame2_tv.loc[2:, "y_pred"], s=20, label="pred")
ax.fill_between(idx2[2:], pred_frame2_tv.loc[2:, "lwr"], pred_frame2_tv.loc[2:, "upr"],
                color="orange", alpha=0.2)
ax.legend()
ax.set_xlabel("time")
ax.set_ylabel("y")
savefig(fig, "dataset2_ss_tv_vs_obs.png")

tv_intercept_arr2 = pred_frame2_tv["intercept"].to_numpy()
tv_slope_arr2 = pred_frame2_tv["slope"].to_numpy()
tv_pred_arr2 = pred_frame2_tv["y_pred"].to_numpy()
create_regression_animation(
    x2,
    y2,
    tv_intercept_arr2,
    tv_slope_arr2,
    tv_pred_arr2,
    "dataset2_ss_tv_animation.gif",
    start=2,
)

# %% [markdown]
"""
### Modeling Known Change Points with Exogenous Dummies

Providing the change point indicators as regressors allows the local level
component to remain smooth and keeps the predictive intervals tight.
"""

# %%
exog2 = period_dummies.copy()
exog2["x"] = x2
ss_mod2_exog = sm.tsa.UnobservedComponents(y2, "llevel", exog=exog2)
ss_res2_exog = ss_mod2_exog.fit(disp=False)
print(ss_res2_exog.summary())

pred_frame2_exog = ss_res2_exog.get_prediction().summary_frame()
pred_frame2_exog.rename(columns={
    "mean": "y_pred",
    "mean_se": "pred_se",
    "mean_ci_lower": "lwr",
    "mean_ci_upper": "upr",
}, inplace=True)
pred_frame2_exog["intercept"] = ss_res2_exog.predicted_state[0, :N2]

mae_ss2_exog = mean_absolute_error(y2[2:], pred_frame2_exog.loc[2:, "y_pred"])
print(f"Dataset2 state space (with change point dummies) MAE = {mae_ss2_exog:.3f}")

fig, ax = plt.subplots(figsize=(7, 3))
ax.scatter(idx2, y2, s=20, alpha=0.8, label="observed")
ax.scatter(idx2[2:], pred_frame2_exog.loc[2:, "y_pred"], s=20, label="pred")
ax.fill_between(idx2[2:], pred_frame2_exog.loc[2:, "lwr"], pred_frame2_exog.loc[2:, "upr"],
                color="orange", alpha=0.2)
ax.legend()
ax.set_xlabel("time")
ax.set_ylabel("y")
savefig(fig, "dataset2_ss_exog_vs_obs.png")

# %% [markdown]
"""
### Heavy-Tailed Evolution (Conceptual)

The original notebook explored Stan code with a Cauchy prior on the state
innovations.  Compiling the Stan model requires `pystan`, which is not bundled
with this environment.  If `pystan` is installed locally, the Stan program can
be evaluated by uncommenting the block below.
"""

# %%
try:
    from pystan import StanModel  # type: ignore
except ModuleNotFoundError:
    StanModel = None
    print("pystan is not available; skipping the Stan compilation demo.")

stan_code = """
    data {
      int N;
      vector[N] Y;
      vector[N] X;
    }
    parameters {
      real mu0;
      real<lower=0> s_mu;
      vector<lower=-pi()/2, upper=pi()/2>[N-1] mu_unif;
      real<lower=0> s_y;
      real beta;
    }
    transformed parameters {
      vector[N] mu;
      mu[1] = mu0;
      for (n in 2:N)
        mu[n] = mu[n-1] + s_mu * tan(mu_unif[n-1]);
    }
    model {
      Y ~ normal(mu + beta * X, s_y);
    }
"""

if StanModel is not None:
    stan_model = StanModel(model_code=stan_code)
    print("Stan model compiled successfully.")


# %% [markdown]
"""
## Dataset 3: Local Level with Additive Shock

An abrupt shock is added after time step 20 to show how the Kalman filter
handles sudden jumps compared with OLS.
"""

# %%
N3 = 50
rng = np.random.default_rng(5)

state3 = A0 + rng.normal(0.0, 0.6, N3).cumsum()
base3 = rng.normal(0.0, 1.0, N3).cumsum()
trend3 = np.repeat((0.1, -0.05, 0.1), (12, 20, 18))
x3 = base3 + trend3
y3 = state3 + 0.8 * x3 + rng.normal(0.0, 0.6, N3)
y3[20:] += 8.0
idx3 = np.arange(N3)

# %%
fig, axes = plt.subplots(2, 1, figsize=(8, 6), sharex=True)
axes[0].scatter(idx3, x3, marker="s")
axes[0].set_ylabel("x")
axes[1].scatter(idx3, y3)
axes[1].set_ylabel("y")
axes[1].set_xlabel("time")
savefig(fig, "dataset3_timeseries.png")

fig, ax = plt.subplots(figsize=(4, 4))
ax.scatter(x3, y3)
ax.set_xlabel("x")
ax.set_ylabel("y")
savefig(fig, "dataset3_scatter.png")

# %%
ols_df3 = online_linear_regression(x3, y3)
ss_mod3 = sm.tsa.UnobservedComponents(y3, "llevel", exog=x3)
ss_res3 = ss_mod3.fit(disp=False)

pred_frame3 = ss_res3.get_prediction().summary_frame()
pred_frame3.rename(columns={
    "mean": "y_pred",
    "mean_ci_lower": "lwr",
    "mean_ci_upper": "upr",
}, inplace=True)

fig, axes = plt.subplots(2, 1, figsize=(8, 6), sharex=True)
axes[0].scatter(idx3, x3)
axes[0].set_ylabel("x")
axes[1].scatter(idx3, y3, s=15, alpha=0.8, label="observed")
axes[1].scatter(idx3[1:], pred_frame3.loc[1:, "y_pred"], marker="s", alpha=0.8,
                label="state space")
axes[1].fill_between(idx3[1:], pred_frame3.loc[1:, "lwr"], pred_frame3.loc[1:, "upr"],
                     color="orange", alpha=0.2)
axes[1].scatter(ols_df3.index, ols_df3["y_pred"], marker="x", label="OLS")
axes[1].legend()
axes[1].set_ylabel("y")
axes[1].set_xlabel("time")
savefig(fig, "dataset3_comparison.png")


# %% [markdown]
"""
## Decomposition Example: Trend + Level + Seasonality

A final simulation shows how the Unobserved Components model can recover trend
and seasonal components from data constructed with known latent pieces.
"""

# %%
N4 = 100
rng = np.random.default_rng(42)

level4 = rng.normal(scale=1.5, size=N4).cumsum()
trend4 = np.repeat(0.4, N4).cumsum()
seasonal4 = 4.0 * np.sin(2 * np.pi * np.arange(N4) / 8) + rng.normal(scale=0.4, size=N4)
noise4 = rng.normal(scale=1.0, size=N4)
y4 = level4 + trend4 + seasonal4 + noise4

fig, axes = plt.subplots(4, 1, figsize=(10, 9), sharex=True)
axes[0].plot(level4)
axes[0].set_ylabel("level")
axes[1].plot(trend4)
axes[1].set_ylabel("trend")
axes[2].plot(seasonal4)
axes[2].set_ylabel("seasonal")
axes[3].plot(y4)
axes[3].set_ylabel("observed")
axes[3].set_xlabel("time")
savefig(fig, "dataset4_components.png")

fig, ax = plt.subplots(figsize=(8, 3))
ax.plot(y4)
ax.set_ylabel("y")
ax.set_xlabel("time")
savefig(fig, "dataset4_series.png")

mod4 = sm.tsa.UnobservedComponents(
    y4,
    level=True,
    trend=True,
    stochastic_level=True,
    stochastic_trend=True,
    seasonal=8,
)
res4 = mod4.fit(disp=False)
print(res4.summary())

fig = res4.plot_components(figsize=(10, 12))
savefig(fig, "dataset4_decomposition.png")

print("Processing complete. Figures and animations are in the artifacts/ folder.")

# %%