Sum of Square Differences (SSD) in numpy/scipy

Question

I'm trying to use Python and Numpy/Scipy to implement an image processing algorithm. The profiler tells me a lot of time is being spent in the following function (called often), which tells me the sum of square differences between two images

def ssd(A,B):
    s = 0
    for i in range(3):
        s += sum(pow(A[:,:,i] - B[:,:,i],2))
    return s

How can I speed this up? Thanks.

Answer 1

Just

s = numpy.sum((A[:,:,0:3]-B[:,:,0:3])**2)

(which I expect is likely just sum((A-B)**2) if the shape is always (,,3))

You can also use the sum method: ((A-B)**2).sum()

Right?

Answer 2

Just to mention that one can also use np.dot:

def ssd(A,B):
  dif = A.ravel() - B.ravel()
  return np.dot( dif, dif )

This might be a bit faster and possibly more accurate than alternatives using np.sum and **2, but doesn't work if you want to compute ssd along a specified axis. In that case, there might be a magical subscript formula using np.einsum.