Question

A question from a complete Python novice.

I have a column array where I need to force certain values to zero depending on a conditional statement applied to another array. I have found two solutions, which both provide the correct answer. But they are both quite time consuming for the larger arrays I typically need (>1E6 elements) – also I suspect that it is poor programming technique. The two versions are:

from numpy import zeros,abs,multiply,array,reshape

def testA(y, f, FC1, FC2):
    c = zeros((len(f),1))
    for n in xrange(len(f)):
        if abs(f[n,0]) >= FC1 and abs(f[n,0]) <= FC2:
            c[n,0] = 1.
    w = multiply(c,y)
    return w

def testB(y, f, FC1, FC2):
    z = [(abs(f[n,0])>=FC1 and abs(f[n,0])<=FC2) for n in xrange(len(f))]
    z = multiply(array(z,dtype=float).reshape(len(f),1), y)
    return z

The input arrays are column arrays as this matches the post processing to be done. The test can be done like:

>>> from numpy.random import normal as randn
>>> fs, N = 1.E3, 2**22
>>> f = fs/N*arange(N).reshape((N,1))
>>> x = randn(size=(N,1))
>>> w1 = testA(x,f,200.,550.)
>>> z1 = testB(x,f,200.,550.)

On my laptop testA takes 18.7 seconds and testB takes 19.3 – both for N=2**22. In testB I also tried to include “z = [None]*len(f)” to preallocate as suggested in another thread but this doesn’t really make any difference.

I have two questions, which I hope to have the same answer:

1. What is the “correct” Python solution to this problem?
2. Is there anything I can do to get the answer faster?

I have deliberately not used any time at all using compiled Python for example – I wanted to have some working code first. Hopefully also something, which is good Python style. I hope to be able to get the execution time for N=2**22 below two seconds or so. This particular operation will be used many times so the execution time does matter.

I apologize in advance if the question is stupid – I haven’t been able to find an answer in the overwhelming amount of not always easily accessible Python documentation or in another thread.

Answer 1

use bool array to access elements in array y:

def testC(y, f, FC1, FC2):
    f2 = abs(f)
    idx = (f2>=FC1) & (f2<=FC2)
    y[~idx] = 0
    return y