Предельная гибкость с эллипсами Эйнсума - PullRequest
Новая
       26

Предельная гибкость с эллипсами Эйнсума

0 голосов
/ 06 апреля 2019

У меня есть вопрос об einsum ellipsis, который, я думал, наверняка будет где-то на StackExchange, но почему-то я не могу его найти.

По сути, у меня есть некоторый код, который выполняет множество матричных и векторных сокращений, используя numpy's einsum. Входными данными обычно являются некоторые параметры, которые затем используются для создания векторов и матриц. Код работает хорошо, но теперь я хотел бы обобщить его, чтобы входные параметры можно было сканировать в определенном диапазоне. Самое приятное, что можно сделать, это сделать их векторами и изменить мои einsum выражения так, чтобы они принимали произвольное количество дополнительных измерений, которые просто переносятся. Этот вопрос состоит в том, чтобы спросить, возможно ли это, и если да, то как.

Так что, на мой взгляд, эта проблема сводится к следующему. Скажем, у меня есть выражение einsum, которое создает умножение матриц, например 

c = np.einsum('ij,jk->ik', a, b)

Теперь я хочу добавить произвольное количество индексов к a и b и просто добавить их в качестве дополнительных индексов в окончательную матрицу, например, 

c = np.einsum('ijabc,jkde->ikabcde', a, b)

Теперь, когда вы делаете это только для одного из a или b, вы можете легко сделать это с помощью многоточия 

c = np.einsum('ij...,jk->ik...', a, b)

Так что мой вопрос в том, можете ли вы иметь несколько эллипсов в einsum, например, 

c = np.einsum('ij...,jk...->ik...', a, b)

Это, конечно, приведет к ошибке, но, надеюсь, из примеров понятно, что я имею в виду.

Поддерживает ли einsum этот тип многолоточечной нотации? Или есть ли другой способ реализовать это без зацикливания?

Я предполагаю, что такого способа не существует, потому что нужно было бы указать einsum, в каком порядке располагать оставшиеся индексы, т. Е. Нужно было бы как-то маркировать эллипсы.

1 Ответ

1 голос
/ 06 апреля 2019

Так как нет осей, которые должны быть выровнены, мы можем просто использовать tensordot, который позволяет осям, не участвующим в уменьшении суммы, быть "растянутыми" с дополнительными rollaxis, как, например, - 

np.rollaxis(np.tensordot(a,b,axes=(1,0)),a.ndim-1,1)

Если вы хотите использовать einsum, мы можем изменить их на 3D так, чтобы последняя ось их была объединена (третья ось объединена в одну), а затем перейти к einsum и, наконец, измените форму, чтобы их ndim-1 фигуры были распределены в выводе, что-то вроде этого - 

shp_a = a.shape
shp_b = b.shape
shp_a[:1] + shp_a[2:]
out_shp = shp_a[:1] + (shp_b[1],) + shp_a[2:] + shp_b[2:]

a3D = a.reshape(shp_a[:2]+(-1,))
b3D = b.reshape(shp_b[:2]+(-1,))
out = np.einsum('ijk,jlm->ilkm',a3D,b3D).reshape(out_shp)

Мы могли бы также сгенерировать соответствующую нотацию einsum непосредственно и, следовательно, пропустить все манипуляции с массивами и, следовательно, сосредоточиться на самой манипуляции со строками, чтобы получить что-то вроде этого - 

import string

def einsum_spreadout(a,b,a_axes,b_axes,a_spread_axis,b_spread_axis):
    from numpy.core import numerictypes as nt

    if isinstance(a_axes, (int, nt.integer)):
        a_axes = (a_axes,)

    if isinstance(b_axes, (int, nt.integer)):
        b_axes = (b_axes,)

    s = string.ascii_letters

    a_str = s[:a.ndim]
    b_str = s[a.ndim:a.ndim+b.ndim]

    b_str_ar = np.frombuffer(b_str,dtype='S1').copy()
    for (i,j) in zip(a_axes,b_axes):
        b_str_ar[j] = a_str[i]
    b_str = ''.join(b_str_ar)    

    out_str = a_str[:a_spread_axis] + b_str[:b_spread_axis]
    out_str += a_str[a_spread_axis:] + b_str[b_spread_axis:]

    out_str_ar = np.frombuffer(out_str,dtype='S1').copy()
    out_str = ''.join(out_str_ar[~np.isin(out_str_ar,np.take(b_str_ar,b_axes))])
    einsum_str = a_str+','+b_str+'->'+out_str

    return np.einsum(einsum_str,a,b)

В нескольких примерах показано, как его использовать - 

>>> a = np.random.rand(3,4,6,7,8)
>>> b = np.random.rand(4,5,9,10)
>>> einsum_spreadout(a,b,a_axes=1,b_axes=0,a_spread_axis=2,b_spread_axis=2).shape
(3, 5, 6, 7, 8, 9, 10)

>>> b = np.random.rand(4,5,6,10)
>>> einsum_spreadout(a,b,a_axes=(1,2),b_axes=(0,2),a_spread_axis=2,b_spread_axis=2).shape
(3, 5, 7, 8, 10)

>>> einsum_spreadout(a,b,a_axes=(1,2),b_axes=(0,2),a_spread_axis=4,b_spread_axis=4).shape
(3, 7, 5, 10, 8)
...