[draft] Passing ax to some plot functions

[balopat]

Take this PR as a draft regarding ax in the plotting API.

A bug: ax is unused currently in plot_tree_rubberband - so I fixed that.

I also passed ax in

• plot_tree_circuit
• plot_tree_flat
• plot_tree_contractions

The logic in all cases is that ax takes precedence over the passed in figsize if available and fig =None is returned.

I find this useful as it helps to create multiplot subplots like in #49.

Open questions:

• would you like to have ax consistently across all plot functions?
• would the figsize / fig logic described above be sufficient?
• I don't like that rubberband shows the plot and there is no way to close it or force it to close (I'm working on WSL and the command line, not in a notebook env and it keeps popping up) - can we modify that too maybe by default turning show() off when ax is passed?
• I guess these would need some testing too ideally, though plotting testing can get a bit hairy so maybe it's not worth it, let me know your thoughts on that too

[jcmgray]

Thanks for the PR! I think yes generally desirable to be able to pass in an ax kwarg, happy to have that added/fixed on any plotting functions.

Yes I agree the logic should be that if an ax is supplied then fig kwargs are ignored and fig=None is returned. The show_and_close decorator should then not automatically pyplot.show(fig) the plot.

At the moment, fig is retrieved from Drawing.fig which actually retrieves it from ax.figure so show_and_close always sees a figure returned and displays it. One could add the ax logic to show_and_close or simply handle setting fig=None in the plotting functions that use Drawing.

Regarding testing, yes ideally just any basic test that shows the functions run without erroring (i.e. don't worry about explicit tests that plot output is as expected etc.). However, I wouldn't consider it required.