Fix/filter static cov component (unit8co#1128)

* simplify and improve univariate_component()

* add unit test

* Update darts/timeseries.py

Co-authored-by: Dennis Bader <[email protected]>

Co-authored-by: Dennis Bader <[email protected]>

Author: hrzn

darts/tests/test_timeseries.py

@@ -105,6 +105,31 @@ def test_integer_indexing(self):
             list(indexed_ts.time_index) == list(pd.RangeIndex(2, 7, step=1))
         )
 
+    def test_univariate_component(self):
+        series = TimeSeries.from_values(np.array([10, 20, 30])).with_columns_renamed(
+            "0", "component"
+        )
+        mseries = concatenate([series] * 3, axis="component")
+        mseries = mseries.with_hierarchy(
+            {"component_1": ["component"], "component_2": ["component"]}
+        )
+
+        static_cov = pd.DataFrame(
+            {"dim0": [1, 2, 3], "dim1": [-2, -1, 0], "dim2": [0.0, 0.1, 0.2]}
+        )
+
+        mseries = mseries.with_static_covariates(static_cov)
+
+        for univ_series in [
+            mseries.univariate_component(1),
+            mseries.univariate_component("component_1"),
+        ]:
+            # hierarchy should be dropped
+            self.assertIsNone(univ_series.hierarchy)
+
+            # only the right static covariate column should be retained
+            self.assertEqual(univ_series.static_covariates.sum().sum(), 1.1)
+
     def test_column_names(self):
         # test the column names resolution
         columns_before = [

darts/timeseries.py

@@ -2702,6 +2702,9 @@ def univariate_component(self, index: Union[str, int]) -> "TimeSeries":
         Retrieve one of the components of the series
         and return it as new univariate ``TimeSeries`` instance.
 
+        This drops the hierarchy (if any), and retains only the relevant static
+        covariates column.
+
         Parameters
         ----------
         index
@@ -2713,11 +2716,8 @@ def univariate_component(self, index: Union[str, int]) -> "TimeSeries":
         TimeSeries
             A new univariate TimeSeries instance.
         """
-        if isinstance(index, int):
-            new_xa = self._xa.isel(component=index).expand_dims(DIMS[1], axis=1)
-        else:
-            new_xa = self._xa.sel(component=index).expand_dims(DIMS[1], axis=1)
-        return self.__class__(new_xa)
+
+        return self[index if isinstance(index, str) else self.components[index]]
 
     def add_datetime_attribute(
         self, attribute, one_hot: bool = False, cyclic: bool = False