Create awkward array of same layout from numpy-compatible function

[jrueb]

I often encounter the problem of wanting to evaluate a function that is numpy-compatible on awkward arrays. Let's take as an example scipy.stats.gamma.pdf. I can do

a = np.array([1, 2])
x = a + 2
p = scipy.stats.gamma.pdf(x, a)

But if a is an awkward array (for example a2 = ak.Array([[[1, 2], []], [None], [[None]]])), I have to find another way.
I found this old comment about unwrapping the array manually with array.layout.contents: #489 (comment)
However, this is not trivial and handling everything correctly can be error-prone if the user is not so experienced with awkward layouts.
What I would like to have is a functionality to easily obtain the awkward array from a function like scipy.stats.gamma.pdf. Using above example, this functionality would easily generate an array equal to ak.Array([[p, []], [None], [[None]]]).
Maybe one can make a function called ak.map_content (or something like this), that does the unwrapping. For example:

a2 = ak.Array([[[1, 2], []], [None], [[None]]])
p2 = ak.map_content(scipy.stats.gamma.pdf, a2 + 2, a2)

Or maybe this functionality already exists? In this case, I am happy to learn more.

[agoose77]

Hi @jrueb, thanks for opening this discussion!

We have a function ak.transform that implements a recursive layout visitor. You can use this to walk down to the final level, whereby all arrays are 1D NumPy arrays:

def gamma_pdf(x, *args, loc=0, scale=1):
    first, *remainder = ak.broadcast_arrays(x, *args, loc, scale)

    def apply(inputs, **kwargs):
        if all(layout.is_numpy for layout in inputs):
            arrays = [x.to_backend_array() for x in inputs]
            return ak.to_layout(
                scipy.stats.gamma.pdf(*arrays)
            )

    return ak.transform(apply, first, *remainder, return_value="original")

Although ak.transform performs broadcasting, ak.broadcast_arrays implements an additional step that permits broadcasting scalars against arrays, so it's useful to start with an explicit broadcast.

The above implementation does not tell you if the transformer did nothing — maybe we should add that as a feature. You can quickly define an outer-scope variable to ensure this (or define a complex if statement to catch the exit condition):

def gamma_pdf(x, *args, loc=0, scale=1):
    first, *remainder = ak.broadcast_arrays(x, *args, loc, scale)

    transformer_did_return = False

    def apply(inputs, **kwargs):
        nonlocal transformer_did_return
        if all(layout.is_numpy for layout in inputs):
            arrays = [x.to_backend_array() for x in inputs]

            transformer_did_return = True
            return ak.to_layout(
                scipy.stats.gamma.pdf(*arrays)
            )

    result = ak.transform(apply, first, *remainder, return_value="original")

    if not transformer_did_return:
        raise TypeError

    return result

Hopefully this helps to answer your question?

[agoose77]

I will ping @jpivarski for a second opinion, but I've not made up my mind on this yet.

With #2595 and another PR to remove the need for ak.broadcast_arrays, we'd have something like

def gamma_pdf(x, *args, loc=0, scale=1):
    def apply(inputs, **kwargs):
        if all(layout.is_numpy for layout in inputs):
            arrays = [x.to_backend_array() for x in inputs]
            return ak.to_layout(
                scipy.stats.gamma.pdf(*arrays)
            )
    return ak.transform(apply, x, *args, loc, scale, return_value="original")

To me, this seems fairly concise, but I can equally see a need to keep users in our L1 "high-level user API" in cases where they're wanting to use these kind of functions in their analyses.

[grst]

A L1-level apply function could be useful for the users of my scverse/scirpy package, see also this discussion. Many of the users are actually python novices who just follow a tutorial, so I wouldn't dare to suggest them writing a custom transform function.

[agoose77]

@grst could you provide some examples of functions like this? In the case of your linked discussion, the new ak.str operations submodule looks like it will be useful!

>>> import awkward as ak
>>> data = ak.Array(
...     [
...         ["TRBC2*00", "TRBC2*00"],
...         ["TRBC1*00", "TRAC*00", "TRBC1*00"],
...         ["TRBC2*00", "TRBC2*00", "TRAC*00"],
...         ["TRBC2*00", "TRAC*00"],
...         ["TRBC2*00", "TRAC*00"],
...     ]
... )
>>> result = ak.str.split_pattern(data, "*")
>>> no_suffix = result[..., 0]
>>> no_suffix.show()
[['TRBC2', 'TRBC2'],
 ['TRBC1', 'TRAC', 'TRBC1'],
 ['TRBC2', 'TRBC2', 'TRAC'],
 ['TRBC2', 'TRAC'],
 ['TRBC2', 'TRAC']]

It sounds like we might benefit from having something like ak.apply_at_axis which uses ak.transform to operate at a given axis.

[grst]

I'll think about it again, but a .str interface likely covers most (if not all) of the things I had in mind. Is this already documented somewhere? I couldn't find that in the API docs

[agoose77]

@grst my English could have been clearer there — the new ak.str submodule is currently under develompent in #2616