Fast Weighted Scatter Matrix Calculation

In this question six months ago, jez was nice enough to help me come up with a fast approximation for the outer product of row differences, i.e.: K = np.zeros((len(X), len(X))) for

Solution 1:

Here is a way to vectorize the calculation you specified. If you do a lot of this kind of thing, then it may be worth learning how to use, "numpy.tensordot". It multiplies all elements according to standard numpy broadcasting, and then it sums over the pairs of axes given with the kwrd, "axes".

Here is the code:

# Importsimport numpy as np
from numpy.random import random

# Original calculation for testing purposes defftrue(A, X):
  ""
  K = np.zeros((len(X), len(X)))
  KA_true = np.zeros((len(X), len(X)))

  for i, Xi inenumerate(X):
    for j, Xj inenumerate(X):
      dij = Xi - Xj
      K += np.outer(dij, dij)
      KA_true += A[i, j] * np.outer(dij, dij) 
  return ftrue

# Better: No Python loops. But, makes a large temporary array.deffbetter(A, X):
  ""
  c = X[:, None, :] - X[None, :, :]
  b = A[:, :, None] * c           # ! BAD ! temporary array size N**3
  KA_better = np.tensordot(b, c, axes = [(0,1),(0,1)])
  return KA_better

# Best way: No Python for loops. No large temporary arraysdeffbest(A, X):
  ""
  KA_best = np.tensordot(A.sum(1)[:,None] * X, X, axes=[(0,), (0,)])
  KA_best += np.tensordot(A.sum(0)[:,None] * X, X, axes=[(0,), (0,)])
  KA_best -= np.tensordot(np.dot(A, X), X, axes=[(0,), (0,)])
  KA_best -= np.tensordot(X, np.dot(A, X), axes=[(0,), (0,)])
  return KA_best


# Test scriptif __name__ == "__main__":

  # Parameters for the computation 
  N = 250
  X = random((N, N))
  A = random((N, N))

  # Print the error
  KA_better = fbetter(A, X)
  KA_best = fbest(A, X)

  # Test against true if array size isn't too bigif N<100:
    KA_true = ftrue(A, X)
    err = abs(KA_better - KA_true).mean()
    msg = "Mean absolute difference (better): {}."print(msg.format(err))

  # Test best against better
  err = abs(KA_best - KA_better).mean()
  msg = "Mean absolute difference (best): {}."print(msg.format(err))

My first attempt (fbetter) made a large temporary array of size NxNxN. The second attempt (fbest) never makes anything bigger than NxN. This works pretty good up to N~1000.

Also, the code runs faster when the output array is smaller.

I have MKL installed so the calls to tensordot are really fast and run in parallel.

Thanks for the question. This was a nice exercise and reminded me how important it is to avoid making large temporary arrays.