How can I return a matplotlib figure from a function?

Question

I need to plot changing molecule numbers against time. But I'm also trying to investigate the effects of parallel processing so I'm trying to avoid writing to global variables. At the moment I have the following two numpy arrays tao_all , contains all the time points to be plotted on the x-axis and popul_num_all which contains the changing molecule numbers to be plotted on the y-axis.

The current code I've got for plotting is as follows:

for i, label in enumerate(['Enzyme', 'Substrate', 'Enzyme-Substrate complex', 'Product']):
    figure1 = plt.plot(tao_all, popul_num_all[:, i], label=label)
plt.legend()
plt.tight_layout()
plt.show()

I need to encapsulate this in a function that takes the above arrays as the input and returns the graph. I've read a couple of other posts on here that say I should write my results to an axis and return the axis? But I can't quite get my head around applying that to my problem?

Cheers

Answer 1

def plot_func(x, y):
    fig,ax = plt.subplots()
    ax.plot(x, y)
    return fig

Usage:

fig = plot_func([1,2], [3,4])

Alternatively you may want to return ax . For details about Figure and Axes see the docs . You can get the axes array from the figure by fig.axes and the figure from the axes by ax.get_figure() .

Answer 2

In addition to above answer, I can suggest you to use matplotlib animation.FuncAnimation method if you are working with the time series and want to make your visualization better.

You can find the details here https://matplotlib.org/api/_as_gen/matplotlib.animation.FuncAnimation.html