feat: named axis for `ak.Array` by pfackeldey · Pull Request #3238 · scikit-hep/awkward

pfackeldey · 2024-09-12T17:52:50Z

Proposal for named axis

This PR addresses #2596.

References for other named axis implementations:

Motivation

As argumented at PyHEP.dev 2023 and by the Harvard NLP group in their "Tensor Considered Harmful" write-up, named axis can be a powerful tool to make code more readable and less error-prone.

Design

`ak.Array` with named axis

Named axis are implemented through a mapping from named axis to positional axis.
named axis are hashables (currently restricted to strings), except for integers (and None) as they are reserved for positional axis.

import typing

AxisName: typing.Alias = typing.Hashable

By default an ak.Array uses positional axis, but named axis can be added to the array in the following ways:

import awkward as ak

# tuple:
#   positional axis: (0, 1)
#   named axis: {"events": 0, "jets": 1}
array = ak.Array([[1, 2], [3], [], [4, 5, 6]], named_axis=("events", "jets"))

# dict:
#   positional axis: (0, 1)
#   named axis:  {"events": 0, "jets": 1}
array = ak.Array([[1, 2], [3], [], [4, 5, 6]], named_axis={"events": 0, "jets": 1})

# the dict interface allows to name single axis, also negative positional axis
array = ak.Array([[1, 2], [3], [], [4, 5, 6]], named_axis={"jets": -1})

# attach axis naming to an existing array
array = ak.Array([[1, 2], [3], [], [4, 5, 6]])
array = ak.with_named_axis(array, ("events", "jets"))
# or
array = ak.with_named_axis(array, {"events": 0, "jets": 1})

The named_axis argument of the constructor of an ak.Array is a either tuple of AxisName or a dict of AxisName to integers.
It is stored in the .attrs attribute of the array with a reserved key "__named_axis__" of type dict[AxisName, int].
The two types of axis can be accessed through the named_axis and positional_axis property (always represented as a tuple):

import awkward as ak

array = ak.Array([[1, 2], [3], [], [4, 5, 6]], named_axis=("events", "jets"))
array.named_axis
>>>  {"events": 0, "jets": 1}
array.positional_axis
>>> (0, 1)

array = ak.Array([[1, 2], [3], [], [4, 5, 6]], named_axis={"jets": -1})
array.named_axis
>>> {"jets": -1}

Named axis in high-level functions

Named axis can be used by all high-level functions, e.g. ak.sum, ak.max, etc.:

import awkward as ak

array = ak.Array([[1, 2], [3], [], [4, 5, 6]], named_axis=("events", "jets"))

# sum over the "jets" axis
sum_jets = ak.sum(array, axis="jets")
>>> ak.Array([3, 3, 0, 15])
sum_jets.named_axis
>>> {"events": 0}

# the `keepdims=True` argument keeps the named axis
sum_jets = ak.sum(array, axis="jets", keepdims=True)
>>> ak.Array([[3], [3], [], [15]])
sum_jets.named_axis
>>> {"events": 0, "jets": 1}

There are different scenarios how named axis are propagated to the resulting array:

Named axis are unchanged e.g. ak.sum(array, axis="jets", keepdims=True) or array ** 2.
Named axis are removed from the resulting array e.g. through reduction ak.sum(array, axis="jets").
Named axis are unified e.g. in binary operations of two ak.Array or broadcasting:

import awkward as ak

array1 = ak.Array([[1, 2], [3, 4]], named_axis=("In", None))
array2 = ak.Array([[5, 6], [7, 8]], named_axis=(None, "Out"))

(array1 + array2).named_axis
>>> {"In": 0, "Out": 1}

Here, checks for matching named axis are performed, the rules are:

ak.Array([1], named_axis=("foo",)) + ak.Array([1], named_axis=("foo",))    # OK
ak.Array([1], named_axis=("foo",)) + ak.Array([1], named_axis=(None,))     # OK
ak.Array([1], named_axis=("foo",)) + ak.Array([1], named_axis=("bar",))    # raise Exception

