python - 在数组中重塑数组

Question

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

通过uproot DATA=[variablename:[array of entries [array of subentries]] 将根文件转换为字典

我创建了一个数组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

Answer 1

注意到 fKarl Knechtel 的第二条评论，“你应该避免自己显式地迭代 Numpy 数组（实际上可以保证有一个内置的 Numpy 东西可以做你想做的事情，而且可能比原生 Python 快得多），”那里是一种通过一次数组编程来做到这一点的方法，但不是在 NumPy 中。Uproot 返回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"等的字段，每条记录都包含float32或的可变长度列表int32，具体取决于字段。我们可以通过将其转换为 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"（您可以通过传递给TTree.arrays或在 Awkward Array 中使用ak.unzip和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 没有记录（除了结构化数组，这也意味着列在内存中是连续的；尴尬数组的“记录”是抽象的）。因此，让我们将我们拥有的内容解压缩到六个单独的数组中。

>>> 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.        ]])]

现在 NumPydstack是合适的。（这使得它们在内存中是连续的，所以如果你愿意，你可以使用 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.        ]]])

Answer 2

观察 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_DATAdict 的各个条目（它们是二维数据）并将它们堆叠在一起，那就太好了。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']
))