fix: re-wrap awkward arrays in highlevel evaluator (#306)

* fix: re-wrap awkward arrays in highlevel evaluator

Closes #296

* Refactor as proposed

Author: nsmith-

src/correctionlib/convert.py

@@ -1,8 +1,10 @@
 """Tools to convert other formats to correctionlib"""
 
+from __future__ import annotations
+
 from collections.abc import Iterable, Sequence
 from numbers import Real
-from typing import TYPE_CHECKING, Any, Optional, Union, cast
+from typing import TYPE_CHECKING, Any, cast
 
 import numpy
 
@@ -30,7 +32,7 @@
 
 def from_uproot_THx(
     path: str,
-    axis_names: Optional[list[str]] = None,
+    axis_names: list[str] | None = None,
     flow: Literal["clamp", "error"] = "error",
 ) -> Correction:
     """Convert a ROOT histogram
@@ -51,17 +53,17 @@ def from_uproot_THx(
 
 
 def from_histogram(
-    hist: "PlottableHistogram",
-    axis_names: Optional[list[str]] = None,
-    flow: Optional[Union[Content, Literal["clamp", "error"]]] = "error",
+    hist: PlottableHistogram,
+    axis_names: list[str] | None = None,
+    flow: Content | Literal["clamp", "error"] | None = "error",
 ) -> Correction:
     """Read any object with PlottableHistogram interface protocol
 
     Interface as defined in
     https://github.com/scikit-hep/uhi/blob/v0.1.1/src/uhi/typing/plottable.py
     """
 
-    def read_axis(axis: "PlottableAxis", pos: int) -> Variable:
+    def read_axis(axis: PlottableAxis, pos: int) -> Variable:
         axtype = "real"
         if len(axis) == 0:
             raise ValueError(f"Zero-length axis {axis}, what to do?")
@@ -83,7 +85,7 @@ def read_axis(axis: "PlottableAxis", pos: int) -> Variable:
     variables = [read_axis(ax, i) for i, ax in enumerate(hist.axes)]
     # Here we could try to optimize the ordering
 
-    def edges(axis: "PlottableAxis") -> list[float]:
+    def edges(axis: PlottableAxis) -> list[float]:
         out = []
         for i, b in enumerate(axis):
             if isinstance(b, (str, int)):
@@ -96,16 +98,16 @@ def edges(axis: "PlottableAxis") -> list[float]:
                 out.append(b[1])
         return out
 
-    def flatten_to(values: "ndarray[Any, Any]", depth: int) -> Iterable[Any]:
+    def flatten_to(values: ndarray[Any, Any], depth: int) -> Iterable[Any]:
         for value in values:
             if depth > 0:
                 yield from flatten_to(value, depth - 1)
             else:
                 yield value
 
     def build_data(
-        values: "ndarray[Any, Any]",
-        axes: Sequence["PlottableAxis"],
+        values: ndarray[Any, Any],
+        axes: Sequence[PlottableAxis],
         variables: list[Variable],
     ) -> Content:
         vartype = variables[0].type
@@ -179,9 +181,9 @@ def build_data(
 
 
 def ndpolyfit(
-    points: list["ndarray[Any, Any]"],
-    values: "ndarray[Any, Any]",
-    weights: "ndarray[Any, Any]",
+    points: list[ndarray[Any, Any]],
+    values: ndarray[Any, Any],
+    weights: ndarray[Any, Any],
     varnames: list[str],
     degree: tuple[int],
 ) -> tuple[Correction, Any]:
@@ -216,7 +218,7 @@ def ndpolyfit(
         raise NotImplementedError(
             "correctionlib Formula not available for more than 4 variables"
         )
-    _degree: "ndarray[Any, Any]" = numpy.array(degree, dtype=int)
+    _degree: ndarray[Any, Any] = numpy.array(degree, dtype=int)
     npoints = len(values)
     powergrid = numpy.ones(shape=(npoints, *(_degree + 1)))
     for i, (x, deg) in enumerate(zip(points, _degree)):

src/correctionlib/highlevel.py

@@ -1,9 +1,11 @@
 """High-level correctionlib objects"""
 
+from __future__ import annotations
+
 import json
 from collections.abc import Iterator, Mapping
 from numbers import Integral
-from typing import TYPE_CHECKING, Any, Callable, Union
+from typing import TYPE_CHECKING, Any, Callable
 
 import numpy
 from packaging import version
@@ -101,7 +103,7 @@ def _call_as_numpy(
 
 def _wrap_awkward(
     func: Callable[..., Any],
-    *args: Union["awkward.Array", "numpy.ndarray[Any, Any]", str, int, float],
+    *args: awkward.Array | numpy.ndarray[Any, Any] | str | int | float,
 ) -> Any:
     from functools import partial
 
@@ -136,14 +138,14 @@ def _wrap_awkward(
 
 def _call_dask_correction(
     correction: Any,
-    *args: Union["numpy.ndarray[Any, Any]", str, int, float],
+    *args: numpy.ndarray[Any, Any] | str | int | float,
 ):
     return _wrap_awkward(correction._base.evalv, *args)
 
 
 def _wrap_dask_awkward(
     correction: Any,
-    *args: Union["numpy.ndarray[Any, Any]", str, int, float],
+    *args: numpy.ndarray[Any, Any] | str | int | float,
 ) -> Any:
     import dask.delayed
     import dask_awkward
@@ -177,14 +179,58 @@ def _wrap_dask_awkward(
     )
 
 
+def _isinstance(arg: Any, clsprefix: str) -> bool:
+    """Return True if arg is an instance of a class with the given prefix
+
+    Avoids importing modules
+    """
+    return str(type(arg)).startswith(f"<class '{clsprefix}.")
+
+
+def _evaluate(
+    corr: Correction | CompoundCorrection,
+    *args: awkward.Array | numpy.ndarray[Any, Any] | str | int | float,
+) -> float | awkward.Array | numpy.ndarray[Any, numpy.dtype[numpy.float64]]:
+    # TODO: create a ufunc with numpy.vectorize in constructor?
+    if any(_isinstance(arg, "dask.array") for arg in args):
+        raise TypeError(
+            "Correctionlib does not yet handle dask.array collections. "
+            "If you require this functionality (i.e. you cannot or do "
+            "not want to use dask_awkward/awkward arrays) please open an "
+            "issue at https://github.com/cms-nanoAOD/correctionlib/issues."
+        )
+    if any(_isinstance(arg, "dask_awkward") for arg in args):
+        return _wrap_dask_awkward(corr, *args)  # type: ignore
+    if any(_isinstance(arg, "awkward") for arg in args):
+        return _wrap_awkward(corr._base.evalv, *args)  # type: ignore
+    if all(isinstance(arg, (str, int, float)) for arg in args):
+        return corr._base.evaluate(*args)  # type: ignore
+
+    # everything else: convert to numpy and broadcast
+    vargs = [
+        numpy.asarray(arg) for arg in args if not isinstance(arg, (str, int, float))
+    ]
+    assert vargs, "should have caught all-scalar case above"
+    bargs = numpy.broadcast_arrays(*vargs)
+    oshape = bargs[0].shape
+    fargs = (arg.flatten() for arg in bargs)
+    out = corr._base.evalv(
+        *(
+            next(fargs) if not isinstance(arg, (str, int, float)) else arg
+            for arg in args
+        )
+    )
+    return out.reshape(oshape)
+
+
 class Correction:
     """High-level correction evaluator object
 
     This class is typically instantiated by accessing a named correction from
     a CorrectionSet object, rather than directly by construction.
     """
 
-    def __init__(self, base: correctionlib._core.Correction, context: "CorrectionSet"):
+    def __init__(self, base: correctionlib._core.Correction, context: CorrectionSet):
         self._base = base
         self._name = base.name
         self._context = context
@@ -218,43 +264,9 @@ def output(self) -> correctionlib._core.Variable:
         return self._base.output
 
     def evaluate(
-        self, *args: Union["numpy.ndarray[Any, Any]", str, int, float]
-    ) -> Union[float, "numpy.ndarray[Any, numpy.dtype[numpy.float64]]"]:
-        # TODO: create a ufunc with numpy.vectorize in constructor?
-        if any(str(type(arg)).startswith("<class 'dask.array.") for arg in args):
-            raise TypeError(
-                "Correctionlib does not yet handle dask.array collections. "
-                "If you require this functionality (i.e. you cannot or do "
-                "not want to use dask_awkward/awkward arrays) please open an "
-                "issue at https://github.com/cms-nanoAOD/correctionlib/issues."
-            )
-        try:
-            vargs = [
-                numpy.asarray(arg)
-                for arg in args
-                if not isinstance(arg, (str, int, float))
-            ]
-        except NotImplementedError:
-            if any(str(type(arg)).startswith("<class 'dask_awkward.") for arg in args):
-                return _wrap_dask_awkward(self, *args)  # type: ignore
-            raise
-        except (ValueError, TypeError):
-            if any(str(type(arg)).startswith("<class 'awkward.") for arg in args):
-                return _wrap_awkward(self._base.evalv, *args)  # type: ignore
-            raise
-
-        if vargs:
-            bargs = numpy.broadcast_arrays(*vargs)
-            oshape = bargs[0].shape
-            fargs = (arg.flatten() for arg in bargs)
-            out = self._base.evalv(
-                *(
-                    next(fargs) if not isinstance(arg, (str, int, float)) else arg
-                    for arg in args
-                )
-            )
-            return out.reshape(oshape)
-        return self._base.evaluate(*args)  # type: ignore
+        self, *args: awkward.Array | numpy.ndarray[Any, Any] | str | int | float
+    ) -> float | awkward.Array | numpy.ndarray[Any, numpy.dtype[numpy.float64]]:
+        return _evaluate(self, *args)
 
 
 class CompoundCorrection:
@@ -265,7 +277,7 @@ class CompoundCorrection:
     """
 
     def __init__(
-        self, base: correctionlib._core.CompoundCorrection, context: "CorrectionSet"
+        self, base: correctionlib._core.CompoundCorrection, context: CorrectionSet
     ):
         self._base = base
         self._name = base.name
@@ -296,50 +308,16 @@ def output(self) -> correctionlib._core.Variable:
         return self._base.output
 
     def evaluate(
-        self, *args: Union["numpy.ndarray[Any, Any]", str, int, float]
-    ) -> Union[float, "numpy.ndarray[Any, numpy.dtype[numpy.float64]]"]:
-        # TODO: create a ufunc with numpy.vectorize in constructor?
-        if any(str(type(arg)).startswith("<class 'dask.array.") for arg in args):
-            raise TypeError(
-                "Correctionlib does not yet handle dask.array collections. "
-                "if you require this functionality (i.e. you cannot or do "
-                "not want to use dask_awkward/awkward arrays) please open an "
-                "issue at https://github.com/cms-nanoAOD/correctionlib/issues."
-            )
-        try:
-            vargs = [
-                numpy.asarray(arg)
-                for arg in args
-                if not isinstance(arg, (str, int, float))
-            ]
-        except NotImplementedError:
-            if any(str(type(arg)).startswith("<class 'dask_awkward.") for arg in args):
-                return _wrap_dask_awkward(self, *args)  # type: ignore
-            raise
-        except (ValueError, TypeError):
-            if any(str(type(arg)).startswith("<class 'awkward.") for arg in args):
-                return _wrap_awkward(self._base.evalv, *args)  # type: ignore
-            raise
-
-        if vargs:
-            bargs = numpy.broadcast_arrays(*vargs)
-            oshape = bargs[0].shape
-            fargs = (arg.flatten() for arg in bargs)
-            out = self._base.evalv(
-                *(
-                    next(fargs) if not isinstance(arg, (str, int, float)) else arg
-                    for arg in args
-                )
-            )
-            return out.reshape(oshape)
-        return self._base.evaluate(*args)  # type: ignore
+        self, *args: awkward.Array | numpy.ndarray[Any, Any] | str | int | float
+    ) -> float | awkward.Array | numpy.ndarray[Any, numpy.dtype[numpy.float64]]:
+        return _evaluate(self, *args)
 
 
 class _CompoundMap(Mapping[str, CompoundCorrection]):
     def __init__(
         self,
         base: Mapping[str, correctionlib._core.CompoundCorrection],
-        context: "CorrectionSet",
+        context: CorrectionSet,
     ):
         self._base = base
         self._context = context