Named axis in indexing

In addition, named axis can be used to select data:

import awkward as ak

array = ak.Array([[1, 2], [3], [], [4, 5, 6]], named_axis=("events", "jets"))

# select the first event
array[{"events": 0, "jets": slice(None)}]
>>> <Array [1, 2] jets:0 type='2 * int64'>

# select the first jet of each event
array[{"events": slice(None), "jets": slice(0, 1)}]
>>> <Array [[1], [3], [], [4]] events:0,jets:1 type='4 * var * int64'>

# mixed positional & named indexing
array[{0: slice(None), "jets": slice(0, 1)}]
>>> <Array [[1], [3], [], [4]] events:0,jets:1 type='4 * var * int64'>

For synthatic sugar use np.s_:

import awkward as ak

array = ak.Array([[1, 2], [3], [], [4, 5, 6]], named_axis=("events", "jets"))

# select the first jet of each event
array[{"events": np.s_[...], "jets": np.s_[0:1]}]
>>> <Array [[1], [3], [], [4]] events:0,jets:1 type='4 * var * int64'>

# or mixed with positional axis
array[{0: np.s_[...], "jets": np.s_[0:1]}]
>>> <Array [[1], [3], [], [4]] events:0,jets:1 type='4 * var * int64'>

This PR has to touch a lot of code and needs to add custom named axis propagation to each high-level operation. Thus, this PR is currently in draft mode.

Looking forward to ideas, thoughts, feedback on this effort!

pfackeldey · 2024-09-13T19:09:37Z

jpivarski · 2024-09-13T19:13:10Z

And all the data types that can be passed into square brackets with __getitem__.

codecov · 2024-09-13T20:47:05Z

Codecov Report

Attention: Patch coverage is 90.83527% with 79 lines in your changes missing coverage. Please review.

Project coverage is 82.27%. Comparing base (b749e49) to head (83c0aa6).
Report is 176 commits behind head on main.

Files with missing lines	Patch %	Lines
src/awkward/highlevel.py	45.91%	53 Missing ⚠️
src/awkward/_namedaxis.py	91.09%	17 Missing ⚠️
src/awkward/_broadcasting.py	92.85%	2 Missing ⚠️
src/awkward/operations/ak_without_named_axis.py	86.66%	2 Missing ⚠️
src/awkward/_regularize.py	92.30%	1 Missing ⚠️
src/awkward/_typing.py	50.00%	1 Missing ⚠️
src/awkward/contents/content.py	97.91%	1 Missing ⚠️
src/awkward/operations/ak_corr.py	87.50%	1 Missing ⚠️
src/awkward/operations/ak_covar.py	87.50%	1 Missing ⚠️

Additional details and impacted files

Files with missing lines	Coverage Δ
src/awkward/__init__.py	`97.14% <100.00%> (+0.08%)`	⬆️
src/awkward/_connect/numexpr.py	`91.76% <100.00%> (+0.97%)`	⬆️
src/awkward/_connect/numpy.py	`92.06% <100.00%> (+0.04%)`	⬆️
src/awkward/_layout.py	`87.56% <100.00%> (+1.31%)`	⬆️
src/awkward/_nplikes/array_like.py	`97.14% <ø> (+27.75%)`	⬆️
src/awkward/_nplikes/typetracer.py	`75.05% <ø> (+0.19%)`	⬆️
src/awkward/_operators.py	`94.91% <ø> (ø)`
src/awkward/contents/numpyarray.py	`90.50% <100.00%> (-1.01%)`	⬇️
src/awkward/operations/__init__.py	`100.00% <100.00%> (ø)`
src/awkward/operations/ak_all.py	`96.66% <100.00%> (+1.01%)`	⬆️
... and 68 more

