Python实现GPR数据3D网格生成for循环过慢，如何优化？

我编写了一段用于生成GPR数据3D网格的Python代码，但其运行耗时极长；此前使用的MATLAB版本同类代码运行速度并未如此缓慢。尝试通过完全向量化改造以移除for循环，却遭遇内存不足问题。请问如何大幅提升代码运行速度？是否应当放弃使用Python？

import pandas as pd
import numpy as np

# 加载CSV文件
csv_file = 'C:/Users/nad/Downloads/Yard/070/interpolated_points.csv'
df = pd.read_csv(csv_file)

# 提取坐标（单位为米）
X_coords = df['X'].values
Z_coords = df['Z'].values
Y_coords = df['Y'].values

# 定义单元格尺寸（单位为米）
cell_size = 0.0048  # 单元格尺寸（米）

# 定义立方体的单元格数量维度
depth_cells = 104
width_cells = depth_cells * 2
length_cells = depth_cells * 2

# 将立方体维度转换为实际单位（米）
depth = depth_cells * cell_size
width = width_cells * cell_size
length = length_cells * cell_size

# 确定整个网格的边界
min_x, max_x = X_coords.min() - width / 2, X_coords.max() + width / 2
min_z, max_z = Z_coords.min(), Z_coords.max() + depth
min_y, max_y = Y_coords.min() - length / 2, Y_coords.max() + length / 2

# 转换为整数
grid_size_x = int((max_x - min_x) / cell_size)
grid_size_z = int((max_z - min_z) / cell_size)
grid_size_y = int((max_y - min_y) / cell_size)

# 创建带有实际坐标的3D网格
x_coords_grid = np.linspace(min_x, max_x, grid_size_x)
z_coords_grid = np.linspace(min_z, max_z, grid_size_z)
y_coords_grid = np.linspace(min_y, max_y, grid_size_y)

# 创建x、y、z坐标的3D网格矩阵
X_grid, Z_grid, Y_grid = np.meshgrid(x_coords_grid, z_coords_grid, y_coords_grid, indexing='ij')

# 加载GPR数据矩阵（假设列数与CSV中的点数一致）
gpr_data = np.loadtxt('C:/Users/nad/Downloads/Yard/070/gpr_data.csv', delimiter=',')

# 定义GPR数据的距离间隔，使用更新后的单元格尺寸作为分辨率
gpr_data_resolution = cell_size  # 设置为单元格尺寸

# 初始化空的3D数组，用于存储每个网格单元的GPR数据总和
gpr_grid_data = np.zeros((grid_size_x, grid_size_z, grid_size_y))

# 对每个点进行向量化距离计算
for point_index, (x_point, z_point, y_point) in enumerate(zip(X_coords, Z_coords, Y_coords)):
    # 计算所有网格单元到当前点的欧氏距离
    distance_grid = 2 * np.sqrt((X_grid - x_point) ** 2 +
                            (Z_grid - z_point) ** 2 +
                            (Y_grid - y_point) ** 2)
    
    # 确定对应此距离的GPR数据行
    gpr_data_rows = (distance_grid / gpr_data_resolution).astype(int)

    # 掩码确保有效的行索引
    valid_mask = (gpr_data_rows >= 0) & (gpr_data_rows < gpr_data.shape[0])
    
    # 提取有效索引对应的GPR数据，无效值设为0
    gpr_data_values = np.zeros_like(distance_grid)
    
    # 提取有效索引的GPR数据
    gpr_data_values[valid_mask] = gpr_data[gpr_data_rows[valid_mask], point_index]

    # 将当前点的GPR数据累加到最终3D网格中
    gpr_grid_data += gpr_data_values

# 现在，gpr_grid_data包含了所有点贡献求和后的每个单元的GPR数据

完全没必要放弃Python，通过以下几个方向优化，就能大幅提升代码运行效率：

1. 避免全网格重复计算，改用索引映射减少运算量 #

当前代码每次循环都对整个3D网格计算欧氏距离，这是最大性能瓶颈。可以先把点坐标转换为网格索引，利用索引相对位置计算距离，避免重复生成巨型3D数组：

# 提前计算所有点对应的网格索引
point_x_idx = ((X_coords - min_x) / cell_size).astype(int)
point_z_idx = ((Z_coords - min_z) / cell_size).astype(int)
point_y_idx = ((Y_coords - min_y) / cell_size).astype(int)

# 预先生成网格索引的广播数组
x_idx_grid = np.arange(grid_size_x)[:, np.newaxis, np.newaxis]
z_idx_grid = np.arange(grid_size_z)[np.newaxis, :, np.newaxis]
y_idx_grid = np.arange(grid_size_y)[np.newaxis, np.newaxis, :]

