From aeb53c4bf81926ca7e810918576ce23d3bcbbd6c Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Jonas=20R=C3=BCbenach?= Date: Thu, 3 Aug 2023 11:28:19 +0200 Subject: [PATCH 1/9] Make rochester_lookup compatible with dask_awkward --- src/coffea/lookup_tools/rochester_lookup.py | 24 +++++++++++++++------ 1 file changed, 17 insertions(+), 7 deletions(-) diff --git a/src/coffea/lookup_tools/rochester_lookup.py b/src/coffea/lookup_tools/rochester_lookup.py index 3e8bf6eab..bc462f4ba 100644 --- a/src/coffea/lookup_tools/rochester_lookup.py +++ b/src/coffea/lookup_tools/rochester_lookup.py @@ -1,4 +1,5 @@ import awkward +import dask_awkward as dak import numpy from coffea.lookup_tools.dense_lookup import dense_lookup @@ -75,7 +76,7 @@ def _error(self, func, *args): newargs = args + (0, 0) default = func(*newargs) - result = numpy.zeros_like(default) + result = awkward.zeros_like(default) for s in range(self._nsets): oneOver = 1.0 / self._members[s] for m in range(self._members[s]): @@ -226,12 +227,21 @@ def _kExtra(self, kpt, eta, nl, u, s=0, m=0): cbN_flat = awkward.flatten(cbN) cbS_flat = awkward.flatten(cbS) - invcdf = awkward.unflatten( - doublecrystalball.ppf( - u_flat, cbA_flat, cbA_flat, cbN_flat, cbN_flat, loc, cbS_flat - ), - counts, - ) + args = (u_flat, cbA_flat, cbA_flat, cbN_flat, cbN_flat, loc, cbS_flat) + + if any(isinstance(arg, dak.Array) for arg in args): + def apply(*args): + args_lz = [awkward.typetracer.length_zero_if_typetracer(arg) for arg in args] + out = awkward.Array(doublecrystalball.ppf(*args_lz)) + if awkward.backend(args[0]) == "typetracer": + out = awkward.Array(out.layout.to_typetracer(forget_length=True), behavior=out.behavior) + return out + + invcdf = dak.map_partitions(apply, *args) + else: + invcdf = doublecrystalball.ppf(*args) + + invcdf = awkward.unflatten(invcdf, counts) x = awkward.where( mask, From e6acf51e3d362947c04e9fc6d2e35201e31f638a Mon Sep 17 00:00:00 2001 From: "pre-commit-ci[bot]" <66853113+pre-commit-ci[bot]@users.noreply.github.com> Date: Thu, 3 Aug 2023 09:35:01 +0000 Subject: [PATCH 2/9] [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci --- src/coffea/lookup_tools/rochester_lookup.py | 10 ++++++++-- 1 file changed, 8 insertions(+), 2 deletions(-) diff --git a/src/coffea/lookup_tools/rochester_lookup.py b/src/coffea/lookup_tools/rochester_lookup.py index bc462f4ba..fece55d6d 100644 --- a/src/coffea/lookup_tools/rochester_lookup.py +++ b/src/coffea/lookup_tools/rochester_lookup.py @@ -230,11 +230,17 @@ def _kExtra(self, kpt, eta, nl, u, s=0, m=0): args = (u_flat, cbA_flat, cbA_flat, cbN_flat, cbN_flat, loc, cbS_flat) if any(isinstance(arg, dak.Array) for arg in args): + def apply(*args): - args_lz = [awkward.typetracer.length_zero_if_typetracer(arg) for arg in args] + args_lz = [ + awkward.typetracer.length_zero_if_typetracer(arg) for arg in args + ] out = awkward.Array(doublecrystalball.ppf(*args_lz)) if awkward.backend(args[0]) == "typetracer": - out = awkward.Array(out.layout.to_typetracer(forget_length=True), behavior=out.behavior) + out = awkward.Array( + out.layout.to_typetracer(forget_length=True), + behavior=out.behavior, + ) return out invcdf = dak.map_partitions(apply, *args) From 5feaff63a5fd03c913a60caa3e01b2195db6b8d9 Mon Sep 17 00:00:00 2001 From: Lindsey Gray Date: Thu, 3 Aug 2023 10:08:53 -0500 Subject: [PATCH 3/9] reactivate rochester corrections tests --- tests/test_lookup_tools.py | 2 -- 1 file changed, 2 deletions(-) diff --git a/tests/test_lookup_tools.py b/tests/test_lookup_tools.py index f4d5a3b30..4b86ec3df 100644 --- a/tests/test_lookup_tools.py +++ b/tests/test_lookup_tools.py @@ -372,8 +372,6 @@ def test_jec_txt_effareas(): def test_rochester(): - pytest.xfail("weird side effect from running other tests... passes by itself") - rochester_data = lookup_tools.txt_converters.convert_rochester_file( "tests/samples/RoccoR2018.txt.gz", loaduncs=True ) From fbfbb0d0357f6d143eae821a2da5b005bdff60c8 Mon Sep 17 00:00:00 2001 From: Lindsey Gray Date: Thu, 3 Aug 2023 10:45:40 -0500 Subject: [PATCH 4/9] hoisted by my own petard --- tests/test_lookup_tools.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/tests/test_lookup_tools.py b/tests/test_lookup_tools.py index 4b86ec3df..a57c6b7cf 100644 --- a/tests/test_lookup_tools.py +++ b/tests/test_lookup_tools.py @@ -388,7 +388,7 @@ def test_rochester(): # test against nanoaod events = NanoEventsFactory.from_root( - os.path.abspath("tests/samples/nano_dimuon.root") + {os.path.abspath("tests/samples/nano_dimuon.root"): "Events"} ).events() data_k = rochester.kScaleDT( From 78d2ed2f0e378772ff212e8a0d0473190ef18142 Mon Sep 17 00:00:00 2001 From: Lindsey Gray Date: Thu, 3 Aug 2023 15:36:23 -0500 Subject: [PATCH 5/9] this time, with feeling --- tests/test_lookup_tools.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/tests/test_lookup_tools.py b/tests/test_lookup_tools.py index a57c6b7cf..7bebd7e18 100644 --- a/tests/test_lookup_tools.py +++ b/tests/test_lookup_tools.py @@ -404,7 +404,7 @@ def test_rochester(): # test against mc events = NanoEventsFactory.from_root( - os.path.abspath("tests/samples/nano_dy.root") + {os.path.abspath("tests/samples/nano_dy.root"): "Events"}, ).events() hasgen = ~np.isnan(ak.fill_none(events.Muon.matched_gen.pt, np.nan)) From 129e2691f14c5b49787186b401726fb18c7c877b Mon Sep 17 00:00:00 2001 From: Lindsey Gray Date: Fri, 4 Aug 2023 13:06:24 -0500 Subject: [PATCH 6/9] daskify tests --- tests/test_lookup_tools.py | 22 ++++++++++++---------- 1 file changed, 12 insertions(+), 10 deletions(-) diff --git a/tests/test_lookup_tools.py b/tests/test_lookup_tools.py index 7bebd7e18..9a6a7954d 100644 --- a/tests/test_lookup_tools.py +++ b/tests/test_lookup_tools.py @@ -388,23 +388,25 @@ def test_rochester(): # test against nanoaod events = NanoEventsFactory.from_root( - {os.path.abspath("tests/samples/nano_dimuon.root"): "Events"} + {os.path.abspath("tests/samples/nano_dimuon.root"): "Events"}, + permit_dask=True, ).events() data_k = rochester.kScaleDT( events.Muon.charge, events.Muon.pt, events.Muon.eta, events.Muon.phi ) - data_k = np.array(ak.flatten(data_k)) + data_k = ak.flatten(data_k).compute().to_numpy() assert all(np.isclose(data_k, official_data_k)) data_err = rochester.kScaleDTerror( events.Muon.charge, events.Muon.pt, events.Muon.eta, events.Muon.phi ) - data_err = np.array(ak.flatten(data_err), dtype=float) + data_err = ak.flatten(data_err).compute().to_numpy() assert all(np.isclose(data_err, official_data_err, atol=1e-8)) # test against mc events = NanoEventsFactory.from_root( {os.path.abspath("tests/samples/nano_dy.root"): "Events"}, + permit_dask=True, ).events() hasgen = ~np.isnan(ak.fill_none(events.Muon.matched_gen.pt, np.nan)) @@ -424,10 +426,10 @@ def test_rochester(): events.Muon.nTrackerLayers[~hasgen], mc_rand[~hasgen], ) - mc_k = np.array(ak.flatten(ak.ones_like(events.Muon.pt))) - hasgen_flat = np.array(ak.flatten(hasgen)) - mc_k[hasgen_flat] = np.array(ak.flatten(mc_kspread)) - mc_k[~hasgen_flat] = np.array(ak.flatten(mc_ksmear)) + mc_k = ak.flatten(ak.ones_like(events.Muon.pt)).compute().to_numpy() + hasgen_flat = ak.flatten(hasgen).compute().to_numpy() + mc_k[hasgen_flat] = ak.flatten(mc_kspread).compute().to_numpy() + mc_k[~hasgen_flat] = ak.flatten(mc_ksmear).compute().to_numpy() assert all(np.isclose(mc_k, official_mc_k)) mc_errspread = rochester.kSpreadMCerror( @@ -445,9 +447,9 @@ def test_rochester(): events.Muon.nTrackerLayers[~hasgen], mc_rand[~hasgen], ) - mc_err = np.array(ak.flatten(ak.ones_like(events.Muon.pt))) - mc_err[hasgen_flat] = np.array(ak.flatten(mc_errspread)) - mc_err[~hasgen_flat] = np.array(ak.flatten(mc_errsmear)) + mc_err = ak.flatten(ak.ones_like(events.Muon.pt)).compute().to_numpy() + mc_err[hasgen_flat] = ak.flatten(mc_errspread).compute().to_numpy() + mc_err[~hasgen_flat] = ak.flatten(mc_errsmear).compute().to_numpy() assert all(np.isclose(mc_err, official_mc_err, atol=1e-8)) From dcbc45b97eddf5c22fc3d7e2d7bae195f27c9143 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Jonas=20R=C3=BCbenach?= Date: Tue, 8 Aug 2023 18:26:57 +0200 Subject: [PATCH 7/9] Fix daskification of test_rochester --- tests/test_lookup_tools.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/tests/test_lookup_tools.py b/tests/test_lookup_tools.py index 9a6a7954d..7fde363de 100644 --- a/tests/test_lookup_tools.py +++ b/tests/test_lookup_tools.py @@ -410,7 +410,7 @@ def test_rochester(): ).events() hasgen = ~np.isnan(ak.fill_none(events.Muon.matched_gen.pt, np.nan)) - mc_rand = ak.unflatten(mc_rand, ak.num(hasgen)) + mc_rand = ak.unflatten(dak.from_awkward(ak.Array(mc_rand), 1), ak.num(hasgen)) mc_kspread = rochester.kSpreadMC( events.Muon.charge[hasgen], events.Muon.pt[hasgen], From 88176f22a853032c519c2b2ef0cb4460bd9549a1 Mon Sep 17 00:00:00 2001 From: Lindsey Gray Date: Wed, 6 Sep 2023 10:09:24 -0500 Subject: [PATCH 8/9] repin to awkward 2.4.1 --- pyproject.toml | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/pyproject.toml b/pyproject.toml index 454ed3319..87b29d513 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -37,7 +37,7 @@ classifiers = [ "Topic :: Utilities", ] dependencies = [ - "awkward>=2.3.3", + "awkward>=2.4.1", "uproot>=5.0.10", "dask[array]>=2023.4.0", "dask-awkward>=2023.7.1,!=2023.8.0", From 5dd6868d98534f960e04c973cba5d06a751534b7 Mon Sep 17 00:00:00 2001 From: Lindsey Gray Date: Wed, 6 Sep 2023 15:00:36 -0500 Subject: [PATCH 9/9] awkward 2.4.2 (just to not clobber main) --- pyproject.toml | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/pyproject.toml b/pyproject.toml index 87b29d513..f5b168c1e 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -37,7 +37,7 @@ classifiers = [ "Topic :: Utilities", ] dependencies = [ - "awkward>=2.4.1", + "awkward>=2.4.2", "uproot>=5.0.10", "dask[array]>=2023.4.0", "dask-awkward>=2023.7.1,!=2023.8.0",