CODEWITHSUNDEEP
pythonarraysnumpyvectorization
Christian Alis
Christian Alis

Reputation: 6806

Call function with internal sum over 2D array

Suppose I have a python function f() that accepts 2 scalar and 1 array_like parameters:

def f(a, b, arr):
    X = a * np.exp(-arr**2 / b)
    return np.sum(a * np.log(X) - arr)

What I want to do is to evaluate f() for different values of a and b while keeping the same arr:

XX, YY = np.meshgrid(A_axis, B_axis)
arr = np.arange(10)
ZZ = np.empty_like(XX)
for i in range(XX.shape[0]):
    for j in range(YY.shape[1]):
        ZZ[i,j] = f(XX[i,j], YY[i,j], arr)

Is there a way to vectorize this? I'm thinking of converting XX, YY and arr into 3D arrays of the same shape but the np.sum() in f() will always return a scalar.

Upvotes: 1

Views: 123

Answers (1)

HYRY
HYRY

Reputation: 97331

  1. Construct an open mesh from xaxis, yaxis, arr data by np.ix_()

  2. Call np.sum() with axis=-1.

Here is the code:

import numpy as np

### original code
def f(a, b, arr):
    X = a * np.exp(-arr**2 / b)
    return np.sum(a * np.log(X) - arr)

A_axis = np.linspace(1, 5, 8)
B_axis = np.linspace(1, 2, 9)
XX, YY = np.meshgrid(A_axis, B_axis)
arr = np.arange(10)
ZZ = np.empty_like(XX)
for i in range(XX.shape[0]):
    for j in range(YY.shape[1]):
        ZZ[i,j] = f(XX[i,j], YY[i,j], arr)

### use broadcast        
def f(a, b, arr):
    X = a * np.exp(-arr**2 / b)
    return np.sum(a * np.log(X) - arr, axis=-1)
B, A, C = np.ix_(B_axis, A_axis, arr)
result = f(A, B, C)

print np.allclose(ZZ, result)

Upvotes: 2

Related Questions