从Python中的嵌套for循环中获取插值

Question

我有一个在 (xz) 平面上具有特定形状的 2D 结构。为了简单起见，我在这里将其设置为圆形。我基本上需要绕 z 轴旋转该结构，我的想法是使用插值函数来做到这一点。

RegularGridInterpolator

（链接到文档）听起来适合这个想法，因为我使用给定（xz）平面中的结构作为插值器的输入，然后，当我绕 z 轴旋转时，计算sqrt(x^2 + y^2) 在每个位置（从顶部看，即沿 z 轴），对应于原始 x 坐标，而 z 仍然是 z。

代码很好，但是当数组变大时（每个方向最多 1000 个点），速度非常慢。这很可能是由于计算插值的嵌套 for 循环造成的。我考虑过使用列表理解来实现这一点，但无法让它与这里的插值器一起使用。所以我的问题是，如何摆脱至少一个 for 循环，也许更多？

这是代码：

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 )

通过2D输出和3D输出可以看出，见下图，它可以工作，但是速度很慢。

Answer 1

如果使用

np.meshgrid

将数组广播到 2D/3D，则可以替换循环。

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