Implement map method for TimeSeries

darts/tests/test_timeseries.py

@@ -431,3 +431,31 @@ def test_short_series_slice(self):
         self.assertEqual(seriesA.time_index()[-1], self.series1.time_index()[0])
         seriesC = self.series1.slice(pd.Timestamp('20130105'), pd.Timestamp('20130105'))
         self.assertEqual(len(seriesC), 1)
+
+    def test_map(self):
+        fn = np.sin
+        series = TimeSeries.from_times_and_values(pd.date_range('20000101', '20000110'), np.random.randn(10, 3))
+
+        df_0 = series.pd_dataframe()
+        df_2 = series.pd_dataframe()
+        df_01 = series.pd_dataframe()
+        df_012 = series.pd_dataframe()
+
+        df_0[[0]] = df_0[[0]].applymap(fn)
+        df_2[[2]] = df_2[[2]].applymap(fn)
+        df_01[[0, 1]] = df_01[[0, 1]].applymap(fn)
+        df_012 = df_012.applymap(fn)
+
+        series_0 = TimeSeries(df_0, 'D')
+        series_2 = TimeSeries(df_2, 'D')
+        series_01 = TimeSeries(df_01, 'D')
+        series_012 = TimeSeries(df_012, 'D')
+
+        self.assertEqual(series_0, series.map(fn, 0))
+        self.assertEqual(series_0, series.map(fn, [0]))
+        self.assertEqual(series_2, series.map(fn, 2))
+        self.assertEqual(series_01, series.map(fn, [0, 1]))
+        self.assertEqual(series_012, series.map(fn, [0, 1, 2]))
+        self.assertEqual(series_012, series.map(fn))
+
+        self.assertNotEqual(series01, series.map(fn))

darts/timeseries.py

@@ -880,6 +880,39 @@ def is_within_range(self,
         index = self.time_index()
         return index[0] <= ts <= index[-1]
 
+
+    def map(self,
+            fn: Callable[[np.number], np.number],
+            cols: Optional[Union[List[int], int]] = None) -> 'TimeSeries':
+        """
+        Applies the function `fn` elementwise to all values in this TimeSeries, or, to only those
+        values in the columns specified by the optional argument `cols`. Returns a new
+        TimeSeries instance.
+
+        Parameters
+        ----------
+        fn
+            A numerical function
+        cols
+            Optionally, an integer or list of integers specifying the column(s) onto which fn should be applied
+
+        Returns
+        -------
+        TimeSeries
+            A new TimeSeries instance
+        """
+        if cols is None:
+            new_dataframe = self._df.applymap(fn)
+        else:
+            if isinstance(cols, int):
+                cols = [cols]
+            raise_if_not(all([0 <= index and index < self.width for index in cols]),
+                     'The indices in `cols` must be between 0 and the number of components of the current '
+                     'TimeSeries instance - 1, {}'.format(self.width - 1), logger)
+            new_dataframe = self.pd_dataframe() 
+            new_dataframe[cols] = new_dataframe[cols].applymap(fn)
+        return TimeSeries(new_dataframe, self.freq_str())
+
     @staticmethod
     def _combine_or_none(df_a: Optional[pd.DataFrame],
                          df_b: Optional[pd.DataFrame],