这是解决此类问题的一般策略.
首先,编写一个小脚本,将循环显式写入两个不同的函数,最后进行测试,确保两个函数完全相同:
import numpy as np
from numpy import newaxis
def explicit(a):
n = a.shape[0]
m = np.zeros_like(a)
for k in range(n):
for i in range(n):
for j in range(n):
m[k,i] += a[i,j] - a[i,k] - a[k,j] + a[k,k]
return m
def implicit(a):
n = a.shape[0]
m = np.zeros_like(a)
for k in range(n):
for i in range(n):
for j in range(n):
m[k,i] += a[i,j] - a[i,k] - a[k,j] + a[k,k]
return m
a = np.random.randn(10,10)
assert np.allclose(explicit(a), implicit(a), atol=1e-10, rtol=0.)
然后,通过编辑隐式来逐步向量化函数,在每一步运行脚本以确保它们继续保持不变:
步骤1
def implicit(a):
n = a.shape[0]
m = np.zeros_like(a)
for k in range(n):
for i in range(n):
m[k,i] = (a[i,:] - a[k,:]).sum() - n*a[i,k] + n*a[k,k]
return m
第2步
def implicit(a):
n = a.shape[0]
m = np.zeros_like(a)
m = - n*a.T + n*np.diag(a)[:,newaxis]
for k in range(n):
for i in range(n):
m[k,i] += (a[i,:] - a[k,:]).sum()
return m
第3步
def implicit(a):
n = a.shape[0]
m = np.zeros_like(a)
m = - n*a.T + n*np.diag(a)[:,newaxis]
m += (a.T[newaxis,...] - a[...,newaxis]).sum(1)
return m
瞧瞧’!最后一个没有循环.为了矢量化这种方程,broadcasting就是要走的路!
警告:确保显式是您想要向量化的等式.我不确定是否也应该总结不依赖于j的术语.