How to flatten arrays, especially for plotting

In a data analysis, it is important to plot your data frequently, and the interactive nature of array-at-a-time functions facilitate that.

However, plotting views your data as a generic set or sequence—the structure of nested lists and records can’t be captured by standard plots. Histograms (including 2-dimensional heatmaps) take input data to be an unordered set, as do scatter plots. Connected-line plots, such as time-series, use the sequential order of the data, but there aren’t many visualizations that show nestedness. (Maybe there will be, in the future.)

As such, these standard plotting routines expect simple structures, either a single flat array (in which the order may be relevant or irrelevant) or several same-length arrays (in which the relative or absolute order is relevant). Encountering an Awkward Array, they may try to call np.asarray on it, which only works if the array can be made rectilinear or they may try to iterate over it in Python, which can be prohibitively slow if the dataset is large.

Scope of destructuring

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),

  • remove record structures,

  • remove missing data

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. Destructuring is an information-losing operation, so your guidance is required to eliminate exactly the structure you want to eliminate, and there are several ways to do that, depending on what you want to do.

After destructuring, you might still need to call np.asarray on the output because the plotting library might not recognize an ak.Array as an array. You’ll probably also want to develop your destructuring on a commandline or a different Jupyter cell from the plotting library function call, to understand what structure the output has without the added complication of the plotting library’s error messages.

import awkward as ak
import numpy as np

ak.flatten with axis=None

As mentioned above, ak.flatten with axis=None is the sledgehammer that turns any array into a 1-dimensional array with no nested lists, no nested records, and no missing data.

array = ak.Array([[{"x": 1.1, "y": [1]}, {"x": None, "y": [1, 2]}], [], [{"x": 3.3, "y": [1, 2, 3]}]])
array
<Array [[{x: 1.1, y: [1]}, ... y: [1, 2, 3]}]] type='3 * var * {"x": ?float64, "...'>
array.type
3 * var * {"x": ?float64, "y": var * int64}
ak.flatten(array, axis=None)
<Array [1.1, 3.3, 1, 1, 2, 1, 2, 3] type='8 * float64'>

Calling this function on an already flat array does nothing, so you don’t have to worry about what state your array had been in before you called it.

ak.flatten(ak.flatten(array, axis=None), axis=None)
<Array [1.1, 3.3, 1, 1, 2, 1, 2, 3] type='8 * float64'>

However, there are a few questions you should be asking yourself:

  • Did the nested lists have special meaning? What does the plot represent if I just concatenate them all?

  • Did the record fields have distinct meanings? In this example, what does it mean to put floating-point x values and nested-list y values in the same bucket of numbers to plot? Does it matter that there are more y values than x values? In most circumstances, you do not want to mix record fields in a plot.

  • It’s likely that we do want to ignore all the missing data, but does dropping them mean that an array representing x-axis values has lost its alignment with an array representing y-axis values?

Selecting record fields

A more controlled way to extract fields from a record is to project them by name.

array = ak.Array([[{"x": 1.1, "y": [1], "z": "one"}, {"x": None, "y": [1, 2], "z": "two"}], [], [{"x": 3.3, "y": [1, 2, 3], "z": "three"}]])
array
<Array [[{x: 1.1, y: [1], ... 3], z: 'three'}]] type='3 * var * {"x": ?float64, ...'>

If we want only the x field, we can ask for it as an attribute (because it’s a valid Python name) or with a string-valued slice:

array.x
<Array [[1.1, None], [], [3.3]] type='3 * var * ?float64'>
array["x"]
<Array [[1.1, None], [], [3.3]] type='3 * var * ?float64'>

This controls the biggest deficiency of ak.flatten with axis=None, the mixing of data with different meanings.

ak.flatten(array.x, axis=None)
<Array [1.1, 3.3] type='2 * float64'>
ak.flatten(array.y, axis=None)
<Array [1, 1, 2, 1, 2, 3] type='6 * int64'>

If some of your fields can be safely flattened—together into one set—and others can’t, you can use a list of strings to pick just the fields you want.

ak.flatten(array[["x", "y"]], axis=None)
<Array [1.1, 3.3, 1, 1, 2, 1, 2, 3] type='8 * float64'>

(Careful! A tuple has a special meaning in slices, which doesn’t apply here.)

array[("x", "y")]
---------------------------------------------------------------------------
ValueError                                Traceback (most recent call last)
/tmp/ipykernel_2307/2503192639.py in <module>
----> 1 array[("x", "y")]

