prompt: We are currently solving the following issue within our repository. Here is the issue text:
  --- BEGIN 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.
  
  --- END ISSUE ---
  
  Below are some code segments, each from a relevant file. One or more of these files may contain bugs.
  --- BEGIN FILES ---
  Path: `lightwood/mixer/sktime.py`
  Content:
  ```
  1 import numpy as np
  2 import pandas as pd
  3 from typing import Dict, Union
  4 from sktime.forecasting.arima import AutoARIMA
  5 
  6 from lightwood.api import dtype
  7 from lightwood.helpers.log import log
  8 from lightwood.mixer.base import BaseMixer
  9 from lightwood.api.types import PredictionArguments
  10 from lightwood.encoder.time_series.helpers.common import get_group_matches
  11 from lightwood.data.encoded_ds import EncodedDs, ConcatedEncodedDs
  12 
  13 
  14 class SkTime(BaseMixer):
  15     forecaster: str
  16     n_ts_predictions: int
  17     target: str
  18     supports_proba: bool
  19 
  20     def __init__(
  21             self, stop_after: int, target: str, dtype_dict: Dict[str, str],
  22             n_ts_predictions: int, ts_analysis: Dict):
  23         super().__init__(stop_after)
  24         self.target = target
  25         dtype_dict[target] = dtype.float
  26         self.model_class = AutoARIMA
  27         self.models = {}
  28         self.n_ts_predictions = n_ts_predictions
  29         self.ts_analysis = ts_analysis
  30         self.forecasting_horizon = np.arange(1, self.n_ts_predictions)
  31         self.cutoff_index = {}  # marks index at which training data stops and forecasting window starts
  32         self.grouped_by = ['__default'] if not ts_analysis['tss'].group_by else ts_analysis['tss'].group_by
  33         self.supports_proba = False
  34         self.stable = True
  35 
  36     def fit(self, train_data: EncodedDs, dev_data: EncodedDs) -> None:
  37         log.info('Started fitting sktime forecaster for array prediction')
  38 
  39         all_subsets = ConcatedEncodedDs([train_data, dev_data])
  40         df = all_subsets.data_frame.sort_values(by=f'__mdb_original_{self.ts_analysis["tss"].order_by[0]}')
  41         data = {'data': df[self.target],
  42                 'group_info': {gcol: df[gcol].tolist()
  43                                for gcol in self.grouped_by} if self.ts_analysis['tss'].group_by else {}}
  44 
  45         for group in self.ts_analysis['group_combinations']:
  46             # many warnings might be thrown inside of statsmodels during stepwise procedure
  47             self.models[group] = self.model_class(suppress_warnings=True)
  48 
  49             if self.grouped_by == ['__default']:
  50                 series_idxs = data['data'].index
  51                 series_data = data['data'].values
  52             else:
  53                 series_idxs, series_data = get_group_matches(data, group)
  54 
  55             if series_data.size > 0:
  56                 series = pd.Series(series_data.squeeze(), index=series_idxs)
  57                 series = series.sort_index(ascending=True)
  58                 series = series.reset_index(drop=True)
  59                 try:
  60                     self.models[group].fit(series)
  61                 except ValueError:
  62                     self.models[group] = self.model_class(deseasonalize=False)
  63                     self.models[group].fit(series)
  64 
  65                 self.cutoff_index[group] = len(series)
  66 
  67             if self.grouped_by == ['__default']:
  68                 break
  69 
  70     def __call__(self, ds: Union[EncodedDs, ConcatedEncodedDs],
  71                  args: PredictionArguments = PredictionArguments()) -> pd.DataFrame:
  72         if args.predict_proba:
  73             log.warning('This model does not output probability estimates')
  74 
  75         length = sum(ds.encoded_ds_lenghts) if isinstance(ds, ConcatedEncodedDs) else len(ds)
  76         ydf = pd.DataFrame(0,  # zero-filled
  77                            index=np.arange(length),
  78                            columns=['prediction'],
  79                            dtype=object)
  80 
  81         data = {'data': ds.data_frame[self.target].reset_index(drop=True),
  82                 'group_info': {gcol: ds.data_frame[gcol].tolist()
  83                                for gcol in self.grouped_by} if self.ts_analysis['tss'].group_by else {}}
  84 
  85         # all_idxs = list(range(length))  # @TODO: substract, and assign empty predictions to remainder
  86 
  87         for group in self.ts_analysis['group_combinations']:
  88 
  89             if self.grouped_by == ['__default']:
  90                 series_idxs = data['data'].index
  91                 series_data = data['data'].values
  92             else:
  93                 series_idxs, series_data = get_group_matches(data, group)
  94 
  95             if series_data.size > 0:
  96                 forecaster = self.models[group] if self.models[group].is_fitted else self.models['__default']
  97 
  98                 series = pd.Series(series_data.squeeze(), index=series_idxs)
  99                 series = series.sort_index(ascending=True)
  100                 series = series.reset_index(drop=True)
  101 
  102                 for idx, _ in enumerate(series.iteritems()):
  103                     ydf['prediction'].iloc[series_idxs[idx]] = forecaster.predict(
  104                         np.arange(idx,  # +cutoff
  105                                   idx + self.n_ts_predictions)).tolist()  # +cutoff
  106 
  107             if self.grouped_by == ['__default']:
  108                 break
  109 
  110         return ydf[['prediction']]
  111
  ```
  --- END FILES ---
  
  Please first localize the bug based on the issue statement, and then generate a patch according to the `git diff` format fenced by three backticks.
  
  Here is an example:
  
  ```diff
  diff --git a/examples/server_async.py b/examples/server_async.py
  --- a/examples/server_async.py
  +++ b/examples/server_async.py
  @@ -313,4 +313,4 @@
   
   
   if __name__ == "__main__":
  -    asyncio.run(run_async_server("."), debug=True)
  +    asyncio.run(run_async_server(), debug=True)
  diff --git a/examples/server_sync.py b/examples/server_sync.py
  --- a/examples/server_sync.py
  +++ b/examples/server_sync.py
  @@ -313,5 +313,5 @@
   
   
   if __name__ == "__main__":
  -    server = run_sync_server(".")
  +    server = run_sync_server()
       server.shutdown()
  
  ```
  
