Features/indexing (#150)

* add support for columns to the TimeSeries object

* add colum support indexing to timeseries

* fix wrong docstring

* refactor indexing, fix docstring, columns as last arg

* clean indexing method

* refactor indexing only based on loc and iloc

* Update darts/timeseries.py

Co-authored-by: Julien Herzen <[email protected]>

* use underlying columns by default

* fix column added on intern _df and use self.freq_str

* fix parameter position in from_times_and_values

* fix the tests to use str columns

* fix docstring timeseries

* remove None check on df that should exists

* add comment for clarifying that _df is a copy

* add separate function to process columns

* adapt map with str col indexing

Co-authored-by: Julien Herzen <[email protected]>
Co-authored-by: TheMP <[email protected]>

Author: 3 people

darts/models/forecasting_model.py

@@ -93,7 +93,7 @@ def _build_forecast_series(self,
 
         time_index = self._generate_new_dates(len(points_preds))
 
-        return TimeSeries.from_times_and_values(time_index, points_preds, self.training_series.freq())
+        return TimeSeries.from_times_and_values(time_index, points_preds, freq=self.training_series.freq())
 
 
 class UnivariateForecastingModel(ForecastingModel):

darts/tests/test_timeseries.py

@@ -347,8 +347,9 @@ def test_getitem(self):
         seriesA: TimeSeries = self.series1.drop_after(pd.Timestamp("20130105"))
         self.assertEqual(self.series1[pd.date_range('20130101', ' 20130104')], seriesA)
         self.assertEqual(self.series1[:4], seriesA)
-        self.assertTrue(self.series1[pd.Timestamp('20130101')].equals(self.series1.pd_dataframe()[:1]))
-        self.assertEqual(self.series1[pd.Timestamp('20130101'):pd.Timestamp('20130105')], seriesA)
+        self.assertTrue(self.series1[pd.Timestamp('20130101')] == TimeSeries(self.series1.pd_dataframe()[:1],
+                                                                             freq=self.series1.freq()))
+        self.assertEqual(self.series1[pd.Timestamp('20130101'):pd.Timestamp('20130104')], seriesA)
 
         with self.assertRaises(IndexError):
             self.series1[pd.date_range('19990101', '19990201')]
@@ -441,21 +442,21 @@ def test_map(self):
         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_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_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(series_01, series.map(fn))

darts/tests/test_timeseries_multivariate.py

@@ -13,18 +13,18 @@ class TimeSeriesMultivariateTestCase(unittest.TestCase):
     times1 = pd.date_range('20130101', '20130110')
     times2 = pd.date_range('20130206', '20130215')
     dataframe1 = pd.DataFrame({
-        0: range(10),
-        1: range(5, 15),
-        2: range(10, 20)
+        "0": range(10),
+        "1": range(5, 15),
+        "2": range(10, 20)
     }, index=times1)
     dataframe2 = pd.DataFrame({
-        0: np.arange(1, 11),
-        1: np.arange(1, 11) * 3,
-        2: np.arange(1, 11) * 5
+        "0": np.arange(1, 11),
+        "1": np.arange(1, 11) * 3,
+        "2": np.arange(1, 11) * 5
     }, index=times1)
     dataframe3 = pd.DataFrame({
-        0: np.arange(1, 11),
-        1: np.arange(11, 21),
+        "0": np.arange(1, 11),
+        "1": np.arange(11, 21),
     }, index=times2)
     series1 = TimeSeries(dataframe1)
     series2 = TimeSeries(dataframe2)
@@ -44,7 +44,7 @@ def test_creation(self):
         # Series cannot be lower than three without passing frequency as argument to constructor
         with self.assertRaises(ValueError):
             TimeSeries(self.dataframe1.iloc[:2, :])
-        TimeSeries(self.dataframe1.iloc[:2, :], 'D')
+        TimeSeries(self.dataframe1.iloc[:2, :], freq='D')
 
     def test_eq(self):
         seriesA = TimeSeries(self.dataframe1)
@@ -73,9 +73,9 @@ def test_rescale(self):
         seriesB = self.series2.rescale_with_value(1)
         self.assertEqual(seriesB, TimeSeries(
             pd.DataFrame({
-                0: np.arange(1, 11),
-                1: np.arange(1, 11),
-                2: np.arange(1, 11)
+                "0": np.arange(1, 11),
+                "1": np.arange(1, 11),
+                "2": np.arange(1, 11)
             }, index=self.dataframe2.index).astype(float)
         ))
 
@@ -119,7 +119,7 @@ def test_stack(self):
             self.series1.stack(self.series3)
         seriesA = self.series1.stack(self.series2)
         dataframeA = pd.concat([self.dataframe1, self.dataframe2], axis=1)
-        dataframeA.columns = range(6)
+        dataframeA.columns = [str(i) for i in range(6)]
         self.assertTrue((seriesA.pd_dataframe() == dataframeA).all().all())
         self.assertEqual(seriesA.values().shape, (len(self.dataframe1),
                                                   len(self.dataframe1.columns) + len(self.dataframe2.columns)))
@@ -130,7 +130,7 @@ def test_univariate_component(self):
         with self.assertRaises(ValueError):
             self.series1.univariate_component(3)
         seriesA = self.series1.univariate_component(1)
-        self.assertTrue(seriesA == TimeSeries.from_times_and_values(self.times1, range(5, 15)))
+        self.assertTrue(seriesA == TimeSeries.from_times_and_values(self.times1, range(5, 15), columns=["1"]))
         seriesB = self.series1.univariate_component(0).stack(seriesA).stack(self.series1.univariate_component(2))
         self.assertTrue(self.series1 == seriesB)