在数组中重塑数组

Reshape Array in Array in Array

我打开了一个包含 2000 个条目的根文件,子条目的数量可变,并且每一列中都有一个不同的变量。可以说我只对其中的 5 个感兴趣。我想将它们放在 np.shape(array)=(2000,250,5) 的数组中。 250 足以包含每个条目的所有子条目。

根文件通过uproot转换成字典 DATA=[变量名:[条目数组[子条目数组]]

我创建了一个数组 np.zeros(2000,250,5) 并用我想要的数据填充它,但这需要大约 500 毫秒,我需要一个可以扩展的解决方案,因为我以后的目标是 100 万个条目。我找到了多种解决方案,但我的最低解决方案约为 300 毫秒

lim_i=len(N_DATA["nTrack"])
i=0
INPUT_ARRAY=np.zeros((lim_i,500,5))
for l in range(len(INPUT_ARRAY)):
    while i < lim_i:
        EVENT=np.zeros((500,5))
        k=0
        lim_k=len(TRACK_DATA["Track_pt"][i])
        while k<lim_k:
            EVENT[k][0]=TRACK_DATA["Track_pt"][i][k]
            EVENT[k][1]=TRACK_DATA["Track_phi"][i][k]
            EVENT[k][2]=TRACK_DATA["Track_eta"][i][k]
            EVENT[k][3]=TRACK_DATA["Track_dxy"][i][k]
            EVENT[k][4]=TRACK_DATA["Track_charge"][i][k]
            k+=1
        INPUT_ARRAY[i]=EVENT
        i+=1
INPUT_ARRAY

观察 1:我们可以直接分配给 INPUT_ARRAY[i] 的适当子数组,而不是创建 EVENT 作为 INPUT_ARRAY[i] 的代理,然后将其复制进去。(我还将以小写形式设置您的变量名称,以遵循正常约定。

lim_i = len(n_data["nTrack"])
i = 0
input_array = np.zeros((lim_i,500,5))
for l in range(len(input_array)):
    while i < lim_i:
        k = 0
        lim_k = len(track_data["Track_pt"][i])
        while k < lim_k:
            input_array[i][k][0] = track_data["Track_pt"][i][k]
            input_array[i][k][1] = track_data["Track_phi"][i][k]
            input_array[i][k][2] = track_data["Track_eta"][i][k]
            input_array[i][k][3] = track_data["Track_dxy"][i][k]
            input_array[i][k][4] = track_data["Track_charge"][i][k]
            k += 1
        i += 1

观察2:我们在最内层循环中所做的赋值具有相同的基本结构。如果我们可以获取 TRACK_DATA 字典的各个条目(二维数据)并将它们堆叠在一起,那就太好了。 Numpy 有一个方便(高效)的内置功能,用于沿三维堆叠二维数据:np.dstack。准备好那个3维数组后,我们就可以机械地从中复制进去了:

track_array = np.dstack((
    track_data['Track_pt'],
    track_data['Track_phi'],
    track_data['Track_eta'],
    track_data['Track_dxy'],
    track_data['Track_charge']
))
lim_i = len(n_data["nTrack"])
i = 0
input_array = np.zeros((lim_i,500,5))
for l in range(len(input_array)):
    while i < lim_i:
        k = 0
        lim_k = len(track_data["Track_pt"][i])
        while k < lim_k:
            input_array[i][k][0] = track_data[i][k][0]
            input_array[i][k][1] = track_data[i][k][1]
            input_array[i][k][2] = track_data[i][k][2]
            input_array[i][k][3] = track_data[i][k][3]
            input_array[i][k][4] = track_data[i][k][4]
            k += 1
        i += 1

观察 3:但是现在,我们最内层循环的目的只是沿着最后一个维度复制 track_data 的整个块。我们可以直接这样做:

track_array = np.dstack((
    track_data['Track_pt'],
    track_data['Track_phi'],
    track_data['Track_eta'],
    track_data['Track_dxy'],
    track_data['Track_charge']
))
lim_i = len(n_data["nTrack"])
i = 0
input_array = np.zeros((lim_i,500,5))
for l in range(len(input_array)):
    while i < lim_i:
        k = 0
        lim_k = len(track_data["Track_pt"][i])
        while k < lim_k:
            input_array[i][k] = track_data[i][k]
            k += 1
        i += 1

观察4:但实际上,相同的推理适用于数组的其他两个维度。显然,我们的目的是复制从 dstack; 产生的整个数组;这已经是一个新的数组,所以我们可以直接使用它。

input_array = np.dstack((
    track_data['Track_pt'],
    track_data['Track_phi'],
    track_data['Track_eta'],
    track_data['Track_dxy'],
    track_data['Track_charge']
))

注意到 fKarl Knechtel 的第二条评论,“你应该避免自己显式迭代 Numpy 数组(实际上可以保证有一个内置的 Numpy 东西可以做你想做的事情,而且可能比原生 Python can)," 有一种方法可以通过一次数组编程来做到这一点,但在 NumPy 中不行。 Uproot returns Awkward Arrays 的原因是因为您需要一种有效处理可变长度数据的方法。

我没有你的文件,但我将从类似的文件开始:

>>> import uproot4
>>> import skhep_testdata
>>> events = uproot4.open(skhep_testdata.data_path("uproot-HZZ.root"))["events"]

此文件中以 "Muon_" 开头的分支与您的曲目具有相同的可变长度结构。 (C++ 类型名是一个动态大小的数组,在 Python 中解释为“锯齿状”。)

>>> events.show(filter_name="Muon_*")
name                 | typename                 | interpretation                
---------------------+--------------------------+-------------------------------
Muon_Px              | float[]                  | AsJagged(AsDtype('>f4'))
Muon_Py              | float[]                  | AsJagged(AsDtype('>f4'))
Muon_Pz              | float[]                  | AsJagged(AsDtype('>f4'))
Muon_E               | float[]                  | AsJagged(AsDtype('>f4'))
Muon_Charge          | int32_t[]                | AsJagged(AsDtype('>i4'))
Muon_Iso             | float[]                  | AsJagged(AsDtype('>f4'))

如果你只是要求这些数组,你会得到它们作为一个尴尬的数组。

>>> muons = events.arrays(filter_name="Muon_*")
>>> muons
<Array [{Muon_Px: [-52.9, 37.7, ... 0]}] type='2421 * {"Muon_Px": var * float32,...'>

为了更好地使用它们,让我们导入 Awkward Array 并从询问其类型开始。

>>> import awkward1 as ak
>>> ak.type(muons)
2421 * {"Muon_Px": var * float32, "Muon_Py": var * float32, "Muon_Pz": var * float32, "Muon_E": var * float32, "Muon_Charge": var * int32, "Muon_Iso": var * float32}

这是什么意思?这意味着您有 2421 条记录,其中包含名为 "Muon_Px" 等的字段,每个记录都包含 float32int32 的可变长度列表,具体取决于字段。我们可以通过将其转换为 Python 列表和字典来查看其中之一。

>>> muons[0].tolist()
{'Muon_Px': [-52.89945602416992, 37.7377815246582],
 'Muon_Py': [-11.654671669006348, 0.6934735774993896],
 'Muon_Pz': [-8.16079330444336, -11.307581901550293],
 'Muon_E': [54.77949905395508, 39.401695251464844],
 'Muon_Charge': [1, -1],
 'Muon_Iso': [4.200153350830078, 2.1510612964630127]}

(您可以通过将 how="zip" 传递给 Awkward Array 中的 TTree.arrays or using ak.unzip and ak.zip 来制作这些记录列表,而不是列表记录,但这与您想要执行的填充相切。 )

问题是列表的长度不同。 NumPy 没有任何函数可以帮助我们,因为它完全处理直线数组。因此,我们需要一个专门针对 Awkward Array 的函数,ak.num.

>>> ak.num(muons)
<Array [{Muon_Px: 2, ... Muon_Iso: 1}] type='2421 * {"Muon_Px": int64, "Muon_Py"...'>

这告诉我们每个列表中每个字段的元素数。为了清楚起见,请看第一个:

>>> ak.num(muons)[0].tolist()
{'Muon_Px': 2, 'Muon_Py': 2, 'Muon_Pz': 2, 'Muon_E': 2, 'Muon_Charge': 2, 'Muon_Iso': 2}

您想将这些不规则列表变成大小相同的规则列表。这就是所谓的“填充”。同样,有一个函数,但我们首先需要获得元素的最大数量,以便我们知道要填充多少。

>>> ak.max(ak.num(muons))
4

所以让我们让它们的长度都为 4。

>>> ak.pad_none(muons, ak.max(ak.num(muons)))
<Array [{Muon_Px: [-52.9, 37.7, ... None]}] type='2421 * {"Muon_Px": var * ?floa...'>

再次,让我们看一下第一个以了解我们拥有的东西。

{'Muon_Px': [-52.89945602416992, 37.7377815246582, None, None],
 'Muon_Py': [-11.654671669006348, 0.6934735774993896, None, None],
 'Muon_Pz': [-8.16079330444336, -11.307581901550293, None, None],
 'Muon_E': [54.77949905395508, 39.401695251464844, None, None],
 'Muon_Charge': [1, -1, None, None],
 'Muon_Iso': [4.200153350830078, 2.1510612964630127, None, None]}

您想用零填充它们,而不是 None,因此我们将缺失值转换为零。

>>> ak.fill_none(ak.pad_none(muons, ak.max(ak.num(muons))), 0)[0].tolist()
{'Muon_Px': [-52.89945602416992, 37.7377815246582, 0.0, 0.0],
 'Muon_Py': [-11.654671669006348, 0.6934735774993896, 0.0, 0.0],
 'Muon_Pz': [-8.16079330444336, -11.307581901550293, 0.0, 0.0],
 'Muon_E': [54.77949905395508, 39.401695251464844, 0.0, 0.0],
 'Muon_Charge': [1, -1, 0, 0],
 'Muon_Iso': [4.200153350830078, 2.1510612964630127, 0.0, 0.0]}

最后,NumPy 没有记录(structured array 除外,这也意味着列在内存中是连续的;Awkward Array 的“记录”是抽象的)。因此,让我们将我们拥有的内容解压缩到六个单独的数组中。

>>> arrays = ak.unzip(ak.fill_none(ak.pad_none(muons, ak.max(ak.num(muons))), 0))
>>> arrays
(<Array [[-52.9, 37.7, 0, 0, ... 23.9, 0, 0, 0]] type='2421 * var * float64'>,
 <Array [[-11.7, 0.693, 0, 0, ... 0, 0, 0]] type='2421 * var * float64'>,
 <Array [[-8.16, -11.3, 0, 0, ... 0, 0, 0]] type='2421 * var * float64'>,
 <Array [[54.8, 39.4, 0, 0], ... 69.6, 0, 0, 0]] type='2421 * var * float64'>,
 <Array [[1, -1, 0, 0], ... [-1, 0, 0, 0]] type='2421 * var * int64'>,
 <Array [[4.2, 2.15, 0, 0], ... [0, 0, 0, 0]] type='2421 * var * float64'>)

请注意,这一行执行从 Uproot (muons) 的初始数据拉取开始的所有操作。我现在不打算分析它,但您会发现这一行比显式循环要快得多。

现在我们所拥有的在语义上等同于六个 NumPy 数组,因此我们将它们转换为 NumPy。 (尝试使用不规则数据这样做会失败。您必须明确填充数据。)

>>> numpy_arrays = [ak.to_numpy(x) for x in arrays]
>>> numpy_arrays
[array([[-52.89945602,  37.73778152,   0.        ,   0.        ],
        [ -0.81645936,   0.        ,   0.        ,   0.        ],
        [ 48.98783112,   0.82756668,   0.        ,   0.        ],
        ...,
        [-29.75678635,   0.        ,   0.        ,   0.        ],
        [  1.14186978,   0.        ,   0.        ,   0.        ],
        [ 23.9132061 ,   0.        ,   0.        ,   0.        ]]),
 array([[-11.65467167,   0.69347358,   0.        ,   0.        ],
        [-24.40425873,   0.        ,   0.        ,   0.        ],
        [-21.72313881,  29.8005085 ,   0.        ,   0.        ],
        ...,
        [-15.30385876,   0.        ,   0.        ,   0.        ],
        [ 63.60956955,   0.        ,   0.        ,   0.        ],
        [-35.66507721,   0.        ,   0.        ,   0.        ]]),
 array([[ -8.1607933 , -11.3075819 ,   0.        ,   0.        ],
        [ 20.19996834,   0.        ,   0.        ,   0.        ],
        [ 11.16828537,  36.96519089,   0.        ,   0.        ],
        ...,
        [-52.66374969,   0.        ,   0.        ,   0.        ],
        [162.17631531,   0.        ,   0.        ,   0.        ],
        [ 54.71943665,   0.        ,   0.        ,   0.        ]]),
 array([[ 54.77949905,  39.40169525,   0.        ,   0.        ],
        [ 31.69044495,   0.        ,   0.        ,   0.        ],
        [ 54.73978806,  47.48885727,   0.        ,   0.        ],
        ...,
        [ 62.39516068,   0.        ,   0.        ,   0.        ],
        [174.20863342,   0.        ,   0.        ,   0.        ],
        [ 69.55621338,   0.        ,   0.        ,   0.        ]]),
 array([[ 1, -1,  0,  0],
        [ 1,  0,  0,  0],
        [ 1, -1,  0,  0],
        ...,
        [-1,  0,  0,  0],
        [-1,  0,  0,  0],
        [-1,  0,  0,  0]]),
 array([[4.20015335, 2.1510613 , 0.        , 0.        ],
        [2.18804741, 0.        , 0.        , 0.        ],
        [1.41282165, 3.38350415, 0.        , 0.        ],
        ...,
        [3.76294518, 0.        , 0.        , 0.        ],
        [0.55081069, 0.        , 0.        , 0.        ],
        [0.        , 0.        , 0.        , 0.        ]])]

现在 NumPy 的 dstack 是合适的。 (这使得它们在内存中是连续的,所以如果你愿意,你可以使用 NumPy 的结构化数组。我会发现跟踪哪个索引意味着哪个变量更容易,但这取决于你。实际上,Xarray 特别擅长跟踪直线数组的元数据。)

>>> import numpy as np
>>> np.dstack(numpy_arrays)
array([[[-52.89945602, -11.65467167,  -8.1607933 ,  54.77949905,
           1.        ,   4.20015335],
        [ 37.73778152,   0.69347358, -11.3075819 ,  39.40169525,
          -1.        ,   2.1510613 ],
        [  0.        ,   0.        ,   0.        ,   0.        ,
           0.        ,   0.        ],
        [  0.        ,   0.        ,   0.        ,   0.        ,
           0.        ,   0.        ]],

       [[ -0.81645936, -24.40425873,  20.19996834,  31.69044495,
           1.        ,   2.18804741],
        [  0.        ,   0.        ,   0.        ,   0.        ,
           0.        ,   0.        ],
        [  0.        ,   0.        ,   0.        ,   0.        ,
           0.        ,   0.        ],
        [  0.        ,   0.        ,   0.        ,   0.        ,
           0.        ,   0.        ]],

       [[ 48.98783112, -21.72313881,  11.16828537,  54.73978806,
           1.        ,   1.41282165],
        [  0.82756668,  29.8005085 ,  36.96519089,  47.48885727,
          -1.        ,   3.38350415],
        [  0.        ,   0.        ,   0.        ,   0.        ,
           0.        ,   0.        ],
        [  0.        ,   0.        ,   0.        ,   0.        ,
           0.        ,   0.        ]],

       ...,

       [[-29.75678635, -15.30385876, -52.66374969,  62.39516068,
          -1.        ,   3.76294518],
        [  0.        ,   0.        ,   0.        ,   0.        ,
           0.        ,   0.        ],
        [  0.        ,   0.        ,   0.        ,   0.        ,
           0.        ,   0.        ],
        [  0.        ,   0.        ,   0.        ,   0.        ,
           0.        ,   0.        ]],

       [[  1.14186978,  63.60956955, 162.17631531, 174.20863342,
          -1.        ,   0.55081069],
        [  0.        ,   0.        ,   0.        ,   0.        ,
           0.        ,   0.        ],
        [  0.        ,   0.        ,   0.        ,   0.        ,
           0.        ,   0.        ],
        [  0.        ,   0.        ,   0.        ,   0.        ,
           0.        ,   0.        ]],

       [[ 23.9132061 , -35.66507721,  54.71943665,  69.55621338,
          -1.        ,   0.        ],
        [  0.        ,   0.        ,   0.        ,   0.        ,
           0.        ,   0.        ],
        [  0.        ,   0.        ,   0.        ,   0.        ,
           0.        ,   0.        ],
        [  0.        ,   0.        ,   0.        ,   0.        ,
           0.        ,   0.        ]]])