[FEATURE] Remove plt.close() from the plotting code or otherwise expose the figure / axes handles

[AA-TK]

Problem

plt.close() prevents subsequent code from interacting with the plot and re-styling the figures.

Proposed Solution

It would be preferable (to me at least) if the plotting code followed standard convention of returning handles to figure and axes.

Alternatives

I tried hacking that line of code out on my end and it didn't result in the figure being accessible using plt.gca() in subsequent code.

Additional Context

No response