Color choice for multivariate timeseries (#2680)

* colors arg added, accept lists of length num_comp

* color or c in kwargs

* check added

* changelog entry added

* refactoring after discussion with Dennis

* CHANGELOG changed

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

* custom_colors instead of multi_color

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

* logic check upgraded

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

* doc corrected

* spelling

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

* change of logic

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

* change of logic follow-up

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

* alpha moved to arguments, error if color and c are used at the same time

* update plotting

---------

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

Author: 3 people

CHANGELOG.md

@@ -13,7 +13,8 @@ but cannot always guarantee backwards compatibility. Changes that may **break co
 
 - Added ONNX support for torch-based models with method `TorchForecastingModel.to_onnx()`. Check out [this example](https://unit8co.github.io/darts/userguide/gpu_and_tpu_usage.html#exporting-model-to-onnx-format-for-inference) from the user guide on how to export and load a model for inference. [#2620](https://github.com/unit8co/darts/pull/2620) by [Antoine Madrona](https://github.com/madtoinou)
 - Made method `ForecastingModel.untrained_model()` public. Use this method to get a new (untrained) model instance created with the same parameters. [#2684](https://github.com/unit8co/darts/pull/2684) by [Timon Erhart](https://github.com/turbotimon)
-- Made it possbile to run the quickstart notebook `00-quickstart.ipynb` locally. [#2691](https://github.com/unit8co/darts/pull/2691) by [Jules Authier](https://github.com/authierj)
+- `TimeSeries.plot()` now supports setting the color for each component in the series. Simply pass a list / sequence of colors with length matching the number of components as parameters "c" or "colors". [#2680](https://github.com/unit8co/darts/pull/2680) by [Jules Authier](https://github.com/authierj)
+- Made it possible to run the quickstart notebook `00-quickstart.ipynb` locally. [#2691](https://github.com/unit8co/darts/pull/2691) by [Jules Authier](https://github.com/authierj)
 
 **Fixed**
 

darts/timeseries.py

@@ -4105,6 +4105,9 @@ def plot(
         label: Optional[Union[str, Sequence[str]]] = "",
         max_nr_components: int = 10,
         ax: Optional[matplotlib.axes.Axes] = None,
+        alpha: Optional[float] = None,
+        color: Optional[Union[str, tuple, Sequence[str, tuple]]] = None,
+        c: Optional[Union[str, tuple, Sequence[str, tuple]]] = None,
         *args,
         **kwargs,
     ) -> matplotlib.axes.Axes:
@@ -4144,8 +4147,16 @@ def plot(
             Optionally, an axis to plot on. If `None`, and `new_plot=False`, will use the current axis. If
             `new_plot=True`, will create a new axis.
         alpha
-             Optionally, set the line alpha for deterministic series, or the confidence interval alpha for
+            Optionally, set the line alpha for deterministic series, or the confidence interval alpha for
             probabilistic series.
+        color
+            Can either be a single color or list of colors. Any matplotlib color is accepted (string, hex string,
+            RGB/RGBA tuple). If a single color and the series only has a single component, it is used as the color
+            for that component. If a single color and the series has multiple components, it is used as the color
+            for each component. If a list of colors with length equal to the number of components in the series, the
+            colors will be mapped to the components in order.
+        c
+            An alias for `color`.
         args
             some positional arguments for the `plot()` method
         kwargs
@@ -4172,40 +4183,63 @@ def plot(
                 logger,
             )
 
-        if new_plot:
-            fig, ax = plt.subplots()
+        if max_nr_components == -1:
+            n_components_to_plot = self.n_components
         else:
-            if ax is None:
-                ax = plt.gca()
+            n_components_to_plot = min(self.n_components, max_nr_components)
 
-        if not any(lw in kwargs for lw in ["lw", "linewidth"]):
-            kwargs["lw"] = 2
-
-        n_components_to_plot = max_nr_components
-        if n_components_to_plot == -1:
-            n_components_to_plot = self.n_components
-        elif self.n_components > max_nr_components:
+        if self.n_components > n_components_to_plot:
             logger.warning(
-                f"Number of components is larger than {max_nr_components} ({self.n_components}). "
-                f"Plotting only the first {max_nr_components} components."
-                f"You can overwrite this in the using the `plot_all_components` argument in plot()"
-                f"Beware that plotting a large number of components may cause performance issues."
+                f"Number of series components ({self.n_components}) is larger than the maximum number of "
+                f"components to plot ({max_nr_components}). Plotting only the first `{max_nr_components}` "
+                f"components. You can adjust the number of components to plot using `max_nr_components`."
             )
 
         if not isinstance(label, str) and isinstance(label, Sequence):
-            raise_if_not(
-                len(label) == self.n_components
-                or (
-                    self.n_components > n_components_to_plot
-                    and len(label) >= n_components_to_plot
+            if len(label) != self.n_components and len(label) != n_components_to_plot:
+                raise_log(
+                    ValueError(
+                        f"The `label` sequence must have the same length as the number of series components "
+                        f"({self.n_components}) or as the number of plotted components ({n_components_to_plot}). "
+                        f"Received length `{len(label)}`."
+                    ),
+                    logger,
+                )
+            custom_labels = True
+        else:
+            custom_labels = False
+
+        if color and c:
+            raise_log(
+                ValueError(
+                    "`color` and `c` must not be used simultaneously, use one or the other."
                 ),
-                "The label argument should have the same length as the number of plotted components "
-                f"({min(self.n_components, n_components_to_plot)}), only {len(label)} labels were provided",
                 logger,
             )
-            custom_labels = True
+        color = color or c
+        if not isinstance(color, (str, tuple)) and isinstance(color, Sequence):
+            if len(color) != self.n_components and len(color) != n_components_to_plot:
+                raise_log(
+                    ValueError(
+                        f"The `color` sequence must have the same length as the number of series components "
+                        f"({self.n_components}) or as the number of plotted components ({n_components_to_plot}). "
+                        f"Received length `{len(label)}`."
+                    ),
+                    logger,
+                )
+            custom_colors = True
         else:
-            custom_labels = False
+            custom_colors = False
+
+        kwargs["alpha"] = alpha
+        if not any(lw in kwargs for lw in ["lw", "linewidth"]):
+            kwargs["lw"] = 2
+
+        if new_plot:
+            fig, ax = plt.subplots()
+        else:
+            if ax is None:
+                ax = plt.gca()
 
         for i, c in enumerate(self._xa.component[:n_components_to_plot]):
             comp_name = str(c.values)
@@ -4219,7 +4253,6 @@ def plot(
             else:
                 central_series = comp.mean(dim=DIMS[2])
 
-            alpha = kwargs["alpha"] if "alpha" in kwargs else None
             if custom_labels:
                 label_to_use = label[i]
             else:
@@ -4230,6 +4263,7 @@ def plot(
                 else:
                     label_to_use = f"{label}_{comp_name}"
             kwargs["label"] = label_to_use
+            kwargs["c"] = color[i] if custom_colors else color
 
             kwargs_central = deepcopy(kwargs)
             if not self.is_deterministic: