[enhancement] Add conditional seasonality as a new feature (#1067)

* Extended config_seasonality with condition_name

*Indicates whether conditional seasonality is applied or not - Type: boolean

* Integrated condition_name into the workflow

*Mask out all the seasonality features that do not belong to the conditioned seasonality

*Ensured conditional seasonality is part of the model's data_params

*Integrated check in check_dataframe to rule out false column names for conditional seasonality

* Added conditional seasonality to test_custom_seasons() in test_integration.py

* Fixed Linters issues

* Fixed Linters issues

* Allow binary input and floats between 0 and 1 in conditional columns

* Allow binary input and floats between 0 and 1 in conditional columns

Author: judussoari

neuralprophet/configure.py

@@ -40,6 +40,7 @@ def init_data_params(
         config_lagged_regressors: Optional[ConfigLaggedRegressors] = None,
         config_regressors=None,
         config_events: Optional[ConfigEvents] = None,
+        config_seasonality: Optional[ConfigSeasonality] = None,
     ):
         if len(df["ID"].unique()) == 1:
             if not self.global_normalization:
@@ -51,6 +52,7 @@ def init_data_params(
             config_lagged_regressors=config_lagged_regressors,
             config_regressors=config_regressors,
             config_events=config_events,
+            config_seasonality=config_seasonality,
             global_normalization=self.global_normalization,
             global_time_normalization=self.global_normalization,
         )
@@ -303,6 +305,7 @@ class Season:
     resolution: int
     period: float
     arg: np_types.SeasonalityArgument
+    condition_name: Optional[str]
 
 
 @dataclass
@@ -315,16 +318,17 @@ class ConfigSeasonality:
     daily_arg: np_types.SeasonalityArgument = "auto"
     periods: OrderedDict = field(init=False)  # contains SeasonConfig objects
     global_local: np_types.SeasonGlobalLocalMode = "local"
+    condition_name: Optional[str] = None
 
     def __post_init__(self):
         if self.reg_lambda > 0 and self.computation == "fourier":
             log.info("Note: Fourier-based seasonality regularization is experimental.")
             self.reg_lambda = 0.001 * self.reg_lambda
         self.periods = OrderedDict(
             {
-                "yearly": Season(resolution=6, period=365.25, arg=self.yearly_arg),
-                "weekly": Season(resolution=3, period=7, arg=self.weekly_arg),
-                "daily": Season(resolution=6, period=1, arg=self.daily_arg),
+                "yearly": Season(resolution=6, period=365.25, arg=self.yearly_arg, condition_name=None),
+                "weekly": Season(resolution=3, period=7, arg=self.weekly_arg, condition_name=None),
+                "daily": Season(resolution=6, period=1, arg=self.daily_arg, condition_name=None),
             }
         )
 
@@ -333,8 +337,8 @@ def __post_init__(self):
             log.error("Invalid global_local mode '{}'. Set to 'global'".format(self.global_local))
             self.global_local = "global"
 
-    def append(self, name, period, resolution, arg):
-        self.periods[name] = Season(resolution=resolution, period=period, arg=arg)
+    def append(self, name, period, resolution, arg, condition_name):
+        self.periods[name] = Season(resolution=resolution, period=period, arg=arg, condition_name=condition_name)
 
 
 @dataclass

neuralprophet/df_utils.py

@@ -10,7 +10,7 @@
 import pandas as pd
 
 if TYPE_CHECKING:
-    from neuralprophet.configure import ConfigEvents, ConfigLaggedRegressors
+    from neuralprophet.configure import ConfigEvents, ConfigLaggedRegressors, ConfigSeasonality
 
 
 log = logging.getLogger("NP.df_utils")
@@ -141,6 +141,7 @@ def data_params_definition(
     config_lagged_regressors: Optional[ConfigLaggedRegressors] = None,
     config_regressors=None,
     config_events: Optional[ConfigEvents] = None,
+    config_seasonality: Optional[ConfigSeasonality] = None,
 ):
     """
     Initialize data scaling values.
@@ -178,6 +179,8 @@ def data_params_definition(
         extra regressors (with known future values) with sub_parameters normalize (bool)
     config_events : configure.ConfigEvents
         user specified events configs
+    config_seasonality : configure.ConfigSeasonality
+        user specified seasonality configs
 
     Returns
     -------
@@ -221,6 +224,13 @@ def data_params_definition(
             if event not in df.columns:
                 raise ValueError(f"Event {event} not found in DataFrame.")
             data_params[event] = ShiftScale()
+    if config_seasonality is not None:
+        for season in config_seasonality.periods:
+            condition_name = config_seasonality.periods[season].condition_name
+            if condition_name is not None:
+                if condition_name not in df.columns:
+                    raise ValueError(f"Seasonality condition {condition_name} not found in DataFrame.")
+                data_params[condition_name] = ShiftScale()
     return data_params
 
 
@@ -230,6 +240,7 @@ def init_data_params(
     config_lagged_regressors: Optional[ConfigLaggedRegressors] = None,
     config_regressors=None,
     config_events: Optional[ConfigEvents] = None,
+    config_seasonality: Optional[ConfigSeasonality] = None,
     global_normalization=False,
     global_time_normalization=False,
 ):
@@ -265,6 +276,8 @@ def init_data_params(
             extra regressors (with known future values)
         config_events : configure.ConfigEvents
             user specified events configs
+        config_seasonality : configure.ConfigSeasonality
+            user specified seasonality configs
         global_normalization : bool
 
             ``True``: sets global modeling training with global normalization
@@ -289,7 +302,7 @@ def init_data_params(
     df, _, _, _ = prep_or_copy_df(df)
     df_merged = df.copy(deep=True).drop("ID", axis=1)
     global_data_params = data_params_definition(
-        df_merged, normalize, config_lagged_regressors, config_regressors, config_events
+        df_merged, normalize, config_lagged_regressors, config_regressors, config_events, config_seasonality
     )
     if global_normalization:
         log.debug(
@@ -300,7 +313,7 @@ def init_data_params(
     for df_name, df_i in df.groupby("ID"):
         df_i.drop("ID", axis=1, inplace=True)
         local_data_params[df_name] = data_params_definition(
-            df_i, normalize, config_lagged_regressors, config_regressors, config_events
+            df_i, normalize, config_lagged_regressors, config_regressors, config_events, config_seasonality
         )
         if global_time_normalization:
             # Overwrite local time normalization data_params with global values (pointer)
@@ -387,7 +400,7 @@ def normalize(df, data_params):
     return df
 
 
-def check_single_dataframe(df, check_y, covariates, regressors, events):
+def check_single_dataframe(df, check_y, covariates, regressors, events, seasonalities):
     """Performs basic data sanity checks and ordering
     as well as prepare dataframe for fitting or predicting.
 
@@ -403,6 +416,8 @@ def check_single_dataframe(df, check_y, covariates, regressors, events):
             regressor column names
         events : list or dict
             event column names
+        seasonalities : list or dict
+            seasonalities column names
 
     Returns
     -------
@@ -451,6 +466,13 @@ def check_single_dataframe(df, check_y, covariates, regressors, events):
             columns.extend(events)
         else:  # treat as dict
             columns.extend(events.keys())
+    if seasonalities is not None:
+        for season in seasonalities.periods:
+            condition_name = seasonalities.periods[season].condition_name
+            if condition_name is not None:
+                if not df[condition_name].isin([True, False]).all() and not df[condition_name].between(0, 1).all():
+                    raise ValueError(f"Condition column {condition_name} must be boolean or numeric between 0 and 1.")
+                columns.append(condition_name)
     for name in columns:
         if name not in df:
             raise ValueError(f"Column {name!r} missing from dataframe")
@@ -470,7 +492,7 @@ def check_single_dataframe(df, check_y, covariates, regressors, events):
     return df, regressors_to_remove
 
 
-def check_dataframe(df, check_y=True, covariates=None, regressors=None, events=None):
+def check_dataframe(df, check_y=True, covariates=None, regressors=None, events=None, seasonalities=None):
     """Performs basic data sanity checks and ordering,
     as well as prepare dataframe for fitting or predicting.
 
@@ -487,6 +509,8 @@ def check_dataframe(df, check_y=True, covariates=None, regressors=None, events=N
             regressor column names
         events : list or dict
             event column names
+        seasonalities : list or dict
+            seasonalities column names
 
     Returns
     -------
@@ -497,7 +521,7 @@ def check_dataframe(df, check_y=True, covariates=None, regressors=None, events=N
     checked_df = pd.DataFrame()
     regressors_to_remove = []
     for df_name, df_i in df.groupby("ID"):
-        df_aux, reg = check_single_dataframe(df_i, check_y, covariates, regressors, events)
+        df_aux, reg = check_single_dataframe(df_i, check_y, covariates, regressors, events, seasonalities)
         df_aux = df_aux.copy(deep=True)
         if len(reg) > 0:
             regressors_to_remove.append(*reg)

neuralprophet/forecaster.py

@@ -409,6 +409,7 @@ def __init__(
             weekly_arg=weekly_seasonality,
             daily_arg=daily_seasonality,
             global_local=season_global_local,
+            condition_name=None,
         )
 
         # Events
@@ -625,13 +626,18 @@ def add_country_holidays(self, country_name, lower_window=0, upper_window=0, reg
         self.config_country_holidays.init_holidays()
         return self
 
-    def add_seasonality(self, name, period, fourier_order):
+    def add_seasonality(self, name, period, fourier_order, condition_name=None):
         """Add a seasonal component with specified period, number of Fourier components, and regularization.
 
         Increasing the number of Fourier components allows the seasonality to change more quickly
         (at risk of overfitting).
         Note: regularization and mode (additive/multiplicative) are set in the main init.
 
+        If condition_name is provided, the dataframe passed to `fit` and
+        `predict` should have a column with the specified condition_name
+        containing only zeros and ones, deciding when to apply seasonality.
+        Floats between 0 and 1 can be used to apply seasonality partially.
+
         Parameters
         ----------
             name : string
@@ -640,17 +646,22 @@ def add_seasonality(self, name, period, fourier_order):
                 number of days in one period.
             fourier_order : int
                 number of Fourier components to use.
-
+            condition_name : string
+                string name of the seasonality condition.
         """
         if self.fitted:
             raise Exception("Seasonality must be added prior to model fitting.")
         if name in ["daily", "weekly", "yearly"]:
             log.error("Please use inbuilt daily, weekly, or yearly seasonality or set another name.")
         # Do not Allow overwriting built-in seasonalities
         self._validate_column_name(name, seasons=True)
+        if condition_name is not None:
+            self._validate_column_name(condition_name)
         if fourier_order <= 0:
             raise ValueError("Fourier Order must be > 0")
-        self.config_seasonality.append(name=name, period=period, resolution=fourier_order, arg="custom")
+        self.config_seasonality.append(
+            name=name, period=period, resolution=fourier_order, condition_name=condition_name, arg="custom"
+        )
         return self
 
     def fit(
@@ -2404,6 +2415,7 @@ def _check_dataframe(self, df, check_y=True, exogenous=True):
             covariates=self.config_lagged_regressors if exogenous else None,
             regressors=self.config_regressors if exogenous else None,
             events=self.config_events if exogenous else None,
+            seasonalities=self.config_seasonality if exogenous else None,
         )
         for reg in regressors_to_remove:
             log.warning(f"Removing regressor {reg} because it is not present in the data.")
@@ -2507,6 +2519,7 @@ def _init_train_loader(self, df, num_workers=0):
             config_lagged_regressors=self.config_lagged_regressors,
             config_regressors=self.config_regressors,
             config_events=self.config_events,
+            config_seasonality=self.config_seasonality,
         )
 
         df = self._normalize(df)

neuralprophet/time_dataset.py

@@ -288,7 +288,7 @@ def _stride_time_features_for_forecasts(x):
     inputs["time"] = time
 
     if config_seasonality is not None:
-        seasonalities = seasonal_features_from_dates(df["ds"], config_seasonality)
+        seasonalities = seasonal_features_from_dates(df, config_seasonality)
         for name, features in seasonalities.items():
             if max_lags == 0:
                 seasonalities[name] = np.expand_dims(features, axis=1)
@@ -608,15 +608,15 @@ def make_regressors_features(df, config_regressors):
     return additive_regressors, multiplicative_regressors
 
 
-def seasonal_features_from_dates(dates, config_seasonality: configure.ConfigSeasonality):
+def seasonal_features_from_dates(df, config_seasonality: configure.ConfigSeasonality):
     """Dataframe with seasonality features.
 
     Includes seasonality features, holiday features, and added regressors.
 
     Parameters
     ----------
-        dates : pd.Series
-            With dates for computing seasonality features
+        df : pd.DataFrame
+            Dataframe with all values
         config_seasonality : configure.ConfigSeasonality
             Configuration for seasonalities
 
@@ -626,6 +626,7 @@ def seasonal_features_from_dates(dates, config_seasonality: configure.ConfigSeas
             Dictionary with keys for each period name containing an np.array
             with the respective regression features. each with dims: (len(dates), 2*fourier_order)
     """
+    dates = df["ds"]
     assert len(dates.shape) == 1
     seasonalities = OrderedDict({})
     # Seasonality features
@@ -639,5 +640,7 @@ def seasonal_features_from_dates(dates, config_seasonality: configure.ConfigSeas
                 )
             else:
                 raise NotImplementedError
+            if period.condition_name is not None:
+                features = features * df[period.condition_name].values[:, np.newaxis]
             seasonalities[name] = features
     return seasonalities

tests/test_integration.py

@@ -236,10 +236,17 @@ def test_custom_seasons():
         batch_size=BATCH_SIZE,
         learning_rate=LR,
     )
-    m = m.add_seasonality(name="quarterly", period=90, fourier_order=5)
+    # conditional seasonality
+    df["ds"] = pd.to_datetime(df["ds"])
+    df["on_season"] = df["ds"].apply(lambda x: x.month in [9, 10, 11, 12, 1])
+    df["off_season"] = df["ds"].apply(lambda x: x.month not in [9, 10, 11, 12, 1])
+    m.add_seasonality(name="on_season", period=7, fourier_order=3, condition_name="on_season")
+    m.add_seasonality(name="off_season", period=7, fourier_order=3, condition_name="off_season")
     log.debug(f"seasonalities: {m.config_seasonality.periods}")
     metrics_df = m.fit(df, freq="D")
     future = m.make_future_dataframe(df, n_historic_predictions=365, periods=365)
+    future["on_season"] = future["ds"].apply(lambda x: x.month in [9, 10, 11, 12, 1])
+    future["off_season"] = future["ds"].apply(lambda x: x.month not in [9, 10, 11, 12, 1])
     forecast = m.predict(df=future)
     log.debug(f"season params: {m.model.season_params.items()}")
     if PLOT: