Resolving issue #260 ( added exogenous variables to Chronos-2 )

timecopilot/models/foundation/chronos.py

@@ -84,6 +84,7 @@ def __init__(
         self.repo_id = repo_id
         self.batch_size = batch_size
         self.alias = alias
+        self.supports_exogenous = "chronos-2" in repo_id
 
     @contextmanager
     def _get_model(self) -> BaseChronosPipeline:
@@ -161,6 +162,76 @@ def _predict(
             fcsts_mean_np = fcsts_mean.numpy()  # type: ignore
             fcsts_quantiles_np = None
         return fcsts_mean_np, fcsts_quantiles_np
+
+
+
+    def _forecast_chronos2_df(
+        self,
+        df: pd.DataFrame,
+        h: int,
+        freq: str | None,
+        qc: QuantileConverter,
+    ) -> pd.DataFrame:
+
+        id_col = "unique_id"
+        ts_col = "ds"
+        target_col = "y"
+
+        required = {id_col, ts_col, target_col}
+        if not required.issubset(df.columns):
+            raise ValueError("missing required columns")
+
+        base_cols = [id_col, ts_col, target_col]
+        exog_cols = []
+        for col in df.columns:
+            if col not in base_cols:
+                exog_cols.append(col)
+
+        cols = base_cols + exog_cols
+        context_df = df[cols].copy()
+        context_df = context_df.sort_values([id_col, ts_col]).reset_index(drop=True)
+
+        # model call
+        with self._get_model() as model:
+            if not isinstance(model, Chronos2Pipeline):
+                raise ValueError("wrong model type")
+
+            pred_df = model.predict_df(
+                context_df,
+                future_df=None,
+                prediction_length=h,
+                quantile_levels=qc.quantiles,
+                id_column=id_col,
+                timestamp_column=ts_col,
+                target=target_col,
+            )
+
+        pred_df = pred_df.rename(columns={id_col: "unique_id", ts_col: "ds"})
+
+        if "predictions" not in pred_df.columns:
+            raise ValueError("predictions column missing")
+
+        # main forecast
+        pred_df[self.alias] = pred_df["predictions"]
+
+        if qc.quantiles is not None:
+            for i, q in enumerate(qc.quantiles):
+                raw_col = str(q)
+                if raw_col in pred_df.columns:
+                    pred_df[f"{self.alias}-q-{int(q * 100)}"] = pred_df[raw_col]
+
+            pred_df = qc.maybe_convert_quantiles_to_level(
+                pred_df,
+                models=[self.alias],
+            )
+
+        # final selection
+        final_cols = ["unique_id", "ds", self.alias]
+        for col in pred_df.columns:
+            if col.startswith(self.alias + "-"):
+                final_cols.append(col)
+
+        return pred_df[final_cols].copy()
 
     def forecast(
         self,
@@ -218,6 +289,15 @@ def forecast(
         """
         freq = self._maybe_infer_freq(df, freq)
         qc = QuantileConverter(level=level, quantiles=quantiles)
+
+        if "chronos-2" in self.repo_id:
+            return self._forecast_chronos2_df(
+                df=df,
+                h=h,
+                freq=freq,
+                qc=qc,
+            )
+
         dataset = TimeSeriesDataset.from_df(df, batch_size=self.batch_size)
         fcst_df = dataset.make_future_dataframe(h=h, freq=freq)
         with self._get_model() as model:

timecopilot/models/utils/forecaster.py

@@ -251,10 +251,15 @@ def cross_validation(
             freq=pd.tseries.frequencies.to_offset(freq),
             step_size=h if step_size is None else step_size,
         )
+        supports_exogenous = getattr(self, "supports_exogenous", False)
         for _, (cutoffs, train, valid) in tqdm(enumerate(splits)):
-            if len(valid.columns) > 3:
+
+            has_exog = len(valid.columns) > 3  # columns beyond unique_id, ds, y
+
+            if has_exog and not supports_exogenous:
                 raise NotImplementedError(
-                    "Cross validation with exogenous variables is not yet supported."
+                    "Cross validation with exogenous variables is not yet supported for "
+                    f"model {getattr(self, 'alias', type(self).__name__)}."
                 )
             y_pred = self.forecast(
                 df=train,