Calculate all configurations of leptons with opposite charge

[cesarecazzaniga]

Dear experts,

I have a question regarding the implementation of the following calculation with awkward arrays.
Let's say I have an array with electrons for each event (with pt, phi, eta, pdgid, charge, etc...), and I want to calculate the invariant mass of all the di-electron pairs with the correct charge configurations. Since the choice of the correct charge configuration is not unique, then I would like to pick up the configuration which minimise a metric f. Let's say the metric f could be the sum of the distances delta R ^2 for each given couple in a given charge configuration (not sure about this metric choice, but I had a try on a c++ code and seems to reasonably work .... but maybe there can be better possibilities).

So, let' say I have 4 electrons: e1+, e2-, e3+, e4- , first I would like to have all the allowed couples (+,-): {(1+,2-) , (3+,4-)} or {(1+,4-),(2-,3+)}.
After that I want to calculate for each of charge configuration the metric f, eg. : f({(1+,2-) , (3+,4-)}) = deltaR(1,2)^2 + deltaR(3,4)^2.
Finally, I want to find the unique charge configuration that minimise the given metric.

Could you please guide me a bit on how to implement this ?

Let me know, and thanks in advance for the help !

[agoose77]

Hi @cesarecazzaniga My understanding of your problem is:

• You have an array of type N * var * {...}, where N is the number of events, and var is the number of electrons per-event.
• You want to find all possible electron-electron pairings per-event that satisfy the charge configuration requirement.
• You want to choose the best of these possible configurations by minimizing a metric.

If I've correctly described your problem, then at a high-level it sounds like you want to generate an array of N * var * var * {...} proposals, evaluate your metric, and then extract the minimum. This last step is quite straightforward, especially with the vector library that provides an implementation of deltaR (and deltaR2) for Momentum4D vectors, e.g.

proposal = ...

metric = np.sum(proposal.slot0.deltaR2(proposal.slot1), axis=-1)
i_best = ak.argmin(metric, keepdims=True, axis=-1)
best = proposal[i_best]

The difficulty here is, I imagine, generating your proposals.

One option is to take all of the permutations over the sequence, and then the valid pairs within that sequence. We don't have a built-in method for generating random permutations, but you could implement Heap's algorithm to compute them with Numba. Still, this is likely not a good solution - for a start, it is explosively combinatorial, so I'd like to clarify that this is what you're trying to do first!

Here's an example of generating the permutations. I probably would not do this in actual code, because of how quick it is to consume all of your RAM

https://gist.github.com/agoose77/0c4ef3ee7582aebf75871e5974f4af72

Note that generating the permutations this way includes repeats - this doesn't account for the fact that the order between pairs does not matter, or that (a, b) is the same pair as (b, a)

[cesarecazzaniga]

Hello, so when running the code with your implementation (@agoose77 ), I run into troubles with some initial charge configurations like

x = ak.Array(
    [
        [
         {"id": 0, "charge": +1}, 
         {"id": 1, "charge": +1},   
        ],
        
        [
            {"id": 0, "charge": -1},
            {"id": 1, "charge": 1},
            {"id": 2, "charge": 1},
            {"id": 3, "charge": -1},
        ],
        
    ]
) 

So, for this case (but there are also others where you have for example in one event only one lepton) I get the following error message at this stage here of your code:

pos_, _ = ak.broadcast_arrays(pos[..., np.newaxis, :], ix.slot0[..., [0]]) neg_, _ = ak.broadcast_arrays(neg[..., np.newaxis, :], ix.slot1[..., [0]])

the error is like:

ValueError: in ListOffsetArray64 attempting to get 0, index out of range

Do you know how can I deal with this ?

[cesarecazzaniga]

By the way ... I have not tried yet the solution proposed by @jpivarski (thanks for also the nice reference code) , maybe I can have a try also with that one ...

[agoose77]

Could you clarify whether there is an upper bound on the number of pairs? If so, the explicit method suggested by @jpivarski may be a clearer option.

In any case, to answer your question:

The purpose of ak.broadcast_arrays here is to "grow" the newly added (via np.newaxis) 1 dimension in pos so that the resulting array can be indexed by ix. Unlike other operations, indexing requires that the structure matches (rather than just be broadcastable), so we have to do this explicitly.

A clearer solution that I didn't think about at the time is to use ak.num instead of a single-value slice, e.g.

pos_, _ = ak.broadcast_arrays(pos[..., np.newaxis, :], ak.num(ix.slot0, axis=-1))
neg_, _ = ak.broadcast_arrays(neg[..., np.newaxis, :], ak.num(ix.slot1, axis=-1))

The result of ak.num(..., axis=-1) actually has one-fewer dimensions than pos[..., np.newaxis, :]. For your example:

>>> pos[..., np.newaxis, :].type
2 * 1 * var * {"id": int64, "charge": int64}
>>> ak.num(ix.slot0, axis=-1)
2 * var * int64

This is fine, though, because ak.broadcast_arrays will left broadcast these so that the 2s line up, i.e. it will add a 1-sized dimension to the end of the ak.num(...) argument:

2 * var * 1 * int64

so that we have:

>>> pos_.type
2 *  var  * var * {"id": int64, "charge": int64}
#     ^ the 1 is grown to var

[cesarecazzaniga]

sorry, I have not specified an upper bound on the number of pairs .... because in my specific case it would be better not to assume it

[jpivarski]

Not having an upper bound on combinations to check means you can't do the explicit combinatorics in my suggestion. Brian's notebooks do not use explicit combinatorics, but they do put an upper limit on the number of particles, so it's more general than the problem it's solving. That's why I've had this "explicit combinatorics" idea in the back of my head for a while.