... and 84 files with indirect coverage changes

…xing named axis propagation

…x named axis propagation in indexing for type tracers

…ak.mean; remove inplace addition of arrays from test

… Ellipsis failed

…tible with branched structures;fix regularize_axis in all highlevel ops

…ranch_depth using inner_shape property

…at depended on broadcasting to work)

pfackeldey · 2024-10-07T16:21:59Z

The PR is ready for review @agoose77 and @jpivarski. If you know why the windows tests are failing, please let me know if you have an idea... (I thought xfail would mark them as expected to fail and thus the test would pass?)

I'm sorry that the PR got so huge, but the largest part are tests related to named axes and translating named axes to positional axes in each high-level function.
I'd be especially happy for your review of named axes propagation in indexing and broadcasting operations as they are the most complicated ones.

agoose77 · 2024-10-07T16:36:29Z

@pfackeldey the xfail is failing because it's ... succeeding (xpass) 🤣

jpivarski · 2024-10-07T18:08:13Z

I'd like to see the rendered documentation, but the Deploy Branch Preview jobs are getting skipped. I don't know why: they were running PRs relatively recently.

Meanwhile, I'm still going over the rest of the PR.

Edit: It's because this is a fork, not a branch. That's okay; the markdown looks good; we'll see the rendered version after merging.

jpivarski

It's hard to review such a large PR, but we went over it at length during development. I think we iterated much more quickly because of that, and so I don't have much to say here, on the PR itself. The spot-check issues below are all minor. Still, we should check in again before the final merger, since everybody working on the codebase needs to be ready to incorporate the update when it comes.

docs/user-guide/how-to-array-properties-named-axis.md

src/awkward/_namedaxis.py

src/awkward/highlevel.py

src/awkward/operations/ak_local_index.py

src/awkward/operations/ak_singletons.py

src/awkward/operations/ak_sort.py

src/awkward/operations/ak_std.py

src/awkward/operations/ak_sum.py

src/awkward/operations/ak_unflatten.py

jpivarski · 2024-10-07T18:45:34Z

Oh, and I should have mentioned that this is very high quality code (type annotations, docstrings, comments)! Thank you!

pfackeldey · 2024-10-07T19:18:08Z

Thank you very much for your review @jpivarski! I'll add your suggestions soon 👍

agoose77

A small review -- I'll pop back tomorrow.

src/awkward/highlevel.py

src/awkward/_namedaxis.py

pfackeldey and others added 4 commits September 12, 2024 13:24

start implementing named axis for awkward array

eb5c188

style: pre-commit fixes

3978f43

add support for named axis for first batch of highlevel functions

556362e

formatting

16a119b

pfackeldey changed the title ~~Feat: named axis for ak.Array~~ feat: named axis for ak.Array Sep 12, 2024

next batch of high-level functions

c96f61e

pfackeldey and others added 3 commits September 13, 2024 16:38

fix type hints & safer control flow when checking for named axis

7f7db85

style: pre-commit fixes

572051e

(hopefully) fix old (<3.10) python type annotation syntax

e57e2c3

pfackeldey and others added 18 commits September 13, 2024 16:47

(hopefully) fix old (<3.10) python type annotation syntax

e338332

next batch of highlevel functions

a07adfa

next batch of highlevel functions

4c51e6b

style: pre-commit fixes

d05517a

update named axis implementation to not use tuples at all; start inde…

a705981

…xing named axis propagation

add named axis propagation for binary ops, some highlevel ops, and fi…

1ded71b

…x named axis propagation in indexing for type tracers

Merge remote-tracking branch 'upstream/main' into feat/named_axes

7158d07

fix keepdims in ak.covar & ak.corr; properly propagate named axis in …

00669a3

…ak.mean; remove inplace addition of arrays from test

add ak.std & ak.var; fix bug in indexing where == comparisons against…

bb97999

… Ellipsis failed

