Avoid using for loop in np.array operations. Python

Question

I have these two arrays:

import numpy as np
a = np.array([0, 10, 20])
b = np.array([20, 30, 40, 50])

I'd like to add both in the following way:

for i in range (len(a)):
   for j in range(len(b)):
      c = a[i] + b[j]
      d = delta(c, dr)

As you see for each iteration I get a value c which I pass through a function delta (see note at the end of the post). The thing is that I want to avoid slow Python "for" loops when the arrays are huge.

One thing I could do would be:

c = np.ravel(a(-1, 1) + b)

Which is much much faster. The problem is that now c is an array, and again I would have to go throw it using a for loop. So, do you have any idea on how I could do this without using a for loop at all.

NOTE: delta is a function I define in the following way:

def delta(r,dr):
   if r >= 0.5*dr and r <= 1.5*dr:
     delta = (5-3*abs(r)/dr-np.sqrt(-3*(1-abs(r)/dr)**2+1))/(6*dr)
   elif r <= 0.5*dr:
     delta = (1+np.sqrt(-3*(r/dr)**2+1))/(3*dr)
   else:
     delta = 0
   return delta

Is function a suitable NumPy ufunc, or a Python function? In the former case, you might be able to optimize; in the latter, you wouldn't. — nneonneo
– nneonneo, Commented Nov 3, 2015 at 16:15
Your delta function takes two variables, not one as in the for loop. Also, it doesn't return anything. Did you intend to return the delta variable? — Frank M
– Frank M, Commented Nov 3, 2015 at 16:59

DilithiumMatrix · Accepted Answer · 2015-11-03 18:16:45Z

3

Using ravel is a good idea. Note that you could also use simple array broadcasting (a[:, np.newaxis] + b[np.newaxis, :]).

For your function, you can improve this a lot because it is composed of only three particular cases. Probably the best approach is to use masking for each of those three sections.

You're starting with:

def delta(r,dr):
   if r >= 0.5*dr and r <= 1.5*dr:
     delta = (5-3*abs(r)/dr-np.sqrt(-3*(1-abs(r)/dr)**2+1))/(6*dr)
   elif r <= 0.5*dr:
     delta = (1+np.sqrt(-3*(r/dr)**2+1))/(3*dr)
   else:
     delta = 0

A common alternative approach would be something like:

def delta(r, dr):
    res = np.zeros_like(r)
    ma = (r >= 0.5*dr) & (r <= 1.5*dr)  # Create first mask
    res[ma] = (5-3*np.abs(r[ma])/dr[ma]-np.sqrt(-3*(1-np.abs(r[ma])/dr[ma])**2+1))/(6*dr[ma])
    ma = (r <= 0.5*dr)    # Create second mask
    res[ma] = (1+np.sqrt(-3*(r[ma]/dr[ma])**2+1))/(3*dr[ma])
    return res

Initializing to zeros handles the final else case. Also I'm assuming np.abs is faster than abs --- but I'm not actually sure...

Edit: for sparse matrices

The same basic idea should apply, but perhaps instead of using a boolean masking array, using the valid indices themselves would be better... e.g. something like:

res = scipy.sparse.coo_matrix(np.shape(r))
ma = np.where((r >= 0.5*dr) & (r <= 1.5*dr))  # Create first mask
res[ma] = ...

edited Nov 3, 2015 at 18:16

answered Nov 3, 2015 at 16:55

DilithiumMatrix

18.8k24 gold badges82 silver badges126 bronze badges

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

5 Comments

Frank M Over a year ago

NIce. I was just writing a nearly identical answer!

Frank M Over a year ago

BTW: and doesn't work between bool arrays. You'll need to use &.

Fernando Bastos García Over a year ago

@DilithiumMatrix That´s just what I needed, it is much faster for huge matrix. However the reason of doing it one by one using a for-loop was to create 3 vector one with the result other with i and other with k and then create a coo-matrix, since most of the elements are zero. Now I have a dense matrix. Any ideas to create a coo_matrix and not write thousands of elements with value zero? Thanks.

DilithiumMatrix Over a year ago

@FernandoBastosGarcía hmm, interesting, honestly I've never used sparse matrices :/ but I've added an edit with what I assume should work... Let me know!

hpaulj Over a year ago

I think you should post a separate sparse matrix question.

Frank M · Accepted Answer · 2015-11-03 17:14:00Z

0

This is the same answer as DilithiumMatrix, but using logical functions that numpy accepts to generate the masks.

import numpy as np
def delta(r, dr):
    res = np.zeros(r.shape)
    mask1 = (r >= 0.5*dr) & (r <= 1.5*dr)
    res[mask1] = \
        (5-3*np.abs(r[mask1])/dr \
        - np.sqrt(-3*(1-np.abs(r[mask1])/dr)**2+1)) \
         /(6*dr)
    mask2 = np.logical_not(mask1) & (r <= 0.5*dr)
    res[mask2] = (1+np.sqrt(-3*(r[mask2]/dr)**2+1))/(3*dr)
    return res

answered Nov 3, 2015 at 17:14

Frank M

1,59015 silver badges16 bronze badges

1 Comment

Fernando Bastos García Over a year ago

please check the comment I write in DilithiumMatrix answer. Thanks

Chinmay Kanchi · Accepted Answer · 2015-11-03 17:11:15Z

-1

Assuming your two arrays (a and b) are not enormous, you could do something like this:

import itertools
a = numpy.array([1,2,3])
b = numpy.array([4,5,6])
c = numpy.sum(list(itertools.product(a, b), 1)
def func(x, y):
    return x*y
numpy.vectorize(func)(c, 10)

Note that with large arrays, this simply won't work - you'll have n**2 elements in c, which means that even for smallish-seeming pairs of arrays, you'll use enormous amounts of memory. For 2 arrays with 100,000 elements each, total memory required will be in the range of 74 GB.

answered Nov 3, 2015 at 17:11

Chinmay Kanchi

66.6k24 gold badges92 silver badges115 bronze badges

1 Comment

Frank M Over a year ago

You don't need the itertools.product to implement a vectorized product. Also, your func doesn't implement the OP's function.

Collectives™ on Stack Overflow

Avoid using for loop in np.array operations. Python

3 Answers 3

5 Comments

1 Comment

1 Comment

Your Answer

Hot Network Questions

Collectives™ on Stack Overflow

3 Answers 3

5 Comments

1 Comment

1 Comment

Your Answer

Sign up or log in

Post as a guest

Related