From 9660c65b22b8384178441ddc114958bb6850e8a2 Mon Sep 17 00:00:00 2001 From: Aditi Singh Date: Thu, 22 Jan 2026 02:00:42 +0530 Subject: [PATCH 1/4] feat: auto-create timestamps in prettify_prediction when test_data is None - Removed NotImplementedError and instead generate timestamps automatically - Uses training data's end_date and frequency to create prediction timestamps - Supports np.ndarray, pd.Series, and pd.DataFrame inputs --- flaml/automl/time_series/ts_data.py | 14 +++++++++++--- 1 file changed, 11 insertions(+), 3 deletions(-) diff --git a/flaml/automl/time_series/ts_data.py b/flaml/automl/time_series/ts_data.py index 625049c601..d9359f2d44 100644 --- a/flaml/automl/time_series/ts_data.py +++ b/flaml/automl/time_series/ts_data.py @@ -243,13 +243,21 @@ def prettify_prediction(self, y_pred: Union[pd.DataFrame, pd.Series, np.ndarray] y_pred[self.time_col] = self.test_data[self.time_col] else: + # Auto-create timestamps when test_data is None if isinstance(y_pred, np.ndarray): - raise ValueError("Can't enrich np.ndarray as self.test_data is None") + y_pred = pd.DataFrame(data=y_pred, columns=self.target_names) elif isinstance(y_pred, pd.Series): assert len(self.target_names) == 1, "Not enough columns in y_pred" y_pred = pd.DataFrame({self.target_names[0]: y_pred}) - # TODO auto-create the timestamps for the time column instead of throwing - raise NotImplementedError("Need a non-None test_data for this to work, for now") + + # Generate timestamps based on training data's end_date and frequency + train_end_date = self.train_data[self.time_col].max() + pred_timestamps = pd.date_range( + start=train_end_date + pd.Timedelta(1, self.frequency), + periods=len(y_pred), + freq=self.frequency, + ) + y_pred[self.time_col] = pred_timestamps assert isinstance(y_pred, pd.DataFrame) assert self.time_col in y_pred.columns From b2f118f0255821fa95ed7e04d50c76e395bfae51 Mon Sep 17 00:00:00 2001 From: Aditi Singh Date: Thu, 22 Jan 2026 13:17:48 +0530 Subject: [PATCH 2/4] fix: timestamp generation in prettify_prediction for empty test_data --- flaml/automl/time_series/ts_data.py | 6 +- test/automl/test_model.py | 107 ++++++++++++++++++++++++++++ 2 files changed, 110 insertions(+), 3 deletions(-) diff --git a/flaml/automl/time_series/ts_data.py b/flaml/automl/time_series/ts_data.py index d9359f2d44..19028b248b 100644 --- a/flaml/automl/time_series/ts_data.py +++ b/flaml/automl/time_series/ts_data.py @@ -253,10 +253,10 @@ def prettify_prediction(self, y_pred: Union[pd.DataFrame, pd.Series, np.ndarray] # Generate timestamps based on training data's end_date and frequency train_end_date = self.train_data[self.time_col].max() pred_timestamps = pd.date_range( - start=train_end_date + pd.Timedelta(1, self.frequency), - periods=len(y_pred), + start=train_end_date, + periods=len(y_pred) + 1, freq=self.frequency, - ) + )[1:] # Skip the first timestamp (train_end_date itself) y_pred[self.time_col] = pred_timestamps assert isinstance(y_pred, pd.DataFrame) diff --git a/test/automl/test_model.py b/test/automl/test_model.py index 10ee6112fc..1a443bb33e 100644 --- a/test/automl/test_model.py +++ b/test/automl/test_model.py @@ -142,6 +142,113 @@ def test_prep(): print(lgbm.feature_importances_) +def test_prettify_prediction_auto_timestamps(): + """Test that prettify_prediction auto-generates timestamps when test_data is None. + + This tests the fix for the TODO that previously raised NotImplementedError. + """ + import pandas as pd + + # Create training data with daily frequency + n = 30 + train_df = DataFrame( + { + "date": pd.date_range(start="2023-01-01", periods=n, freq="D"), + "value": np.random.randn(n), + } + ) + + # Create TimeSeriesDataset without test_data + tsds = TimeSeriesDataset(train_df, time_col="date", target_names="value") + + # Verify test_data is empty (not None, but empty DataFrame per __init__) + assert len(tsds.test_data) == 0 + + # Test with np.ndarray input + pred_steps = 5 + y_pred_array = np.random.randn(pred_steps) + result_array = tsds.prettify_prediction(y_pred_array) + + assert isinstance(result_array, pd.DataFrame) + assert "date" in result_array.columns + assert "value" in result_array.columns + assert len(result_array) == pred_steps + # Verify timestamps start from training end + 1 period + expected_start = pd.date_range(start=train_df["date"].max(), periods=2, freq="D")[1] + assert result_array["date"].iloc[0] == expected_start + + # Test with pd.Series input + y_pred_series = pd.Series(np.random.randn(pred_steps)) + result_series = tsds.prettify_prediction(y_pred_series) + + assert isinstance(result_series, pd.DataFrame) + assert "date" in result_series.columns + assert "value" in result_series.columns + assert len(result_series) == pred_steps + assert result_series["date"].iloc[0] == expected_start + + # Test with pd.DataFrame input + y_pred_df = pd.DataFrame({"value": np.random.randn(pred_steps)}) + result_df = tsds.prettify_prediction(y_pred_df) + + assert isinstance(result_df, pd.DataFrame) + assert "date" in result_df.columns + assert "value" in result_df.columns + assert len(result_df) == pred_steps + assert result_df["date"].iloc[0] == expected_start + + # Verify the generated timestamps follow the correct frequency + expected_dates = pd.date_range(start=expected_start, periods=pred_steps, freq="D") + pd.testing.assert_index_equal( + pd.DatetimeIndex(result_array["date"]), + expected_dates, + check_names=False, + ) + + print("test_prettify_prediction_auto_timestamps passed!") + + +def test_prettify_prediction_auto_timestamps_monthly(): + """Test auto-timestamp generation with monthly frequency.""" + import pandas as pd + + # Create training data with monthly frequency + n = 24 + train_df = DataFrame( + { + "date": pd.date_range(start="2022-01-01", periods=n, freq="MS"), + "value": np.random.randn(n), + } + ) + + tsds = TimeSeriesDataset(train_df, time_col="date", target_names="value") + assert len(tsds.test_data) == 0 + + pred_steps = 6 + y_pred = np.random.randn(pred_steps) + result = tsds.prettify_prediction(y_pred) + + assert isinstance(result, pd.DataFrame) + assert len(result) == pred_steps + + # For monthly frequency, verify timestamps are monthly + expected_dates = pd.date_range( + start=train_df["date"].max(), + periods=pred_steps + 1, + freq="MS", + )[1:] # Skip first (which is train_end_date) + + pd.testing.assert_index_equal( + pd.DatetimeIndex(result["date"]), + expected_dates, + check_names=False, + ) + + print("test_prettify_prediction_auto_timestamps_monthly passed!") + + if __name__ == "__main__": test_lrl2() test_prep() + test_prettify_prediction_auto_timestamps() + test_prettify_prediction_auto_timestamps_monthly() From 14e33e3ca06f511617f697b58b3353ac03038968 Mon Sep 17 00:00:00 2001 From: Aditi Singh Date: Fri, 23 Jan 2026 01:03:32 +0530 Subject: [PATCH 3/4] test: simplify and relocate auto-timestamp tests per review feedback --- test/automl/test_forecast.py | 114 +++++++++++++++++++++++++++++++++++ test/automl/test_model.py | 107 -------------------------------- 2 files changed, 114 insertions(+), 107 deletions(-) diff --git a/test/automl/test_forecast.py b/test/automl/test_forecast.py index c77431e167..bfc2145d68 100644 --- a/test/automl/test_forecast.py +++ b/test/automl/test_forecast.py @@ -724,6 +724,120 @@ def test_log_training_metric_ts_models(): assert automl.best_estimator == estimator +def test_prettify_prediction_auto_timestamps_data_types(): + """Test auto-timestamp generation with different input data types (orthogonal test). + + This tests the fix for the TODO that previously raised NotImplementedError. + Tests np.ndarray, pd.Series, and pd.DataFrame inputs with daily frequency. + """ + from flaml.automl.time_series import TimeSeriesDataset + + # Create training data with daily frequency + n = 30 + train_df = pd.DataFrame( + { + "date": pd.date_range(start="2023-01-01", periods=n, freq="D"), + "value": np.random.randn(n), + } + ) + tsds = TimeSeriesDataset(train_df, time_col="date", target_names="value") + assert len(tsds.test_data) == 0 + + pred_steps = 5 + expected_start = pd.date_range(start=train_df["date"].max(), periods=2, freq="D")[1] + + # Test np.ndarray + result = tsds.prettify_prediction(np.random.randn(pred_steps)) + assert isinstance(result, pd.DataFrame) + assert len(result) == pred_steps + assert result["date"].iloc[0] == expected_start + + # Test pd.Series + result = tsds.prettify_prediction(pd.Series(np.random.randn(pred_steps))) + assert isinstance(result, pd.DataFrame) + assert len(result) == pred_steps + assert result["date"].iloc[0] == expected_start + + # Test pd.DataFrame + result = tsds.prettify_prediction(pd.DataFrame({"value": np.random.randn(pred_steps)})) + assert isinstance(result, pd.DataFrame) + assert len(result) == pred_steps + assert result["date"].iloc[0] == expected_start + + +def test_prettify_prediction_auto_timestamps_frequencies(): + """Test auto-timestamp generation with different frequencies (orthogonal test). + + Tests daily and monthly frequencies with np.ndarray input. + """ + from flaml.automl.time_series import TimeSeriesDataset + + pred_steps = 6 + + # Test daily frequency + train_df_daily = pd.DataFrame( + { + "date": pd.date_range(start="2023-01-01", periods=30, freq="D"), + "value": np.random.randn(30), + } + ) + tsds_daily = TimeSeriesDataset(train_df_daily, time_col="date", target_names="value") + result = tsds_daily.prettify_prediction(np.random.randn(pred_steps)) + expected_dates = pd.date_range(start=train_df_daily["date"].max(), periods=pred_steps + 1, freq="D")[1:] + pd.testing.assert_index_equal(pd.DatetimeIndex(result["date"]), expected_dates, check_names=False) + + # Test monthly frequency + train_df_monthly = pd.DataFrame( + { + "date": pd.date_range(start="2022-01-01", periods=24, freq="MS"), + "value": np.random.randn(24), + } + ) + tsds_monthly = TimeSeriesDataset(train_df_monthly, time_col="date", target_names="value") + result = tsds_monthly.prettify_prediction(np.random.randn(pred_steps)) + expected_dates = pd.date_range(start=train_df_monthly["date"].max(), periods=pred_steps + 1, freq="MS")[1:] + pd.testing.assert_index_equal(pd.DatetimeIndex(result["date"]), expected_dates, check_names=False) + + +def test_auto_timestamps_e2e(budget=3): + """E2E test: train a model and predict without explicit test_data timestamps. + + This showcases the improvement from this PR - users can now make predictions + without providing explicit test data timestamps. + """ + try: + import statsmodels # noqa: F401 + except ImportError: + print("statsmodels not installed, skipping E2E test") + return + + # Create training data + n = 100 + train_df = pd.DataFrame( + { + "ds": pd.date_range(start="2020-01-01", periods=n, freq="D"), + "y": np.sin(np.linspace(0, 10, n)) + np.random.randn(n) * 0.1, + } + ) + + # Train model + automl = AutoML() + automl.fit( + dataframe=train_df, + label="y", + period=10, + task="ts_forecast", + time_budget=budget, + estimator_list=["arima"], + ) + + # Predict using steps (no explicit test_data) - this is the key improvement + y_pred = automl.predict(10) + assert y_pred is not None + assert len(y_pred) == 10 + print("E2E test passed: model trained and predicted without explicit test_data!") + + if __name__ == "__main__": # test_forecast_automl(60) # test_multivariate_forecast_num(5) diff --git a/test/automl/test_model.py b/test/automl/test_model.py index 1a443bb33e..10ee6112fc 100644 --- a/test/automl/test_model.py +++ b/test/automl/test_model.py @@ -142,113 +142,6 @@ def test_prep(): print(lgbm.feature_importances_) -def test_prettify_prediction_auto_timestamps(): - """Test that prettify_prediction auto-generates timestamps when test_data is None. - - This tests the fix for the TODO that previously raised NotImplementedError. - """ - import pandas as pd - - # Create training data with daily frequency - n = 30 - train_df = DataFrame( - { - "date": pd.date_range(start="2023-01-01", periods=n, freq="D"), - "value": np.random.randn(n), - } - ) - - # Create TimeSeriesDataset without test_data - tsds = TimeSeriesDataset(train_df, time_col="date", target_names="value") - - # Verify test_data is empty (not None, but empty DataFrame per __init__) - assert len(tsds.test_data) == 0 - - # Test with np.ndarray input - pred_steps = 5 - y_pred_array = np.random.randn(pred_steps) - result_array = tsds.prettify_prediction(y_pred_array) - - assert isinstance(result_array, pd.DataFrame) - assert "date" in result_array.columns - assert "value" in result_array.columns - assert len(result_array) == pred_steps - # Verify timestamps start from training end + 1 period - expected_start = pd.date_range(start=train_df["date"].max(), periods=2, freq="D")[1] - assert result_array["date"].iloc[0] == expected_start - - # Test with pd.Series input - y_pred_series = pd.Series(np.random.randn(pred_steps)) - result_series = tsds.prettify_prediction(y_pred_series) - - assert isinstance(result_series, pd.DataFrame) - assert "date" in result_series.columns - assert "value" in result_series.columns - assert len(result_series) == pred_steps - assert result_series["date"].iloc[0] == expected_start - - # Test with pd.DataFrame input - y_pred_df = pd.DataFrame({"value": np.random.randn(pred_steps)}) - result_df = tsds.prettify_prediction(y_pred_df) - - assert isinstance(result_df, pd.DataFrame) - assert "date" in result_df.columns - assert "value" in result_df.columns - assert len(result_df) == pred_steps - assert result_df["date"].iloc[0] == expected_start - - # Verify the generated timestamps follow the correct frequency - expected_dates = pd.date_range(start=expected_start, periods=pred_steps, freq="D") - pd.testing.assert_index_equal( - pd.DatetimeIndex(result_array["date"]), - expected_dates, - check_names=False, - ) - - print("test_prettify_prediction_auto_timestamps passed!") - - -def test_prettify_prediction_auto_timestamps_monthly(): - """Test auto-timestamp generation with monthly frequency.""" - import pandas as pd - - # Create training data with monthly frequency - n = 24 - train_df = DataFrame( - { - "date": pd.date_range(start="2022-01-01", periods=n, freq="MS"), - "value": np.random.randn(n), - } - ) - - tsds = TimeSeriesDataset(train_df, time_col="date", target_names="value") - assert len(tsds.test_data) == 0 - - pred_steps = 6 - y_pred = np.random.randn(pred_steps) - result = tsds.prettify_prediction(y_pred) - - assert isinstance(result, pd.DataFrame) - assert len(result) == pred_steps - - # For monthly frequency, verify timestamps are monthly - expected_dates = pd.date_range( - start=train_df["date"].max(), - periods=pred_steps + 1, - freq="MS", - )[1:] # Skip first (which is train_end_date) - - pd.testing.assert_index_equal( - pd.DatetimeIndex(result["date"]), - expected_dates, - check_names=False, - ) - - print("test_prettify_prediction_auto_timestamps_monthly passed!") - - if __name__ == "__main__": test_lrl2() test_prep() - test_prettify_prediction_auto_timestamps() - test_prettify_prediction_auto_timestamps_monthly() From b481f8a408b768ab46b593f592c34a4df3afce9c Mon Sep 17 00:00:00 2001 From: Aditi Singh Date: Sat, 24 Jan 2026 23:32:50 +0530 Subject: [PATCH 4/4] style: fix formatting and respond to PR feedback --- auto_create_timestamps_pr.md | 17 +++++++++++++++++ flaml/automl/time_series/ts_data.py | 4 +++- test/automl/test_forecast.py | 12 ++++++++---- 3 files changed, 28 insertions(+), 5 deletions(-) create mode 100644 auto_create_timestamps_pr.md diff --git a/auto_create_timestamps_pr.md b/auto_create_timestamps_pr.md new file mode 100644 index 0000000000..7dc121ac08 --- /dev/null +++ b/auto_create_timestamps_pr.md @@ -0,0 +1,17 @@ +# Auto-Create Timestamps in `prettify_prediction()` When `test_data` is None + +## Why are these changes needed? + +Currently, the `TimeSeriesDataset.prettify_prediction()` method in `flaml/automl/time_series/ts_data.py` throws a `NotImplementedError` when `test_data` is `None`. +This is frustrating for users who want to make predictions without providing explicit test data timestamps. + +**This PR implements automatic timestamp generation** by: + +1. Using the training data's end date as the starting point. +2. Generating future timestamps based on the inferred frequency. +3. Supporting `np.ndarray`, `pd.Series`, and `pd.DataFrame`. + +## Checks + +- [x] Pre-commit linting (black, ruff). +- [x] Added regression tests demonstrating the fix. diff --git a/flaml/automl/time_series/ts_data.py b/flaml/automl/time_series/ts_data.py index 19028b248b..e6d872bedb 100644 --- a/flaml/automl/time_series/ts_data.py +++ b/flaml/automl/time_series/ts_data.py @@ -256,7 +256,9 @@ def prettify_prediction(self, y_pred: Union[pd.DataFrame, pd.Series, np.ndarray] start=train_end_date, periods=len(y_pred) + 1, freq=self.frequency, - )[1:] # Skip the first timestamp (train_end_date itself) + )[ + 1: + ] # Skip the first timestamp (train_end_date itself) y_pred[self.time_col] = pred_timestamps assert isinstance(y_pred, pd.DataFrame) diff --git a/test/automl/test_forecast.py b/test/automl/test_forecast.py index bfc2145d68..1f3137f7c2 100644 --- a/test/automl/test_forecast.py +++ b/test/automl/test_forecast.py @@ -725,10 +725,13 @@ def test_log_training_metric_ts_models(): def test_prettify_prediction_auto_timestamps_data_types(): - """Test auto-timestamp generation with different input data types (orthogonal test). + """Test auto-timestamp generation with different input data types. - This tests the fix for the TODO that previously raised NotImplementedError. - Tests np.ndarray, pd.Series, and pd.DataFrame inputs with daily frequency. + Before this PR fix, calling prettify_prediction() with test_data=None raised: + - ValueError for np.ndarray: "Can't enrich np.ndarray as self.test_data is None" + - NotImplementedError for pd.Series/DataFrame: "Need a non-None test_data for this to work" + + This test verifies the fix works for np.ndarray, pd.Series, and pd.DataFrame inputs. """ from flaml.automl.time_series import TimeSeriesDataset @@ -766,8 +769,9 @@ def test_prettify_prediction_auto_timestamps_data_types(): def test_prettify_prediction_auto_timestamps_frequencies(): - """Test auto-timestamp generation with different frequencies (orthogonal test). + """Test auto-timestamp generation with different frequencies. + Before this PR fix, this would raise NotImplementedError when test_data is None. Tests daily and monthly frequencies with np.ndarray input. """ from flaml.automl.time_series import TimeSeriesDataset