Return the figure plotted by TFT explainer variable selection

[souryadey]

Is your feature request related to a current problem? Please describe.
The method TFTExplainer.plot_variable_selection() does not return the figure or axes where it plots. This means that the figure or axes cannot be used downstream, e.g. for saving them or closing them.

Describe proposed solution
Please return the figure fig or the axes axes from the method. Source code is here.

Describe potential alternatives
A clear and concise description of any alternative solutions or existing features that might solve it.

Additional context
The TFTExplainer.plot_attention() method returns the axes, which is helpful. Similar behavior for TFTExplainer.plot_variable_selection() will be appreciated.

[dennisbader]

Hi @souryadey and thanks for raising this issue. The difference compared to plot_attention() is that for variable selection, we create a new figure per time series that was passed to TFTExplainer.explain() (e.g. a list of TimeSeries as series would mean multiple figures). That is why even if we returned the ax in the current implementation, it would only contain the information of the last series. We could of course return a list of axes / figures in case of multi-series.

If we are okay with that, we can go ahead with a PR. Or could we find a better solution?

In any case, would you like to contribute to this with a PR? :)

[souryadey]

Thank you for your reply. Yes, returning a list of axes / figures would be a great solution. That way they can be used downstream.