From 9821513de39446787146be641368be447837b665 Mon Sep 17 00:00:00 2001 From: cmoore24-24 <66485982+cmoore24-24@users.noreply.github.com> Date: Wed, 7 Jun 2023 08:11:35 -0500 Subject: [PATCH] feat: Adding EnergyCorrelator from Contrib (#216) * beginning work for energy correlator addition * Update _multievent.py * energy correlator with array output * fixed typo * adjusting argument positions * Adding to the ECF help message * dask bit for ECFs * removing unneeded lines * Changing default arguments, fix typo in gend2 * Update _ext.cpp Another typo * adding ECF test * fixing test issues * Changing some formatting * Removing debug additions * Updated DACS to match regular ACS * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * linting --------- Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com> Co-authored-by: Lindsey Gray --- src/_ext.cpp | 73 ++++++++++++++++++++++++++++++++ src/fastjet/__init__.py | 21 +++++++++ src/fastjet/_generalevent.py | 14 ++++++ src/fastjet/_multievent.py | 11 +++++ src/fastjet/_pyjet.py | 21 +++++++++ src/fastjet/_singleevent.py | 12 ++++++ tests/test_002-exclusive_jets.py | 63 +++++++++++++++++++++++++++ 7 files changed, 215 insertions(+) diff --git a/src/_ext.cpp b/src/_ext.cpp index 47ed6ae7..837fa477 100644 --- a/src/_ext.cpp +++ b/src/_ext.cpp @@ -13,6 +13,7 @@ #include #include #include +#include #include #include @@ -1580,6 +1581,78 @@ PYBIND11_MODULE(_ext, m) { Returns: pt, eta, phi, m of inclusive jets. )pbdoc") + .def("to_numpy_energy_correlators", + [](const output_wrapper ow, const int n_jets = 1, const double beta = 1, double npoint = 0, int angles = 0, double alpha = 0, std::string func = "generalized") { + auto css = ow.cse; + int64_t len = css.size(); + + std::transform(func.begin(), func.end(), func.begin(), + [](unsigned char c){ return std::tolower(c); }); + auto energy_correlator = std::shared_ptr>(nullptr); + if ( func == "ratio" ) { + energy_correlator = std::make_shared(npoint, beta); } + else if ( func == "doubleratio" ) { + energy_correlator = std::make_shared(npoint, beta); } + else if ( func == "c1" ) { + energy_correlator = std::make_shared(beta);} + else if ( func == "c2" ) { + energy_correlator = std::make_shared(beta);} + else if ( func == "d2" ) { + energy_correlator = std::make_shared(beta);} + else if ( func == "generalized" ) { + energy_correlator = std::make_shared(angles, npoint, beta);} + else if (func == "generalizedd2") { + energy_correlator = std::make_shared(alpha, beta);} + else if (func == "nseries") { + energy_correlator = std::make_shared(npoint, beta);} + else if (func == "n2") { + energy_correlator = std::make_shared(beta);} + else if (func == "n3") { + energy_correlator = std::make_shared(beta);} + else if (func == "mseries") { + energy_correlator = std::make_shared(npoint, beta);} + else if (func == "m2") { + energy_correlator = std::make_shared(beta);} + else if (func == "cseries") { + energy_correlator = std::make_shared(npoint, beta);} + else if (func == "useries") { + energy_correlator = std::make_shared(npoint, beta);} + else if (func == "u1") { + energy_correlator = std::make_shared(beta);} + else if (func == "u2") { + energy_correlator = std::make_shared(beta);} + else if (func == "u3") { + energy_correlator = std::make_shared(beta);} + else if (func == "generic") { + energy_correlator = std::make_shared(npoint, beta);} // The generic energy correlator is not normalized; i.e. does not use a momentum fraction when being calculated. + + std::vector ECF_vec; + + for (unsigned int i = 0; i < css.size(); i++){ // iterate through events + auto jets = css[i]->exclusive_jets(n_jets); + int size = css[i]->exclusive_jets(n_jets).size(); + + for (unsigned int j = 0; j < jets.size(); j++){ + auto ecf_result = energy_correlator->result(jets[j]); // + ECF_vec.push_back(ecf_result); + } + } + + auto ECF = py::array(ECF_vec.size(), ECF_vec.data()); + + return ECF; + }, R"pbdoc( + Calculates the energy correlators for each jet in each event. + Args: + n_jets: number of exclusive subjets. + beta: beta parameter for energy correlators. + npoint: n-point specification for ECFs. Also used to determine desired n-point function for all series classes. + angles: number of angles for generalized energy correlators. + alpha: alpha parameter for generalized D2. + func: energy correlator function to use. + Returns: + Energy correlators for each jet in each event. + )pbdoc") .def("to_numpy_exclusive_njet_lund_declusterings", [](const output_wrapper ow, const int n_jets = 0) { auto css = ow.cse; diff --git a/src/fastjet/__init__.py b/src/fastjet/__init__.py index 2fd8ce22..b75450e8 100644 --- a/src/fastjet/__init__.py +++ b/src/fastjet/__init__.py @@ -355,6 +355,27 @@ def exclusive_jets_constituents(self, njets: int = 10) -> ak.Array: raise AssertionError() + def exclusive_jets_energy_correlator( + self, + njets: int = 1, + beta: int = 1, + npoint: int = 0, + angles: int = 0, + alpha=0, + func="generalized", + ) -> ak.Array: + """Returns the energy correlator of each exclusive jet. + + Args: + njets (int): The number of jets it was clustered to. + n_point (int): The number of points in the correlator. + angle: The number of angles to be used in the correlator (if angle != n_point, ECFG is used). + beta: The beta value for the correlator. + + Returns: + awkward.highlevel.Array: Returns an Awkward Array of the same type as the input. + """ + def exclusive_jets_lund_declusterings(self, njets: int = 10) -> ak.Array: """Returns the Lund declustering Delta and k_T parameters from exclusive n_jets. diff --git a/src/fastjet/_generalevent.py b/src/fastjet/_generalevent.py index 65ee5cac..60bf671b 100644 --- a/src/fastjet/_generalevent.py +++ b/src/fastjet/_generalevent.py @@ -439,6 +439,20 @@ def exclusive_jets_constituent_index(self, njets): res = ak.Array(self._replace_multi()) return res + def exclusive_jets_energy_correlator( + self, njets=1, n_point=0, angle: int = 0, beta=1, alpha=0, func="generalized" + ): + if njets <= 0: + raise ValueError("Njets cannot be <= 0") + + self._out = [] + self._input_flag = 0 + for i in range(len(self._clusterable_level)): + np_results = self._results[i].to_numpy_energy_correlators() + self._out.append(ak.Array(ak.contents.NumpyArray(np_results[0]))) + res = ak.Array(self._replace_multi()) + return res + def exclusive_jets_lund_declusterings(self, njets): if njets <= 0: raise ValueError("Njets cannot be <= 0") diff --git a/src/fastjet/_multievent.py b/src/fastjet/_multievent.py index 36b3de83..ee809749 100644 --- a/src/fastjet/_multievent.py +++ b/src/fastjet/_multievent.py @@ -192,6 +192,17 @@ def exclusive_jets_constituent_index(self, njets): out = ak.Array(ak.contents.ListOffsetArray(ak.index.Index64(off), out.layout)) return out + def exclusive_jets_energy_correlator( + self, njets=1, beta=1, npoint=0, angles=0, alpha=0, func="generalized" + ): + if njets <= 0: + raise ValueError("Njets cannot be <= 0") + np_results = self._results.to_numpy_energy_correlators( + njets, beta, npoint, angles, alpha, func + ) + out = ak.Array(ak.contents.NumpyArray(np_results)) + return out + def exclusive_jets_lund_declusterings(self, njets): if njets <= 0: raise ValueError("Njets cannot be <= 0") diff --git a/src/fastjet/_pyjet.py b/src/fastjet/_pyjet.py index 41cfa9bc..4c6c8dd6 100644 --- a/src/fastjet/_pyjet.py +++ b/src/fastjet/_pyjet.py @@ -143,6 +143,13 @@ def exclusive_jets_constituent_index(self, njets=10): def exclusive_jets_constituents(self, njets=10): return self._internalrep.exclusive_jets_constituents(njets) + def exclusive_jets_energy_correlator( + self, njets=1, beta=1, npoint=0, angles=0, alpha=0, func="generalized" + ): + return self._internalrep.exclusive_jets_energy_correlator( + njets, beta, npoint, angles, alpha, func + ) + def exclusive_jets_lund_declusterings(self, njets=10): return self._internalrep.exclusive_jets_lund_declusterings(njets) @@ -421,6 +428,20 @@ def exclusive_jets_constituent_index(self, njets=10): def exclusive_jets_constituents(self, njets=10): return _dak_dispatch(self, "exclusive_jets_constituents", njets=njets) + def exclusive_jets_energy_correlator( + self, njets=1, beta=1, npoint=0, angles=0, alpha=0, func="generalized" + ): + return _dak_dispatch( + self, + "exclusive_jets_energy_correlator", + njets=njets, + beta=beta, + npoint=npoint, + angles=angles, + alpha=alpha, + func=func, + ) + def exclusive_jets_lund_declusterings(self, njets=10): return _dak_dispatch(self, "exclusive_jets_lund_declusterings", njets=njets) diff --git a/src/fastjet/_singleevent.py b/src/fastjet/_singleevent.py index bdb57b3b..37ffa29d 100644 --- a/src/fastjet/_singleevent.py +++ b/src/fastjet/_singleevent.py @@ -180,6 +180,18 @@ def exclusive_jets_constituent_index(self, njets): out = ak.Array(ak.contents.ListOffsetArray(ak.index.Index64(off), out.layout)) return out[0] + def exclusive_jets_energy_correlator( + self, njets=1, beta=1, npoint=0, angles=0, alpha=0, func="generalized" + ): + if njets <= 0: + raise ValueError("Njets cannot be <= 0") + + np_results = self._results.to_numpy_energy_correlators( + njets, beta, npoint, angles, alpha, func + ) + out = ak.Array(ak.contents.NumpyArray(np_results)) + return out[0] + def exclusive_jets_lund_declusterings(self, njets): if njets <= 0: raise ValueError("Njets cannot be <= 0") diff --git a/tests/test_002-exclusive_jets.py b/tests/test_002-exclusive_jets.py index f6564a04..817fc3c1 100644 --- a/tests/test_002-exclusive_jets.py +++ b/tests/test_002-exclusive_jets.py @@ -174,6 +174,69 @@ def test_exclusive_lund_declustering_multi(): assert ak.all(is_close) +def test_exclusive_energy_correlator(): + array = ak.Array( + [ + {"px": 1.2, "py": 3.2, "pz": 5.4, "E": 2.5, "ex": 0.78}, + {"px": 1.25, "py": 3.15, "pz": 5.4, "E": 2.4, "ex": 0.78}, + {"px": 1.4, "py": 3.15, "pz": 5.4, "E": 2.0, "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}, + ], + with_name="Momentum4D", + ) + + jetdef = fastjet.JetDefinition(fastjet.cambridge_algorithm, 0.8) + cluster = fastjet._pyjet.AwkwardClusterSequence(array, jetdef) + + ec1 = cluster.exclusive_jets_energy_correlator(func="generic", npoint=1) + ec2 = cluster.exclusive_jets_energy_correlator(func="generic", npoint=2) + ecg2 = cluster.exclusive_jets_energy_correlator( + func="generalized", npoint=2, angles=1 + ) + + is_close = ak.ravel( + ak.isclose(ak.Array([ec2 / ec1 / ec1]), ak.Array([ecg2]), rtol=1e-12, atol=0) + ) + + assert ak.all(is_close) + + +def test_exclusive_energy_correlator_multi(): + array = ak.Array( + [ + [ + {"px": 1.2, "py": 3.2, "pz": 5.4, "E": 2.5, "ex": 0.78}, + {"px": 1.25, "py": 3.15, "pz": 5.4, "E": 2.4, "ex": 0.78}, + {"px": 1.4, "py": 3.15, "pz": 5.4, "E": 2.0, "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}, + ], + [ + {"px": 1.2, "py": 3.2, "pz": 5.4, "E": 2.5, "ex": 0.78}, + {"px": 1.25, "py": 3.15, "pz": 5.4, "E": 2.4, "ex": 0.78}, + {"px": 1.4, "py": 3.15, "pz": 5.4, "E": 2.0, "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}, + ], + ], + with_name="Momentum4D", + ) + + jetdef = fastjet.JetDefinition(fastjet.cambridge_algorithm, 0.8) + cluster = fastjet._pyjet.AwkwardClusterSequence(array, jetdef) + + ec1 = cluster.exclusive_jets_energy_correlator(func="generic", npoint=1) + ec2 = cluster.exclusive_jets_energy_correlator(func="generic", npoint=2) + ecg2 = cluster.exclusive_jets_energy_correlator( + func="generalized", npoint=2, angles=1 + ) + + is_close = ak.ravel(ak.isclose((ec2 / ec1 / ec1), ecg2, rtol=1e-12, atol=0)) + + assert ak.all(is_close) + + def test_exclusive_constituents_multi(): array = ak.Array( [