~/python3.8/lib/python3.8/site-packages/awkward/highlevel.py in __getitem__(self, where)
    991         """
    992         if not hasattr(self, "_tracers"):
--> 993             tmp = ak._util.wrap(self.layout[where], self._behavior)
    994         else:
    995             tmp = ak._connect._jax.jax_utils._jaxtracers_getitem(self, where)

ValueError: cannot slice NumpyArray by field name

(https://github.com/scikit-hep/awkward-1.0/blob/1.4.0/src/libawkward/array/NumpyArray.cpp#L1290)

If you have records inside of records, you can extract them with nested projection if they have common names.

array = ak.Array([
    {"x": {"up": 1, "down": -1}, "y": {"up": 1.1, "down": -1.1}},
    {"x": {"up": 2, "down": -2}, "y": {"up": 2.2, "down": -2.2}},
    {"x": {"up": 3, "down": -3}, "y": {"up": 3.3, "down": -3.3}},
    {"x": {"up": 4, "down": -4}, "y": {"up": 4.4, "down": -4.4}},
])
array
<Array [{x: {up: 1, down: -1, ... down: -4.4}}] type='4 * {"x": {"up": int64, "d...'>
ak.flatten(array[["x", "y"], "up"], axis=None)
<Array [1, 2, 3, 4, 1.1, 2.2, 3.3, 4.4] type='8 * float64'>

ak.flatten for one axis

Since axis=None is so dangerous, the default value of ak.flatten is axis=1. This flattens only the first nested dimension.

ak.flatten(ak.Array([[0, 1, 2], [], [3, 4], [5], [6, 7, 8, 9]]))
<Array [0, 1, 2, 3, 4, 5, 6, 7, 8, 9] type='10 * int64'>

It also removes missing values in the axis that is being flattened because flattening considers a missing list like an empty list.

ak.flatten(ak.Array([[0, 1, 2], None, [3, 4], [5], [6, 7, 8, 9]]))
<Array [0, 1, 2, 3, 4, 5, 6, 7, 8, 9] type='10 * int64'>

It does not flatten or remove missing values from any other axis.

ak.flatten(ak.Array([[[0, 1, 2, 3, 4]], [], [[5], [6, 7, 8, 9]]]))
<Array [[0, 1, 2, 3, 4], [5], [6, 7, 8, 9]] type='3 * var * int64'>
ak.flatten(ak.Array([[[0, 1, 2, None]], [], [[5], [6, 7, 8, 9]]]))
<Array [[0, 1, 2, None], [5], [6, 7, 8, 9]] type='3 * var * ?int64'>

Moreover, you can’t flatten already-flat data because a 1-dimensional array does not have an axis=1. (axis starts counting at 0.)

ak.flatten(ak.Array([1, 2, 3, 4, 5]))
---------------------------------------------------------------------------
ValueError                                Traceback (most recent call last)
/tmp/ipykernel_2307/175060114.py in <module>
----> 1 ak.flatten(ak.Array([1, 2, 3, 4, 5]))

~/python3.8/lib/python3.8/site-packages/awkward/operations/structure.py in flatten(array, axis, highlevel, behavior)
   1896 
   1897     else:
-> 1898         out = layout.flatten(axis)
   1899 
   1900     return ak._util.maybe_wrap_like(out, array, behavior, highlevel)

ValueError: axis out of range for flatten

(https://github.com/scikit-hep/awkward-1.0/blob/1.4.0/src/libawkward/array/NumpyArray.cpp#L1663)

axis=0 is a valid option for ak.flatten, but since there can’t be any lists at this level, it only removes missing values.

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

Selecting one element from each list

Flattening removes list structure without removing values. Often, you want to do the opposite of that: you want to plot one element from each list. This makes the plot “aware” of your list structure.

This kind of operation is usually just a slice.

array = ak.Array([[0, 1, 2], [3, 4], [5], [6, 7, 8, 9]])
array
<Array [[0, 1, 2], [3, 4, ... 5], [6, 7, 8, 9]] type='4 * var * int64'>
array[:, 0]
<Array [0, 3, 5, 6] type='4 * int64'>

The above syntax selects all lists from the array (axis=0) and the first element from each list (axis=1). We could have as easily selected the last:

array[:, -1]
<Array [2, 4, 5, 9] type='4 * int64'>

A plot made from ak.flatten(array) would be a plot of all numbers with no knowledge of lists; a plot made from array[:, 0] would be a plot of lists, as represented by the first element in each. It depends on what you want to plot.

What if you get this error?

array = ak.Array([[0, 1, 2], [], [3, 4], [5], [6, 7, 8, 9]])
array
<Array [[0, 1, 2], [], ... [5], [6, 7, 8, 9]] type='5 * var * int64'>
array[:, 0]
---------------------------------------------------------------------------
ValueError                                Traceback (most recent call last)
/tmp/ipykernel_2307/566312820.py in <module>
----> 1 array[:, 0]

~/python3.8/lib/python3.8/site-packages/awkward/highlevel.py in __getitem__(self, where)
    991         """
    992         if not hasattr(self, "_tracers"):
