[Python] Speeding up 3D interpolation

I am using python to perform the homotopy formula to obtain a vector potential A of a magnetic field B at points x in 3D space:

A(x) = integral_L=0,1 (B(Lx)cross(Lx)dL)

where cross represents the cross product

Here B(x) is only known on a discrete regular grid (constant resolution in each direction, but the resolution in x,y,z are not necessarily the same) and so interpolation is required to perform the line integral.

I was able to vectorize the loop over every point in space x with the following code and using RegularGridInterpolator:

#stuff related to 3D linear interpolation
print("Setting up interpolation stuff")

interpolator_x = RegularGridInterpolator((X-x0,Y-y0,Z-z0),Bx, method='linear')
interpolator_y = RegularGridInterpolator((X-x0,Y-y0,Z-z0),By, method='linear')
interpolator_z = RegularGridInterpolator((X-x0,Y-y0,Z-z0),Bz, method='linear')
print("created interpolator")

#find an A from B using homotopy
nlamb = 101
lamb=np.arange(nlamb)/(nlamb-1)
dlamb = lamb[1]-lamb[0]
print("looping over lambda")
start = time.time()
dxi = x3D-x0
dyi = y3D-y0
dzi = z3D-z0
for il in np.arange(nlamb):
print(il, '/', nlamb)
X1 = lamb[il]*dxi
Y1 = lamb[il]*dyi
Z1 = lamb[il]*dzi
Bx_values = interpolator_x((X1,Y1,Z1))
By_values = interpolator_y((X1,Y1,Z1))
Bz_values = interpolator_z((X1,Y1,Z1))
Ax_hom = Ax_hom + (((By_values*Z1)-(Bz_values*Y1))*dlamb)
Ay_hom = Ay_hom + (((Bz_values*X1)-(Bx_values*Z1))*dlamb)
Az_hom = Az_hom + (((Bx_values*Y1)-(By_values*X1))*dlamb)
print("time to do loop over lambda :"+str(time.time()-start))


For my test data of 128x128*82 points this takes about 110s.

I though I should be able to vectorize the integral also. ChatGPT gave me this

#chaGPT attempt to vectorize over lambda - acually slower
start = time.time()
lamb = np.linspace(0, 1, nlamb)
X1 = lamb[:, None, None, None] * (x3D - x0)
Y1 = lamb[:, None, None, None] * (y3D - y0)
Z1 = lamb[:, None, None, None] * (z3D - z0)

# Reshape to (N, points) for interpolation
points = np.stack([X1, Y1, Z1], axis=-1).reshape(-1, 3)

# Interpolate all at once
Bx_values = interpolator_x(points).reshape(nlamb, *x3D.shape)
By_values = interpolator_y(points).reshape(nlamb, *x3D.shape)
Bz_values = interpolator_z(points).reshape(nlamb, *x3D.shape)

# Compute Ax, Ay, Az using broadcasting
Ax_hom = np.sum(((By_values * Z1) - (Bz_values * Y1)) * dlamb, axis=0)
Ay_hom = np.sum(((Bz_values * X1) - (Bx_values * Z1)) * dlamb, axis=0)
Az_hom = np.sum(((Bx_values * Y1) - (By_values * X1)) * dlamb, axis=0)
print("trying to vectorize integral over lambda:"+str(time.time()-start))


However, for the test grid is is actually slowr (~140s), and for larger grids it stays slower than the loop over lambda version.

This is because the interpolation is the slowest part of the calculation. So, I think my problem is "How can I speed up this 3D interpolation"?

Thanks

Continue reading...