feat: Softdrop

.github/workflows/ci.yml

@@ -78,7 +78,7 @@ jobs:
  with:
  submodules: recursive
 
- - uses: docker/setup-qemu-action@v2.1.0
+ - uses: docker/setup-qemu-action@v2.2.0
  if: matrix.arch != 'auto64'
  with:
  platforms: all
@@ -89,7 +89,7 @@ jobs:
  brew install make gcc automake swig gmp mpfr boost
  export PATH="/usr/local/opt/make/libexec/gnubin:$PATH"
 
- - uses: pypa/cibuildwheel@v2.13.0
+ - uses: pypa/cibuildwheel@v2.13.1
  env:
  CIBW_ARCHS: ${{ matrix.arch }}
  CIBW_BUILD: cp${{ matrix.python }}-*

.github/workflows/wheels.yml

@@ -49,7 +49,7 @@ jobs:
  arch: aarch64
 
  steps:
- - uses: docker/setup-qemu-action@v2.1.0
+ - uses: docker/setup-qemu-action@v2.2.0
  if: matrix.arch != 'auto64'
  with:
  platforms: arm64
@@ -64,7 +64,7 @@ jobs:
  brew install make gcc automake swig gmp mpfr boost
  export PATH="/usr/local/opt/make/libexec/gnubin:$PATH"
 
- - uses: pypa/cibuildwheel@v2.13.0
+ - uses: pypa/cibuildwheel@v2.13.1
  env:
  CIBW_ARCHS: ${{ matrix.arch }}
  CIBW_BUILD: cp${{ matrix.python }}-*

.pre-commit-config.yaml

@@ -24,7 +24,7 @@ repos:
  - id: isort
 
 - repo: https://github.com/asottile/pyupgrade
- rev: v3.4.0
+ rev: v3.6.0
  hooks:
  - id: pyupgrade
  args: ["--py36-plus"]

setup.cfg

@@ -58,6 +58,7 @@ test =
  uproot>=5
  dask>=2023.4.0;python_version>"3.7"
  dask-awkward>=2023.4.2;python_version>"3.7"
+ distributed>=2023.4.0;python_version>"3.7"
 
 [tool:pytest]
 addopts = -vv -rs -Wd

src/_ext.cpp

@@ -15,12 +15,15 @@
 #include <fastjet/PseudoJet.hh>
 #include <fastjet/contrib/EnergyCorrelator.hh>
 #include <fastjet/contrib/LundGenerator.hh>
+#include <fastjet/contrib/SoftDrop.hh>
 
 #include <pybind11/numpy.h>
 #include <pybind11/operators.h>
 #include <pybind11/pybind11.h>
 #include <pybind11/stl.h>
 
+#include <iostream>
+
 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<PseudoJet> * subtractor = 0,*/ double mu_cut = std::numeric_limits<double>::infinity()){
+
+ auto css = ow.cse;
+ std::vector<double> consts_groomed_px;
+ std::vector<double> consts_groomed_py;
+ std::vector<double> consts_groomed_pz;
+ std::vector<double> consts_groomed_E;
+ std::vector<int> nconstituents;
+ std::vector<double> jet_groomed_pt;
+ std::vector<double> jet_groomed_eta;
+ std::vector<double> jet_groomed_phi;
+ std::vector<double> 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", 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); });
@@ -1632,7 +1735,6 @@ PYBIND11_MODULE(_ext, m) {
 
  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]); //
@@ -1683,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++){

src/fastjet/__init__.py

@@ -189,14 +189,7 @@
 
 
 class JetDefinition(JetDefinitionNoCast):
- def __init__(
- self,
- jet_algorithm_in,
- R_in,
- recomb_scheme_in=0,
- strategy_in=1,
- nparameters_in=1,
- ):
+ def __init__(self, *args, **kwargs):
  r"""
 
  `JetDefinition(JetAlgorithm jet_algorithm_in, double R_in, RecombinationScheme
@@ -206,6 +199,14 @@ def __init__(
  how algorithically to run it).
 
  """
+
+ R_in = kwargs.pop("R_in", None)
+ as_kwargs = False
+ if R_in is None:
+ R_in = args[1]
+ else:
+ as_kwargs = True
+
  if not isinstance(R_in, (float, int)):
  raise ValueError(
  f"R_in should be a real number, got {R_in} of type {type(R_in)}"
@@ -214,9 +215,23 @@ def __init__(
  if isinstance(R_in, int):
  R_in = float(R_in)
 
- super().__init__(
- jet_algorithm_in, R_in, recomb_scheme_in, strategy_in, nparameters_in
- )
+ new_args = args
+ new_kwargs = kwargs
+ if as_kwargs:
+ new_kwargs = kwargs.copy()
+ new_kwargs["R_in"] = R_in
+ else:
+ new_args = (args[0], R_in, *args[2:])
+
+ self.args = new_args
+ self.kwargs = new_kwargs
+ super().__init__(*new_args, **kwargs)
+
+ def __setstate__(self, state):
+ self.__init__(*state["args"], **state["kwargs"])
+
+ def __getstate__(self):
+ return {"args": self.args, "kwargs": self.kwargs}
 
 
 class ClusterSequence: # The super class
@@ -354,16 +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 = -1,
- alpha=0,
- func="generalized",
- normalized=True,
+ 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.
 
@@ -376,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.

src/fastjet/_generalevent.py

@@ -438,6 +438,46 @@ 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, angles: int = -1, beta=1, alpha=0, func="generalized", normalized=True,

src/fastjet/_multievent.py

@@ -191,6 +191,40 @@ 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=-1, alpha=0, func="generalized", normalized=True,