diff --git a/src/_ext.cpp b/src/_ext.cpp index e181ecc0..651b6aa7 100644 --- a/src/_ext.cpp +++ b/src/_ext.cpp @@ -186,11 +186,14 @@ PYBIND11_MODULE(_ext, m) { auto bufparid = parid.request(); int *ptrid = (int *)bufparid.ptr; - auto eventoffsets = py::array(py::buffer_info(nullptr, sizeof(int), py::format_descriptor::value, 1, {len}, {sizeof(int)})); + auto eventoffsets = py::array(py::buffer_info(nullptr, sizeof(int), py::format_descriptor::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::value, 1, {jk}, {sizeof(int)})); auto bufjetoffsets = jetoffsets.request(); int *ptrjetoffsets = (int *)bufjetoffsets.ptr; @@ -298,6 +301,78 @@ PYBIND11_MODULE(_ext, m) { Returns: pt, eta, phi, m of inclusive jets. )pbdoc") + .def("to_numpy_exclusive_njet_with_constituents", + [](const output_wrapper ow, const int n_jets = 0) { + auto css = ow.cse; + int64_t len = css.size(); + auto jk = 0; + auto sizepar = 0; + + for(int i = 0; i < len; i++){ + jk += css[i]->exclusive_jets(n_jets).size(); + sizepar += css[i]->n_particles(); + } + jk++; + + auto parid = py::array(py::buffer_info(nullptr, sizeof(int), py::format_descriptor::value, 1, {sizepar}, {sizeof(int)})); + auto bufparid = parid.request(); + int *ptrid = (int *)bufparid.ptr; + + auto eventoffsets = py::array(py::buffer_info(nullptr, sizeof(int), py::format_descriptor::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::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 idx = css[i]->particle_jet_indices(jets); + int64_t sizz = css[i]->n_particles(); + auto prev = ptrjetoffsets[jetidx-1]; + + for (unsigned int j = 0; j < jets.size(); j++){ + ptrjetoffsets[jetidx] = jets[j].constituents().size() + prev; + prev = ptrjetoffsets[jetidx]; + jetidx++; + } + for(int k = 0; k < size; k++){ // iterate through jets in event + for(int j = 0; j ::value, 1, {jk}, {sizeof(int)})); auto bufparid = parid.request(); int *ptrid = (int *)bufparid.ptr; - auto eventoffsets = py::array(py::buffer_info(nullptr, sizeof(int), py::format_descriptor::value, 1, {len}, {sizeof(int)})); + auto eventoffsets = py::array(py::buffer_info(nullptr, sizeof(int), py::format_descriptor::value, 1, {len+1}, {sizeof(int)})); auto bufeventoffsets = eventoffsets.request(); int *ptreventoffsets = (int *)bufeventoffsets.ptr; size_t eventidx = 0; + ptreventoffsets[eventidx] = 0; + eventidx++; size_t idxh = 0; auto eventprev = 0; for (unsigned int i = 0; i < css.size(); i++){ diff --git a/src/fastjet/__init__.py b/src/fastjet/__init__.py index 16f68b5a..eee470a2 100644 --- a/src/fastjet/__init__.py +++ b/src/fastjet/__init__.py @@ -286,6 +286,18 @@ def constituent_index(self, min_pt: float = 0) -> ak.Array: """ raise AssertionError() + def exclusive_jets_constituents_index(self, njets: int = 10) -> ak.Array: + """Returns the index of the constituent of each exclusive jet. + + 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 constituents(self, min_p: float = 0) -> ak.Array: """Returns the particles that make up each Jet. @@ -297,6 +309,18 @@ def constituents(self, min_p: float = 0) -> ak.Array: """ raise AssertionError() + def exclusive_jets_constituents(self, njets: int = 10) -> ak.Array: + """Returns the particles that make up each exclusive jet. + + 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. diff --git a/src/fastjet/_generalevent.py b/src/fastjet/_generalevent.py index 133f0870..5dab7794 100644 --- a/src/fastjet/_generalevent.py +++ b/src/fastjet/_generalevent.py @@ -608,7 +608,36 @@ def constituents(self, min_pt): self._input_flag = 0 for i in range(len(self._clusterable_level)): np_results = self._results[i].to_numpy_with_constituents(min_pt) - off = np.insert(np_results[-1], 0, 0) + off = np_results[-1] + out = ak.Array( + ak.layout.ListOffsetArray64( + ak.layout.Index64(np_results[0]), + ak.layout.NumpyArray(np_results[1]), + ), + behavior=self.data.behavior, + ) + outputs_to_inputs = ak.Array( + ak.layout.ListOffsetArray64(ak.layout.Index64(off), out.layout) + ) + shape = ak.num(outputs_to_inputs) + total = np.sum(shape) + duplicate = ak.unflatten(np.zeros(total, np.int64), shape) + prepared = self._clusterable_level[i][:, np.newaxis][duplicate] + self._out.append(prepared[outputs_to_inputs]) + res = ak.Array(self._replace_multi()) + return res + + def exclusive_jets_constituents(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]), @@ -632,7 +661,33 @@ def constituent_index(self, min_pt): self._input_flag = 0 for i in range(len(self._clusterable_level)): np_results = self._results[i].to_numpy_with_constituents(min_pt) - off = np.insert(np_results[-1], 0, 0) + off = np_results[-1] + out = ak.Array( + ak.layout.ListOffsetArray64( + ak.layout.Index64(np_results[0]), + ak.layout.NumpyArray(np_results[1]), + ), + 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 exclusive_jets_constituent_index(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]), @@ -818,7 +873,7 @@ def unique_history_order(self): self._input_flag = 0 for i in range(len(self._clusterable_level)): np_results = self._results[i].to_numpy_unique_history_order() - off = np.insert(np_results[-1], 0, 0) + off = np_results[-1] self._out.append( ak.Array( ak.layout.ListOffsetArray64( diff --git a/src/fastjet/_multievent.py b/src/fastjet/_multievent.py index 4942388e..7552b9d3 100644 --- a/src/fastjet/_multievent.py +++ b/src/fastjet/_multievent.py @@ -175,7 +175,21 @@ def exclusive_jets_ycut(self, ycut): def constituent_index(self, min_pt): np_results = self._results.to_numpy_with_constituents(min_pt) - off = np.insert(np_results[-1], 0, 0) + off = np_results[-1] + out = ak.Array( + ak.layout.ListOffsetArray64( + ak.layout.Index64(np_results[0]), ak.layout.NumpyArray(np_results[1]) + ) + ) + out = ak.Array(ak.layout.ListOffsetArray64(ak.layout.Index64(off), out.layout)) + return out + + def exclusive_jets_constituent_index(self, njets): + if njets <= 0: + raise ValueError("Njets cannot be <= 0") + + np_results = self._results.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.NumpyArray(np_results[1]) @@ -186,7 +200,7 @@ def constituent_index(self, min_pt): def unique_history_order(self): np_results = self._results.to_numpy_unique_history_order() - off = np.insert(np_results[-1], 0, 0) + off = np_results[-1] out = ak.Array( ak.layout.ListOffsetArray64( ak.layout.Index64(off), ak.layout.NumpyArray(np_results[0]) @@ -232,6 +246,17 @@ def constituents(self, min_pt): prepared = self.data[:, np.newaxis][duplicate] return prepared[outputs_to_inputs] + def exclusive_jets_constituents(self, njets): + if njets <= 0: + raise ValueError("Njets cannot be <= 0") + + outputs_to_inputs = self.exclusive_jets_constituent_index(njets) + shape = ak.num(outputs_to_inputs) + total = np.sum(shape) + duplicate = ak.unflatten(np.zeros(total, np.int64), shape) + prepared = self.data[:, np.newaxis][duplicate] + return prepared[outputs_to_inputs] + def exclusive_subjets(self, data, dcut, nsub): try: px = data.px diff --git a/src/fastjet/_pyjet.py b/src/fastjet/_pyjet.py index ca3d33d5..8db4a571 100644 --- a/src/fastjet/_pyjet.py +++ b/src/fastjet/_pyjet.py @@ -206,6 +206,12 @@ def constituent_index(self, min_pt=0): def constituents(self, min_pt=0): return self._internalrep.constituents(min_pt) + def exclusive_jets_constituent_index(self, njets=10): + return self._internalrep.exclusive_jets_constituent_index(njets) + + def exclusive_jets_constituents(self, njets=10): + return self._internalrep.exclusive_jets_constituents(njets) + def exclusive_dmerge(self, njets=10): return self._internalrep.exclusive_dmerge(njets) diff --git a/src/fastjet/_singleevent.py b/src/fastjet/_singleevent.py index 62926374..25efe05a 100644 --- a/src/fastjet/_singleevent.py +++ b/src/fastjet/_singleevent.py @@ -163,7 +163,21 @@ def exclusive_jets_ycut(self, ycut): def constituent_index(self, min_pt): np_results = self._results.to_numpy_with_constituents(min_pt) - off = np.insert(np_results[-1], 0, 0) + off = np_results[-1] + out = ak.Array( + ak.layout.ListOffsetArray64( + ak.layout.Index64(np_results[0]), ak.layout.NumpyArray(np_results[1]) + ) + ) + out = ak.Array(ak.layout.ListOffsetArray64(ak.layout.Index64(off), out.layout)) + return out[0] + + def exclusive_jets_constituent_index(self, njets): + if njets <= 0: + raise ValueError("Njets cannot be <= 0") + + np_results = self._results.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.NumpyArray(np_results[1]) @@ -179,7 +193,28 @@ def unique_history_order(self): def constituents(self, min_pt): np_results = self._results.to_numpy_with_constituents(min_pt) - off = np.insert(np_results[-1], 0, 0) + off = np_results[-1] + out = ak.Array( + ak.layout.ListOffsetArray64( + ak.layout.Index64(np_results[0]), ak.layout.NumpyArray(np_results[1]) + ) + ) + outputs_to_inputs = ak.Array( + ak.layout.ListOffsetArray64(ak.layout.Index64(off), out.layout) + ) + shape = ak.num(outputs_to_inputs) + total = np.sum(shape) + duplicate = ak.unflatten(np.zeros(total, np.int64), shape) + prepared = self.data[:, np.newaxis][duplicate] + return prepared[outputs_to_inputs][0] + + def exclusive_jets_constituents(self, njets): + if njets <= 0: + raise ValueError("Njets cannot be <= 0") + + np_results = self._results.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.NumpyArray(np_results[1]) diff --git a/tests/test_002-exclusive_jets.py b/tests/test_002-exclusive_jets.py index 206d7ed4..414fd3bb 100644 --- a/tests/test_002-exclusive_jets.py +++ b/tests/test_002-exclusive_jets.py @@ -64,6 +64,73 @@ def test_exclusive_multi(): assert multi_exclusive_dcut == cluster.exclusive_jets(dcut=0.0001).to_list() +def test_exclusive_constituents_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}, + ] + ) + jetdef = fastjet.JetDefinition(fastjet.antikt_algorithm, 0.6) + cluster = fastjet._pyjet.AwkwardClusterSequence(array, jetdef) + + constituent_output = [ + [ + {"px": 32.2, "py": 64.21, "pz": 543.34, "E": 24.12}, + {"px": 32.45, "py": 63.21, "pz": 543.14, "E": 24.56}, + ], + [{"px": 1.2, "py": 3.2, "pz": 5.4, "E": 2.5}], + ] + assert constituent_output == cluster.exclusive_jets_constituents(2).to_list() + constituent_index_output = [[1, 2], [0]] + assert ( + constituent_index_output + == cluster.exclusive_jets_constituent_index(2).to_list() + ) + + +def test_exclusive_constituents_multi(): + 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}, + ], + [ + {"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}, + ], + ] + ) + jetdef = fastjet.JetDefinition(fastjet.antikt_algorithm, 0.6) + cluster = fastjet._pyjet.AwkwardClusterSequence(array, jetdef) + + constituents_output = [ + [ + [ + {"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": 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}], + ], + ] + assert cluster.exclusive_jets_constituents(2).to_list() == constituents_output + constituent_index_out = [[[1, 2], [0]], [[1, 2], [0]]] + assert ( + constituent_index_out == cluster.exclusive_jets_constituent_index(2).to_list() + ) + + def test_exclusive_ycut(): array = ak.Array( [