First, I think you mean you have

jet = ak.Array({
    "constituents": ak.Array([[0, 1, 3, 4], [2]]),
    "energy": ak.Array([1.2, 3.4])
})

because I'd expect the "constituents" indexes to be 0-based, not 1-based. But even if it is 1-based, just start by subtracting 1.

>>> jet.constituents - 1
<Array [[0, 1, 3, 4], [2]] type='2 * var * int64'>

The biggest problem here is that these indexes are nested one level deeper than the particles_p4 that you want to slice. You want the 0, 1, 3, 4, and also the 2, in your jet.constituents to be indexes in the not-nested list, particles_p4.

If we just arbitrarily flatten them (axis=-1 means to squash the last/deepest dimension):

>>> ak.flatten(jet.constituents, axis=-1)
<Array [0, 1, 3, 4, 2] type='5 * int64'>

these indexes are exactly what you'd need to apply to particles_p4. Here, I'm using the current (2.x) version of Awkward Array, so that I can use .show(), but the integer-array slice works in any version of Awkward Array.

>>> particles_p4[ak.flatten(jet.constituents, axis=-1)].show(type=True)
type: 5 * Momentum4D[
    x: int64,
    y: int64,
    z: int64,
    tau: int64
]
[{x: 1, y: 1, z: 1, tau: 1},
 {x: 2, y: 2, z: 2, tau: 2},
 {x: 4, y: 4, z: 4, tau: 4},
 {x: 5, y: 5, z: 5, tau: 5},
 {x: 3, y: 3, z: 3, tau: 3}]

If we take that as a partial solution, all we need to do now is put the nested structure back into the result.

ak.flatten has an opposite, ak.unflatten, which takes a flat array and adds nestedness from an array of list lengths. You can get the list lengths from the original jet.constituents with ak.num. Again, I'll use axis=-1 so that this answer will generalize to deeper nestings.

>>> lengths = ak.num(jet.constituents, axis=-1)
>>> lengths
<Array [4, 1] type='2 * int64'>

>>> rearranged = particles_p4[ak.flatten(jet.constituents, axis=-1)]
>>> rearranged
<MomentumArray4D [{x: 1, y: 1, z: 1, tau: 1}, ..., {...}] type='5 * Momentu...'>

>>> result = ak.unflatten(rearranged, lengths, axis=-1)
>>> result.show(type=True)
type: 2 * var * Momentum4D[
    x: int64,
    y: int64,
    z: int64,
    tau: int64
]
[[{x: 1, y: 1, z: 1, tau: 1}, {x: 2, ...}, ..., {x: 5, y: 5, z: 5, tau: 5}],
 [{x: 3, y: 3, z: 3, tau: 3}]]

For the bonus round, if all of the above arrays (particles_p4 and jet) were arrays of lists, where each list represents one event, rather than an array representing one event, then the above would hold. I'm taking it as a given that the length of the particles_p4_by_event is equal to the length of the jet_by_event arrays, and the values of jet_by_event.constituents are indexes within each event in particles_p4_by_event (not global indexes; each event should restart at zero). That is, all of your arrays agree on how many events there are, and each event is handled individually, with no cross-over between events.

