3
\$\begingroup\$

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)    
\$\endgroup\$
1
\$\begingroup\$

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.

\$\endgroup\$

Your Answer

By clicking “Post Your Answer”, you agree to our terms of service, privacy policy and cookie policy

Not the answer you're looking for? Browse other questions tagged or ask your own question.