Efficiently avoid large intermediate array in computation?

user_2056067 0
Viewed 163

I have two arrays p and alpha, and I want to <=-compare each element of p with each element of alpha, and then aggregate (count) over the first axis. My code:

s = np.sum(np.less_equal.outer(p, alpha), axis=0)

p is at least one- but possibly multidimensional, and can have dimensions like 100 × 1000000. alpha is one-dimensional and has typically 100 to 1000 elements.

The problem is that np.less_equal.outer creates an intermediate array, which in the worst case can be of size 100 × 1000000 × 1000 = 1011 elements, close to 1 TB, far beyond my memory capacity.

My approach is to split the computation up along the first axis:

s = np.zeros(shape=(*p.shape[1:], len(alpha)), dtype='int64')
for pr in p:
    s += np.less_equal.outer(pr, alpha)

That seems to work, but I'm wondering whether NumPy has tools to make this more efficient (vectorized)?

numpy-ndarray numpy python
2 Answers

In numpy increasing the blocksize may help, but it is also quite straight-forward to implement the algorithm in compiled code. I used Numba here, but with a few changes it should also work in Cython.

First Try

import numpy as np
import numba as nb
@nb.njit(fastmath=True)
def outer_cust(p,alpha):
    res=np.zeros((p.shape[1],alpha.shape[0]),dtype=np.int64)
    for i in range(p.shape[0]):
        for j in range(p.shape[1]):
            for k in range(alpha.shape[0]):
                if p[i,j]<=alpha[k]:
                    res[j,k]+=1
                else:
                    res[j,k]+=0
    return res

This method does not use any temporary arrays. Nevertheless the memory access isn't aligned (assuming C-contiguous arrays). Therefor the performance will be not really good, but already faster than the numpy algorithm.

Second Try

@nb.njit(fastmath=True,parallel=True,cache=True)
def outer_cust_2(p,alpha):
    #p should be fortran ordered to avoid memory copy
    p_T=np.ascontiguousarray(p.T)
    res=np.empty((p.shape[1],alpha.shape[0]),dtype=np.int64)
    for j in nb.prange(p_T.shape[0]):
        for k in range(alpha.shape[0]):
            acc=0
            for i in range(p_T.shape[1]):
                if p_T[j,i]<=alpha[k]:
                    acc+=1
            res[j,k]=acc
    return res

In this method a memory copy is needed if the input is c-ordered. But the memory access is aligned which offers a high speedup, especially on large arrays. It is also easily parallelize-able.

Timings

#Smaller arrays
p=np.random.rand(1000,10000)
alpha=np.random.rand(1000)
%timeit outer_cust_np(p,alpha)
#20 s ± 580 ms per loop (mean ± std. dev. of 7 runs, 1 loop each)
%timeit outer_cust(p,alpha)
#5.56 s ± 61.6 ms per loop (mean ± std. dev. of 7 runs, 1 loop each)
%timeit outer_cust_2(p,alpha)
#166 ms ± 1.19 ms per loop (mean ± std. dev. of 7 runs, 10 loops each)

#Maximum size
p=np.random.rand(1000,1000000)
alpha=np.random.rand(1000)
%timeit outer_cust_np(p,alpha)
#too slow to measure
%timeit outer_cust(p,alpha)
#too slow to measure
%timeit outer_cust_2(p,alpha)
#24.2 s ± 1.52 s per loop (mean ± std. dev. of 7 runs, 1 loop each)
user_40457740

It looks to me like you are generating data for some downstream computation. If that particular downstream computation can be done in parts, then you can just go column by column, save the result to disk, and merge the results later.

Another approach is to get the required parts of this large matrix on demand during the downstream computation using indexed operations on p and alpha.

I would prefer the second approach for simplicity reasons.

user_31340040
Related