Here is a general strategy for solving this kind of problem.
First write a small script, a loop written explicitly in two different functions, and a test at the end to verify that the two functions are exactly the same:
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.)
Then vectorize the function step by step by editing implicit, running the script at each step to make sure that they remain the same:
Step 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
Step 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
Step 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
Et voila '! . , broadcasting - !
: , explicit - , . , , j.