Answer from Jim Pivarski on Stack Overflow
🌐
Awkward-array
awkward-array.org › doc › main › user-guide › how-to-restructure-flatten.html
How to flatten arrays, especially for plotting — Awkward Array 2.13.0 documentation
To destructure an array for plotting, you’ll want to · remove nested lists, definitely for variable-length ones (”var *” in the type string) and possibly for regular ones as well (”N *” in the type string, where N is an integer), ... There are two functions that are responsible for flattening arrays: ak.flatten() with axis=None; and ak.ravel(); but you don’t want to apply them without thinking, because structure is important to the meaning of your data and you want to be able to interpret the plot.
🌐
GitHub
github.com › scikit-hep › awkward › issues › 704
How to flatten awkward arrays (e.g. for plotting) · Issue #704 · scikit-hep/awkward
February 5, 2021 - The page https://awkward-array.org/how-to-restructure-flatten.html looks like it should explain exactly what I need, but it is empty so far. Hence this issue.
Author: scikit-hep
🌐
Awkward-array
awkward-array.org › how-to-restructure-flatten.html
How to flatten arrays, especially for plotting
To destructure an array for plotting, you’ll want to · remove nested lists, definitely for variable-length ones (”var *” in the type string) and possibly for regular ones as well (”N *” in the type string, where N is an integer), ... There is a function that does all of these things in one call, ak.flatten with axis=None, but you don’t want to apply that without thinking because structure is important to the meaning of your data and you want to be able to interpret the plot.
🌐
Awkward-array
awkward-array.org › doc › main › reference › generated › ak.unflatten.html
ak.unflatten — Awkward Array 2.13.0 documentation
>>> original = ak.Array([[0, 1, 2], [], [3, 4], [5], [6, 7, 8, 9]]) >>> counts = ak.num(original) >>> array = ak.flatten(original) >>> counts <Array [3, 0, 2, 1, 4] type='5 * int64'> >>> array <Array [0, 1, 2, 3, 4, 5, 6, 7, 8, 9] type='10 * int64'> >>> ak.unflatten(array, counts) <Array [[0, 1, 2], [], [3, ...], [5], [6, 7, 8, 9]] type='5 * var * int64'>
Top answer
1 of 1
1

First, I think you mean you have

jet = ak.Array({
    "constituents": ak.Array([[0, 1, 3, 4], [2]]),
    "energy": ak.Array([1.2, 3.4])
})

because I'd expect the "constituents" indexes to be 0-based, not 1-based. But even if it is 1-based, just start by subtracting 1.

>>> jet.constituents - 1
<Array [[0, 1, 3, 4], [2]] type='2 * var * int64'>

The biggest problem here is that these indexes are nested one level deeper than the particles_p4 that you want to slice. You want the 0, 1, 3, 4, and also the 2, in your jet.constituents to be indexes in the not-nested list, particles_p4.

If we just arbitrarily flatten them (axis=-1 means to squash the last/deepest dimension):

>>> ak.flatten(jet.constituents, axis=-1)
<Array [0, 1, 3, 4, 2] type='5 * int64'>

these indexes are exactly what you'd need to apply to particles_p4. Here, I'm using the current (2.x) version of Awkward Array, so that I can use .show(), but the integer-array slice works in any version of Awkward Array.

>>> particles_p4[ak.flatten(jet.constituents, axis=-1)].show(type=True)
type: 5 * Momentum4D[
    x: int64,
    y: int64,
    z: int64,
    tau: int64
]
[{x: 1, y: 1, z: 1, tau: 1},
 {x: 2, y: 2, z: 2, tau: 2},
 {x: 4, y: 4, z: 4, tau: 4},
 {x: 5, y: 5, z: 5, tau: 5},
 {x: 3, y: 3, z: 3, tau: 3}]

If we take that as a partial solution, all we need to do now is put the nested structure back into the result.

ak.flatten has an opposite, ak.unflatten, which takes a flat array and adds nestedness from an array of list lengths. You can get the list lengths from the original jet.constituents with ak.num. Again, I'll use axis=-1 so that this answer will generalize to deeper nestings.

>>> lengths = ak.num(jet.constituents, axis=-1)
>>> lengths
<Array [4, 1] type='2 * int64'>

>>> rearranged = particles_p4[ak.flatten(jet.constituents, axis=-1)]
>>> rearranged
<MomentumArray4D [{x: 1, y: 1, z: 1, tau: 1}, ..., {...}] type='5 * Momentu...'>

>>> result = ak.unflatten(rearranged, lengths, axis=-1)
>>> result.show(type=True)
type: 2 * var * Momentum4D[
    x: int64,
    y: int64,
    z: int64,
    tau: int64
]
[[{x: 1, y: 1, z: 1, tau: 1}, {x: 2, ...}, ..., {x: 5, y: 5, z: 5, tau: 5}],
 [{x: 3, y: 3, z: 3, tau: 3}]]

For the bonus round, if all of the above arrays (particles_p4 and jet) were arrays of lists, where each list represents one event, rather than an array representing one event, then the above would hold. I'm taking it as a given that the length of the particles_p4_by_event is equal to the length of the jet_by_event arrays, and the values of jet_by_event.constituents are indexes within each event in particles_p4_by_event (not global indexes; each event should restart at zero). That is, all of your arrays agree on how many events there are, and each event is handled individually, with no cross-over between events.

