How to restructure arrays with zip/unzip and project#

Hide code cell content
%config InteractiveShell.ast_node_interactivity = "last_expr_or_assign"

Unzipping an array of records#

As discussed in How to create arrays of records, in addition to primitive types like numpy.float64 and numpy.datetime64, Awkward Arrays can also contain records. These records are formed from a fixed number of optionally named fields.

import awkward as ak
import numpy as np

records = ak.Array(
    [
        {"x": 1, "y": 1.1, "z": "one"},
        {"x": 2, "y": 2.2, "z": "two"},
        {"x": 3, "y": 3.3, "z": "three"},
        {"x": 4, "y": 4.4, "z": "four"},
        {"x": 5, "y": 5.5, "z": "five"},
    ]
)
[{x: 1, y: 1.1, z: 'one'},
 {x: 2, y: 2.2, z: 'two'},
 {x: 3, y: 3.3, z: 'three'},
 {x: 4, y: 4.4, z: 'four'},
 {x: 5, y: 5.5, z: 'five'}]
----------------------------
type: 5 * {
    x: int64,
    y: float64,
    z: string
}

Although it is useful to be able to create arrays from a sequence of records (as arrays of structures), Awkward Array implements arrays as structures of arrays. It is therefore more natural to think about arrays in terms of their fields. In the above example, we have created an array of records from a list of dictionaries. We can see that the x field of records contains five numpy.int64 values:

records.x
[1,
 2,
 3,
 4,
 5]
---------------
type: 5 * int64

If we wanted to look at each of the fields of records, we could pull them out individually from the array:

records.y
[1.1,
 2.2,
 3.3,
 4.4,
 5.5]
-----------------
type: 5 * float64
records.z
['one',
 'two',
 'three',
 'four',
 'five']
----------------
type: 5 * string

Clearly, for arrays with a large number of fields, retrieving each field in this manner would become tedious rather quickly. ak.unzip() can be used to directly build a tuple of the field arrays:

ak.unzip(records)
(<Array [1, 2, 3, 4, 5] type='5 * int64'>,
 <Array [1.1, 2.2, 3.3, 4.4, 5.5] type='5 * float64'>,
 <Array ['one', 'two', 'three', 'four', 'five'] type='5 * string'>)

Records are not required to have field names. A record without field names is known as a “tuple”, e.g.

tuples = ak.Array(
    [
        (1, 1.1, "one"),
        (2, 2.2, "two"),
        (3, 3.3, "three"),
        (4, 4.4, "four"),
        (5, 5.5, "five"),
    ]
)
[(1, 1.1, 'one'),
 (2, 2.2, 'two'),
 (3, 3.3, 'three'),
 (4, 4.4, 'four'),
 (5, 5.5, 'five')]
-------------------
type: 5 * (
    int64,
    float64,
    string
)

If we unzip an array of tuples, we obtain the same result as for records:

ak.unzip(tuples)
(<Array [1, 2, 3, 4, 5] type='5 * int64'>,
 <Array [1.1, 2.2, 3.3, 4.4, 5.5] type='5 * float64'>,
 <Array ['one', 'two', 'three', 'four', 'five'] type='5 * string'>)

ak.unzip() can be combined with ak.fields() to build a mapping from field name to field array:

dict(zip(ak.fields(records), ak.unzip(records)))
{'x': <Array [1, 2, 3, 4, 5] type='5 * int64'>,
 'y': <Array [1.1, 2.2, 3.3, 4.4, 5.5] type='5 * float64'>,
 'z': <Array ['one', 'two', 'three', 'four', 'five'] type='5 * string'>}

For tuples, the field names will be strings corresponding to the field index:

dict(zip(ak.fields(tuples), ak.unzip(tuples)))
{'0': <Array [1, 2, 3, 4, 5] type='5 * int64'>,
 '1': <Array [1.1, 2.2, 3.3, 4.4, 5.5] type='5 * float64'>,
 '2': <Array ['one', 'two', 'three', 'four', 'five'] type='5 * string'>}

Zipping together arrays#

Because Awkward Arrays unzip into distinct arrays, it is reasonable to ask whether the reverse is possible, i.e. given the following arrays

age = ak.Array([18, 32, 87, 55])
name = ak.Array(["Dorit", "Caitlin", "Theodor", "Albano"]);

can we form an array of records? The ak.zip() function provides a way to join compatible arrays into a single array of records:

