Fix/list indexing length leq 2 (#2857)

* fix list indexing with length less than or equal to 2

* update changelog

Author: dennisbader

CHANGELOG.md

@@ -22,6 +22,7 @@ but cannot always guarantee backwards compatibility. Changes that may **break co
 - Fixed a bug in `SKLearnModel.get_estimator()` for univariate quantile models that use `multi_models=False` , where using `quantile` did not return the correct fitted quantile model / estimator. [#2838](https://github.com/unit8co/darts/pull/2838) by [Dennis Bader](https://github.com/dennisbader).
 - Fixed a bug in `LightGBMModel` and `CatBoostModel` when using component-specific lags and categorical features, where certain lag scenarios could result in incorrect categorical feature declaration. [#2852](https://github.com/unit8co/darts/pull/2852) by [Dennis Bader](https://github.com/dennisbader).
 - Fixed a bug in `darts.utils.timeseries_generation.sine_timeseries()`, where the returned series ignored the specified `dtype`. [#2856](https://github.com/unit8co/darts/pull/2856) by [Dennis Bader](https://github.com/dennisbader).
+- Fixed a bug in `TimeSeries.__getitem__()`, where indexing with a list of integers of `length <= 2` resulted in an error. [#2857](https://github.com/unit8co/darts/pull/2857) by [Dennis Bader](https://github.com/dennisbader).
 - Removed `darts/tests` and `examples` from the Darts package distribution. These are only required for internal testing. [#2854](https://github.com/unit8co/darts/pull/2854) by [Dennis Bader](https://github.com/dennisbader).
 
 **Dependencies**

darts/tests/test_timeseries.py

@@ -229,8 +229,8 @@ def test_integer_range_indexing(self):
         with pytest.raises(ValueError) as exc:
             _ = series_int[[2, 3, 6]]
         assert str(exc.value).startswith(
-            "Could not convert integer index to a `pandas.RangeIndex`. "
-            "Found non-unique step sizes/frequencies: `{1, 3}`. "
+            "Cannot index a `TimeSeries` with a list of integers with non-constant step sizes. "
+            "Observed step sizes: `{1, 3}`"
         )
 
         # check integer indexing features when series index does not start at 0
@@ -901,16 +901,6 @@ def test_getitem_datetime_index(self):
         assert series.freq == 3 * self.series1.freq
         assert series == self.series1[: 3 * 3 : 3]
 
-        series = self.series1[[0, 3, 6]]
-        assert series.freq == 3 * self.series1.freq
-        assert series == self.series1[: 3 * 3 : 3]
-
-        with pytest.raises(ValueError) as exc:
-            _ = self.series1[[0, 2, 6]]
-        assert str(exc.value).startswith(
-            "The time index is missing the `freq` attribute, and the frequency could not be directly inferred"
-        )
-
         # not all dates in index
         with pytest.raises(KeyError):
             self.series1[pd.date_range("19990101", "19990201")]
@@ -921,6 +911,8 @@ def test_getitem_datetime_index(self):
         with pytest.raises(ValueError):
             self.series1[::-1]
 
+        # list of integers is tested below
+
     def test_getitem_integer_index(self):
         freq = 3
         start = 1
@@ -972,6 +964,92 @@ def test_getitem_integer_index(self):
         with pytest.raises(KeyError):
             _ = series[pd.RangeIndex(start, stop=end + 2 * freq, step=freq)]
 
+        # list of integers is tested below
+
+    @pytest.mark.parametrize("is_dti", [True, False])
+    def test_getitem_list_of_integers(self, is_dti):
+        if is_dti:
+            # datetime-indexed series
+            series = self.series1
+        else:
+            # range-indexed series
+            idx_int = generate_index(start=1, length=len(self.series1), freq=3)
+            series = TimeSeries(times=idx_int, values=self.series1.values())
+
+        # list of integers (relative to the start) with constant step size
+        # length 3; multiple steps of size 1 starting at 0
+        series_new = series[[0, 1, 2]]
+        assert series_new.freq == series.freq
+        assert series_new == series[:3]
+
+        # length 3; multiple steps of size 1 starting after 0
+        series_new = series[[1, 2, 3]]
+        assert series_new.freq == series.freq
+        assert series_new == series[1:4]
+
+        # length 3; multiple steps starting at 0
+        series_new = series[[0, 3, 6]]
+        assert series_new.freq == 3 * series.freq
+        assert series_new == series[: 3 * 3 : 3]
+
+        # length 3; multiple steps starting after 0
+        series_new = series[[2, 4, 6]]
+        assert series_new.freq == 2 * series.freq
+        assert series_new == series[2:8:2]
+
+        # length 2; 1 step
+        series_new = series[[0, 3]]
+        assert series_new.freq == 3 * series.freq
+        assert series_new == series[: 2 * 3 : 3]
+
+        # length 1; no steps
+        series_new = series[[0]]
+        assert series_new.freq == series.freq
+        assert series_new == series[:1]
+
+        series_new = series[0]
+        assert series_new.freq == series.freq
+        assert series_new == series[:1]
+
+        # list of integers (relative to the end) with constant step size
+        n_steps = len(series)
+        # length 3; multiple steps
+        series_new = series[[-n_steps, -n_steps + 3, -n_steps + 6]]
+        assert series_new.freq == 3 * series.freq
+        assert series_new == series[: 3 * 3 : 3]
+
+        # length 2; 1 step
+        series_new = series[[-n_steps, -n_steps + 3]]
+        assert series_new.freq == 3 * series.freq
+        assert series_new == series[: 2 * 3 : 3]
+
+        # length 1; no steps
+        series_new = series[[-n_steps]]
+        assert series_new.freq == series.freq
+        assert series_new == series[:1]
+
+        # different step sizes
+        with pytest.raises(ValueError) as exc:
+            _ = series[[0, 2, 6]]
+        assert str(exc.value).startswith(
+            "Cannot index a `TimeSeries` with a list of integers with non-constant step sizes. "
+            "Observed step sizes: `{2, 4}`."
+        )
+
+        # decreasing step sizes
+        with pytest.raises(ValueError) as exc:
+            _ = series[[3, 2, 1]]
+        assert str(exc.value).startswith(
+            "Indexing a `TimeSeries` with a list of integers with `step<=0` is not possible"
+        )
+
+        # constant step sizes
+        with pytest.raises(ValueError) as exc:
+            _ = series[[3, 3]]
+        assert str(exc.value).startswith(
+            "Indexing a `TimeSeries` with a list of integers with `step<=0` is not possible"
+        )
+
     def test_getitem_frequency_inferrence(self):
         ts = self.series1
         assert ts.freq == "D"

darts/timeseries.py

@@ -5737,12 +5737,30 @@ def _check_range():
                     metadata=self.metadata,
                 )
             elif all(isinstance(i, (int, np.int64)) for i in key):
-                return self.__class__(
-                    times=self._time_index[key],
-                    values=self._values[key],
-                    components=self.components,
-                    **self._attrs,
-                )
+                # convert list of integers to slice (must have constant step size)
+                step_sizes = set(right - left for left, right in zip(key[:-1], key[1:]))
+                if len(step_sizes) > 1:
+                    raise_log(
+                        ValueError(
+                            f"Cannot index a `TimeSeries` with a list of integers with non-constant step sizes. "
+                            f"Observed step sizes: `{step_sizes}`."
+                        ),
+                        logger,
+                    )
+                elif len(step_sizes) == 1:
+                    step_size = step_sizes.pop()
+                else:
+                    step_size = 1
+
+                if step_size <= 0:
+                    raise_log(
+                        ValueError(
+                            "Indexing a `TimeSeries` with a list of integers with `step<=0` is not "
+                            "possible since `TimeSeries` must have a monotonically increasing time index."
+                        ),
+                        logger=logger,
+                    )
+                return self[key[0] : key[-1] + step_size : step_size]
 
             elif all(isinstance(t, pd.Timestamp) for t in key):
                 _check_dt()