Substitute deep loops in Python

Question

Unfortunately, I often have while loops in my Python code that cause my programs to slow down significantly.

The following is an example of a while loop (shape = (1000,1000,3)):

i = 0
j = 0
while i < arr.shape[0]:
    while j < arr.shape[1]:
        if arr[i,j,0] <= 5 and arr[i,j,0] > 0:
            arr[i,j,:] = 1
        else:
            arr[i,j,:] = 0

        j = j + 1
    j = 0
    i = i + 1

-->

def f(x):
        return 1 if x <= 5and x > 0 else 0

f = np.vectorize(f)
arr= f(arr)

EDIT: another while loop

i = 0
j = 0
while i < arr1.shape[0]:
    while j < arr1.shape[1]:

        if arr1[i, j, 0] == 0 and arr1[i, j, 1] == 0 and arr1[i, j, 2] == 0:
            arr1[i, j, :] = arr1[i, j, :]
        else:
            arr1[i, j, :] = arr2[i, j, :]

        j = j + 1
    j = 0
    i = i + 1

Is there a way to speed things up? I'm not sure how.

Since you tagged numpy you can look [here](https://stackoverflow.com/questions/7701429/efficient-evaluation-of-a-function-at-every-cell-of-a-numpy-array) for your answer or hints — Arpit Solanki, Jan 30 '18 at 14:30
this doesn't answer your question. but if you are stuck with loops (and even if you aren't) consider using `numba` to JIT-compile. you will likely get large speed-ups. — jpp, Jan 30 '18 at 14:40
`arr[i,j,:] = int(0 < arr[i,j,0] <= 5)` should speed it up a little (chained comparison + direct bool assignment — Jean-François Fabre, Jan 30 '18 at 14:51
Maybe have a look at [numpy.where](https://docs.scipy.org/doc/numpy/reference/generated/numpy.where.html) — kvantour, Jan 30 '18 at 15:03
An example using numba on such tasks https://stackoverflow.com/a/45403017/4045774 — max9111, Jan 30 '18 at 15:28
`for i in range(arr.shape[0]):` is the preferred way of looping over a range of indices. No difference in speed, just cleaner code. — hpaulj, Jan 30 '18 at 17:33
@jp_data_analysis, @max9111, I don't think we should be suggesting `numba` at this stage. The OP needs to learn to work with basic array masking and indexing first. The use of `while` loops suggests he's Python novice. — hpaulj, Jan 30 '18 at 17:37
@hpaulj, fair point. it should, as you say, be a last resort. BUT after you have vectorised, i still feel it can be used for benchmarking. — jpp, Jan 30 '18 at 17:39

lightalchemist · Accepted Answer · 2018-03-04T16:24:17.520

EDIT In my hurry, my previous answer was incorrect. Thanks to @gboffi for pointing it out.

Part 1

def original(arr):
    i = 0
    j = 0
    while i < arr.shape[0]:
        while j < arr.shape[1]:
            if arr[i,j,0] <= 5 and arr[i,j,0] > 0:
                arr[i,j,:] = 1
            else:
                arr[i,j,:] = 0

            j = j + 1
        j = 0
        i = i + 1

    return arr

def vectorized(arr):
    mask = (0 < arr[:, :, 0]) & (arr[:, :, 0] <= 5)
    arr[mask] = 1
    arr[~mask] = 0
    return arr

def vectorized2(arr):
    """Works only if assigning 0 and 1s"""
    mask = (0 < arr[:, :, 0]) & (arr[:, :, 0] <= 5)
    mask = np.dstack([mask] * arr.shape[2])
    return mask.astype(np.float32)

Benchcmark

The following uses arr from the next benchmark. I did not have time to perform extensive tests on this so I suggest that you run the vectorized versions below and use np.allclose to compare the result from your original code for various cases of arr that you can verify.

In [101]: %timeit v1 = vectorized(arr.copy())
1000 loops, best of 3: 312 µs per loop

In [102]: %timeit v2 = original(arr.copy())
100 loops, best of 3: 10.4 ms per loop

In [103]: np.allclose(v1, v2)
Out[103]: True

In [108]: %timeit v3 = vectorized2(arr.copy())
10000 loops, best of 3: **83.3 µs** per loop

In [110]: v3 = vectorized2(arr.copy())

In [111]: np.allclose(v1, v3)
Out[111]: True

Part 2

def vectorized(arr1, arr2):
    mask = np.all(arr1 == 0, axis=2)
    mask = np.dstack([mask] * arr1.shape[2])
    return np.where(mask, arr1, arr2)

def original(arr1, arr2):
    i = 0
    j = 0
    while i < arr1.shape[0]:
        while j < arr1.shape[1]:

            if arr1[i, j, 0] == 0 and arr1[i, j, 1] == 0 and arr1[i, j, 2] == 0:
                arr1[i, j, :] = arr1[i, j, :]
            else:
                arr1[i, j, :] = arr2[i, j, :]

            j = j + 1

        j = 0
        i = i + 1

    return arr1

Benchmark for second loop

# Prepare data
m = 100
n = 100
d = 3
np.random.seed(0)
arr = np.random.randint(0, 11, size=(m, n, d))

true_mask = np.random.randint(0, 2, size=(m, n, 1), dtype=np.bool)
true_mask = np.dstack([true_mask] * d)

arr[true_mask] = 0

arr1 = arr.copy()
arr2 = -1 * np.ones_like(arr1)

In [84]: v1 = vectorized(arr1, arr2)

In [85]: v2 = original(arr1, arr2)

In [86]: np.allclose(v1, v2)
Out[86]: True

In [87]: %timeit v1 = vectorized(arr1, arr2)
1000 loops, best of 3: 284 µs per loop

In [88]: %timeit v2 = original(arr1, arr2)
100 loops, best of 3: 12.6 ms per loop

Maybe later. I'm rushing off somewhere now. If after 2 hrs no one has solved it I will answer it. Bye. — lightalchemist, Jan 30 '18 at 15:16
@freddykrueger I've updated my answer with benchmarks and fixed a bug in the original answer. — lightalchemist, Jan 30 '18 at 19:39

score 0 · Answer 2 · answered Jan 30 '18 at 15:53

Let's generate some test data

In [61]: np.random.seed(0) ; a = np.random.randint(10, size=(3,4,2))-3

In [62]: a
Out[62]: 
array([[[ 2, -3],
        [ 0,  0],
        [ 4,  6],
        [ 0,  2]],

       [[-1,  1],
        [ 4,  3],
        [ 5,  5],
        [-2,  3]],

       [[ 4,  4],
        [ 5, -2],
        [ 2,  6],
        [ 5,  6]]])

According to the example in the question, each row whose 1st element is in [1…5] becomes a row of ones, all the other rows become rows of zeros.

Let's generate a mask that reflects the test the OP is performing (only the 1st element of each row in the above display is involved)

In [63]: t = np.logical_and(0<a[:,:,0],a[:,:,0]<=5)

and use this mask to address the rows... here we have to add the colon operator to complete the addressing

In [64]: a[ t,:] = 1
In [65]: a[~t,:] = 0

Let's display the result

In [66]: a
Out[66]: 
array([[[1, 1],
        [0, 0],
        [1, 1],
        [0, 0]],

       [[0, 0],
        [1, 1],
        [1, 1],
        [0, 0]],

       [[1, 1],
        [1, 1],
        [1, 1],
        [1, 1]]])

If I interpret correctly the OP example, this is what they want.

I have to say that an answer apparently similar to mine

In [67]: np.random.seed(0) ; a = np.random.randint(10, size=(3,4,2))-3
In [68]: t = (0 < a) & (a <= 5)
In [69]: a[t] = 1
In [70]: a[~t]= 0
In [71]: a
Out[71]: 
array([[[1, 0],
        [0, 0],
        [1, 0],
        [0, 1]],

       [[0, 1],
        [1, 1],
        [1, 1],
        [0, 1]],

       [[1, 1],
        [1, 0],
        [1, 0],
        [1, 0]]])

gives a different result and the OP showed their appreciation for it... who knows?

Substitute deep loops in Python

2 Answers2

Part 1

Part 2