Create BaseModel.py

BaseModel.py

@@ -0,0 +1,83 @@
+import pandas as pd
+from prophet import Prophet
+from prophet.diagnostics import cross_validation, performance_metrics
+
+class GeoMagModel:
+    """
+    A class for geomagnetic data forecasting using Prophet.
+    """
+
+    def __init__(self, changepoint_prior_scale=0.1, weekly_seasonality=True):
+        """
+        Initialize the GeoMagModel with Prophet configuration.
+        
+        Args:
+            changepoint_prior_scale (float): Flexibility of changepoint detection.
+            weekly_seasonality (bool): Whether to include weekly seasonality.
+        """
+        self.changepoint_prior_scale = changepoint_prior_scale
+        self.weekly_seasonality = weekly_seasonality
+        self.model = None
+
+    @staticmethod
+    def prepare_for_prophet(df):
+        """
+        Prepare the DataFrame for Prophet by renaming columns as required.
+        
+        Args:
+            df (pd.DataFrame): The cleaned DataFrame with 'timestamp' and 'Dst' columns.
+        
+        Returns:
+            pd.DataFrame: A DataFrame with 'ds' (timestamp) and 'y' (Dst) columns for Prophet.
+        """
+        return df.rename(columns={'timestamp': 'ds', 'Dst': 'y'})
+
+    def train(self, df):
+        """
+        Train the Prophet model on the given DataFrame.
+        
+        Args:
+            df (pd.DataFrame): The DataFrame prepared for Prophet.
+        """
+        df = self.prepare_for_prophet(df)
+        self.model = Prophet(interval_width=0.70, changepoint_prior_scale=self.changepoint_prior_scale)
+
+        if self.weekly_seasonality:
+            self.model.add_seasonality(name='weekly', period=7, fourier_order=5, prior_scale=10)
+
+        self.model.fit(df)
+
+    def forecast(self, periods=6, freq='h'):
+        """
+        Make a future forecast with the trained model.
+        
+        Args:
+            periods (int): The number of future periods to forecast. Defaults to 6.
+            freq (str): The frequency of the forecast ('h' for hours). Defaults to 'h'.
+        
+        Returns:
+            pd.DataFrame: The forecast DataFrame with 'ds', 'yhat', 'yhat_lower', and 'yhat_upper' columns.
+        """
+        if self.model is None:
+            raise ValueError("Model has not been trained. Call train() before forecast().")
+        
+        future_dates = self.model.make_future_dataframe(periods=periods, freq=freq)
+        forecast = self.model.predict(future_dates)
+        return forecast[['ds', 'yhat', 'yhat_lower', 'yhat_upper']]
+
+    def cross_validate(self, initial='384 hours', period='12 hours', horizon='6 hours'):
+        """
+        Perform cross-validation on the trained Prophet model.
+        
+        Args:
+            initial (str): Initial training period for cross-validation.
+            period (str): Frequency of making predictions.
+            horizon (str): Forecast horizon for each prediction.
+        
+        Returns:
+            pd.DataFrame: Cross-validation results including metrics.
+        """
+        if self.model is None:
+            raise ValueError("Model has not been trained. Call train() before cross_validate().")
+        
+        return cross_validation(self.model, initial=initial, period=period, horizon=horizon)