How can I use a 2D array of boolean rows to filter another 2D array?
<p>I have some data in a (3, m) array.</p>

<p>I have another array of masks in (n, 3) shape.  The rows of this mask are boolean filters that need to be applied on the data array before performing some function.  Is there a vectorized way to apply the filter and compute the function?</p>

<p>Here's an example using a loop for clarity, assuming the function is a mean().  I'd like to do this using purely Numpy (without the list comprehension).</p>

<p>(Obviously, the sizes of the arrays are much larger in reality.)</p>

<pre><code>import numpy as np

data = np.array([
       [ 0,  1,  2,  3],
       [ 4,  5,  6,  7],
       [ 8,  9, 10, 11]
    ])

masks = np.array([
        [True, True, False],
        [False, True, False],
        [False, True, True],
        [True, False, False],
        [True, False, True]
    ])

means = np.array([data[mask].mean(axis=0) for mask in masks])

# means
array([[ 2.,  3.,  4.,  5.],
       [ 4.,  5.,  6.,  7.],
       [ 6.,  7.,  8.,  9.],
       [ 0.,  1.,  2.,  3.],
       [ 4.,  5.,  6.,  7.]])
</code></pre>

<p>This feels a bit crude and messy, but it does work without a loop.</p>

<p>There are two main tasks:</p>

<ul>
<li>expand <code>data</code> so it can be indexed with <code>masks</code> - from (5,4) to (5,3,4) </li>
<li>apply <code>means</code> to groups of rows; the closest I can find is <code>np.sum.reduceat</code>.</li>
</ul>

<hr>

<p>construct <code>reduceat</code> indices:</p>

<pre><code>In [253]: cnt = masks.sum(axis=1)
In [254]: cnt1=np.concatenate(([0],np.cumsum(cnt)[:-1]))
In [255]: cnt
Out[255]: array([2, 1, 2, 1, 2])   # True count per row
In [256]: cnt1
Out[256]: array([0, 2, 3, 5, 6])   # reduceat index positions
</code></pre>

<p>expand <code>data</code> and <code>mask</code>:</p>

<pre><code>In [257]: mdata=data[None,...].repeat(masks.shape[0],0)[masks,:]
</code></pre>

<p><code>add</code> rows and divide by the row count for each group</p>

<pre><code>In [258]: np.add.reduceat(mdata,cnt1,0)/cnt[:,None]
Out[258]: 
array([[ 2.,  3.,  4.,  5.],
       [ 4.,  5.,  6.,  7.],
       [ 6.,  7.,  8.,  9.],
       [ 0.,  1.,  2.,  3.],
       [ 4.,  5.,  6.,  7.]])
</code></pre>

<p>In case it helps:</p>

<pre><code>In [263]: mdata
Out[263]: 
array([[ 0,  1,  2,  3],
       [ 4,  5,  6,  7],
       [ 4,  5,  6,  7],
       [ 4,  5,  6,  7],
       [ 8,  9, 10, 11],
       [ 0,  1,  2,  3],
       [ 0,  1,  2,  3],
       [ 8,  9, 10, 11]])
</code></pre>

<p>A possibly better way to get this <code>mdata</code> is</p>

<pre><code>In [285]: data[np.where(masks)[1],:]
Out[285]: 
array([[ 0,  1,  2,  3],
       [ 4,  5,  6,  7],
       [ 4,  5,  6,  7],
       [ 4,  5,  6,  7],
       [ 8,  9, 10, 11],
       [ 0,  1,  2,  3],
       [ 0,  1,  2,  3],
       [ 8,  9, 10, 11]])
</code></pre>

<p>The <code>where(...)[1]</code> are the column positions of the True in <code>masks</code>, which are the rows that we want to select from <code>data</code>.</p>

<p>===========================</p>

<p><code>@capitalistcuttle</code> creates a (5,3,4) array as well, but avoids the need for <code>reduceat</code> by zeroing out the <code>False</code> rows.  That way the can enter into the <code>mean</code> or <code>sum</code> without affecting value.  That reminds me of how masked arrays perform tasks like this.  They <code>fill</code> the masked values with a value like 0 or 1 that doesn't affect the calculation.</p>

<p>Inspired by that here's a MaskedArray solution</p>

<p>Expand both <code>data</code> and <code>masks</code> to the (5,3,4) size:</p>

<pre><code>In [322]: data1=data[None,:,:].repeat(5,0)
In [323]: masks1=masks[:,:,None].repeat(4,-1)
In [324]: data1.shape, masks1.shape
Out[324]: ((5, 3, 4), (5, 3, 4))
</code></pre>

<p>Make masked array from that:</p>

<pre><code>In [325]: madata=np.ma.MaskedArray(data1,~masks1)
In [326]: madata
Out[326]: 
masked_array(data =
 [[[0 1 2 3]
  [4 5 6 7]
  [-- -- -- --]]

 [[-- -- -- --]
  [4 5 6 7]
  [-- -- -- --]]
 ...
 [[0 1 2 3]
  [-- -- -- --]
  [8 9 10 11]]],
             mask =
 [[[False False False False]
  [False False False False]
  [ True  True  True  True]]

 [[ True  True  True  True]
  [False False False False]
  [ True  True  True  True]]
  ...],
       fill_value = 999999)
</code></pre>

<p>Now we can simply use the <code>mean</code> method, letting it take care of 0 fill and adjusting for the number of valid rows.</p>

<pre><code>In [327]: madata.mean(axis=1)
Out[327]: 
masked_array(data =
 [[2.0 3.0 4.0 5.0]
 [4.0 5.0 6.0 7.0]
 [6.0 7.0 8.0 9.0]
 [0.0 1.0 2.0 3.0]
 [4.0 5.0 6.0 7.0]],
             mask =
 [[False False False False]
 [False False False False]
 [False False False False]
 [False False False False]
 [False False False False]],
       fill_value = 1e+20)
</code></pre>

<p>That the <code>.data</code> attribute to convert back to regular array.</p>

<p>This MaskedArray approach is probably slower, since it creates a larger array, but it may be more general - it can be used operations, as long as they are defined in <code>np.ma</code> or its methods.</p>

python performance numpy vectorization