How to filter with ragged arrays#
import awkward as ak
import numpy as np
What is awkward indexing?#
One of the most powerful features of NumPy is the expressiveness of its indexing system. A NumPy array can be sliced in many different ways, such as with a single integer, or an array of integers. Awkward Array implements most of these indexing styles, but adds an additional variant: awkward indexing.
Consider the following ragged array:
array = ak.Array(
[
[
[0.0, 1.1, 2.2],
[3.3, 4.4, 5.5, 6.6],
[7.7],
],
[],
[
[8.8, 9.9, 10.10, 11.11, 12.12],
],
]
)
array
[[[0, 1.1, 2.2], [3.3, 4.4, 5.5, 6.6], [7.7]], [], [[8.8, 9.9, 10.1, 11.1, 12.1]]] ---------------------------------------------- backend: cpu nbytes: 176 B type: 3 * var * var * float64
We can easily pull out the first two items with a simple slice
array[..., :2]
[[[0, 1.1], [3.3, 4.4], [7.7]], [], [[8.8, 9.9]]] ------------------------------- backend: cpu nbytes: 128 B type: 3 * var * var * float64
But what if we wanted to pull out a different number of items for each sublist, e.g. to produce the following array:
[[[], [3.3], [7.7]],
[],
[[10.10, 11.11, 12.12]]]
----------------------------------------------
type: 3 * var * var * float64
To produce this result, we need awkward indexing.
Building an awkward index#
Awkward indexing requires an index array that
has a structure matching the array being sliced up to (but not including) the final dimension of the index
has at least one ragged (
var) dimension or contain missing values
By structure, we mean the number of sublists in each dimension, which can be seen with ak.num():
axis=0 has a single list of three items:
ak.num(array, axis=0)
array(3)
axis=1 has three lists, the first with three items, the second with zero items, the third with a single item:
ak.num(array, axis=1)
[3, 0, 1] --------------- backend: cpu nbytes: 24 B type: 3 * int64
To put this more simply, the final dimension of the awkward index is used to pull items out of the array. Therefore, Awkward needs the preceeding dimensions to line up!
Recall that we wanted to pull out the following result from array using awkward indexing:
[[[], [3.3], [7.7]],
[],
[[10.10, 11.11, 12.12]]]
----------------------------------------------
type: 3 * var * var * float64
It’s clear that we want to pull specific items out of the final dimension of the array. Let’s find out where these particular items are located in their sublists. Awkward Array provides a special function ak.local_index() to find the index of each item in the array
ak.local_index(["x", "y", "z"])
[0, 1, 2] --------------- backend: cpu nbytes: 24 B type: 3 * int64
The word “local” refers to the way that ak.local_index() computes the index of each item relative to the sublist in which it is found. e.g. for a two-dimensional array:
ak.local_index(
[
["up", "charm", "top"],
["down", "strange"],
["bottom"],
]
)
[[0, 1, 2], [0, 1], [0]] --------------------- backend: cpu nbytes: 80 B type: 3 * var * int64
ak.local_index() also takes an axis parameter, but here we only need the default axis=-1. It can be seen that this local index has exactly the same structure as array.
array
[[[0, 1.1, 2.2], [3.3, 4.4, 5.5, 6.6], [7.7]], [], [[8.8, 9.9, 10.1, 11.1, 12.1]]] ---------------------------------------------- backend: cpu nbytes: 176 B type: 3 * var * var * float64
ak.local_index(array)
[[[0, 1, 2], [0, 1, 2, 3], [0]], [], [[0, 1, 2, 3, 4]]] -------------------------------- backend: cpu nbytes: 176 B type: 3 * var * var * int64
To create our awkward index, all we need to do is create an array like ak.local_index(array), but with only the local indices that we want to keep, i.e.
index = ak.Array(
[
[[], [0], [0]],
[],
[[2, 3, 4]],
]
)
We can see that this array matches the leading structure of array, and has at least one var dimension
index.type.show()
3 * var * var * int64
Let’s see what slicing array with this awkward index looks like:
array[index]
[[[], [3.3], [7.7]], [], [[10.1, 11.1, 12.1]]] ----------------------------- backend: cpu nbytes: 112 B type: 3 * var * var * float64
Clearly this index produces the result that we were aiming for!
Indexing with argmin and argmax#
Awkward indexing is especially useful when combined with the positional ak.argmin() and ak.argmax() reducers. These functions accept an keepdims=True argument that can be used to keep the same number of dimensions as the original array. There is also a mask_identity argument is explained in Indexing with missing values. For now, we will set it to False.
array = ak.Array(
[
[10, 3, 2, 9],
[4, 5, 5, 12, 6],
[8, 9, -1],
]
)
array
[[10, 3, 2, 9], [4, 5, 5, 12, 6], [8, 9, -1]] --------------------- backend: cpu nbytes: 128 B type: 3 * var * int64
With keepdims=False, all reducers collapse a dimension of the original array:
ak.argmin(array, axis=1, keepdims=False, mask_identity=False)
[2, 0, 2] --------------- backend: cpu nbytes: 24 B type: 3 * int64
If we try and use this index to slice array, it will likely not produce the result we might initially expect:
array[ak.argmin(array, axis=1, keepdims=False, mask_identity=False)]
[[8, 9, -1], [10, 3, 2, 9], [8, 9, -1]] --------------------- backend: cpu nbytes: 144 B type: 3 * var * int64
Instead of pulling out the smallest items in array along axis=1, we have simply re-arranged the sublists of array along axis=0. Our index has only a single dimension, so for each value in ak.argmin(array, axis=-1), Awkward pulls out the corresponding item from array. We want to pull values out of the second dimension, so our index array needs to be two dimensional.
Let’s now look at what happens with keepdims=True. The result is a two dimensional, fully regular array, with no missing values:
ak.argmin(array, axis=-1, keepdims=True, mask_identity=False)
[[2], [0], [2]] ------------------- backend: cpu nbytes: 24 B type: 3 * 1 * int64
Before we can use this as an index array, we need to convert at least one dimension to a ragged dimension. This follows from rule (2) described in Building an awkward index.
ak.from_regular(
ak.argmin(array, axis=-1, keepdims=True, mask_identity=False)
)
[[2], [0], [2]] --------------------- backend: cpu nbytes: 56 B type: 3 * var * int64
We can now use this array to index into array:
array[
ak.from_regular(
ak.argmin(array, axis=-1, keepdims=True, mask_identity=False)
)
]
[[2], [4], [-1]] --------------------- backend: cpu nbytes: 56 B type: 3 * var * int64
it produces the expected result!
Filtering with booleans#
As described in Building an awkward index, Awkward Array’s awkward indexing is a generalisation of the advanced indexing supported by NumPy. It is therefore reasonable to ask whether Awkward supports awkward indexing with
boolean values, selecting only values for which the index is True.
Let’s create an array of integers:
numbers = ak.Array(
[
[0, 1, 2, 3],
[4, 5, 6],
[8, 9, 10, 11, 12],
]
)
We can use awkward indexing to keep only the even values. Let’s generate a boolean mask with the same structure as numbers. In order for there to be a single boolean value for each item in numbers, the filter array must have exactly the same number of elements. Ufuncs, such as np.mod(), are powerful tools for generating boolean masks, as they directly preserve the exact structure of the original array:
is_even = (numbers % 2) == 0
is_even
[[True, False, True, False], [True, False, True], [True, False, True, False, True]] ---------------------------------- backend: cpu nbytes: 44 B type: 3 * var * bool
numbers
[[0, 1, 2, 3], [4, 5, 6], [8, 9, 10, 11, 12]] --------------------- backend: cpu nbytes: 128 B type: 3 * var * int64
Now we can use is_even to slice numbers:
numbers[is_even]
[[0, 2], [4, 6], [8, 10, 12]] --------------------- backend: cpu nbytes: 88 B type: 3 * var * int64
Note that this is different to what would happen with NumPy’s boolean indexing:
numbers_np = np.array(
[
[0, 1, 2, 3],
[4, 5, 6, 7],
[8, 9, 10, 11],
]
)
numbers_np[(numbers_np % 2) == 0]
array([ 0, 2, 4, 6, 8, 10])
NumPy, lacking a ragged array structure, has to flatten the result whereas Awkward Array preserves the number of dimensions in the result.
numbers[
[[True, False, True, False],
[False],
[False, True, False]]
]
[[0, 2], [], [9]] --------------------- backend: cpu nbytes: 56 B type: 3 * var * int64