people = ak.zip({"age": age, "name": name})
[{age: 18, name: 'Dorit'},
 {age: 32, name: 'Caitlin'},
 {age: 87, name: 'Theodor'},
 {age: 55, name: 'Albano'}]
----------------------------
type: 4 * {
    age: int64,
    name: string
}

Similarly, we could also build an array of tuples by passing a sequence of arrays:

ak.zip([age, name])
[(18, 'Dorit'),
 (32, 'Caitlin'),
 (87, 'Theodor'),
 (55, 'Albano')]
-----------------
type: 4 * (
    int64,
    string
)

Zipping and unzipping arrays is a lightweight operation, and so you should not hesitate to zip together arrays if it makes sense for the problem at hand. One of the benefits of combining arrays into an array of records is that slicing and masking operations are applied to all fields, e.g.

people[age > 35]
[{age: 87, name: 'Theodor'},
 {age: 55, name: 'Albano'}]
----------------------------
type: 2 * {
    age: int64,
    name: string
}

Arrays with different dimensions#

So far, we’ve looked at simple arrays with the same dimension in each field. It is actually possible to build arrays with fields of different dimensions, e.g.

x = ak.Array(
    [
        103,
        450,
        33,
        4,
    ]
)

digits_of_x = ak.Array(
    [
        [1, 0, 3],
        [4, 5, 0],
        [3, 3],
        [4],
    ]
)
x_and_digits = ak.zip({"x": x, "digits": digits_of_x})
[[{x: 103, digits: 1}, {x: 103, digits: 0}, {x: 103, digits: 3}],
 [{x: 450, digits: 4}, {x: 450, digits: 5}, {x: 450, digits: 0}],
 [{x: 33, digits: 3}, {x: 33, digits: 3}],
 [{x: 4, digits: 4}]]
-----------------------------------------------------------------
type: 4 * var * {
    x: int64,
    digits: int64
}

The type of this array is

x_and_digits.type
ArrayType(ListType(RecordType([NumpyType('int64'), NumpyType('int64')], ['x', 'digits'])), 4, None)

Note that the x field has changed type:

x.type
ArrayType(NumpyType('int64'), 4, None)
x_and_digits.x.type
ArrayType(ListType(NumpyType('int64')), 4, None)

In zipping the two arrays together, the x has been broadcast against digits_of_x. Sometimes you might want to limit the broadcasting to a particular depth (dimension). This can be done by passing the depth_limit parameter:

x_and_digits = ak.zip({"x": x, "digits": digits_of_x}, depth_limit=1)
[{x: 103, digits: [1, 0, 3]},
 {x: 450, digits: [4, 5, 0]},
 {x: 33, digits: [3, 3]},
 {x: 4, digits: [4]}]
-----------------------------
type: 4 * {
    x: int64,
    digits: var * int64
}

Now the x field has a single dimension

x_and_digits.x.type
ArrayType(NumpyType('int64'), 4, None)

Arrays with different dimension lengths#

What happens if we zip together arrays with the same dimensions, but different lengths in each dimensions?

x_and_y = ak.Array(
    [
        [103, 903],
        [450, 83],
        [33, 8],
        [4, 109],
    ]
)

digits_of_x_and_y = ak.Array(
    [
        [1, 0, 3, 9, 0, 3],
        [4, 5, 0, 8, 3],
        [3, 3, 8],
        [4, 1, 0, 9],
    ]
)

ak.zip({"x_and_y": x_and_y, "digits": digits_of_x_and_y})
---------------------------------------------------------------------------
ValueError                                Traceback (most recent call last)
Cell In[21], line 19
      1 x_and_y = ak.Array(
      2     [
      3         [103, 903],
   (...)
      7     ]
      8 )
     10 digits_of_x_and_y = ak.Array(
     11     [
     12         [1, 0, 3, 9, 0, 3],
   (...)
     16     ]
     17 )
---> 19 ak.zip({"x_and_y": x_and_y, "digits": digits_of_x_and_y})

File ~/micromamba/envs/awkward-docs/lib/python3.11/site-packages/awkward/_dispatch.py:64, in named_high_level_function.<locals>.dispatch(*args, **kwargs)
     62 # Failed to find a custom overload, so resume the original function
     63 try:
---> 64     next(gen_or_result)
     65 except StopIteration as err:
     66     return err.value

