From 907b4320067e668e9388e653f0c78e3cfa9bb853 Mon Sep 17 00:00:00 2001 From: Lindsey Gray Date: Fri, 16 Jun 2023 08:50:25 -0500 Subject: [PATCH 1/2] test dask uniformly with distributed.Client --- tests/test_008-dask.py | 133 +++++++++++++++++++++-------------------- 1 file changed, 68 insertions(+), 65 deletions(-) diff --git a/tests/test_008-dask.py b/tests/test_008-dask.py index 2f1a703f..d381e18c 100644 --- a/tests/test_008-dask.py +++ b/tests/test_008-dask.py @@ -63,32 +63,35 @@ def test_multi(): def test_single(): - array = ak.Array( - [ - {"px": 1.2, "py": 3.2, "pz": 5.4, "E": 2.5, "ex": 0.78}, - {"px": 32.2, "py": 64.21, "pz": 543.34, "E": 24.12, "ex": 0.35}, - {"px": 32.45, "py": 63.21, "pz": 543.14, "E": 24.56, "ex": 0.0}, + from distributed import Client + + with Client() as _: + array = ak.Array( + [ + {"px": 1.2, "py": 3.2, "pz": 5.4, "E": 2.5, "ex": 0.78}, + {"px": 32.2, "py": 64.21, "pz": 543.34, "E": 24.12, "ex": 0.35}, + {"px": 32.45, "py": 63.21, "pz": 543.14, "E": 24.56, "ex": 0.0}, + ] + ) + darray = dak.from_awkward(array, 1) + jetdef = fastjet.JetDefinition(fastjet.antikt_algorithm, 0.6) + cluster = fastjet._pyjet.DaskAwkwardClusterSequence(darray, jetdef) + inclusive_jets_out = [ + {"px": 1.2, "py": 3.2, "pz": 5.4, "E": 2.5}, + {"px": 64.65, "py": 127.41999999999999, "pz": 1086.48, "E": 48.68}, ] - ) - darray = dak.from_awkward(array, 1) - jetdef = fastjet.JetDefinition(fastjet.antikt_algorithm, 0.6) - cluster = fastjet._pyjet.DaskAwkwardClusterSequence(darray, jetdef) - inclusive_jets_out = [ - {"px": 1.2, "py": 3.2, "pz": 5.4, "E": 2.5}, - {"px": 64.65, "py": 127.41999999999999, "pz": 1086.48, "E": 48.68}, - ] - - assert inclusive_jets_out == cluster.inclusive_jets().compute().to_list() - constituent_output = [ - [{"px": 1.2, "py": 3.2, "pz": 5.4, "E": 2.5}], - [ - {"px": 32.2, "py": 64.21, "pz": 543.34, "E": 24.12}, - {"px": 32.45, "py": 63.21, "pz": 543.14, "E": 24.56}, - ], - ] - assert constituent_output == cluster.constituents().compute().to_list() - constituent_index_output = [[0], [1, 2]] - assert constituent_index_output == cluster.constituent_index().compute().to_list() + + assert inclusive_jets_out == cluster.inclusive_jets().compute().to_list() + constituent_output = [ + [{"px": 1.2, "py": 3.2, "pz": 5.4, "E": 2.5}], + [ + {"px": 32.2, "py": 64.21, "pz": 543.34, "E": 24.12}, + {"px": 32.45, "py": 63.21, "pz": 543.14, "E": 24.56}, + ], + ] + assert constituent_output == cluster.constituents().compute().to_list() + constituent_index_output = [[0], [1, 2]] + assert constituent_index_output == cluster.constituent_index().compute().to_list() def test_inclusive_from_file(): @@ -96,44 +99,44 @@ def test_inclusive_from_file(): import uproot from dask_awkward.lib.testutils import assert_eq + from distributed import Client + + with Client() as _: + data_dir = Path(__file__).parent / "samples" + testfile = data_dir / "pfnano_skim.root" + + devents = uproot.dask({testfile: "Events"}) + + dpfcands = dak.zip( + { + "pt": devents.PFCands_pt, + "eta": devents.PFCands_eta, + "phi": devents.PFCands_phi, + "mass": devents.PFCands_mass, + }, + with_name="Momentum4D", + behavior=vector.backends.awkward.behavior, + ) + + djetdef = fastjet.JetDefinition(fastjet.antikt_algorithm, 0.4) + dcluseq = fastjet.ClusterSequence(dpfcands, djetdef) + + events = uproot.open({testfile: "Events"}).arrays( + ["PFCands_pt", "PFCands_eta", "PFCands_phi", "PFCands_mass"] + ) + + pfcands = ak.zip( + { + "pt": events.PFCands_pt, + "eta": events.PFCands_eta, + "phi": events.PFCands_phi, + "mass": events.PFCands_mass, + }, + with_name="Momentum4D", + behavior=vector.backends.awkward.behavior, + ) + + jetdef = fastjet.JetDefinition(fastjet.antikt_algorithm, 0.4) + cluseq = fastjet.ClusterSequence(pfcands, jetdef) - vector = pytest.importorskip("vector") - - data_dir = Path(__file__).parent / "samples" - testfile = data_dir / "pfnano_skim.root" - - devents = uproot.dask({testfile: "Events"}) - - dpfcands = dak.zip( - { - "pt": devents.PFCands_pt, - "eta": devents.PFCands_eta, - "phi": devents.PFCands_phi, - "mass": devents.PFCands_mass, - }, - with_name="Momentum4D", - behavior=vector.backends.awkward.behavior, - ) - - djetdef = fastjet.JetDefinition(fastjet.antikt_algorithm, 0.4) - dcluseq = fastjet.ClusterSequence(dpfcands, djetdef) - - events = uproot.open({testfile: "Events"}).arrays( - ["PFCands_pt", "PFCands_eta", "PFCands_phi", "PFCands_mass"] - ) - - pfcands = ak.zip( - { - "pt": events.PFCands_pt, - "eta": events.PFCands_eta, - "phi": events.PFCands_phi, - "mass": events.PFCands_mass, - }, - with_name="Momentum4D", - behavior=vector.backends.awkward.behavior, - ) - - jetdef = fastjet.JetDefinition(fastjet.antikt_algorithm, 0.4) - cluseq = fastjet.ClusterSequence(pfcands, jetdef) - - assert_eq(dcluseq.inclusive_jets(), cluseq.inclusive_jets()) + assert_eq(dcluseq.inclusive_jets(), cluseq.inclusive_jets()) From a260dfa58a0e59e3c4a6a11deb748afe4ce92ca7 Mon Sep 17 00:00:00 2001 From: "pre-commit-ci[bot]" <66853113+pre-commit-ci[bot]@users.noreply.github.com> Date: Fri, 16 Jun 2023 13:50:37 +0000 Subject: [PATCH 2/2] [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci --- tests/test_008-dask.py | 6 ++++-- 1 file changed, 4 insertions(+), 2 deletions(-) diff --git a/tests/test_008-dask.py b/tests/test_008-dask.py index d381e18c..3d716321 100644 --- a/tests/test_008-dask.py +++ b/tests/test_008-dask.py @@ -91,7 +91,9 @@ def test_single(): ] assert constituent_output == cluster.constituents().compute().to_list() constituent_index_output = [[0], [1, 2]] - assert constituent_index_output == cluster.constituent_index().compute().to_list() + assert ( + constituent_index_output == cluster.constituent_index().compute().to_list() + ) def test_inclusive_from_file(): @@ -100,7 +102,7 @@ def test_inclusive_from_file(): import uproot from dask_awkward.lib.testutils import assert_eq from distributed import Client - + with Client() as _: data_dir = Path(__file__).parent / "samples" testfile = data_dir / "pfnano_skim.root"