|
12 | 12 | # permissions and limitations under the License. |
13 | 13 |
|
14 | 14 | import pytest |
| 15 | +import numpy as np |
15 | 16 |
|
16 | | -from gluonts.ext.rotbaum import TreeEstimator |
| 17 | +from gluonts.ext.rotbaum import TreeEstimator, TreePredictor |
17 | 18 |
|
18 | 19 | from gluonts.testutil.dummy_datasets import make_dummy_datasets_with_features |
| 20 | +from gluonts.dataset.common import ListDataset |
19 | 21 |
|
20 | 22 | # TODO: Add support for categorical and dynamic features. |
21 | 23 |
|
@@ -59,3 +61,68 @@ def test_rotbaum_smoke(datasets): |
59 | 61 | predictor = estimator.train(dataset_train) |
60 | 62 | forecasts = list(predictor.predict(dataset_test)) |
61 | 63 | assert len(forecasts) == len(dataset_test) |
| 64 | + |
| 65 | + |
| 66 | +def test_short_history_item_pred(): |
| 67 | + prediction_length = 7 |
| 68 | + freq = "D" |
| 69 | + |
| 70 | + dataset = ListDataset( |
| 71 | + data_iter=[ |
| 72 | + { |
| 73 | + "start": "2017-10-11", |
| 74 | + "item_id": "item_1", |
| 75 | + "target": np.array( |
| 76 | + [ |
| 77 | + 1.0, |
| 78 | + 9.0, |
| 79 | + 2.0, |
| 80 | + 0.0, |
| 81 | + 0.0, |
| 82 | + 1.0, |
| 83 | + 5.0, |
| 84 | + 3.0, |
| 85 | + 4.0, |
| 86 | + 2.0, |
| 87 | + 0.0, |
| 88 | + 0.0, |
| 89 | + 1.0, |
| 90 | + 6.0, |
| 91 | + ] |
| 92 | + ), |
| 93 | + "feat_static_cat": np.array([0.0, 0.0], dtype=float), |
| 94 | + "past_feat_dynamic_real": np.array( |
| 95 | + [ |
| 96 | + [1.0222e06 for i in range(14)], |
| 97 | + [750.0 for i in range(14)], |
| 98 | + ] |
| 99 | + ), |
| 100 | + }, |
| 101 | + { |
| 102 | + "start": "2017-10-11", |
| 103 | + "item_id": "item_2", |
| 104 | + "target": np.array([7.0, 0.0, 0.0, 23.0, 13.0]), |
| 105 | + "feat_static_cat": np.array([0.0, 1.0], dtype=float), |
| 106 | + "past_feat_dynamic_real": np.array( |
| 107 | + [[0 for i in range(5)], [750.0 for i in range(5)]] |
| 108 | + ), |
| 109 | + }, |
| 110 | + ], |
| 111 | + freq=freq, |
| 112 | + ) |
| 113 | + |
| 114 | + predictor = TreePredictor( |
| 115 | + freq=freq, |
| 116 | + prediction_length=prediction_length, |
| 117 | + quantiles=[0.1, 0.5, 0.9], |
| 118 | + max_n_datapts=50000, |
| 119 | + method="QuantileRegression", |
| 120 | + use_past_feat_dynamic_real=True, |
| 121 | + use_feat_dynamic_real=False, |
| 122 | + use_feat_dynamic_cat=False, |
| 123 | + use_feat_static_real=False, |
| 124 | + cardinality="auto", |
| 125 | + ) |
| 126 | + predictor = predictor.train(dataset) |
| 127 | + forecasts = list(predictor.predict(dataset)) |
| 128 | + assert forecasts[1].quantile(0.5).shape[0] == prediction_length |
0 commit comments