Skip to content
Open
Prev Previous commit
Next Next commit
Update ARIMA time series Model
  • Loading branch information
ai-nerdy committed Oct 6, 2025
commit a130bf811395172d4530b4ad875030ba96cfeb55
42 changes: 20 additions & 22 deletions machine_learning/arima.py
Original file line number Diff line number Diff line change
Expand Up @@ -5,7 +5,7 @@

>>> import numpy as np
>>> series = np.array([1, 2, 3, 4, 5, 6, 7, 8, 9, 10])
>>> model = ARIMAModel(p=2, d=1, q=0)
>>> model = ARIMAModel(ar_order=2, diff_order=1, ma_order=0)
>>> model.fit(series)
ARIMAModel(...)
>>> model.predict(series, n_periods=2)
Expand All @@ -17,24 +17,24 @@


class ARIMAModel:
def __init__(self, p: int = 1, d: int = 0, q: int = 0) -> None:
def __init__(self, ar_order: int = 1, diff_order: int = 0, ma_order: int = 0) -> None:
"""Initialize ARIMA model.
Args:
p: AR order
d: Differencing order
q: MA order (not used in this implementation)
ar_order: Autoregressive order (p)
diff_order: Differencing order (d)
ma_order: Moving average order (q, not used in this implementation)
"""
self.p = p
self.d = d
self.q = q
self.ar_order = ar_order
self.diff_order = diff_order
self.ma_order = ma_order
self.coef_: Optional[np.ndarray] = None
self.resid_: Optional[np.ndarray] = None

def difference(self, series: np.ndarray, order: int) -> np.ndarray:
def difference(self, time_series: np.ndarray, order: int) -> np.ndarray:
"""Apply differencing to make series stationary."""
for _ in range(order):
series = np.diff(series)
return series
time_series = np.diff(time_series)
return time_series

def fit(self, time_series: np.ndarray) -> "ARIMAModel":
"""Fit ARIMA model to the given time series.
Expand All @@ -44,18 +44,16 @@ def fit(self, time_series: np.ndarray) -> "ARIMAModel":
self
>>> import numpy as np
>>> series = np.array([1, 2, 3, 4, 5, 6, 7, 8, 9, 10])
>>> model = ARIMAModel(p=2, d=1, q=0)
>>> model = ARIMAModel(ar_order=2, diff_order=1, ma_order=0)
>>> model.fit(series)
ARIMAModel(...)
"""
y = np.asarray(time_series)
y_diff = self.difference(y, self.d)
y_diff = self.difference(y, self.diff_order)
# Build lagged feature matrix
feature_matrix = np.column_stack(
[np.roll(y_diff, i) for i in range(1, self.p + 1)]
)
feature_matrix = feature_matrix[self.p :]
target = y_diff[self.p :]
feature_matrix = np.column_stack([np.roll(y_diff, i) for i in range(1, self.ar_order + 1)])
feature_matrix = feature_matrix[self.ar_order:]
target = y_diff[self.ar_order:]
# Add intercept
feature_matrix = np.hstack(
[np.ones((feature_matrix.shape[0], 1)), feature_matrix]
Expand All @@ -74,17 +72,17 @@ def predict(self, time_series: np.ndarray, n_periods: int = 1) -> np.ndarray:
1D numpy array of forecasted values
>>> import numpy as np
>>> series = np.array([1, 2, 3, 4, 5, 6, 7, 8, 9, 10])
>>> model = ARIMAModel(p=2, d=1, q=0)
>>> model = ARIMAModel(ar_order=2, diff_order=1, ma_order=0)
>>> model.fit(series)
ARIMAModel(...)
>>> model.predict(series, n_periods=2)
array([10.99999999, 12.00000001])
"""
y = np.asarray(time_series)
y_pred = list(y[-self.p :])
y_pred = list(y[-self.ar_order:])
for _ in range(n_periods):
# Build feature vector for prediction
features = [1] + y_pred[-self.p :][::-1]
features = [1] + y_pred[-self.ar_order:][::-1]
next_val = np.dot(features, self.coef_)
y_pred.append(next_val)
return np.array(y_pred[self.p :])
return np.array(y_pred[self.ar_order:])