Using List Comprehension to Make an Array to Plot

Question

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

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.

this is going down a bad (complex and error prone) path. 1) try to programatically set your ranges. 2) you're mixing list comprehensions (super slow) with numpy (faster, separate package). That's ok, but you would be better served by using one or the other — anon01
– anon01, Commented Oct 4, 2020 at 5:38

Cameron Chandler · Accepted Answer · 2020-10-04 05:36:49Z

1

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

answered Oct 4, 2020 at 5:36

Cameron Chandler

4183 silver badges10 bronze badges

Sign up to request clarification or add additional context in comments.

Comments

Luis Garcia · Accepted Answer · 2020-10-04 06:23:48Z

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

Collectives™ on Stack Overflow

Using List Comprehension to Make an Array to Plot

2 Answers 2

Comments

Comments

Your Answer

Linked

Hot Network Questions

Collectives™ on Stack Overflow

2 Answers 2

Comments

Comments

Your Answer

Sign up or log in

Post as a guest

Linked

Related