CODEWITHSUNDEEP
pythonnumpymatplotlib
rmiller415
rmiller415

Reputation: 43

Using List Comprehension to Make an Array to Plot

I was asked to write some code to make a triangle wave and plot it. I originally wrote it using if/elif logic inside of a for loop. This works well and gives me the correct results, but it's so slow. I rewrote the code using list comprehension, but the plot ends up being misshapen and incorrect at times. I think the issue is the addition and subtraction in my code. Here is the code that worked:

import matplotlib.pyplot as plt
import numpy as np

f = np.linspace(-4,4,8000)
t = np.linspace(-4,4,8000)
for i in range(f.size):
    if f[i] < -3:
        f[i] = (1 - np.sqrt((f[i] + 4)**2))
    elif -3 <= f[i] < -1:
        f[i] = (1 - np.sqrt((f[i] + 2)**2))
    elif -1<= f[i] < 1:
        f[i] = (1- np.sqrt(f[i]**2))
    elif 1 <= f[i] < 3:
        f[i] = (1 - np.sqrt((f[i] - 2)**2))
    elif f[i] >= 3:
        f[i] = (1 - np.sqrt((f[i] - 4)**2))
    i += 1
    
plt.plot(t,f)
plt.show()

The correct plot:

The correct plot:

The rewritten code is:

import matplotlib.pyplot as plt
import numpy as np

f = np.linspace(-4,4,8000)
t = np.linspace(-4,4,8000)
f = [(1 - np.sqrt((i + 4)**2)) if i <= -3 else i for i in f]
f = [(1 - np.sqrt((i + 2)**2)) if -3 <= i <= -1 else i for i in f]
f = [(1 - np.sqrt(i**2)) if -1 <= i <= 1 else i for i in f]
f = [(1 - np.sqrt((i - 2)**2)) if 1 <= i <= 3 else i for i in f]
f = [(1 - np.sqrt((i - 4)**2)) if i >= 3 else i for i in f]

plt.plot(t,f)
plt.show()

The problem image I get is: enter image description here

It seems like the problem is the 1 - sqrt stuff. I'm not sure how to correct this problem.

Edit for clarification: I am plotting the triangle wave, but this is more about the generation of the discrete points to learn about scipy's discrete fft package for a scientific computing course. Generating the plot only verifies that I have good points to feed into the fft. Also, after thinking about it, my concern with performance was unnecessary. I think taking the original path I was on will get the results I wanted and future data points will be generated for me by instrumentation, so if it takes 2 seconds instead of .2 seconds to make it work I will be fine with that.

Upvotes: 1

Views: 296

Answers (2)

Luis Garcia
Luis Garcia

Reputation: 26

You have a problem with the first two array compression and with the start number.

import matplotlib.pyplot as plt
import numpy as np

f = np.linspace(-4,3.9999,8000)
t = np.linspace(-4,3.9999,8000)
f = [(1 - np.sqrt((i + 5)**2)) if i <= -3 else i for i in f]
f = [(1 - np.sqrt((i + 3)**2)) if -3 <= i <= -1 else i for i in f]
f = [(1 - np.sqrt(i**2)) if -1 <= i <= 1 else i for i in f]
f = [(1 - np.sqrt((i - 2)**2)) if 1 <= i <= 3 else i for i in f]
f = [(1 - np.sqrt((i - 4)**2)) if i >= 3 else i for i in f]
plt.plot(t,f)
plt.show()

By the way. If you want to plot a triangle wave, you don't need to plot each point, you can create an array from -4 to 4(inclusive) and apply just one array compression:

import matplotlib.pyplot as plt

x = range(-4, 5)
y = range(-4, 5)
y = [0 if e%2 == 1 else 1 for e in y]
plt.plot(x,  y)
plt.show()

Plot image

Upvotes: 0

Cameron Chandler
Cameron Chandler

Reputation: 418

I would suggest that your first method is preferable. You may like to consider using np.abs() in preference of np.sqrt(x**2) as this will slightly speed things up. Alternatively, you can just reduce the 8000 to a more reasonable number like 100. On my computer your code runs in 200ms, how much faster do you need it?

You may like to consider the following code, adapted from Triangle wave shaped array in Python

from scipy import signal
import numpy as np
import matplotlib.pyplot as plt
t = np.linspace(-4, 4, 8000)
f = (1 - signal.sawtooth(2 * np.pi * 0.5 * t, 0.5))/2
plt.plot(t, f)

It runs 10x as fast

Upvotes: 1

Related Questions