Trigonometric Subplots with Matplotlib

Question

The code consists of 4 functions that result in a figure with 3 subplots:

• a sine wave
• a cosine wave
• and the superposition of both (a bichromatic wave).

I have shown in 4 ways, the first two are basically plot the same way, but the generation of the waves are different. I also show different methods, by adding subplots directly using plt.subplots(3, 1), also by adding one at a time using plt.subplot(3,1,i), and the last function uses the sine and cosine plot data to create the bichromatic wave.

I would like to know what to improve for the code, and also what to improve for the methods in using matplotlib ..?

---

import matplotlib.pyplot as plt
import math
pi = math.pi


ex_1 = "1. Using lists (comprehension) to contain A*sin(a*x), B*cos(b*x), and the superposition. \
Using plt.subplots to created the plots."

def trig_subplots_1(x=[], amplitude=(1, 1), frequency=(1, 1)):

    sine = [amplitude[0]*math.sin(frequency[0]*i) for i in x]
    cosine = [amplitude[1]*math.cos(frequency[1]*i) for i in x]
    superposition = [s+c for s,c in zip(sine, cosine)]

    fig, ax = plt.subplots(3,1)

    ax[0].plot(x, sine)
    ax[0].set_title('Sine wave')

    ax[1].plot(x, cosine)
    ax[1].set_title('Cosine wave')

    ax[2].plot(x, superposition)
    ax[2].set_title('Bichromatic wave')

    plt.tight_layout()
    plt.show()


ex_2 = "2. Generate A*sin(a*x), B*cos(b*x), and the superposition by list.append in a for-loop. \
Using plt.subplots to created the plots."

def trig_subplots_2(x=[], amplitude=(1, 1), frequency=(1, 1)):

    sine, cosine, superposition = [], [], []
    for i in x:
        sine.append( amplitude[0]*math.sin(frequency[0]*i) )
        cosine.append( amplitude[1]*math.cos(frequency[1]*i) )
        superposition.append( sine[-1] + cosine[-1] )

    fig, ax = plt.subplots(3,1)

    ax[0].plot(x, sine)
    ax[0].set_title('Sine wave')

    ax[1].plot(x, cosine)
    ax[1].set_title('Cosine wave')

    ax[2].plot(x, superposition)
    ax[2].set_title('Bichromatic wave')

    plt.tight_layout()
    plt.show()


ex_3 = "3. Similar as no.1, but creating each subplot axes by using plt.subplot(nrows, ncols, index)."

def trig_subplots_3(x=[], amplitude=(1, 1), frequency=(1, 1)):

    sine = [amplitude[0]*math.sin(frequency[0]*i) for i in x]
    cosine = [amplitude[1]*math.cos(frequency[1]*i) for i in x]
    superposition = [s+c for s,c in zip(sine, cosine)]

    plt.subplot(3,1,1)
    plt.title('Sine wave')
    plt.plot(x, sine)

    plt.subplot(3,1,2)
    plt.title('Cosine wave')
    plt.plot(x, cosine)

    plt.subplot(3,1,3)
    plt.title('Bichromatic wave')
    plt.plot(x, superposition)

    plt.tight_layout()
    plt.show()


ex_4 = "4. Similar as no.1, but the superposition wave is generated after plotting the sine and cosine then and \
get the y values from each plot using .get_ydata()."

def trig_subplots_4(x=[], amplitude=(1, 1), frequency=(1, 1)):

    sine = [amplitude[0]*math.sin(frequency[0]*i) for i in x]
    cosine = [amplitude[1]*math.cos(frequency[1]*i) for i in x]

    fig, ax = plt.subplots(3,1)

    sine_plot = ax[0].plot(x, sine)
    ax[0].set_title('Sine wave')

    cosine_plot = ax[1].plot(x, cosine)
    ax[1].set_title('Cosine wave')

    superposition = sine_plot[0].get_ydata() + cosine_plot[0].get_ydata()
    ax[2].plot(x, superposition)
    ax[2].set_title('Bichromatic wave')

    plt.tight_layout()
    plt.show()



x = [i*0.01*pi for i in range(1000)]
amplitude = (1, 1)
frequency = (1, 0.75)

print(ex_1)
trig_subplots_1(x, amplitude, frequency)
print(ex_2)
trig_subplots_2(x, amplitude, frequency)    
print(ex_3)
trig_subplots_3(x, amplitude, frequency)    
print(ex_4)
trig_subplots_4(x, amplitude, frequency)    

Answer 1

Your code is repeating a little bit. For example,

ax[0].plot(x, sine)
ax[0].set_title('Sine wave')

ax[1].plot(x, cosine)
ax[1].set_title('Cosine wave')

ax[2].plot(x, superposition)
ax[2].set_title('Bichromatic wave')

This can be rewritten using a helper function (or alternatively a for-loop), which may look something like this:

def plot_subroutine(axis, xdata, ydata, title):
    axis.plot(xdata, ydata)
    axis.set_title(title)

y1 = ... # sine 
y2 = ... # cosine
y3 = ... # bichromatic
yn = (y1, y2, y3)

title1 = 'Sine Wave'
title2 = 'Cosine Wave'
title3 = 'Bichromatic Wave'
titles = (title1, title2, title3)

fig, ax = plt.subplots(3,1)

Then instead of calling each plot subroutine individually (as below)

plot_subroutine(ax[0], x, yn[0])
plot_subroutine(ax[1], x, yn[1])
plot_subroutine(ax[2], x, yn[2])

You can instead do something like

if len(ax) == len(yn):
    for i in range(len(ax)):
        plot_subroutine(ax[i], x, yn[i])
plt.show()

The for-loop has the advantage that you can later change the number of output plots without changing too much of your original code (if say, you wanted to output 2 or 4 plots, or perhaps change the ydata to something entirely different).

This may or may not be relevant, but it may be nice to associate each individual plot with a specific legend label, marker style, and/or color, all of which can be handled using matplotlib. This can be especially useful for overlayed plots (especially considering that your data consists of sin(x) and cos(x) over the same x-interval.