.map() improvements (#2911)

* .map() improvements

* changelog update

* example fix

* again example fix

* codecov changes

* add shape check and update docs

* update docs

---------

Co-authored-by: dennisbader <[email protected]>

Author: jakubchlapek

CHANGELOG.md

@@ -11,6 +11,8 @@ but cannot always guarantee backwards compatibility. Changes that may **break co
 
 **Improved**
 
+- 🔴 Improved the performance of the `TimeSeries.map()` method for functions that take two arguments. The mapping is now applied on the entire time index and values array which requires users to reshape the time index explicitly within the function. See more information in the `TimeSeries.map()` method documentation. [#2911](https://github.com/unit8co/darts/pull/2911) by [Jakub Chłapek](https://github.com/jakubchlapek)
+
 **Fixed**
 
 **Dependencies**

darts/tests/dataprocessing/transformers/test_mappers.py

@@ -18,11 +18,11 @@ def inverse_func(x):
 
     @staticmethod
     def ts_func(ts, x):
-        return x - ts.month
+        return x - ts.month.values.reshape(-1, 1, 1)
 
     @staticmethod
     def inverse_ts_func(ts, x):
-        return x + ts.month
+        return x + ts.month.values.reshape(-1, 1, 1)
 
     plus_ten = Mapper(func.__func__)
     plus_ten_invertible = InvertibleMapper(func.__func__, inverse_func.__func__)

darts/tests/test_timeseries.py

@@ -1670,7 +1670,7 @@ def test_map_with_timestamp(self):
         zeroes = zeroes.with_columns_renamed("constant", "linear")
 
         def function(ts, x):
-            return x - ts.month
+            return x - ts.month.values.reshape(-1, 1, 1)
 
         new_series = series.map(function)
         assert new_series == zeroes
@@ -1695,6 +1695,20 @@ def add(x, y, z):
         with pytest.raises(ValueError):
             series.map(ufunc_add)
 
+    def test_map_fn_not_callable(self):
+        series = linear_timeseries(length=3)
+        with pytest.raises(TypeError) as exc:
+            series.map(fn=1)
+        assert str(exc.value) == "fn must be a callable"
+
+    def test_map_fn_wrong_output_shape(self):
+        series = linear_timeseries(length=3)
+        with pytest.raises(ValueError) as exc:
+            series.map(fn=lambda x: np.concatenate([x] * 2, axis=1))
+        assert str(exc.value) == (
+            "fn must return an array of shape `(3, 1, 1)`. Received shape `(3, 2, 1)`"
+        )
+
     def test_gaps(self):
         times1 = pd.date_range("20130101", "20130110")
         times2 = pd.date_range("20120101", "20210301", freq=freqs["QE"])

darts/timeseries.py

@@ -3662,16 +3662,16 @@ def is_within_range(self, ts: Union[pd.Timestamp, int]) -> bool:
     def map(
         self,
         fn: Union[
-            Callable[[np.number], np.number],
-            Callable[[Union[pd.Timestamp, int], np.number], np.number],
+            Callable[[np.ndarray], np.ndarray],
+            Callable[[Union[pd.DatetimeIndex, pd.RangeIndex], np.ndarray], np.ndarray],
         ],
     ) -> Self:  # noqa: E501
         """Return a new series with the function `fn` applied to the values of this series.
 
         If `fn` takes 1 argument it is simply applied on the values array of shape `(time, n_components, n_samples)`.
-        If `fn` takes 2 arguments, it is applied repeatedly on the `(ts, value[ts])` tuples, where `ts` denotes a
-        timestamp value, and `value[ts]` denotes the array of values at this timestamp, of shape
-        `(n_components, n_samples)`.
+        If `fn` takes 2 arguments, it is applied on the `(ts, values)` tuple, where `ts` denotes the
+        series' time index, and `values` denotes the series' array of values, of shape
+        `(n_timestamps, n_components, n_samples)`. Timestamp index's shape should be `(n, 1, 1)`;
 
         Parameters
         ----------
@@ -3687,9 +3687,36 @@ def map(
         -------
         TimeSeries
             A new series with the function `fn` applied to the values.
+
+        Examples
+        --------
+        >>> from darts import TimeSeries
+        >>> from darts.utils.utils import generate_index
+        >>> # create a simple TimeSeries
+        >>> series = TimeSeries.from_times_and_values(
+        >>>     times=generate_index("2020-01-01", length=3, freq="D"),
+        >>>     values=range(3),
+        >>> )
+        >>> # map function on values only
+        >>> def fn1(values):
+        >>>     return values / 3.
+        >>>
+        >>> series.map(fn1).values()
+        array([[0.        ],
+               [0.33333333],
+               [0.66666667]])
+        >>>
+        >>> # map function on time index and values
+        >>> def fn2(times, values):
+        >>>     return values / times.days_in_month.values.reshape(-1, 1, 1)
+        >>>
+        >>> series.map(fn2).values()
+        array([[0.        ],
+               [0.03225806],
+               [0.06451613]])
         """
         if not isinstance(fn, Callable):
-            raise_log(TypeError("fn should be callable"), logger)
+            raise_log(TypeError("fn must be a callable"), logger)
 
         if isinstance(fn, np.ufunc):
             if fn.nin == 1 and fn.nout == 1:
@@ -3716,30 +3743,18 @@ def map(
 
         if num_args == 1:  # apply fn on values directly
             values = fn(self._values)
-        elif num_args == 2:  # map function uses timestamp f(timestamp, x)
-            # go over shortest amount of iterations, either over time steps or components and samples
-            if self.n_timesteps <= self.n_components * self.n_samples:
-                new_vals = np.vstack([
-                    np.expand_dims(
-                        fn(self.time_index[i], self._values[i, :, :]), axis=0
-                    )
-                    for i in range(self.n_timesteps)
-                ])
-            else:
-                new_vals = np.stack(
-                    [
-                        np.column_stack([
-                            fn(self.time_index, self._values[:, i, j])
-                            for j in range(self.n_samples)
-                        ])
-                        for i in range(self.n_components)
-                    ],
-                    axis=1,
-                )
-            values = new_vals
-
+        elif num_args == 2:
+            # apply function on (times, values)
+            values = fn(self._time_index, self._values)
         else:
-            raise_log(ValueError("fn must have either one or two arguments"), logger)
+            raise_log(ValueError("fn must accept either one or two arguments"), logger)
+
+        if values.shape != self.shape:
+            raise_log(
+                ValueError(
+                    f"fn must return an array of shape `{self.shape}`. Received shape `{values.shape}`"
+                )
+            )
 
         return self.__class__(
             times=self._time_index,

examples/00-quickstart.ipynb

@@ -420,7 +420,7 @@
     }
    ],
    "source": [
-    "series.map(lambda ts, x: x / ts.days_in_month).plot();"
+    "series.map(lambda ts, x: x / ts.days_in_month.values.reshape(-1, 1, 1)).plot();"
    ]
   },
   {