python - 尴尬的数组间隔添加属性

Question

我想从根文件中提取数据并使其成形，以一个 numpy 数组/张量结束，以将其填充到神经网络中。我已经能够通过填充获得我想要的轨迹数据，将其转换为一个 numpy 数组，但我想用它们所来自的喷气机的数据来扩展我的数组。所以我有所有航迹的信息，每架喷气机的信息以及它们对应的航迹间隔。我的第一个本能是构建一个轨道形状的数组，并使用类似的东西np.dstack来合并这两个。

import uproot4 as uproot
import numpy as np
import awkward1 as ak

def ak_into_np(ak_array):
    data=np.dstack([ak.to_numpy(x) for x in ak_array])
    return data    

def get_data(filename,padding_size):
        f=uproot.open(filename)
        events= f["btagana/ttree;1"]
        track_data=events.arrays(filter_name=["Track_pt","Track_phi","Track_eta","Track_dxy","Track_dz","Track_charge"])
        jet_interval=events.arrays(filter_name=["Jet_nFirstTrack","Jet_nLastTrack"])
        jet_interval=jet_interval["Jet_nLastTrack"]-jet_interval["Jet_nFirstTrack"]
        
        jet_data=events.arrays(filter_name=["Jet_pt","Jet_phi","Jet_eta"])
        arrays_track=ak.unzip(ak.fill_none(ak.pad_none(track_data, padding_size), 0))
        arrays_interval=ak.unzip(ak.fill_none(ak.pad_none(jet_interval,padding_size),0))
        arrays_jet=ak.unzip(ak.fill_none(ak.pad_none(jet_data,padding_size),0))
        track=ak_into_np(arrays_track)
        jet=ak_into_np(arrays_jet)
        interval=ak_into_np(arrays_interval)
        
        return track,jet,interval

这是我到目前为止的地方。出于效率原因，我希望在进入 numpy 之前能够在尴尬中实现这一点。我在 numpy 中尝试了以下方法：

def extend(track,jet,interval):
    events,tracks,varstrack=(np.shape(track))
    events,jets,varsjet=np.shape(jet)
    jet_into_track_data=[]
    for i in range(events):
        dataloop=[]
        for k in range(jets):
            if interval[i][k][0]!=0 :
                dataloop.append(np.broadcast_to(jet[i][k],(interval[i][k][0],varsjet)))
            else 
    jet_into_track_data.append(dataloop)
        
        
    
            
    return jet_into_track_data

但它已经花费了大约 3 秒，甚至没有达到我仅 2000 个事件的目标。目的基本是[track_variables] ->[track_variables,jet_variables if track is in intervall]，要存起来[(event1)[[track_1],...,[track_padding_size]],...,(eventn)[[track_1],...,[track_padding_size]]]

Answer 1

我看不到您的原始数据的结构，也不清楚所需的最终状态是什么，但我可以举一个受上述启发的示例，您可以适应。另外，我将忽略填充，因为这只会使事情复杂化。在完成组合数学之前，您可能希望推迟填充。

tracks和下面jets来自同一组事件（即数组具有相同的长度并且它们是锯齿状的，每个列表中的轨道数量不同，每个列表中的喷射器数量不同）。由于jets以某种方式派生自tracks，因此喷射器的数量严格少于轨道，并且我认为在您的问题中，它们之间的联系是这样的，即每个喷射器对应于一个连续的、非重叠的集合tracks（“间隔”）。

真实的轨道和喷气式飞机将具有比这些更多的属性——这只是一个最简单的例子。我已经给出了轨道，"id"所以我们可以区分它们，但是这些喷气机只有包含"start"索引和排他"stop"索引。

如果您的曲目没有标识符，您可以使用ak.local_index添加它们。

>>> import awkward1 as ak
>>> tracks = ak.Array([[{"id": 0}, {"id": 1}, {"id": 2}],
...                    [],
...                    [{"id": 0}, {"id": 1}]])
>>> jets = ak.Array([[{"start": 0, "stop": 2}, {"start": 2, "stop": 3}],
...                  [],
...                  [{"start": 1, "stop": 2}, {"start": 0, "stop": 1}]])

如果您在每个事件之间都有所有组合tracks，jets那么您可以使用切片来选择匹配的组合。当您有不精确的匹配（您必须与 ΔR 或其他东西匹配）时，这特别有用。ak.cartesian函数生成组合列表，并按第nested=True一个参数对结果进行分组：

