Creating plotting function with arbitrary number of input functions

Question

I am having a bit of trouble making a general plotting function that plots the real and imaginary parts of a function that is provided. I want to generalise my function so that it can accept an arbitrary number of function inputs and then plot them with a legend. Hence I should be able to call the function using:

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

and the function responds accordingly. What is the most compact way of doing this? Furthermore if the legend could be placed outside the two subplots (since the labels apply to both plots) that would be preferred.

At the moment I just have the following for manually plotting two two functions:

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)

Ok sorry about that I didn't see that post, I posted an answer to my own question in case someone is looking for specifically this. Is that right or should I remove the question all together? — Dipole, Sep 08 '14 at 13:39
You can leave the question. The idea is that the duplicate questions remain, increasing the chances that future users can find a question relevant to their needs, but all the relevant *answers* are concentrated in one place. — chepner, Sep 08 '14 at 13:41

score 1 · Answer 1 · edited May 23 '17 at 11:49

You should be able to make use of args or kwargs in your function declaration (see e.g. *args and **kwargs? for more information). For example, to pass an arbitrary number of functions use:

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))

You could solve the label issue with 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))

which you can call like:

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

Hmm, I guess I should check for comments after editing and before posting my own answer! — radpotato, Sep 08 '14 at 13:42
+1 No thanks, I guess we ended up with the same conclusion! I appreciate the help anyways! — Dipole, Sep 08 '14 at 13:43

score 0 · Answer 2 · answered Sep 08 '14 at 13:38

Thanks to @chepner I got this

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)

Creating plotting function with arbitrary number of input functions

2 Answers2