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 <lindsey.gray@gmail.com>

src/_ext.cpp

@@ -13,6 +13,7 @@
 #include <fastjet/GhostedAreaSpec.hh>
 #include <fastjet/JetDefinition.hh>
 #include <fastjet/PseudoJet.hh>
+#include <fastjet/contrib/EnergyCorrelator.hh>
 #include <fastjet/contrib/LundGenerator.hh>
 
 #include <pybind11/numpy.h>
@@ -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<fastjet::FunctionOfPseudoJet<double>>(nullptr);
+        if ( func == "ratio" ) {
+          energy_correlator = std::make_shared<fastjet::contrib::EnergyCorrelatorRatio>(npoint, beta); }
+        else if ( func == "doubleratio" ) {
+          energy_correlator = std::make_shared<fastjet::contrib::EnergyCorrelatorDoubleRatio>(npoint, beta); }
+        else if ( func == "c1" ) {
+          energy_correlator = std::make_shared<fastjet::contrib::EnergyCorrelatorC1>(beta);}
+        else if ( func == "c2" ) {
+          energy_correlator = std::make_shared<fastjet::contrib::EnergyCorrelatorC2>(beta);}
+        else if ( func == "d2" ) {
+          energy_correlator = std::make_shared<fastjet::contrib::EnergyCorrelatorD2>(beta);}
+        else if ( func == "generalized" ) {
+          energy_correlator = std::make_shared<fastjet::contrib::EnergyCorrelatorGeneralized>(angles, npoint, beta);}
+        else if (func == "generalizedd2") {
+          energy_correlator = std::make_shared<fastjet::contrib::EnergyCorrelatorGeneralizedD2>(alpha, beta);}
+        else if (func == "nseries") {
+          energy_correlator = std::make_shared<fastjet::contrib::EnergyCorrelatorNseries>(npoint, beta);}
+        else if (func == "n2") {
+          energy_correlator = std::make_shared<fastjet::contrib::EnergyCorrelatorN2>(beta);}
+        else if (func == "n3") {
+          energy_correlator = std::make_shared<fastjet::contrib::EnergyCorrelatorN3>(beta);}
+        else if (func == "mseries") {
+          energy_correlator = std::make_shared<fastjet::contrib::EnergyCorrelatorMseries>(npoint, beta);}
+        else if (func == "m2") {
+          energy_correlator = std::make_shared<fastjet::contrib::EnergyCorrelatorM2>(beta);}
+        else if (func == "cseries") {
+          energy_correlator = std::make_shared<fastjet::contrib::EnergyCorrelatorCseries>(npoint, beta);}
+        else if (func == "useries") {
+          energy_correlator = std::make_shared<fastjet::contrib::EnergyCorrelatorUseries>(npoint, beta);}
+        else if (func == "u1") {
+          energy_correlator = std::make_shared<fastjet::contrib::EnergyCorrelatorU1>(beta);}
+        else if (func == "u2") {
+          energy_correlator = std::make_shared<fastjet::contrib::EnergyCorrelatorU2>(beta);}
+        else if (func == "u3") {
+          energy_correlator = std::make_shared<fastjet::contrib::EnergyCorrelatorU3>(beta);}
+        else if (func == "generic") {
+          energy_correlator = std::make_shared<fastjet::contrib::EnergyCorrelator>(npoint, beta);} // The generic energy correlator is not normalized; i.e. does not use a momentum fraction when being calculated.
+
+        std::vector<double> 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;

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.
 

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")

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")

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)
 

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")

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(
         [