Spaces:
Sleeping
Sleeping
| """SARIMA forecast for BBCA closing prices. | |
| The order comes from the notebook's auto_arima search: | |
| SARIMA(0,1,0)(2,1,0)[5], where 5 is one trading week. | |
| """ | |
| from __future__ import annotations | |
| import warnings | |
| from dataclasses import dataclass | |
| import numpy as np | |
| import pandas as pd | |
| from statsmodels.tsa.statespace.sarimax import SARIMAX | |
| ORDER = (0, 1, 0) | |
| SEASONAL_ORDER = (2, 1, 0, 5) | |
| FIT_WINDOW = 250 # about one year of trading days | |
| MAX_HORIZON = 10 | |
| class Forecast: | |
| dates: pd.DatetimeIndex | |
| mean: np.ndarray | |
| lower: np.ndarray | |
| upper: np.ndarray | |
| base_date: pd.Timestamp # last observation the forecast starts from | |
| base_price: float | |
| base_is_live: bool | |
| fit_points: int | |
| def fit(values: np.ndarray): | |
| """Fit SARIMA on a plain array (no date index, so no frequency warnings).""" | |
| with warnings.catch_warnings(): | |
| warnings.simplefilter("ignore") | |
| model = SARIMAX( | |
| np.asarray(values, dtype=float), | |
| order=ORDER, | |
| seasonal_order=SEASONAL_ORDER, | |
| enforce_stationarity=False, | |
| enforce_invertibility=False, | |
| ) | |
| return model.fit(disp=False) | |
| def next_trading_days(after: pd.Timestamp, n: int) -> pd.DatetimeIndex: | |
| """Weekdays after `after`. IDX holidays are not removed.""" | |
| return pd.bdate_range(after + pd.Timedelta(days=1), periods=n) | |
| def forecast( | |
| closes: pd.Series, | |
| horizon: int, | |
| live_price: float | None = None, | |
| live_date: pd.Timestamp | None = None, | |
| ) -> Forecast: | |
| horizon = int(min(max(horizon, 1), MAX_HORIZON)) | |
| series = closes.iloc[-FIT_WINDOW:] | |
| values = series.to_numpy(dtype=float) | |
| base_date = series.index[-1] | |
| base_is_live = False | |
| if live_price is not None and live_date is not None: | |
| # Treat the running price as today's close. The forecast then moves | |
| # with the market during trading hours. | |
| values = np.append(values, live_price) | |
| base_date = live_date | |
| base_is_live = True | |
| result = fit(values) | |
| frame = result.get_forecast(steps=horizon).summary_frame(alpha=0.05) | |
| return Forecast( | |
| dates=next_trading_days(base_date, horizon), | |
| mean=frame["mean"].to_numpy(), | |
| lower=frame["mean_ci_lower"].to_numpy(), | |
| upper=frame["mean_ci_upper"].to_numpy(), | |
| base_date=base_date, | |
| base_price=float(values[-1]), | |
| base_is_live=base_is_live, | |
| fit_points=len(values), | |
| ) | |