orbit
orbit copied to clipboard
KTR predict method doesn't consider single knot case.
Describe the bug
In the case when user want to fit a flat level (and build customized trend as regressors) by specifying level_segments=0
, the KTR model is able to fit a single knot for the level component. However, the predict function produces IndexError because the underlying code assumes KTRModel.knots_tp_levels
to have at least 2 components.
To Reproduce
import orbit
from orbit.models import KTR
from orbit.diagnostics.plot import plot_predicted_data, plot_predicted_components
from orbit.utils.dataset import load_electricity_demand
df = load_electricity_demand()
date_col = 'date'
response_col = 'electricity'
df[response_col] = np.log(df[response_col])
test_size=365
train_df=df[:-test_size]
test_df=df[-test_size:]
ktr = KTR(
response_col=response_col,
date_col=date_col,
seed=2021,
estimator='pyro-svi',
level_segments=0,
# bootstrap sampling to capture uncertainties
n_bootstrap_draws=1e4,
# pyro training config
num_steps=301,
message=100,
)
ktr.fit(train_df)
ktr.predict(test_df)
The code above gives IndexError: index -2 is out of bounds for axis 0 with size 1
.
Expected behavior Separate code paths for single and multiple level knots.
Screenshots If applicable, add screenshots to help explain your problem.
Environment (please complete the following information):
- OS: Linux
- Python Version: 3.10.8
- Versions of Major Dependencies (
pandas
,scikit-learn
,cython
):pandas==1.5.3
,scikit-learn==1.2.1
Additional context Could we separate level and trend? I.e. Level is the dynamic intercept part, and trend is the dynamic slope part, which can extend into the future smoothly.