diff --git a/src/awkward/_connect/numpy.py b/src/awkward/_connect/numpy.py index 15273bd334..b971a89f03 100644 --- a/src/awkward/_connect/numpy.py +++ b/src/awkward/_connect/numpy.py @@ -23,6 +23,7 @@ from awkward._typing import Iterator from awkward._util import Sentinel from awkward.contents.numpyarray import NumpyArray +from awkward.units import get_unit_registry # NumPy 1.13.1 introduced NEP13, without which Awkward ufuncs won't work, which # would be worse than lacking a feature: it would cause unexpected output. @@ -211,6 +212,49 @@ def _array_ufunc_signature(ufunc, inputs): return tuple(signature) +def _array_ufunc_custom_units(ufunc, inputs, kwargs, behavior): + registry = get_unit_registry() + if registry is None: + return None + + # Check if we have units + for x in inputs: + if isinstance(x, ak.contents.Content) and x.parameter("__units__"): + break + elif isinstance(x, registry.Quantity): + break + # Exit now, if not! + else: + return None + + # Wrap all Awkward Arrays with + nextinputs = [] + for x in inputs: + if isinstance(x, ak.contents.Content): + nextinputs.append( + registry.Quantity( + ak.with_parameter(x, "__units__", None, behavior=behavior), + x.parameter("__units__"), + ) + ) + else: + nextinputs.append(x) + out = ufunc(*nextinputs, **kwargs) + if not isinstance(out, tuple): + out = (out,) + + nextout = [] + for qty in out: + assert isinstance(qty, registry.Quantity) + assert isinstance(qty.magnitude, ak.Array) + nextout.append( + ak.with_parameter( + qty.magnitude, "__units__", str(qty.units), highlevel=False + ) + ) + return tuple(nextout) + + def array_ufunc(ufunc, method, inputs, kwargs): if method != "__call__" or len(inputs) == 0 or "out" in kwargs: return NotImplemented @@ -221,12 +265,35 @@ def array_ufunc(ufunc, method, inputs, kwargs): inputs = _array_ufunc_custom_cast(inputs, behavior, backend) def action(inputs, **ignore): + # Do we have any units in the mix? If so, delegate to `pint` to perform + # the ufunc dispatch. This will re-enter `array_ufunc`, but without units + # NOTE: there's nothing preventing us from handling units for non-NumpyArray + # contents, but for now we restrict ourselves to NumpyArray (in the + # NumpyArray constructor). By running _before_ the custom machinery, + # custom user ufuncs can avoid needing to worry about units + out = _array_ufunc_custom_units(ufunc, inputs, kwargs, behavior) + if out is not None: + return out + signature = _array_ufunc_signature(ufunc, inputs) custom = find_ufunc(behavior, signature) - # Do we have a custom ufunc (an override of the given ufunc)? + # Do we have a custom specific ufunc (an override of the given ufunc)? if custom is not None: return _array_ufunc_adjust(custom, inputs, kwargs, behavior) + # Do we have a custom generic ufunc override (a function that accepts _all_ ufuncs)? + for x in inputs: + if not isinstance(x, ak.contents.Content): + continue + apply_ufunc = find_ufunc_generic(ufunc, x, behavior) + if apply_ufunc is None: + continue + out = _array_ufunc_adjust_apply( + apply_ufunc, ufunc, method, inputs, kwargs, behavior + ) + if out is not None: + return out + if ufunc is numpy.matmul: raise NotImplementedError( "matrix multiplication (`@` or `np.matmul`) is not yet implemented for Awkward Arrays" @@ -259,17 +326,6 @@ def action(inputs, **ignore): return (NumpyArray(result, backend=backend, parameters=parameters),) - # Do we have a custom generic ufunc override (a function that accepts _all_ ufuncs)? - for x in inputs: - if isinstance(x, ak.contents.Content): - apply_ufunc = find_ufunc_generic(ufunc, x, behavior) - if apply_ufunc is not None: - out = _array_ufunc_adjust_apply( - apply_ufunc, ufunc, method, inputs, kwargs, behavior - ) - if out is not None: - return out - if all( x.parameter("__array__") is not None or x.parameter("__record__") is not None @@ -295,7 +351,7 @@ def action(inputs, **ignore): return None - if sum(int(isinstance(x, ak.contents.Content)) for x in inputs) == 1: + if sum(int(isinstance(x, ak.contents.Content)) for x in inputs) == 1 and 0: where = None for i, x in enumerate(inputs): if isinstance(x, ak.contents.Content): diff --git a/src/awkward/contents/content.py b/src/awkward/contents/content.py index 340424dbcb..15e062915d 100644 --- a/src/awkward/contents/content.py +++ b/src/awkward/contents/content.py @@ -102,6 +102,13 @@ def _init(self, parameters: dict[str, Any] | None, backend: Backend): ) ) else: + if not self.is_numpy and parameters.get("__units__") is not None: + raise TypeError( + '{} is not allowed to have parameters["__units__"] != None'.format( + type(self).__name__, + ) + ) + if not self.is_list and parameters.get("__array__") in ( "string", "bytestring", diff --git a/src/awkward/units.py b/src/awkward/units.py new file mode 100644 index 0000000000..82a36908c1 --- /dev/null +++ b/src/awkward/units.py @@ -0,0 +1,59 @@ +# BSD 3-Clause License; see https://github.com/scikit-hep/awkward-1.0/blob/main/LICENSE + +from __future__ import annotations + +from threading import RLock + +_lock = RLock() +_unit_registry = None +_checked_for_pint = False + + +def set_unit_registry(registry): + global _unit_registry + with _lock: + _unit_registry = registry + + +def get_unit_registry(): + with _lock: + _register_if_available() + return _unit_registry + + +def _register_if_available(): + global _checked_for_pint, _unit_registry + with _lock: + if _checked_for_pint: + return + + try: + import pint + except ModuleNotFoundError: + return + else: + _unit_registry = pint.UnitRegistry() + finally: + _checked_for_pint = True + + +def register_and_check(): + """ + Build `pint` unit registry + """ + try: + import pint # noqa: F401 + + except ModuleNotFoundError: + raise ModuleNotFoundError( + """install the 'pint' package with: + + python3 -m pip install pint + + or + + conda install -c conda-forge pint + """ + ) from None + + _register_if_available()