Plot a piecewise-defined function

Question

I would like to plot the following function:

$$ \begin{align} \Lambda(\delta\tau) &\equiv\ \chi(\delta\tau, 0) = \frac{1}{T_i} \int_0^{T_i} a(t_0+t')a(t_0+t'+\delta\tau) dt' \\ &= \begin{cases} 1 - \frac{\left|\delta\tau\right|}{\tau_c}, &\left|\delta\tau\right| \le \tau_c(1+\frac{\tau_c}{T_i}) \\ -\frac{\tau_c}{T_i}, &\left|\delta\tau\right| \gt \tau_c(1+\frac{\tau_c}{T_i}) \\ \end{cases} \end{align} $$ which represents a simple triangular shape. There is a conditional statement. My implementation uses a for loop as follows:

def waf_delay(delay_increment):
 for d in delay_increment:
 if np.abs(d) <= delay_chip*(1+delay_chip/integration_time):
 yield 1 - np.abs(d)/delay_chip
 else:
 yield -delay_chip/integration_time;
integration_time = 1e-3 # seconds
delay_chip = 1/1.023e6 # seconds
x=np.arange(-5.0, 5.0, 0.1) 
y=list(waf_delay(x))
plt.plot(x, y)
plt.show()

Is there a more correct way to transform an array based on a condition rather than just looping through it? Instead of having something like this:

def f(x_array):
 for x in x_array:
 if np.abs(x) <= 3:
 yield 1 - x/3
 else:
 yield 0
x=np.arange(-5.0, 5.0, 0.1) 
y=list(f(x))
plt.plot(x, y)
plt.show()

I would like to write something like this:

def f(x):
 if np.abs(x) <= 3:
 yield 1 - x/3
 else:
 yield 0
x=np.arange(-5.0, 5.0, 0.1) 
plt.plot(x, f(x))
plt.show()

that could take an array.

Answer 1

There are two ways to solve this problem. The first one is numpy.where, which can take two arrays and it will choose from one wherever a condition is true and from the other wherever it is false. This only works if your piecewise function has only two possible states (as is the case here):

def waf_delay(delays):
 return np.where(np.abs(delays) <= delay_chip*(1+delay_chip/integration_time),
 1 - np.abs(delays)/delay_chip,
 -delay_chip/integration_time)

Another more general possibility is to use numpy.piecewise, but that is probably overkill here:

def f1(x):
 return 1 - np.abs(x)/delay_chip
def f2(x):
 return -delay_chip/integration_time
cut_off = delay_chip*(1+delay_chip/integration_time)
y = np.piecewise(x, [np.abs(x) <= cut_off, np.abs(x) > cut_off], [f1, f2])

Note that in both cases no for d in delays is needed, because all functions used are vecotorized (the basic arithmetic operations for numpy arrays are and so is numpy.abs).

• Blasco: Very clean. This is exactly what I was expecting to be possible to do but I wasn't sure how. Do you have any recommendation on where to learn this algorithms?
• Graipher: @WooWapDaBug Looking at SO questions/answers about numpy and reading the documentation helped me to learn a lot about neat tricks. Also having to use it somewhere obviously helps, but I don't know of any tutorial that gets into the advanced stuff in a good way...
• Blasco: Graipher Very true. I try my best to learn and to contribute from SO. Thank you very much for sharing your knowledge. It is very valuable