Repository: mindsdb/lightwood

Issue:
Incorrect `Sktime` forecasting horizon starting point
* Lightwood version: 1.5.0

From a few internal tests, it seems the `sktime` time series mixer is not emitting forecasts from the end of the validation dataset, but from the training dataset instead, leading to predictions that will be incorrectly displaced.


Discussion:

Resolving patch:
diff --git a/lightwood/mixer/sktime.py b/lightwood/mixer/sktime.py
--- a/lightwood/mixer/sktime.py
+++ b/lightwood/mixer/sktime.py
@@ -32,6 +32,7 @@ def __init__(
         self.grouped_by = ['__default'] if not ts_analysis['tss'].group_by else ts_analysis['tss'].group_by
         self.supports_proba = False
         self.stable = True
+        self.prepared = False
 
     def fit(self, train_data: EncodedDs, dev_data: EncodedDs) -> None:
         log.info('Started fitting sktime forecaster for array prediction')
@@ -67,10 +68,26 @@ def fit(self, train_data: EncodedDs, dev_data: EncodedDs) -> None:
             if self.grouped_by == ['__default']:
                 break
 
+    def partial_fit(self, train_data: EncodedDs, dev_data: EncodedDs) -> None:
+        """
+        Note: sktime asks for "specification of the time points for which forecasts are requested",
+        and this mixer complies by assuming forecasts will start immediately after the last observed
+        value.
+
+        Because of this, `partial_fit` ensures that both `dev` and `test` splits are used to fit the AutoARIMA model.
+
+        Due to how lightwood implements the `update` procedure, expected inputs are (for a train-dev-test split):
+
+        :param dev_data: original `test` split (used to validate and select model if ensemble is `BestOf`)
+        :param train_data: includes original `train` and `dev` split
+        """  # noqa
+        self.fit(dev_data, train_data)
+        self.prepared = True
+
     def __call__(self, ds: Union[EncodedDs, ConcatedEncodedDs],
                  args: PredictionArguments = PredictionArguments()) -> pd.DataFrame:
         if args.predict_proba:
-            log.warning('This model does not output probability estimates')
+            log.warning('This mixer does not output probability estimates')
 
         length = sum(ds.encoded_ds_lenghts) if isinstance(ds, ConcatedEncodedDs) else len(ds)
         ydf = pd.DataFrame(0,  # zero-filled
@@ -101,8 +118,7 @@ def __call__(self, ds: Union[EncodedDs, ConcatedEncodedDs],
 
                 for idx, _ in enumerate(series.iteritems()):
                     ydf['prediction'].iloc[series_idxs[idx]] = forecaster.predict(
-                        np.arange(idx,  # +cutoff
-                                  idx + self.n_ts_predictions)).tolist()  # +cutoff
+                        np.arange(idx, idx + self.n_ts_predictions)).tolist()
 
             if self.grouped_by == ['__default']:
                 break


Test patch:
diff --git a/tests/integration/advanced/test_timeseries.py b/tests/integration/advanced/test_timeseries.py
--- a/tests/integration/advanced/test_timeseries.py
+++ b/tests/integration/advanced/test_timeseries.py
@@ -144,3 +144,56 @@ def test_2_time_series_classification(self):
 
         predictor.learn(train)
         predictor.predict(test)
+
+    def test_3_time_series_sktime_mixer(self):
+        """
+        Tests `sktime` mixer individually, as it has a special notion of absolute
+        temporal timestamps that we need to ensure are being used correctly. In
+        particular, given a train-dev-test split, any forecasts coming from a sktime
+        mixer should start from the latest observed data in the entire dataset.
+        """  # noqa
+
+        from sklearn.metrics import r2_score
+        from scipy import signal
+        from lightwood.api.high_level import (
+            ProblemDefinition,
+            json_ai_from_problem,
+            code_from_json_ai,
+            predictor_from_code,
+        )
+
+        # synth square wave
+        tsteps = 100
+        target = 'Value'
+        t = np.linspace(0, 1, tsteps, endpoint=False)
+        ts = signal.sawtooth(2 * np.pi * 5 * t, width=0.5)
+        df = pd.DataFrame(columns=['Time', target])
+        df['Time'] = t
+        df[target] = ts
+
+        train = df[:int(len(df) * 0.8)]
+        test = df[int(len(df) * 0.8):]
+
+        pdef = ProblemDefinition.from_dict({'target': target,
+                                            'time_aim': 10,
+                                            'timeseries_settings': {
+                                                'order_by': ['Time'],
+                                                'window': 5,
+                                                'nr_predictions': 20
+                                            }})
+
+        json_ai = json_ai_from_problem(df, problem_definition=pdef)
+        json_ai.outputs[target].mixers = [{
+            "module": "SkTime",
+            "args": {
+                "stop_after": "$problem_definition.seconds_per_mixer",
+                "n_ts_predictions": "$problem_definition.timeseries_settings.nr_predictions",
+            }}]
+
+        code = code_from_json_ai(json_ai)
+        predictor = predictor_from_code(code)
+
+        predictor.learn(train)
+        ps = predictor.predict(test)
+
+        assert r2_score(ps['truth'].values, ps['prediction'].iloc[0]) >= 0.95