fix: fix #238 – wrong handling of offsets (#274)

* add test for masking in exclusive_constituents

* First attempt to fix the issue. Implemented starts & stops.

* should use the pointer

* start/stop are smaller than offsets array

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

* add back missing clustering invocation

* probably not needed, the code here could be substantially sped up.

* remove insertion of leading zeros in starts/stops

* remove insertion of leading zeros in starts/stops

* remove insertion of leading zeros in starts/stops

* let's see what's going on in this...

* specifically use int64_t for parid

* remove prints

* consistency

* visually inspect outputs...

* the expected testing output was wrong 0/1 are flipped in index 3

* remove prints

* better test

* remove unnecessary to_packed

* revert int64 changes, not necessary

---------

Co-authored-by: Lindsey Gray <[email protected]>
Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>

Author: 3 people

src/_ext.cpp

@@ -33,9 +33,9 @@ typedef struct {
   PyObject *next;
 } SwigPyObject;
 
-template <typename T>
-T swigtocpp(py::object obj) { // unwraps python object to get the cpp pointer
-                              // from the swig bindings
+template <typename T> T swigtocpp(py::object obj) {
+  // unwraps python object to get the cpp pointer
+  // from the swig bindings
   auto upointer = obj.attr("this").ptr();
   auto swigpointer = reinterpret_cast<SwigPyObject *>(upointer);
   auto objpointervoid = swigpointer->ptr;
@@ -59,47 +59,44 @@ output_wrapper interfacemulti(
     py::array_t<double, py::array::c_style | py::array::forcecast> pyi,
     py::array_t<double, py::array::c_style | py::array::forcecast> pzi,
     py::array_t<double, py::array::c_style | py::array::forcecast> Ei,
-    py::array_t<int, py::array::c_style | py::array::forcecast> offsets,
+    py::array_t<int, py::array::c_style | py::array::forcecast> starts,
+    py::array_t<int, py::array::c_style | py::array::forcecast> stops,
     py::object jetdef) {
-  py::buffer_info infooff = offsets.request();
+  // requesting buffer information of the input
+  py::buffer_info infostarts = starts.request();
+  py::buffer_info infostops = stops.request();
   py::buffer_info infopx = pxi.request();
-  py::buffer_info infopy =
-      pyi.request(); // requesting buffer information of the input
+  py::buffer_info infopy = pyi.request();
   py::buffer_info infopz = pzi.request();
   py::buffer_info infoE = Ei.request();
 
-  auto offptr = static_cast<int *>(infooff.ptr);
+  // pointers to the initial values
+  auto startsptr = static_cast<int *>(infostarts.ptr);
+  auto stopsptr = static_cast<int *>(infostops.ptr);
   auto pxptr = static_cast<double *>(infopx.ptr);
-  auto pyptr =
-      static_cast<double *>(infopy.ptr); // pointer to the initial value
+  auto pyptr = static_cast<double *>(infopy.ptr);
   auto pzptr = static_cast<double *>(infopz.ptr);
   auto Eptr = static_cast<double *>(infoE.ptr);
 
-  int dimoff = infooff.shape[0];
+  int dimoff = infostarts.shape[0];
   output_wrapper ow;
-  std::vector<double> nevents;
-  std::vector<double> offidx;
-  std::vector<double> constphi;
-  std::vector<double> idx;
-  std::vector<double> idxo;
-  for (int i = 0; i < dimoff - 1; i++) {
+  for (int i = 0; i < dimoff; i++) {
     std::vector<fj::PseudoJet> particles;
-    for (int j = *offptr; j < *(offptr + 1); j++) {
+    for (int j = *startsptr; j < *stopsptr; j++) {
       particles.push_back(fj::PseudoJet(*pxptr, *pyptr, *pzptr, *Eptr));
       pxptr++;
       pyptr++;
       pzptr++;
       Eptr++;
     }
-
     std::vector<fj::PseudoJet> jets;
     auto jet_def = swigtocpp<fj::JetDefinition *>(jetdef);
     std::shared_ptr<std::vector<fj::PseudoJet>> pj =
         std::make_shared<std::vector<fj::PseudoJet>>(particles);
     std::shared_ptr<fastjet::ClusterSequence> cs =
         std::make_shared<fastjet::ClusterSequence>(*pj, *jet_def);
-    auto j = cs->inclusive_jets();
-    offptr++;
+    startsptr++;
+    stopsptr++;
     ow.cse.push_back(cs);
     ow.parts.push_back(pj);
   }

src/fastjet/_generalevent.py

@@ -23,14 +23,15 @@ def __init__(self, data, jetdef):
                 behavior=self._clusterable_level[i].behavior,
                 attrs=self._clusterable_level[i].attrs,
             )
-            px, py, pz, E, offsets = self.extract_cons(self._clusterable_level[i])
+            px, py, pz, E, starts, stops = self.extract_cons(self._clusterable_level[i])
             px = self.correct_byteorder(px)
             py = self.correct_byteorder(py)
             pz = self.correct_byteorder(pz)
             E = self.correct_byteorder(E)
-            offsets = self.correct_byteorder(offsets)
+            starts = self.correct_byteorder(starts)
+            stops = self.correct_byteorder(stops)
             self._results.append(
-                fastjet._ext.interfacemulti(px, py, pz, E, offsets, jetdef)
+                fastjet._ext.interfacemulti(px, py, pz, E, starts, stops, jetdef)
             )
 
     def _check_listoffset_subtree(self, data):
@@ -360,9 +361,9 @@ def extract_cons(self, array):
         E = np.asarray(
             ak.Array(array.layout.content, behavior=array.behavior, attrs=array.attrs).E
         )
-        off = np.asarray(array.layout.stops)
-        off = np.insert(off, 0, 0)
-        return px, py, pz, E, off
+        starts = np.asarray(array.layout.starts)
+        stops = np.asarray(array.layout.stops)
+        return px, py, pz, E, starts, stops
 
     def _replace_multi(self):
         self._mod_data = self.data

src/fastjet/_multievent.py

@@ -9,15 +9,17 @@
 class _classmultievent:
     def __init__(self, data, jetdef):
         self.jetdef = jetdef
-        # data = ak.Array(data.layout.to_ListOffsetArray64(True)) If I do this vector can't convert the coordinates
         self.data = data
-        px, py, pz, E, offsets = self.extract_cons(self.data)
+        px, py, pz, E, starts, stops = self.extract_cons(self.data)
         px = self.correct_byteorder(px)
         py = self.correct_byteorder(py)
         pz = self.correct_byteorder(pz)
         E = self.correct_byteorder(E)
-        offsets = self.correct_byteorder(offsets)
-        self._results = fastjet._ext.interfacemulti(px, py, pz, E, offsets, jetdef)
+        starts = self.correct_byteorder(starts)
+        stops = self.correct_byteorder(stops)
+        self._results = fastjet._ext.interfacemulti(
+            px, py, pz, E, starts, stops, jetdef
+        )
 
     def _check_record(self, data):
         return data.layout.is_record or data.layout.is_numpy
@@ -34,9 +36,9 @@ def extract_cons(self, array):
         py = np.asarray(ak.Array(array.layout.content, behavior=array.behavior).py)
         pz = np.asarray(ak.Array(array.layout.content, behavior=array.behavior).pz)
         E = np.asarray(ak.Array(array.layout.content, behavior=array.behavior).E)
-        off = np.asarray(array.layout.stops)
-        off = np.insert(off, 0, 0)
-        return px, py, pz, E, off
+        starts = np.asarray(array.layout.starts)
+        stops = np.asarray(array.layout.stops)
+        return px, py, pz, E, starts, stops
 
     def single_to_jagged(self, array):
         single = ak.Array(
@@ -52,7 +54,9 @@ def single_to_jagged(self, array):
                     ["px", "py", "pz", "E"],
                     parameters={"__record__": "Momentum4D"},
                 ),
-            )
+            ),
+            behavior=array.behavior,
+            attrs=array.attrs,
         )
         return single
 
@@ -94,6 +98,7 @@ def inclusive_jets(self, min_pt):
                 ),
             ),
             behavior=self.data.behavior,
