How would I use Dask to perform parallel operations on slices of NumPy arrays?
<p>I have a numpy array of coordinates of size n_slice x 2048 x 3, where n_slice is in the tens of thousands.  I want to apply the following operation on each 2048 x 3 slice separately</p>

<pre><code>import numpy as np
from scipy.spatial.distance import pdist

# load coor from a binary xyz file, dcd format

n_slice, n_coor, _ = coor.shape
r = np.arange(n_coor)
dist = np.zeros([n_slice, n_coor, n_coor])

# this loop is what I want to parallelize, each slice is completely independent
for i in xrange(n_slice): 
    dist[i, r[:, None] &lt; r] = pdist(coor[i])
</code></pre>

<p>I tried using Dask by making <code>coor</code> a <code>dask.array</code>,</p>

<pre><code>import dask.array as da
dcoor = da.from_array(coor, chunks=(1, 2048, 3))
</code></pre>

<p>but simply replacing <code>coor</code> by <code>dcoor</code> will not expose the parallelism.  I could see setting up parallel threads to run for each slice but how do I leverage Dask to handle the parallelism?</p>

<p>Here is the parallel implementation using <code>concurrent.futures</code></p>

<pre><code>import concurrent.futures
import multiprocessing

n_cpu = multiprocessing.cpu_count()

def get_dist(coor, dist, r):
    dist[r[:, None] &lt; r] = pdist(coor)

# load coor from a binary xyz file, dcd format

n_slice, n_coor, _ = coor.shape
r = np.arange(n_coor)
dist = np.zeros([n_slice, n_coor, n_coor])

with concurrent.futures.ThreadPoolExecutor(max_workers=n_cpu) as executor:
    for i in xrange(n_slice):
        executor.submit(get_dist, cool[i], dist[i], r)
</code></pre>

<p>It is possible this problem is not well suited to Dask since there are no inter-chunk computations.</p>

<h3><code>map_blocks</code></h3>

<p>The <a href="http://dask.pydata.org/en/latest/array-api.html#dask.array.core.map_blocks" rel="nofollow">map_blocks</a> method may be helpful:</p>

<pre><code>dcoor.map_blocks(pdist)
</code></pre>

<h3>Uneven arrays</h3>

<p>It looks like you're doing a bit of fancy slicing to insert particular values into particular locations of an output array.  This will probably be awkward to do with dask.arrays.  Instead, I recommend making a function that produces a numpy array</p>

<pre><code>def myfunc(chunk):
    values = pdist(chunk[0, :, :])
    output = np.zeroes((2048, 2048))
    r = np.arange(2048)
    output[r[:, None] &lt; r] = values
    return output

dcoor.map_blocks(myfunc)
</code></pre>

<h3><code>delayed</code></h3>

<p>Worst case scenario you can always use <a href="http://dask.pydata.org/en/latest/delayed.html" rel="nofollow">dask.delayed</a></p>

<pre><code>from dask import delayed, compute
coor2 = delayed(coor)
slices = [coor2[i] for i in range(coor.shape[0])]
slices2 = [delayed(pdist)(slice) for slice in slices]
results = compute(*slices2)
</code></pre>

python arrays numpy parallel-processing dask