Is it possible to plot within user-defined function with python and matplotlib?

Question

What I want to do is define a function which contain plotting sentences. Like this:

import matplotlib.pyplot as plt

def myfun(args, ax):
    #...do some calculation with args
    ax.plot(...)
    ax.axis(...)

fig.plt.figure()
ax1=fig.add_subplot(121)
ax2=fig.add_subplot(122)
para=[[args1,ax1],[args2,ax2]]
map(myfun, para)

I found that the myfun is called. If I add plt.show() in myfun, it can plot in the correct subplot, but nothing in the other one. And, if plt.show() is added in the end, nothing but two pairs of axis are plotted. I think the problem is that the figure is not transferred to the main function successfully. Is it possible to do something like this with python and matplotlib? Thanks!

Answer 1

Function that is called via map should have only one parameter.

import matplotlib.pyplot as plt

def myfun(args):
    data, ax = args
    ax.plot(*data)

fig = plt.figure()
ax1 = fig.add_subplot(121)
ax2 = fig.add_subplot(122)
para = [
    [[[1,2,3],[1,2,3]],ax1],
    [[[1,2,3],[3,2,1]],ax2],
]
map(myfun, para)
plt.show()

If you want to keep your function signature, use itertools.starmap.

import itertools
import matplotlib.pyplot as plt

def myfun(data, ax):
    ax.plot(*data)

fig = plt.figure()
ax1 = fig.add_subplot(121)
ax2 = fig.add_subplot(122)
para = [
    [[[1,2,3],[1,2,3]],ax1],
    [[[1,2,3],[3,2,1]],ax2],
]
list(itertools.starmap(myfun, para)) # list is need to iterator to be consumed.
plt.show()