Loopless chi-square grid search in Python numpy

Question

Best-fit linear parameters A and B (y=Ax+b) correspond to the minimum of the chi-square function over these parameters. I want to do a brute force grid search for the global chi-square minimum (guaranteed because 2-parameter linear chi-square is a paraboloid) and have achieved it with 3 nested loops (below) but want to avoid loops (i.e., vectorize using Numpy's array broadcasting properties).

Chi-square (weighted least squares) is defined as:

Chi-square(k,j) = sum (y[i]-(A[k]*x[i]+B[j]))/yerr[i])^2

Below is Python Numpy code that fills a 100 x 100 grid with chi-square values over the 10,000 combinations of A and B parameter values (100 values each). There are three data arrays: x, y and yerr.

Thanks for any help towards a loopless version of a 2-parameter linear chi-square grid search in Python Numpy.

Note I want to do a grid search and not use scipy.optimize.minimize -- thanks!

Keith

# create parameter grid
a = np.linspace(80,120,100)
b = np.linspace(10,40,100)
A,B = np.meshgrid(a,b)

# calculate chi-square over parameter grid
chi2=np.zeros((100,100))

for k in range(100):
    for j in range(100):
        for i in range (len(y)):
            chi2a = ((y[i]-a[k]*x[i]-b[j])/yerr[i])**2;
            chi2[k,j]+=chi2a;

Indentation for the last line looks wrong. Also, shouldn't it be chi2[k,j]+= chi2a instead? — Divakar
Yes to both comments (now edited). Thank you for catching this! — Carey
Think we need two more tab indentations there, as we need to sum update for each computation of chi2a. — Divakar
Also, range(1,100) should be range(100) and son on to cover all elements. — Divakar
Yes, thank you! Just tested the script with your corrections and it appears to run fine now. Thanks again. — Carey

Divakar Divakar · Accepted Answer · 2017-05-28T09:18:23

Here's one vectorized approach making use of NumPy broadcasting -

subs = (y-a[:,None,None]*x-b[:,None])/yerr
chi2 = (subs**2).sum(2)

Here's another and possibly faster one with np.einsum -

chi2 = np.einsum('ijk,ijk->ij',subs,subs) #subs from previous one

Loopless chi-square grid search in Python numpy

1 Answers