Элегантное выражение для двумерного точечного произведения двух матриц

У меня есть два массива с двумя номерами с такими же размерами, A и B, и я пытаюсь вычислить их произведение в виде строки. Я мог бы сделать:

np.sum(A * B, axis=1)

Есть ли другой способ сделать это, чтобы numpy делал продукт с меткой точки за один шаг, а не два? Может быть, с tensordot?

Ответ 1

Это хорошее приложение для numpy.einsum.

a = np.random.randint(0, 5, size=(6, 4))
b = np.random.randint(0, 5, size=(6, 4))

res1 = np.einsum('ij, ij->i', a, b)
res2 = np.sum(a*b, axis=1)

print(res1)
# [18  6 20  9 16 24]

print(np.allclose(res1, res2))
# True

einsum также имеет тенденцию быть немного быстрее.

a = np.random.normal(size=(5000, 1000))
b = np.random.normal(size=(5000, 1000))

%timeit np.einsum('ij, ij->i', a, b)
# 100 loops, best of 3: 8.4 ms per loop

%timeit np.sum(a*b, axis=1)
# 10 loops, best of 3: 28.4 ms per loop

Ответ 2

Еще быстрее inner1d от numpy.core.umath_tests:

введите описание изображения здесь


Код для воспроизведения сюжета:

import numpy
from numpy.core.umath_tests import inner1d
import perfplot


perfplot.show(
        setup=lambda n: (numpy.random.rand(n, 3), numpy.random.rand(n, 3)),
        kernels=[
            lambda a: numpy.sum(a[0]*a[1], axis=1),
            lambda a: numpy.einsum('ij, ij->i', a[0], a[1]),
            lambda a: inner1d(a[0], a[1])
            ],
        labels=['sum', 'einsum', 'inner1d'],
        n_range=[2**k for k in range(20)],
        xlabel='len(a), len(b)',
        logx=True,
        logy=True
        )

Ответ 3

Несмотря на то, что для умеренных размеров данных он значительно медленнее, я бы использовал

np.diag(A.dot(B.T))

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

Для большинства людей, которые приходят на ваш код, это будет более понятным, чем einsum, а также не требует, чтобы вы нарушили некоторые лучшие практики, вложив ваши вычисления в мини-строку DSL, чтобы служить аргументом для некоторых вызов функции.

Я согласен, что вычислять недиагональные элементы стоит избегать для больших случаев. Для меня было бы действительно очень важно заботиться об этом, и компромисс для оплаты ужасной цены выражения вычисления во встроенной строке в einsum довольно серьезный.