diff --git a/src/_ext.cpp b/src/_ext.cpp index 837fa477..2fb72d04 100644 --- a/src/_ext.cpp +++ b/src/_ext.cpp @@ -15,12 +15,15 @@ #include #include #include +#include #include #include #include #include +#include + namespace fj = fastjet; namespace py = pybind11; using namespace pybind11::literals; @@ -1581,10 +1584,110 @@ PYBIND11_MODULE(_ext, m) { Returns: pt, eta, phi, m of inclusive jets. )pbdoc") + .def("to_numpy_softdrop_grooming", + [](const output_wrapper ow, const int n_jets = 1, double beta = 0, double symmetry_cut = 0.1, + std::string symmetry_measure = "scalar_z", double R0 = 0.8, std::string recursion_choice = "larger_pt", + /*const FunctionOfPseudoJet * subtractor = 0,*/ double mu_cut = std::numeric_limits::infinity()){ + + auto css = ow.cse; + std::vector consts_groomed_px; + std::vector consts_groomed_py; + std::vector consts_groomed_pz; + std::vector consts_groomed_E; + std::vector nconstituents; + std::vector jet_groomed_pt; + std::vector jet_groomed_eta; + std::vector jet_groomed_phi; + std::vector jet_groomed_m; + + fastjet::contrib::RecursiveSymmetryCutBase::SymmetryMeasure sym_meas = fastjet::contrib::RecursiveSymmetryCutBase::SymmetryMeasure::scalar_z; + if (symmetry_measure == "scalar_z") { + sym_meas = fastjet::contrib::RecursiveSymmetryCutBase::SymmetryMeasure::scalar_z; + } + else if (symmetry_measure == "vector_z") { + sym_meas = fastjet::contrib::RecursiveSymmetryCutBase::SymmetryMeasure::vector_z; + } + else if (symmetry_measure == "y") { + sym_meas = fastjet::contrib::RecursiveSymmetryCutBase::SymmetryMeasure::y; + } + else if (symmetry_measure == "theta_E") { + sym_meas = fastjet::contrib::RecursiveSymmetryCutBase::SymmetryMeasure::theta_E; + } + else if (symmetry_measure == "cos_theta_E") { + sym_meas = fastjet::contrib::RecursiveSymmetryCutBase::SymmetryMeasure::cos_theta_E; + } + + fastjet::contrib::RecursiveSymmetryCutBase::RecursionChoice rec_choice = fastjet::contrib::RecursiveSymmetryCutBase::RecursionChoice::larger_pt; + if (recursion_choice == "larger_pt") { + rec_choice = fastjet::contrib::RecursiveSymmetryCutBase::RecursionChoice::larger_pt; + } + else if (recursion_choice == "larger_mt") { + rec_choice = fastjet::contrib::RecursiveSymmetryCutBase::RecursionChoice::larger_mt; + } + else if (recursion_choice == "larger_m") { + rec_choice = fastjet::contrib::RecursiveSymmetryCutBase::RecursionChoice::larger_m; + } + else if (recursion_choice == "larger_E") { + rec_choice = fastjet::contrib::RecursiveSymmetryCutBase::RecursionChoice::larger_E; + } + + fastjet::contrib::SoftDrop* sd = new fastjet::contrib::SoftDrop(beta, symmetry_cut, sym_meas, R0, mu_cut, rec_choice/*, subtractor*/); + + for (unsigned int i = 0; i < css.size(); i++){ // iterate through events + auto jets = css[i]->exclusive_jets(n_jets); + for (unsigned int j = 0; j < jets.size(); j++){ + auto soft = sd->result(jets[j]); + nconstituents.push_back(soft.constituents().size()); + jet_groomed_pt.push_back(soft.pt()); + jet_groomed_eta.push_back(soft.eta()); + jet_groomed_phi.push_back(soft.phi()); + jet_groomed_m.push_back(soft.m()); + for (unsigned int k = 0; k < soft.constituents().size(); k++){ + consts_groomed_px.push_back(soft.constituents()[k].px()); + consts_groomed_py.push_back(soft.constituents()[k].py()); + consts_groomed_pz.push_back(soft.constituents()[k].pz()); + consts_groomed_E.push_back(soft.constituents()[k].E()); + } + } + } + + auto consts_px = py::array(consts_groomed_px.size(), consts_groomed_px.data()); + auto consts_py = py::array(consts_groomed_py.size(), consts_groomed_py.data()); + auto consts_pz = py::array(consts_groomed_pz.size(), consts_groomed_pz.data()); + auto consts_E = py::array(consts_groomed_E.size(), consts_groomed_E.data()); + auto eventsize = py::array(nconstituents.size(), nconstituents.data()); + auto jet_pt = py::array(jet_groomed_pt.size(), jet_groomed_pt.data()); + auto jet_eta = py::array(jet_groomed_eta.size(), jet_groomed_eta.data()); + auto jet_phi = py::array(jet_groomed_phi.size(), jet_groomed_phi.data()); + auto jet_m = py::array(jet_groomed_m.size(), jet_groomed_m.data()); + + return std::make_tuple( + consts_px, + consts_py, + consts_pz, + consts_E, + eventsize, + jet_pt, + jet_eta, + jet_phi, + jet_m + ); + }, R"pbdoc( + Performs softdrop pruning on jets. + Args: + n_jets: number of exclusive subjets. + beta: softdrop beta parameter. + symmetry_cut: softdrop symmetry cut value. + symmetry_measure: Which symmetry measure to use, found in RecursiveSymmetryCutBase.hh + R0: softdrop R0 parameter. + recursion_choice: Which recursion choice to use, found in RecursiveSymmetryCutBase.hh + subtractor: an optional pointer to a pileup subtractor (ignored if zero) + Returns: + Returns an array of values from the jet after it has been groomed by softdrop. + )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") { + [](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", bool normalized = true) { auto css = ow.cse; - int64_t len = css.size(); std::transform(func.begin(), func.end(), func.begin(), [](unsigned char c){ return std::tolower(c); }); @@ -1623,14 +1726,15 @@ PYBIND11_MODULE(_ext, m) { energy_correlator = std::make_shared(beta);} else if (func == "u3") { energy_correlator = std::make_shared(beta);} - else if (func == "generic") { + else if (func == "generic" && normalized == false) { 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. + else if (func == "generic" && normalized == true) { + energy_correlator = std::make_shared(angles, npoint, beta);} //Using the Generalized class with angles=-1 returns a generic ECF that has been normalized 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]); // @@ -1681,14 +1785,12 @@ PYBIND11_MODULE(_ext, m) { 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++){ diff --git a/src/fastjet/__init__.py b/src/fastjet/__init__.py index ad81f1bb..85ced84e 100644 --- a/src/fastjet/__init__.py +++ b/src/fastjet/__init__.py @@ -369,15 +369,27 @@ def exclusive_jets_constituents(self, njets: int = 10) -> ak.Array: """ raise AssertionError() + + def exclusive_jets_softdrop_grooming( + self, njets: int = 1, beta: int = 0, symmetry_cut: float = 0.1, symmetry_measure="scalar_z", R0: float = 0.8, recursion_choice="larger_pt", + #subtractor: int = 0, + mu_cut: float = float('inf'), + ) -> ak.Array: + """Performs softdrop pruning on jets. + Args: + n_jets: number of exclusive subjets. + beta: softdrop beta parameter. + symmetry_cut: softdrop symmetry cut value. + symmetry_measure: Which symmetry measure to use, found in RecursiveSymmetryCutBase.hh + R0: softdrop R0 parameter. + recursion_choice: Which recursion choice to use, found in RecursiveSymmetryCutBase.hh + subtractor: an optional pointer to a pileup subtractor (ignored if zero) + Returns: + Returns an array of values from the jet after it has been groomed by softdrop.""" + raise AssertionError() def exclusive_jets_energy_correlator( - self, - njets: int = 1, - beta: int = 1, - npoint: int = 0, - angles: int = 0, - alpha=0, - func="generalized", + self, njets: int = 1, beta: int = 1, npoint: int = 0, angles: int = -1, alpha=0, func="generalized", normalized=True, ) -> ak.Array: """Returns the energy correlator of each exclusive jet. @@ -390,6 +402,7 @@ def exclusive_jets_energy_correlator( Returns: awkward.highlevel.Array: Returns an Awkward Array of the same type as the input. """ + 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. diff --git a/src/fastjet/_generalevent.py b/src/fastjet/_generalevent.py index 60bf671b..cc79c6f4 100644 --- a/src/fastjet/_generalevent.py +++ b/src/fastjet/_generalevent.py @@ -438,9 +438,49 @@ def exclusive_jets_constituent_index(self, njets): ) res = ak.Array(self._replace_multi()) return res + + def exclusive_jets_softdrop_grooming( + self, njets=1, beta = 0.0, symmetry_cut = 0.1, symmetry_measure = "scalar_z", R0 = 0.8, recursion_choice = "larger_pt", + #subtractor = 0, + mu_cut = float('inf'), + ): + 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.to_numpy_softdrop_grooming( + njets, beta, symmetry_cut, symmetry_measure, R0, recursion_choice, #subtractor, + mu_cut, + ) + + px = ak.unflatten(ak.Array(ak.contents.NumpyArray(np_results[0])), ak.Array(ak.contents.NumpyArray(np_results[4])), highlevel=False) + py = ak.unflatten(ak.Array(ak.contents.NumpyArray(np_results[1])), ak.Array(ak.contents.NumpyArray(np_results[4])), highlevel=False) + pz = ak.unflatten(ak.Array(ak.contents.NumpyArray(np_results[2])), ak.Array(ak.contents.NumpyArray(np_results[4])), highlevel=False) + E = ak.unflatten(ak.Array(ak.contents.NumpyArray(np_results[3])), ak.Array(ak.contents.NumpyArray(np_results[4])), highlevel=False) + jetpt = ak.Array(ak.contents.NumpyArray(np_results[5])) + jeteta = ak.Array(ak.contents.NumpyArray(np_results[6])) + jetphi = ak.Array(ak.contents.NumpyArray(np_results[7])) + jetmass = ak.Array(ak.contents.NumpyArray(np_results[8])) + + self._out.append(ak.zip({ + "constituents": + ak.zip( + {"px": px, "py": py, "pz": pz, "E": E}, depth_limit=2), + "msoftdrop": jetmass, + "ptsoftdrop": jetpt, + "etasoftdrop": jeteta, + "phisoftdrop": jetphi,}, + depth_limit=1 + ) + ) + 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" + self, njets=1, n_point=0, angles: int = -1, beta=1, alpha=0, func="generalized", normalized=True, ): if njets <= 0: raise ValueError("Njets cannot be <= 0") diff --git a/src/fastjet/_multievent.py b/src/fastjet/_multievent.py index ee809749..c46b13a7 100644 --- a/src/fastjet/_multievent.py +++ b/src/fastjet/_multievent.py @@ -191,14 +191,48 @@ def exclusive_jets_constituent_index(self, njets): ) out = ak.Array(ak.contents.ListOffsetArray(ak.index.Index64(off), out.layout)) return out + + def exclusive_jets_softdrop_grooming( + self, njets=1, beta = 0.0, symmetry_cut = 0.1, symmetry_measure = "scalar_z", R0 = 0.8, recursion_choice = "larger_pt", + #subtractor = 0, + mu_cut = float('inf'), + ): + if njets <= 0: + raise ValueError("Njets cannot be <= 0") + np_results = self._results.to_numpy_softdrop_grooming( + njets, beta, symmetry_cut, symmetry_measure, R0, recursion_choice, #subtractor, + mu_cut, + ) + + px = ak.unflatten(ak.Array(ak.contents.NumpyArray(np_results[0])), ak.Array(ak.contents.NumpyArray(np_results[4])), highlevel=False) + py = ak.unflatten(ak.Array(ak.contents.NumpyArray(np_results[1])), ak.Array(ak.contents.NumpyArray(np_results[4])), highlevel=False) + pz = ak.unflatten(ak.Array(ak.contents.NumpyArray(np_results[2])), ak.Array(ak.contents.NumpyArray(np_results[4])), highlevel=False) + E = ak.unflatten(ak.Array(ak.contents.NumpyArray(np_results[3])), ak.Array(ak.contents.NumpyArray(np_results[4])), highlevel=False) + jetpt = ak.Array(ak.contents.NumpyArray(np_results[5])) + jeteta = ak.Array(ak.contents.NumpyArray(np_results[6])) + jetphi = ak.Array(ak.contents.NumpyArray(np_results[7])) + jetmass = ak.Array(ak.contents.NumpyArray(np_results[8])) + + out = ak.zip({ + "constituents": + ak.zip( + {"px": px, "py": py, "pz": pz, "E": E}, depth_limit=2), + "msoftdrop": jetmass, + "ptsoftdrop": jetpt, + "etasoftdrop": jeteta, + "phisoftdrop": jetphi,}, + depth_limit=1 + ) + + return out def exclusive_jets_energy_correlator( - self, njets=1, beta=1, npoint=0, angles=0, alpha=0, func="generalized" + self, njets=1, beta=1, npoint=0, angles=-1, alpha=0, func="generalized", normalized=True, ): if njets <= 0: raise ValueError("Njets cannot be <= 0") np_results = self._results.to_numpy_energy_correlators( - njets, beta, npoint, angles, alpha, func + njets, beta, npoint, angles, alpha, func, normalized, ) out = ak.Array(ak.contents.NumpyArray(np_results)) return out diff --git a/src/fastjet/_pyjet.py b/src/fastjet/_pyjet.py index 4c6c8dd6..376b2114 100644 --- a/src/fastjet/_pyjet.py +++ b/src/fastjet/_pyjet.py @@ -142,12 +142,22 @@ 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_softdrop_grooming( + self, njets=1, beta = 0.0, symmetry_cut = 0.1, symmetry_measure = "scalar_z", R0 = 0.8, recursion_choice = "larger_pt", + #subtractor = 0, + mu_cut = float('inf'), + ): + return self._internalrep.exclusive_jets_softdrop_grooming( + njets, beta, symmetry_cut, symmetry_measure, R0, recursion_choice, #subtractor, + mu_cut, + ) def exclusive_jets_energy_correlator( - self, njets=1, beta=1, npoint=0, angles=0, alpha=0, func="generalized" + self, njets=1, beta=1, npoint=0, angles=-1, alpha=0, func="generalized", normalized=True ): return self._internalrep.exclusive_jets_energy_correlator( - njets, beta, npoint, angles, alpha, func + njets, beta, npoint, angles, alpha, func, normalized, ) def exclusive_jets_lund_declusterings(self, njets=10): @@ -428,8 +438,26 @@ 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_softdrop_grooming( + self, njets=1, beta = 0.0, symmetry_cut = 0.1, symmetry_measure = "scalar_z", R0 = 0.8, recursion_choice = "larger_pt", + #subtractor = 0, + mu_cut = float('inf'), + ): + return _dak_dispatch( + self, + "exclusive_jets_softdrop_grooming", + njets=njets, + beta=beta, + symmetry_cut=symmetry_cut, + symmetry_measure=symmetry_measure, + R0=R0, + recursion_choice=recursion_choice, + #subtractor=subtractor, + mu_cut=mu_cut, + ) + def exclusive_jets_energy_correlator( - self, njets=1, beta=1, npoint=0, angles=0, alpha=0, func="generalized" + self, njets=1, beta=1, npoint=0, angles=-1, alpha=0, func="generalized", normalized=False, ): return _dak_dispatch( self, @@ -440,6 +468,7 @@ def exclusive_jets_energy_correlator( angles=angles, alpha=alpha, func=func, + normalized=normalized, ) def exclusive_jets_lund_declusterings(self, njets=10): diff --git a/src/fastjet/_singleevent.py b/src/fastjet/_singleevent.py index 37ffa29d..dbe30fa8 100644 --- a/src/fastjet/_singleevent.py +++ b/src/fastjet/_singleevent.py @@ -180,14 +180,48 @@ 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_softdrop_grooming( + self, njets=1, beta = 0.0, symmetry_cut = 0.1, symmetry_measure = "scalar_z", R0 = 0.8, recursion_choice = "larger_pt", + #subtractor = 0, + mu_cut = float('inf'), + ): + if njets <= 0: + raise ValueError("Njets cannot be <= 0") + + np_results = self._results.to_numpy_softdrop_grooming( + njets, beta, symmetry_cut, symmetry_measure, R0, recursion_choice, #subtractor, + mu_cut, + ) + + px = ak.unflatten(ak.Array(ak.contents.NumpyArray(np_results[0])), ak.Array(ak.contents.NumpyArray(np_results[4])), highlevel=False) + py = ak.unflatten(ak.Array(ak.contents.NumpyArray(np_results[1])), ak.Array(ak.contents.NumpyArray(np_results[4])), highlevel=False) + pz = ak.unflatten(ak.Array(ak.contents.NumpyArray(np_results[2])), ak.Array(ak.contents.NumpyArray(np_results[4])), highlevel=False) + E = ak.unflatten(ak.Array(ak.contents.NumpyArray(np_results[3])), ak.Array(ak.contents.NumpyArray(np_results[4])), highlevel=False) + jetpt = ak.Array(ak.contents.NumpyArray(np_results[5])) + jeteta = ak.Array(ak.contents.NumpyArray(np_results[6])) + jetphi = ak.Array(ak.contents.NumpyArray(np_results[7])) + jetmass = ak.Array(ak.contents.NumpyArray(np_results[8])) + + out = ak.zip({ + "constituents": + ak.zip( + {"px": px, "py": py, "pz": pz, "E": E}, depth_limit=2), + "msoftdrop": jetmass, + "ptsoftdrop": jetpt, + "etasoftdrop": jeteta, + "phisoftdrop": jetphi,}, + depth_limit=1 + ) + return out[0] + def exclusive_jets_energy_correlator( - self, njets=1, beta=1, npoint=0, angles=0, alpha=0, func="generalized" + self, njets=1, beta=1, npoint=0, angles=-1, alpha=0, func="generalized", normalized=True, ): if njets <= 0: raise ValueError("Njets cannot be <= 0") np_results = self._results.to_numpy_energy_correlators( - njets, beta, npoint, angles, alpha, func + njets, beta, npoint, angles, alpha, func, normalized, ) out = ak.Array(ak.contents.NumpyArray(np_results)) return out[0]