diff --git a/python/prophet/diagnostics.py b/python/prophet/diagnostics.py index a754691..6a45db6 100644 --- a/python/prophet/diagnostics.py +++ b/python/prophet/diagnostics.py @@ -58,7 +58,7 @@ def generate_cutoffs(df, horizon, initial, period): return list(reversed(result)) -def cross_validation(model, horizon, period=None, initial=None, parallel=None, cutoffs=None, disable_tqdm=False): +def cross_validation(model, horizon, period=None, initial=None, parallel=None, cutoffs=None, disable_tqdm=False, extra_output_columns=None): """Cross-Validation for time series. Computes forecasts from historical cutoff points, which user can input. @@ -82,8 +82,6 @@ def cross_validation(model, horizon, period=None, initial=None, parallel=None, c cross validation. If not provided, they are generated as described above. parallel : {None, 'processes', 'threads', 'dask', object} - disable_tqdm: if True it disables the progress bar that would otherwise show up when parallel=None - How to parallelize the forecast computation. By default no parallelism is used. @@ -109,6 +107,10 @@ def cross_validation(model, horizon, period=None, initial=None, parallel=None, c for args in zip(*iterables) ] return results + + disable_tqdm: if True it disables the progress bar that would otherwise show up when parallel=None + extra_output_columns: A String or List of Strings e.g. 'trend' or ['trend']. + Additional columns to 'yhat' and 'ds' to be returned in output. Returns ------- @@ -120,10 +122,15 @@ def cross_validation(model, horizon, period=None, initial=None, parallel=None, c df = model.history.copy().reset_index(drop=True) horizon = pd.Timedelta(horizon) - predict_columns = ['ds', 'yhat'] + if model.uncertainty_samples: predict_columns.extend(['yhat_lower', 'yhat_upper']) + + if extra_output_columns is not None: + if isinstance(extra_output_columns, str): + extra_output_columns = [extra_output_columns] + predict_columns.extend([c for c in extra_output_columns if c not in predict_columns]) # Identify largest seasonality period period_max = 0. diff --git a/python/prophet/tests/test_diagnostics.py b/python/prophet/tests/test_diagnostics.py index 6dcfc3e..fda4896 100644 --- a/python/prophet/tests/test_diagnostics.py +++ b/python/prophet/tests/test_diagnostics.py @@ -101,6 +101,19 @@ class TestCrossValidation: _ = diagnostics.cross_validation(m, *args) # check single forecast function called expected number of times assert n_calls == forecasts + + @pytest.mark.parametrize("extra_output_columns", ["trend", ["trend"]]) + def test_check_extra_output_columns_cross_validation(self, ts_short, backend, extra_output_columns): + m = Prophet(stan_backend=backend) + m.fit(ts_short) + df_cv = diagnostics.cross_validation( + m, + horizon="1 days", + period="1 days", + initial="140 days", + extra_output_columns=extra_output_columns + ) + assert "trend" in df_cv.columns @pytest.mark.parametrize("growth", ["logistic", "flat"]) def test_cross_validation_logistic_or_flat_growth(self, growth, ts_short, backend):