+            attrs=self.data.attrs,
         )
 
     def unclustered_particles(self):
@@ -114,6 +119,7 @@ def unclustered_particles(self):
                 ),
             ),
             behavior=self.data.behavior,
+            attrs=self.data.attrs,
         )
 
     def exclusive_jets(self, n_jets, dcut):
@@ -145,6 +151,7 @@ def exclusive_jets(self, n_jets, dcut):
                 ),
             ),
             behavior=self.data.behavior,
+            attrs=self.data.attrs,
         )
 
     def exclusive_jets_up_to(self, n_jets):
@@ -168,6 +175,7 @@ def exclusive_jets_up_to(self, n_jets):
                 ),
             ),
             behavior=self.data.behavior,
+            attrs=self.data.attrs,
         )
 
     def exclusive_jets_ycut(self, ycut):
@@ -188,6 +196,7 @@ def exclusive_jets_ycut(self, ycut):
                 ),
             ),
             behavior=self.data.behavior,
+            attrs=self.data.attrs,
         )
 
     def constituent_index(self, min_pt):
@@ -278,6 +287,8 @@ def exclusive_jets_softdrop_grooming(
                 "pzsoftdrop": jetpz,
             },
             depth_limit=1,
+            behavior=self.data.behavior,
+            attrs=self.data.attrs,
         )
         return out
 
