Uso de numba para la similitud de coseno entre un vector y filas en un matix

Encontré esta idea usando numba para un cálculo rápido de la similitud del coseno

import numba @numba.jit(target='cpu', nopython=True) def fast_cosine(u, v): m = u.shape[0] udotv = 0 u_norm = 0 v_norm = 0 for i in range(m): if (np.isnan(u[i])) or (np.isnan(v[i])): continue udotv += u[i] * v[i] u_norm += u[i] * u[i] v_norm += v[i] * v[i] u_norm = np.sqrt(u_norm) v_norm = np.sqrt(v_norm) if (u_norm == 0) or (v_norm == 0): ratio = 1.0 else: ratio = udotv / (u_norm * v_norm) return ratio 

Los resultados parecen prometedores (en mi máquina 500ns vs. solo 200us sin jit decorator).

Me gustaría usar numba para paralelizar este cálculo entre un vector u y una matriz candidata M , es decir, el coseno en cada fila.

Ejemplo:

 def fast_cosine_matrix(u, M): """ Return array of cosine similarity between u and rows in M >>> import numpy as np >>> u = np.random.rand(100) >>> M = np.random.rand(10, 100) >>> fast_cosine_matrix(u, M) """ 

Una forma es simplemente reescribir con la segunda entrada de una matriz. Pero obtengo un NotImplementedError si trato de iterar sobre las filas de una matriz. Voy a intentar simplemente usando rebanadas.

Pensé en usar vectorize pero no puedo hacerlo funcionar.

Alternativa: hacer un UFunc generalizado con numba.

 @numba.guvectorize(["void(float64[:], float64[:], float64[:])"], "(n),(n)->()", target='parallel') def fast_cosine_gufunc(u, v, result): m = u.shape[0] udotv = 0 u_norm = 0 v_norm = 0 for i in range(m): if (np.isnan(u[i])) or (np.isnan(v[i])): continue udotv += u[i] * v[i] u_norm += u[i] * u[i] v_norm += v[i] * v[i] u_norm = np.sqrt(u_norm) v_norm = np.sqrt(v_norm) if (u_norm == 0) or (v_norm == 0): ratio = 1.0 else: ratio = udotv / (u_norm * v_norm) result[:] = ratio u = np.random.rand(100) M = np.random.rand(100000, 100) fast_cosine_gufunc(u, M[0,:]) fast_cosine_gufunc(u, M) 

Solución reescribiéndolo un poco:

 import numpy as np import numba @numba.jit(target='cpu', nopython=True, parallel=True) def fast_cosine_matrix(u, M): scores = np.zeros(M.shape[0]) for i in numba.prange(M.shape[0]): v = M[i] m = u.shape[0] udotv = 0 u_norm = 0 v_norm = 0 for j in range(m): if (np.isnan(u[j])) or (np.isnan(v[j])): continue udotv += u[j] * v[j] u_norm += u[j] * u[j] v_norm += v[j] * v[j] u_norm = np.sqrt(u_norm) v_norm = np.sqrt(v_norm) if (u_norm == 0) or (v_norm == 0): ratio = 1.0 else: ratio = udotv / (u_norm * v_norm) scores[i] = ratio return scores u = np.random.rand(100) M = np.random.rand(100000, 100) fast_cosine_matrix(u, M)