Rochester corrections not compatible with dask_awkward

[jrueb]

Describe the bug
The Rochester corrections of given by coffea.lookup_tools.rochester_lookup.rochester_lookup are not compatible with dask_awkward. In particular, trying to execute kSmearMC on dask_awkward arrays will give a NotImplementedError. The reason seems to be that it uses scipy.stats functionality (through doublecrystalball).

To Reproduce

import awkward as ak
import dask_awkward as dak
import coffea.lookup_tools


rochester_data = coffea.lookup_tools.txt_converters.convert_rochester_file(
    "tests/samples/RoccoR2018.txt.gz", loaduncs=True
)
rochester = coffea.lookup_tools.rochester_lookup.rochester_lookup(rochester_data)

x = ak.Array([[1.0]])
x = dak.from_awkward(x, 1)
mc_ksmear = rochester.kSmearMC(x, x, x, x, x, x)

Where tests/samples/RoccoR2018.txt.gz can be found in the coffea repository.

Desktop (please complete the following information):

• Coffea Version 2023.7.0rc0

[lgray]

@ssrothman is working on a correctionlib version of the rochester corrections so maybe this ends up fixing the issue in general.

For the current version in coffea I intend to get it working with dask awkward. I can likely knock it out this month but there are other things that need attention first. I suspect that the fix is not that bad, probably just wrapping some stuff in map_partitions.

As a library developer it may be a useful exercise for you to implement that part ( and submit it as a PR to coffea ;-) ).