@@ -452,6 +463,7 @@ def exclusive_subjets_up_to(self, data, nsub):
                 ),
             ),
             behavior=self.data.behavior,
+            attrs=self.data.attrs,
         )
 
     def exclusive_subdmerge(self, data, nsub):
@@ -554,6 +566,7 @@ def childless_pseudojets(self):
                 ),
             ),
             behavior=self.data.behavior,
+            attrs=self.data.attrs,
         )
 
     def jets(self):
@@ -574,6 +587,7 @@ def jets(self):
                 ),
             ),
             behavior=self.data.behavior,
+            attrs=self.data.attrs,
         )
 
     def get_parents(self, data):
@@ -600,6 +614,7 @@ def get_parents(self, data):
                 ),
             ),
             behavior=self.data.behavior,
+            attrs=self.data.attrs,
         )
 
     def get_child(self, data):
@@ -626,4 +641,5 @@ def get_child(self, data):
                 ),
             ),
             behavior=self.data.behavior,
+            attrs=self.data.attrs,
         )

src/fastjet/_singleevent.py

@@ -10,13 +10,16 @@ class _classsingleevent:
     def __init__(self, data, jetdef):
         self.jetdef = jetdef
         self.data = self.single_to_jagged(data)
-        px, py, pz, E, offsets = self.extract_cons(self.data)
+        px, py, pz, E, starts, stops = self.extract_cons(self.data)
         px = self.correct_byteorder(px)
         py = self.correct_byteorder(py)
         pz = self.correct_byteorder(pz)
         E = self.correct_byteorder(E)
-        offsets = self.correct_byteorder(offsets)
-        self._results = fastjet._ext.interfacemulti(px, py, pz, E, offsets, jetdef)
+        starts = self.correct_byteorder(starts)
+        stops = self.correct_byteorder(stops)
+        self._results = fastjet._ext.interfacemulti(
+            px, py, pz, E, starts, stops, jetdef
+        )
 
     def correct_byteorder(self, data):
         if data.dtype.byteorder == "=":
@@ -36,9 +39,9 @@ def extract_cons(self, array):
         py = np.asarray(ak.Array(array.layout.content, behavior=array.behavior).py)
         pz = np.asarray(ak.Array(array.layout.content, behavior=array.behavior).pz)
         E = np.asarray(ak.Array(array.layout.content, behavior=array.behavior).E)
-        off = np.asarray(array.layout.stops)
-        off = np.insert(off, 0, 0)
-        return px, py, pz, E, off
+        starts = np.asarray(array.layout.starts)
+        stops = np.asarray(array.layout.stops)
+        return px, py, pz, E, starts, stops
 
     def _check_record(self, data):
         return data.layout.is_record or data.layout.is_numpy

tests/test_002-exclusive_jets.py

@@ -534,6 +534,43 @@ def test_exclusive_constituents_multi():
     )
 
 
+def test_exclusive_constituents_masking_multi():
+    array = ak.Array(
+        [
+            [
+                {"px": -0.181, "py": -0.798, "pz": -0.652, "E": 1.06},
+                {"px": 0.377, "py": 0.116, "pz": -0.0749, "E": 0.425},
+            ],
+            [
+                {"px": 0.54, "py": -0.65, "pz": -0.00527, "E": 0.857},
+                {"px": 0.253, "py": -0.0342, "pz": -0.0731, "E": 0.3},
+            ],
+            [
+                {"px": 0.294, "py": 0.254, "pz": -0.259, "E": 0.467},
+                {"px": 4.65, "py": -1.88, "pz": -3.29, "E": 6},
+            ],
+            [
+                {"px": 1.45, "py": -0.179, "pz": -0.876, "E": 1.71},
+                {"px": 12.8, "py": -1.8, "pz": -7.18, "E": 14.8},
+            ],
+            [
+                {"px": -3.55, "py": -1.64, "pz": -0.0941, "E": 3.91},
+                {"px": -1.33, "py": -1.03, "pz": 0.147, "E": 1.7},
+            ],
+        ],
+        with_name="Momentum4D",
+    )
+    mask = ak.Array([True, True, False, True, True])
+    jetdef = fastjet.JetDefinition(fastjet.kt_algorithm, 1)
+    result_all = fastjet.ClusterSequence(
+        array, jetdef
+    ).exclusive_jets_constituent_index(njets=2)
+    result_mask = fastjet.ClusterSequence(
+        array[mask], jetdef
+    ).exclusive_jets_constituent_index(njets=2)
+    assert ak.all(result_mask == result_all[mask])
+
+
 def test_exclusive_ycut():
     array = ak.Array(
         [