Functions that plot should take an ax argument

[davidgilbertson]

🚀 Feature

Currently, functions like trainer.tuner.lr_find().plot() will generate a new Matplotlib figure and axes. I believe best practice is for such functions to accept an ax so that the user may pass in their own ax.

Motivation

1. Consistency with other packages. E.g. Pandas plot, Statsmodels plot functions, etc.
2. This allows a user to have a figure with multiple axes, and fill one of those axes with the output of, say, lr_find().plot() (not super compelling in the case of LR finder)
3. I use a Matplotlib wrapper that adds functionality, like reusing the same window over multiple runs, scroll-to-zoom, and tooltips-on-hover. But this doesn't work if Lightning creates its own figure and axes each time.

Pitch

Just add ax as an optional arg to any function that plots, which I believe is just the LR finder ATM.

Since the function returns a figure, you can pluck that from the ax if provided, so it wouldn't be a breaking change. Something like:

if ax is None:
    import matplotlib.pyplot as plt

    fig, ax = plt.subplots()
else:
    fig, ax = ax, ax.figure

Alternatives

Can't think of any.

Additional context

There's a page somewhere in the Matplotlib docs that describes general best practices for writing functions that produce charts but I can't find it!

cc @Borda @akihironitta @rohitgr7

[awaelchli]

@rohitgr7 What do you think, should we add this? @davidgilbertson are you interested in contributing this?

[davidgilbertson]

I've got a full schedule at the moment, sorry.

Actually after I logged this I remembered that I've monkey-patched pyplot.subplots so that even if a package doesn't accept an axes object, it will use the one I want anyway. So for me there's no need, but I still think it's a good practice in general.

[bipinKrishnan]

@awaelchli I am new to the lightning codebase, can I work on this issue? Just to confirm, the solution would be to change this line:
https://github.com/Lightning-AI/lightning/blob/d5003b1c07fda783f651a732c86ad48656be42c1/src/pytorch_lightning/tuner/lr_finder.py#L150

to this(as suggested by @davidgilbertson):

if ax is None:
    fig, ax = plt.subplots()
else:
    fig, ax = ax.figure, ax

[davidgilbertson]

Oh, I just realised my initial snippet got the wrong order, good catch @bipinKrishnan . Actually it can just be:

if ax is None:
    fig, ax = plt.subplots()
else:
    fig = ax.figure

And in case you didn't know, this parameter can be typed as ax: Axes using from matplotlib.axes import Axes

[awaelchli]

Yes @bipinKrishnan, feel free to open a PR and we will be very happy to review it!