From b079fa594efe0c7b0d14ef805f918b951af7fd54 Mon Sep 17 00:00:00 2001
From: ourownstory <ourownstory@users.noreply.github.com>
Date: Wed, 11 Sep 2024 16:37:49 -0700
Subject: [PATCH 01/30] remove stackers from TimeNet init

---
 neuralprophet/time_net.py | 16 ++++++++++++----
 1 file changed, 12 insertions(+), 4 deletions(-)

diff --git a/neuralprophet/time_net.py b/neuralprophet/time_net.py
index 096da0e67..6dbb0afca 100644
--- a/neuralprophet/time_net.py
+++ b/neuralprophet/time_net.py
@@ -62,10 +62,10 @@ def __init__(
         num_seasonalities_modelled: int = 1,
         num_seasonalities_modelled_dict: dict = None,
         meta_used_in_model: bool = False,
-        train_components_stacker: Optional[ComponentStacker] = None,
-        val_components_stacker: Optional[ComponentStacker] = None,
-        test_components_stacker: Optional[ComponentStacker] = None,
-        predict_components_stacker: Optional[ComponentStacker] = None,
+        # train_components_stacker: Optional[ComponentStacker] = None,
+        # val_components_stacker: Optional[ComponentStacker] = None,
+        # test_components_stacker: Optional[ComponentStacker] = None,
+        # predict_components_stacker: Optional[ComponentStacker] = None,
     ):
         """
         Parameters
@@ -143,6 +143,14 @@ def __init__(
         # General
         self.config_model = config_model
         self.config_model.n_forecasts = n_forecasts
+        # refactor components stacker to be a dict of stackers for train, val, test, predict
+        # self.components_stacker = {
+        #     "train": train_components_stacker,
+        #     "val": val_components_stacker,
+        #     "test": test_components_stacker,
+        #     "predict": predict_components_stacker,
+        # }
+
         self.train_components_stacker = train_components_stacker
         self.val_components_stacker = val_components_stacker
         self.test_components_stacker = test_components_stacker

From 299593c00c5c81a8c3f18cc2282ecd12b4f1bb15 Mon Sep 17 00:00:00 2001
From: ourownstory <ourownstory@users.noreply.github.com>
Date: Wed, 11 Sep 2024 16:42:59 -0700
Subject: [PATCH 02/30] refactor stackers to dict

---
 neuralprophet/forecaster.py |  4 +--
 neuralprophet/time_net.py   | 54 ++++++++++++++++++++-----------------
 2 files changed, 31 insertions(+), 27 deletions(-)

diff --git a/neuralprophet/forecaster.py b/neuralprophet/forecaster.py
index 22aa9c58f..8650456ec 100644
--- a/neuralprophet/forecaster.py
+++ b/neuralprophet/forecaster.py
@@ -2128,7 +2128,7 @@ def predict_seasonal_components(self, df: pd.DataFrame, quantile: float = 0.5):
         prev_n_lags = self.config_ar.n_lags
         prev_max_lags = self.config_model.max_lags
         prev_n_forecasts = self.config_model.n_forecasts
-        prev_predict_components_stacker = self.model.predict_components_stacker
+        prev_predict_components_stacker = self.model.components_stacker["predict"]
 
         self.config_model.max_lags = 0
         self.config_ar.n_lags = 0
@@ -2198,7 +2198,7 @@ def predict_seasonal_components(self, df: pd.DataFrame, quantile: float = 0.5):
         self.config_ar.n_lags = prev_n_lags
         self.config_model.max_lags = prev_max_lags
         self.config_model.n_forecasts = prev_n_forecasts
-        self.model.predict_components_stacker = prev_predict_components_stacker
+        self.model.components_stacker["predict"] = prev_predict_components_stacker
 
         return df
 
diff --git a/neuralprophet/time_net.py b/neuralprophet/time_net.py
index 6dbb0afca..1647e3c56 100644
--- a/neuralprophet/time_net.py
+++ b/neuralprophet/time_net.py
@@ -144,17 +144,21 @@ def __init__(
         self.config_model = config_model
         self.config_model.n_forecasts = n_forecasts
         # refactor components stacker to be a dict of stackers for train, val, test, predict
-        # self.components_stacker = {
-        #     "train": train_components_stacker,
-        #     "val": val_components_stacker,
-        #     "test": test_components_stacker,
-        #     "predict": predict_components_stacker,
-        # }
-
-        self.train_components_stacker = train_components_stacker
-        self.val_components_stacker = val_components_stacker
-        self.test_components_stacker = test_components_stacker
-        self.predict_components_stacker = predict_components_stacker
+        self.components_stacker = {
+            "train": None,
+            "val": None,
+            "test": None,
+            "predict": None,
+        }
+
+        # self.train_components_stacker = train_components_stacker
+        # self.val_components_stacker = val_components_stacker
+        # self.test_components_stacker = test_components_stacker
+        # self.predict_components_stacker = predict_components_stacker
+        # self.components_stacker["train"] = None
+        # self.components_stacker["val"] = None
+        # self.components_stacker["test"] = None
+        # self.components_stacker["predict"] = None
 
         # Lightning Config
         self.config_train = config_train
@@ -327,13 +331,13 @@ def ar_weights(self) -> torch.Tensor:
 
     def set_components_stacker(self, components_stacker, mode):
         if mode == "train":
-            self.train_components_stacker = components_stacker
+            self.components_stacker["train"] = components_stacker
         if mode == "val":
-            self.val_components_stacker = components_stacker
+            self.components_stacker["val"] = components_stacker
         if mode == "test":
-            self.test_components_stacker = components_stacker
+            self.components_stacker["test"] = components_stacker
         if mode == "predict":
-            self.predict_components_stacker = components_stacker
+            self.components_stacker["predict"] = components_stacker
 
     def get_covar_weights(self, covar_input=None) -> torch.Tensor:
         """
@@ -781,15 +785,15 @@ def training_step(self, batch, batch_idx):
 
         epoch_float = self.trainer.current_epoch + batch_idx / float(self.train_steps_per_epoch)
         self.train_progress = epoch_float / float(self.config_train.epochs)
-        targets = self.train_components_stacker.unstack_component("targets", batch_tensor=inputs_tensor)
-        time = self.train_components_stacker.unstack_component("time", batch_tensor=inputs_tensor)
+        targets = self.components_stacker["train"].unstack_component("targets", batch_tensor=inputs_tensor)
+        time = self.components_stacker["train"].unstack_component("time", batch_tensor=inputs_tensor)
         # Global-local
         if self.meta_used_in_model:
             meta_name_tensor = torch.tensor([self.id_dict[i] for i in meta["df_name"]], device=self.device)
         else:
             meta_name_tensor = None
         # Run forward calculation
-        predicted, _ = self.forward(inputs_tensor, self.train_components_stacker, meta_name_tensor)
+        predicted, _ = self.forward(inputs_tensor, self.components_stacker["train"], meta_name_tensor)
         # Store predictions in self for later network visualization
         self.train_epoch_prediction = predicted
         # Calculate loss
@@ -826,15 +830,15 @@ def training_step(self, batch, batch_idx):
 
     def validation_step(self, batch, batch_idx):
         inputs_tensor, meta = batch
-        targets = self.val_components_stacker.unstack_component("targets", batch_tensor=inputs_tensor)
-        time = self.val_components_stacker.unstack_component("time", batch_tensor=inputs_tensor)
+        targets = self.components_stacker["val"].unstack_component("targets", batch_tensor=inputs_tensor)
+        time = self.components_stacker["val"].unstack_component("time", batch_tensor=inputs_tensor)
         # Global-local
         if self.meta_used_in_model:
             meta_name_tensor = torch.tensor([self.id_dict[i] for i in meta["df_name"]], device=self.device)
         else:
             meta_name_tensor = None
         # Run forward calculation
-        predicted, _ = self.forward(inputs_tensor, self.val_components_stacker, meta_name_tensor)
+        predicted, _ = self.forward(inputs_tensor, self.components_stacker["val"], meta_name_tensor)
         # Calculate loss
         loss, reg_loss = self.loss_func(time, predicted, targets)
         # Metrics
@@ -848,15 +852,15 @@ def validation_step(self, batch, batch_idx):
 
     def test_step(self, batch, batch_idx):
         inputs_tensor, meta = batch
-        targets = self.test_components_stacker.unstack_component("targets", batch_tensor=inputs_tensor)
-        time = self.test_components_stacker.unstack_component("time", batch_tensor=inputs_tensor)
+        targets = self.components_stacker["test"].unstack_component("targets", batch_tensor=inputs_tensor)
+        time = self.components_stacker["test"].unstack_component("time", batch_tensor=inputs_tensor)
         # Global-local
         if self.meta_used_in_model:
             meta_name_tensor = torch.tensor([self.id_dict[i] for i in meta["df_name"]], device=self.device)
         else:
             meta_name_tensor = None
         # Run forward calculation
-        predicted, _ = self.forward(inputs_tensor, self.test_components_stacker, meta_name_tensor)
+        predicted, _ = self.forward(inputs_tensor, self.components_stacker["test"], meta_name_tensor)
         # Calculate loss
         loss, reg_loss = self.loss_func(time, predicted, targets)
         # Metrics
