Lund de-clusterings for exclusive jets

src/_ext.cpp

@@ -1579,6 +1579,83 @@ PYBIND11_MODULE(_ext, m) {
         Returns:
           pt, eta, phi, m of inclusive jets.
       )pbdoc")
+      .def("to_numpy_exclusive_njet_lund_declusterings",
+      [](const output_wrapper ow, const int n_jets = 0) {
+        auto css = ow.cse;
+        int64_t len = css.size();
+        auto jk = 0;
+
+        for(int i = 0; i < len; i++){
+          jk += css[i]->exclusive_jets(n_jets).size();
+        }
+        jk++;
+
+        std::vector<double> Delta_vec;
+        std::vector<double> kt_vec;
+
+        auto eventoffsets = py::array(py::buffer_info(nullptr, sizeof(int), py::format_descriptor<int>::value, 1, {len+1}, {sizeof(int)}));
+        auto bufeventoffsets = eventoffsets.request();
+        int *ptreventoffsets = (int *)bufeventoffsets.ptr;
+        size_t eventidx = 0;
+
+        ptreventoffsets[eventidx] = 0;
+        eventidx++;
+
+        auto jetoffsets = py::array(py::buffer_info(nullptr, sizeof(int), py::format_descriptor<int>::value, 1, {jk}, {sizeof(int)}));
+        auto bufjetoffsets = jetoffsets.request();
+        int *ptrjetoffsets = (int *)bufjetoffsets.ptr;
+        size_t jetidx = 0;
+
+        size_t idxh = 0;
+        ptrjetoffsets[jetidx] = 0;
+        jetidx++;
+        auto eventprev = 0;
+
+        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();
+          auto prev = ptrjetoffsets[jetidx-1];
+
+          for (unsigned int j = 0; j < jets.size(); j++){
+            // adapted from https://github.com/fdreyer/LundPlane/blob/master/LundGenerator.cc
+            PseudoJet pair, j1, j2;
+            pair = jets[j];
+            int splittings = 0;
+            while (pair.has_parents(j1, j2)) {
+              if (j1.pt2() < j2.pt2()) std::swap(j1,j2);
+              double Delta = j1.delta_R(j2);
+              Delta_vec.push_back(Delta);
+              kt_vec.push_back(j2.pt() * Delta);
+              pair = j1;
+              splittings++;
+            }
+
+            ptrjetoffsets[jetidx] = splittings + prev;
+            prev = ptrjetoffsets[jetidx];
+            jetidx++;
+          }
+
+          ptreventoffsets[eventidx] = jets.size() + eventprev;
+          eventprev = ptreventoffsets[eventidx];
+          eventidx++;
+        }
+
+        auto Deltas = py::array(Delta_vec.size(), Delta_vec.data());
+        auto kts = py::array(kt_vec.size(), kt_vec.data());
+
+        return std::make_tuple(
+            jetoffsets,
+            Deltas,
+            kts,
+            eventoffsets
+          );
+      }, "n_jets"_a = 0, R"pbdoc(
+        Calculates the Lund declustering Delta and k_T parameters from exclusive n_jets and converts them to numpy arrays.
+        Args:
+          n_jets: Number of exclusive subjets. Default: 0.
+        Returns:
+          jet offsets, splitting Deltas, kts, and event offsets.
+      )pbdoc")
       .def("to_numpy_unclustered_particles",
       [](const output_wrapper ow) {
         auto css = ow.cse;

src/fastjet/__init__.py

@@ -321,6 +321,18 @@ def exclusive_jets_constituents(self, njets: int = 10) -> ak.Array:
 
         raise AssertionError()
 
+    def exclusive_jets_lund_declusterings(self, njets: int = 10) -> ak.Array:
+        """Returns the Lund declustering Delta and k_T parameters from exclusive n_jets.
+
+        Args:
+            njets (int): The number of jets it was clustered to.
+
+        Returns:
+            awkward.highlevel.Array: Returns an Awkward Array of the same type as the input.
+        """
+
+        raise AssertionError()
+
     def exclusive_dmerge(self, njets: int = 10) -> Union[ak.Array, float]:
         """Returns the dmin corresponding to the recombination that went from n+1 to n jets.
 

src/fastjet/_generalevent.py

@@ -703,6 +703,38 @@ def exclusive_jets_constituent_index(self, njets):
         res = ak.Array(self._replace_multi())
         return res
 
+    def exclusive_jets_lund_declusterings(self, njets):
+        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_exclusive_njet_with_constituents(
+                njets
+            )
+            off = np_results[-1]
+            out = ak.Array(
+                ak.layout.ListOffsetArray64(
+                    ak.layout.Index64(np_results[0]),
+                    ak.layout.RecordArray(
+                        (
+                            ak.layout.NumpyArray(np_results[1]),
+                            ak.layout.NumpyArray(np_results[2]),
+                        ),
+                        ("Delta", "kt"),
+                    ),
+                ),
+                behavior=self.data.behavior,
+            )
+            self._out.append(
+                ak.Array(
+                    ak.layout.ListOffsetArray64(ak.layout.Index64(off), out.layout)
+                )
+            )
+        res = ak.Array(self._replace_multi())
+        return res
+
     def unclustered_particles(self):
         self._out = []
         self._input_flag = 0

src/fastjet/_multievent.py

@@ -198,6 +198,27 @@ def exclusive_jets_constituent_index(self, njets):
         out = ak.Array(ak.layout.ListOffsetArray64(ak.layout.Index64(off), out.layout))
         return out
 
+    def exclusive_jets_lund_declusterings(self, njets):
+        if njets <= 0:
+            raise ValueError("Njets cannot be <= 0")
+
+        np_results = self._results.to_numpy_exclusive_njet_lund_declusterings(njets)
+        off = np_results[-1]
+        out = ak.Array(
+            ak.layout.ListOffsetArray64(
+                ak.layout.Index64(np_results[0]),
+                ak.layout.RecordArray(
+                    (
+                        ak.layout.NumpyArray(np_results[1]),
+                        ak.layout.NumpyArray(np_results[2]),
+                    ),
+                    ("Delta", "kt"),
+                ),
+            )
+        )
+        out = ak.Array(ak.layout.ListOffsetArray64(ak.layout.Index64(off), out.layout))
+        return out
+
     def unique_history_order(self):
         np_results = self._results.to_numpy_unique_history_order()
         off = np_results[-1]

src/fastjet/_pyjet.py

@@ -212,6 +212,9 @@ 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_lund_declusterings(self, njets=10):
+        return self._internalrep.exclusive_jets_lund_declusterings(njets)
+
     def exclusive_dmerge(self, njets=10):
         return self._internalrep.exclusive_dmerge(njets)
 

src/fastjet/_singleevent.py

@@ -186,6 +186,27 @@ def exclusive_jets_constituent_index(self, njets):
         out = ak.Array(ak.layout.ListOffsetArray64(ak.layout.Index64(off), out.layout))
         return out[0]
 
+    def exclusive_jets_lund_declusterings(self, njets):
+        if njets <= 0:
+            raise ValueError("Njets cannot be <= 0")
+
+        np_results = self._results.to_numpy_exclusive_njet_lund_declusterings(njets)
+        off = np_results[-1]
+        out = ak.Array(
+            ak.layout.ListOffsetArray64(
+                ak.layout.Index64(np_results[0]),
+                ak.layout.RecordArray(
+                    (
+                        ak.layout.NumpyArray(np_results[1]),
+                        ak.layout.NumpyArray(np_results[2]),
+                    ),
+                    ("Delta", "kt"),
+                ),
+            )
+        )
+        out = ak.Array(ak.layout.ListOffsetArray64(ak.layout.Index64(off), out.layout))
+        return out[0]
+
     def unique_history_order(self):
         np_results = self._results.to_numpy_unique_history_order()
         out = ak.Array(ak.layout.NumpyArray(np_results[0]))

tests/test_002-exclusive_jets.py

@@ -90,6 +90,80 @@ def test_exclusive_constituents_single():
     )
 
 
+def test_exclusive_lund_declustering_single():
+    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)
+
+    lds = cluster.exclusive_jets_lund_declusterings(2)
+
+    lund_declustering_output = [
+        [
+            {"Delta": 0.08755181299980186, "kt": 0.30179987478357223},
+            {"Delta": 0.019481226884377707, "kt": 0.06602095529127928},
+        ],
+        [{"Delta": 0.014750342295225208, "kt": 1.0480537658466145}],
+    ]
+
+    assert lund_declustering_output == lds.to_list()
+
+
+def test_exclusive_lund_declustering_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)
+
+    lds = cluster.exclusive_jets_lund_declusterings(2)
+
+    lund_declustering_output = [
+        [
+            [
+                {"Delta": 0.08755181299980186, "kt": 0.30179987478357223},
+                {"Delta": 0.019481226884377707, "kt": 0.06602095529127928},
+            ],
+            [{"Delta": 0.014750342295225208, "kt": 1.0480537658466145}],
+        ],
+        [
+            [
+                {"Delta": 0.08755181299980186, "kt": 0.30179987478357223},
+                {"Delta": 0.019481226884377707, "kt": 0.06602095529127928},
+            ],
+            [{"Delta": 0.014750342295225208, "kt": 1.0480537658466145}],
+        ],
+    ]
+
+    assert lund_declustering_output == lds.to_list()
+
+
 def test_exclusive_constituents_multi():
     array = ak.Array(
         [