@@ -372,11 +350,11 @@ def __init__(self, data: Any):
         self._base = correctionlib._core.CorrectionSet.from_string(self._data)
 
     @classmethod
-    def from_file(cls, filename: str) -> "CorrectionSet":
+    def from_file(cls, filename: str) -> CorrectionSet:
         return cls(open_auto(filename))
 
     @classmethod
-    def from_string(cls, data: str) -> "CorrectionSet":
+    def from_string(cls, data: str) -> CorrectionSet:
         return cls(data)
 
     def __getstate__(self) -> dict[str, Any]:

tests/test_issue296.py

@@ -0,0 +1,24 @@
+import awkward as ak
+import numpy as np
+
+import correctionlib.schemav2 as cs
+
+
+def test_issue_296():
+    corr = cs.Correction(
+        name="test",
+        version=1,
+        inputs=[
+            cs.Variable(name="x", type="int"),
+        ],
+        output=cs.Variable(name="out", type="real"),
+        data=cs.HashPRNG(nodetype="hashprng", inputs=["x"], distribution="normal"),
+    ).to_evaluator()
+
+    x = ak.Array([1, 2, 3, 4, 5])
+    result = corr.evaluate(x)
+    assert isinstance(result, ak.Array)
+
+    x = np.array([1, 2, 3, 4, 5])
+    result = corr.evaluate(x)
+    assert isinstance(result, np.ndarray)