🌐
Awkward-array
awkward-array.org › doc › 2.7 › _sources › user-guide › how-to-restructure-pad.md.txt
Awkward-array
If you have many levels to flatten at once, you can use `axis=None`: ```{code-cell} ipython3 ak.flatten(ak.Array([[[[[[1.1, 2.2, 3.3]]], [[[4.4, 5.5]]]]]]), axis=None) ``` However, be aware that {func}`ak.flatten` with `axis=None` will also merge all fields of a record, which is usually undesirable, and the order might not be what you expect.
Find elsewhere
🌐
Cms-opendata-workshop
cms-opendata-workshop.github.io › workshop2022-lesson-cpp-root-python › 08-awkward
Using awkward arrays to analyze HEP data – ROOT with C++ and python
August 2, 2022 - So can we use manipulate this object like a numpy array? Yes! If we’re careful about accessing the array properly. ... When we histogram, however, we need to make use of the awkward.flatten function.
🌐
Awkward-array
awkward-array.org › doc › main › getting-started › index.html
Getting started — Awkward Array 2.13.0 documentation
Such an array can be used in numerical ... lists and the data type to reflect the fact that no values are missing. ak.flatten() removes missing values in the process of flattening nested lists (it treats None like [])....
Top answer
1 of 1
1

When I mentioned in a comment that np.searchsorted is where you should be looking, I hadn't noticed that myDict includes every consecutive integer as a key. Having a dense lookup table like this would allow faster algorithms, which also happen to be simpler in Awkward Array.

So, assuming that there's a key in myDict for each integer from 0 up to some value, you can equally well represent the lookup table as

>>> lookup = ak.Array([myDict[i] for i in range(len(myDict))])
>>> lookup
<Array [19.5, 34.1, 10.9] type='3 * float64'>

