Optimize Python: Large arrays, memory problems
<p>I'm having a speed problem running a python / numypy code. I don't know how to make it faster, maybe someone else?</p>

<p>Assume there is a surface with two triangulation, one fine (..._fine) with M points, one coarse with N points. Also, there's data on the coarse mesh at every point (N floats). I'm trying to do the following:</p>

<p>For every point on the fine mesh, find the k closest points on coarse mesh and get mean value. Short: interpolate data from coarse to fine. </p>

<p>My code right now goes like that. With large data (in my case M = 2e6, N = 1e4) the code runs about 25 minutes, guess due to the explicit for loop not going into numpy. Any ideas how to solve that one with smart indexing? M x N arrays blowing the RAM..</p>

<pre><code>import numpy as np

p_fine.shape =&gt; m x 3
p.shape =&gt; n x 3

data_fine = np.empty((m,))
for i, ps in enumerate(p_fine):
    data_fine[i] = np.mean(data_coarse[np.argsort(np.linalg.norm(ps-p,axis=1))[:k]])
</code></pre>

<p>Cheers!</p>

<p><strong>Approach #1</strong></p>

<p>We are working with large sized datasets and memory is an issue, so I will try to optimize the computations within the loop. Now, we can use <a href="http://docs.scipy.org/doc/numpy/reference/generated/numpy.einsum.html" rel="nofollow"><code>np.einsum</code></a> to replace <code>np.linalg.norm</code> part and <a href="http://docs.scipy.org/doc/numpy/reference/generated/numpy.argpartition.html" rel="nofollow"><code>np.argpartition</code></a> in place of actual sorting with <code>np.argsort</code>, like so -</p>

<pre><code>out = np.empty((m,))
for i, ps in enumerate(p_fine):
    subs = ps-p
    sq_dists = np.einsum('ij,ij-&gt;i',subs,subs)
    out[i] = data_coarse[np.argpartition(sq_dists,k)[:k]].sum()
out = out/k
</code></pre>

<p><strong>Approach #2</strong></p>

<p>Now, as another approach we can also use <a href="http://docs.scipy.org/doc/scipy/reference/generated/scipy.spatial.distance.cdist.html" rel="nofollow"><code>Scipy's cdist</code></a> for a fully vectorized solution, like so -</p>

<pre><code>from scipy.spatial.distance import cdist
out = data_coarse[np.argpartition(cdist(p_fine,p),k,axis=1)[:,:k]].mean(1)
</code></pre>

<p>But, since we are memory bound here, we can perform these operations in chunks. Basically, we would get chunks of rows from that tall array <code>p_fine</code> that has millions of rows and use <code>cdist</code> and thus at each iteration get chunks of output elements instead of just one scalar. With this, we would cut the loop count by the length of that chunk.</p>

<p>So, finally we would have an implementation like so -</p>

<pre><code>out = np.empty((m,))
L = 10 # Length of chunk (to be used as a param)
num_iter = m//L
for j in range(num_iter):
    p_fine_slice = p_fine[L*j:L*j+L]
    out[L*j:L*j+L] = data_coarse[np.argpartition(cdist\
                           (p_fine_slice,p),k,axis=1)[:,:k]].mean(1)
</code></pre>

<p><strong>Runtime test</strong></p>

<p>Setup -</p>

<pre><code># Setup inputs
m,n = 20000,100
p_fine = np.random.rand(m,3)
p = np.random.rand(n,3)
data_coarse = np.random.rand(n)
k = 5

def original_approach(p,p_fine,m,n,k):
    data_fine = np.empty((m,))
    for i, ps in enumerate(p_fine):
        data_fine[i] = np.mean(data_coarse[np.argsort(np.linalg.norm\
                                                 (ps-p,axis=1))[:k]])
    return data_fine

def proposed_approach(p,p_fine,m,n,k):    
    out = np.empty((m,))
    for i, ps in enumerate(p_fine):
        subs = ps-p
        sq_dists = np.einsum('ij,ij-&gt;i',subs,subs)
        out[i] = data_coarse[np.argpartition(sq_dists,k)[:k]].sum()
    return out/k

def proposed_approach_v2(p,p_fine,m,n,k,len_per_iter):
    L = len_per_iter
    out = np.empty((m,))    
    num_iter = m//L
    for j in range(num_iter):
        p_fine_slice = p_fine[L*j:L*j+L]
        out[L*j:L*j+L] = data_coarse[np.argpartition(cdist\
                               (p_fine_slice,p),k,axis=1)[:,:k]].sum(1)
    return out/k
</code></pre>

<p>Timings -</p>

<pre><code>In [134]: %timeit original_approach(p,p_fine,m,n,k)
1 loops, best of 3: 1.1 s per loop

In [135]: %timeit proposed_approach(p,p_fine,m,n,k)
1 loops, best of 3: 539 ms per loop

In [136]: %timeit proposed_approach_v2(p,p_fine,m,n,k,len_per_iter=100)
10 loops, best of 3: 63.2 ms per loop

In [137]: %timeit proposed_approach_v2(p,p_fine,m,n,k,len_per_iter=1000)
10 loops, best of 3: 53.1 ms per loop

In [138]: %timeit proposed_approach_v2(p,p_fine,m,n,k,len_per_iter=2000)
10 loops, best of 3: 63.8 ms per loop
</code></pre>

<p>So, there's about <strong><code>2x</code></strong> improvement with the first proposed approach  and <strong><code>20x</code></strong> over the original approach with the second one at the sweet spot with the <code>len_per_iter</code> param set at <code>1000</code>. Hopefully this will bring down your 25 minutes runtime to little over a minute. Not bad I guess!</p>

python arrays performance numpy large-data