Быстрое внутреннее произведение двух двумерных масок в масках

Question

Моя проблема заключается в следующем. У меня есть два массива X и Y формы n, p где p >> n (например, n = 50, p = 10000).

У меня также есть маска mask (1-й массив логических значений размера p) относительно p, малая плотность (например, np.mean(mask) равно 0,05).

Я пытаюсь как можно быстрее вычислить внутреннее произведение X и Y относительно mask: вывод inner представляет собой массив формы n, n и такой, что inner[i, j] = np.sum(X[i, np.logical_not(mask)] * Y[j, np.logical_not(mask)]).

Я пытался использовать библиотеку numpy.ma, но она довольно медленная для моего использования:

import numpy as np
import numpy.ma as ma
n, p = 50, 10000
density = 0.05
mask = np.array(np.random.binomial(1, density, size=p), dtype=np.bool_)
mask_big = np.ones(n)[:, None] * mask[None, :]
X = np.random.randn(n, p)
Y = np.random.randn(n, p)
X_ma = ma.array(X, mask=mask_big)
Y_ma = ma.array(Y, mask=mask_big)

Но тогда на моей машине X_ma.dot(Y_ma.T) примерно в 5 раз медленнее, чем X.dot(Y.T) ...

Начнем с того, что я думаю, что проблема в том, что .dot не знает, что маска относится только к p, но я не знаю, возможно ли использовать эту информацию.

Я ищу способ выполнить вычисление, не будучи намного медленнее, чем наивная точка.

Большое спасибо!

Answer 1

Мы можем использовать matrix-multiplication с замаскированными версиями и без них, поскольку замаскированное вычитание из полной версии дает нам желаемый результат -

inner = X.dot(Y.T)-X[:,mask].dot(Y[:,mask].T)

Или просто использовать обращенную маску, хотя будет медленнеедля спарси mask -

inner = X[:,~mask].dot(Y[:,~mask].T)

Сроки -

In [34]: np.random.seed(0)
    ...: p,n = 10000,50
    ...: X = np.random.rand(n,p)
    ...: Y = np.random.rand(n,p)
    ...: mask = np.random.rand(p)>0.95

In [35]: mask.mean()
Out[35]: 0.0507

In [36]: %timeit X.dot(Y.T)-X[:,mask].dot(Y[:,mask].T)
100 loops, best of 3: 2.54 ms per loop

In [37]: %timeit X[:,~mask].dot(Y[:,~mask].T)
100 loops, best of 3: 4.1 ms per loop

In [39]: %%timeit
    ...: inner = np.empty((n,n))
    ...: for i in range(X.shape[0]):
    ...:     for j in range(X.shape[0]):
    ...:         inner[i, j] = np.sum(X[i, ~mask] * Y[j, ~mask])
1 loop, best of 3: 302 ms per loop