如何无循环将按波长分组的2D DataFrame转为3D NumPy数组？

问题背景 #

现有2D DataFrame conc，记录不同波长wl下4个layer的gas1、gas2、gas3浓度数据。原循环方法因数据量过大效率低下，需无循环的高效实现方案。

数据结构 #

conc =  
        wl    gas1  gas2  gas3  layer
    0   5000  10    13    250    1
    1   5000  20    14    260    2
    2   5000  30    15    270    3
    3   5000  40    16    280    4
    4   5001  50    17    290    1
    5   5001  60    18    300    2
    6   5001  70    19    310    3
    7   5001  80    20    320    4
    ...
    497 5125  20    25    650    1
    498 5125  35    15    550    2
    499 5125  55    30    750    3
    500 5125  95    21    650    4

原循环代码（低效） #

wls = set(conc.loc[:,"wl"])
new_3D_array = np.zeros((len(wls), 4, 3))   # 4 layers, 3 gases

for k, wl in enumerate(wls):
    sub_array = conc[conc.loc[:,"wl"]==wl]
    new_3D_array[k,:,:] = sub_array.loc[:,["gas1", "gas2", "gas3"]]

期望输出 #

[[[ 10.  13. 250.]
  [ 20.  14. 260.]
  [ 30.  15. 270.]
  [ 40.  16. 280.]]

 [[ 50.  17. 290.]
  [ 60.  18. 300.]
  [ 70.  19. 310.]
  [ 80.  20. 320.]]
 ....

 [[ 20.  25. 650.]
  [ 35.  15. 550.]
  [ 55.  30. 750.]
  [ 95.  21. 650.]]]

高效无循环实现方案 #

方案1：直接Reshape（效率最高，需数据结构符合要求）

如果数据严格按wl分组，每个wl对应连续4行，且layer按1-4顺序排列，直接用numpy的reshape操作：

import numpy as np

# 先确保数据按wl和layer排序（如果原数据未排序）
conc_sorted = conc.sort_values(["wl", "layer"])
# 提取气体浓度数据
gas_values = conc_sorted[["gas1", "gas2", "gas3"]].values
# 重塑为3D数组：(波长数, 层数, 气体数)
new_3D_array = gas_values.reshape(-1, 4, 3)

此方法完全依赖numpy底层优化，无额外计算开销，速度最快。

方案2：GroupBy + Stack（通用型，无需严格连续）

如果数据分组不连续或排序混乱，用pandas的groupby结合numpy.stack实现：

import numpy as np
import pandas as pd

conc_sorted = conc.sort_values(["wl", "layer"])
# 按wl分组，将每组的气体数据转为数组，再堆叠成3D结构
new_3D_array = np.stack(
    conc_sorted.groupby("wl")[["gas1", "gas2", "gas3"]].apply(np.array).values
)

该方法自动处理分组，兼容性强，效率远高于循环。

方案3：透视表转换（灵活但效率稍低）

通过透视表重构数据后调整维度：

import numpy as np
import pandas as pd

# 构建透视表：行=wl，列=layer，值=各气体浓度
pivot_df = conc.pivot_table(
    index="wl", 
    columns="layer", 
    values=["gas1", "gas2", "gas3"]
)
# 重塑并转置得到目标结构
new_3D_array = pivot_df.values.reshape(len(pivot_df), 3, 4).transpose(0, 2, 1)

适合需要中间透视表进行其他分析的场景，效率略低于前两种方案。

内容的提问来源于stack exchange，提问作者Dr. Paprika