How to resolve TypingErrors when the property of an ak.Array subclass relies on Numba-compiled functions?

Question

I've perused the docs and demo notebooks in the awkward1 repo (and it's entirely possible I missed something obvious), but I've strayed into unfamiliar territory and would like some help.

Suppose I have points in a polar coordinate frame that I provide to users as an array of records using the awkward1 library. Some users may prefer to work with the points in a Cartesian coordinate frame, so I can accommodate them by adding behaviors, and the following example works.

import awkward1 as ak
import numpy as np


class Polar2Cartesian(object):
    """Mixin class to convert from polar to Cartesian."""
    @property
    def x(self):
        return self.r * np.cos(self.theta)

    # And similarly for y, but this is just an example


# Register the behaviors with awkward
class Point(ak.Record, Polar2Cartesian):
    pass

class PointArray(ak.Array, Polar2Cartesian):
    pass

ak.behavior['Point'] = Point
ak.behavior['*', 'Point'] = PointArray

# Define the points
r = np.array([2.53, 0.29, 3.18])
theta = np.array([np.pi/4, np.pi, -np.pi/3])
points = ak.zip({'r': r, 'theta': theta}, with_name='Point')

# x-component of the first point
points[0].x == r[0] * np.cos(theta[0]) # True

# x-components of all points
points.x == r * np.cos(theta) # <Array [True, True, True] type='3 * bool'>

Now suppose the conversion must use a Numba-compiled function. I've defined one in the following example, but in principle it could come from a third-party library over which I have no control. The following works for individual records, but raises a TypingError over the array of records.

import awkward1 as ak
import numba as nb
import numpy as np


@nb.njit
def cos(x):
    """A contrived example appears."""
    return np.cos(x)

class Polar2Cartesian(object):
    """Mixin class to convert from polar to Cartesian."""
    @property
    def x(self):
        return self.r * cos(self.theta) # uses jitted function here

    # And similarly for y, but this is just an example

# Define and register the behaviors like before, and create
# the same array of Points as in the previous example

# This still works
points[0].x == r[0] * np.cos(theta[0]) # True

# This now raises a TypingError
points.x == r * np.cos(theta)

Parsing the traceback reveals the reason for the TypingError:

TypingError: can't resolve ufunc cos for types (awkward1.ArrayView(awkward1.NumpyArrayType(array(float64, 1d, A), none, {}), None, ()),)

That's fair, as it would be unreasonable to expect Numba to know how to use this type without help from the developer.

A workaround I have found is to call np.asarray() within the definition of Polar2Cartesian.x, i.e. cos(np.asarray(self.theta)). Although cos returns a NumPy array, Polar2Cartesian.x ends up becoming an awkward Array due to the subsequent multiplication with self.r. In general, one could convert the return value to an awkward record or array as needed.

But is this solution "awkward1-approved", or should I be going the route where I provide typer and lower functions for theta and x as additional behaviors? If the latter, could someone walk me through how to correctly write the typer and lower functions?

At this point, I wonder if the scope of my question title is too broad. If someone has a better suggestion, please let me know.