--> 993             tmp = ak._util.wrap(self.layout[where], self._behavior)
    994         else:
    995             tmp = ak._connect._jax.jax_utils._jaxtracers_getitem(self, where)

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

(https://github.com/scikit-hep/awkward-1.0/blob/1.4.0/src/cpu-kernels/awkward_NumpyArray_getitem_next_at.cpp#L21)

It says that it can’t get element 0 of one of the lists, and that’s because this array contains an empty list.

One way to deal with that is to take a range-slice, rather than ask for an individual element from each list.

array[:, :1]
<Array [[0], [], [3], [5], [6]] type='5 * var * int64'>

But this array still has structure, so you can flatten it as an additional step.

ak.flatten(array[:, :1])
<Array [0, 3, 5, 6] type='4 * int64'>

Alternatively, you may want to attack the problem head-on: the issue is that some lists have too few elements, so why not remove those lists with an explicit slice? The ak.num function tells us the length of each nested list.

ak.num(array)
<Array [3, 0, 2, 1, 4] type='5 * int64'>
ak.num(array) > 0
<Array [True, False, True, True, True] type='5 * bool'>

Slicing the first dimension with this would ensure that the second dimension always has the element we seek.

array[ak.num(array) > 0, 0]
<Array [0, 3, 5, 6] type='4 * int64'>

The same applies if we’re taking the last element:

array[ak.num(array) > 0, -1]
<Array [2, 4, 5, 9] type='4 * int64'>

You can also do fancy things, requesting both the first and last element of each list, as long as it doesn’t run afoul of slicing rules (which were constrained to match NumPy’s in cases that overlap).

array[ak.num(array) > 0, [0, -1]]   # these two arrays have different lengths, can't be broadcasted as in NumPy advanced slicing
---------------------------------------------------------------------------
ValueError                                Traceback (most recent call last)
/tmp/ipykernel_2307/2672928886.py in <module>
----> 1 array[ak.num(array) > 0, [0, -1]]   # these two arrays have different lengths, can't be broadcasted as in NumPy advanced slicing

~/python3.8/lib/python3.8/site-packages/awkward/highlevel.py in __getitem__(self, where)
    991         """
    992         if not hasattr(self, "_tracers"):
--> 993             tmp = ak._util.wrap(self.layout[where], self._behavior)
    994         else:
    995             tmp = ak._connect._jax.jax_utils._jaxtracers_getitem(self, where)

ValueError: cannot broadcast arrays in slice

(https://github.com/scikit-hep/awkward-1.0/blob/1.4.0/src/libawkward/Slice.cpp#L1026)
array[ak.num(array) > 0][:, [0, -1]]   # so just put them in different slices
<Array [[0, 2], [3, 4], [5, 5], [6, 9]] type='4 * 2 * int64'>

And then flatten the result (if necessary—the shape is regular; some plotting libraries would interpret it as a single set of numbers).

ak.flatten(array[ak.num(array) > 0][:, [0, -1]])
<Array [0, 2, 3, 4, 5, 5, 6, 9] type='8 * int64'>

Aggregating each list

Reductions should be familiar to users of SQL and Pandas; after grouping data by some quantity, one must apply some aggregating operation on each group to get one number for each group. The one-element slices of the previous section are like SQL’s FIRST_VALUE and LAST_VALUE, which is a special case of reducing.

The architypical aggregation function is “sum,” which reduces a list by adding up its values. ak.sum and its relatives, ak.prod (product/multiplication), ak.mean, etc., are all reducers in Awkward Array.

Following NumPy, their default axis is None, but for this application, you’ll need to specify an explicit axis.

array = ak.Array([[0, 1, 2], [], [3, 4], [5], [6, 7, 8, 9]])
array
<Array [[0, 1, 2], [], ... [5], [6, 7, 8, 9]] type='5 * var * int64'>
ak.sum(array, axis=1)
<Array [3, 0, 7, 5, 30] type='5 * int64'>

Some of these are not defined for empty lists, so you’ll need to either replace the missing values with ak.fill_none or flatten them.

ak.mean(array, axis=1)
<Array [1, None, 3.5, 5, 7.5] type='5 * ?float64'>
ak.fill_none(ak.mean(array, axis=1), 0)   # fill with zero
<Array [1, 0, 3.5, 5, 7.5] type='5 * float64'>
ak.fill_none(ak.mean(array, axis=1), ak.mean(array))   # fill with the mean of all
<Array [1, 4.5, 3.5, 5, 7.5] type='5 * float64'>
ak.flatten(ak.mean(array, axis=1), axis=0)
<Array [1, 3.5, 5, 7.5] type='4 * float64'>

Each of these has a different effect: filling with 0 puts an identifiable value in the plot (a peak at 0 if it’s a histogram), filling with the overall mean imputes a value in missing cases, flattening away the missing values reduces the number of entries in the plot. Each of these has a different meaning when interpreting your plot!

Minimizing/maximizing over each list

Minimizing and maximizing are also reducers, ak.min and ak.max (and ak.ptp for the peak-to-peak difference between the minimum and maximum).

They deserve their own section because they are an important case.

array = ak.Array([[0, 2, 1], [], [4, 3], [5], [8, 6, 7, 9]])
array
<Array [[0, 2, 1], [], ... [5], [8, 6, 7, 9]] type='5 * var * int64'>
ak.min(array, axis=1)
<Array [0, None, 3, 5, 6] type='5 * ?int64'>
ak.max(array, axis=1)
<Array [2, None, 4, 5, 9] type='5 * ?int64'>

As before, they aren’t defined for empty lists, so you’ll have to choose a method to eliminate the missing values.

Sometimes, you want the “top N” elements from each list, rather than the “top 1.” Awkward Array doesn’t (yet) have a function for the “top N” elements, but it can be done with ak.sort and a slice.

ak.sort(array, axis=1)
<Array [[0, 1, 2], [], ... [5], [6, 7, 8, 9]] type='5 * var * int64'>
ak.sort(array, axis=1)[:, -2:]
<Array [[1, 2], [], [3, 4], [5], [8, 9]] type='5 * var * int64'>

We still have work to do: some of these lists are shorter than the 2 elements we asked for. What should be done with them? Eliminate all lists with fewer than two elements?

ak.sort(array[ak.num(array) >= 2], axis=1)[:, -2:]
<Array [[1, 2], [3, 4], [8, 9]] type='3 * var * int64'>

Or just concatenate everything so that we don’t lose the lists with only one value (5 in this example)?

ak.flatten(ak.sort(array, axis=1)[:, -2:])
<Array [1, 2, 3, 4, 5, 8, 9] type='7 * int64'>

Minimizing/maximizing lists of records

Unlike numbers, records do not have an ordering: you cannot call ak.min on an array of records. But usually, what you want to do instead is to find the minimum or maximum of some quantity calculated from the records and pick records (or record fields) from that.

array = ak.Array([
    [{"x": 2, "y": 2, "z": 2.2}, {"x": 1, "y": 1, "z": 1.1}, {"x": 3, "y": 3, "z": 3.3}],
    [],
    [{"x": 5, "y": 5, "z": 5.5}, {"x": 4, "y": 4, "z": 4.4}],
    [{"x": 7, "y": 7, "z": 7.7}, {"x": 9, "y": 9, "z": 9.9}, {"x": 8, "y": 8, "z": 8.8}, {"x": 6, "y": 6, "z": 6.6}],
])
array
<Array [[{x: 2, y: 2, z: 2.2, ... z: 6.6}]] type='4 * var * {"x": int64, "y": in...'>

The ak.argmin and ak.argmax functions return the integer index where the minimum or maximum of some numeric formula can be found.

np.sqrt(array.x**2 + array.y**2)
<Array [[2.83, 1.41, 4.24, ... 11.3, 8.49]] type='4 * var * float64'>
ak.argmax(np.sqrt(array.x**2 + array.y**2), axis=1)
<Array [2, None, 0, 1] type='4 * ?int64'>

These integer indexes can be used as slices if they don’t eliminate a dimension, which can be requested via keepdims=True. This makes a length-1 list for each reduced output.

maximize_by = ak.argmax(np.sqrt(array.x**2 + array.y**2), axis=1, keepdims=True)
maximize_by
<Array [[2], [None], [0], [1]] type='4 * var * ?int64'>

Applying this to the original array, we get the “best” record in each list, according to maximize_by.

array[maximize_by]
<Array [[{x: 3, y: 3, z: 3.3, ... z: 9.9}]] type='4 * var * ?{"x": int64, "y": i...'>
array[maximize_by].tolist()
[[{'x': 3, 'y': 3, 'z': 3.3}],
 [None],
 [{'x': 5, 'y': 5, 'z': 5.5}],
 [{'x': 9, 'y': 9, 'z': 9.9}]]

This still has list structures and missing values, so it’s ready for ak.flatten, assuming that we extract the appropriate record field to plot.

ak.flatten(array[maximize_by].z, axis=None)
<Array [3.3, 5.5, 9.9] type='3 * float64'>

Concatenating independently restructured arrays

Sometimes, what you want to do can’t be a single expression. Suppose we have this data:

array = ak.Array([[{"x": 1.1, "y": [1]}, {"x": 2.2, "y": [1, 2]}], [], [{"x": 3.3, "y": [1, 2, 3]}]])
array
<Array [[{x: 1.1, y: [1]}, ... y: [1, 2, 3]}]] type='3 * var * {"x": float64, "y...'>

and we want to combine all x values and the maximum y value in a plot. This requires a different expression on array.x from array.y.

ak.flatten(array.x)
<Array [1.1, 2.2, 3.3] type='3 * float64'>
ak.flatten(ak.max(array.y, axis=2), axis=None)
<Array [1, 2, 3] type='3 * int64'>

To get all of these into one array (because the plotting function only accepts one argument), you’ll need to ak.concatenate them.

ak.concatenate([
    ak.flatten(array.x),
    ak.flatten(ak.max(array.y, axis=2), axis=None),
])
<Array [1.1, 2.2, 3.3, 1, 2, 3] type='6 * float64'>

Maintaining alignment between arrays with missing values

Dropping missing values with ak.flatten doesn’t keep track of where they were removed. This is a problem if the plotting library takes separate sequences for the x-axis and y-axis, and these must be aligned.

Instead of ak.flatten, you can use ak.is_none.

array = ak.Array([
    {"x": 1, "y": 5.5},
    {"x": 2, "y": 3.3},
    {"x": None, "y": 2.2},
    {"x": 4, "y": None},
    {"x": 5, "y": 1.1},
])
array
<Array [{x: 1, y: 5.5}, ... {x: 5, y: 1.1}] type='5 * {"x": ?int64, "y": ?float64}'>
ak.is_none(array.x)
<Array [False, False, True, False, False] type='5 * bool'>
ak.is_none(array.y)
<Array [False, False, False, True, False] type='5 * bool'>
to_keep = ~(ak.is_none(array.x) | ak.is_none(array.y))
to_keep
<Array [True, True, False, False, True] type='5 * bool'>
array.x[to_keep], array.y[to_keep]
(<Array [1, 2, 5] type='3 * ?int64'>,
 <Array [5.5, 3.3, 1.1] type='3 * ?float64'>)

Actually drawing structure

If need be, you can change the plotter to match the data.

array = ak.Array([
    [{"x": 1, "y": 3.3}, {"x": 2, "y": 1.1}, {"x": 3, "y": 2.2}],
    [],
    [{"x": 4, "y": 5.5}, {"x": 5, "y": 4.4}],
    [{"x": 5, "y": 1.1}, {"x": 4, "y": 3.3}, {"x": 2, "y": 5.5}, {"x": 1, "y": 4.4}],
])
array
<Array [[{x: 1, y: 3.3}, ... x: 1, y: 4.4}]] type='4 * var * {"x": int64, "y": f...'>
import matplotlib.pyplot as plt
import matplotlib.path
import matplotlib.patches

fig, ax = plt.subplots()

for line in array:
    if len(line) > 0:
        vertices = np.dstack([np.asarray(line.x), np.asarray(line.y)])[0]
        codes = [matplotlib.path.Path.MOVETO] + [matplotlib.path.Path.LINETO] * (len(line) - 1)
        path = matplotlib.path.Path(vertices, codes)
        ax.add_patch(matplotlib.patches.PathPatch(path, facecolor="none"))

ax.set_xlim(0, 6)
ax.set_ylim(0, 6);
_images/how-to-restructure-flatten_109_0.png

(The above example assumes that len(array) is small enough to iterate over in Python, but vectorizes over each list in the array. It was adapted from the Matplotlib path tutorial.)