def apply_mask(a, indices, axis):
    magic_index = [np.arange(i) for i in indices.shape]
magic_index = np.ix_(*magic_index)
magic_index = magic_index[:axis] + (indices,) + magic_index[axis:]
return a[magic_index]
numpy 3d to 2d transformation based on 2d mask array