Mus
Mus

Reputation: 397

Plot multiple functions with the same properties in matplotlib

I try to visualise multiple function coming from the single parameter. I need to do some kind of loop over parameter. I would like to do assign the same colour, legend etc to all plotted functions of given parameter.

The problem is that all plots I try, matplotlib assign different colours and always give a label.

Basicaly I would like to achive something like:

import numpy as np
import matplotlib.pyplot as plt
def plot2():    
    fig, ax = plt.subplots()
    x = np.arange(0,10,0.1)
    ax.plot(x,1*np.sin(x),'b-')
    ax.plot(x,1*np.cos(x),'b-',label='trig a={}'.format(1))
    ax.plot(x,2*np.sin(x),'g-')
    ax.plot(x,2*np.cos(x),'g-',label='trig a={}'.format(2))
    ax.plot(x,3*np.sin(x),'r-')
    ax.plot(x,3*np.cos(x),'r-',label='trig a={}'.format(3))
    ax.legend()

but with function like:

def plotTrig():
    fig, ax = plt.subplots()
    x = np.arange(0,10,0.1)
    for a in [1,2,3]:
        ax.plot(x,a*np.sin(x),x,a*np.cos(x),label='trig a={}'.format(a))
    ax.legend()

above is only simplified example. In practice I have much more functions and parameters, so solution with looping over colors is not very helpful

Upvotes: 2

Views: 3309

Answers (1)

StefanS
StefanS

Reputation: 1770

I think now I understand what you want. You'll never run out of colors, because matplotlib supports a wide range of color definitions. Any legal HTML-name, any RGB-triple, ...

I'm don't know how to conditionally set the label for artists, so that part of the following (the if) is a hack that can be improved by someone with more knowledge about the inner working of matplotlib.

import numpy as np
import matplotlib.pyplot as plt

def my_sin(x, a):
    return a * np.sin(x)

def my_cos(x, a):
    return a * np.cos(x)

def my_tanh(x, a):
    return np.tanh(x / a - 1)

def plotTrig(x, data, colors, parameters):
    fig, ax = plt.subplots()
    for ind, a in enumerate(parameters):
        for name, func in data.iteritems():
            if (name == 'sin'):  # or any other
                ax.plot(x, func(x, a), '-',
                        color=colors[ind],
                        label='trig a={}'.format(a))
            else:
                ax.plot(x, func(x, a), '-',
                        color=colors[ind])
    ax.legend()


if __name__ == '__main__':
    # prepare data
    x = np.arange(0,10,0.1)
    data = {}  # dictionary to hold the values
    data['sin'] = my_sin
    data['cos'] = my_cos
    data['tanh'] = my_tanh
    # list to hold the colors for each parameter
    colors = ['burlywood', 'r', '#0000FF', '0.25', (0.75, 0, 0.75)]
    # parameters
    parameters = [1, 2, 3, 4, 5]
    plotTrig(x, data, colors, parameters)
    plt.show()

The idea is to put the different functions in a container, so we can iterate over them (a list would have worked as well) and then use the same color for every function, but different color for each parameter. The label is only added to one function with the hacky if statement.

I could have done it a lot simpler if I'd just set the dictionary values to the results of the functions:

data['sin'] = np.sin(x)

and then plot with

ax.plot(x, a * func, '-',...

to reproduce your example, but then your parameters could only be applied to the result of the function. This way you can use them in any way you can express as a function.

Result: Result

Upvotes: 3

Related Questions