在numpy中依赖于2个数组的函数的向量化

Question

I'm trying to vectorize a function that consist of a loop.

The original function is:

def error(X, Y, m, c):
    total = 0
    for i in range(20):
        total += (Y[i]-(m*X[i]+c))**2
    return total 

I've tried the following but It doesn't work:

def error(X, Y, m, c):
    errorVector = np.array([(y-(m*x+c))**2 for (x,y) in (X,Y)])
    total = errorVector.sum()
    return total

How can I vectorize the function?

Answer 1

This is one way, assuming X and Y have first dimension of length 20.

def error(X, Y, m, c):
    total = 0
    for i in range(20):
        total += (Y[i]-(m*X[i]+c))**2
    return total 

def error_vec(X, Y, m, c):
    return np.sum((Y - (m*X + c))**2)

m, c = 3, 4
X = np.arange(20)
Y = np.arange(20)

assert error(X, Y, m, c) == error_vec(X, Y, m, c)

Answer 2

To complement @jpp's answer (which assumes that X and Y both have the shape (20, ...) ), here's an exact equivalent of your error function:

def error(X, Y, m, c):
  return np.sum((Y[:20] - (m * X[:20] + c)) ** 2)