How to filter lists within arrays using ragged slicing#
import awkward as ak
import numpy as np
What is ragged slicing?#
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: ragged 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]]] ---------------------------------------------- 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]]] ------------------------------- 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 ragged indexing.
Building a ragged index#
Ragged 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.
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)
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] --------------- type: 3 * int64
To put this more simply, the final dimension of the ragged 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 ragged 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] --------------- 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]] --------------------- 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]]] ---------------------------------------------- type: 3 * var * var * float64
ak.local_index(array)
[[[0, 1, 2], [0, 1, 2, 3], [0]], [], [[0, 1, 2, 3, 4]]] -------------------------------- type: 3 * var * var * int64
To create our ragged 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 ragged index looks like:
array[index]
[[[], [3.3], [7.7]], [], [[10.1, 11.1, 12.1]]] ----------------------------- type: 3 * var * var * float64
Clearly this index produces the result that we were aiming for!
Indexing with argmin and argmax#
Ragged 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.
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]] --------------------- type: 3 * var * int64
Without keepdims=True, all reducers collapse a dimension of the original array
ak.argmin(array, axis=1)
[2, 0, 2] ---------------- 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)]
[[8, 9, -1], [10, 3, 2, 9], [8, 9, -1]] ----------------------------- type: 3 * option[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:
ak.argmin(array, axis=-1, keepdims=True)
[[2], [0], [2]] -------------------- type: 3 * 1 * ?int64
array[ak.argmin(array, axis=-1, keepdims=True)]
[[2], [4], [-1]] ---------------------- type: 3 * var * ?int64
This now produces the expected result!
Filtering with missing sublists#
Ragged indexing supports using None in place of empty sublists within an index. For example
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]] --------------------- type: 4 * var * int64
Let’s use build a ragged index to pull some values out of array. Rather than using empty lists, we can use None to mask out sublists that we don’t care about:
array[
[
[0, 1],
None,
[],
[2],
],
]
[[10, 3], None, [], [-1]] ----------------------------- type: 4 * option[var * int64]
If we compare this with simply providing an empty sublist,
array[
[
[0, 1],
[],
[],
[2],
],
]
[[10, 3], [], [], [-1]] --------------------- type: 4 * var * int64
we can see that the None value introduces an
Filtering with booleans#
Awkward Array’s ragged indexing is a generalisation of the advanced indexing supported by NumPy. It is therefore reasonable to ask whether Awkward supports ragged 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 ragged 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 are powerful means of generating boolean masks, as they directly preserve the exact structure of the original array:
is_even = (numbers % 2) == 0
numbers
[[0, 1, 2, 3], [4, 5, 6], [8, 9, 10, 11, 12]] --------------------- type: 3 * var * int64
is_even
[[True, False, True, False], [True, False, True], [True, False, True, False, True]] ---------------------------------- type: 3 * var * bool
Now we can use is_even to slice numbers:
numbers[is_even]
[[0, 2], [4, 6], [8, 10, 12]] --------------------- 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.