Add docstrings for functions in evaluate_ar_model.py and fit_ar_model.py.

Author: Lenr4

src/lennart_epp/analysis/evaluate_ar_model.py

@@ -5,23 +5,63 @@
 
 
 def _compute_residuals(df: pd.DataFrame, model_results: dict) -> np.ndarray:
+    """Compute residuals for an autoregressive (AR) model.
+
+    Args:
+        df (pd.DataFrame): The input DataFrame containing the 'close_price' column.
+        model_results (dict): A dictionary containing the fitted AR parameters.
+
+    Returns:
+        np.ndarray: An array of residuals.
+    """
     fitted_values = _predict_ar(df, model_results)
     return df["close_price"].iloc[len(df) - len(fitted_values) :] - fitted_values
 
 
 def _calculate_aic(residuals: np.ndarray, p: int) -> float:
+    """Calculate the Akaike Information Criterion (AIC) for model evaluation.
+
+    Args:
+        residuals (np.ndarray): An array of residuals from the AR model.
+        p (int): The number of autoregressive parameters in the model.
+
+    Returns:
+        float: The computed AIC value.
+    """
     n = len(residuals)
     rss = np.sum(residuals**2)
     return n * np.log(rss / n) + 2 * (p + 1)
 
 
 def _calculate_bic(residuals: np.ndarray, p: int) -> float:
+    """Calculate the Bayesian Information Criterion (BIC) for model evaluation.
+
+    Args:
+        residuals (np.ndarray): An array of residuals from the AR model.
+        p (int): The number of autoregressive parameters in the model.
+
+    Returns:
+        float: The computed BIC value.
+    """
     n = len(residuals)
     rss = np.sum(residuals**2)
     return n * np.log(rss / n) + np.log(n) * (p + 1)
 
 
 def _predict_ar(df: pd.DataFrame, model_results: dict) -> np.ndarray:
+    """Generate 1 step predictions (fitted values) using an autoregressive (AR) model.
+
+    The function applies AR model coefficients to past values of 'close_price'
+    to make predictions.
+
+    Args:
+        df (pd.DataFrame): The input DataFrame containing the 'close_price' column.
+        model_results (dict): A dictionary containing model coefficients under
+                              'integrated_coefficients'.
+
+    Returns:
+        np.ndarray: An array of predicted values based on the AR model.
+    """
     integrated_coeff = model_results["integrated_coefficients"][
         "coefficient"
     ].to_numpy()
@@ -39,6 +79,25 @@ def _predict_ar(df: pd.DataFrame, model_results: dict) -> np.ndarray:
 def evaluate_ar_models(
     df: pd.DataFrame, max_p: int = 15, criterion: str = "aic"
 ) -> dict:
+    """Evaluate multiple Autoregressive (AR) models and select the best ones.
+
+    This function fits AR models for different lag values (p) up to `max_p`,
+    computes the AIC and BIC scores for each model, and returns the top-performing
+    models based on the selected criterion.
+
+    Args:
+        df (pd.DataFrame): The input DataFrame containing the 'close_price' column.
+        max_p (int, optional): The maximum number of autoregressive lags to test.
+                               Defaults to 15.
+        criterion (str, optional): The selection criterion for ranking models.
+                                   Can be "aic" or "bic". Defaults to "aic".
+
+    Returns:
+        dict: A dictionary containing:
+            - **top_models** (list[dict]): The top 3 sorted by the given criterion.
+            - **model_metrics** (list[dict]): A list of evaluated models with metrics.
+            - **metadata** (dict): Information about evaluation, `max_p`, `criterion`.
+    """
     results = []
 
     for p in range(1, max_p + 1):

src/lennart_epp/analysis/fit_ar_model.py

@@ -6,19 +6,53 @@
 
 def _check_stationarity(
     df: pd.DataFrame, column: str, significance: float = 0.05
-) -> tuple[bool, float]:
+) -> tuple[bool, float, float]:
+    """Perform the Augmented Dickey-Fuller (ADF) test to check stationarity.
+
+    Args:
+        df (pd.DataFrame): The DataFrame containing the time series data.
+        column (str): The name of the column to test for stationarity.
+        significance (float, optional): The significance level for the test.
+
+    Returns:
+        tuple:
+            - bool: True if the series is stationary
+            - float: The p-value from the ADF test.
+            - float: The ADF test statistic.
+    """
     adf_test = ADF(df[column].dropna())
     p_value = adf_test.pvalue
+    test_statistic_adf = adf_test.stat
 
-    return p_value < significance, p_value
+    return p_value < significance, p_value, test_statistic_adf
 
 
 def _difference_series(df: pd.DataFrame, column: str) -> pd.DataFrame:
