|
import pandas as pd |
|
from prophet import Prophet |
|
from prophet.diagnostics import cross_validation |
|
|
|
class GeoMagModel: |
|
""" |
|
A class for forecasting geomagnetic data using the Prophet model. |
|
""" |
|
|
|
def __init__(self, changepoint_prior_scale=0.1, weekly_seasonality=True): |
|
""" |
|
Initialize the GeoMagModel with Prophet configuration. |
|
|
|
Args: |
|
changepoint_prior_scale (float): Controls 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_data(df): |
|
""" |
|
Prepare the DataFrame for Prophet by renaming columns. |
|
|
|
Args: |
|
df (pd.DataFrame): Input DataFrame with 'timestamp' and 'Dst' columns. |
|
|
|
Returns: |
|
pd.DataFrame: A DataFrame with 'ds' (timestamp) and 'y' (Dst) columns. |
|
""" |
|
if not {'timestamp', 'Dst'}.issubset(df.columns): |
|
raise ValueError("DataFrame must contain 'timestamp' and 'Dst' columns.") |
|
return df.rename(columns={'timestamp': 'ds', 'Dst': 'y'}) |
|
|
|
def train(self, df): |
|
""" |
|
Train the Prophet model using the provided DataFrame. |
|
|
|
Args: |
|
df (pd.DataFrame): DataFrame prepared for Prophet with 'ds' and 'y' columns. |
|
""" |
|
df = self.prepare_data(df) |
|
|
|
self.model = Prophet( |
|
interval_width=0.7, |
|
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=1, freq='H'): |
|
""" |
|
Generate future forecasts with the trained model. |
|
|
|
Args: |
|
periods (int): Number of future periods to forecast. Default is 1. |
|
freq (str): Frequency of the forecast (e.g., 'H' for hours). Default is 'H'. |
|
|
|
Returns: |
|
pd.DataFrame: Forecast DataFrame with columns: 'ds', 'yhat', 'yhat_lower', 'yhat_upper'. |
|
""" |
|
if self.model is None: |
|
raise ValueError("Model is not 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='36 hours', period='12 hours', horizon='1 hours'): |
|
""" |
|
Perform cross-validation on the trained 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 with metrics. |
|
""" |
|
if self.model is None: |
|
raise ValueError("Model is not trained. Call train() before cross_validate().") |
|
|
|
return cross_validation(self.model, initial=initial, period=period, horizon=horizon) |
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|