Populating NumPy array with most performance

Question

I want to populate a matrix by function f() which consumes arrays a, b, c and d:

A nested loop is possible but I'm looking for a faster way. I tried np.fromfunction with no luck. Function f has a conditional so the solution should preferably support conditionals. Example function:

def f(a,b,c,c):
    return a+b+c+d if a==b else a*b*c*d

How np.fromfunction failed:

>>> a = np.array([1,2,3,4,5])
>>> b = np.array([10,20,30])
>>> def f(i,j): return a[i] * b[j]
>>> np.fromfunction(f, (3,5))
Traceback (most recent call last):
  File "<stdin>", line 1, in <module>
  File "C:\Users\Anaconda3\lib\site-packages\numpy\core\numeric.py", line 1853, in fromfunction
    return function(*args, **kwargs)
  File "<stdin>", line 1, in fun
IndexError: arrays used as indices must be of integer (or boolean) type

Without more details about the function, it's going to be hard to suggest the correct approach. With the information you've given, fromfunction or nested loops seems the best way. What did you try that you say you had no luck with? — pho
– pho, Commented Feb 7, 2023 at 9:56
I've added an example function if it helps. The reason I had no luck with fromfunction is because the indices are already broadcasted which cannot be used to pick values from 1d arrays . why this fails is shown above — XYZ
– XYZ, Commented Feb 7, 2023 at 10:01
An if condition could be vectorized so you don't need a call to fromfunction, which would have to call your function for every index. The fromfunction that failed includes an indexing operation. Is this a part of the function you actually want? Your example could be recreated without fromfunction by adding an axis to b and broadcasting: b[:, None] * a gives you the 3x5 array where the (i, j) element is a multiplication of a[j] and b[i] — pho
– pho, Commented Feb 7, 2023 at 18:42
Also your first fun doesn't really make sense: do you want to return the sum (or product) of the indices (which is what a, b c, d) are, or do you want to return the values at those indices in the arrays? Did you mean def fun(i, j): return a[i]+b[i]+c[i]+d[j] if i==j else a[i]*b[i]*c[i]*d[j] — pho
– pho, Commented Feb 7, 2023 at 18:59

jakevdp · Accepted Answer · 2023-02-07 13:27:56Z

1

The reason the function fails is that np.fromfunction passes floating-point values, which are not valid as indices. You can modify your function like this to make it work:

def fun(i,j):
  return a[j.astype(int)] * b[i.astype(int)]

print(np.fromfunction(fun, (3,5)))

[[ 10  20  30  40  50]
 [ 20  40  60  80 100]
 [ 30  60  90 120 150]]

answered Feb 7, 2023 at 13:27

jakevdp

88.5k15 gold badges140 silver badges177 bronze badges

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

Comments

pho · Accepted Answer · 2023-02-07 23:50:12Z

Jake has explained why your fromfunction approach fails. However, you don't need fromfunction for your example. You could simply add an axis to b and have numpy broadcast the shapes:

a = np.array([1,2,3,4,5])
b = np.array([10,20,30])
def fun(i,j): return a[j.astype(int)] * b[i.astype(int)]

f1 = np.fromfunction(fun, (3, 5))
f2 = b[:, None] * a

(f1 == f2).all() # True

Extending this to the function you showed that contains an if condition, you could just split the if into two operations in sequence: creating an array given by the if expression, and overwriting the relevant parts by the else expression.

a = np.array([1, 2, 3, 4, 5])
b = np.array([5, 4, 3, 2, 1])
c = np.array([100, 200, 300, 400, 500])

d = np.array([0, 1, 2, 3])

# Calculate the values at all indices as the product
result = d[:, None] * (a * b * c)
# array([[   0,    0,    0,    0,    0],
#        [ 500, 1600, 2700, 3200, 2500],
#        [1000, 3200, 5400, 6400, 5000],
#        [1500, 4800, 8100, 9600, 7500]])

# Calculate sum 
sum_arr = d[:, None] + (a + b + c)
# array([[106, 206, 306, 406, 506],
#        [107, 207, 307, 407, 507],
#        [108, 208, 308, 408, 508],
#        [109, 209, 309, 409, 509]])

# Set diagonal elements (i==j) to sum:
np.fill_diagonal(result, np.diag(sum_arr))

which gives the following result:

array([[ 106,    0,    0,    0,    0],
       [ 500,  207, 2700, 3200, 2500],
       [1000, 3200,  308, 6400, 5000],
       [1500, 4800, 8100,  409, 7500]])

Collectives™ on Stack Overflow

Populating NumPy array with most performance

2 Answers 2

Comments

Comments

Your Answer

Hot Network Questions

Collectives™ on Stack Overflow

2 Answers 2

Comments

Comments

Your Answer

Sign up or log in

Post as a guest

Related