How to zip with integer keys i.e. how to concatenate branches to form an array

[wiso]

Sorry, for the naive question.
I have a ROOT file with branches like

PMT1_E: float
PMT2_E: float
...
PMT8_E: float

so I have 8 branches of type float. I would like to structure my data so that I can access it as

events[0].PMT.E[0]  # first event, energy (E) of the first PMT

or events[0].PMT[0].E

and I also want to be able to do ak.sum(events.PMT.E), which should return a list (one entry for each event) of floats.

I am simply doing

PMT_E = ak.zip([events['PMT%d_E' % i] for i in range(1, 8)])

The problem is that now I can do PMT_E[0]["0"] where the last index is a string (and I can't do ak.sum). I see that zip requires keys to be strings. Maybe there are other functions?

[ikrommyd]

Are there more variables per PMT than energy? Do you have PMT1_X, PMT1_Y PMT1_Z.... for example?

[ikrommyd]

If I'm understanding what you are saying with the ak.sum(events.PMT.E), you want it to sum the energies for all PMTs correct?
I can probably recreate some mock data that look like what you're saying where I have 3 PMTs and they each have 3 values X,Y and Z

import awkward as ak
import numpy as np
import uproot

np.random.seed(42)
events_data = {}
for i in range(1, 4):
    events_data[f'PMT{i}_X'] = np.random.normal(0, 1, 1000).astype(np.float32)
    events_data[f'PMT{i}_Y'] = np.random.normal(0, 1, 1000).astype(np.float32)
    events_data[f'PMT{i}_Z'] = np.random.normal(0, 1, 1000).astype(np.float32)
events = ak.Array(events_data)

with uproot.recreate("mock_data.root") as f:
    f["tree"] = events_data

Now you probably want to stack the PMTs into different "columns". I would do something like this probably

import awkward as ak
import uproot
import numpy as np

events_flat = uproot.open("mock_data.root")["tree"].arrays(how=dict)

# Restructure: stack PMT data so that events.PMT.X[:, pmt_index] gives X for that PMT
PMT_X = ak.Array(np.column_stack([events_flat['PMT1_X'], events_flat['PMT2_X'], events_flat['PMT3_X']]))
PMT_Y = ak.Array(np.column_stack([events_flat['PMT1_Y'], events_flat['PMT2_Y'], events_flat['PMT3_Y']]))
PMT_Z = ak.Array(np.column_stack([events_flat['PMT1_Z'], events_flat['PMT2_Z'], events_flat['PMT3_Z']]))

events = ak.zip({"PMT": ak.zip({"X": PMT_X, "Y": PMT_Y, "Z": PMT_Z})})

Now you get this structure

In [25]: events.PMT
Out[25]: <Array [[{X: 0.497, Y: 1.4, ...}, ...], ...] type='1000 * 3 * {X: float32, ...'>

In [26]: events.PMT.fields
Out[26]: ['X', 'Y', 'Z']

In [27]: events.PMT.X
Out[27]: <Array [[0.497, -1.91, -1.11], ..., [0.573, ...]] type='1000 * 3 * float32'>

So you can get the value of X of the first PMT by doing

In [29]: events.PMT.X[:,0]
Out[29]: <Array [0.497, -0.138, 0.648, ..., 0.641, -0.571, 0.573] type='1000 * float32'>

and you can sum X across PMTs with

In [30]: ak.sum(events.PMT.X, axis=1)
Out[30]: <Array [-2.53, -1.63, -0.708, ..., -0.207, 1.07, 1.33] type='1000 * float32'>

Does that create a good structure for you to work with?

[wiso]

Thank you very much. np.column_stack is the answer.

Yes, I have multiple fields (your X, Y, Z), but that is easy to handle with zip. Here my solution

PMT_E = np.column_stack([events[f'PMT{i}'] for i in range(1, 8 + 1)])
PMT_ADC = np.column_stack([events[f'PMT{i}_adc'] for i in range(1, 8 + 1)])
PMT = ak.zip({"E": PMT_E, "ADC": PMT_ADC})
events['PMT'] = PMT

[ikrommyd]

As far as I'm aware, it's probably best in your problem to keep the different PMTs as "columns" and not as records that you can access with a getitem like ["0"], ["1"] etc. The reason is that since you wanna sum across PMTs, I think this structure makes sense. Otherwise if you wanna keep different PMTs as different fields of a record, you'd probably need to define a custom "summation behavior" about how to "sum" PMTs together. We do this for things like 4-vector objects for example and we can sum them together but I think it over-complicates your situation. Maybe I'm being stupid here and there's something I'm missing but this how I'd do it.

[wiso]

with the solution you proposed I can easily do

ak.sum(events.PMT.E, axis=1)

which is what I need.

[ikrommyd]

My logic here was similar to how we structure physics objects like events.Electron.pt gets you the array of pt for all electrons. events.Electron.eta gets you eta. events.Electron.pt[:, 0] gets you the pt of the leading electron in every event. So variables/attributes/fields or whatever you wanna call them go on the outside like events.Electron.something and form a record and the different values of use fields for each electron go on the inside accessible by index. This is what I did to your PMTs too.

[wiso]

yes, exactly. I can do

events.PMT[:, 0]

to get the first PMT for all the events. There is a little caveat that we name the first PMT as 1 (0 has a special meaning), but it is not a big deal. Using a string as an index will let me use "1" as the first index, but everything else is more complicated. I can simply add a function, which tells me the name of the cell, e.g. index 0 -> "Cell 1". The point is: how to find the index?

Assume I have

def cell_name(cell_index): return f"Cell {cell_index + 1}"

I would like to call

events[0].PMT[0].name

and get "Cell 1"

[wiso]

Since you are expert of coffea, do you think it makes sense to write a custom schema?

[ikrommyd]

Uhmm....maybe? 🤣 If it's something more complicated it's not a bad idea. However schemas are more complicated than what you're just doing here so it's possible that it's not worth the effort. You can decide that.

[wiso]

sure, I was hoping in some help, since I can't find documentation about schemas in cofeea.

[ikrommyd]

It's not something external that people typically do. You can try to understand the code I guess here https://github.com/scikit-hep/coffea/tree/master/src/coffea/nanoevents/schemas. I'm warning you though....if may be an effort that is not really required if you can just load your PMT file and do this restructuring with column stacks, looks and zips.

[ikrommyd]

The schemas are complictaed and contain a lot of internal jargon which is why we're trying to make this hopefully easier at some poing in awkward-zipper: https://github.com/maxymnaumchyk/awkward-zipper