如何从稀疏格式数据高效构建3D numpy矩阵？

我有一份外部程序生成的pandas DataFrame，以“稀疏矩阵”形式存储(x,y,z)坐标对应的温度值——并非所有坐标组合都有对应值（比如(0,-0.25,0)和(5.1,0.0,-0.21)缺失）。需要将其重组为3D numpy.ndarray。之前尝试用scipy.sparse.csr_matrix（仅支持2D），通过循环堆叠2D矩阵再转密集矩阵，可行但速度极慢，求更高效的实现方案。

数据示例 #

In []: print(data)
Out[]:
          x     y     z        T
0      0.05 -0.25  0.00  1516.61
1      0.10 -0.25  0.00  1572.59
2      0.15 -0.25  0.00  1590.58
3      0.20 -0.25  0.00  1615.54
4      0.25 -0.25  0.00  1629.69
    ...   ...   ...      ...
18647  4.25  0.00 -0.21  1813.54
18648  4.30  0.00 -0.21  1821.97
18649  4.35  0.00 -0.21  1814.80
18650  4.40  0.00 -0.21  1822.60
18651  4.45  0.00 -0.21  1813.72

[18652 rows x 4 columns]

In []: x = data.x.unique()

In []: y = data.y.unique()

In []: z = data.z.unique()

In []: print(len(x)*len(y)*len(z))
Out[]: 25894

尝试的低效解决方案（硬编码部分数值） #

def get_dense_matrix(data,tol=1e-6):
    x = np.arange(data.x.min(),data.x.max()+2*tol,0.05)
    y = np.arange(data.y.min(),data.y.max()+2*tol,0.05)
    z = np.arange(data.z.min(),data.z.max()+2*tol,0.05)
    T = np.zeros((len(x),len(y),len(z)))
    for i,_z in enumerate(z):
        zdata = data[np.abs(data.z-_z)<tol]
        row = np.round((zdata.x+zdata.x.min())*20)
        col = np.round((zdata.y+zdata.y.min())*20)
        T_ = zdata['T']
        A = csr_matrix((T_, (row, col)), shape=(len(x), len(y)))
        A = A.toarray()
        T[:,:,i] = A
    return np.expand_dims(T,0)

高效实现方案 #

核心思路 #

放弃循环处理单一层级的方式，用向量化操作+坐标索引映射直接完成3D数组赋值，避免稀疏矩阵转换和循环开销，所有步骤基于numpy/pandas的原生向量化接口，速度提升显著。

代码实现 #

import numpy as np
import pandas as pd

def get_dense_matrix_fast(data, tol=1e-6):
    # 获取排序后的唯一坐标值，确保网格顺序一致
    x_vals = np.sort(data['x'].unique())
    y_vals = np.sort(data['y'].unique())
    z_vals = np.sort(data['z'].unique())
    
    # 将原始坐标映射为3D数组的索引
    x_idx = np.searchsorted(x_vals, data['x'], side='left')
    y_idx = np.searchsorted(y_vals, data['y'], side='left')
    z_idx = np.searchsorted(z_vals, data['z'], side='left')
    
    # 处理浮点精度误差，筛选出匹配正确的坐标
    x_match = np.abs(x_vals[x_idx] - data['x']) < tol
    y_match = np.abs(y_vals[y_idx] - data['y']) < tol
    z_match = np.abs(z_vals[z_idx] - data['z']) < tol
    valid_mask = x_match & y_match & z_match
    
    # 初始化3D数组，缺失值默认填充0，可按需改为NaN
    dense_T = np.zeros((len(x_vals), len(y_vals), len(z_vals)))
    
    # 直接通过索引赋值有效温度数据
    dense_T[x_idx[valid_mask], y_idx[valid_mask], z_idx[valid_mask]] = data['T'][valid_mask]
    
    # 保持与原方案一致的输出维度（添加一个前置维度）
    return np.expand_dims(dense_T, axis=0)

关键优化点 #

1. 向量化替代循环：所有坐标映射、数据赋值操作均为numpy向量化操作，比Python循环快10~100倍
2. 移除稀疏矩阵转换：直接通过索引赋值，省去中间稀疏矩阵的创建与转密集过程，减少内存和计算开销
3. 通用坐标映射：用np.searchsorted替代硬编码步长计算，适配任意间隔的坐标数据，避免原方案中*20这类依赖特定步长的代码

可选极致优化（坐标无精度误差场景） #

如果原始坐标值是精确的（浮点误差可忽略），可以用pd.factorize更快完成坐标到索引的映射，省去精度校验步骤：

x_idx, x_vals = pd.factorize(data['x'], sort=True)
y_idx, y_vals = pd.factorize(data['y'], sort=True)
z_idx, z_vals = pd.factorize(data['z'], sort=True)

dense_T = np.zeros((len(x_vals), len(y_vals), len(z_vals)))
dense_T[x_idx, y_idx, z_idx] = data['T']
return np.expand_dims(dense_T, axis=0)

内容的提问来源于stack exchange，提问作者F. Heath