Merge pull request #5 from nimish/hot-fix

Author: nimish

src/tsgam_estimator/tsgam_estimator.py

@@ -1688,7 +1688,9 @@ def _fit_ar_model(self, X_array: ndarray, y: ndarray, time_indices: ndarray) ->
             self.ar_noise_loc_ = None
             self.ar_noise_scale_ = None
 
-    def predict(self, X: pd.DataFrame) -> ndarray:
+    def predict(self, X: pd.DataFrame,
+                remove_periodic : bool = False, remove_exogenous : bool = False, 
+                remove_trend : bool = False) -> ndarray:
         """
         Predict target values for new data.
 
@@ -1739,7 +1741,7 @@ def predict(self, X: pd.DataFrame) -> ndarray:
         timestamps, X_array = self._ensure_timestamp_index(X)
 
         # Prediction data must be regularly spaced with no gaps
-        self._validate_frequency(timestamps, self.freq_)
+        self._validate_frequency(timestamps, self.freq_, allow_gaps=True)
 
         # Convert timestamps to indices using stored reference
         time_indices = self._timestamps_to_indices(timestamps, self.time_reference_)
@@ -1754,7 +1756,7 @@ def predict(self, X: pd.DataFrame) -> ndarray:
         predictions = np.full(len(X_array), constant_value)
 
         # Add exogenous terms if present
-        if self.config.exog_config:
+        if self.config.exog_config and not remove_exogenous:
             for ix, exog_cfg in enumerate(self.config.exog_config):
                 exog_var = X_array[:, ix]
 
@@ -1798,7 +1800,7 @@ def predict(self, X: pd.DataFrame) -> ndarray:
                 predictions += exog_pred
 
         # Add Fourier terms if present
-        if self.config.multi_periodic_config:
+        if self.config.multi_periodic_config and not remove_periodic:
             # Check for NaN in time_indices
             if np.any(np.isnan(time_indices)):
                 raise ValueError("Time indices contain NaN. Check timestamp conversion.")
@@ -1866,7 +1868,8 @@ def predict(self, X: pd.DataFrame) -> ndarray:
             predictions += fourier_contrib
 
         # Add trend term if present
-        if self.config.trend_config is not None and self.config.trend_config.trend_type != TrendType.NONE and 'trend' in self.variables_:
+        if (self.config.trend_config is not None and self.config.trend_config.trend_type != TrendType.NONE 
+            and 'trend' in self.variables_ and not remove_trend):
             trend = self.variables_['trend'].value
             if trend is None:
                 raise ValueError("Trend coefficients are None. Model may not have converged.")
@@ -1906,7 +1909,8 @@ def predict(self, X: pd.DataFrame) -> ndarray:
                     trend_extended[n_periods_fit:] = trend[-1]
 
                 trend = trend_extended
-
+            elif n_periods_pred < n_periods_fit:
+                trend = trend[:T_pred.shape[1]]
             # Add trend term to predictions
             predictions += T_pred @ trend