-    df[f"diff_{column}"] = df[column].diff().dropna()
-    return df
+    """Apply first differencing to a column in a DataFrame.
+
+    Args:
+        df (pd.DataFrame): The dataframe containing the time series.
+        column (str): The column to be differenced.
+
+    Returns:
+        pd.DataFrame: A new DataFrame with the differenced column.
+    """
+    df_copy = df.copy()
+    df_copy[f"diff_{column}"] = df_copy[column].diff().dropna()
+    return df_copy[[f"diff_{column}"]]
 
 
 def _create_lagged_features(df: pd.DataFrame, column: str, p: int) -> pd.DataFrame:
+    """Generate lagged features for an autoregressive model.
+
+    Args:
+        df (pd.DataFrame): The input DataFrame containing the time series data.
+        column (str): The column for which lagged features should be created.
+        p (int): The number of lagged periods to generate.
+
+    Returns:
+        pd.DataFrame: A DataFrame with the original column and its lagged features.
+    """
     for lag in range(1, p + 1):
         df[f"{column}_lag{lag}"] = df[column].shift(lag)
 
@@ -28,6 +62,16 @@ def _create_lagged_features(df: pd.DataFrame, column: str, p: int) -> pd.DataFra
 
 
 def _ar_model(df: pd.DataFrame, column: str, p: int) -> np.ndarray:
+    """Estimate autoregressive (AR) model parameters using the least squares method.
+
+    Args:
+        df (pd.DataFrame): The DataFrame with time series data and lagged features.
+        column (str): The target column for the autoregressive model.
+        p (int): The order (number of lags) of the AR model.
+
+    Returns:
+        np.ndarray: An array of estimated coefficients, including the intercept term.
+    """
     x = df[[f"{column}_lag{i}" for i in range(1, p + 1)]].to_numpy()
     y = df[column].to_numpy()
 
@@ -41,23 +85,28 @@ def _ar_model(df: pd.DataFrame, column: str, p: int) -> np.ndarray:
 def _integrate_ar_coefficients(
     diff_coefficients: np.ndarray, *, differenced: bool
 ) -> pd.DataFrame:
+    """Convert differenced AR model coefficients to integrated form.
+
+    Args:
+        diff_coefficients (np.ndarray): The coefficients from the differenced AR model.
+        differenced (bool): Whether the model was fitted on differenced data.
+
+    Returns:
+        pd.DataFrame: A DataFrame with integrated coefficients and corresponding lags.
+    """
     if not differenced:
         integrated_coeff = diff_coefficients
     else:
-        integrated_coeff = np.zeros(len(diff_coefficients) + 1)  # Platz für AR(p+1)
-        integrated_coeff[0] = diff_coefficients[0]  # Intercept bleibt gleich
+        integrated_coeff = np.zeros(len(diff_coefficients) + 1)
+        integrated_coeff[0] = diff_coefficients[0]
 
-        # Erster AR-Koeffizient (vom differenzierten Model)
         integrated_coeff[1] = 1 + diff_coefficients[1]
 
-        # Nachfolgende AR-Koeffizienten
         for i in range(2, len(diff_coefficients)):
             integrated_coeff[i] = diff_coefficients[i - 1] - diff_coefficients[i]
 
-        # Zusätzliches Lag (durch Integration)
         integrated_coeff[-1] = -diff_coefficients[-1]
 
-    # DataFrame für besseren Überblick
     integrated_coeff_df = pd.DataFrame(
         {
             "coefficient": integrated_coeff,
@@ -71,9 +120,24 @@ def _integrate_ar_coefficients(
     return integrated_coeff_df
 
 
-def fit_ar_model(df: pd.DataFrame, column: str = "close_price", p: int = 3) -> dict:
-    """Fitte ein AR(p)-Modell und speichere differenzierte & originale Koeffizienten."""
-    is_stationary, p_value = _check_stationarity(df, column)
+def fit_ar_model(df: pd.DataFrame, column: str = "close_price", p: int = 1) -> dict:
+    """Fit an autoregressive (AR) model of order p.
+
+    Args:
+        df (pd.DataFrame): The DataFrame containing the time series data.
+        column (str, optional): The target column to model. Defaults to "close_price".
+        p (int, optional): The order of the AR model. Defaults to 1.
+
+    Returns:
+        dict: A dictionary containing:
+            - "coefficients" (np.ndarray): Estimated coefficients of the AR(p) model.
+            - "integrated_coefficients" (pd.DataFrame): Integrated coefficients.
+            - "lag_order" (int): The order of the AR model (p).
+            - "p_value" (float): The p-value from the stationarity test.
+            - "differenced" (bool): Whether the series was differenced before fitting.
+
+    """
+    is_stationary, p_value, test_statistic_adf = _check_stationarity(df, column)
     differenced = False
 
     if not is_stationary: