Эффективное точечное произведение от плотности до плотности

Question

Я хотел бы эффективно вычислить скалярное произведение двух матриц низкого ранга A и B только для конкретных индексов. Как правило, A и B имеют форму (100, 10000), и я хотел бы получить только 1% от записей A.T @ B.

Я уже пробовал использовать библиотеку Numba, но, похоже, она работает намного медленнее, чем вычисление произведения плотных точек np.dot (A.T, B).

Answer 1

Вы можете использовать einsum:

import numpy as np

def direct():
    return (A.T@B)[iy, ix]

def einsum():
    return np.einsum('ij,ij->j',A[:,iy],B[:,ix])

def make(M,N,p):
    global A,B,iy,ix
    A,B = (np.random.randint(0,10,(M,N)) for _ in "AB")
    iy,ix = np.random.randint(0,N,(2,int(N*N*p)))

from time import time

make(100,10000,0.01)
T=[]
T.append(time())
d = direct()
T.append(time())
e = einsum()
T.append(time())
print("results equal: ",np.allclose(d,e))
print("direct {:.6f} einsum {:.6f}".format(*np.diff(T)))

Пример прогона:

results equal:  True
direct 18.463711 einsum 1.947425

Answer 2

Если вы хотите получить подмножество полученного массива точечного произведения, просто возьмите подмножество, прежде чем использовать точку. Т.е., если вам нужна матрица 10х10 «левого верхнего» вывода, просто сделайте

np.dot(A.T[:10,:], B[:, :10])

Если вам нужны определенные индексы, вы можете использовать более причудливую индексацию. Например, если вам нужны индексы 3, 5 и 29, вы можете сделать что-то вроде этого:

indices = np.array([3, 5, 29]).reshape(-1, 1)
inner_all = np.arange(A.shape[0]).reshape(-1, 1)
result = np.dot(A.T[indices, inner_all.T], B[inner_all, indices.T])

Если вы, только для первой строки, хотите столбцы 1, 3 и 5:

rows = np.array([1]).reshape(-1, 1)
columns = rows = np.array([1, 3, 5]).reshape(-1, 1)
inner_all = np.arange(A.shape[0]).reshape(-1, 1)
result = np.dot(A.T[rows, inner_all.T], B[inner_all, columns.T])