I want to draw a square wave in python, but i cant do that, because I cant check the value of x in f function:

import matplotlib.pyplot as plt
import numpy as np
import math


I = float(input())
t = float(input())
T = float(input())


def f(x):
    if x % T <= (T / 2):
        return I
    else:
        return -I


x = np.linspace(0, t, 1000)
plt.plot(x, f(x))
plt.xlim(0, t)
plt.ylim(-10, 10)
plt.axhline(color="black")
plt.show()

how can i do that?

CodePudding user response：

I guess you are looking for numpy.where:

import matplotlib.pyplot as plt
import numpy as np

ampl = 2
wl = 3
xmax = 9

def square_wave(x, wave_length, amplitude):
    return np.where(
        x % wave_length <= wave_length / 2, amplitude, -amplitude
    )

x = np.linspace(0, xmax, 100)
plt.plot(x, square_wave(x=x, wave_length=wl, amplitude=ampl))
plt.axhline(color="black")
plt.show()

CodePudding user response：

The problem you have, it that you are applying f on the whole array x. Which is a good idea generally speaking, since the whole point of numpy is to fasten computation by doing exactly that (python is a very slow language, but numpy is not written in python. So every batch of computation that numpy can do for you is way faster).

So what you experience here is the broadcasting and vectorization capabilities of numpy: a code, f, that you, obviously, wrote with scalar in mind, almost works with argument x being an array.

x%T is an array whose each element is a x[i]%T for i in range(len(x)).
x%T<=T/2 is an array of booleans whose each element is true or false, depending on whether x[i]%T<=T/2 or not.

And there it stop workings. Because you can use an array of boolean as a condition to a if. A condition of a if must be a single truth value: either it is true, or it isn't. And what you have is an array of that, some true, some false. Hence the error you got:

ValueError: The truth value of an array with more than one element is ambiguous.

Now you have 3 solutions, from the easiest to the one needing the most adaptation, but also from the slowest to the fastest.

1. You could just call f on scalar, as intended. So build an array y such as y[i]=f(x[i]) and pass that to plot

y = np.zeros_like(x)
for i in range(len(x)):
    y[i]=f(x[i])
plt.plot(x,y)

Your f function, as is, works well. It is just the calling that was wrong. It is made to be called with scalar, you called it with vector. With this for loop, I correct that, and call it (n times) with scalar.

2. You could let numpy do that with so called vectorize function.

vectorize turns a function meant to be called on a scalar (or other things, but what I am interested in is this case) into a function that can be called with arrays of any shape. It returns then an array of the same shape, whose all elements are f(z) for z being any element of x

So in your case you could

plt.plot(x, np.vectorize(f)(x))

But that is a false good idea. All that it does it the same as my 1st version, but it does the for loop for you.

Still, it is faster. And still quite easy to write (at least in such a canonical case. Using vectorize can be quite a challenge sometimes).

3. Vectorize

But most of the times, and your case is no exception, you can vectorize the function yourself. That is rewrite f to process not just one scalar but a vector of scalars. We have seen in preamble that half of the job is already done: you have computed an array of booleans telling you if a return element should be I or -I.

For example

def f(x):
    return np.where(x%T<=T/2, I, -I)

plt.plot(x, f(x))

does what you want. x%T<=T/2 is the array of booleans we mentionned before. And np.where returns an array of the same shape whole values are taken from 2nd or 3rd argument.

Another method, may more common in maths could be

def f(x):
    return (x%T<=T/2)*2*I -I

(based on the fact that in numpy, a boolean, used in arithmetic operation, is 0 or 1)

Timing

The 1st solution (using a for loop), on my computer, takes 308 ms per run.

The 2nd one (using vectorize), 159 ms per run. So twice faster. Because for loop is done in numpy not in python. But, still, the code inside the for loop is still python code.

The 3rd one (np.where or the arithmetic version), takes 13 ms per run.

So, you see, it is not really an option. You need to use the 3rd one. Generally speaking, when using numpy, you must avoid for loop at all cost, and try to think to each operation as an operation on an array, not on elements that are in it. I provided you 3 options to help understand what happened, and what needed to be done. But in reality, only of those 3 solutions is valid (even if you don't care of computation time, it is better to start thinking that way when using numpy)