尷尬的數組在間隔中添加屬性

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

我們可以從那里開始，選擇介於"start"和"stop"值之間的"start" "id" "stop"值。 我開始寫一個解決方案，但是切片變得有點復雜，生成所有笛卡爾組合的計算成本比這個問題嚴格需要的要多（雖然沒有寫一個 for 循環那么昂貴！），以及笛卡爾積對於近似匹配比您擁有的精確索引更有用。

相反，讓我們在Numba 中編寫一個 for 循環，這是一個用於 Python 的即時編譯器。 Numba 在它可以編譯的 Python 中受到限制（在編譯時必須知道所有類型），但它可以識別只讀的 Awkward Arrays 和 append-only ak.ArrayBuilder 。

這是一個循環，它只考慮同一事件中的tracks和jets ，在軌道上循環，並將與每個track匹配的第一個jet放入輸出 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 中是允許的——所以很高興知道是哪個更改導致它如此抱怨。

希望這些示例可以幫助您入門！