Update ARIMA time series Model

Author: ai-nerdy

machine_learning/arima.py

@@ -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)
@@ -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.
@@ -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]
@@ -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:])