>>> all_combinations = ak.cartesian([tracks, jets], nested=True)
>>> all_combinations.tolist()
[[[({'id': 0}, {'start': 0, 'stop': 2}),
   ({'id': 0}, {'start': 2, 'stop': 3})],
  [({'id': 1}, {'start': 0, 'stop': 2}),
   ({'id': 1}, {'start': 2, 'stop': 3})],
  [({'id': 2}, {'start': 0, 'stop': 2}),
   ({'id': 2}, {'start': 2, 'stop': 3})]],
 [[]],
 [[({'id': 0}, {'start': 1, 'stop': 2}),
   ({'id': 0}, {'start': 0, 'stop': 1})],
  [({'id': 1}, {'start': 1, 'stop': 2}),
   ({'id': 1}, {'start': 0, 'stop': 1})]]]

我们可以从那里开始，选择介于和值"id"之间的值。我开始写一个解决方案，但是切片变得有点复杂，生成所有笛卡尔组合的计算成本比这个问题严格需要的要高（尽管没有写一个 for 循环那么昂贵！），以及笛卡尔积对于近似匹配比您拥有的精确索引更有用。"start""stop"

相反，让我们在 Numba 中编写一个 for 循环，Numba是 Python 的即时编译器。Numba 在可以编译的 Python 中受到限制（在编译时必须知道所有类型），但它可以识别只读的 Awkward Arrays 和只追加的ak.ArrayBuilder。

这是一个只考虑tracks和jets在同一事件中的循环，循环遍历轨道，并将jet匹配每个轨道的第一个track放入输出 ArrayBuilder。

>>> import numba as nb
>>> @nb.njit
... def match(tracks_in_events, jets_in_events, output):
...     for tracks, jets in zip(tracks_in_events, jets_in_events):
...         output.begin_list()
...         for track in tracks:
...             for jet in jets:
...                 if jet.start <= track.id < jet.stop:
...                     output.append(jet)   # at most one
...                     break
...         output.end_list()
...     return output
... 
>>> builder = match(tracks, jets, ak.ArrayBuilder())
>>> builder.snapshot().tolist()
[[{'start': 0, 'stop': 2}, {'start': 0, 'stop': 2}, {'start': 2, 'stop': 3}],
 [],
 [{'start': 2, 'stop': 3}, {'start': 0, 'stop': 2}]]

请注意，这些喷气机对象被复制以匹配适当的轨道。（这个“复制”实际上只是一个指针，而不是真正的副本。）要将其附加到轨道，您可以分配它：

>>> tracks["jet"] = builder.snapshot()
>>> tracks.tolist()
[[{'id': 0, 'jet': {'start': 0, 'stop': 2}},
  {'id': 1, 'jet': {'start': 0, 'stop': 2}},
  {'id': 2, 'jet': {'start': 2, 'stop': 3}}],
 [],
 [{'id': 0, 'jet': {'start': 2, 'stop': 3}},
  {'id': 1, 'jet': {'start': 0, 'stop': 2}}]]

在这里，我假设您想将喷气机附加到每个轨道 - 也许您想将所有关联轨道的集合附加到每个喷气机：

>>> @nb.njit
... def match2(tracks_in_events, jets_in_events, output):
...     for tracks, jets in zip(tracks_in_events, jets_in_events):
...         output.begin_list()
...         for jet in jets:
...             output.begin_list()
...             for track in tracks:
...                 if jet.start <= track.id < jet.stop:
...                     output.append(track)    # all tracks
...             output.end_list()
...         output.end_list()
...     return output
... 
>>> jets["tracks"] = match2(tracks, jets, ak.ArrayBuilder()).snapshot()
>>> jets.tolist()
[[{'start': 0, 'stop': 2, 'tracks': [
      {'id': 0, 'jet': {'start': 0, 'stop': 2}},
      {'id': 1, 'jet': {'start': 0, 'stop': 2}}]},
  {'start': 2, 'stop': 3, 'tracks': [
      {'id': 2, 'jet': {'start': 2, 'stop': 3}}]}],
 [],
 [{'start': 1, 'stop': 2, 'tracks': [
      {'id': 1, 'jet': {'start': 0, 'stop': 2}}]},
  {'start': 0, 'stop': 1, 'tracks': [
      {'id': 0, 'jet': {'start': 0, 'stop': 2}}]}]]

（因为我两者都做了，所以现在有两种方式的链接。）将喷气机连接到轨道，而不是轨道到喷气机，增加了一些轨道可能与任何喷气机无关的复杂性，在这种情况下你必须考虑可能丢失的数据。（提示：为“不匹配”和“是匹配”将它们设为零或一个喷气机列表，然后使用ak.firsts将空列表转换为无。）

或者，您可以将 Numba 编译函数的输出设为纯 NumPy 数组。Numba 知道很多 NumPy 函数。构建 Numba 编译函数的提示：从不做任何事情的数据的最小循环开始，在添加输出时测试运行它。当 Numba 无法识别某些东西的类型时，它会抱怨很多输出——这在非编译的 Python 中是允许的——所以很高兴知道是哪个更改导致它抱怨这么多。

希望这些示例可以帮助您入门！