File ~/micromamba/envs/awkward-docs/lib/python3.11/site-packages/awkward/operations/ak_zip.py:151, in zip(arrays, depth_limit, parameters, with_name, right_broadcast, optiontype_outside_record, highlevel, behavior, attrs)
    148     yield arrays
    150 # Implementation
--> 151 return _impl(
    152     arrays,
    153     depth_limit,
    154     parameters,
    155     with_name,
    156     right_broadcast,
    157     optiontype_outside_record,
    158     highlevel,
    159     behavior,
    160     attrs,
    161 )

File ~/micromamba/envs/awkward-docs/lib/python3.11/site-packages/awkward/operations/ak_zip.py:241, in _impl(arrays, depth_limit, parameters, with_name, right_broadcast, optiontype_outside_record, highlevel, behavior, attrs)
    238     else:
    239         return None
--> 241 out = ak._broadcasting.broadcast_and_apply(
    242     layouts, action, right_broadcast=right_broadcast
    243 )
    244 assert isinstance(out, tuple) and len(out) == 1
    245 out = out[0]

File ~/micromamba/envs/awkward-docs/lib/python3.11/site-packages/awkward/_broadcasting.py:1140, in broadcast_and_apply(inputs, action, depth_context, lateral_context, allow_records, left_broadcast, right_broadcast, numpy_to_regular, regular_to_jagged, function_name, broadcast_parameters_rule)
   1138 backend = backend_of(*inputs, coerce_to_common=False)
   1139 isscalar = []
-> 1140 out = apply_step(
   1141     backend,
   1142     broadcast_pack(inputs, isscalar),
   1143     action,
   1144     0,
   1145     depth_context,
   1146     lateral_context,
   1147     {
   1148         "allow_records": allow_records,
   1149         "left_broadcast": left_broadcast,
   1150         "right_broadcast": right_broadcast,
   1151         "numpy_to_regular": numpy_to_regular,
   1152         "regular_to_jagged": regular_to_jagged,
   1153         "function_name": function_name,
   1154         "broadcast_parameters_rule": broadcast_parameters_rule,
   1155     },
   1156 )
   1157 assert isinstance(out, tuple)
   1158 return tuple(broadcast_unpack(x, isscalar) for x in out)

File ~/micromamba/envs/awkward-docs/lib/python3.11/site-packages/awkward/_broadcasting.py:1118, in apply_step(backend, inputs, action, depth, depth_context, lateral_context, options)
   1116     return result
   1117 elif result is None:
-> 1118     return continuation()
   1119 else:
   1120     raise AssertionError(result)

File ~/micromamba/envs/awkward-docs/lib/python3.11/site-packages/awkward/_broadcasting.py:1087, in apply_step.<locals>.continuation()
   1085 # Any non-string list-types?
   1086 elif any(x.is_list and not is_string_like(x) for x in contents):
-> 1087     return broadcast_any_list()
   1089 # Any RecordArrays?
   1090 elif any(x.is_record for x in contents):

File ~/micromamba/envs/awkward-docs/lib/python3.11/site-packages/awkward/_broadcasting.py:623, in apply_step.<locals>.broadcast_any_list()
    620         nextinputs.append(x)
    621         nextparameters.append(NO_PARAMETERS)
--> 623 outcontent = apply_step(
    624     backend,
    625     nextinputs,
    626     action,
    627     depth + 1,
    628     copy.copy(depth_context),
    629     lateral_context,
    630     options,
    631 )
    632 assert isinstance(outcontent, tuple)
    633 parameters = parameters_factory(nextparameters, len(outcontent))

File ~/micromamba/envs/awkward-docs/lib/python3.11/site-packages/awkward/_broadcasting.py:1118, in apply_step(backend, inputs, action, depth, depth_context, lateral_context, options)
   1116     return result
   1117 elif result is None:
-> 1118     return continuation()
   1119 else:
   1120     raise AssertionError(result)

File ~/micromamba/envs/awkward-docs/lib/python3.11/site-packages/awkward/_broadcasting.py:1087, in apply_step.<locals>.continuation()
   1085 # Any non-string list-types?
   1086 elif any(x.is_list and not is_string_like(x) for x in contents):
-> 1087     return broadcast_any_list()
   1089 # Any RecordArrays?
   1090 elif any(x.is_record for x in contents):

File ~/micromamba/envs/awkward-docs/lib/python3.11/site-packages/awkward/_broadcasting.py:672, in apply_step.<locals>.broadcast_any_list()
    670 for x, x_is_string in zip(inputs, input_is_string):
    671     if isinstance(x, listtypes) and not x_is_string:
