diff --git a/flaml/automl/automl.py b/flaml/automl/automl.py index 1daccaebe1..9aef4e8eb1 100644 --- a/flaml/automl/automl.py +++ b/flaml/automl/automl.py @@ -882,15 +882,25 @@ def predict( if estimator is None: logger.warning("No estimator is trained. Please run fit with enough budget.") return None + orig_index = getattr(X, "index", None) + if callable(orig_index): + orig_index = None X = self._state.task.preprocess(X, self._transformer) y_pred = estimator.predict(X, **pred_kwargs) if isinstance(y_pred, np.ndarray) and y_pred.ndim > 1 and isinstance(y_pred, np.ndarray): y_pred = y_pred.flatten() if self._label_transformer: - return self._label_transformer.inverse_transform(Series(y_pred.astype(int))) - else: - return y_pred + y_pred = self._label_transformer.inverse_transform(Series(y_pred.astype(int))) + if ( + orig_index is not None + and not self._state.task.is_ts_forecastpanel() + and isinstance(y_pred, (Series, DataFrame)) + and len(y_pred) == len(orig_index) + ): + y_pred = y_pred.copy() + y_pred.index = orig_index + return y_pred def predict_proba(self, X, **pred_kwargs): """Predict the probability of each class from features, only works for diff --git a/flaml/automl/task/time_series_task.py b/flaml/automl/task/time_series_task.py index 9f16840891..9219ff2759 100644 --- a/flaml/automl/task/time_series_task.py +++ b/flaml/automl/task/time_series_task.py @@ -1,3 +1,4 @@ +import datetime import logging import time from typing import List @@ -126,13 +127,17 @@ def validate_data( else: target_names = label + time_col_inferred = False if self.time_col is None: + time_col_inferred = True if isinstance(X_train_all, pd.DataFrame): assert dataframe is None, "One of dataframe and X arguments must be None" - self.time_col = X_train_all.columns[0] + if len(X_train_all.columns) > 0: + self.time_col = X_train_all.columns[0] elif dataframe is not None: assert X_train_all is None, "One of dataframe and X arguments must be None" - self.time_col = dataframe.columns[0] + if len(dataframe.columns) > 0: + self.time_col = dataframe.columns[0] else: self.time_col = "ds" @@ -141,17 +146,68 @@ def validate_data( if X_train_all is not None: assert y_train_all is not None, "If X_train_all is not None, y_train_all must also be" assert dataframe is None, "If X_train_all is provided, dataframe must be None" + if isinstance(X_train_all, pd.DataFrame): + needs_promotion = self._needs_datetime_index_promotion( + X_train_all, + time_col=self.time_col, + is_inferred=time_col_inferred, + allow_index_fallback=False, + ) + X_train_all, y_train_all = self._align_y_to_X_and_promote( + X_train_all, + y_train_all, + needs_promotion=needs_promotion, + ) + if needs_promotion: + X_train_all, promoted_time_col = self._promote_datetime_index_if_needed( + X_train_all, + time_col=self.time_col, + is_inferred=time_col_inferred, + allow_index_fallback=False, + reserved_cols=target_names, + ) + if promoted_time_col is not None: + self.time_col = promoted_time_col + if isinstance(y_train_all, (pd.DataFrame, pd.Series)): + y_train_all = y_train_all.reset_index(drop=True) + if not time_col_inferred and self.time_col not in X_train_all.columns: + raise ValueError( + f"For '{TS_FORECAST}' task, time column {self.time_col} must contain timestamp values." + ) dataframe = TimeSeriesDataset.to_dataframe(X_train_all, y_train_all, target_names, self.time_col) elif dataframe is not None: assert label is not None, "A label or list of labels must be provided." assert isinstance(dataframe, pd.DataFrame), "dataframe must be a pandas DataFrame" assert label in dataframe.columns, f"{label} must a column name in dataframe" + needs_promotion = self._needs_datetime_index_promotion( + dataframe, + time_col=self.time_col, + is_inferred=time_col_inferred, + allow_index_fallback=False, + ) + if needs_promotion: + dataframe, promoted_time_col = self._promote_datetime_index_if_needed( + dataframe, + time_col=self.time_col, + is_inferred=time_col_inferred, + allow_index_fallback=False, + reserved_cols=target_names, + ) + if promoted_time_col is not None: + self.time_col = promoted_time_col + if not time_col_inferred and self.time_col not in dataframe.columns: + raise ValueError( + f"For '{TS_FORECAST}' task, time column {self.time_col} must contain timestamp values." + ) else: raise ValueError("Must supply either X_train_all and y_train_all, or dataframe and label") + if self.time_col is None: + self.time_col = "ds" + try: - dataframe.loc[:, self.time_col] = pd.to_datetime(dataframe[self.time_col]) + dataframe[self.time_col] = pd.to_datetime(dataframe[self.time_col]) except Exception: raise ValueError( f"For '{TS_FORECAST}' task, time column {self.time_col} must contain timestamp values." @@ -161,6 +217,33 @@ def validate_data( if X_val is not None: assert y_val is not None, "If X_val is not None, y_val must also be" + if isinstance(X_val, pd.DataFrame): + needs_promotion_val = self._needs_datetime_index_promotion( + X_val, + time_col=self.time_col, + is_inferred=False, + allow_index_fallback=True, + ) + X_val, y_val = self._align_y_to_X_and_promote( + X_val, + y_val, + needs_promotion=needs_promotion_val, + ) + if needs_promotion_val: + X_val, promoted_time_col_val = self._promote_datetime_index_if_needed( + X_val, + time_col=self.time_col, + is_inferred=False, + allow_index_fallback=True, + reserved_cols=target_names, + ) + if promoted_time_col_val is not None: + if isinstance(y_val, (pd.DataFrame, pd.Series)): + y_val = y_val.reset_index(drop=True) + if self.time_col not in X_val.columns: + raise ValueError( + f"For '{TS_FORECAST}' task, time column {self.time_col} must contain timestamp values." + ) val_df = TimeSeriesDataset.to_dataframe(X_val, y_val, target_names, self.time_col) val_len = len(val_df) else: @@ -366,9 +449,11 @@ def _preprocess(self, X, transformer=None): X = pd.DataFrame( dict( [ - (transformer._str_columns[idx], X[idx]) - if isinstance(X[0], List) - else (transformer._str_columns[idx], [X[idx]]) + ( + (transformer._str_columns[idx], X[idx]) + if isinstance(X[0], List) + else (transformer._str_columns[idx], [X[idx]]) + ) for idx in range(len(X)) ] ) @@ -385,10 +470,168 @@ def _preprocess(self, X, transformer=None): X = transformer.transform(X) return X + @staticmethod + def _is_valid_timestamp_col(df, col): + """Check if col in df is a valid datetime-like column without numeric nanosecond coercion.""" + if col is None or not isinstance(df, pd.DataFrame) or col not in df.columns: + return False + series = df[col] + if pd.api.types.is_datetime64_any_dtype(series): + return True + if pd.api.types.is_numeric_dtype(series) or pd.api.types.is_bool_dtype(series): + return False + non_null = series.dropna() + if non_null.empty: + return False + first_val = non_null.iloc[0] + if not isinstance(first_val, (str, datetime.date, datetime.datetime, pd.Timestamp, np.datetime64)): + return False + try: + import warnings + + with warnings.catch_warnings(): + warnings.simplefilter("ignore") + converted = pd.to_datetime(series, errors="raise") + return not converted.isna().all() + except Exception: + return False + + @classmethod + def _needs_datetime_index_promotion(cls, df, time_col=None, is_inferred=False, allow_index_fallback=False): + """Determine whether df's DatetimeIndex needs to be promoted to a column.""" + if not isinstance(df, pd.DataFrame) or not isinstance(df.index, pd.DatetimeIndex): + return False + if allow_index_fallback: + return time_col is not None and time_col not in df.columns + if is_inferred: + if time_col is not None and cls._is_valid_timestamp_col(df, time_col): + return False + return True + return False + + @classmethod + def _align_y_to_X_and_promote(cls, X, y, needs_promotion=False): + """Align y to X before index promotion or dataset construction. + + When needs_promotion is True, X's DatetimeIndex will be promoted to the time column. + In that case, if y also has a DatetimeIndex, validate that index labels match and + align y to X by label. If y has a RangeIndex (or is a Series without DatetimeIndex), + pair y with X positionally. + + When needs_promotion is False (e.g. X already has a valid timestamp column), + preserve the baseline positional contract for Series (y.index = X.index). + For DataFrame targets, align by label if unique labels match, or raise ValueError. + """ + if not isinstance(X, pd.DataFrame) or not isinstance(y, (pd.DataFrame, pd.Series)): + return X, y + + if len(X) != len(y): + raise ValueError(f"Target length ({len(y)}) does not match feature length ({len(X)}).") + + # When X's DatetimeIndex is being promoted to the time column: + if needs_promotion: + if isinstance(X.index, pd.DatetimeIndex) and isinstance(y.index, pd.DatetimeIndex): + if not y.index.equals(X.index): + if not X.index.is_unique or not y.index.is_unique: + raise ValueError("Duplicate index labels are not supported.") + try: + if not X.index.difference(y.index).empty or not y.index.difference(X.index).empty: + raise ValueError("Target index labels do not match feature index labels.") + except TypeError: + raise ValueError("Target index labels do not match feature index labels.") + y = y.reindex(X.index) + return X, y + # Positional pairing when promotion is needed and y has RangeIndex / non-DatetimeIndex + y = y.copy() + y.index = X.index + return X, y + + # When promotion is NOT needed (e.g. X already has a timestamp column): + if isinstance(y, pd.Series): + y = y.copy() + y.index = X.index + return X, y + elif isinstance(y, pd.DataFrame): + if not y.index.equals(X.index): + if not X.index.is_unique or not y.index.is_unique: + raise ValueError("Duplicate index labels are not supported.") + try: + if not X.index.difference(y.index).empty or not y.index.difference(X.index).empty: + raise ValueError("Target index labels do not match feature index labels.") + except TypeError: + raise ValueError("Target index labels do not match feature index labels.") + y = y.reindex(X.index) + return X, y + + return X, y + + @classmethod + def _promote_datetime_index_if_needed( + cls, df, time_col=None, is_inferred=False, allow_index_fallback=False, reserved_cols=None + ): + """Promote a DatetimeIndex to a guaranteed-unique column if needed. + + Args: + df: A pandas DataFrame. + time_col: The current timestamp column name (if known). + is_inferred: Whether time_col was inferred rather than explicitly specified. + allow_index_fallback: Whether to allow falling back to df.index when time_col is absent. + reserved_cols: Additional column names to avoid colliding with (e.g. target names). + + Returns: + Tuple of (df_processed, promoted_col_name or None). + """ + if not cls._needs_datetime_index_promotion( + df, + time_col=time_col, + is_inferred=is_inferred, + allow_index_fallback=allow_index_fallback, + ): + return df, None + + if is_inferred: + base_name = "ds" if df.index.name is None else df.index.name + promoted_col = base_name + # Guarantee unique column name avoiding collisions with existing columns and reserved names + reserved = set(reserved_cols) if reserved_cols else set() + target_col = promoted_col + i = 1 + while target_col in df.columns or target_col in reserved: + target_col = f"{promoted_col}_{i}" + i += 1 + else: + target_col = time_col + + df_out = df.copy() + df_out.insert(0, target_col, df.index) + df_out = df_out.reset_index(drop=True) + logger.info(f"Timestamp column not found, promoted DataFrame index as '{target_col}'.") + return df_out, target_col + def preprocess(self, X, transformer=None): + orig_index = getattr(X, "index", None) + if isinstance(X, pd.DataFrame): + X, _ = self._promote_datetime_index_if_needed( + X, + time_col=self.time_col, + is_inferred=False, + allow_index_fallback=True, + reserved_cols=self.target_names, + ) if isinstance(X, (pd.DataFrame, np.ndarray, pd.Series)): X = normalize_ts_data(X.copy(), self.target_names, self.time_col) - return self._preprocess(X, transformer) + if ( + isinstance(X, pd.DataFrame) + and self.time_col in X.columns + and not pd.api.types.is_datetime64_any_dtype(X[self.time_col]) + and self._is_valid_timestamp_col(X, self.time_col) + ): + X = X.copy() + X[self.time_col] = pd.to_datetime(X[self.time_col]) + X = self._preprocess(X, transformer) + if orig_index is not None and isinstance(X, pd.DataFrame) and len(X) == len(orig_index): + X.index = orig_index + return X elif isinstance(X, int): return X else: diff --git a/flaml/automl/time_series/ts_data.py b/flaml/automl/time_series/ts_data.py index b1327fe4ad..811de063b7 100644 --- a/flaml/automl/time_series/ts_data.py +++ b/flaml/automl/time_series/ts_data.py @@ -539,6 +539,8 @@ def normalize_ts_data(X_train_all, target_names, time_col, y_train_all=None): if y_train_all is None: return X_train_all else: + if len(y_train_all) != len(X_train_all): + raise ValueError(f"Target length ({len(y_train_all)}) does not match feature length ({len(X_train_all)}).") if isinstance(y_train_all, np.ndarray): # TODO: will need to revisit this when doing multivariate y y_train_all = pd.DataFrame( @@ -549,6 +551,19 @@ def normalize_ts_data(X_train_all, target_names, time_col, y_train_all=None): elif isinstance(y_train_all, pd.Series): y_train_all = pd.DataFrame(y_train_all) y_train_all.index = X_train_all.index + elif isinstance(y_train_all, pd.DataFrame): + if not y_train_all.index.equals(X_train_all.index): + if not X_train_all.index.is_unique or not y_train_all.index.is_unique: + raise ValueError("Duplicate index labels are not supported.") + try: + if ( + not X_train_all.index.difference(y_train_all.index).empty + or not y_train_all.index.difference(X_train_all.index).empty + ): + raise ValueError("Target index labels do not match feature index labels.") + except TypeError: + raise ValueError("Target index labels do not match feature index labels.") + y_train_all = y_train_all.reindex(X_train_all.index) dataframe = pd.concat([X_train_all, y_train_all], axis=1) diff --git a/flaml/automl/time_series/ts_model.py b/flaml/automl/time_series/ts_model.py index 445cf337b8..80ef57cb74 100644 --- a/flaml/automl/time_series/ts_model.py +++ b/flaml/automl/time_series/ts_model.py @@ -226,6 +226,8 @@ def predict(self, X: Union[TimeSeriesDataset, DataFrame], **kwargs): } ) ) + if hasattr(X, "index") and len(out) == len(X.index): + out.index = X.index return out else: @@ -321,10 +323,23 @@ def predict(self, X, **kwargs): X = X.rename(columns={self.time_col: "ds"}) if self._model is not None: + input_index = X.index.copy() + input_times = pd.Index(pd.to_datetime(X["ds"])) if "ds" in X.columns else None X = self._preprocess(X) forecast = self._model.predict(X, **kwargs) out = forecast["yhat"] + if input_times is not None and "ds" in forecast.columns: + forecast_times = pd.Index(pd.to_datetime(forecast["ds"])) + if not forecast_times.equals(input_times): + if not forecast_times.is_unique or not input_times.is_unique: + raise ValueError("Prophet cannot restore prediction order for duplicate timestamps.") + positions = forecast_times.get_indexer(input_times) + if (positions < 0).any(): + raise ValueError("Prophet forecast timestamps do not match prediction input timestamps.") + out = out.iloc[positions] out.name = self.target_names[0] + if len(out) == len(input_index): + out.index = input_index return out else: @@ -356,13 +371,16 @@ def predict(self, X, **kwargs) -> pd.Series: if isinstance(X, DataFrame): if X.shape[0] == 0: return pd.Series([], name=self.target_names[0], dtype=float) - start = X[self.time_col].iloc[0] - end = X[self.time_col].iloc[-1] + requested_dates = pd.DatetimeIndex(pd.to_datetime(X[self.time_col])) + start = requested_dates.min() + end = requested_dates.max() exog = self._preprocess(X[self.regressors]).values if len(self.regressors) else None + in_sample = self.train_end_date is not None and end <= self.train_end_date + if in_sample and exog is not None: + exog = exog[requested_dates.argsort(kind="stable")] if self.train_end_date is not None and start > self.train_end_date: first_forecast_date = self.train_end_date + pd.tseries.frequencies.to_offset(self.frequency) forecast_dates = pd.date_range(start=first_forecast_date, end=end, freq=self.frequency) - requested_dates = pd.DatetimeIndex(X[self.time_col]) positions = forecast_dates.get_indexer(requested_dates) if (positions < 0).any() or (positions[1:] <= positions[:-1]).any(): raise ValueError( @@ -388,6 +406,15 @@ def predict(self, X, **kwargs) -> pd.Series: forecast = self._model.predict(start=start, end=end, exog=exog, **kwargs) else: forecast = self._model.predict(start=start, end=end, **kwargs) + if in_sample: + forecast_dates = forecast.index + if not isinstance(forecast_dates, pd.DatetimeIndex): + forecast_dates = pd.date_range(start=start, periods=len(forecast), freq=self.frequency) + positions = forecast_dates.get_indexer(requested_dates) + if (positions < 0).any(): + raise ValueError("Prediction timestamps are not covered by the model forecast.") + forecast = forecast.iloc[positions] + forecast.index = requested_dates else: raise ValueError( "X needs to be either a pandas Dataframe with dates as the first column" diff --git a/test/automl/test_forecast.py b/test/automl/test_forecast.py index e5c262c3c5..1caea219cb 100644 --- a/test/automl/test_forecast.py +++ b/test/automl/test_forecast.py @@ -222,6 +222,71 @@ def test_numpy_large(): ) +def test_statsmodels_in_sample_forecast_restores_requested_timestamp_order(): + from flaml.automl.time_series.ts_data import TimeSeriesDataset + from flaml.automl.time_series.ts_model import StatsModelsEstimator + + dates = pd.date_range("2026-01-01", periods=4, freq="D") + train_data = pd.DataFrame({"ds": dates, "y": range(4)}) + dataset = TimeSeriesDataset(train_data, time_col="ds", target_names="y") + estimator = StatsModelsEstimator() + estimator.fit(dataset) + estimator.regressors = [] + estimator.enrich = lambda X: X + estimator._model = Mock() + estimator._model.predict.return_value = pd.Series([10.0, 20.0, 30.0, 40.0], index=dates) + requested_dates = dates[[0, 2, 1, 3]].rename("ds") + + forecast = estimator.predict(pd.DataFrame({"ds": requested_dates})) + + estimator._model.predict.assert_called_once_with(start=dates[0], end=dates[-1]) + expected = pd.Series([10.0, 30.0, 20.0, 40.0], index=requested_dates, name="y") + pd.testing.assert_series_equal(forecast, expected) + + +def test_statsmodels_automl_predict_restores_unsorted_timestamps_with_real_model(): + from types import SimpleNamespace + + from statsmodels.tsa.arima.model import ARIMA as StatsmodelsARIMA + + from flaml.automl.time_series.ts_data import TimeSeriesDataset + from flaml.automl.time_series.ts_model import StatsModelsEstimator + + dates = pd.date_range("2026-01-01", periods=12, freq="D") + values = np.array([2.0, 9.0, 1.0, 8.0, 3.0, 12.0, 4.0, 11.0, 5.0, 15.0, 7.0, 17.0]) + target = pd.Series(values, index=dates, name="y") + train_data = pd.DataFrame({"ds": dates, "y": values}) + dataset = TimeSeriesDataset(train_data, time_col="ds", target_names="y") + + estimator = StatsModelsEstimator() + estimator.fit(dataset) + estimator.regressors = [] + estimator._model = StatsmodelsARIMA(target, order=(1, 0, 0)).fit() + + requested_dates = dates[[0, 2, 1, 3]].rename("ds") + X = pd.DataFrame({"ds": requested_dates}, index=requested_dates) + task = SimpleNamespace( + preprocess=lambda frame, transformer: frame, + is_ts_forecastpanel=lambda: False, + ) + automl = AutoML() + automl._state = SimpleNamespace(task=task) + automl._transformer = None + automl._trained_estimator = estimator + automl._label_transformer = None + + actual = automl.predict(X) + + chronological = estimator._model.predict(start=dates[0], end=dates[3]) + assert not np.isclose(chronological.iloc[1], chronological.iloc[2]) + expected = pd.Series( + chronological.reindex(requested_dates).to_numpy(), + index=X.index, + name="y", + ) + pd.testing.assert_series_equal(actual, expected) + + def test_statsmodels_future_forecast_uses_training_boundary(): import pandas as pd diff --git a/test/automl/test_ts_forecast_datetime_index.py b/test/automl/test_ts_forecast_datetime_index.py new file mode 100644 index 0000000000..2ce7c94906 --- /dev/null +++ b/test/automl/test_ts_forecast_datetime_index.py @@ -0,0 +1,741 @@ +"""Regression tests for ts_forecast input where the timestamp is the DataFrame index. + +FLAML expects the timestamp to live in a *column* (`time_col`). When a user passes timestamps +as a `DatetimeIndex`, `validate_data` promotes the index to a column, preserves explicit index +names, ensures label-based alignment for DatetimeIndex targets, and maintains backwards +compatibility for explicit columns and NumPy array inputs across training, validation, and prediction. +""" + +from types import SimpleNamespace + +import numpy as np +import pandas as pd +import pytest + +from flaml import AutoML + + +def _make_train_df(periods=120): + return pd.DataFrame( + {"y": np.sin(np.arange(periods) / 6) * 10 + 50}, + index=pd.date_range("2018-01-01", periods=periods, freq="MS"), + ) + + +def _fit(automl, **kwargs): + settings = { + "task": "ts_forecast", + "label": "y", + "period": 12, + "time_budget": 5, + "estimator_list": ["lgbm"], + "metric": "mape", + "log_file_name": False, + "verbose": 0, + } + settings.update(kwargs) + automl.fit(**settings) + + +def test_unnamed_datetime_index_becomes_the_time_column(): + automl = AutoML() + _fit(automl, dataframe=_make_train_df()) + assert automl._state.task.time_col == "ds" + + +def test_named_datetime_index_keeps_its_name(): + df = _make_train_df() + df.index.name = "month" + automl = AutoML() + _fit(automl, dataframe=df) + assert automl._state.task.time_col == "month" + + +def test_promoted_index_keeps_every_training_row(): + """Pre-promotion the label column was read as epoch nanoseconds, collapsing the 120 real + timestamps onto a handful of 1970-01-01 values and dropping the duplicates.""" + df = _make_train_df() + automl = AutoML() + _fit(automl, dataframe=df) + assert automl._state.data_size[0] == len(df) + + +def test_explicit_time_col_still_wins_over_the_index(): + df = _make_train_df() + df["ts"] = df.index + automl = AutoML() + _fit(automl, dataframe=df, time_col="ts") + assert automl._state.task.time_col == "ts" + + +def test_genuine_missing_timestamp_still_raises(): + df = pd.DataFrame({"note": ["not-a-date"] * 120, "y": np.arange(120, dtype=float)}) + automl = AutoML() + with pytest.raises(ValueError) as exc: + _fit(automl, dataframe=df) + assert "must contain timestamp values" in str(exc.value) + + +def test_timestamp_column_input_is_unchanged(): + df = _make_train_df().reset_index().rename(columns={"index": "ds"}) + automl = AutoML() + _fit(automl, dataframe=df, time_col="ds") + assert automl._state.task.time_col == "ds" + assert len(automl.predict(df[["ds"]].tail(12))) == 12 + + +def test_index_name_collision_resolves_to_unique_column(): + """When dataframe already has a column named 'ds' (or 'index') that is not datetime, + promoting an unnamed DatetimeIndex must not overwrite or conflict with the existing column.""" + df = _make_train_df() + df["ds"] = np.arange(len(df), dtype=float) # Collision candidate + df["ds_1"] = np.arange(len(df), dtype=float) # Double collision candidate + automl = AutoML() + _fit(automl, dataframe=df) + # Target column should be uniquely chosen as ds_2 avoiding collision + assert automl._state.task.time_col == "ds_2" + assert "ds" in automl._feature_names_in_ + assert "ds_1" in automl._feature_names_in_ + assert "ds_2" in automl._feature_names_in_ + assert automl.data_size_full == len(df) + + +def test_validation_data_with_datetime_index(): + """Verify that validation data passed as DataFrame with DatetimeIndex is properly promoted.""" + df_train = _make_train_df(periods=100) + val_index = pd.date_range("2026-05-01", periods=20, freq="MS") + df_val = pd.DataFrame({"y": np.sin(np.arange(100, 120) / 6) * 10 + 50}, index=val_index) + X_val = pd.DataFrame(index=val_index) + y_val = df_val["y"] + automl = AutoML() + _fit( + automl, + dataframe=df_train, + X_val=X_val, + y_val=y_val, + ) + assert automl._state.task.time_col == "ds" + assert automl._state.eval_method == "holdout" + assert len(automl.predict(X_val)) == len(val_index) + + +def test_validation_data_with_dataframe_target_and_datetime_index(): + """Verify that validation data with DataFrame y_val and DatetimeIndex preserves aligned indexes.""" + df_train = _make_train_df(periods=100) + val_index = pd.date_range("2026-05-01", periods=20, freq="MS") + df_val = pd.DataFrame({"y": np.sin(np.arange(100, 120) / 6) * 10 + 50}, index=val_index) + X_val = pd.DataFrame(index=val_index) + y_val = df_val[["y"]] # DataFrame target with DatetimeIndex + automl = AutoML() + _fit( + automl, + dataframe=df_train, + X_val=X_val, + y_val=y_val, + ) + assert automl._state.task.time_col == "ds" + assert automl._state.eval_method == "holdout" + assert len(automl.predict(X_val)) == len(val_index) + + +def test_xtrain_ytrain_with_datetime_index(): + """Verify that (X_train, y_train) inputs with DatetimeIndex are promoted and work seamlessly.""" + train_idx = pd.date_range("2018-01-01", periods=100, freq="MS") + val_idx = pd.date_range("2026-05-01", periods=20, freq="MS") + X_train = pd.DataFrame({"feat": np.arange(100, dtype=float)}, index=train_idx) + y_train = pd.Series(np.sin(np.arange(100) / 6) * 10 + 50, index=train_idx, name="y") + X_val = pd.DataFrame({"feat": np.arange(100, 120, dtype=float)}, index=val_idx) + y_val = pd.DataFrame({"y": np.sin(np.arange(100, 120) / 6) * 10 + 50}, index=val_idx) + automl = AutoML() + _fit(automl, X_train=X_train, y_train=y_train, X_val=X_val, y_val=y_val) + assert automl._state.task.time_col == "ds" + assert len(automl.predict(X_val)) == 20 + + +def test_predict_with_datetime_index(): + """Verify that future prediction data passed with DatetimeIndex works without explicit time_col.""" + df = _make_train_df(periods=120) + automl = AutoML() + _fit(automl, dataframe=df) + + # Future prediction input using DatetimeIndex + future_index = pd.date_range("2028-01-01", periods=12, freq="MS") + future_df = pd.DataFrame(index=future_index) + preds = automl.predict(future_df) + assert len(preds) == 12 + + +def test_training_with_reordered_dataframe_target_aligns_by_label(): + """Verify that when y_train is passed with shuffled timestamps, values align by label rather than position.""" + train_idx = pd.date_range("2018-01-01", periods=60, freq="MS") + X_train = pd.DataFrame({"feat": np.arange(60, dtype=float)}, index=train_idx) + # Shuffled index for y_train + shuffled_idx = train_idx[::-1] + y_values = np.sin(np.arange(60)[::-1] / 6) * 10 + 50 + y_train = pd.DataFrame({"y": y_values}, index=shuffled_idx) + + automl = AutoML() + _fit(automl, X_train=X_train, y_train=y_train, keep_search_state=True) + + # When y_train is reordered, AutoML must correctly align y by index rather than crashing or distorting data + preds = automl.predict(X_train.tail(12)) + assert len(preds) == 12 + assert automl._state.task.time_col == "ds" + expected_y = y_train.reindex(train_idx)["y"].values + actual_y = automl._state.X_train_all.train_data["y"].values + np.testing.assert_allclose(actual_y, expected_y) + + +def test_validation_with_reordered_dataframe_target_aligns_by_label(): + """Verify that when y_val is passed with shuffled timestamps, values align by label rather than position.""" + df_train = _make_train_df(periods=80) + val_idx = pd.date_range("2024-09-01", periods=20, freq="MS") + X_val = pd.DataFrame(index=val_idx) + # Shuffled index for y_val + shuffled_val_idx = val_idx[::-1] + y_val = pd.DataFrame({"y": (np.arange(20)[::-1] + 100.0)}, index=shuffled_val_idx) + + automl = AutoML() + _fit(automl, dataframe=df_train, X_val=X_val, y_val=y_val, keep_search_state=True) + + preds = automl.predict(X_val) + assert len(preds) == 20 + assert automl._state.eval_method == "holdout" + expected_val_y = y_val.reindex(val_idx)["y"].values + actual_val_y = automl._state.X_val.test_data["y"].values + np.testing.assert_allclose(actual_val_y, expected_val_y) + + +def test_mismatched_target_index_raises_value_error(): + """Verify that passing targets with non-matching labels raises ValueError instead of silent corruption.""" + train_idx = pd.date_range("2018-01-01", periods=60, freq="MS") + mismatched_idx = pd.date_range("2019-01-01", periods=60, freq="MS") + X_train = pd.DataFrame({"feat": np.arange(60, dtype=float)}, index=train_idx) + y_train = pd.DataFrame({"y": np.arange(60, dtype=float)}, index=mismatched_idx) + + automl = AutoML() + with pytest.raises(ValueError, match="Target index labels do not match feature index labels"): + _fit(automl, X_train=X_train, y_train=y_train) + + +def test_training_with_range_index_series_target(): + """Verify positional pairing when X_train has DatetimeIndex and y_train is a RangeIndex Series.""" + train_idx = pd.date_range("2018-01-01", periods=60, freq="MS") + X_train = pd.DataFrame({"feat": np.arange(60, dtype=float)}, index=train_idx) + y_train = pd.Series(np.sin(np.arange(60) / 6) * 10 + 50, name="y") # Default RangeIndex + + automl = AutoML() + _fit(automl, X_train=X_train, y_train=y_train) + + assert automl._state.task.time_col == "ds" + assert automl._state.data_size[0] == len(X_train) + preds = automl.predict(X_train.tail(12)) + assert len(preds) == 12 + + +def test_validation_with_range_index_series_target(): + """Verify positional pairing when validation data has DatetimeIndex X_val and RangeIndex Series y_val.""" + df_train = _make_train_df(periods=80) + val_idx = pd.date_range("2024-09-01", periods=20, freq="MS") + X_val = pd.DataFrame({"feat": np.arange(80, 100, dtype=float)}, index=val_idx) + y_val = pd.Series(np.sin(np.arange(80, 100) / 6) * 10 + 50, name="y") # Default RangeIndex + + automl = AutoML() + _fit(automl, dataframe=df_train, X_val=X_val, y_val=y_val) + + assert automl._state.eval_method == "holdout" + preds = automl.predict(X_val) + assert len(preds) == 20 + + +def test_training_and_validation_with_range_index_series_targets(): + """Verify positional pairing when both training and validation targets are RangeIndex Series.""" + train_idx = pd.date_range("2018-01-01", periods=80, freq="MS") + val_idx = pd.date_range("2024-09-01", periods=20, freq="MS") + X_train = pd.DataFrame({"feat": np.arange(80, dtype=float)}, index=train_idx) + y_train = pd.Series(np.sin(np.arange(80) / 6) * 10 + 50, name="y") + X_val = pd.DataFrame({"feat": np.arange(80, 100, dtype=float)}, index=val_idx) + y_val = pd.Series(np.sin(np.arange(80, 100) / 6) * 10 + 50, name="y") + + automl = AutoML() + _fit(automl, X_train=X_train, y_train=y_train, X_val=X_val, y_val=y_val) + + assert automl._state.task.time_col == "ds" + assert automl._state.eval_method == "holdout" + preds = automl.predict(X_val) + assert len(preds) == 20 + + +def test_training_with_reordered_series_target_aligns_by_label(): + """Verify that when y_train is a Series with shuffled timestamps, values align by label.""" + train_idx = pd.date_range("2018-01-01", periods=60, freq="MS") + X_train = pd.DataFrame({"feat": np.arange(60, dtype=float)}, index=train_idx) + shuffled_idx = train_idx[::-1] + y_values = np.sin(np.arange(60)[::-1] / 6) * 10 + 50 + y_train = pd.Series(y_values, index=shuffled_idx, name="y") + + automl = AutoML() + _fit(automl, X_train=X_train, y_train=y_train, keep_search_state=True) + + preds = automl.predict(X_train.tail(12)) + assert len(preds) == 12 + assert automl._state.task.time_col == "ds" + expected_y = y_train.reindex(train_idx).values + actual_y = automl._state.X_train_all.train_data["y"].values + np.testing.assert_allclose(actual_y, expected_y) + + +def test_training_with_range_index_x_datetime_col_and_datetime_index_series_target(): + """Verify positional pairing when X_train is RangeIndex with a datetime column and y_train is DatetimeIndex Series.""" + dates = pd.date_range("2020-01-01", periods=60, freq="MS") + X_train = pd.DataFrame({"ds": dates, "feat": np.arange(60, dtype=float)}) + y_train = pd.Series(np.sin(np.arange(60) / 6) * 10 + 50, index=dates, name="y") + + automl = AutoML() + _fit(automl, X_train=X_train, y_train=y_train, time_col="ds") + + assert automl._state.task.time_col == "ds" + assert automl._state.data_size[0] == len(X_train) + preds = automl.predict(X_train.tail(12)) + assert len(preds) == 12 + + +def test_validation_with_range_index_x_datetime_col_and_datetime_index_series_target(): + """Verify positional pairing when X_val is RangeIndex with a datetime column and y_val is DatetimeIndex Series.""" + train_dates = pd.date_range("2018-01-01", periods=80, freq="MS") + val_dates = pd.date_range("2024-09-01", periods=20, freq="MS") + X_train = pd.DataFrame({"ds": train_dates, "feat": np.arange(80, dtype=float)}) + y_train = pd.Series(np.sin(np.arange(80) / 6) * 10 + 50, name="y") + X_val = pd.DataFrame({"ds": val_dates, "feat": np.arange(80, 100, dtype=float)}) + y_val = pd.Series(np.sin(np.arange(80, 100) / 6) * 10 + 50, index=val_dates, name="y") + + automl = AutoML() + _fit(automl, X_train=X_train, y_train=y_train, X_val=X_val, y_val=y_val, time_col="ds") + + assert automl._state.eval_method == "holdout" + preds = automl.predict(X_val) + assert len(preds) == 20 + + +def test_index_only_training_input_with_separate_target(): + """Verify that X_train with 0 feature columns and DatetimeIndex promotes correctly without IndexError.""" + dates = pd.date_range("2018-01-01", periods=120, freq="MS") + X_train = pd.DataFrame(index=dates) + y_train = pd.Series(np.sin(np.arange(120) / 6) * 10 + 50, index=dates, name="y") + automl = AutoML() + _fit(automl, X_train=X_train, y_train=y_train) + assert automl._state.task.time_col == "ds" + assert len(automl.predict(X_train.tail(12))) == 12 + + +def test_length_mismatch_raises_value_error(): + """Verify that when X has fewer rows than y, extra rows in y are NOT silently truncated.""" + train_idx = pd.date_range("2018-01-01", periods=60, freq="MS") + extra_idx = pd.date_range("2018-01-01", periods=61, freq="MS") + X_train = pd.DataFrame({"feat": np.arange(60, dtype=float)}, index=train_idx) + y_train = pd.Series(np.arange(61, dtype=float), index=extra_idx, name="y") + automl = AutoML() + with pytest.raises(ValueError, match="does not match feature length"): + _fit(automl, X_train=X_train, y_train=y_train) + + +def test_validation_length_mismatch_raises_value_error(): + """Verify that validation data with mismatched X_val and y_val raises ValueError.""" + df_train = _make_train_df(periods=80) + val_idx = pd.date_range("2024-09-01", periods=20, freq="MS") + val_extra_idx = pd.date_range("2024-09-01", periods=21, freq="MS") + X_val = pd.DataFrame(index=val_idx) + y_val = pd.Series(np.arange(21, dtype=float), index=val_extra_idx, name="y") + automl = AutoML() + with pytest.raises(ValueError, match="does not match feature length"): + _fit(automl, dataframe=df_train, X_val=X_val, y_val=y_val) + + +def test_unnamed_datetime_index_colliding_with_target_named_ds(): + """Verify that when target is named 'ds', an unnamed index promotes to 'ds_1' to avoid duplicate columns.""" + dates = pd.date_range("2018-01-01", periods=120, freq="MS") + X_train = pd.DataFrame({"feat": np.arange(120, dtype=float)}, index=dates) + y_train = pd.Series(np.sin(np.arange(120) / 6) * 10 + 50, index=dates, name="ds") + automl = AutoML() + _fit(automl, X_train=X_train, y_train=y_train, label="ds") + assert automl._state.task.time_col == "ds_1" + assert len(automl.predict(X_train.tail(12))) == 12 + + +def test_named_datetime_index_colliding_with_target_named_y(): + """Verify that when index and target are both named 'y', index promotes to 'y_1' to avoid duplicate columns.""" + dates = pd.date_range("2018-01-01", periods=120, freq="MS") + dates.name = "y" + X_train = pd.DataFrame({"feat": np.arange(120, dtype=float)}, index=dates) + y_train = pd.Series(np.sin(np.arange(120) / 6) * 10 + 50, index=dates, name="y") + automl = AutoML() + _fit(automl, X_train=X_train, y_train=y_train, label="y") + assert automl._state.task.time_col == "y_1" + assert len(automl.predict(X_train.tail(12))) == 12 + + +def test_valid_date_string_column_preserved_over_incidental_datetime_index(): + """Verify that a valid ISO date string column is preserved over an incidental DatetimeIndex during inference and fit.""" + dates_str = pd.date_range("2018-01-01", periods=120, freq="MS").strftime("%Y-%m-%d").tolist() + incidental_idx = pd.date_range("2030-01-01", periods=120, freq="MS") + df = pd.DataFrame( + { + "ds": dates_str, + "y": np.sin(np.arange(120) / 6) * 10 + 50, + }, + index=incidental_idx, + ) + automl = AutoML() + # Inferred time_col (no explicit time_col provided) + _fit(automl, dataframe=df) + assert automl._state.task.time_col == "ds" + # Ensure the timeline corresponds to 2018, not 2030 (120 monthly periods end in 2027) + assert automl.model.end_date.year == 2027 + future_df = pd.DataFrame({"ds": pd.date_range("2028-01-01", periods=12, freq="MS").strftime("%Y-%m-%d")}) + assert len(automl.predict(future_df)) == 12 + + +def test_positional_series_contract_preserved_when_timestamp_column_present(): + """Verify that when X has an existing timestamp column, Series target is paired positionally even if X has a DatetimeIndex.""" + from flaml.automl.task.time_series_task import TimeSeriesTask + + # X has explicit timestamp column 'ds' in order, but an incidental reverse DatetimeIndex + dates = pd.date_range("2020-01-01", periods=5, freq="D") + X = pd.DataFrame( + {"ds": dates, "feat": [1.0, 2.0, 3.0, 4.0, 5.0]}, + index=dates[::-1], + ) + y = pd.Series([10.0, 20.0, 30.0, 40.0, 50.0], index=dates[::-1]) + + needs_promotion = TimeSeriesTask._needs_datetime_index_promotion(X, time_col="ds", is_inferred=False) + assert not needs_promotion + + X_aligned, y_aligned = TimeSeriesTask._align_y_to_X_and_promote(X, y, needs_promotion=needs_promotion) + # Target values must stay in original order (positional pairing), NOT re-indexed + assert list(y_aligned.values) == [10.0, 20.0, 30.0, 40.0, 50.0] + + +def test_label_mismatch_with_existing_nan_raises_value_error(): + """Verify that index label mismatch is rejected even when dropped row contains NaN (which defeated NaN-count check).""" + from flaml.automl.task.time_series_task import TimeSeriesTask + + # X has Jan 1 and Jan 2; y has Jan 2 and Jan 3 (with value NaN) + x_idx = pd.date_range("2020-01-01", periods=2, freq="D") + y_idx = pd.date_range("2020-01-02", periods=2, freq="D") + + X = pd.DataFrame({"feat": [1.0, 2.0]}, index=x_idx) + y = pd.Series([10.0, np.nan], index=y_idx) + + # When promotion is needed, mismatched index labels must raise ValueError + with pytest.raises(ValueError, match="Target index labels do not match feature index labels"): + TimeSeriesTask._align_y_to_X_and_promote(X, y, needs_promotion=True) + + +def test_predict_restores_datetime_index_preventing_assignment_nans(): + """Verify that predictions on future_df with DatetimeIndex retain future_df.index, preventing assignment NaNs.""" + df = _make_train_df(periods=120) + automl = AutoML() + _fit(automl, dataframe=df) + + future_index = pd.date_range("2028-01-01", periods=12, freq="MS") + future_df = pd.DataFrame(index=future_index) + + preds = automl.predict(future_df) + assert isinstance(preds, (pd.Series, pd.DataFrame, np.ndarray)) + # When assigned back to future_df, pandas aligns by index. If index was RangeIndex, this produced all NaNs. + future_df["prediction"] = preds + assert not future_df["prediction"].isna().any() + assert len(future_df["prediction"]) == 12 + + +def test_predict_restores_datetime_index_for_estimator_that_resets_output_index(): + """Verify that when an estimator resets its prediction index (like Orbit or Prophet), + AutoML.predict() restores the caller's input DatetimeIndex, preventing assignment NaNs.""" + df = _make_train_df(periods=120) + automl = AutoML() + _fit(automl, dataframe=df) + + future_index = pd.date_range("2028-01-01", periods=12, freq="MS") + future_df = pd.DataFrame(index=future_index) + + # Wrap estimator to reset index to RangeIndex (simulating Orbit or Prophet) + orig_predict = automl._trained_estimator.predict + + def reset_predict(X, **kwargs): + preds = orig_predict(X, **kwargs) + return preds.reset_index(drop=True) if hasattr(preds, "reset_index") else preds + + automl._trained_estimator.predict = reset_predict + + preds = automl.predict(future_df) + assert isinstance(preds, (pd.Series, pd.DataFrame)) + assert preds.index.equals(future_index) + + future_df["prediction"] = preds + assert not future_df["prediction"].isna().any() + assert len(future_df["prediction"]) == 12 + + +def test_is_valid_timestamp_col_boundary_cases(): + """Directly test _is_valid_timestamp_col with various dtypes.""" + from flaml.automl.task.time_series_task import TimeSeriesTask + + df = pd.DataFrame( + { + "num_int": [1, 2, 3], + "num_float": [1.0, 2.0, 3.0], + "num_str": ["100", "200", "300"], + "invalid_str": ["not-a-date", "foo", "bar"], + "iso_str": ["2020-01-01", "2020-01-02", "2020-01-03"], + "dt64": pd.date_range("2020-01-01", periods=3), + "bool_col": [True, False, True], + } + ) + + assert not TimeSeriesTask._is_valid_timestamp_col(df, "num_int") + assert not TimeSeriesTask._is_valid_timestamp_col(df, "num_float") + assert not TimeSeriesTask._is_valid_timestamp_col(df, "num_str") + assert not TimeSeriesTask._is_valid_timestamp_col(df, "invalid_str") + assert not TimeSeriesTask._is_valid_timestamp_col(df, "bool_col") + assert not TimeSeriesTask._is_valid_timestamp_col(df, "non_existent_col") + assert TimeSeriesTask._is_valid_timestamp_col(df, "iso_str") + assert TimeSeriesTask._is_valid_timestamp_col(df, "dt64") + + +def test_reordered_dataframe_target_with_existing_time_col_aligns_by_value(): + """Verify that when X has an existing timestamp column, DataFrame target with reversed index + aligns by index label to preserve target-feature mapping (reproducing maintainer thinkall's case).""" + n = 40 + X = pd.DataFrame( + { + "ds": pd.date_range("2020-01-01", periods=n, freq="D"), + "feat": np.arange(n), + } + ) + y = pd.DataFrame({"y": np.arange(n, dtype=float)[::-1]}, index=np.arange(n)[::-1]) + + automl = AutoML() + _fit( + automl, + X_train=X, + y_train=y, + period=4, + time_budget=3, + keep_search_state=True, + ) + head_df = automl._state.X_train_all.train_data.head() + assert list(head_df["y"].values) == [0.0, 1.0, 2.0, 3.0, 4.0] + + +def test_reordered_dataframe_target_with_datetime_index_promotion_aligns_by_value(): + """Verify that when DatetimeIndex is promoted to time_col, DataFrame target with reversed + DatetimeIndex aligns by timestamp to preserve target-feature mapping.""" + dates = pd.date_range("2020-01-01", periods=40, freq="D") + X = pd.DataFrame({"feat": np.arange(40)}, index=dates) + y = pd.DataFrame({"y": np.arange(40, dtype=float)[::-1]}, index=dates[::-1]) + + automl = AutoML() + _fit( + automl, + X_train=X, + y_train=y, + period=4, + time_budget=3, + keep_search_state=True, + ) + head_df = automl._state.X_train_all.train_data.head() + assert list(head_df["y"].values) == [0.0, 1.0, 2.0, 3.0, 4.0] + + +def test_explicit_missing_time_col_raises_error_even_with_datetime_index(): + """Verify that explicitly specifying a non-existent time_col raises ValueError rather than + silently promoting DatetimeIndex during training (Copilot review discussion_r4308583028).""" + dates = pd.date_range("2020-01-01", periods=40, freq="D") + X = pd.DataFrame({"feat": np.arange(40)}, index=dates) + y = pd.Series(np.arange(40, dtype=float), index=dates, name="y") + + automl = AutoML() + with pytest.raises(ValueError, match="must contain timestamp values"): + _fit(automl, X_train=X, y_train=y, time_col="non_existent_time_col") + + # Also test dataframe input with explicit missing time_col + df = pd.DataFrame({"feat": np.arange(40), "y": np.arange(40, dtype=float)}, index=dates) + automl_df = AutoML() + with pytest.raises(ValueError, match="must contain timestamp values"): + _fit(automl_df, dataframe=df, time_col="non_existent_time_col") + + +def test_mismatched_dataframe_target_labels_raise_value_error(): + """Verify that DataFrame targets with mismatched labels raise ValueError instead of silent corruption.""" + dates = pd.date_range("2020-01-01", periods=40, freq="D") + wrong_dates = pd.date_range("2021-01-01", periods=40, freq="D") + X = pd.DataFrame({"ds": dates, "feat": np.arange(40)}) + y = pd.DataFrame({"y": np.arange(40, dtype=float)}, index=wrong_dates) + + automl = AutoML() + with pytest.raises(ValueError, match="Target index labels do not match feature index labels"): + _fit(automl, X_train=X, y_train=y, time_col="ds") + + +def test_explicit_time_col_missing_from_features_rejected_even_if_present_in_target_dataframe(): + """Verify that an explicitly specified time_col must be present in features X, and cannot be + satisfied by a numeric column in DataFrame target y (reproducing maintainer thinkall's case).""" + dates = pd.date_range("2020-01-01", periods=40, freq="D") + X = pd.DataFrame({"feat": np.arange(40)}, index=dates) + y = pd.DataFrame( + {"y": 3 + 7 * np.arange(40, dtype=float), "ds": np.arange(40)}, + index=dates, + ) + + automl = AutoML() + with pytest.raises(ValueError, match="must contain timestamp values"): + _fit( + automl, + X_train=X, + y_train=y, + label="y", + time_col="ds", + period=4, + keep_search_state=True, + ) + + +def test_numpy_array_training_input_with_explicit_time_col(): + """Verify that a NumPy array X_train with timestamps in column 0 is accepted when time_col is explicit.""" + dates = pd.date_range("2020-01-01", periods=40, freq="D") + X_arr = np.column_stack([dates.astype(object), np.arange(40, dtype=float)]) + y_arr = np.arange(40, dtype=float) + + automl = AutoML() + _fit( + automl, + X_train=X_arr, + y_train=y_arr, + time_col="ds", + period=4, + ) + assert automl._state.task.time_col == "ds" + assert len(automl.predict(X_arr[-12:])) == 12 + + +def test_numpy_array_validation_input_accepted(): + """Verify that a NumPy array X_val with timestamps in column 0 is accepted for holdout validation.""" + dates = pd.date_range("2020-01-01", periods=40, freq="D") + df_train = pd.DataFrame({"ds": dates[:30], "y": np.arange(30, dtype=float)}) + X_val_arr = dates[30:].values.reshape(-1, 1) + y_val_arr = np.arange(30, 40, dtype=float) + + automl = AutoML() + _fit( + automl, + dataframe=df_train, + label="y", + X_val=X_val_arr, + y_val=y_val_arr, + period=4, + ) + assert automl._state.eval_method == "holdout" + assert len(automl.predict(X_val_arr)) == 10 + + +def test_numpy_date_string_array_can_train_and_predict(): + """Verify ISO date strings are converted for NumPy array training and prediction.""" + train_dates = pd.date_range("2018-01-01", periods=120, freq="MS") + X_train = train_dates.strftime("%Y-%m-%d").to_numpy(dtype=object) + y_train = 10 + np.sin(np.arange(len(train_dates)) * 2 * np.pi / 12) + + automl = AutoML() + _fit(automl, X_train=X_train, y_train=y_train, time_col="ds") + + future_dates = pd.date_range("2028-01-01", periods=12, freq="MS") + X_future = future_dates.strftime("%Y-%m-%d").to_numpy(dtype=object) + predictions = automl.predict(X_future) + + assert len(predictions) == len(future_dates) + assert np.isfinite(np.asarray(predictions)).all() + + +@pytest.mark.parametrize("name", ["", 0, "when"]) +def test_explicit_index_name_prediction_form_equivalence(name): + """Verify that explicit index names (including empty string '' and integer 0) are preserved + as the time column, and raw-index future frames produce identical predictions to reset_index() frames.""" + df = _make_train_df(periods=120) + df.index.name = name + automl = AutoML() + _fit(automl, dataframe=df) + assert automl._state.task.time_col == name + + future_index = pd.date_range("2028-01-01", periods=12, freq="MS", name=name) + future_raw = pd.DataFrame(index=future_index) + future_reset = future_raw.reset_index() + + pred_raw = automl.predict(future_raw) + pred_reset = automl.predict(future_reset) + + assert len(pred_raw) == 12 + assert len(pred_reset) == 12 + np.testing.assert_allclose(pred_raw.values, pred_reset.values) + + +def _make_predict_only_automl(predictions, is_panel=False): + automl = AutoML() + automl._state = SimpleNamespace( + task=SimpleNamespace( + preprocess=lambda X, transformer: X, + is_ts_forecastpanel=lambda: is_panel, + ) + ) + automl._transformer = None + automl._label_transformer = None + automl._trained_estimator = SimpleNamespace(predict=lambda X, **kwargs: predictions.copy()) + return automl + + +def test_predict_handles_list_input_when_estimator_returns_series(): + expected = pd.Series([10.0, 20.0], index=[4, 5], name="y") + automl = _make_predict_only_automl(expected) + + actual = automl.predict([[1.0], [2.0]]) + + pd.testing.assert_series_equal(actual, expected) + + +def test_predict_preserves_panel_group_time_index(): + panel_index = pd.MultiIndex.from_tuples( + [("store_b", pd.Timestamp("2024-01-01")), ("store_a", pd.Timestamp("2024-01-01"))], + names=["store", "ds"], + ) + expected = pd.Series([2.0, 1.0], index=panel_index, name="y") + automl = _make_predict_only_automl(expected, is_panel=True) + + actual = automl.predict(pd.DataFrame({"store": ["store_a", "store_b"]})) + + pd.testing.assert_series_equal(actual, expected) + + +def test_prophet_predict_restores_unsorted_input_timestamp_order(monkeypatch): + from flaml.automl.time_series.ts_model import Prophet + + dates = pd.DatetimeIndex(["2028-03-01", "2028-01-01", "2028-02-01"], name="ds") + X = pd.DataFrame({"ds": dates}, index=dates) + + class SortingProphetModel: + def predict(self, frame, **kwargs): + sorted_dates = frame["ds"].sort_values() + return pd.DataFrame( + { + "ds": sorted_dates.to_numpy(), + "yhat": sorted_dates.dt.month.to_numpy(dtype=float), + } + ) + + estimator = Prophet() + estimator.time_col = "ds" + estimator.target_names = ["y"] + estimator._model = SortingProphetModel() + monkeypatch.setattr(estimator, "enrich", lambda frame: frame.copy()) + + actual = estimator.predict(X) + expected = pd.Series([3.0, 1.0, 2.0], index=dates, name="y") + + pd.testing.assert_series_equal(actual, expected)