Functions that plot should take an `ax` argument

🚀 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]

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

[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