Is there a dask akward equivalent to np.interp()?

[green-cabbage]

Hello, I was wondering if there's a way to do linear interpolation using dask_awkward, similar or identical to np.interp(), per se.
I am trying to apply NNLOPS SF weight to my ggH signal model, and the given weight data is in the form of ROOT.TGraph, not the usual ROOT.TH. The way we generate the SF is to linearly interpolate the y value from two closest Higgs.pt points (one left, one right) in the TGraph, and this interpolated y value becomes the SF for that ggH signal event.

I currently have this NNLOPS weight implemented eagerly and saved separately to other lazy weights, but it is leading to extra complications later on (ie normalization of btag shape weights, which requires information from all weights), and it would be just more elegant if I could lazily calculate the NNLOPS weights as well.

If there are better way to lazily implement the NNLOPS, I am all ears! Thank you in advance!

[lgray]

For this one in particular you should probably wrap np.interp in dak.map_partitions (in fact that would be a good PR for this repository).

[green-cabbage]

I tried using dak.map_partitions, and it's giving me a TypeError("cannot compare unknown lengths against known values") statement. My quick eager code works perfectly fine:

xp = np.array([1, 2, 3])
fp = np.array([3, 2, 0])
np.interp(jets.pt.compute(), xp, fp)

this returns [[0], [0]] (jets.pt has just two rows). So then I tried implementing dak.map_partitions: dak_test =dak.map_partitions(np.interp, jets.pt, xp, fp), and it gives me this long error statement

---------------------------------------------------------------------------
TypeError                                 Traceback (most recent call last)
Cell In[91], line 2
      1 # %debug
----> 2 dak_test =dak.map_partitions(np.interp, jets.pt, xp, fp)

