Лучший способ расширить размерность массива numpy для удовлетворения различных требований

Question

У меня есть маска m формы [bz], а словарь d содержит много разных ndarrays, например, d['s'].shape=(bz, 84, 84, 4), d['r'].shape=(bz, 1) и т. Д. Все имеют одно и то же первое измерение размера bz, ноостальное может отличаться. Я хочу правильно расширить размер m, чтобы умножить m на значения в d. Например, я хочу, чтобы m имел форму [bz, 1, 1, 1], чтобы умножить d['s'] и [bz, 1] для d['r']. Я могу придумать противное решение для цикла while следующим образом.

for k, v in d.items():
    if m.shape != v.shape:
        reshaped_m = m.copy()
        while len(reshaped_m.shape) < len(v.shape):
            reshaped_m = reshaped_m[..., None]
        d[k] = v * reshaped_m

Я ищу лучшее решение, спасибо.

Answer 1

Вы можете просто сделать -

for k, v in d.items():
    d[k] = (v.T*m).T

Итак, мы в основном толкаем первую ось до конца, так что v становится транслируемым против m. Затем умножьте на m. Наконец, толкая последнюю ось назад к передней части.

Могут быть различные способы перестановки этих осей, но transposing является самым простым.

Если маска имеет больше измерений, чем 1, нам тоже нужно транспонировать m. Следовательно. для этого случая это будет d[k] = (v.T*m.T).T.

---

Другой способ будет изменить форму m, а затем умножить на v -

for k, v in d.items():
    d[k] =  v*m.reshape(m.shape + (1,)*(v.ndim-m.ndim))

---

Идругой с np.einsum -

for k, v in d.items():
    d[k] = np.einsum('i...,i->i...',v,m)