Un-вещание Numpy массивы

Question

В большой базе кода, я использую np.broadcast_to транслировать массивы (только с помощью простых примеров здесь):

In [1]: x = np.array([1,2,3]) 

In [2]: y = np.broadcast_to(x, (2,1,3)) 

In [3]: y.shape 
Out[3]: (2, 1, 3) 

В другом месте в коде, я использую сторонние функции, которые могут работать в векторном виде на массивах Numpy, но это не ufuncs. Эти функции не понимают вещание, а это означает, что вызов такой функции на массивах, как y, неэффективен. Такие решения, как Numpy's vectorize, не очень хороши либо потому, что, пока они понимают вещание, они вводят цикл for над элементами массива, который тогда очень неэффективен.

В идеале, что бы я хотел сделать, это иметь функцию, которую мы можем назвать, например. unbroadcast, который возвращает массив с минимальной формой, который может быть передан обратно в полный размер, если это необходимо. Так например .:

In [4]: z = unbroadcast(y) 

In [5]: z.shape 
Out[5]: (1, 1, 3) 

Затем я могу запустить функции сторонних на z, а затем транслировать результат обратно в y.shape.

Есть ли способ реализовать unbroadcast, который опирается на публичный API Numpy? Если нет, есть ли какие-либо хаки, которые приведут к желаемому результату?

Answer 1

Это, вероятно, соответствует вашему собственному решению, только немного больше построен -в. Он использует as_strided в numpy.lib.stride_tricks:

import numpy as np 
from numpy.lib.stride_tricks import as_strided 

x = np.arange(16).reshape(2,1,8,1) # shape (2,1,8,1) 
y = np.broadcast_to(x,(2,3,8,5)) # shape (2,3,8,5) broadcast 

def unbroadcast(arr): 
    #determine unbroadcast shape 
    newshape = np.where(np.array(arr.strides) == 0,1,arr.shape) # [2,1,8,1], thanks to @Divakar 
    return as_strided(arr,shape=newshape) # strides are automatically set here 

z = unbroadcast(x) 
np.all(z==x) # is True 

---

Обратите внимание, что в моем оригинальном ответе я не определить функцию, и в результате z массива был (64,0,8,0) в strides, в то время как вход имеет (64,64,8,8). В текущей версии возвращаемый массив z имеет одинаковые шаги до x, я думаю, что передача и возвращение массива заставляет создать копию. Во всяком случае, мы всегда могли бы установить шаги вручную в as_strided, чтобы получить идентичные массивы при любых обстоятельствах, но это не кажется необходимым в приведенной выше настройке.

Answer 2

У меня есть возможное решение, поэтому опубликуйте его здесь (однако, если у кого-то есть лучший, пожалуйста, не стесняйтесь отвечать тоже!). Одним из решений является проверка strides аргумент массива, который будет 0 по транслируемых размеры:

def unbroadcast(array): 
    slices = [] 
    for i in range(array.ndim): 
     if array.strides[i] == 0: 
      slices.append(slice(0, 1)) 
     else: 
      slices.append(slice(None)) 
    return array[slices] 

Это дает:

In [14]: unbroadcast(y).shape 
Out[14]: (1, 1, 3)