Getting interpolation out of nested for loops in python

Question

I have a 2D structure with a certain shape in the (xz) plane. For simplicity, I set it here to be of circular shape. I basically need to rotate that structure around the z axis and my idea was to do that with an interpolation function. RegularGridInterpolator (link to the docs) sounds suitable for this having the idea in mind that I use the structure in the given (xz) plane as input to the interpolater, and then, when I rotate around the z axis, calculate sqrt(x^2 + y^2) at each position (looking from the top, i.e. along the z-axis), corresponding to the original x coordinate, and z simply is still z.

The code is fine, but it is very slow when the arrays become large (up to 1000 points in each direction). That is very likely due to the nested for loops, in which the interpolation is calculated. I thought about using list comprehensions for that, but can't get it to work with the interpolator here. So my question is, how to get rid of at least one of the for loops, maybe more?

Here is the code:

import matplotlib.pyplot as plt
from mayavi import mlab
import numpy as np
import scipy.interpolate as interp

def make_simple_2Dplot( data2plot, xVals, yVals, N_contLevels=8 ):
    fig, ax = plt.subplots()

    # the necessity for a transposed arrays confuses me...
    ax.contourf(x_arr, z_arr, data2plot.T)
    ax.set_aspect('equal')
    ax.set_xlabel('x')
    ax.set_ylabel('z')

    plt.show()

def make_simple_3Dplot( data2plot, xVals, yVals, zVals, N_contLevels=8 ):
    contLevels  = np.linspace( np.amin(data2plot), 
                               np.amax(data2plot),
                               N_contLevels)[1:].tolist()
    fig1    = mlab.figure( bgcolor=(1,1,1), fgcolor=(0,0,0),size=(800,600))
    contPlot = mlab.contour3d( data2plot, contours=contLevels,
                               transparent=True, opacity=.4,
                               figure=fig1
                             )
    mlab.xlabel('x')
    mlab.ylabel('y')
    mlab.zlabel('z')

    mlab.show()

x_min, z_min = 0, 0
x_max, z_max = 10, 10
Nx = 100
Nz = 50

x_arr = np.linspace(x_min, x_max, Nx)
z_arr = np.linspace(z_min, z_max, Nz)

# center of circle in 2D
xc, zc = 5, 5

# radius of circle
rc = 2

# make 2D circle
data_2D = np.zeros( (Nx,Nz) )
for ii in range(Nx):
    for kk in range(Nz):
        if np.sqrt((x_arr[ii]-xc)**2 + (z_arr[kk]-zc)**2) < rc:
            data_2D[ii,kk] = 1

# interpolation function to make 3D object
circle_xz = interp.RegularGridInterpolator( (x_arr,z_arr), data_2D, 
                                            bounds_error=False,
                                            fill_value=0
                                          )

# coordinate arrays for 3D data
y_min = -x_max
y_max = x_max
Ny = 100
x_arr_3D = np.linspace(-x_max, x_max, Nx)
y_arr_3D = np.linspace(y_min, y_max, Ny)
z_arr_3D = np.linspace(z_min, z_max, Nz)

# make 3D circle
data_3D = np.zeros( (Nx, Ny, Nz) )
for ii in range(Nx):
    for jj in range(Ny):
        # calculate R corresponding to x in (xz) plane
        R = np.sqrt(x_arr_3D[ii]**2 + y_arr_3D[jj]**2)
        for kk in range(Nz):
            # hiding the interpolator deep in the nested for loop
            # is probably not very clever
            data_3D[ii,jj,kk] = circle_xz( (R, z_arr_3D[kk]) )

make_simple_2Dplot( data_2D, x_arr, z_arr, N_contLevels=8 )
make_simple_3Dplot( data_3D, x_arr_3D, y_arr_3D, z_arr_3D )

As can be seen by the 2D output and the 3D output, see below, it works, but it is very slow.

Answer 1

So my question is, how to get rid of at least one of the for loops, maybe more?

You can eliminate all the loops if you use np.meshgrid to broadcast the arrays to 2D/3D.

# make 2D circle
x, z = np.meshgrid(x_arr, z_arr, indexing='ij')
data_2D = np.sqrt((x-xc)**2 + (z-zc)**2) < rc
# data_2D = data_2D.astype(np.float64) if you want the plot to look the same

To see what is going on, look at the first 3 x 3 subarray of x and z:

x[:3, :3]
# array([[0.       , 0.       , 0.       ],
#        [0.1010101, 0.1010101, 0.1010101],
#        [0.2020202, 0.2020202, 0.2020202]])
z[:3, :3]
# array([[0.        , 0.20408163, 0.40816327],
#        [0.        , 0.20408163, 0.40816327],
#        [0.        , 0.20408163, 0.40816327]])

These 2D arrays represent every combination of values in x_arr and z_arr, then the operations are done elementwise on the resulting arrays. The second line produces a boolean array (e.g. Trues instead of 1s), but you can convert that to float64 if you want the plot to look like it did before.

You can use the same approach for the 3D plot; meshgrid accepts an arbitrary number of arguments.

# make 3D circle
x, y, z = np.meshgrid(x_arr_3D, y_arr_3D, z_arr_3D, indexing='ij')
R = np.sqrt(x**2 + y**2)
points = np.vstack((R.ravel(), z.ravel())).T
data_3D = circle_xz(points).reshape(R.shape)

ravel and vstack are used is because the __call__ method of the RegularGridInterpolator accepts m x n arrays of points at which to perform the interpolation, where m is the number of points (500000) and n is the number of dimensions (2). After evaluation, we reshape back to the original shape.