1

I am very new to awkward arrays (and python in general) and just want to know how I could add the first elements of each array in an ak.array together.

E.g my_array = [[1,2,3], [4,5,6], [7,8,9].....]

and I want 1 + 4 + 7...

idk_u_pick_a_name
  • 11
  • 1

2 Answers2

1

You can use the following code:

import awkward as ak

my_array = [[1, 2, 3], [4, 5, 6], [7, 8, 9]]

array = ak.Array(my_array)

#Use ak.firsts to extract the first elements of each sub-array
first_elements = ak.firsts(array)

#The ak.sum() function to add the first elements together
result = ak.sum(first_elements)

print(result)
Jim Pivarski
  • 5,568
  • 2
  • 35
  • 47
moaz ahmad
  • 174
  • 5
  • If you want to generalize beyond [ak.firsts](https://awkward-array.org/doc/main/reference/generated/ak.firsts.html) to the second or third, you can use a slice like `array[:, 0]` instead of `ak.firsts(array)` and then generalize to `array[:, 1]`, `array[:, 2]`, etc. One thing that [ak.firsts](https://awkward-array.org/doc/main/reference/generated/ak.firsts.html) does differently is that it inserts `None` when a first element does not exist, rather than raising an `IndexError`. – Jim Pivarski Jun 26 '23 at 17:49
0

If you only want the sum of first elements from all nested lists, then @moaz ahmad's answer is best. (See also my comment on it.)

But if this is a stepping stone toward the sum over all firsts, the sum over all seconds, the sum over all thirds, etc., then what you're describing is an axis=0 summation. (See the documentation for ak.sum.)

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

The first element of the result is 1 + 4 + 7 = 12, the second is 2 + 5 + 8 = 15, and the third is 3 + 6 + 9 = 18. You could take the first element of this result, but @moaz ahmad's answer involves less unnecessary summation.

This also works if the lengths of the lists in the array are not all equal to each other, as they are in your example. (Arrays of equal-length lists are regular 2-dimensional arrays, and NumPy is a better tool for that.) Consider the following irregular example:

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

The first element of the result is 1 + 3 = 4 (the None gets skipped; only relevant if your dataset has missing values), the second element is 2 + 4 + 5 = 11, the third element is 3 + 6 = 9, and the remaining elements are 7, 8, and 9 because they line up with blank spaces.

Jim Pivarski
  • 5,568
  • 2
  • 35
  • 47