The problem of picking values at 0, 1, and 2 becomes just an array-slice. (This array-slice is an O(n) algorithm for array length n, unlike np.searchsorted, which would be O(n log n). That's the cost of having sparse lookup keys.)

The problem, however, is that myArray is nested and lookup is not. We can give lookup the same depth as myArray by slicing it up:

>>> multilookup = lookup[np.newaxis][np.zeros(len(myArray), np.int64)]
>>> multilookup
<Array [[19.5, 34.1, 10.9, ... 34.1, 10.9]] type='2 * 3 * float64'>
>>> multilookup.tolist()
[[19.5, 34.1, 10.9], [19.5, 34.1, 10.9]]

And then multilookup[myArray] is exactly what you want:

>>> multilookup[myArray]
<Array [[19.5, 34.1], [10.9, 34.1, 19.5]] type='2 * var * float64'>

The lookup had to be duplicated because each list within myArray uses global indexes in the whole lookup. If the memory involved in creating multilookup is prohibitive, you could instead break myArray down to match it:

>>> flattened, num = ak.flatten(myArray), ak.num(myArray)
>>> flattened
<Array [0, 1, 2, 1, 0] type='5 * int64'>
>>> num
<Array [2, 3] type='2 * int64'>
>>> lookup[flattened]
<Array [19.5, 34.1, 10.9, 34.1, 19.5] type='5 * float64'>
>>> ak.unflatten(lookup[flattened], nums)
<Array [[19.5, 34.1], [10.9, 34.1, 19.5]] type='2 * var * float64'>

If your keys are not dense from 0 up to some integer, then you'll have to use np.searchsorted:

>>> keys = ak.Array(myDict.keys())
>>> values = ak.Array([myDict[key] for key in keys])
>>> keys
<Array [0, 1, 2] type='3 * int64'>
>>> values
<Array [19.5, 34.1, 10.9] type='3 * float64'>

In this case, the keys are trivial because it is dense. When using np.searchsorted, you have to explicitly cast the flat Awkward Arrays as NumPy (for now; we're looking to fix that).

>>> lookup_index = np.searchsorted(np.asarray(keys), np.asarray(flattened), side="left")
>>> lookup_index
array([0, 1, 2, 1, 0])

Then we pass it through the trivial keys (which doesn't change it, in this case) before passing it to the values.

>>> keys[lookup_index]
<Array [0, 1, 2, 1, 0] type='5 * int64'>
>>> values[keys[lookup_index]]
<Array [19.5, 34.1, 10.9, 34.1, 19.5] type='5 * float64'>
>>> ak.unflatten(values[keys[lookup_index]], num)
<Array [[19.5, 34.1], [10.9, 34.1, 19.5]] type='2 * var * float64'>

But the thing I was waffling about in yesterday's comment was that you have to do this on the flattened form of myArray (flattened) and reintroduce the structure later ak.unflatten, as above. But perhaps we should wrap np.searchsorted as ak.searchsorted to recognize a fully structured Awkward Array in the second argument, at least. (It has to be unstructured to be in the first argument.)

🌐
GitHub
github.com › scikit-hep › awkward › discussions › 532
Returning Awkward Arrays vs returning NumPy arrays from Awkward functions · scikit-hep/awkward · Discussion #532
There are good reasons to have multiple ways of expressing the same array, with and without an IndexedArray, for instance (which delays rearrangement), but those reasons are technical and would be very distracting in a data analysis. However, this is a difference in layout that could make the difference between an Awkward Array and a NumPy array. Finally, we do make sure that any function with the same name as NumPy has the same behavior. However, NumPy doesn't have a np.flatten function.
Author: scikit-hep
🌐
Coderz Column
coderzcolumn.com › tutorials › python › awkward-array-guide-to-work-with-tree-like-nested-variable-sized-datasets
Awkward Array: Guide to Work with JSON-like Nested, Variable-sized Datasets using Numpy-like Idioms by Sunny Solanki
November 20, 2021 - The ravel() function works exactly like flatten() function and can be used to flatten an array. ... We can concatenate more than one awkward array using concatenate() function.
🌐
Readthedocs
awkward-array.readthedocs.io › en › latest › _auto › ak.unflatten.html
ak.unflatten — Awkward Array documentation
>>> original = ak.Array([[0, 1, 2], [], [3, 4], [5], [6, 7, 8, 9]]) >>> counts = ak.num(original) >>> array = ak.flatten(original) >>> counts <Array [3, 0, 2, 1, 4] type='5 * int64'> >>> array <Array [0, 1, 2, 3, 4, 5, 6, 7, 8, 9] type='10 * int64'> >>> ak.unflatten(array, counts) <Array [[0, ...
🌐
Tessl
tessl.io › registry › tessl › pypi-awkward › files › docs › integration.md
2.8.0 • pypi-awkward • tessl • Registry • Tessl
January 29, 2026 - Parameters: - array: Array to flatten Returns: tuple of (leaves, tree_def) for JAX pytree operations """ def jax_pytree_unflatten(tree_def, leaves): """ Reconstruct awkward array from JAX pytree components.
🌐
GitHub
github.com › scikit-hep › awkward › issues › 3033
ak.flatten on dask_awkward BitMaskedArray triggers awkward error · Issue #3033 · scikit-hep/awkward
February 24, 2024 - import awkward as ak import dask_awkward as dak esrc = ak.from_parquet("join_out_flat-part2.parquet") dsrc = dak.from_parquet("join_out_flat-part2.parquet") #awkward OKAY ak.flatten(esrc.electron_pt, axis=None) ak.flatten(esrc.electron_pt, axis=0) ak.flatten(esrc.electron_pt, axis=1) ak.flatten(esrc.electron_pt, axis=-1) #dask_awkward array OKAY ak.flatten(dsrc.electron_pt, axis=0).compute() ak.flatten(dsrc.electron_pt, axis=None).compute() ak.flatten(dsrc.electron_pt.compute(), axis=0) ak.flatten(dsrc.electron_pt.compute(), axis=None) #dask_awkward array ERROR in awkward ak.flatten(dsrc.electron_pt, axis=1) #errors before .compute() ak.flatten(dsrc.electron_pt, axis=-1) #errors before .compute() ak.flatten(dsrc.electron_pt.compute(), axis=1)
Author: scikit-hep