Visualization functions should optionally accept `ax` as argument

[ma-sadeghi]

(for more granular control)

[jgostick]

sem and xray current accept axis. Are you suggesting it would be more convenient to shorten it? Or are you suggesting this argument should be added to other functions? As for the former, I think ax implies that you're passing in a mpl axis object, since fig, ax = plt.subplots() is now pretty universally used. If the latter, which ones did you have in mind?

[ma-sadeghi]

Yeah, I meant an mpl axis. Specifically, I'm proposing that all visualization functions accept an optional argument ax, which if not given will be created internally, and return the ax object rather than the figure object. I've seen this pattern in many places (arviz for example).

The benefit is that you have more control, e.g., you can resize the figure: fig, ax = plt.subplots(figsize=...), or maybe you want to overlay some other stuff on the same axis, or maybe you have a multi-row/column figure and you want a particular entry to be painted by the function, etc.

You can do all this using the current approach too, it's just extra typing.