如何在Python3中加速3D矩阵的赋值操作？

背景 #

我使用Python3处理Fortran90输出的3D数据集以实现可视化。Fortran将3D矩阵以2D文本格式输出，每行结构为：

Value,  x, y, z

示例：

123443.0,  1, 1, 1
   123343.0,  1, 1, 2
   134554.0,  1, 1, 3

其中Value是3D矩阵的元素值，x、y、z对应元素在3D矩阵中的位置。我原本的读取和赋值代码如下：

input_path="/user_path/3D_CUBE_data.txt"

#READ DATASET
read_1=np.loadtxt(input_path)
read_2 = pd.DataFrame(read_1, columns=["data","x","y","z"])

print("input_shape=", read_2.shape)

#CREATE A 3D MATRIX
#(dis, tra, rows)
dis = 41
tra = 18
rows = 4096

data_1 = np.zeros( shape = (dis, tra, rows) )


x_l = list(range(1,dis+1))
y_l = list(range(1,tra+1))
z_l = list(range(1,rows+1))

#ASSIGN VALUES IN THE 3D MATRIX
for y in y_l :
   for x in x_l :
       line_1 = read_2[(read_2["x"]==int(x)) & (read_2["y"]==int(y))]

       data_1[x-1,y-1,:] = line_1.loc[:,["data"]].T

我认为嵌套的两层for循环是赋值过程速度慢的主要原因，希望找到加速这个过程的方法。

优化方案 #

1. 利用NumPy高级索引直接赋值（最推荐） #

直接提取x、y、z的0-based索引，一次性将数据值映射到3D矩阵对应位置，完全规避Python循环：

import numpy as np

input_path="/user_path/3D_CUBE_data.txt"

# 读取数据
read_1 = np.loadtxt(input_path, delimiter=",")
# 拆分出数据和索引列
data_vals = read_1[:, 0]
# 转换为0-based索引（原数据是1-based）
x_indices = (read_1[:, 1] - 1).astype(int)
y_indices = (read_1[:, 2] - 1).astype(int)
z_indices = (read_1[:, 3] - 1).astype(int)

# 创建3D矩阵
dis = 41
tra = 18
rows = 4096
data_1 = np.zeros((dis, tra, rows))

# 一次性完成赋值
data_1[x_indices, y_indices, z_indices] = data_vals

该方法依托NumPy底层的C实现，速度比Python循环快数个数量级。

2. 用Pandas分组操作替代嵌套循环 #

若需保留DataFrame处理流程，可按x、y分组后批量赋值：

import numpy as np
import pandas as pd

input_path="/user_path/3D_CUBE_data.txt"

read_2 = pd.read_csv(input_path, sep=",", header=None, names=["data","x","y","z"], skipinitialspace=True)

dis = 41
tra = 18
rows = 4096
data_1 = np.zeros((dis, tra, rows))

# 按x、y分组后直接赋值
for (x_val, y_val), group in read_2.groupby(["x", "y"]):
    x_idx = int(x_val) - 1
    y_idx = int(y_val) - 1
    # 确保分组内z值是连续有序的（与原文件输出顺序一致）
    data_1[x_idx, y_idx, :] = group["data"].values

分组操作的效率优于嵌套循环，但仍不及纯NumPy索引方案。

3. 优化数据读取步骤 #

针对大文件，替换np.loadtxt为更高效的读取方法，减少整体运行时间：

# 用Pandas读取，自动处理逗号分隔和空格
read_2 = pd.read_csv(input_path, sep=",", header=None, names=["data","x","y","z"], skipinitialspace=True)

# 或用NumPy的genfromtxt
read_1 = np.genfromtxt(input_path, delimiter=",")

4. 提前指定数据类型 #

读取时直接指定各列的数据类型，避免后续类型转换开销：

read_1 = np.loadtxt(
    input_path,
    dtype=[('data', 'float64'), ('x', 'int32'), ('y', 'int32'), ('z', 'int32')],
    delimiter=','
)
data_vals = read_1['data']
x_indices = read_1['x'] - 1
y_indices = read_1['y'] - 1
z_indices = read_1['z'] - 1

内容的提问来源于stack exchange，提问作者Changwan Sun