延伸numpy面具

Question

I want to mask a numpy array a with mask . The mask doesn't have exactly the same shape as a , but it is possible to mask a anyway (I guess because of the additional dimension being 1-dimensional (broadcasting?)).

a.shape
>>> (3, 9, 31, 2, 1)
mask.shape
>>> (3, 9, 31, 2)
masked_a = ma.masked_array(a, mask)

The same logic however, does not apply to array b which has 5 elements in its last dimension.

ext_mask = mask[..., np.newaxis] # extending or not extending has same effect
ext_mask.shape
>>> (3, 9, 31, 2, 1)

b.shape
>>> (3, 9, 31, 2, 5)
masked_b = ma.masked_array(b, ext_mask)
>>> numpy.ma.core.MaskError: Mask and data not compatible: data size is 8370, mask size is 1674.

How can I create a (3, 9, 31, 2, 5) mask from a (3, 9, 31, 2) mask by expanding any True value in the last dimension of the (3, 9, 31, 2) mask to [True, True, True, True, True] (and False respectively)?

Answer 1

This gives the desired result:

masked_b = ma.masked_array(*np.broadcast(b, ext_mask))

I have not profiled this method, but it should be faster than allocating a new mask. According to the documentation , no data is copied:

These arrays are views on the original arrays. They are typically not contiguous. Furthermore, more than one element of a broadcasted array may refer to a single memory location. If you need to write to the arrays, make copies first.

It is possible to verify the no-copying behavior:

bb, mb = np.broadcast(b, ext_mask)
print(mb.shape)       # (3, 9, 31, 2, 5) - same shape as b
print(mb.base.shape)  # (3, 9, 31, 2) - the shape of the original mask
print(mb.strides)     # (558, 62, 2, 1, 0) - that's how it works: 0 stride

Pretty impressive how the numpy developers implemented broadcasting. Values are repeated by using a stride of 0 along the last dimension. Whow!

Edit

I compared the speed of broadcasting and allocating with this code:

import numpy as np
from numpy import ma

a = np.random.randn(30, 90, 31, 2, 1)
b = np.random.randn(30, 90, 31, 2, 5)

mask = np.random.randn(30, 90, 31, 2) > 0
ext_mask = mask[..., np.newaxis]

def broadcasting(a=a, b=b, ext_mask=ext_mask):
    mb1 = ma.masked_array(*np.broadcast_arrays(b, ext_mask))

def allocating(a=a, b=b, ext_mask=ext_mask):
    m2 = np.empty(b.shape, dtype=bool)
    m2[:] = ext_mask
    mb2 = ma.masked_array(b, m2)

Broadcasting is clearly faster than allocating, here:

    # array size: (30, 90, 31, 2, 5)

In [23]: %timeit broadcasting()
The slowest run took 10.39 times longer than the fastest. This could mean that an intermediate result is being cached.
10000 loops, best of 3: 39.4 µs per loop

In [24]: %timeit allocating()
The slowest run took 4.86 times longer than the fastest. This could mean that an intermediate result is being cached.
1000 loops, best of 3: 982 µs per loop

Note that I had to increase array size for the difference in speed to become apparent. With the original array dimensions allocating was slightly faster than broadcasting:

    # array size: (3, 9, 31, 2, 5)

In [28]: %timeit broadcasting()
The slowest run took 9.36 times longer than the fastest. This could mean that an intermediate result is being cached.
10000 loops, best of 3: 39 µs per loop

In [29]: %timeit allocating()
The slowest run took 9.22 times longer than the fastest. This could mean that an intermediate result is being cached.
10000 loops, best of 3: 32.6 µs per loop

The broadcasting solution's runtime seems not to depend on array size.