golden_diff: 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 @@
           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 @@
               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 @@
   
                   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
  
verification_info: {"golden_diff": "diff --git a/lightwood/mixer/sktime.py b/lightwood/mixer/sktime.py\n--- a/lightwood/mixer/sktime.py\n+++ b/lightwood/mixer/sktime.py\n@@ -32,6 +32,7 @@\n         self.grouped_by = ['__default'] if not ts_analysis['tss'].group_by else ts_analysis['tss'].group_by\n         self.supports_proba = False\n         self.stable = True\n+        self.prepared = False\n \n     def fit(self, train_data: EncodedDs, dev_data: EncodedDs) -> None:\n         log.info('Started fitting sktime forecaster for array prediction')\n@@ -67,10 +68,26 @@\n             if self.grouped_by == ['__default']:\n                 break\n \n+    def partial_fit(self, train_data: EncodedDs, dev_data: EncodedDs) -> None:\n+        \"\"\"\n+        Note: sktime asks for \"specification of the time points for which forecasts are requested\",\n+        and this mixer complies by assuming forecasts will start immediately after the last observed\n+        value.\n+\n+        Because of this, `partial_fit` ensures that both `dev` and `test` splits are used to fit the AutoARIMA model.\n+\n+        Due to how lightwood implements the `update` procedure, expected inputs are (for a train-dev-test split):\n+\n+        :param dev_data: original `test` split (used to validate and select model if ensemble is `BestOf`)\n+        :param train_data: includes original `train` and `dev` split\n+        \"\"\"  # noqa\n+        self.fit(dev_data, train_data)\n+        self.prepared = True\n+\n     def __call__(self, ds: Union[EncodedDs, ConcatedEncodedDs],\n                  args: PredictionArguments = PredictionArguments()) -> pd.DataFrame:\n         if args.predict_proba:\n-            log.warning('This model does not output probability estimates')\n+            log.warning('This mixer does not output probability estimates')\n \n         length = sum(ds.encoded_ds_lenghts) if isinstance(ds, ConcatedEncodedDs) else len(ds)\n         ydf = pd.DataFrame(0,  # zero-filled\n@@ -101,8 +118,7 @@\n \n                 for idx, _ in enumerate(series.iteritems()):\n                     ydf['prediction'].iloc[series_idxs[idx]] = forecaster.predict(\n-                        np.arange(idx,  # +cutoff\n-                                  idx + self.n_ts_predictions)).tolist()  # +cutoff\n+                        np.arange(idx, idx + self.n_ts_predictions)).tolist()\n \n             if self.grouped_by == ['__default']:\n                 break\n", "issue": "Incorrect `Sktime` forecasting horizon starting point\n* Lightwood version: 1.5.0\r\n\r\nFrom 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.\n", "before_files": [{"content": "import numpy as np\nimport pandas as pd\nfrom typing import Dict, Union\nfrom sktime.forecasting.arima import AutoARIMA\n\nfrom lightwood.api import dtype\nfrom lightwood.helpers.log import log\nfrom lightwood.mixer.base import BaseMixer\nfrom lightwood.api.types import PredictionArguments\nfrom lightwood.encoder.time_series.helpers.common import get_group_matches\nfrom lightwood.data.encoded_ds import EncodedDs, ConcatedEncodedDs\n\n\nclass SkTime(BaseMixer):\n    forecaster: str\n    n_ts_predictions: int\n    target: str\n    supports_proba: bool\n\n    def __init__(\n            self, stop_after: int, target: str, dtype_dict: Dict[str, str],\n            n_ts_predictions: int, ts_analysis: Dict):\n        super().__init__(stop_after)\n        self.target = target\n        dtype_dict[target] = dtype.float\n        self.model_class = AutoARIMA\n        self.models = {}\n        self.n_ts_predictions = n_ts_predictions\n        self.ts_analysis = ts_analysis\n        self.forecasting_horizon = np.arange(1, self.n_ts_predictions)\n        self.cutoff_index = {}  # marks index at which training data stops and forecasting window starts\n        self.grouped_by = ['__default'] if not ts_analysis['tss'].group_by else ts_analysis['tss'].group_by\n        self.supports_proba = False\n        self.stable = True\n\n    def fit(self, train_data: EncodedDs, dev_data: EncodedDs) -> None:\n        log.info('Started fitting sktime forecaster for array prediction')\n\n        all_subsets = ConcatedEncodedDs([train_data, dev_data])\n        df = all_subsets.data_frame.sort_values(by=f'__mdb_original_{self.ts_analysis[\"tss\"].order_by[0]}')\n        data = {'data': df[self.target],\n                'group_info': {gcol: df[gcol].tolist()\n                               for gcol in self.grouped_by} if self.ts_analysis['tss'].group_by else {}}\n\n        for group in self.ts_analysis['group_combinations']:\n            # many warnings might be thrown inside of statsmodels during stepwise procedure\n            self.models[group] = self.model_class(suppress_warnings=True)\n\n            if self.grouped_by == ['__default']:\n                series_idxs = data['data'].index\n                series_data = data['data'].values\n            else:\n                series_idxs, series_data = get_group_matches(data, group)\n\n            if series_data.size > 0:\n                series = pd.Series(series_data.squeeze(), index=series_idxs)\n                series = series.sort_index(ascending=True)\n                series = series.reset_index(drop=True)\n                try:\n                    self.models[group].fit(series)\n                except ValueError:\n                    self.models[group] = self.model_class(deseasonalize=False)\n                    self.models[group].fit(series)\n\n                self.cutoff_index[group] = len(series)\n\n            if self.grouped_by == ['__default']:\n                break\n\n    def __call__(self, ds: Union[EncodedDs, ConcatedEncodedDs],\n                 args: PredictionArguments = PredictionArguments()) -> pd.DataFrame:\n        if args.predict_proba:\n            log.warning('This model does not output probability estimates')\n\n        length = sum(ds.encoded_ds_lenghts) if isinstance(ds, ConcatedEncodedDs) else len(ds)\n        ydf = pd.DataFrame(0,  # zero-filled\n                           index=np.arange(length),\n                           columns=['prediction'],\n                           dtype=object)\n\n        data = {'data': ds.data_frame[self.target].reset_index(drop=True),\n                'group_info': {gcol: ds.data_frame[gcol].tolist()\n                               for gcol in self.grouped_by} if self.ts_analysis['tss'].group_by else {}}\n\n        # all_idxs = list(range(length))  # @TODO: substract, and assign empty predictions to remainder\n\n        for group in self.ts_analysis['group_combinations']:\n\n            if self.grouped_by == ['__default']:\n                series_idxs = data['data'].index\n                series_data = data['data'].values\n            else:\n                series_idxs, series_data = get_group_matches(data, group)\n\n            if series_data.size > 0:\n                forecaster = self.models[group] if self.models[group].is_fitted else self.models['__default']\n\n                series = pd.Series(series_data.squeeze(), index=series_idxs)\n                series = series.sort_index(ascending=True)\n                series = series.reset_index(drop=True)\n\n                for idx, _ in enumerate(series.iteritems()):\n                    ydf['prediction'].iloc[series_idxs[idx]] = forecaster.predict(\n                        np.arange(idx,  # +cutoff\n                                  idx + self.n_ts_predictions)).tolist()  # +cutoff\n\n            if self.grouped_by == ['__default']:\n                break\n\n        return ydf[['prediction']]\n", "path": "lightwood/mixer/sktime.py"}], "after_files": [{"content": "import numpy as np\nimport pandas as pd\nfrom typing import Dict, Union\nfrom sktime.forecasting.arima import AutoARIMA\n\nfrom lightwood.api import dtype\nfrom lightwood.helpers.log import log\nfrom lightwood.mixer.base import BaseMixer\nfrom lightwood.api.types import PredictionArguments\nfrom lightwood.encoder.time_series.helpers.common import get_group_matches\nfrom lightwood.data.encoded_ds import EncodedDs, ConcatedEncodedDs\n\n\nclass SkTime(BaseMixer):\n    forecaster: str\n    n_ts_predictions: int\n    target: str\n    supports_proba: bool\n\n    def __init__(\n            self, stop_after: int, target: str, dtype_dict: Dict[str, str],\n            n_ts_predictions: int, ts_analysis: Dict):\n        super().__init__(stop_after)\n        self.target = target\n        dtype_dict[target] = dtype.float\n        self.model_class = AutoARIMA\n        self.models = {}\n        self.n_ts_predictions = n_ts_predictions\n        self.ts_analysis = ts_analysis\n        self.forecasting_horizon = np.arange(1, self.n_ts_predictions)\n        self.cutoff_index = {}  # marks index at which training data stops and forecasting window starts\n        self.grouped_by = ['__default'] if not ts_analysis['tss'].group_by else ts_analysis['tss'].group_by\n        self.supports_proba = False\n        self.stable = True\n        self.prepared = False\n\n    def fit(self, train_data: EncodedDs, dev_data: EncodedDs) -> None:\n        log.info('Started fitting sktime forecaster for array prediction')\n\n        all_subsets = ConcatedEncodedDs([train_data, dev_data])\n        df = all_subsets.data_frame.sort_values(by=f'__mdb_original_{self.ts_analysis[\"tss\"].order_by[0]}')\n        data = {'data': df[self.target],\n                'group_info': {gcol: df[gcol].tolist()\n                               for gcol in self.grouped_by} if self.ts_analysis['tss'].group_by else {}}\n\n        for group in self.ts_analysis['group_combinations']:\n            # many warnings might be thrown inside of statsmodels during stepwise procedure\n            self.models[group] = self.model_class(suppress_warnings=True)\n\n            if self.grouped_by == ['__default']:\n                series_idxs = data['data'].index\n                series_data = data['data'].values\n            else:\n                series_idxs, series_data = get_group_matches(data, group)\n\n            if series_data.size > 0:\n                series = pd.Series(series_data.squeeze(), index=series_idxs)\n                series = series.sort_index(ascending=True)\n                series = series.reset_index(drop=True)\n                try:\n                    self.models[group].fit(series)\n                except ValueError:\n                    self.models[group] = self.model_class(deseasonalize=False)\n                    self.models[group].fit(series)\n\n                self.cutoff_index[group] = len(series)\n\n            if self.grouped_by == ['__default']:\n                break\n\n    def partial_fit(self, train_data: EncodedDs, dev_data: EncodedDs) -> None:\n        \"\"\"\n        Note: sktime asks for \"specification of the time points for which forecasts are requested\",\n        and this mixer complies by assuming forecasts will start immediately after the last observed\n        value.\n\n        Because of this, `partial_fit` ensures that both `dev` and `test` splits are used to fit the AutoARIMA model.\n\n        Due to how lightwood implements the `update` procedure, expected inputs are (for a train-dev-test split):\n\n        :param dev_data: original `test` split (used to validate and select model if ensemble is `BestOf`)\n        :param train_data: includes original `train` and `dev` split\n        \"\"\"  # noqa\n        self.fit(dev_data, train_data)\n        self.prepared = True\n\n    def __call__(self, ds: Union[EncodedDs, ConcatedEncodedDs],\n                 args: PredictionArguments = PredictionArguments()) -> pd.DataFrame:\n        if args.predict_proba:\n            log.warning('This mixer does not output probability estimates')\n\n        length = sum(ds.encoded_ds_lenghts) if isinstance(ds, ConcatedEncodedDs) else len(ds)\n        ydf = pd.DataFrame(0,  # zero-filled\n                           index=np.arange(length),\n                           columns=['prediction'],\n                           dtype=object)\n\n        data = {'data': ds.data_frame[self.target].reset_index(drop=True),\n                'group_info': {gcol: ds.data_frame[gcol].tolist()\n                               for gcol in self.grouped_by} if self.ts_analysis['tss'].group_by else {}}\n\n        # all_idxs = list(range(length))  # @TODO: substract, and assign empty predictions to remainder\n\n        for group in self.ts_analysis['group_combinations']:\n\n            if self.grouped_by == ['__default']:\n                series_idxs = data['data'].index\n                series_data = data['data'].values\n            else:\n                series_idxs, series_data = get_group_matches(data, group)\n\n            if series_data.size > 0:\n                forecaster = self.models[group] if self.models[group].is_fitted else self.models['__default']\n\n                series = pd.Series(series_data.squeeze(), index=series_idxs)\n                series = series.sort_index(ascending=True)\n                series = series.reset_index(drop=True)\n\n                for idx, _ in enumerate(series.iteritems()):\n                    ydf['prediction'].iloc[series_idxs[idx]] = forecaster.predict(\n                        np.arange(idx, idx + self.n_ts_predictions)).tolist()\n\n            if self.grouped_by == ['__default']:\n                break\n\n        return ydf[['prediction']]\n", "path": "lightwood/mixer/sktime.py"}]}