Numpy matrix operations

Question

I want to compute the following values for all i and j:

M_ki = Sum[A_ij - A_ik - A_kj + A_kk, 1 <= j <= n]

How can I do it using Numpy (Python) without an explicit loop?

Thanks!

Answer 1

Here is a general strategy for solving this kind of problems.

First, write a small script, with the loop written explicitly in two different functions, and a test at the end making sure 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 continue staying 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'! No loops in the last one. To vectorize that kind of equations, broadcasting is the way to go!

Warning: make sure that explicit is the equation that you wanted to vectorize. I wasn't sure if the terms that do not depend on j should also be summed over.