add ak.(arg)combinations and ak.(arg)cartesian; make named axis compa…

8720a05

…tible with branched structures;fix regularize_axis in all highlevel ops

avoid touching shape too much for purelist_depth, minmax_depth, and b…

1af4376

…ranch_depth using inner_shape property

ak.without_named_axis: allow ak.Records

01b459c

ak.with_named_axis: add check to validate the given named axis mapping

c970d06

fix doc strings and remove obsolete functions in _namedaxis.py module

f5f9495

update Slicer doc string

51bc6d6

docs: add documentation page for named axes

3bb8efa

Merge branch 'main' into feat/named_axes

5774280

propagate named axis through broadcasting; add more highlevel ops (th…

a63c4c3

…at depended on broadcasting to work)

Merge branch 'main' into feat/named_axes

6c376f0

pfackeldey requested review from agoose77 and jpivarski October 7, 2024 16:22

make xfail marker not strict

a2f6fcb

jpivarski approved these changes Oct 7, 2024

View reviewed changes

agoose77 added 5 commits October 8, 2024 09:55

chore: enable mypy on namedaxis

72565a9

fix: use attribute name in type hints

e0fa11b

refactor: rename _NamedAxisKey

74f6950

fix: define NAMED_AXIS_KEY as a literal

edc8290

fix: appease mypy

197de34

agoose77 reviewed Oct 8, 2024

View reviewed changes

pfackeldey added 4 commits October 8, 2024 08:56

highlevel getitem: less (un)wrapping

053f018

highlevel: improve repr and .show for named_axis

4bbd5b1

named_axis: fix typo and avoid dictionary copies where possible

f515ea7

named_axis: improve doc string of _prettify_named_axes

413c86c

jpivarski mentioned this pull request Oct 8, 2024

We should control the type of ak.Array.attrs, and it should be immutable #3277

Closed

pfackeldey added 6 commits October 9, 2024 08:54

highlevel: fix instance check for attrs in __init__

e3fc685

add named_axis to jupyter repr

894d1d0

fix docs

fb001b4

connect _neg2pos_axis with maybe_posaxis

df924c0

named_axis: simplify type hint for named axis in attrs mapping

595b240

regularize axis makes sure now that its type is either int or None

83c0aa6

pfackeldey merged commit cd7d7f6 into scikit-hep:main Oct 11, 2024

pfackeldey mentioned this pull request Nov 21, 2024

Add support for annotated axis names #2596

Closed

NicoWeio mentioned this pull request Apr 30, 2025

Inaccurate typing / documentation of axis parameters #3488

Closed

Conversation

pfackeldey commented Sep 12, 2024 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Proposal for named axis

Motivation

Design

ak.Array with named axis

Named axis in high-level functions

Named axis in indexing

Uh oh!

pfackeldey commented Sep 13, 2024 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Progress

general

broadcasting

slicing

Unary and binary operations

high-level functions

Independent of named axis: improvements / bugs found that are fixed by this PR aswell:

Uh oh!

jpivarski commented Sep 13, 2024

Uh oh!

codecov bot commented Sep 13, 2024 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Codecov Report

Uh oh!

pfackeldey commented Oct 7, 2024 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

agoose77 commented Oct 7, 2024 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

jpivarski commented Oct 7, 2024 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

jpivarski left a comment

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

jpivarski commented Oct 7, 2024

Uh oh!

pfackeldey commented Oct 7, 2024

Uh oh!

agoose77 left a comment

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

3 participants

pfackeldey commented Sep 12, 2024 •

edited

Loading

`ak.Array` with named axis

pfackeldey commented Sep 13, 2024 •

edited

Loading

codecov bot commented Sep 13, 2024 •

edited

Loading

pfackeldey commented Oct 7, 2024 •

edited

Loading

agoose77 commented Oct 7, 2024 •

edited

Loading

jpivarski commented Oct 7, 2024 •

edited

Loading