Расширение маски numpy

Question

Я хочу замаскировать массив numpy a с mask. Маска не имеет точно такой же формы, как a, но в любом случае можно маскировать a (я думаю, из-за того, что дополнительное измерение является одномерным (широковещательным?)).

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

Та же логика, однако, не относится к массиву b, который имеет 5 элементов в своем последнем измерении.

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. 

Как я могу создать (3, 9, 31, 2, 5) маски из (3, 9, 31, 2) маски, расширив любое True значения в последнем измерении (3, 9, 31, 2) маски к [True, True, True, True, True] (и False соответственно)?

Answer 1

Это дает желаемый результат:

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

Я не профилированный этот метод, но он должен быть быстрее, чем выделять новую маску. Согласно documentation, никакие данные не копируются:

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.

можно проверить не-копирование поведения:

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 

Впечатляет, как разработчики Numpy реализованы вещания. Значения повторяются с использованием шага 0 по последнему измерению. Whow!

Редактировать

Я сравнил скорость вещания и распределения с этим кодом:

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 явно быстрее, чем распределение, здесь:

# 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 

Обратите внимание, что я пришлось увеличить размер массива, чтобы разница в скорости стала очевидной. С исходных размерами массива распределения было немного быстрее, чем вещание: среда

# 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 

радиовещательной Solution, кажется, не зависит от размера массива.