创建具有任意输入函数数量的绘图函数

Question

我在制作一个通用的绘图函数时遇到了一些麻烦，这个函数需要能够绘制一个给定函数的实部和虚部。我想让这个函数更通用，可以接受任意数量的函数输入，然后把它们绘制出来，并加上图例。因此，我希望能用以下方式调用这个函数：

tester(0.9, func1)
tester(0.9, func1, func2)
tester(0.9, func1, func2, func3, …)

然后这个函数就能相应地工作。有没有什么更简洁的方法来实现这个？另外，如果图例能放在两个子图的外面（因为标签适用于两个图），那就更好了。

目前，我只是用下面的代码手动绘制两个函数：

import numpy as np
import matplotlib.pyplot as plt  

def tester(radius, func1, func2):

    theta = np.linspace(0,2.1*np.pi,1000)
    z = radius*np.exp(1j*theta)
    w1 = func1(z)
    w2 = func2(z)

    fig = plt.figure()
    ax1 = fig.add_subplot(211)
    ax2 = fig.add_subplot(212)

    ax1.plot(theta, np.real(w1), label='%s' % func1)
    ax2.plot(theta, np.imag(w1))

    ax1.plot(theta, np.real(w2), label='%s' % func2)
    ax2.plot(theta, np.imag(w2))

    ax1.legend() 

    ax1.set_ylabel('Real Component')
    ax2.set_ylabel('Imag Component')
    ax2.set_xlabel(r'Polar $\theta$ at $r=$%.2f' % radius)

    plt.show()
    return 0

def func1(z):
    return np.sqrt(z)

def func2(z):
    return np.sqrt(z**2-1)

tester(0.9, func1, func2)

Answer 1

感谢@chepner，我得到了这个

import numpy as np
import matplotlib.pyplot as plt  


def tester(r, **kwargs):

    theta = np.linspace(0,2.1*np.pi,1000)
    z = r*np.exp(1j*theta)

    fig = plt.figure()
    ax1 = fig.add_subplot(211)
    ax2 = fig.add_subplot(212)

    for k,v in kwargs.iteritems():
        ax1.plot(theta, np.real(v(z)), label='%s' % k)
        ax2.plot(theta, np.imag(v(z)))

    ax1.legend()
    ax1.set_ylabel('Real Component')
    ax2.set_ylabel('Imag Component')
    ax2.set_xlabel(r'Polar $\theta$ at $r=$%.2f' % r)
    plt.show()
    return 0

def func1(z):
    return np.sqrt(z)

def func2(z):
    return np.sqrt(z**2-1)

tester(0.9, func1= func1, func2=func2)

Answer 2

你可以在函数声明时使用 args 或 kwargs（想了解更多可以看看这个链接：*args 和 **kwargs?）。比如，如果你想传递任意数量的函数，可以这样做：

def testuser(radius, *args):

    theta = np.linspace(0,2.1*np.pi,1000)
    z = radius*np.exp(1j*theta)

    fig = plt.figure()
    ax1 = fig.add_subplot(211)
    ax2 = fig.add_subplot(212)

    for i, arg in enumerate(args):
        w = arg(z)
        ax1.plot(theta, np.real(w), label="Function %s" % i)
        ax2.plot(theta, np.imag(w))

你可以用 kwargs 来解决标签的问题：

def testuser(radius, **kwargs):
    # Insert rest of code ...
    for i, arg in enumerate(kwargs):
        w = args[arg](z)
        ax1.plot(theta, np.real(w), label=arg)
        ax2.plot(theta, np.imag(w))

然后你可以像这样调用它：

testuser(0.9, function1=func1, function2=func2, plotlabel=functiondeclaration)