diff --git a/darts/tests/test_timeseries.py b/darts/tests/test_timeseries.py index 31967ca193..b874370aa0 100644 --- a/darts/tests/test_timeseries.py +++ b/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)) diff --git a/darts/timeseries.py b/darts/timeseries.py index d126838eca..baf658f8d8 100644 --- a/darts/timeseries.py +++ b/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],