How could I apply a function of multiple irregular fields without iterating along the axis?

[imdinu]

I would like to be able to do array operations across different fields but I can't figure out how. Let's take the mean as an example. Here I have irregular size arrays for the x and y fields across different entries, but for each individual entry the length of those two arrays is the same. How could I compute the element-wise mean between "x" and "y" for each entry?

data = [
    {"x": np.array([1, 2, 3]), "y": np.array([2, 3, 4])}, 
    {"x": np.array([5, 6]), "y": np.array([6, 7])}, 
    {"x": np.array([0, 0, 0, 0]), "y":np.array([1, 1, 1, 1])}
]
arr = ak.Array(data)

The the desired result would be:

xy_mean = [
    {"avg": [1.5, 2.5, 3.5]},
    {"avg": [5.5, 6.5]},
    {"avg": [0.5, 0.5, 0.5, 0.5]}
]

Note: ak.mean(arr) throws ValueError for any axis argument value except None

Is there any way of achieving this without iterating over each entry? Perhaps some clever uses of ak.behavior?

[imdinu]

Easy solution actually

Ok, I have tinkered around a bit more and realised the answer is really straight forward. I could just use vectorized addition and division:

# vectorized operations
xy_mean = (arr["x"]+arr["y"])/2

This is easily readable and scales very well compared to iterating over each row:

# iterative approach
def iterate(arr):
    return ak.Array(
        [np.mean([arr["x"][i], arr["y"][i]], axis=0) 
        for i in range(len(arr.layout))])

Time comparison

Suggestion

Wouldn't it make sense to have this be the behaviour of ak.mean(arr, axis=1) in this case?

[jpivarski]

This is a good solution, and you've shown what a difference vectorization makes!

As for making it the behavior of ak.mean, none of the Awkward functions compute across fields—record fields are not considered a dimension. What axis=1 does mean is to average over (variable length) lists within an array. It wouldn't be possible for it to mean both.

If you have a situation in which it would be convenient to not have to manually write the names of the fields (maybe because you don't know them, it because they're different in different datasets and you want to write general code), there's ak.fields to give you a list of key names and ak.unzip to operate over the fields themselves.

sum(ak.unzip(arr))/len(ak.fields(arr))

(Untested: I hope that Python's sum will take the first two items to start the summation and not need or assume an identity element, since it's not zero in this case.)

[imdinu]

Thank you for your response. I understand the point about fields not being a considered a dimension. Sometimes it's a bit hard to wrap my head around those irregular data structures 😅.

I have tested your suggestion and it works just fine. It appears to be roughly 20% slower than doing the operations explicitly, but both methods have the same order of complexity so it is not a noticeable difference. (Both run in a about a few milliseconds for 20K entries)

[agoose77]

I suspect you can marginally improve the performance by specifying a start value for sum; it uses 0 by default for the start value, implying n+1 summations. To be clear, this is not an important optimization; the overhead of one additional field addition doesn't scale with the array elements, and is only "significant" for a small number fields.

Nevertheless, for fun, you can do this in a single line using the new walrus operator in Python 3.8+, or with the functional libraries operator and functools:

Walrus

sum((x := ak.unzip(arr))[1:], start=x[0]) / len(x)

operator and functools

functools.reduce(operator.add, ak.unzip(arr)) / len(ak.fields(arr))

In this case, I don't think the one-liners necessarily improve readability, so I'd just introduce a function

def sum_fieldwise(arr):
  head, *tail = ak.unzip(arr)
  return sum(tail, start=head)