#!/usr/bin/env python3
# tests/test_stats_model.py
import pytest
import pandas as pd
import numpy as np
from timeseries_compute.stats_model import ModelARIMA
[docs]
@pytest.fixture
def sample_data():
"""gen ts data for arima tests"""
np.random.seed(42) # fixed seed
n = 50 # sample size
# ar(1) process parameters
auto_coef = 0.7 # how much past values affect future (memory strength)
noise_level_a = 1.0 # noise magnitude for series A (more volatile)
noise_level_b = 0.5 # noise magnitude for series B (less volatile)
mean_value = 0 # centered around zero
# initialize empty series
ar_series_a = np.zeros(n)
ar_series_b = np.zeros(n)
# generate two ar(1) series with different noise levels
for i in range(1, n):
# formula: current = (memory_factor * previous) + random_noise
# series A: higher noise
noise_a = np.random.normal(mean_value, noise_level_a)
ar_series_a[i] = auto_coef * ar_series_a[i - 1] + noise_a
# series B: lower noise
noise_b = np.random.normal(mean_value, noise_level_b)
ar_series_b[i] = auto_coef * ar_series_b[i - 1] + noise_b
# create timeseries dataframe
dates = pd.date_range(start="2023-01-01", periods=n, freq="D")
data = {"A": ar_series_a, "B": ar_series_b}
return pd.DataFrame(data, index=dates)
[docs]
def test_model_arima_initialization(sample_data):
"""test arima init"""
# arima components
ar_terms = 1 # p: autoregressive lags
diff_terms = 0 # d: differencing level
ma_terms = 0 # q: moving average lags
pred_steps = 2 # forecast horizon
model = ModelARIMA(
data=sample_data, order=(ar_terms, diff_terms, ma_terms), steps=pred_steps
)
assert model.data.equals(sample_data)
assert model.order == (ar_terms, diff_terms, ma_terms)
assert model.steps == pred_steps
[docs]
def test_model_arima_fit(sample_data):
"""test fit method"""
# ar(1) model specs
ar_lag = 1 # num of autoregressive lags
diff_level = 0 # no differencing needed
ma_lag = 0 # no moving average component
model = ModelARIMA(data=sample_data, order=(ar_lag, diff_level, ma_lag), steps=2)
fits = model.fit()
assert isinstance(fits, dict)
assert "A" in fits
assert "B" in fits
assert hasattr(fits["A"], "params")
assert hasattr(fits["B"], "params")
[docs]
def test_model_arima_summary(sample_data):
"""test summary output"""
# model order parameters
p_val = 1 # autoregressive order
d_val = 0 # integration order
q_val = 0 # moving average order
model = ModelARIMA(data=sample_data, order=(p_val, d_val, q_val), steps=2)
model.fit()
summaries = model.summary()
assert isinstance(summaries, dict)
assert "A" in summaries
assert "B" in summaries
assert isinstance(summaries["A"], str)
assert isinstance(summaries["B"], str)
[docs]
def test_model_arima_forecast(sample_data):
"""test forecast result"""
# model settings
model_order = (1, 0, 0) # ar order, diff order, ma order
forecast_horizon = 2 # steps to predict
model = ModelARIMA(data=sample_data, order=model_order, steps=forecast_horizon)
model.fit()
forecasts = model.forecast()
assert isinstance(forecasts, dict)
assert "A" in forecasts
assert "B" in forecasts
# For multi-step forecasts (steps > 1), we should get lists
assert isinstance(forecasts["A"], list)
assert isinstance(forecasts["B"], list)
assert len(forecasts["A"]) == forecast_horizon
assert len(forecasts["B"]) == forecast_horizon
# test multi-step forecasting - ensure more than 1 data point is forecasted
multi_step_horizon = 5 # forecast 5 steps ahead
multi_model = ModelARIMA(
data=sample_data, order=model_order, steps=multi_step_horizon
)
multi_model.fit()
multi_forecasts = multi_model.forecast()
# verify that we get forecasts for multiple steps
assert isinstance(multi_forecasts, dict)
assert "A" in multi_forecasts
assert "B" in multi_forecasts
# Should get lists with the correct number of forecast steps
assert isinstance(multi_forecasts["A"], list)
assert isinstance(multi_forecasts["B"], list)
assert len(multi_forecasts["A"]) == multi_step_horizon
assert len(multi_forecasts["B"]) == multi_step_horizon
# Test single-step forecast returns a float
single_step_model = ModelARIMA(data=sample_data, order=model_order, steps=1)
single_step_model.fit()
single_forecasts = single_step_model.forecast()
assert isinstance(single_forecasts["A"], float)
assert isinstance(single_forecasts["B"], float)