Add per-timestep TFT variable importance plot

CHANGELOG.md

@@ -11,6 +11,7 @@ but cannot always guarantee backwards compatibility. Changes that may **break co
 
 **Improved**
 
+- Added `TFTExplainer.plot_variable_selection_over_time()` to visualize the per-timestep encoder/decoder variable importances added in [#3170](https://github.com/unit8co/darts/pull/3170). [#3174](https://github.com/unit8co/darts/pull/3174) by [exactml](https://github.com/exactml).
 - Added support for per-timestep (non-aggregated) encoder and decoder variable importances in `TFTExplainer`, exposed as `TimeSeries` via `TFTExplainabilityResult.get_encoder_importance_over_time()` and `get_decoder_importance_over_time()`. [#3170](https://github.com/unit8co/darts/pull/3170) by [exactml](https://github.com/exactml).
 - Calling `TFTModel.fit_from_dataset()` on a dataset that does not have future covariates now raises an informative exception. [#3149](https://github.com/unit8co/darts/pull/3149) by [YOON KIWOONG](https://github.com/kiwoongyoon).
 - 🔴 Percentage and range-based metrics (`ape`, `mape`, `sape`, `smape`, `wmape`, `ope`, `arre`, `marre`, `coefficient_of_variation`) no longer raise a hard `ValueError` when the denominator is exactly zero. A new `zero_division` parameter controls the behavior: [#3122](https://github.com/unit8co/darts/pull/3122) by [Mahimn](https://github.com/mahimn01).

darts/explainability/tft_explainer.py

@@ -11,6 +11,10 @@
   - decoder importance: future part of future covariates
   - static covariates importance: the numeric and categorical static covariates importance
 
+- :func:`plot_variable_selection_over_time() <TFTExplainer.plot_variable_selection_over_time>` plots the same
+  encoder / decoder variable importances as :func:`plot_variable_selection() <TFTExplainer.plot_variable_selection>`,
+  but per individual timestep instead of aggregated over the whole input chunk / forecast horizon.
+
 - :func:`plot_attention() <TFTExplainer.plot_attention>` plots the transformer attention that the `TFTModel` applies
   on the given past and future input. The attention is aggregated over all attention heads.
 
@@ -341,6 +345,105 @@ def plot_variable_selection(
             return plotted_figures[0]
         return plotted_figures
 
+    def plot_variable_selection_over_time(
+        self,
+        expl_result: TFTExplainabilityResult,
+        show_index_as: Literal["relative", "time"] = "relative",
+        max_nr_components: int = 10,
+        fig_size=None,
+        max_nr_series: int = 5,
+        show_plot: bool = True,
+    ) -> Figure | list[Figure]:
+        """Plots the variable selection (feature importance) of the `TFTModel` per individual timestep.
+
+        Unlike :func:`plot_variable_selection() <TFTExplainer.plot_variable_selection>`, which aggregates
+        importances over the whole input chunk / forecast horizon, this shows how each variable's importance
+        evolves at every timestep. The figure includes two subplots:
+
+        - encoder importance over time: importance of each encoder variable at every input chunk timestep
+        - decoder importance over time: importance of each decoder variable at every forecast horizon timestep
+
+        Parameters
+        ----------
+        expl_result
+            A `TFTExplainabilityResult` object. Corresponds to the output of :func:`explain() <TFTExplainer.explain>`.
+        show_index_as
+            The type of index to be shown on the x-axis. One of ("relative", "time").
+            If "relative", will plot the x-axis from `(-input_chunk_length, output_chunk_length - 1)`. `0` corresponds
+            to the first prediction point.
+            If "time", will plot the x-axis with the actual time index (or range index) of the corresponding
+            `TFTExplainabilityResult`.
+        max_nr_components
+            The maximum number of variables to plot per subplot. -1 means all variables will be plotted.
+        fig_size
+            The size of the figure to be plotted.
+        max_nr_series
+            The maximum number of plots to show in case `expl_result` was computed on multiple series.
+        show_plot
+            Whether to show the plot.
+
+        Returns
+        -------
+        Figure | list[Figure]
+            The matplotlib figures used for plotting. Returns a single Figure when explaining a single series,
+            and a list of Figures when explaining multiple series.
+        """
+        encoder_importance_ot = expl_result.get_encoder_importance_over_time()
+        decoder_importance_ot = expl_result.get_decoder_importance_over_time()
+        if not isinstance(encoder_importance_ot, list):
+            encoder_importance_ot = [encoder_importance_ot]
+            decoder_importance_ot = [decoder_importance_ot]
+
+        plotted_figures: list[Figure] = []
+        for idx, (enc_imp_ot, dec_imp_ot) in enumerate(
+            zip(encoder_importance_ot, decoder_importance_ot)
+        ):
+            if show_index_as == "relative":
+                enc_imp_ot = TimeSeries(
+                    times=generate_index(start=-len(enc_imp_ot), end=-1),
+                    values=enc_imp_ot.values(copy=False),
+                    components=enc_imp_ot.components,
+                )
+                dec_imp_ot = TimeSeries(
+                    times=generate_index(start=0, end=len(dec_imp_ot) - 1),
+                    values=dec_imp_ot.values(copy=False),
+                    components=dec_imp_ot.components,
+                )
+                x_label = "Index relative to first prediction point"
+            elif show_index_as == "time":
+                x_label = "Time index"
+            else:
+                raise_log(
+                    ValueError("`show_index_as` must either be 'relative', or 'time'.")
+                )
+
+            fig, axes = plt.subplots(nrows=2, figsize=fig_size)
+            for ax, imp_ot, ax_title in zip(
+                axes,
+                [enc_imp_ot, dec_imp_ot],
+                [
+                    "Encoder variable importance over time",
+                    "Decoder variable importance over time",
+                ],
+            ):
+                imp_ot.plot(max_nr_components=max_nr_components, ax=ax)
+                ax.set_title(ax_title)
+                ax.set_xlabel(x_label)
+                ax.set_ylabel("Importance in %")
+                ax.legend(bbox_to_anchor=(1.05, 1), loc="upper left")
+
+            fig.tight_layout()
+            if show_plot:
+                plt.show()
+            plotted_figures.append(fig)
+
+            if idx + 1 == max_nr_series:
+                break
+
+        if len(plotted_figures) == 1:
+            return plotted_figures[0]
+        return plotted_figures
+
     def plot_attention(
         self,
         expl_result: TFTExplainabilityResult,

darts/tests/explainability/test_tft_explainer.py

@@ -456,6 +456,112 @@ def test_variable_selection_explanation(self, n_series, mpl_safe_plotting):
             assert isinstance(fig, matplotlib.figure.Figure)
             assert len(fig.get_axes()) == 3
 
+    @pytest.mark.parametrize("n_series", [1, 2])
+    def test_variable_selection_over_time_plotting(self, n_series, mpl_safe_plotting):
+        """Test plotting of per-timestep encoder/decoder variable importance
+        (`TFTExplainer.plot_variable_selection_over_time`)."""
+        model = self.helper_create_model(use_encoders=True, add_relative_idx=True)
+        series, pc, fc = self.helper_get_input(series_option="multivariate")
+        model.fit(series, past_covariates=pc, future_covariates=fc)
+        explainer = TFTExplainer(model)
+        results = explainer.explain(
+            foreground_series=series if n_series == 1 else [series] * 2,
+            foreground_past_covariates=pc if n_series == 1 else [pc] * 2,
+            foreground_future_covariates=fc if n_series == 1 else [fc] * 2,
+        )
+
+        # with `use_encoders=True` and `add_relative_idx=True` on a bivariate target with one past and one
+        # future covariate, the encoder has 9 variables and the decoder has 4 (see `enc_expected`/`dec_expected`
+        # in `test_variable_selection_explanation` above for exactly which ones).
+        n_enc_vars, n_dec_vars = 9, 4
+
+        def _check_plot(**kwargs) -> list[matplotlib.figure.Figure]:
+            figs = explainer.plot_variable_selection_over_time(results, **kwargs)
+            if n_series == 1:
+                figs = [figs]
+            for fig in figs:
+                assert isinstance(fig, matplotlib.figure.Figure)
+                # one subplot for encoder importance over time, one for decoder importance over time
+                enc_ax, dec_ax = fig.get_axes()
+                assert enc_ax.get_title() == "Encoder variable importance over time"
+                assert dec_ax.get_title() == "Decoder variable importance over time"
+            return figs
+
+        # by default, every variable gets its own line in each subplot
+        for fig in _check_plot(show_index_as="relative"):
+            enc_ax, dec_ax = fig.get_axes()
+            assert len(enc_ax.get_lines()) == n_enc_vars
+            assert len(dec_ax.get_lines()) == n_dec_vars
+
+        # `max_nr_components` caps how many variables (lines) are drawn per subplot
+        for fig in _check_plot(show_index_as="relative", max_nr_components=3):
+            enc_ax, dec_ax = fig.get_axes()
+            assert len(enc_ax.get_lines()) == 3
+            assert len(dec_ax.get_lines()) == 3
+
+        # `show_index_as="time"` plots the same data against the actual time index instead of a relative one
+        _check_plot(show_index_as="time")
+
+        with pytest.raises(ValueError, match="`show_index_as` must either be"):
+            _check_plot(show_index_as="invalid")
+
+    @pytest.mark.slow
+    def test_variable_selection_over_time_plotting_stress(self, mpl_safe_plotting):
+        """Stress test `plot_variable_selection_over_time` with a long input/output chunk length and many
+        covariate components -- a long daily `input_chunk_length` and >20 variables are exactly the conditions
+        (raised in https://github.com/unit8co/darts/issues/2685) under which dense/overlapping x-axis date
+        labels and an unreadable legend become a real risk, rather than a theoretical one."""
+        freq = "D"
+        icl, ocl = 60, 20
+        n_pc, n_fc = 15, 15
+        length = icl + ocl + 20
+
+        series = tg.linear_timeseries(length=length, freq=freq)
+        rng = np.random.RandomState(42)
+        pc = TimeSeries(
+            times=series.time_index,
+            values=rng.randn(length, n_pc),
+        )
+        fc_times = tg.linear_timeseries(length=length + ocl, freq=freq).time_index
+        fc = TimeSeries(times=fc_times, values=rng.randn(length + ocl, n_fc))
+
+        model = TFTModel(
+            input_chunk_length=icl,
+            output_chunk_length=ocl,
+            n_epochs=1,
+            add_encoders={"cyclic": {"past": ["month"], "future": ["month"]}},
+            add_relative_index=True,
+            random_state=42,
+            **tfm_kwargs,
+        )
+        model.fit(series, past_covariates=pc, future_covariates=fc)
+        explainer = TFTExplainer(model)
+        results = explainer.explain()
+
+        enc_imp_ot = results.get_encoder_importance_over_time()
+        dec_imp_ot = results.get_decoder_importance_over_time()
+        # comfortably more variables than the default `max_nr_components=10`, and long enough encoder /
+        # decoder windows to previously have produced crowded date labels
+        assert enc_imp_ot.n_components > 20
+        assert dec_imp_ot.n_components > 10
+        assert len(enc_imp_ot) == icl
+        assert len(dec_imp_ot) == ocl
+
+        # default `max_nr_components=10` keeps the plot readable even with 30+ available variables
+        fig = explainer.plot_variable_selection_over_time(results, show_index_as="time")
+
+        enc_ax, dec_ax = fig.get_axes()
+        assert len(enc_ax.get_lines()) == 10
+        assert len(dec_ax.get_lines()) == 10
+
+        # the actual point of the stress test: even over a 60/20-timestep *daily* date range, matplotlib's
+        # date locator keeps the number of rendered x-axis tick labels small, instead of emitting one label
+        # per timestep and producing an unreadable, overlapping axis
+        fig.canvas.draw()
+        for ax in (enc_ax, dec_ax):
+            xticklabels = [t.get_text() for t in ax.get_xticklabels() if t.get_text()]
+            assert len(xticklabels) <= 15, xticklabels
+
     @pytest.mark.parametrize("n_series", [1, 2])
     def test_attention_explanation(self, n_series, mpl_safe_plotting):
         """Test attention (feature importance) explanation results and plotting."""