feat: expose contrib::Njettiness (#336)

lgray · pre-commit-ci[bot] · web-flow · commit 51621b52b4d9 · 2025-03-21T18:47:10.000-05:00
* expose contrib::Njettiness * get shapes right and add test * lint * address first review comments * improve docstring * make njettiness work on input jet constituents instead * njettiness test better demonstrates usage * typo * add in dask dispatch * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * update tests with a real jet and expected taus * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci --------- Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>
diff --git a/src/_ext.cpp b/src/_ext.cpp
@@ -17,6 +17,7 @@
 #include <fastjet/PseudoJet.hh>
 #include <fastjet/contrib/EnergyCorrelator.hh>
 #include <fastjet/contrib/LundGenerator.hh>
+#include <fastjet/contrib/Njettiness.hh>
 #include <fastjet/contrib/SoftDrop.hh>
 
 #include <pybind11/numpy.h>
@@ -28,6 +29,58 @@ namespace fj = fastjet;
 namespace py = pybind11;
 using namespace pybind11::literals;
 
+// adapted from
+// https://github.com/cms-svj/SVJProduction/blob/Run3/interface/NjettinessHelper.h
+namespace njettiness {
+enum MeasureDefinition_t {
+  NormalizedMeasure = 0,     // (beta,R0)
+  UnnormalizedMeasure,       // (beta)
+  OriginalGeometricMeasure,  // (beta)
+  NormalizedCutoffMeasure,   // (beta,R0,Rcutoff)
+  UnnormalizedCutoffMeasure, // (beta,Rcutoff)
+  GeometricCutoffMeasure,    // (beta,Rcutoff)
+  N_MEASURE_DEFINITIONS
+};
+enum AxesDefinition_t {
+  KT_Axes = 0,
+  CA_Axes,
+  AntiKT_Axes, // (axAxesR0)
+  WTA_KT_Axes,
+  WTA_CA_Axes,
+  Manual_Axes,
+  OnePass_KT_Axes,
+  OnePass_CA_Axes,
+  OnePass_AntiKT_Axes, // (axAxesR0)
+  OnePass_WTA_KT_Axes,
+  OnePass_WTA_CA_Axes,
+  OnePass_Manual_Axes,
+  MultiPass_Axes,
+  N_AXES_DEFINITIONS
+};
+const std::unordered_map<std::string, MeasureDefinition_t>
+    measure_def_names_to_enum = {
+        {"NormalizedMeasure", NormalizedMeasure},
+        {"UnnormalizedMeasure", UnnormalizedMeasure},
+        {"OriginalGeometricMeasure", OriginalGeometricMeasure},
+        {"NormalizedCutoffMeasure", NormalizedCutoffMeasure},
+        {"UnnormalizedCutoffMeasure", UnnormalizedCutoffMeasure},
+        {"GeometricCutoffMeasure", GeometricCutoffMeasure}};
+const std::unordered_map<std::string, AxesDefinition_t> axis_def_names_to_enum =
+    {{"KT_Axes", KT_Axes},
+     {"CA_Axes", CA_Axes},
+     {"AntiKT_Axes", AntiKT_Axes},
+     {"WTA_KT_Axes", WTA_KT_Axes},
+     {"WTA_CA_Axes", WTA_CA_Axes},
+     {"Manual_Axes", Manual_Axes},
+     {"OnePass_KT_Axes", OnePass_KT_Axes},
+     {"OnePass_CA_Axes", OnePass_CA_Axes},
+     {"OnePass_AntiKT_Axes", OnePass_AntiKT_Axes},
+     {"OnePass_WTA_KT_Axes", OnePass_WTA_KT_Axes},
+     {"OnePass_WTA_CA_Axes", OnePass_WTA_CA_Axes},
+     {"OnePass_Manual_Axes", OnePass_Manual_Axes},
+     {"MultiPass_Axes", MultiPass_Axes}};
+} // namespace njettiness
+
 typedef struct {
   PyObject_HEAD void *ptr;
   void *ty;
@@ -2315,6 +2368,93 @@ PYBIND11_MODULE(_ext, m) {
           None.
         Returns:
           pt, eta, phi, m of inclusive jets.
+      )pbdoc")
+    .def("to_numpy_njettiness",
+      [](
+         const output_wrapper ow,
+         const std::string& measure_definition,
+         const std::string& axes_definition,
+         const std::vector<unsigned int>& njets,
+         const double beta,
+         const double R0,
+         const double Rcutoff,
+         const int nPass,
+         const double akAxesR0
+      ) {
+        auto maybe_measdef = njettiness::measure_def_names_to_enum.find(measure_definition);
+        const auto measdefenum = maybe_measdef == njettiness::measure_def_names_to_enum.end() ? njettiness::NormalizedMeasure : maybe_measdef->second;
+
+        auto maybe_axesdef = njettiness::axis_def_names_to_enum.find(axes_definition);
+        const auto axesdefenum = maybe_axesdef == njettiness::axis_def_names_to_enum.end() ? njettiness::KT_Axes : maybe_axesdef->second;
+
+        // Get the measure definition
+        fastjet::contrib::NormalizedMeasure          normalizedMeasure        (beta, R0);
+        fastjet::contrib::UnnormalizedMeasure        unnormalizedMeasure      (beta);
+        fastjet::contrib::OriginalGeometricMeasure   geometricMeasure         (beta);
+        fastjet::contrib::NormalizedCutoffMeasure    normalizedCutoffMeasure  (beta, R0, Rcutoff);
+        fastjet::contrib::UnnormalizedCutoffMeasure  unnormalizedCutoffMeasure(beta, Rcutoff);
+
+        fastjet::contrib::MeasureDefinition const * measureDef = 0;
+        switch ( measdefenum ) {
+          case njettiness::UnnormalizedMeasure         : measureDef = &unnormalizedMeasure; break;
+          case njettiness::OriginalGeometricMeasure    : measureDef = &geometricMeasure; break;
+          case njettiness::NormalizedCutoffMeasure     : measureDef = &normalizedCutoffMeasure; break;
+          case njettiness::UnnormalizedCutoffMeasure   : measureDef = &unnormalizedCutoffMeasure; break;
+          case njettiness::NormalizedMeasure : default : measureDef = &normalizedMeasure; break;
+        }
+
+        // Get the axes definition
+        fastjet::contrib::KT_Axes             kt_axes;
+        fastjet::contrib::CA_Axes             ca_axes;
+        fastjet::contrib::AntiKT_Axes         antikt_axes        (akAxesR0);
+        fastjet::contrib::WTA_KT_Axes         wta_kt_axes;
+        fastjet::contrib::WTA_CA_Axes         wta_ca_axes;
+        fastjet::contrib::OnePass_KT_Axes     onepass_kt_axes;
+        fastjet::contrib::OnePass_CA_Axes     onepass_ca_axes;
+        fastjet::contrib::OnePass_AntiKT_Axes onepass_antikt_axes(akAxesR0);
+        fastjet::contrib::OnePass_WTA_KT_Axes onepass_wta_kt_axes;
+        fastjet::contrib::OnePass_WTA_CA_Axes onepass_wta_ca_axes;
+        fastjet::contrib::MultiPass_Axes      multipass_axes     (nPass);
+
+        fastjet::contrib::AxesDefinition const * axesDef = 0;
+        switch ( axesdefenum ) {
+          case  njettiness::KT_Axes : default   : axesDef = &kt_axes; break;
+          case  njettiness::CA_Axes             : axesDef = &ca_axes; break;
+          case  njettiness::AntiKT_Axes         : axesDef = &antikt_axes; break;
+          case  njettiness::WTA_KT_Axes         : axesDef = &wta_kt_axes; break;
+          case  njettiness::WTA_CA_Axes         : axesDef = &wta_ca_axes; break;
+          case  njettiness::OnePass_KT_Axes     : axesDef = &onepass_kt_axes; break;
+          case  njettiness::OnePass_CA_Axes     : axesDef = &onepass_ca_axes; break;
+          case  njettiness::OnePass_AntiKT_Axes : axesDef = &onepass_antikt_axes; break;
+          case  njettiness::OnePass_WTA_KT_Axes : axesDef = &onepass_wta_kt_axes; break;
+          case  njettiness::OnePass_WTA_CA_Axes : axesDef = &onepass_wta_ca_axes; break;
+          case  njettiness::MultiPass_Axes      : axesDef = &multipass_axes; break;
+        }
+
+        auto routine = std::make_shared<fastjet::contrib::Njettiness>(*axesDef, *measureDef);
+
+        const auto& constituents = ow.parts;
+        std::vector<double> taus;
+
+        for (size_t i = 0; i < constituents.size(); ++i) {
+            for(size_t k = 0; k < njets.size(); ++k) {
+              auto tau = routine->getTau(njets[k], *constituents[i]);
+              taus.push_back(tau);
+            }
+        }
+
+        auto taus_out = py::array(taus.size(), taus.data());
+        taus_out.resize({taus.size()/njets.size(), njets.size()});
+
+        return std::make_tuple(
+          taus_out
+        );
+      }, R"pbdoc(
+        Calculates njettiness values from inputs and converts them to numpy arrays.
+        Args:
+          None.
+        Returns:
+          the <njets>-tuple of njettiness values for all found jets, and their offsets
       )pbdoc");
   py::class_<ClusterSequence>(m, "ClusterSequence")
       .def(py::init<const std::vector<PseudoJet> &, const JetDefinition &,
diff --git a/src/fastjet/_generalevent.py b/src/fastjet/_generalevent.py
@@ -5,6 +5,8 @@
 
 import fastjet._ext  # noqa: F401, E402
 
+_default_taus_njettiness = [1, 2, 3, 4]
+
 
 class _classgeneralevent:
     def __init__(self, data, jetdef):
@@ -708,6 +710,54 @@ def exclusive_jets_softdrop_grooming(
         )
         return res
 
+    def njettiness(
+        self,
+        measure_definition="NormalizedMeasure",
+        axes_definition="OnePass_KT_Axes",
+        njets=_default_taus_njettiness,
+        beta=1.0,
+        R0=0.8,
+        Rcutoff=None,
+        nPass=None,
+        akAxesR0=None,
+    ):
+        if isinstance(njets, (int, float)):
+            njets = [njets]
+        if len(njets) == 0:
+            raise ValueError("Must provide at least one njets!")
+        if any(njet <= 0 for njet in njets):
+            raise ValueError("Requested njets must be > 0!")
+
+        double_max = 999.0
+        int_max = 999
+
+        self._out = []
+        self._input_flag = 0
+        for i in range(len(self._clusterable_level)):
+            np_results = self._results[i].to_numpy_njettiness(
+                measure_definition,
+                axes_definition,
+                njets,
+                beta,
+                R0,
+                Rcutoff or double_max,
+                nPass or int_max,
+                akAxesR0 or double_max,
+            )
+            self._out.append(
+                ak.Array(
+                    np_results[0],
+                    behavior=self.data.behavior,
+                    attrs=self.data.attrs,
+                ),
+            )
+        res = ak.Array(
+            self._replace_multi(),
+            behavior=self.data.behavior,
+            attrs=self.data.attrs,
+        )
+        return res
+
     def exclusive_jets_energy_correlator(
         self,
         njets=1,
diff --git a/src/fastjet/_multievent.py b/src/fastjet/_multievent.py
@@ -5,6 +5,8 @@
 
 import fastjet._ext  # noqa: F401, E402
 
+_default_taus_njettiness = [1, 2, 3, 4]
+
 
 class _classmultievent:
     def __init__(self, data, jetdef):
@@ -296,6 +298,42 @@ def exclusive_jets_softdrop_grooming(
         )
         return out
 
+    def njettiness(
+        self,
+        measure_definition="NormalizedMeasure",
+        axes_definition="OnePass_KT_Axes",
+        njets=_default_taus_njettiness,
+        beta=1.0,
+        R0=0.8,
+        Rcutoff=None,
+        nPass=None,
+        akAxesR0=None,
+    ):
+        if isinstance(njets, (int, float)):
+            njets = [njets]
+        if len(njets) == 0:
+            raise ValueError("Must provide at least one njets!")
+        if any(njet <= 0 for njet in njets):
+            raise ValueError("Requested njets must be > 0!")
+
+        double_max = 999.0
+        int_max = 999
+
+        np_results = self._results.to_numpy_njettiness(
+            measure_definition,
+            axes_definition,
+            njets,
+            beta,
+            R0,
+            Rcutoff or double_max,
+            nPass or int_max,
+            akAxesR0 or double_max,
+        )
+        out = ak.Array(
+            ak.contents.NumpyArray(np_results[0]),
+        )
+        return out
+
     def exclusive_jets_energy_correlator(
         self,
         njets=1,
diff --git a/src/fastjet/_pyjet.py b/src/fastjet/_pyjet.py
@@ -10,6 +10,9 @@
 __all__ = ("__version__",)
 
 
+_default_taus_njettiness = [1, 2, 3, 4]
+
+
 class AwkwardClusterSequence(ClusterSequence):
     def __init__(self, data, jetdef):
         if not isinstance(data, ak.Array):
@@ -167,6 +170,28 @@ def exclusive_jets_softdrop_grooming(
             mu_cut,
         )
 
+    def njettiness(
+        self,
+        measure_definition="NormalizedMeasure",
+        axes_definition="OnePass_KT_Axes",
+        njets=_default_taus_njettiness,
+        beta=1.0,
+        R0=0.8,
+        Rcutoff=None,
+        nPass=None,
+        akAxesR0=None,
+    ):
+        return self._internalrep.njettiness(
+            measure_definition=measure_definition,
+            axes_definition=axes_definition,
+            njets=njets,
+            beta=beta,
+            R0=R0,
+            Rcutoff=Rcutoff,
+            nPass=nPass,
+            akAxesR0=akAxesR0,
+        )
+
     def exclusive_jets_energy_correlator(
         self,
         njets=1,
@@ -492,6 +517,30 @@ def exclusive_jets_softdrop_grooming(
             mu_cut=mu_cut,
         )
 
+    def njettiness(
+        self,
+        measure_definition="NormalizedMeasure",
+        axes_definition="OnePass_KT_Axes",
+        njets=_default_taus_njettiness,
+        beta=1.0,
+        R0=0.8,
+        Rcutoff=None,
+        nPass=None,
+        akAxesR0=None,
+    ):
+        return _dak_dispatch(
+            self,
+            "njettiness",
+            measure_definition=measure_definition,
+            axes_definition=axes_definition,
+            njets=njets,
+            beta=beta,
+            R0=R0,
+            Rcutoff=Rcutoff,
+            nPass=nPass,
+            akAxesR0=akAxesR0,
+        )
+
     def exclusive_jets_energy_correlator(
         self,
         njets=1,
diff --git a/src/fastjet/_singleevent.py b/src/fastjet/_singleevent.py
@@ -5,6 +5,8 @@
 
 import fastjet._ext  # noqa: F401, E402
 
+_default_taus_njettiness = [1, 2, 3, 4]
+
 
 class _classsingleevent:
     def __init__(self, data, jetdef):
@@ -272,6 +274,42 @@ def exclusive_jets_softdrop_grooming(
         )
         return out[0]
 
+    def njettiness(
+        self,
+        measure_definition="NormalizedMeasure",
+        axes_definition="OnePass_KT_Axes",
+        njets=_default_taus_njettiness,
+        beta=1.0,
+        R0=0.8,
+        Rcutoff=None,
+        nPass=None,
+        akAxesR0=None,
+    ):
+        if isinstance(njets, (int, float)):
+            njets = [njets]
+        if len(njets) == 0:
+            raise ValueError("Must provide at least one njets!")
+        if any(njet <= 0 for njet in njets):
+            raise ValueError("Requested njets must be > 0!")
+
+        double_max = 999.0
+        int_max = 999
+
+        np_results = self._results.to_numpy_njettiness(
+            measure_definition,
+            axes_definition,
+            njets,
+            beta,
+            R0,
+            Rcutoff or double_max,
+            nPass or int_max,
+            akAxesR0 or double_max,
+        )
+        out = ak.Array(
+            ak.contents.NumpyArray(np_results[0]),
+        )
+        return out
+
     def exclusive_jets_energy_correlator(
         self,
         njets=1,
diff --git a/tests/test_002-exclusive_jets.py b/tests/test_002-exclusive_jets.py