@@ -880,7 +884,7 @@ def predict_step(self, batch, batch_idx, dataloader_idx=0):
         # Run forward calculation
         prediction, components = self.forward(
             inputs_tensor,
-            self.predict_components_stacker,
+            self.components_stacker["predict"],
             meta_name_tensor,
             self.compute_components_flag,
             predict_mode=True,

From 6f9a0e001b456494b2c56c9125d907453f9ff2d1 Mon Sep 17 00:00:00 2001
From: ourownstory <ourownstory@users.noreply.github.com>
Date: Wed, 11 Sep 2024 16:47:25 -0700
Subject: [PATCH 03/30] simplify and document set_compunents_stacker

---
 neuralprophet/time_net.py | 22 +++++++++++++---------
 1 file changed, 13 insertions(+), 9 deletions(-)

diff --git a/neuralprophet/time_net.py b/neuralprophet/time_net.py
index 1647e3c56..ae8efa73c 100644
--- a/neuralprophet/time_net.py
+++ b/neuralprophet/time_net.py
@@ -329,15 +329,19 @@ def ar_weights(self) -> torch.Tensor:
             if isinstance(layer, nn.Linear):
                 return layer.weight
 
-    def set_components_stacker(self, components_stacker, mode):
-        if mode == "train":
-            self.components_stacker["train"] = components_stacker
-        if mode == "val":
-            self.components_stacker["val"] = components_stacker
-        if mode == "test":
-            self.components_stacker["test"] = components_stacker
-        if mode == "predict":
-            self.components_stacker["predict"] = components_stacker
+    def set_components_stacker(self, stacker, mode):
+        """Set the components stacker for the given mode.
+        Parameters
+        ----------
+        stacker : ComponentStacker
+            The components stacker to be set.
+        mode : str
+            The mode for which the components stacker is to be set
+            options: ["train", "val", "test", "predict"]
+        """
+        modes = ["train", "val", "test", "predict"]
+        assert mode in modes, f"mode must be one of {modes}"
+        self.components_stacker[mode] = stacker
 
     def get_covar_weights(self, covar_input=None) -> torch.Tensor:
         """

From 7580c711b7cdd74f1fdd6e5ccedbe846a0ca3cb9 Mon Sep 17 00:00:00 2001
From: ourownstory <ourownstory@users.noreply.github.com>
Date: Wed, 11 Sep 2024 16:49:16 -0700
Subject: [PATCH 04/30] refactor set_components_stacker arg to stacker

---
 neuralprophet/forecaster.py | 4 ++--
 neuralprophet/time_net.py   | 2 +-
 2 files changed, 3 insertions(+), 3 deletions(-)

diff --git a/neuralprophet/forecaster.py b/neuralprophet/forecaster.py
index 8650456ec..1dcd278c1 100644
--- a/neuralprophet/forecaster.py
+++ b/neuralprophet/forecaster.py
@@ -1275,9 +1275,9 @@ def fit(
         if not self.fitted:
             self.model = self._init_model()
 
-        self.model.set_components_stacker(components_stacker=train_components_stacker, mode="train")
+        self.model.set_components_stacker(stacker=train_components_stacker, mode="train")
         if validation_enabled:
-            self.model.set_components_stacker(components_stacker=val_components_stacker, mode="val")
+            self.model.set_components_stacker(stacker=val_components_stacker, mode="val")
 
         # Find suitable learning rate if not set
         if self.config_train.learning_rate is None:
diff --git a/neuralprophet/time_net.py b/neuralprophet/time_net.py
index ae8efa73c..a0794ad66 100644
--- a/neuralprophet/time_net.py
+++ b/neuralprophet/time_net.py
@@ -333,7 +333,7 @@ def set_components_stacker(self, stacker, mode):
         """Set the components stacker for the given mode.
         Parameters
         ----------
-        stacker : ComponentStacker
+        components_stacker : ComponentStacker
             The components stacker to be set.
         mode : str
             The mode for which the components stacker is to be set

From 24ac07e4935bed850b312047e43b7312a7811357 Mon Sep 17 00:00:00 2001
From: ourownstory <ourownstory@users.noreply.github.com>
Date: Wed, 11 Sep 2024 16:57:58 -0700
Subject: [PATCH 05/30] add docstring to forward and introduce mode flag
 instead of passing components_stacker

---
 neuralprophet/time_net.py | 47 ++++++++++++++++++++++++++-------------
 1 file changed, 32 insertions(+), 15 deletions(-)

diff --git a/neuralprophet/time_net.py b/neuralprophet/time_net.py
index a0794ad66..e4d475391 100644
--- a/neuralprophet/time_net.py
+++ b/neuralprophet/time_net.py
@@ -536,14 +536,29 @@ def forward_covar_net(self, covariates):
     def forward(
         self,
         input_tensor: torch.Tensor,
-        components_stacker=ComponentStacker,
+        mode: str,
         meta: Dict = None,
         compute_components_flag: bool = False,
         predict_mode: bool = False,
     ) -> torch.Tensor:
-        """This method defines the model forward pass."""
+        """This method defines the model forward pass.
+        Parameters
+        ----------
+            input_tensor : torch.Tensor
+                Input tensor of dims (batch, n_lags + n_forecasts, n_features)
+            mode : str operation mode ["train", "val", "test", "predict"]
+            meta : dict Static features of the time series
+            compute_components_flag : bool Whether to return the individual components of the model
+            predict_mode : bool Whether the model is in prediction mode
+        Returns
+        -------
+            torch.Tensor Forecast tensor of dims (batch, n_forecasts, n_quantiles)
+            dict of components of the model if compute_components_flag is True,
+                each of dims (batch, n_forecasts, n_quantiles)
 
-        time_input = components_stacker.unstack_component(component_name="time", batch_tensor=input_tensor)
+        """
+
+        time_input = self.components_stacker[mode].unstack_component(component_name="time", batch_tensor=input_tensor)
         # Handle meta argument
         if meta is None and self.meta_used_in_model:
             name_id_dummy = self.id_list[0]
@@ -573,7 +588,7 @@ def forward(
         # Unpack and process seasonalities
         seasonalities_input = None
         if self.config_seasonality and self.config_seasonality.periods:
-            seasonalities_input = components_stacker.unstack_component(
+            seasonalities_input = self.components_stacker[mode].unstack_component(
                 component_name="seasonalities", batch_tensor=input_tensor
             )
             s = self.seasonality(s=seasonalities_input, meta=meta)
@@ -587,15 +602,15 @@ def forward(
         additive_events_input = None
         multiplicative_events_input = None
         if self.events_dims is not None:
-            if "additive_events" in components_stacker.feature_indices:
-                additive_events_input = components_stacker.unstack_component(
+            if "additive_events" in self.components_stacker[mode].feature_indices:
+                additive_events_input = self.components_stacker[mode].unstack_component(
                     component_name="additive_events", batch_tensor=input_tensor
                 )
                 additive_events = self.scalar_features_effects(additive_events_input, self.event_params["additive"])
                 additive_components_nonstationary += additive_events
                 components["additive_events"] = additive_events
-            if "multiplicative_events" in components_stacker.feature_indices:
-                multiplicative_events_input = components_stacker.unstack_component(
+            if "multiplicative_events" in self.components_stacker[mode].feature_indices:
+                multiplicative_events_input = self.components_stacker[mode].unstack_component(
                     component_name="multiplicative_events", batch_tensor=input_tensor
                 )
                 multiplicative_events = self.scalar_features_effects(
@@ -607,15 +622,15 @@ def forward(
         # Unpack and process regressors
         additive_regressors_input = None
         multiplicative_regressors_input = None
-        if "additive_regressors" in components_stacker.feature_indices:
-            additive_regressors_input = components_stacker.unstack_component(
+        if "additive_regressors" in self.components_stacker[mode].feature_indices:
+            additive_regressors_input = self.components_stacker[mode].unstack_component(
                 component_name="additive_regressors", batch_tensor=input_tensor
             )
             additive_regressors = self.future_regressors(additive_regressors_input, "additive")
             additive_components_nonstationary += additive_regressors
             components["additive_regressors"] = additive_regressors
-        if "multiplicative_regressors" in components_stacker.feature_indices:
-            multiplicative_regressors_input = components_stacker.unstack_component(
+        if "multiplicative_regressors" in self.components_stacker[mode].feature_indices:
+            multiplicative_regressors_input = self.components_stacker[mode].unstack_component(
                 component_name="multiplicative_regressors", batch_tensor=input_tensor
             )
             multiplicative_regressors = self.future_regressors(multiplicative_regressors_input, "multiplicative")
@@ -624,8 +639,10 @@ def forward(
 
         # Unpack and process lags
         lags_input = None
-        if "lags" in components_stacker.feature_indices:
-            lags_input = components_stacker.unstack_component(component_name="lags", batch_tensor=input_tensor)
+        if "lags" in self.components_stacker[mode].feature_indices:
+            lags_input = self.components_stacker[mode].unstack_component(
+                component_name="lags", batch_tensor=input_tensor
+            )
             nonstationary_components = (
                 trend[:, : self.n_lags, 0]
                 + additive_components_nonstationary[:, : self.n_lags, 0]
@@ -639,7 +656,7 @@ def forward(
         # Unpack and process covariates
         covariates_input = None
         if self.config_lagged_regressors and self.config_lagged_regressors.regressors is not None:
-            covariates_input = components_stacker.unstack_component(
+            covariates_input = self.components_stacker[mode].unstack_component(
                 component_name="lagged_regressors", batch_tensor=input_tensor
             )
             covariates = self.forward_covar_net(covariates=covariates_input)

From 37e3ab7b60cdb187a4583a7fe274bfa5e9a86621 Mon Sep 17 00:00:00 2001
From: ourownstory <ourownstory@users.noreply.github.com>
Date: Wed, 11 Sep 2024 17:20:18 -0700
Subject: [PATCH 06/30] refactor include components

---
 neuralprophet/forecaster.py |  4 +--
 neuralprophet/time_net.py   | 54 +++++++++++++------------------------
 2 files changed, 21 insertions(+), 37 deletions(-)

diff --git a/neuralprophet/forecaster.py b/neuralprophet/forecaster.py
index 1dcd278c1..8a06a7af8 100644
--- a/neuralprophet/forecaster.py
+++ b/neuralprophet/forecaster.py
@@ -3002,7 +3002,7 @@ def _predict_raw(self, df, df_name, include_components=False):
         # Pass the include_components flag to the model
         if include_components:
             self.model.set_compute_components(include_components)
-            self.model.set_covar_weights(self.model.get_covar_weights())
+            # self.model.set_covar_weights(self.model.get_covar_weights())
         # Compute the predictions and components (if requested)
         result = self.trainer.predict(self.model, loader)
         # unstack the prediction and components
@@ -3066,7 +3066,7 @@ def _predict_raw(self, df, df_name, include_components=False):
                 elif multiplicative:
                     # output absolute value of respective additive component
                     components[name] = value * trend * scale_y  # type: ignore
-
+            self.model.reset_compute_components()
         else:
             components = None
 
diff --git a/neuralprophet/time_net.py b/neuralprophet/time_net.py
index e4d475391..46fe8c3a7 100644
--- a/neuralprophet/time_net.py
+++ b/neuralprophet/time_net.py
@@ -55,17 +55,12 @@ def __init__(
         n_forecasts: int = 1,
         n_lags: int = 0,
         ar_layers: Optional[List[int]] = [],
-        compute_components_flag: bool = False,
         metrics: Optional[np_types.CollectMetricsMode] = {},
         id_list: List[str] = ["__df__"],
         num_trends_modelled: int = 1,
         num_seasonalities_modelled: int = 1,
         num_seasonalities_modelled_dict: dict = None,
         meta_used_in_model: bool = False,
-        # train_components_stacker: Optional[ComponentStacker] = None,
-        # val_components_stacker: Optional[ComponentStacker] = None,
-        # test_components_stacker: Optional[ComponentStacker] = None,
-        # predict_components_stacker: Optional[ComponentStacker] = None,
     ):
         """
         Parameters
@@ -98,8 +93,6 @@ def __init__(
                 ----
                 The default value is ``[]``, which initializes no hidden layers.
 
-            compute_components_flag : bool
-                Flag whether to compute the components of the model or not.
             metrics : dict
                 Dictionary of torchmetrics to be used during training and for evaluation.
             id_list : list
@@ -143,27 +136,18 @@ def __init__(
         # General
         self.config_model = config_model
         self.config_model.n_forecasts = n_forecasts
-        # refactor components stacker to be a dict of stackers for train, val, test, predict
+
+        # Components stackers to unpack the input tensor
         self.components_stacker = {
             "train": None,
             "val": None,
             "test": None,
             "predict": None,
         }
-
-        # self.train_components_stacker = train_components_stacker
-        # self.val_components_stacker = val_components_stacker
-        # self.test_components_stacker = test_components_stacker
-        # self.predict_components_stacker = predict_components_stacker
-        # self.components_stacker["train"] = None
-        # self.components_stacker["val"] = None
-        # self.components_stacker["test"] = None
-        # self.components_stacker["predict"] = None
-
         # Lightning Config
         self.config_train = config_train
         self.config_normalization = config_normalization
-        self.compute_components_flag = compute_components_flag
+        self.include_components = False  # flag to indicate if we are in include_components mode, set in prodiction mode by set_compute_components
         self.config_model = config_model
 
         # Manual optimization: we are responsible for calling .backward(), .step(), .zero_grad().
@@ -538,8 +522,6 @@ def forward(
         input_tensor: torch.Tensor,
         mode: str,
         meta: Dict = None,
-        compute_components_flag: bool = False,
-        predict_mode: bool = False,
     ) -> torch.Tensor:
         """This method defines the model forward pass.
         Parameters
@@ -548,12 +530,10 @@ def forward(
                 Input tensor of dims (batch, n_lags + n_forecasts, n_features)
             mode : str operation mode ["train", "val", "test", "predict"]
             meta : dict Static features of the time series
-            compute_components_flag : bool Whether to return the individual components of the model
-            predict_mode : bool Whether the model is in prediction mode
         Returns
         -------
             torch.Tensor Forecast tensor of dims (batch, n_forecasts, n_quantiles)
-            dict of components of the model if compute_components_flag is True,
+            dict of components of the model if self.include_components is True,
                 each of dims (batch, n_forecasts, n_quantiles)
 
         """
@@ -673,10 +653,12 @@ def forward(
         prediction = predictions_nonstationary + additive_components
 
         # Correct crossing quantiles
-        prediction_with_quantiles = self._compute_quantile_forecasts_from_diffs(prediction, predict_mode)
+        prediction_with_quantiles = self._compute_quantile_forecasts_from_diffs(
+            diffs=prediction, predict_mode=mode != "train"
+        )
 
         # Compute components if required
-        if compute_components_flag:
+        if self.include_components:
             components = self.compute_components(
                 time_input,
                 seasonalities_input,
@@ -783,8 +765,12 @@ def compute_components(
                     )
         return components
 
-    def set_compute_components(self, compute_components_flag):
-        self.compute_components_flag = compute_components_flag
+    def set_compute_components(self, include_components):
+        self.prev_include_components = self.include_components
+        self.include_components = include_components
+
+    def reset_compute_components(self):
+        self.include_components = self.prev_include_components
 
     def loss_func(self, time, predicted, targets):
         loss = None
@@ -814,7 +800,7 @@ def training_step(self, batch, batch_idx):
         else:
             meta_name_tensor = None
         # Run forward calculation
-        predicted, _ = self.forward(inputs_tensor, self.components_stacker["train"], meta_name_tensor)
+        predicted, _ = self.forward(inputs_tensor, mode="train", meta=meta_name_tensor)
         # Store predictions in self for later network visualization
         self.train_epoch_prediction = predicted
         # Calculate loss
@@ -859,7 +845,7 @@ def validation_step(self, batch, batch_idx):
         else:
             meta_name_tensor = None
         # Run forward calculation
-        predicted, _ = self.forward(inputs_tensor, self.components_stacker["val"], meta_name_tensor)
+        predicted, _ = self.forward(inputs_tensor, mode="val", meta=meta_name_tensor)
         # Calculate loss
         loss, reg_loss = self.loss_func(time, predicted, targets)
         # Metrics
@@ -881,7 +867,7 @@ def test_step(self, batch, batch_idx):
         else:
             meta_name_tensor = None
         # Run forward calculation
-        predicted, _ = self.forward(inputs_tensor, self.components_stacker["test"], meta_name_tensor)
+        predicted, _ = self.forward(inputs_tensor, mode="test", meta=meta_name_tensor)
         # Calculate loss
         loss, reg_loss = self.loss_func(time, predicted, targets)
         # Metrics
@@ -905,10 +891,8 @@ def predict_step(self, batch, batch_idx, dataloader_idx=0):
         # Run forward calculation
         prediction, components = self.forward(
             inputs_tensor,
-            self.components_stacker["predict"],
-            meta_name_tensor,
-            self.compute_components_flag,
-            predict_mode=True,
+            mode="predict",
+            meta=meta_name_tensor,
         )
         return prediction, components
 

From 27615cc1cbe807c7a3f8dc3da1ade7733f7ddbd3 Mon Sep 17 00:00:00 2001
From: ourownstory <ourownstory@users.noreply.github.com>
Date: Wed, 11 Sep 2024 17:22:55 -0700
Subject: [PATCH 07/30] fix covar_weights

---
 neuralprophet/forecaster.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/neuralprophet/forecaster.py b/neuralprophet/forecaster.py
index 8a06a7af8..3e18b8bfd 100644
--- a/neuralprophet/forecaster.py
+++ b/neuralprophet/forecaster.py
@@ -3002,7 +3002,7 @@ def _predict_raw(self, df, df_name, include_components=False):
         # Pass the include_components flag to the model
         if include_components:
             self.model.set_compute_components(include_components)
-            # self.model.set_covar_weights(self.model.get_covar_weights())
+            self.model.set_covar_weights(self.model.get_covar_weights())
         # Compute the predictions and components (if requested)
         result = self.trainer.predict(self.model, loader)
         # unstack the prediction and components

From bcf572631a0b6e5dccf9b79c4c4dac7bb5466360 Mon Sep 17 00:00:00 2001
From: ourownstory <ourownstory@users.noreply.github.com>
Date: Wed, 11 Sep 2024 17:25:31 -0700
Subject: [PATCH 08/30] simplify unstack components

---
 neuralprophet/utils_time_dataset.py | 3 +++
 1 file changed, 3 insertions(+)

diff --git a/neuralprophet/utils_time_dataset.py b/neuralprophet/utils_time_dataset.py
index de09b4d9b..f5e6b37e8 100644
--- a/neuralprophet/utils_time_dataset.py
+++ b/neuralprophet/utils_time_dataset.py
@@ -30,6 +30,9 @@ def __init__(
         self.feature_indices = feature_indices
         self.config_seasonality = config_seasonality
         self.lagged_regressor_config = lagged_regressor_config
+        self.unstack_component_func = {
+            "targets": self.unstack_targets,
+        }
 
     def unstack_component(self, component_name, batch_tensor):
         """

From 264d5539537922401305d223f78541a6513d4633 Mon Sep 17 00:00:00 2001
From: ourownstory <ourownstory@users.noreply.github.com>
Date: Wed, 11 Sep 2024 17:27:25 -0700
Subject: [PATCH 09/30] use dict to index unstack component functions

---
 neuralprophet/utils_time_dataset.py | 30 ++++++++++-------------------
 1 file changed, 10 insertions(+), 20 deletions(-)

diff --git a/neuralprophet/utils_time_dataset.py b/neuralprophet/utils_time_dataset.py
index f5e6b37e8..c37cf0284 100644
--- a/neuralprophet/utils_time_dataset.py
+++ b/neuralprophet/utils_time_dataset.py
@@ -32,6 +32,14 @@ def __init__(
         self.lagged_regressor_config = lagged_regressor_config
         self.unstack_component_func = {
             "targets": self.unstack_targets,
+            "time": self.unstack_time,
+            "seasonalities": self.unstack_seasonalities,
+            "lagged_regressors": self.unstack_lagged_regressors,
+            "lags": self.unstack_lags,
+            "additive_events": self.unstack_additive_events,
+            "multiplicative_events": self.unstack_multiplicative_events,
+            "additive_regressors": self.unstack_additive_regressors,
+            "multiplicative_regressors": self.unstack_multiplicative_regressors,
         }
 
     def unstack_component(self, component_name, batch_tensor):
@@ -44,26 +52,8 @@ def unstack_component(self, component_name, batch_tensor):
         Returns:
             Various: The output of the specific unstackion function.
         """
-        if component_name == "targets":
-            return self.unstack_targets(batch_tensor)
-        elif component_name == "time":
-            return self.unstack_time(batch_tensor)
-        elif component_name == "seasonalities":
-            return self.unstack_seasonalities(batch_tensor)
-        elif component_name == "lagged_regressors":
-            return self.unstack_lagged_regressors(batch_tensor)
-        elif component_name == "lags":
-            return self.unstack_lags(batch_tensor)
-        elif component_name == "additive_events":
-            return self.unstack_additive_events(batch_tensor)
-        elif component_name == "multiplicative_events":
-            return self.unstack_multiplicative_events(batch_tensor)
-        elif component_name == "additive_regressors":
-            return self.unstack_additive_regressors(batch_tensor)
-        elif component_name == "multiplicative_regressors":
-            return self.unstack_multiplicative_regressors(batch_tensor)
-        else:
-            raise ValueError(f"Unknown component name: {component_name}")
+        assert component_name in self.unstack_component_func, f"Unknown component name: {component_name}"
+        return self.unstack_component_func[component_name](batch_tensor)
 
     def unstack_targets(self, batch_tensor):
         targets_start_idx, targets_end_idx = self.feature_indices["targets"]

From 1404fab9cf37d534c1b2d2e4ffe77ae769ecd470 Mon Sep 17 00:00:00 2001
From: ourownstory <ourownstory@users.noreply.github.com>
Date: Wed, 11 Sep 2024 17:29:46 -0700
Subject: [PATCH 10/30] remove unused import

---
 neuralprophet/forecaster.py | 1 -
 1 file changed, 1 deletion(-)

diff --git a/neuralprophet/forecaster.py b/neuralprophet/forecaster.py
index 3e18b8bfd..c5f7efa84 100644
--- a/neuralprophet/forecaster.py
+++ b/neuralprophet/forecaster.py
@@ -44,7 +44,6 @@
 from neuralprophet.plot_model_parameters_plotly import plot_parameters as plot_parameters_plotly
 from neuralprophet.plot_utils import get_valid_configuration, log_warning_deprecation_plotly, select_plotting_backend
 from neuralprophet.uncertainty import Conformal
-from neuralprophet.utils_time_dataset import ComponentStacker
 
 log = logging.getLogger("NP.forecaster")
 

From c83fead0ce88ffb89b1712dca24c5dc9d409ace8 Mon Sep 17 00:00:00 2001
From: ourownstory <ourownstory@users.noreply.github.com>
Date: Wed, 11 Sep 2024 17:32:35 -0700
Subject: [PATCH 11/30] fix component_stacker

---
 neuralprophet/forecaster.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/neuralprophet/forecaster.py b/neuralprophet/forecaster.py
index c5f7efa84..9695b6776 100644
--- a/neuralprophet/forecaster.py
+++ b/neuralprophet/forecaster.py
@@ -2137,7 +2137,7 @@ def predict_seasonal_components(self, df: pd.DataFrame, quantile: float = 0.5):
         df = _check_dataframe(self, df, check_y=False, exogenous=False)
         df = _normalize(df=df, config_normalization=self.config_normalization)
         for df_name, df_i in df.groupby("ID"):
-            feature_unstackor = ComponentStacker(
+            feature_unstackor = utils_time_dataset.ComponentStacker(
                 n_lags=0,
                 max_lags=0,
                 n_forecasts=1,

From f106f824b665dc7a915eafe031335f7279e775ce Mon Sep 17 00:00:00 2001
From: ourownstory <ourownstory@users.noreply.github.com>
Date: Wed, 11 Sep 2024 17:41:09 -0700
Subject: [PATCH 12/30] convert to dataclass

---
 neuralprophet/utils_time_dataset.py | 49 +++++++++++++++--------------
 1 file changed, 25 insertions(+), 24 deletions(-)

diff --git a/neuralprophet/utils_time_dataset.py b/neuralprophet/utils_time_dataset.py
index c37cf0284..7351a6e9c 100644
--- a/neuralprophet/utils_time_dataset.py
+++ b/neuralprophet/utils_time_dataset.py
@@ -1,35 +1,36 @@
 from collections import OrderedDict
+from dataclasses import dataclass, default_factory, field
+from typing import Dict, List, Optional, Union
 
 import torch
 
+from neuralprophet.configure_components import LaggedRegressors, Seasonality
 
+
+@dataclass
 class ComponentStacker:
-    def __init__(
-        self,
-        n_lags,
-        n_forecasts,
-        max_lags,
-        feature_indices={},
-        config_seasonality=None,
-        lagged_regressor_config=None,
-    ):
-        """
-        Initializes the ComponentStacker with the necessary parameters.
+    """
+    ComponentStacker is a utility class that helps in stacking and unstacking the different components of the time series data.
+    Args:
+        n_lags (int): Number of lags used in the model.
+        n_forecasts (int): Number of forecasts to be made.
+        max_lags (int): Maximum number of lags used in the model.
+        feature_indices (dict): A dictionary containing the start and end indices of different features in the tensor.
+        config_seasonality (object, optional): Configuration object that defines the seasonality periods.
+        lagged_regressor_config (dict, optional): Configuration dictionary that defines the lagged regressors and their properties.
+    """
 
-        Args:
-            n_lags (int): Number of lags used in the model.
-            n_forecasts (int): Number of forecasts to be made.
-            max_lags (int): Maximum number of lags used in the model.
-            feature_indices (dict): A dictionary containing the start and end indices of different features in the tensor.
-            config_seasonality (object, optional): Configuration object that defines the seasonality periods.
-            lagged_regressor_config (dict, optional): Configuration dictionary that defines the lagged regressors and their properties.
+    n_lags: int
+    n_forecasts: int
+    max_lags: int
+    feature_indices: dict = field(default_factory=dict)
+    config_seasonality: Optional[Seasonality] = None
+    lagged_regressor_config: Optional[LaggedRegressors] = None
+
+    def __post_init__(self):
+        """
+        Initializes mappings to comonent stacking and unstacking functions.
         """
-        self.n_lags = n_lags
-        self.n_forecasts = n_forecasts
-        self.max_lags = max_lags
-        self.feature_indices = feature_indices
-        self.config_seasonality = config_seasonality
-        self.lagged_regressor_config = lagged_regressor_config
         self.unstack_component_func = {
             "targets": self.unstack_targets,
             "time": self.unstack_time,

From 63a64dd8628d8063f7bcdbf104b15b1acb27ad47 Mon Sep 17 00:00:00 2001
From: ourownstory <ourownstory@users.noreply.github.com>
Date: Wed, 11 Sep 2024 17:47:57 -0700
Subject: [PATCH 13/30] simply stack function names

---
 neuralprophet/utils_time_dataset.py | 43 ++++++++++++++++++-----------
 1 file changed, 27 insertions(+), 16 deletions(-)

diff --git a/neuralprophet/utils_time_dataset.py b/neuralprophet/utils_time_dataset.py
index 7351a6e9c..bdb8db23b 100644
--- a/neuralprophet/utils_time_dataset.py
+++ b/neuralprophet/utils_time_dataset.py
@@ -1,10 +1,10 @@
 from collections import OrderedDict
-from dataclasses import dataclass, default_factory, field
-from typing import Dict, List, Optional, Union
+from dataclasses import dataclass, field
+from typing import Optional
 
 import torch
 
-from neuralprophet.configure_components import LaggedRegressors, Seasonality
+from neuralprophet.configure_components import LaggedRegressors, Seasonalities
 
 
 @dataclass
@@ -24,14 +24,27 @@ class ComponentStacker:
     n_forecasts: int
     max_lags: int
     feature_indices: dict = field(default_factory=dict)
-    config_seasonality: Optional[Seasonality] = None
+    config_seasonality: Optional[Seasonalities] = None
     lagged_regressor_config: Optional[LaggedRegressors] = None
+    stack_func: dict = field(init=False)
+    unstack_func: dict = field(init=False)
 
     def __post_init__(self):
         """
         Initializes mappings to comonent stacking and unstacking functions.
         """
-        self.unstack_component_func = {
+        self.stack_func = {
+            "targets": self.stack_targets,
+            "time": self.stack_time,
+            "seasonalities": self.stack_seasonalities,
+            "lagged_regressors": self.stack_lagged_regressors,
+            "lags": self.stack_lags,
+            "additive_events": self.stack_additive_events,
+            "multiplicative_events": self.stack_multiplicative_events,
+            "additive_regressors": self.stack_additive_regressors,
+            "multiplicative_regressors": self.stack_multiplicative_regressors,
+        }
+        self.unstack_func = {
             "targets": self.unstack_targets,
             "time": self.unstack_time,
             "seasonalities": self.unstack_seasonalities,
@@ -53,8 +66,8 @@ def unstack_component(self, component_name, batch_tensor):
         Returns:
             Various: The output of the specific unstackion function.
         """
-        assert component_name in self.unstack_component_func, f"Unknown component name: {component_name}"
-        return self.unstack_component_func[component_name](batch_tensor)
+        assert component_name in self.unstack_func, f"Unknown component name: {component_name}"
+        return self.unstack_func[component_name](batch_tensor)
 
     def unstack_targets(self, batch_tensor):
         targets_start_idx, targets_end_idx = self.feature_indices["targets"]
@@ -180,7 +193,7 @@ def stack_lags_component(self, df_tensors, feature_list, current_idx, n_lags):
             return current_idx + 1
         return current_idx
 
-    def stack_targets_component(self, df_tensors, feature_list, current_idx):
+    def stack_targets(self, df_tensors, feature_list, current_idx):
         """
         Stack the targets feature.
         """
@@ -191,7 +204,7 @@ def stack_targets_component(self, df_tensors, feature_list, current_idx):
             return current_idx + 1
         return current_idx
 
-    def stack_lagged_regerssors_component(self, df_tensors, feature_list, current_idx, config_lagged_regressors):
+    def stack_lagged_regerssors(self, df_tensors, feature_list, current_idx, config_lagged_regressors):
         """
         Stack the lagged regressor features.
         """
@@ -210,7 +223,7 @@ def stack_lagged_regerssors_component(self, df_tensors, feature_list, current_id
             return current_idx + num_features
         return current_idx
 
-    def stack_additive_events_component(
+    def stack_additive_events(
         self,
         df_tensors,
         feature_list,
@@ -233,7 +246,7 @@ def stack_additive_events_component(
             return current_idx + additive_events_tensor.size(1)
         return current_idx
 
-    def stack_multiplicative_events_component(
+    def stack_multiplicative_events(
         self, df_tensors, feature_list, current_idx, multiplicative_event_and_holiday_names
     ):
         """
@@ -251,7 +264,7 @@ def stack_multiplicative_events_component(
             return current_idx + multiplicative_events_tensor.size(1)
         return current_idx
 
-    def stack_additive_regressors_component(self, df_tensors, feature_list, current_idx, additive_regressors_names):
+    def stack_additive_regressors(self, df_tensors, feature_list, current_idx, additive_regressors_names):
         """
         Stack the additive regressor features.
         """
@@ -267,9 +280,7 @@ def stack_additive_regressors_component(self, df_tensors, feature_list, current_
             return current_idx + additive_regressors_tensor.size(1)
         return current_idx
 
-    def stack_multiplicative_regressors_component(
-        self, df_tensors, feature_list, current_idx, multiplicative_regressors_names
-    ):
+    def stack_multiplicative_regressors(self, df_tensors, feature_list, current_idx, multiplicative_regressors_names):
         """
         Stack the multiplicative regressor features.
         """
@@ -285,7 +296,7 @@ def stack_multiplicative_regressors_component(
             return current_idx + len(multiplicative_regressors_names)
         return current_idx
 
-    def stack_seasonalities_component(self, feature_list, current_idx, config_seasonality, seasonalities):
+    def stack_seasonalities(self, feature_list, current_idx, config_seasonality, seasonalities):
         """
         Stack the seasonality features.
         """

From 86e37a1f5a668ab901c3e7926e969832348ab0e2 Mon Sep 17 00:00:00 2001
From: ourownstory <ourownstory@users.noreply.github.com>
Date: Wed, 11 Sep 2024 17:49:43 -0700
Subject: [PATCH 14/30] fix references

---
 neuralprophet/time_dataset.py | 18 +++++++++---------
 1 file changed, 9 insertions(+), 9 deletions(-)

diff --git a/neuralprophet/time_dataset.py b/neuralprophet/time_dataset.py
index 84be0dc05..eb8b86dc0 100644
--- a/neuralprophet/time_dataset.py
+++ b/neuralprophet/time_dataset.py
@@ -123,30 +123,30 @@ def stack_all_features(self):
         current_idx = 0
 
         # Call individual stacking functions
-        current_idx = self.components_stacker.stack_trend_component(self.df_tensors, feature_list, current_idx)
-        current_idx = self.components_stacker.stack_targets_component(self.df_tensors, feature_list, current_idx)
+        current_idx = self.components_stacker.stack_trend(self.df_tensors, feature_list, current_idx)
+        current_idx = self.components_stacker.stack_targets(self.df_tensors, feature_list, current_idx)
 
-        current_idx = self.components_stacker.stack_lags_component(
+        current_idx = self.components_stacker.stack_lags(
             self.df_tensors, feature_list, current_idx, self.config_ar.n_lags
         )
-        current_idx = self.components_stacker.stack_lagged_regerssors_component(
+        current_idx = self.components_stacker.stack_lagged_regerssors(
             self.df_tensors, feature_list, current_idx, self.config_lagged_regressors
         )
-        current_idx = self.components_stacker.stack_additive_events_component(
+        current_idx = self.components_stacker.stack_additive_events(
             self.df_tensors, feature_list, current_idx, self.additive_event_and_holiday_names
         )
-        current_idx = self.components_stacker.stack_multiplicative_events_component(
+        current_idx = self.components_stacker.stack_multiplicative_events(
             self.df_tensors, feature_list, current_idx, self.multiplicative_event_and_holiday_names
         )
-        current_idx = self.components_stacker.stack_additive_regressors_component(
+        current_idx = self.components_stacker.stack_additive_regressors(
             self.df_tensors, feature_list, current_idx, self.additive_regressors_names
         )
-        current_idx = self.components_stacker.stack_multiplicative_regressors_component(
+        current_idx = self.components_stacker.stack_multiplicative_regressors(
             self.df_tensors, feature_list, current_idx, self.multiplicative_regressors_names
         )
 
         if self.config_seasonality is not None and hasattr(self.config_seasonality, "periods"):
-            current_idx = self.components_stacker.stack_seasonalities_component(
+            current_idx = self.components_stacker.stack_seasonalities(
                 feature_list, current_idx, self.config_seasonality, self.seasonalities
             )
 

From f59c6b9f0ef72f55e808080e0dd630b15767026f Mon Sep 17 00:00:00 2001
From: ourownstory <ourownstory@users.noreply.github.com>
Date: Wed, 11 Sep 2024 17:52:49 -0700
Subject: [PATCH 15/30] fix trend/time

---
 neuralprophet/time_dataset.py       | 2 +-
 neuralprophet/utils_time_dataset.py | 2 +-
 2 files changed, 2 insertions(+), 2 deletions(-)

diff --git a/neuralprophet/time_dataset.py b/neuralprophet/time_dataset.py
index eb8b86dc0..49ed8d9a2 100644
--- a/neuralprophet/time_dataset.py
+++ b/neuralprophet/time_dataset.py
@@ -123,7 +123,7 @@ def stack_all_features(self):
         current_idx = 0
 
         # Call individual stacking functions
-        current_idx = self.components_stacker.stack_trend(self.df_tensors, feature_list, current_idx)
+        current_idx = self.components_stacker.stack_time(self.df_tensors, feature_list, current_idx)
         current_idx = self.components_stacker.stack_targets(self.df_tensors, feature_list, current_idx)
 
         current_idx = self.components_stacker.stack_lags(
diff --git a/neuralprophet/utils_time_dataset.py b/neuralprophet/utils_time_dataset.py
index bdb8db23b..1dcf02fef 100644
--- a/neuralprophet/utils_time_dataset.py
+++ b/neuralprophet/utils_time_dataset.py
@@ -173,7 +173,7 @@ def unstack_multiplicative_regressors(self, batch_tensor):
             regressors_start_idx, regressors_end_idx = self.feature_indices["multiplicative_regressors"]
             return batch_tensor[:, regressors_start_idx : regressors_end_idx + 1].unsqueeze(1)
 
-    def stack_trend_component(self, df_tensors, feature_list, current_idx):
+    def stack_time(self, df_tensors, feature_list, current_idx):
         """
         Stack the trend (time) feature.
         """

From 7ce4fcbd8200e3c6f38189c74506c595e7bbef48 Mon Sep 17 00:00:00 2001
From: ourownstory <ourownstory@users.noreply.github.com>
Date: Wed, 11 Sep 2024 17:54:09 -0700
Subject: [PATCH 16/30] fix pre-existing typo

---
 neuralprophet/time_dataset.py       | 2 +-
 neuralprophet/utils_time_dataset.py | 2 +-
 2 files changed, 2 insertions(+), 2 deletions(-)

diff --git a/neuralprophet/time_dataset.py b/neuralprophet/time_dataset.py
index 49ed8d9a2..e89143d62 100644
--- a/neuralprophet/time_dataset.py
+++ b/neuralprophet/time_dataset.py
@@ -129,7 +129,7 @@ def stack_all_features(self):
         current_idx = self.components_stacker.stack_lags(
             self.df_tensors, feature_list, current_idx, self.config_ar.n_lags
         )
-        current_idx = self.components_stacker.stack_lagged_regerssors(
+        current_idx = self.components_stacker.stack_lagged_regressors(
             self.df_tensors, feature_list, current_idx, self.config_lagged_regressors
         )
         current_idx = self.components_stacker.stack_additive_events(
diff --git a/neuralprophet/utils_time_dataset.py b/neuralprophet/utils_time_dataset.py
index 1dcf02fef..4f4fc4f59 100644
--- a/neuralprophet/utils_time_dataset.py
+++ b/neuralprophet/utils_time_dataset.py
@@ -204,7 +204,7 @@ def stack_targets(self, df_tensors, feature_list, current_idx):
             return current_idx + 1
         return current_idx
 
-    def stack_lagged_regerssors(self, df_tensors, feature_list, current_idx, config_lagged_regressors):
+    def stack_lagged_regressors(self, df_tensors, feature_list, current_idx, config_lagged_regressors):
         """
         Stack the lagged regressor features.
         """

From eb3204779dac991b998a1b1f91b804ef4b3a6440 Mon Sep 17 00:00:00 2001
From: ourownstory <ourownstory@users.noreply.github.com>
Date: Wed, 11 Sep 2024 18:06:00 -0700
Subject: [PATCH 17/30] improve seasonality stacker

---
 neuralprophet/time_dataset.py       | 10 +++++++---
 neuralprophet/utils_time_dataset.py | 26 +++++++++++++++++++++++---
 2 files changed, 30 insertions(+), 6 deletions(-)

diff --git a/neuralprophet/time_dataset.py b/neuralprophet/time_dataset.py
index e89143d62..8128b606c 100644
--- a/neuralprophet/time_dataset.py
+++ b/neuralprophet/time_dataset.py
@@ -147,7 +147,10 @@ def stack_all_features(self):
 
         if self.config_seasonality is not None and hasattr(self.config_seasonality, "periods"):
             current_idx = self.components_stacker.stack_seasonalities(
-                feature_list, current_idx, self.config_seasonality, self.seasonalities
+                self.df_tensors,
+                feature_list,
+                current_idx,
+                self.config_seasonality,
             )
 
         # Concatenate all features into one big tensor
@@ -155,7 +158,7 @@ def stack_all_features(self):
 
     def calculate_seasonalities(self):
         """Computes Fourier series components with the specified frequency and order."""
-        self.seasonalities = OrderedDict({})
+        seasonalities = OrderedDict({})
         dates = self.df_tensors["ds"]
         t = (dates - torch.tensor(datetime(1900, 1, 1).timestamp())).float() / (3600 * 24.0)
 
@@ -191,7 +194,8 @@ def compute_fourier_features(t, period, resolution):
                 if period.condition_name is not None:
                     condition_values = self.df_tensors[period.condition_name].unsqueeze(1)
                     features *= condition_values
-                self.seasonalities[name] = features
+                seasonalities[name] = features
+        self.df_tensors["seasonalities"] = seasonalities
 
     def __getitem__(self, index):
         """Overrides parent class method to get an item at index.
diff --git a/neuralprophet/utils_time_dataset.py b/neuralprophet/utils_time_dataset.py
index 4f4fc4f59..2065a2b93 100644
--- a/neuralprophet/utils_time_dataset.py
+++ b/neuralprophet/utils_time_dataset.py
@@ -64,11 +64,30 @@ def unstack_component(self, component_name, batch_tensor):
             component_name (str): The name of the component to unstack.
 
         Returns:
-            Various: The output of the specific unstackion function.
+            Various: The output of the specific unstacking function.
         """
         assert component_name in self.unstack_func, f"Unknown component name: {component_name}"
         return self.unstack_func[component_name](batch_tensor)
 
+    def stack_component(self, component_name, df_tensors, feature_list, current_idx, kwargs):
+        """
+        Routes the unstackion process to the appropriate function based on the component name.
+
+        Args:
+            component_name (str): The name of the component to stack.
+            df_tensors
+            feature_list
+            current_idx
+            kwargs for specific component, mostly component configuration
+
+        Returns:
+            current_idx: the current index in the stack of features.
+        """
+        assert component_name in self.stack_func, f"Unknown component name: {component_name}"
+        return self.stack_func[component_name](
+            df_tensors=df_tensors, feature_list=feature_list, current_idx=current_idx, **kwargs
+        )
+
     def unstack_targets(self, batch_tensor):
         targets_start_idx, targets_end_idx = self.feature_indices["targets"]
         if self.max_lags > 0:
@@ -182,7 +201,7 @@ def stack_time(self, df_tensors, feature_list, current_idx):
         self.feature_indices["time"] = (current_idx, current_idx)
         return current_idx + 1
 
-    def stack_lags_component(self, df_tensors, feature_list, current_idx, n_lags):
+    def stack_lags(self, df_tensors, feature_list, current_idx, n_lags):
         """
         Stack the lags feature.
         """
@@ -296,10 +315,11 @@ def stack_multiplicative_regressors(self, df_tensors, feature_list, current_idx,
             return current_idx + len(multiplicative_regressors_names)
         return current_idx
 
-    def stack_seasonalities(self, feature_list, current_idx, config_seasonality, seasonalities):
+    def stack_seasonalities(self, df_tensors, feature_list, current_idx, config_seasonality):
         """
         Stack the seasonality features.
         """
+        seasonalities = df_tensors["seasonalities"]
         if config_seasonality and config_seasonality.periods:
             for seasonality_name, features in seasonalities.items():
                 seasonal_tensor = features

From 5f71d436ee6bc472ed8bbf9a6cf5b0e69a124de1 Mon Sep 17 00:00:00 2001
From: ourownstory <ourownstory@users.noreply.github.com>
Date: Wed, 11 Sep 2024 18:09:10 -0700
Subject: [PATCH 18/30] revert seasonalities

---
 neuralprophet/time_dataset.py       | 10 +++-------
 neuralprophet/utils_time_dataset.py |  4 ++--
 2 files changed, 5 insertions(+), 9 deletions(-)

diff --git a/neuralprophet/time_dataset.py b/neuralprophet/time_dataset.py
index 8128b606c..e89143d62 100644
--- a/neuralprophet/time_dataset.py
+++ b/neuralprophet/time_dataset.py
@@ -147,10 +147,7 @@ def stack_all_features(self):
 
         if self.config_seasonality is not None and hasattr(self.config_seasonality, "periods"):
             current_idx = self.components_stacker.stack_seasonalities(
-                self.df_tensors,
-                feature_list,
-                current_idx,
-                self.config_seasonality,
+                feature_list, current_idx, self.config_seasonality, self.seasonalities
             )
 
         # Concatenate all features into one big tensor
@@ -158,7 +155,7 @@ def stack_all_features(self):
 
     def calculate_seasonalities(self):
         """Computes Fourier series components with the specified frequency and order."""
-        seasonalities = OrderedDict({})
+        self.seasonalities = OrderedDict({})
         dates = self.df_tensors["ds"]
         t = (dates - torch.tensor(datetime(1900, 1, 1).timestamp())).float() / (3600 * 24.0)
 
@@ -194,8 +191,7 @@ def compute_fourier_features(t, period, resolution):
                 if period.condition_name is not None:
                     condition_values = self.df_tensors[period.condition_name].unsqueeze(1)
                     features *= condition_values
-                seasonalities[name] = features
-        self.df_tensors["seasonalities"] = seasonalities
+                self.seasonalities[name] = features
 
     def __getitem__(self, index):
         """Overrides parent class method to get an item at index.
diff --git a/neuralprophet/utils_time_dataset.py b/neuralprophet/utils_time_dataset.py
index 2065a2b93..88b6fc39d 100644
--- a/neuralprophet/utils_time_dataset.py
+++ b/neuralprophet/utils_time_dataset.py
@@ -84,6 +84,7 @@ def stack_component(self, component_name, df_tensors, feature_list, current_idx,
             current_idx: the current index in the stack of features.
         """
         assert component_name in self.stack_func, f"Unknown component name: {component_name}"
+        # this is currently not working for seasonalities
         return self.stack_func[component_name](
             df_tensors=df_tensors, feature_list=feature_list, current_idx=current_idx, **kwargs
         )
@@ -315,11 +316,10 @@ def stack_multiplicative_regressors(self, df_tensors, feature_list, current_idx,
             return current_idx + len(multiplicative_regressors_names)
         return current_idx
 
-    def stack_seasonalities(self, df_tensors, feature_list, current_idx, config_seasonality):
+    def stack_seasonalities(self, feature_list, current_idx, config_seasonality, seasonalities):
         """
         Stack the seasonality features.
         """
-        seasonalities = df_tensors["seasonalities"]
         if config_seasonality and config_seasonality.periods:
             for seasonality_name, features in seasonalities.items():
                 seasonal_tensor = features

From d79008766c7a9d60a13417d2e84e8b953a5f9335 Mon Sep 17 00:00:00 2001
From: ourownstory <ourownstory@users.noreply.github.com>
Date: Wed, 11 Sep 2024 18:12:01 -0700
Subject: [PATCH 19/30] rename stack/unstack

---
 neuralprophet/forecaster.py         | 4 ++--
 neuralprophet/utils_time_dataset.py | 7 ++++---
 2 files changed, 6 insertions(+), 5 deletions(-)

diff --git a/neuralprophet/forecaster.py b/neuralprophet/forecaster.py
index 9695b6776..2681f9e1f 100644
--- a/neuralprophet/forecaster.py
+++ b/neuralprophet/forecaster.py
@@ -2168,12 +2168,12 @@ def predict_seasonal_components(self, df: pd.DataFrame, quantile: float = 0.5):
                     meta_name_tensor = None
                 elif self.model.config_seasonality.global_local in ["local", "glocal"]:
                     meta = OrderedDict()
-                    time_input = feature_unstackor.unstack_component("time", inputs_tensor)
+                    time_input = feature_unstackor.unstack("time", inputs_tensor)
                     meta["df_name"] = [df_name for _ in range(time_input.shape[0])]
                     meta_name_tensor = torch.tensor([self.model.id_dict[i] for i in meta["df_name"]])  # type: ignore
                 else:
                     meta_name_tensor = None
-                seasonalities_input = feature_unstackor.unstack_component("seasonalities", inputs_tensor)
+                seasonalities_input = feature_unstackor.unstack("seasonalities", inputs_tensor)
                 for name in self.config_seasonality.periods:
                     features = seasonalities_input[name]
                     quantile_index = self.config_model.quantiles.index(quantile)
diff --git a/neuralprophet/utils_time_dataset.py b/neuralprophet/utils_time_dataset.py
index 88b6fc39d..2fbc23604 100644
--- a/neuralprophet/utils_time_dataset.py
+++ b/neuralprophet/utils_time_dataset.py
@@ -56,7 +56,7 @@ def __post_init__(self):
             "multiplicative_regressors": self.unstack_multiplicative_regressors,
         }
 
-    def unstack_component(self, component_name, batch_tensor):
+    def unstack(self, component_name, batch_tensor):
         """
         Routes the unstackion process to the appropriate function based on the component name.
 
@@ -69,7 +69,7 @@ def unstack_component(self, component_name, batch_tensor):
         assert component_name in self.unstack_func, f"Unknown component name: {component_name}"
         return self.unstack_func[component_name](batch_tensor)
 
-    def stack_component(self, component_name, df_tensors, feature_list, current_idx, kwargs):
+    def stack(self, component_name, df_tensors, feature_list, current_idx, kwargs):
         """
         Routes the unstackion process to the appropriate function based on the component name.
 
@@ -316,10 +316,11 @@ def stack_multiplicative_regressors(self, df_tensors, feature_list, current_idx,
             return current_idx + len(multiplicative_regressors_names)
         return current_idx
 
-    def stack_seasonalities(self, feature_list, current_idx, config_seasonality, seasonalities):
+    def stack_seasonalities(self, feature_list, current_idx, config_seasonality, seasonalities, df_tensors=None):
         """
         Stack the seasonality features.
         """
+        # TODO conform to other stack functions, using df_tensors
         if config_seasonality and config_seasonality.periods:
             for seasonality_name, features in seasonalities.items():
                 seasonal_tensor = features

From 348b3d73fcb0d534c72720495cfc3109e4f5726c Mon Sep 17 00:00:00 2001
From: ourownstory <ourownstory@users.noreply.github.com>
Date: Wed, 11 Sep 2024 18:13:13 -0700
Subject: [PATCH 20/30] update timenet

---
 neuralprophet/time_net.py | 30 ++++++++++++++----------------
 1 file changed, 14 insertions(+), 16 deletions(-)

diff --git a/neuralprophet/time_net.py b/neuralprophet/time_net.py
index 46fe8c3a7..2d8a2d474 100644
--- a/neuralprophet/time_net.py
+++ b/neuralprophet/time_net.py
@@ -538,7 +538,7 @@ def forward(
 
         """
 
-        time_input = self.components_stacker[mode].unstack_component(component_name="time", batch_tensor=input_tensor)
+        time_input = self.components_stacker[mode].unstack(component_name="time", batch_tensor=input_tensor)
         # Handle meta argument
         if meta is None and self.meta_used_in_model:
             name_id_dummy = self.id_list[0]
@@ -568,7 +568,7 @@ def forward(
         # Unpack and process seasonalities
         seasonalities_input = None
         if self.config_seasonality and self.config_seasonality.periods:
-            seasonalities_input = self.components_stacker[mode].unstack_component(
+            seasonalities_input = self.components_stacker[mode].unstack(
                 component_name="seasonalities", batch_tensor=input_tensor
             )
             s = self.seasonality(s=seasonalities_input, meta=meta)
@@ -583,14 +583,14 @@ def forward(
         multiplicative_events_input = None
         if self.events_dims is not None:
             if "additive_events" in self.components_stacker[mode].feature_indices:
-                additive_events_input = self.components_stacker[mode].unstack_component(
+                additive_events_input = self.components_stacker[mode].unstack(
                     component_name="additive_events", batch_tensor=input_tensor
                 )
                 additive_events = self.scalar_features_effects(additive_events_input, self.event_params["additive"])
                 additive_components_nonstationary += additive_events
                 components["additive_events"] = additive_events
             if "multiplicative_events" in self.components_stacker[mode].feature_indices:
-                multiplicative_events_input = self.components_stacker[mode].unstack_component(
+                multiplicative_events_input = self.components_stacker[mode].unstack(
                     component_name="multiplicative_events", batch_tensor=input_tensor
                 )
                 multiplicative_events = self.scalar_features_effects(
@@ -603,14 +603,14 @@ def forward(
         additive_regressors_input = None
         multiplicative_regressors_input = None
         if "additive_regressors" in self.components_stacker[mode].feature_indices:
-            additive_regressors_input = self.components_stacker[mode].unstack_component(
+            additive_regressors_input = self.components_stacker[mode].unstack(
                 component_name="additive_regressors", batch_tensor=input_tensor
             )
             additive_regressors = self.future_regressors(additive_regressors_input, "additive")
             additive_components_nonstationary += additive_regressors
             components["additive_regressors"] = additive_regressors
         if "multiplicative_regressors" in self.components_stacker[mode].feature_indices:
-            multiplicative_regressors_input = self.components_stacker[mode].unstack_component(
+            multiplicative_regressors_input = self.components_stacker[mode].unstack(
                 component_name="multiplicative_regressors", batch_tensor=input_tensor
             )
             multiplicative_regressors = self.future_regressors(multiplicative_regressors_input, "multiplicative")
@@ -620,9 +620,7 @@ def forward(
         # Unpack and process lags
         lags_input = None
         if "lags" in self.components_stacker[mode].feature_indices:
-            lags_input = self.components_stacker[mode].unstack_component(
-                component_name="lags", batch_tensor=input_tensor
-            )
+            lags_input = self.components_stacker[mode].unstack(component_name="lags", batch_tensor=input_tensor)
             nonstationary_components = (
                 trend[:, : self.n_lags, 0]
                 + additive_components_nonstationary[:, : self.n_lags, 0]
@@ -636,7 +634,7 @@ def forward(
         # Unpack and process covariates
         covariates_input = None
         if self.config_lagged_regressors and self.config_lagged_regressors.regressors is not None:
-            covariates_input = self.components_stacker[mode].unstack_component(
+            covariates_input = self.components_stacker[mode].unstack(
                 component_name="lagged_regressors", batch_tensor=input_tensor
             )
             covariates = self.forward_covar_net(covariates=covariates_input)
@@ -792,8 +790,8 @@ def training_step(self, batch, batch_idx):
 
         epoch_float = self.trainer.current_epoch + batch_idx / float(self.train_steps_per_epoch)
         self.train_progress = epoch_float / float(self.config_train.epochs)
-        targets = self.components_stacker["train"].unstack_component("targets", batch_tensor=inputs_tensor)
-        time = self.components_stacker["train"].unstack_component("time", batch_tensor=inputs_tensor)
+        targets = self.components_stacker["train"].unstack("targets", batch_tensor=inputs_tensor)
+        time = self.components_stacker["train"].unstack("time", batch_tensor=inputs_tensor)
         # Global-local
         if self.meta_used_in_model:
             meta_name_tensor = torch.tensor([self.id_dict[i] for i in meta["df_name"]], device=self.device)
@@ -837,8 +835,8 @@ def training_step(self, batch, batch_idx):
 
     def validation_step(self, batch, batch_idx):
         inputs_tensor, meta = batch
-        targets = self.components_stacker["val"].unstack_component("targets", batch_tensor=inputs_tensor)
-        time = self.components_stacker["val"].unstack_component("time", batch_tensor=inputs_tensor)
+        targets = self.components_stacker["val"].unstack("targets", batch_tensor=inputs_tensor)
+        time = self.components_stacker["val"].unstack("time", batch_tensor=inputs_tensor)
         # Global-local
         if self.meta_used_in_model:
             meta_name_tensor = torch.tensor([self.id_dict[i] for i in meta["df_name"]], device=self.device)
@@ -859,8 +857,8 @@ def validation_step(self, batch, batch_idx):
 
     def test_step(self, batch, batch_idx):
         inputs_tensor, meta = batch
-        targets = self.components_stacker["test"].unstack_component("targets", batch_tensor=inputs_tensor)
-        time = self.components_stacker["test"].unstack_component("time", batch_tensor=inputs_tensor)
+        targets = self.components_stacker["test"].unstack("targets", batch_tensor=inputs_tensor)
+        time = self.components_stacker["test"].unstack("time", batch_tensor=inputs_tensor)
         # Global-local
         if self.meta_used_in_model:
             meta_name_tensor = torch.tensor([self.id_dict[i] for i in meta["df_name"]], device=self.device)

From 83137f29889b1fbd3b0e7440fba581b1a087f883 Mon Sep 17 00:00:00 2001
From: ourownstory <ourownstory@users.noreply.github.com>
Date: Wed, 11 Sep 2024 18:15:02 -0700
Subject: [PATCH 21/30] use stack function

---
 neuralprophet/time_dataset.py | 8 ++++++--
 1 file changed, 6 insertions(+), 2 deletions(-)

diff --git a/neuralprophet/time_dataset.py b/neuralprophet/time_dataset.py
index e89143d62..acc1376e2 100644
--- a/neuralprophet/time_dataset.py
+++ b/neuralprophet/time_dataset.py
@@ -126,8 +126,12 @@ def stack_all_features(self):
         current_idx = self.components_stacker.stack_time(self.df_tensors, feature_list, current_idx)
         current_idx = self.components_stacker.stack_targets(self.df_tensors, feature_list, current_idx)
 
-        current_idx = self.components_stacker.stack_lags(
-            self.df_tensors, feature_list, current_idx, self.config_ar.n_lags
+        current_idx = self.components_stacker.stack(
+            component_name="lags",
+            df_tensors=self.df_tensors,
+            feature_list=feature_list,
+            current_idx=current_idx,
+            n_lags=self.config_ar.n_lags,
         )
         current_idx = self.components_stacker.stack_lagged_regressors(
             self.df_tensors, feature_list, current_idx, self.config_lagged_regressors

From 218a43dea85df01144b988402131ee3e2f0bbeaf Mon Sep 17 00:00:00 2001
From: ourownstory <ourownstory@users.noreply.github.com>
Date: Wed, 11 Sep 2024 18:15:54 -0700
Subject: [PATCH 22/30] kwargs

---
 neuralprophet/utils_time_dataset.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/neuralprophet/utils_time_dataset.py b/neuralprophet/utils_time_dataset.py
index 2fbc23604..01f748ea8 100644
--- a/neuralprophet/utils_time_dataset.py
+++ b/neuralprophet/utils_time_dataset.py
@@ -69,7 +69,7 @@ def unstack(self, component_name, batch_tensor):
         assert component_name in self.unstack_func, f"Unknown component name: {component_name}"
         return self.unstack_func[component_name](batch_tensor)
 
-    def stack(self, component_name, df_tensors, feature_list, current_idx, kwargs):
+    def stack(self, component_name, df_tensors, feature_list, current_idx, **kwargs):
         """
         Routes the unstackion process to the appropriate function based on the component name.
 

From ca7c52ef3403635d03bca89fa5f7d57f3f987220 Mon Sep 17 00:00:00 2001
From: ourownstory <ourownstory@users.noreply.github.com>
Date: Wed, 11 Sep 2024 18:24:35 -0700
Subject: [PATCH 23/30] use stacker abstraction

---
 neuralprophet/time_dataset.py       | 23 ++++++++++++++++++++++-
 neuralprophet/utils_time_dataset.py | 14 ++++++--------
 2 files changed, 28 insertions(+), 9 deletions(-)

diff --git a/neuralprophet/time_dataset.py b/neuralprophet/time_dataset.py
index acc1376e2..b72b32abc 100644
--- a/neuralprophet/time_dataset.py
+++ b/neuralprophet/time_dataset.py
@@ -119,9 +119,30 @@ def stack_all_features(self):
         Stack all features into one large tensor by calling individual stacking methods.
         """
         feature_list = []
-
         current_idx = 0
 
+        components_args: dict = {
+            "time": {},
+            "targets": {},
+            "lags": {"n_lags": self.config_ar.n_lags},
+            "lagged_regressors": {"config": self.config_lagged_regressors},
+            "additive_events": {"additive_event_and_holiday_names": self.additive_event_and_holiday_names},
+            "multiplicative_events": {
+                "multiplicative_event_and_holiday_names": self.multiplicative_event_and_holiday_names
+            },
+            "additive_regressors": {"additive_regressors_names": self.additive_regressors_names},
+            "multiplicative_regressors": {"multiplicative_regressors_names": self.multiplicative_regressors_names},
+            "seasonalities": {"config": self.config_seasonality, "seasonalities": self.seasonalities},
+        }
+        for component, args in components_args.items():
+            current_idx = self.components_stacker.stack(
+                component_name=component,
+                df_tensors=self.df_tensors,
+                feature_list=feature_list,
+                current_idx=current_idx,
+                **args,
+            )
+
         # Call individual stacking functions
         current_idx = self.components_stacker.stack_time(self.df_tensors, feature_list, current_idx)
         current_idx = self.components_stacker.stack_targets(self.df_tensors, feature_list, current_idx)
diff --git a/neuralprophet/utils_time_dataset.py b/neuralprophet/utils_time_dataset.py
index 01f748ea8..4e4ac9b3d 100644
--- a/neuralprophet/utils_time_dataset.py
+++ b/neuralprophet/utils_time_dataset.py
@@ -224,18 +224,16 @@ def stack_targets(self, df_tensors, feature_list, current_idx):
             return current_idx + 1
         return current_idx
 
-    def stack_lagged_regressors(self, df_tensors, feature_list, current_idx, config_lagged_regressors):
+    def stack_lagged_regressors(self, df_tensors, feature_list, current_idx, config):
         """
         Stack the lagged regressor features.
         """
-        if config_lagged_regressors is not None and config_lagged_regressors.regressors is not None:
-            lagged_regressor_tensors = [
-                df_tensors[name].unsqueeze(-1) for name in config_lagged_regressors.regressors.keys()
-            ]
+        if config is not None and config.regressors is not None:
+            lagged_regressor_tensors = [df_tensors[name].unsqueeze(-1) for name in config.regressors.keys()]
             stacked_lagged_regressor_tensor = torch.cat(lagged_regressor_tensors, dim=-1)
             feature_list.append(stacked_lagged_regressor_tensor)
             num_features = stacked_lagged_regressor_tensor.size(-1)
-            for i, name in enumerate(config_lagged_regressors.regressors.keys()):
+            for i, name in enumerate(config.regressors.keys()):
                 self.feature_indices[f"lagged_regressor_{name}"] = (
                     current_idx + i,
                     current_idx + i + 1,
@@ -316,12 +314,12 @@ def stack_multiplicative_regressors(self, df_tensors, feature_list, current_idx,
             return current_idx + len(multiplicative_regressors_names)
         return current_idx
 
-    def stack_seasonalities(self, feature_list, current_idx, config_seasonality, seasonalities, df_tensors=None):
+    def stack_seasonalities(self, df_tensors, feature_list, current_idx, config, seasonalities):
         """
         Stack the seasonality features.
         """
         # TODO conform to other stack functions, using df_tensors
-        if config_seasonality and config_seasonality.periods:
+        if config and config.periods:
             for seasonality_name, features in seasonalities.items():
                 seasonal_tensor = features
                 feature_list.append(seasonal_tensor)

From 5357f8eb3f808c4210c9ef108ed858855ff58d01 Mon Sep 17 00:00:00 2001
From: ourownstory <ourownstory@users.noreply.github.com>
Date: Wed, 11 Sep 2024 18:25:58 -0700
Subject: [PATCH 24/30] simplify names

---
 neuralprophet/time_dataset.py | 78 +++++++++++++++++------------------
 1 file changed, 38 insertions(+), 40 deletions(-)

diff --git a/neuralprophet/time_dataset.py b/neuralprophet/time_dataset.py
index b72b32abc..abb7b1596 100644
--- a/neuralprophet/time_dataset.py
+++ b/neuralprophet/time_dataset.py
@@ -121,59 +121,57 @@ def stack_all_features(self):
         feature_list = []
         current_idx = 0
 
-        components_args: dict = {
+        component_args: dict = {
             "time": {},
             "targets": {},
             "lags": {"n_lags": self.config_ar.n_lags},
             "lagged_regressors": {"config": self.config_lagged_regressors},
-            "additive_events": {"additive_event_and_holiday_names": self.additive_event_and_holiday_names},
-            "multiplicative_events": {
-                "multiplicative_event_and_holiday_names": self.multiplicative_event_and_holiday_names
-            },
-            "additive_regressors": {"additive_regressors_names": self.additive_regressors_names},
-            "multiplicative_regressors": {"multiplicative_regressors_names": self.multiplicative_regressors_names},
+            "additive_events": {"names": self.additive_event_and_holiday_names},
+            "multiplicative_events": {"names": self.multiplicative_event_and_holiday_names},
+            "additive_regressors": {"names": self.additive_regressors_names},
+            "multiplicative_regressors": {"names": self.multiplicative_regressors_names},
             "seasonalities": {"config": self.config_seasonality, "seasonalities": self.seasonalities},
         }
-        for component, args in components_args.items():
+        for component_name, args in component_args.items():
             current_idx = self.components_stacker.stack(
-                component_name=component,
+                component_name=component_name,
                 df_tensors=self.df_tensors,
                 feature_list=feature_list,
                 current_idx=current_idx,
                 **args,
             )
 
-        # Call individual stacking functions
-        current_idx = self.components_stacker.stack_time(self.df_tensors, feature_list, current_idx)
-        current_idx = self.components_stacker.stack_targets(self.df_tensors, feature_list, current_idx)
-
-        current_idx = self.components_stacker.stack(
-            component_name="lags",
-            df_tensors=self.df_tensors,
-            feature_list=feature_list,
-            current_idx=current_idx,
-            n_lags=self.config_ar.n_lags,
-        )
-        current_idx = self.components_stacker.stack_lagged_regressors(
-            self.df_tensors, feature_list, current_idx, self.config_lagged_regressors
-        )
-        current_idx = self.components_stacker.stack_additive_events(
-            self.df_tensors, feature_list, current_idx, self.additive_event_and_holiday_names
-        )
-        current_idx = self.components_stacker.stack_multiplicative_events(
-            self.df_tensors, feature_list, current_idx, self.multiplicative_event_and_holiday_names
-        )
-        current_idx = self.components_stacker.stack_additive_regressors(
-            self.df_tensors, feature_list, current_idx, self.additive_regressors_names
-        )
-        current_idx = self.components_stacker.stack_multiplicative_regressors(
-            self.df_tensors, feature_list, current_idx, self.multiplicative_regressors_names
-        )
-
-        if self.config_seasonality is not None and hasattr(self.config_seasonality, "periods"):
-            current_idx = self.components_stacker.stack_seasonalities(
-                feature_list, current_idx, self.config_seasonality, self.seasonalities
-            )
+        # # Call individual stacking functions
+        # current_idx = self.components_stacker.stack_time(self.df_tensors, feature_list, current_idx)
+        # current_idx = self.components_stacker.stack_targets(self.df_tensors, feature_list, current_idx)
+
+        # current_idx = self.components_stacker.stack(
+        #     component_name="lags",
+        #     df_tensors=self.df_tensors,
+        #     feature_list=feature_list,
+        #     current_idx=current_idx,
+        #     n_lags=self.config_ar.n_lags,
+        # )
+        # current_idx = self.components_stacker.stack_lagged_regressors(
+        #     self.df_tensors, feature_list, current_idx, self.config_lagged_regressors
+        # )
+        # current_idx = self.components_stacker.stack_additive_events(
+        #     self.df_tensors, feature_list, current_idx, self.additive_event_and_holiday_names
+        # )
+        # current_idx = self.components_stacker.stack_multiplicative_events(
+        #     self.df_tensors, feature_list, current_idx, self.multiplicative_event_and_holiday_names
+        # )
+        # current_idx = self.components_stacker.stack_additive_regressors(
+        #     self.df_tensors, feature_list, current_idx, self.additive_regressors_names
+        # )
+        # current_idx = self.components_stacker.stack_multiplicative_regressors(
+        #     self.df_tensors, feature_list, current_idx, self.multiplicative_regressors_names
+        # )
+
+        # if self.config_seasonality is not None and hasattr(self.config_seasonality, "periods"):
+        #     current_idx = self.components_stacker.stack_seasonalities(
+        #         feature_list, current_idx, self.config_seasonality, self.seasonalities
+        #     )
 
         # Concatenate all features into one big tensor
         self.all_features = torch.cat(feature_list, dim=1)  # Concatenating along the second dimension

From f209d97ed2439ccc521b8f320ec429c79d66eafc Mon Sep 17 00:00:00 2001
From: ourownstory <ourownstory@users.noreply.github.com>
Date: Wed, 11 Sep 2024 18:26:57 -0700
Subject: [PATCH 25/30] names

---
 neuralprophet/utils_time_dataset.py | 40 ++++++++++-------------------
 1 file changed, 14 insertions(+), 26 deletions(-)

diff --git a/neuralprophet/utils_time_dataset.py b/neuralprophet/utils_time_dataset.py
index 4e4ac9b3d..a3456baf8 100644
--- a/neuralprophet/utils_time_dataset.py
+++ b/neuralprophet/utils_time_dataset.py
@@ -241,19 +241,13 @@ def stack_lagged_regressors(self, df_tensors, feature_list, current_idx, config)
             return current_idx + num_features
         return current_idx
 
-    def stack_additive_events(
-        self,
-        df_tensors,
-        feature_list,
-        current_idx,
-        additive_event_and_holiday_names,
-    ):
+    def stack_additive_events(self, df_tensors, feature_list, current_idx, names):
         """
         Stack the additive event and holiday features.
         """
-        if additive_event_and_holiday_names:
+        if names:
             additive_events_tensor = torch.cat(
-                [df_tensors[name].unsqueeze(-1) for name in additive_event_and_holiday_names],
+                [df_tensors[name].unsqueeze(-1) for name in names],
                 dim=1,
             )
             feature_list.append(additive_events_tensor)
@@ -264,16 +258,12 @@ def stack_additive_events(
             return current_idx + additive_events_tensor.size(1)
         return current_idx
 
-    def stack_multiplicative_events(
-        self, df_tensors, feature_list, current_idx, multiplicative_event_and_holiday_names
-    ):
+    def stack_multiplicative_events(self, df_tensors, feature_list, current_idx, names):
         """
         Stack the multiplicative event and holiday features.
         """
-        if multiplicative_event_and_holiday_names:
-            multiplicative_events_tensor = torch.cat(
-                [df_tensors[name].unsqueeze(-1) for name in multiplicative_event_and_holiday_names], dim=1
-            )
+        if names:
+            multiplicative_events_tensor = torch.cat([df_tensors[name].unsqueeze(-1) for name in names], dim=1)
             feature_list.append(multiplicative_events_tensor)
             self.feature_indices["multiplicative_events"] = (
                 current_idx,
@@ -282,14 +272,12 @@ def stack_multiplicative_events(
             return current_idx + multiplicative_events_tensor.size(1)
         return current_idx
 
-    def stack_additive_regressors(self, df_tensors, feature_list, current_idx, additive_regressors_names):
+    def stack_additive_regressors(self, df_tensors, feature_list, current_idx, names):
         """
         Stack the additive regressor features.
         """
-        if additive_regressors_names:
-            additive_regressors_tensor = torch.cat(
-                [df_tensors[name].unsqueeze(-1) for name in additive_regressors_names], dim=1
-            )
+        if names:
+            additive_regressors_tensor = torch.cat([df_tensors[name].unsqueeze(-1) for name in names], dim=1)
             feature_list.append(additive_regressors_tensor)
             self.feature_indices["additive_regressors"] = (
                 current_idx,
@@ -298,20 +286,20 @@ def stack_additive_regressors(self, df_tensors, feature_list, current_idx, addit
             return current_idx + additive_regressors_tensor.size(1)
         return current_idx
 
-    def stack_multiplicative_regressors(self, df_tensors, feature_list, current_idx, multiplicative_regressors_names):
+    def stack_multiplicative_regressors(self, df_tensors, feature_list, current_idx, names):
         """
         Stack the multiplicative regressor features.
         """
-        if multiplicative_regressors_names:
+        if names:
             multiplicative_regressors_tensor = torch.cat(
-                [df_tensors[name].unsqueeze(-1) for name in multiplicative_regressors_names], dim=1
+                [df_tensors[name].unsqueeze(-1) for name in names], dim=1
             )  # Shape: [batch_size, num_multiplicative_regressors, 1]
             feature_list.append(multiplicative_regressors_tensor)
             self.feature_indices["multiplicative_regressors"] = (
                 current_idx,
-                current_idx + len(multiplicative_regressors_names) - 1,
+                current_idx + len(names) - 1,
             )
-            return current_idx + len(multiplicative_regressors_names)
+            return current_idx + len(names)
         return current_idx
 
     def stack_seasonalities(self, df_tensors, feature_list, current_idx, config, seasonalities):

From 5a06d398d2da9f26047c909003b767be92b19939 Mon Sep 17 00:00:00 2001
From: ourownstory <ourownstory@users.noreply.github.com>
Date: Wed, 11 Sep 2024 18:28:58 -0700
Subject: [PATCH 26/30] fix seasons

---
 neuralprophet/time_dataset.py | 6 ++++--
 1 file changed, 4 insertions(+), 2 deletions(-)

diff --git a/neuralprophet/time_dataset.py b/neuralprophet/time_dataset.py
index abb7b1596..cf0be6ffb 100644
--- a/neuralprophet/time_dataset.py
+++ b/neuralprophet/time_dataset.py
@@ -130,8 +130,10 @@ def stack_all_features(self):
             "multiplicative_events": {"names": self.multiplicative_event_and_holiday_names},
             "additive_regressors": {"names": self.additive_regressors_names},
             "multiplicative_regressors": {"names": self.multiplicative_regressors_names},
-            "seasonalities": {"config": self.config_seasonality, "seasonalities": self.seasonalities},
         }
+        if self.config_seasonality is not None and hasattr(self.config_seasonality, "periods"):
+            component_args["seasonalities"] = {"config": self.config_seasonality, "seasonalities": self.seasonalities}
+
         for component_name, args in component_args.items():
             current_idx = self.components_stacker.stack(
                 component_name=component_name,
@@ -170,7 +172,7 @@ def stack_all_features(self):
 
         # if self.config_seasonality is not None and hasattr(self.config_seasonality, "periods"):
         #     current_idx = self.components_stacker.stack_seasonalities(
-        #         feature_list, current_idx, self.config_seasonality, self.seasonalities
+        #         None, feature_list, current_idx, self.config_seasonality, self.seasonalities
         #     )
 
         # Concatenate all features into one big tensor

From 74186a6aa1b809ef3a5ab1069c618310898886b9 Mon Sep 17 00:00:00 2001
From: ourownstory <ourownstory@users.noreply.github.com>
Date: Wed, 11 Sep 2024 18:35:30 -0700
Subject: [PATCH 27/30] explicit update of feature_list and no double returns

---
 neuralprophet/time_dataset.py       | 34 +----------------------------
 neuralprophet/utils_time_dataset.py | 32 +++++++++++++--------------
 2 files changed, 17 insertions(+), 49 deletions(-)

diff --git a/neuralprophet/time_dataset.py b/neuralprophet/time_dataset.py
index cf0be6ffb..3bd04d086 100644
--- a/neuralprophet/time_dataset.py
+++ b/neuralprophet/time_dataset.py
@@ -135,7 +135,7 @@ def stack_all_features(self):
             component_args["seasonalities"] = {"config": self.config_seasonality, "seasonalities": self.seasonalities}
 
         for component_name, args in component_args.items():
-            current_idx = self.components_stacker.stack(
+            feature_list, current_idx = self.components_stacker.stack(
                 component_name=component_name,
                 df_tensors=self.df_tensors,
                 feature_list=feature_list,
@@ -143,38 +143,6 @@ def stack_all_features(self):
                 **args,
             )
 
-        # # Call individual stacking functions
-        # current_idx = self.components_stacker.stack_time(self.df_tensors, feature_list, current_idx)
-        # current_idx = self.components_stacker.stack_targets(self.df_tensors, feature_list, current_idx)
-
-        # current_idx = self.components_stacker.stack(
-        #     component_name="lags",
-        #     df_tensors=self.df_tensors,
-        #     feature_list=feature_list,
-        #     current_idx=current_idx,
-        #     n_lags=self.config_ar.n_lags,
-        # )
-        # current_idx = self.components_stacker.stack_lagged_regressors(
-        #     self.df_tensors, feature_list, current_idx, self.config_lagged_regressors
-        # )
-        # current_idx = self.components_stacker.stack_additive_events(
-        #     self.df_tensors, feature_list, current_idx, self.additive_event_and_holiday_names
-        # )
-        # current_idx = self.components_stacker.stack_multiplicative_events(
-        #     self.df_tensors, feature_list, current_idx, self.multiplicative_event_and_holiday_names
-        # )
-        # current_idx = self.components_stacker.stack_additive_regressors(
-        #     self.df_tensors, feature_list, current_idx, self.additive_regressors_names
-        # )
-        # current_idx = self.components_stacker.stack_multiplicative_regressors(
-        #     self.df_tensors, feature_list, current_idx, self.multiplicative_regressors_names
-        # )
-
-        # if self.config_seasonality is not None and hasattr(self.config_seasonality, "periods"):
-        #     current_idx = self.components_stacker.stack_seasonalities(
-        #         None, feature_list, current_idx, self.config_seasonality, self.seasonalities
-        #     )
-
         # Concatenate all features into one big tensor
         self.all_features = torch.cat(feature_list, dim=1)  # Concatenating along the second dimension
 
diff --git a/neuralprophet/utils_time_dataset.py b/neuralprophet/utils_time_dataset.py
index a3456baf8..fbafaaa4a 100644
--- a/neuralprophet/utils_time_dataset.py
+++ b/neuralprophet/utils_time_dataset.py
@@ -200,7 +200,7 @@ def stack_time(self, df_tensors, feature_list, current_idx):
         time_tensor = df_tensors["t"].unsqueeze(-1)  # Shape: [T, 1]
         feature_list.append(time_tensor)
         self.feature_indices["time"] = (current_idx, current_idx)
-        return current_idx + 1
+        return feature_list, current_idx + 1
 
     def stack_lags(self, df_tensors, feature_list, current_idx, n_lags):
         """
@@ -210,8 +210,8 @@ def stack_lags(self, df_tensors, feature_list, current_idx, n_lags):
             lags_tensor = df_tensors["y_scaled"].unsqueeze(-1)
             feature_list.append(lags_tensor)
             self.feature_indices["lags"] = (current_idx, current_idx)
-            return current_idx + 1
-        return current_idx
+            current_idx = current_idx + 1
+        return feature_list, current_idx
 
     def stack_targets(self, df_tensors, feature_list, current_idx):
         """
@@ -221,8 +221,8 @@ def stack_targets(self, df_tensors, feature_list, current_idx):
             targets_tensor = df_tensors["y_scaled"].unsqueeze(-1)
             feature_list.append(targets_tensor)
             self.feature_indices["targets"] = (current_idx, current_idx)
-            return current_idx + 1
-        return current_idx
+            current_idx = current_idx + 1
+        return feature_list, current_idx
 
     def stack_lagged_regressors(self, df_tensors, feature_list, current_idx, config):
         """
@@ -238,8 +238,8 @@ def stack_lagged_regressors(self, df_tensors, feature_list, current_idx, config)
                     current_idx + i,
                     current_idx + i + 1,
                 )
-            return current_idx + num_features
-        return current_idx
+            current_idx = current_idx + num_features
+        return feature_list, current_idx
 
     def stack_additive_events(self, df_tensors, feature_list, current_idx, names):
         """
@@ -255,8 +255,8 @@ def stack_additive_events(self, df_tensors, feature_list, current_idx, names):
                 current_idx,
                 current_idx + additive_events_tensor.size(1) - 1,
             )
-            return current_idx + additive_events_tensor.size(1)
-        return current_idx
+            current_idx = current_idx + additive_events_tensor.size(1)
+        return feature_list, current_idx
 
     def stack_multiplicative_events(self, df_tensors, feature_list, current_idx, names):
         """
@@ -269,8 +269,8 @@ def stack_multiplicative_events(self, df_tensors, feature_list, current_idx, nam
                 current_idx,
                 current_idx + multiplicative_events_tensor.size(1) - 1,
             )
-            return current_idx + multiplicative_events_tensor.size(1)
-        return current_idx
+            current_idx = current_idx + multiplicative_events_tensor.size(1)
+        return feature_list, current_idx
 
     def stack_additive_regressors(self, df_tensors, feature_list, current_idx, names):
         """
@@ -283,8 +283,8 @@ def stack_additive_regressors(self, df_tensors, feature_list, current_idx, names
                 current_idx,
                 current_idx + additive_regressors_tensor.size(1) - 1,
             )
-            return current_idx + additive_regressors_tensor.size(1)
-        return current_idx
+            current_idx = current_idx + additive_regressors_tensor.size(1)
+        return feature_list, current_idx
 
     def stack_multiplicative_regressors(self, df_tensors, feature_list, current_idx, names):
         """
@@ -299,8 +299,8 @@ def stack_multiplicative_regressors(self, df_tensors, feature_list, current_idx,
                 current_idx,
                 current_idx + len(names) - 1,
             )
-            return current_idx + len(names)
-        return current_idx
+            current_idx = current_idx + len(names)
+        return feature_list, current_idx
 
     def stack_seasonalities(self, df_tensors, feature_list, current_idx, config, seasonalities):
         """
@@ -316,4 +316,4 @@ def stack_seasonalities(self, df_tensors, feature_list, current_idx, config, sea
                     current_idx + seasonal_tensor.size(1),
                 )
                 current_idx += seasonal_tensor.size(1)
-        return current_idx
+        return feature_list, current_idx

From d3219d9acce3ec10584d58ffe6f3a13d3ffe142d Mon Sep 17 00:00:00 2001
From: ourownstory <ourownstory@users.noreply.github.com>
Date: Wed, 11 Sep 2024 18:47:47 -0700
Subject: [PATCH 28/30] conform seasonalities

---
 neuralprophet/time_dataset.py       | 8 +++++---
 neuralprophet/utils_time_dataset.py | 6 +++---
 2 files changed, 8 insertions(+), 6 deletions(-)

diff --git a/neuralprophet/time_dataset.py b/neuralprophet/time_dataset.py
index 3bd04d086..4c25a6be5 100644
--- a/neuralprophet/time_dataset.py
+++ b/neuralprophet/time_dataset.py
@@ -104,6 +104,7 @@ def __init__(
         self.df["ds"] = self.df["ds"].apply(lambda x: x.timestamp())  # Convert to Unix timestamp in seconds
         self.df_tensors["ds"] = torch.tensor(self.df["ds"].values, dtype=torch.int64)
 
+        self.seasonalities = None
         if self.config_seasonality is not None and hasattr(self.config_seasonality, "periods"):
             self.calculate_seasonalities()
 
@@ -121,6 +122,9 @@ def stack_all_features(self):
         feature_list = []
         current_idx = 0
 
+        # Add seasonalities to df_tensors, this needs to be done after create_sample2index_map, before stacking.
+        self.df_tensors["seasonalities"] = self.seasonalities
+
         component_args: dict = {
             "time": {},
             "targets": {},
@@ -130,10 +134,8 @@ def stack_all_features(self):
             "multiplicative_events": {"names": self.multiplicative_event_and_holiday_names},
             "additive_regressors": {"names": self.additive_regressors_names},
             "multiplicative_regressors": {"names": self.multiplicative_regressors_names},
+            "seasonalities": {"config": self.config_seasonality},
         }
-        if self.config_seasonality is not None and hasattr(self.config_seasonality, "periods"):
-            component_args["seasonalities"] = {"config": self.config_seasonality, "seasonalities": self.seasonalities}
-
         for component_name, args in component_args.items():
             feature_list, current_idx = self.components_stacker.stack(
                 component_name=component_name,
diff --git a/neuralprophet/utils_time_dataset.py b/neuralprophet/utils_time_dataset.py
index fbafaaa4a..a2f8a3733 100644
--- a/neuralprophet/utils_time_dataset.py
+++ b/neuralprophet/utils_time_dataset.py
@@ -302,13 +302,13 @@ def stack_multiplicative_regressors(self, df_tensors, feature_list, current_idx,
             current_idx = current_idx + len(names)
         return feature_list, current_idx
 
-    def stack_seasonalities(self, df_tensors, feature_list, current_idx, config, seasonalities):
+    def stack_seasonalities(self, df_tensors, feature_list, current_idx, config):
         """
         Stack the seasonality features.
         """
-        # TODO conform to other stack functions, using df_tensors
+        # if config is not None and hasattr(config, "periods"):
         if config and config.periods:
-            for seasonality_name, features in seasonalities.items():
+            for seasonality_name, features in df_tensors["seasonalities"].items():
                 seasonal_tensor = features
                 feature_list.append(seasonal_tensor)
                 self.feature_indices[f"seasonality_{seasonality_name}"] = (

From d37b534bfefdc4724fac46eca146194c3128b18c Mon Sep 17 00:00:00 2001
From: ourownstory <ourownstory@users.noreply.github.com>
Date: Wed, 11 Sep 2024 18:59:45 -0700
Subject: [PATCH 29/30] move stack_all to stacker

---
 neuralprophet/forecaster.py         |  4 ----
 neuralprophet/time_dataset.py       | 19 +++----------------
 neuralprophet/utils_time_dataset.py | 25 ++++++++++++++++++++++++-
 3 files changed, 27 insertions(+), 21 deletions(-)

diff --git a/neuralprophet/forecaster.py b/neuralprophet/forecaster.py
index 2681f9e1f..6f14e2684 100644
--- a/neuralprophet/forecaster.py
+++ b/neuralprophet/forecaster.py
@@ -1209,7 +1209,6 @@ def fit(
             max_lags=self.config_model.max_lags,
             config_seasonality=self.config_seasonality,
             lagged_regressor_config=self.config_lagged_regressors,
-            feature_indices={},
         )
         dataset = self._create_dataset(df, predict_mode=False, components_stacker=train_components_stacker)
         # Determine the max_number of epochs
@@ -1252,7 +1251,6 @@ def fit(
                 n_forecasts=self.config_model.n_forecasts,
                 config_seasonality=self.config_seasonality,
                 lagged_regressor_config=self.config_lagged_regressors,
-                feature_indices={},
             )
             dataset_val = self._create_dataset(df_val, predict_mode=False, components_stacker=val_components_stacker)
             loader_val = DataLoader(dataset_val, batch_size=min(1024, len(dataset_val)), shuffle=False, drop_last=False)
@@ -1490,7 +1488,6 @@ def test(self, df: pd.DataFrame, verbose: bool = True):
             max_lags=self.config_model.max_lags,
             config_seasonality=self.config_seasonality,
             lagged_regressor_config=self.config_lagged_regressors,
-            feature_indices={},
         )
         dataset = self._create_dataset(df, predict_mode=False, components_stacker=components_stacker)
         self.model.set_components_stacker(components_stacker, mode="test")
@@ -2988,7 +2985,6 @@ def _predict_raw(self, df, df_name, include_components=False):
             max_lags=self.config_model.max_lags,
             config_seasonality=self.config_seasonality,
             lagged_regressor_config=self.config_lagged_regressors,
-            feature_indices={},
         )
         dataset = self._create_dataset(df, predict_mode=True, components_stacker=components_stacker)
         self.model.set_components_stacker(components_stacker, mode="predict")
diff --git a/neuralprophet/time_dataset.py b/neuralprophet/time_dataset.py
index 4c25a6be5..1044d63eb 100644
--- a/neuralprophet/time_dataset.py
+++ b/neuralprophet/time_dataset.py
@@ -111,20 +111,16 @@ def __init__(
         # Construct index map
         self.sample2index_map, self.length = self.create_sample2index_map(self.df, self.df_tensors)
 
+        # Stack all features into one large tensor
         self.components_stacker = components_stacker
-
-        self.stack_all_features()
+        self.all_features = self.stack_all_features()
 
     def stack_all_features(self):
         """
         Stack all features into one large tensor by calling individual stacking methods.
         """
-        feature_list = []
-        current_idx = 0
-
         # Add seasonalities to df_tensors, this needs to be done after create_sample2index_map, before stacking.
         self.df_tensors["seasonalities"] = self.seasonalities
-
         component_args: dict = {
             "time": {},
             "targets": {},
@@ -136,17 +132,8 @@ def stack_all_features(self):
             "multiplicative_regressors": {"names": self.multiplicative_regressors_names},
             "seasonalities": {"config": self.config_seasonality},
         }
-        for component_name, args in component_args.items():
-            feature_list, current_idx = self.components_stacker.stack(
-                component_name=component_name,
-                df_tensors=self.df_tensors,
-                feature_list=feature_list,
-                current_idx=current_idx,
-                **args,
-            )
 
-        # Concatenate all features into one big tensor
-        self.all_features = torch.cat(feature_list, dim=1)  # Concatenating along the second dimension
+        return self.components_stacker.stack_all_features(self.df_tensors, component_args)
 
     def calculate_seasonalities(self):
         """Computes Fourier series components with the specified frequency and order."""
diff --git a/neuralprophet/utils_time_dataset.py b/neuralprophet/utils_time_dataset.py
index a2f8a3733..4aa756f44 100644
--- a/neuralprophet/utils_time_dataset.py
+++ b/neuralprophet/utils_time_dataset.py
@@ -84,11 +84,34 @@ def stack(self, component_name, df_tensors, feature_list, current_idx, **kwargs)
             current_idx: the current index in the stack of features.
         """
         assert component_name in self.stack_func, f"Unknown component name: {component_name}"
-        # this is currently not working for seasonalities
         return self.stack_func[component_name](
             df_tensors=df_tensors, feature_list=feature_list, current_idx=current_idx, **kwargs
         )
 
+    def stack_all_features(self, df_tensors, component_args):
+        """
+        Stack all features into one large tensor by calling individual stacking methods.
+        Concatenation along dimension second dimension (dim=1).
+
+        Args:
+            df_tensors: Dictionary containing the tensors for different features.
+            component_args: Dictionary containing the configuration of different components.
+        """
+        feature_list = []
+        current_idx = 0
+
+        for component_name, args in component_args.items():
+            feature_list, current_idx = self.stack(
+                component_name=component_name,
+                df_tensors=df_tensors,
+                feature_list=feature_list,
+                current_idx=current_idx,
+                **args,
+            )
+
+        # Concatenate all features into one big tensor
+        return torch.cat(feature_list, dim=1)  # Concatenating along the second dimension
+
     def unstack_targets(self, batch_tensor):
         targets_start_idx, targets_end_idx = self.feature_indices["targets"]
         if self.max_lags > 0:

From a2f5c4ac4406ec3d1db26d6a17586d6a582efdb5 Mon Sep 17 00:00:00 2001
From: ourownstory <ourownstory@users.noreply.github.com>
Date: Wed, 11 Sep 2024 19:10:09 -0700
Subject: [PATCH 30/30] ruff

---
 neuralprophet/time_net.py | 1 -
 1 file changed, 1 deletion(-)

diff --git a/neuralprophet/time_net.py b/neuralprophet/time_net.py
index 2d8a2d474..e6fd8d7e0 100644
--- a/neuralprophet/time_net.py
+++ b/neuralprophet/time_net.py
@@ -21,7 +21,6 @@
     reg_func_trend,
     reg_func_trend_glocal,
 )
-from neuralprophet.utils_time_dataset import ComponentStacker
 from neuralprophet.utils_torch import init_parameter, interprete_model
 
 log = logging.getLogger("NP.time_net")