--> 672         next_content = broadcast_to_offsets_avoiding_carry(x, offsets)
    673         nextinputs.append(next_content)
    674         nextparameters.append(x._parameters)

File ~/micromamba/envs/awkward-docs/lib/python3.11/site-packages/awkward/_broadcasting.py:372, in broadcast_to_offsets_avoiding_carry(list_content, offsets)
    370         return list_content.content[:next_length]
    371     else:
--> 372         return list_content._broadcast_tooffsets64(offsets).content
    373 elif isinstance(list_content, ListArray):
    374     # Is this list contiguous?
    375     if index_nplike.array_equal(
    376         list_content.starts.data[1:], list_content.stops.data[:-1]
    377     ):
    378         # Does this list match the offsets?

File ~/micromamba/envs/awkward-docs/lib/python3.11/site-packages/awkward/contents/listoffsetarray.py:412, in ListOffsetArray._broadcast_tooffsets64(self, offsets)
    407     next_content = self._content[this_start:]
    409 if index_nplike.known_data and not index_nplike.array_equal(
    410     this_zero_offsets, offsets.data
    411 ):
--> 412     raise ValueError("cannot broadcast nested list")
    414 return ListOffsetArray(
    415     offsets, next_content[: offsets[-1]], parameters=self._parameters
    416 )

ValueError: cannot broadcast nested list