循环内用索引差计算距离，省去浮点坐标运算的冗余：

for point_idx, (x_idx, z_idx, y_idx) in enumerate(zip(point_x_idx, point_z_idx, point_y_idx)):
    # 先计算索引距离，再转换为实际距离
    idx_dist = np.sqrt((x_idx_grid - x_idx)**2 + (z_idx_grid - z_idx)**2 + (y_idx_grid - y_idx)**2)
    distance_grid = 2 * idx_dist * cell_size
    # 后续逻辑保持不变，但运算量大幅降低

2. 分块批量处理，平衡向量化与内存占用 #

完全向量化会生成(N_points, grid_size_x, grid_size_z, grid_size_y)的巨型数组，必然内存溢出。采用小批量向量化，既利用numpy的向量化优势，又控制内存占用：

batch_size = 10  # 根据自身内存调整大小
for i in range(0, len(X_coords), batch_size):
    batch_x = X_coords[i:i+batch_size]
    batch_z = Z_coords[i:i+batch_size]
    batch_y = Y_coords[i:i+batch_size]
    batch_gpr = gpr_data[:, i:i+batch_size]
    
    # 利用广播批量计算距离
    dists = 2 * np.sqrt((X_grid[np.newaxis, ...] - batch_x[:, np.newaxis, np.newaxis, np.newaxis])**2 +
                        (Z_grid[np.newaxis, ...] - batch_z[:, np.newaxis, np.newaxis, np.newaxis])**2 +
                        (Y_grid[np.newaxis, ...] - batch_y[:, np.newaxis, np.newaxis, np.newaxis])**2)
    rows = (dists / gpr_data_resolution).astype(int)
    valid = (rows >=0) & (rows < gpr_data.shape[0])
    
    # 批量提取GPR值并累加
    batch_vals = np.zeros_like(dists)
    batch_vals[valid] = batch_gpr[rows[valid], np.repeat(np.arange(batch_size), np.sum(valid, axis=(1,2,3)))]
    gpr_grid_data += batch_vals.sum(axis=0)

3. 用Numba编译加速循环 #

如果循环逻辑无法简化，用Numba的JIT编译将Python循环转换为机器码，速度接近C：

from numba import jit

@jit(nopython=True, parallel=True)  # 启用多线程并行
def compute_gpr_grid(X_coords, Z_coords, Y_coords, gpr_data, cell_size, min_x, min_z, min_y, grid_size_x, grid_size_z, grid_size_y, gpr_data_resolution):
    gpr_grid_data = np.zeros((grid_size_x, grid_size_z, grid_size_y))
    max_gpr_rows = gpr_data.shape[0]
    for point_idx in range(len(X_coords)):
        x_point = X_coords[point_idx]
        z_point = Z_coords[point_idx]
        y_point = Y_coords[point_idx]
        # 直接遍历网格单元，避免生成巨型数组
        for i in range(grid_size_x):
            x_grid = min_x + i * cell_size
            dx_sq = (x_grid - x_point) ** 2
            for j in range(grid_size_z):
                z_grid = min_z + j * cell_size
                dz_sq = (z_grid - z_point) ** 2
                for k in range(grid_size_y):
                    y_grid = min_y + k * cell_size
                    dy_sq = (y_grid - y_point) ** 2
                    distance = 2 * np.sqrt(dx_sq + dz_sq + dy_sq)
                    row = int(distance / gpr_data_resolution)
                    if 0 <= row < max_gpr_rows:
                        gpr_grid_data[i,j,k] += gpr_data[row, point_idx]
    return gpr_grid_data

# 调用加速后的函数
gpr_grid_data = compute_gpr_grid(X_coords, Z_coords, Y_coords, gpr_data, cell_size, min_x, min_z, min_y, grid_size_x, grid_size_z, grid_size_y, gpr_data_resolution)

4. 优化数据加载与内存使用 #

• 用np.load替代np.loadtxt加载GPR数据：先将csv转成npy格式，加载速度提升数倍：

  # 仅需执行一次：np.save('gpr_data.npy', gpr_data)
  gpr_data = np.load('gpr_data.npy')
  

• 降低数组精度：如果业务允许，用float32代替float64，减少一半内存占用：

  gpr_data = gpr_data.astype(np.float32)
  gpr_grid_data = np.zeros((grid_size_x, grid_size_z, grid_size_y), dtype=np.float32)
  

内容的提问来源于stack exchange，提问作者Nihat Alperen Dayanır