File [/depot/cms/kernels/coffea2024/lib/python3.10/site-packages/dask_awkward/lib/core.py:2107](https://cms.geddes.rcac.purdue.edu/depot/cms/kernels/coffea2024/lib/python3.10/site-packages/dask_awkward/lib/core.py#line=2106), in map_partitions(base_fn, label, token, meta, output_divisions, traverse, *args, **kwargs)
   2099         arg_lens_for_repackers.append(1)
   2101 fn = ArgsKwargsPackedFunction(
   2102     base_fn,
   2103     arg_repackers,
   2104     kwarg_repacker,
   2105     arg_lens_for_repackers,
   2106 )
-> 2107 return _map_partitions(
   2108     fn,
   2109     *arg_flat_deps_expanded,
   2110     *kwarg_flat_deps,
   2111     label=label,
   2112     token=token,
   2113     meta=meta,
   2114     output_divisions=output_divisions,
   2115 )

File [/depot/cms/kernels/coffea2024/lib/python3.10/site-packages/dask_awkward/lib/core.py:1943](https://cms.geddes.rcac.purdue.edu/depot/cms/kernels/coffea2024/lib/python3.10/site-packages/dask_awkward/lib/core.py#line=1942), in _map_partitions(fn, label, token, meta, output_divisions, *args, **kwargs)
   1935 lay = partitionwise_layer(
   1936     fn,
   1937     name,
   1938     *args,
   1939     **kwargs,
   1940 )
   1942 if meta is None:
-> 1943     meta = map_meta(fn, *args, **kwargs)
   1945 hlg = HighLevelGraph.from_collections(
   1946     name,
   1947     lay,
   1948     dependencies=deps,
   1949 )
   1951 if len(dak_arrays) == 0:

File [/depot/cms/kernels/coffea2024/lib/python3.10/site-packages/dask_awkward/lib/core.py:2467](https://cms.geddes.rcac.purdue.edu/depot/cms/kernels/coffea2024/lib/python3.10/site-packages/dask_awkward/lib/core.py#line=2466), in map_meta(fn, *deps)
   2463 def map_meta(fn: Callable | ArgsKwargsPackedFunction, *deps: Any) -> ak.Array | None:
   2464     # NOTE: fn is assumed to be a *packed* function
   2465     #       as defined up in map_partitions. be careful!
   2466     try:
-> 2467         meta = fn(*to_meta(deps))
   2468         return meta
   2469     except Exception as err:
   2470         # if compute-unknown-meta is False then we don't care about
   2471         # this failure and we return None.

File [/depot/cms/kernels/coffea2024/lib/python3.10/site-packages/dask_awkward/lib/core.py:1903](https://cms.geddes.rcac.purdue.edu/depot/cms/kernels/coffea2024/lib/python3.10/site-packages/dask_awkward/lib/core.py#line=1902), in ArgsKwargsPackedFunction.__call__(self, *args_deps_expanded)
   1901     len_args += n_args
   1902 kwargs = self.kwarg_repacker(args_deps_expanded[len_args:])[0]
-> 1903 return self.fn(*args, **kwargs)

File [/depot/cms/kernels/coffea2024/lib/python3.10/site-packages/awkward/highlevel.py:1527](https://cms.geddes.rcac.purdue.edu/depot/cms/kernels/coffea2024/lib/python3.10/site-packages/awkward/highlevel.py#line=1526), in Array.__array_function__(self, func, types, args, kwargs)
   1513 def __array_function__(self, func, types, args, kwargs):
   1514     """
   1515     Intercepts attempts to pass this Array to those NumPy functions other
   1516     than universal functions that have an Awkward equivalent.
   (...)
   1525     See also #__array_ufunc__.
   1526     """
-> 1527     return ak._connect.numpy.array_function(
   1528         func, types, args, kwargs, behavior=self._behavior, attrs=self._attrs
   1529     )

File [/depot/cms/kernels/coffea2024/lib/python3.10/site-packages/awkward/_connect/numpy.py:109](https://cms.geddes.rcac.purdue.edu/depot/cms/kernels/coffea2024/lib/python3.10/site-packages/awkward/_connect/numpy.py#line=108), in array_function(func, types, args, kwargs, behavior, attrs)
    106 unique_backends = frozenset(_find_backends(all_arguments))
    107 backend = common_backend(unique_backends)
--> 109 rectilinear_args = tuple(_to_rectilinear(x, backend) for x in args)
    110 rectilinear_kwargs = {k: _to_rectilinear(v, backend) for k, v in kwargs.items()}
    111 result = func(*rectilinear_args, **rectilinear_kwargs)

File [/depot/cms/kernels/coffea2024/lib/python3.10/site-packages/awkward/_connect/numpy.py:109](https://cms.geddes.rcac.purdue.edu/depot/cms/kernels/coffea2024/lib/python3.10/site-packages/awkward/_connect/numpy.py#line=108), in <genexpr>(.0)
    106 unique_backends = frozenset(_find_backends(all_arguments))
    107 backend = common_backend(unique_backends)
--> 109 rectilinear_args = tuple(_to_rectilinear(x, backend) for x in args)
    110 rectilinear_kwargs = {k: _to_rectilinear(v, backend) for k, v in kwargs.items()}
    111 result = func(*rectilinear_args, **rectilinear_kwargs)

File [/depot/cms/kernels/coffea2024/lib/python3.10/site-packages/awkward/_connect/numpy.py:78](https://cms.geddes.rcac.purdue.edu/depot/cms/kernels/coffea2024/lib/python3.10/site-packages/awkward/_connect/numpy.py#line=77), in _to_rectilinear(arg, backend)
     69     # Otherwise, cast to layout and convert
     70     else:
     71         layout = ak.to_layout(
     72             arg,
     73             allow_record=False,
   (...)
     76             string_policy="error",
     77         )
---> 78         return layout.to_backend(backend).to_backend_array(allow_missing=True)
     79 elif isinstance(arg, tuple):
     80     return tuple(_to_rectilinear(x, backend) for x in arg)

File [/depot/cms/kernels/coffea2024/lib/python3.10/site-packages/awkward/contents/content.py:1021](https://cms.geddes.rcac.purdue.edu/depot/cms/kernels/coffea2024/lib/python3.10/site-packages/awkward/contents/content.py#line=1020), in Content.to_backend_array(self, allow_missing, backend)
   1019 else:
   1020     backend = regularize_backend(backend)
-> 1021 return self._to_backend_array(allow_missing, backend)

File [/depot/cms/kernels/coffea2024/lib/python3.10/site-packages/awkward/contents/listoffsetarray.py:2078](https://cms.geddes.rcac.purdue.edu/depot/cms/kernels/coffea2024/lib/python3.10/site-packages/awkward/contents/listoffsetarray.py#line=2077), in ListOffsetArray._to_backend_array(self, allow_missing, backend)
   2076     return buffer.view(np.dtype(("S", max_count)))
   2077 else:
-> 2078     return self.to_RegularArray()._to_backend_array(allow_missing, backend)

File [/depot/cms/kernels/coffea2024/lib/python3.10/site-packages/awkward/contents/regulararray.py:1300](https://cms.geddes.rcac.purdue.edu/depot/cms/kernels/coffea2024/lib/python3.10/site-packages/awkward/contents/regulararray.py#line=1299), in RegularArray._to_backend_array(self, allow_missing, backend)
   1294 # ShapeItem is a defined type, but some nplikes don't map onto the entire space; e.g.
   1295 # NumPy never has `None` shape items. We require that if a shape-item is used between nplikes
   1296 # they both be the same "known-shape-ness".
   1297 assert (
   1298     self._backend.index_nplike.known_data == self._backend.nplike.known_data
   1299 )
-> 1300 return self._backend.nplike.reshape(
   1301     out[
   1302         : self._backend.nplike.shape_item_as_index(
   1303             self._length * self._size
   1304         )
   1305     ],
   1306     shape,
   1307 )

File [/depot/cms/kernels/coffea2024/lib/python3.10/site-packages/awkward/_nplikes/typetracer.py:1169](https://cms.geddes.rcac.purdue.edu/depot/cms/kernels/coffea2024/lib/python3.10/site-packages/awkward/_nplikes/typetracer.py#line=1168), in TypeTracer.reshape(self, x, shape, copy)
   1167 new_shape = [*shape]
   1168 for i, item in enumerate(shape):
-> 1169     if item == -1:
   1170         new_shape[i] = size // new_size
   1171         break

File /depot/cms/kernels/coffea2024/lib/python3.10/site-packages/awkward/_nplikes/shape.py:70, in UnknownLength.__eq__(self, other)
     68     return True
     69 else:
---> 70     raise TypeError("cannot compare unknown lengths against known values")

TypeError: cannot compare unknown lengths against known values

I know this forum isn't here to debug my problems, but since I am not familiar with how dask awkward wraps awkward and numpy functions, I would really appreciate just some starting hints to get me started. Thank you!

[lgray]

It means you probably need to wrap it in some code to handle the typetracers. Put in a feature request on dask-awkward and we can make it happen.

[green-cabbage]

I have migrated this discussion to this issue, and @lgray kindly laid out a custom solution to my problem and it works, though I hope the feature would be added to the official package!