CODEWITHSUNDEEP
pythonnumpyreducelarge-data
Josephine Moeller
Josephine Moeller

Reputation: 934

Function application and reduction on large arrays

I have two numpy arrays, X and Y whose shapes are X.shape == (m,d) and Y.shape == (n,d), where m, n, and d are non-trivial sizes. I need to make a third array Z whose shape is Z.shape == (m,n).

An element Z[i,j] is the result of taking f(X[i,k],Y[j,k]) for k in range(d) and then summing over all k, for some non-linear f.

The obvious way to do this is to do this: 

Z = numpy.zeros((m,n), dtype = numpy.float64)
for i in range(m):
    for j in range(n):
        Z[i,j] += (f(X[i,:],Y[j,:])).sum() # I can compose f from ufuncs

but what I'm really asking is whether there's some kind of clever broadcasting trick that I can use to compute Z that will:

  1. take advantage of numpy's optimizations if possible
  2. do this without putting an array of shape (n,m,d) in memory (n*m doubles will fit in memory, but n*m*d doubles won't)

Does anyone know of a way to do this? Thanks in advance.

Upvotes: 2

Views: 81

Answers (1)

Bi Rico
Bi Rico

Reputation: 25813

Here is the solution you don't want, I've included it because I believe this is the "canonical" solution to your problem.

# A simple function of x, y
def f(x, y):
    return 2*x + 3*y**2

x = x.reshape((m, 1, d))
y = y.reshape((1, n, d))
temp = f(x, y)
Z = temp.sum(2)

If you want to avoid creating the temporary array temp, which is quite large, you could try looping over the d dimension. In some cases the overhead of the following loop will be quite small and you'll get almost the same performance, with much less memory usage.

Z = np.zeros((m, n))
for i in range(d):
    Z += f(x[:, :, i], y[:, :, i])

Let me know if that helps.

Upvotes: 1

Related Questions