This error occurred while calling

    ak.zip(
        {'x_and_y': <Array [[103, 903], [450, 83], [33, ...], [4, 109]] type=...
    )

Arrays which cannot be broadcast against each other will raise a ValueError. In this case, we want to stop broadcasting at the first dimension (depth_limit=1)

ak.zip({"x_and_y": x_and_y, "digits": digits_of_x_and_y}, depth_limit=1)
[{x_and_y: [103, 903], digits: [1, 0, 3, ..., 0, 3]},
 {x_and_y: [450, 83], digits: [4, 5, 0, 8, 3]},
 {x_and_y: [33, 8], digits: [3, 3, 8]},
 {x_and_y: [4, 109], digits: [4, 1, 0, 9]}]
-----------------------------------------------------
type: 4 * {
    x_and_y: var * int64,
    digits: var * int64
}

Projecting arrays#

Sometimes we are interested only in a subset of the fields of an array. For example, imagine that we have an array of coordinates on the \(\hat{x}\hat{y}\) plane:

triangle = ak.Array(
    [
        {"x": 1, "y": 6, "z": 0},
        {"x": 2, "y": 7, "z": 0},
        {"x": 3, "y": 8, "z": 0},
    ]
)
[{x: 1, y: 6, z: 0},
 {x: 2, y: 7, z: 0},
 {x: 3, y: 8, z: 0}]
--------------------
type: 3 * {
    x: int64,
    y: int64,
    z: int64
}

If we know that these points should lie on a plane, then we might wish to discard the \(\hat{z}\) coordinate. We can do this by slicing only the \(\hat{x}\) and \(\hat{y}\) fields:

triangle_2d = triangle[["x", "y"]]
[{x: 1, y: 6},
 {x: 2, y: 7},
 {x: 3, y: 8}]
--------------
type: 3 * {
    x: int64,
    y: int64
}

Note that the key passed to the subscript operator is a list ["x", "y"], not a tuple. Awkward Array recognises the list to mean “take both the "x" and "y" fields”.

Projections can be combined with array slicing and masking, e.g.

triangle_2d_first_2 = triangle[:2, ["x", "y"]]
[{x: 1, y: 6},
 {x: 2, y: 7}]
--------------
type: 2 * {
    x: int64,
    y: int64
}

Let’s now consider an array of triangles, i.e. a polygon:

triangles = ak.Array(
    [
        [
            {"x": 1, "y": 6, "z": 0},
            {"x": 2, "y": 7, "z": 0},
            {"x": 3, "y": 8, "z": 0},
        ],
        [
            {"x": 4, "y": 9, "z": 0},
            {"x": 5, "y": 10, "z": 0},
            {"x": 6, "y": 11, "z": 0},
        ],
    ]
)
[[{x: 1, y: 6, z: 0}, {x: 2, y: 7, z: 0}, {x: 3, y: 8, z: 0}],
 [{x: 4, y: 9, z: 0}, {x: 5, y: 10, z: 0}, {x: 6, y: 11, z: 0}]]
----------------------------------------------------------------
type: 2 * var * {
    x: int64,
    y: int64,
    z: int64
}

We can combine an int index 0 with a str projection to view the "x" coordinates of the first triangle vertices

triangles[0, "x"]
[1,
 2,
 3]
---------------
type: 3 * int64

We could even ignore the first vertex of each triangle

triangles[0, 1:, "x"]
[2,
 3]
---------------
type: 2 * int64

Projections commute (to the left) with other indices to produce the same result as their “natural” position. This means that the above projection could also be written as

triangles[0, "x", 1:]
[2,
 3]
---------------
type: 2 * int64

or even

triangles["x", 0, 1:]
[2,
 3]
---------------
type: 2 * int64

For columnar Awkward Arrays, there is no performance difference between any of these approaches; projecting the records of an array just changes its metadata, rather than invoking any loops over the data.

Projecting records-of-records#

The records of an array can themselves contain records

polygon = ak.Array(
    [
        {
            "vertex": [
                {"x": 1, "y": 6, "z": 0},
                {"x": 2, "y": 7, "z": 0},
                {"x": 3, "y": 8, "z": 0},
            ],
            "normal": [
                {"x": 0.164, "y": 0.986, "z": 0.0},
                {"x": 0.275, "y": 0.962, "z": 0.0},
                {"x": 0.351, "y": 0.936, "z": 0.0},
            ],
            "n_vertex": 3,
        },
        {
            "vertex": [
                {"x": 4, "y": 9, "z": 0},
                {"x": 5, "y": 10, "z": 0},
                {"x": 6, "y": 11, "z": 0},
                {"x": 7, "y": 12, "z": 0},
            ],
            "normal": [
                {"x": 0.406, "y": 0.914, "z": 0.0},
                {"x": 0.447, "y": 0.894, "z": 0.0},
                {"x": 0.470, "y": 0.878, "z": 0.0},
                {"x": 0.504, "y": 0.864, "z": 0.0},
            ],
            "n_vertex": 4,
        },
    ]
)
[{vertex: [{x: 1, y: 6, z: 0}, ..., {...}], normal: [...], n_vertex: 3},
 {vertex: [{x: 4, y: 9, z: 0}, ..., {...}], normal: [...], n_vertex: 4}]
------------------------------------------------------------------------
type: 2 * {
    vertex: var * {
        x: int64,
        y: int64,
        z: int64
    },
    normal: var * {
        x: float64,
        y: float64,
        z: float64
    },
    n_vertex: int64
}

Naturally we can access the "vertex" field with the . operator:

polygon.vertex
[[{x: 1, y: 6, z: 0}, {x: 2, y: 7, z: 0}, {x: 3, y: 8, z: 0}],
 [{x: 4, y: 9, z: 0}, {x: 5, y: 10, z: 0}, {...}, {x: 7, y: 12, z: 0}]]
-----------------------------------------------------------------------
type: 2 * var * {
    x: int64,
    y: int64,
    z: int64
}

We can view the "x" field of the vertex array with an additional lookup

polygon.vertex.x
[[1, 2, 3],
 [4, 5, 6, 7]]
---------------------
type: 2 * var * int64

The . operator represents the simplest slice of a single string, i.e.

polygon["vertex"]
[[{x: 1, y: 6, z: 0}, {x: 2, y: 7, z: 0}, {x: 3, y: 8, z: 0}],
 [{x: 4, y: 9, z: 0}, {x: 5, y: 10, z: 0}, {...}, {x: 7, y: 12, z: 0}]]
-----------------------------------------------------------------------
type: 2 * var * {
    x: int64,
    y: int64,
    z: int64
}

The slice corresponding to the nested lookup .vertex.x is given by a tuple of str:

polygon[("vertex", "x")]
[[1, 2, 3],
 [4, 5, 6, 7]]
---------------------
type: 2 * var * int64

It is even possible to combine multiple and single projections. Let’s project the "x" field of the "vertex" and "normal" fields:

polygon[["vertex", "normal"], "x"]
[{vertex: [1, 2, 3], normal: [0.164, ..., 0.351]},
 {vertex: [4, 5, 6, 7], normal: [0.406, ..., 0.504]}]
-----------------------------------------------------
type: 2 * {
    vertex: var * int64,
    normal: var * float64
}