使用任意数量的输入函数创建绘图功能

时间:2014-09-08 13:18:55

标签: python matplotlib plot

我在绘制一个通用绘图函数时遇到了一些麻烦,该绘图函数绘制了所提供函数的实部和虚部。我想概括我的函数,以便它可以接受任意数量的函数输入,然后用图例绘制它们。因此我应该能够使用:

调用该函数
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)

2 个答案:

答案 0 :(得分:1)

您应该能够在函数声明中使用args或kwargs(有关详细信息,请参阅*args and **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)

答案 1 :(得分:0)

感谢@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)