jetgo 0.1.0__py3-none-any.whl

jetgo/__init__.py

@@ -0,0 +1,20 @@
+"""
+    @file:              __init__.py
+    @Author:            Maxence Larose
+
+    @Creation Date:     09/2026
+    @Last modification: 09/2026
+
+    @Description:       This file defines package-level metadata (version, author, license, maintainer) for the
+                        ``jetgo`` package. It deliberately exposes no classes: the public API lives in the
+                        subpackages, so that importing ``jetgo`` never pulls in Pythia8 or FastJet.
+"""
+
+__version__ = "0.1.0"
+
+__author__ = "Maxence Larose"
+__credits__ = ["Maxence Larose"]
+__license__ = "BSD-3-Clause"
+__maintainer__ = "Maxence Larose"
+__email__ = "maxence.larose@stonybrook.edu"
+__status__ = "Production"

jetgo/_pythia.py

@@ -0,0 +1,63 @@
+"""
+    @file:              _pythia.py
+    @Author:            Maxence Larose
+
+    @Creation Date:     06/2026
+    @Last modification: 09/2026
+
+    @Description:       This file resolves and re-exports the Pythia8 module, accounting for the different module
+                        name used by the pip-installed ``pythia8mc`` package versus a self-compiled ``pythia8``
+                        build.
+
+                        Pythia8 is an optional dependency of this package: only event generation and particle
+                        selection need it at runtime, so observables and taggers import without it. Importing
+                        this module on a machine without either build raises with an actionable message rather
+                        than a bare ModuleNotFoundError from somewhere deep in the import graph.
+"""
+
+import importlib
+
+# The module is named pythia8 when users compile it themselves, and pythia8mc when installed via pip.
+_CANDIDATE_MODULE_NAMES = ("pythia8mc", "pythia8")
+
+_NOT_INSTALLED_MESSAGE = (
+    "Pythia8 is required for event generation but is not installed. Install it with "
+    "`pip install jetgo[pythia]`, which pulls in the `pythia8mc` wheel, or compile Pythia8 yourself with its "
+    "Python interface enabled so that `import pythia8` works. Note that the `pythia8mc` wheels are Linux-only; "
+    "on another platform, a self-compiled build is the only option. The rest of jetgo, including every "
+    "observable and tagger, imports and runs without Pythia8."
+)
+
+
+class _Pythia8NotInstalled:
+    """
+    Stands in for the Pythia8 module when neither build is installed.
+
+    Importing a module must not fail merely because an optional dependency is absent, or nothing that
+    merely *mentions* Pythia8 could be imported either, including the modules whose only use of it is a
+    type annotation. Touching anything on this stand-in raises instead, so the error arrives when event
+    generation is actually attempted and names the way to fix it.
+    """
+
+    def __getattr__(self, name: str):
+        raise ModuleNotFoundError(_NOT_INSTALLED_MESSAGE)
+
+    def __bool__(self) -> bool:
+        return False
+
+
+def _resolve():
+    for name in _CANDIDATE_MODULE_NAMES:
+        try:
+            return importlib.import_module(name)
+        except ModuleNotFoundError:
+            continue
+
+    return _Pythia8NotInstalled()
+
+
+pythia8 = _resolve()
+
+#: Whether a working Pythia8 build was found. Lets callers branch on availability without having to
+#: import Pythia8 themselves, and lets the test suite skip what genuinely needs a generator.
+PYTHIA8_AVAILABLE = not isinstance(pythia8, _Pythia8NotInstalled)

jetgo/kinematics.py

@@ -0,0 +1,133 @@
+"""
+    @file:              kinematics.py
+    @Author:            Maxence Larose
+
+    @Creation Date:     07/2026
+    @Last modification: 09/2026
+
+    @Description:       This file defines the shared kinematic helpers used across the package: the ``delta_r``
+                        function, so every part of the codebase that needs a ΔR (the EEC observable, a tagger
+                        matching a jet to a parton, ...) agrees on the same definition, and the encoding of a
+                        particle's Pythia8 index and charge into a ``PseudoJet``'s user index.
+
+                        FastJet's ``PseudoJet`` carries no particle identity of its own, only a four-momentum and
+                        a single integer slot, ``user_index()``. This package uses that slot to carry two facts
+                        about the particle a ``PseudoJet`` was built from: its position in the Pythia8 event
+                        record, and whether it is electrically charged. The magnitude is the Pythia8 index and
+                        the sign is the charge, positive for charged and negative for neutral.
+
+                        That encoding works because index 0 in a Pythia8 event record is a reserved placeholder
+                        and never a real particle, so no real particle has an index whose sign is meaningless.
+                        ``ParticleSelector`` writes the encoding and observables read it back. Anyone writing
+                        their own ``Observable`` subclass should read it through the functions below rather than
+                        touching ``user_index()`` directly, since the encoding is an implementation detail that
+                        these three functions define.
+"""
+
+import math
+
+import fastjet as fj
+
+
+def delta_r(
+    a: fj.PseudoJet,
+    b: fj.PseudoJet,
+    use_pseudorapidity: bool = False
+) -> float:
+    """
+    Angular separation ΔR between two four-vectors.
+
+    Parameters
+    ----------
+    a : fj.PseudoJet
+        First four-vector.
+    b : fj.PseudoJet
+        Second four-vector.
+    use_pseudorapidity : bool, default=False
+        If True, ΔR is computed using the pseudorapidity η instead of the (true, mass-dependent) rapidity y, i.e.
+        ΔR = sqrt((Δη)² + (Δφ)²) instead of ΔR = sqrt((Δy)² + (Δφ)²).
+
+    Returns
+    -------
+    delta_r : float
+        ΔR = sqrt((Δy or Δη)² + (Δφ)²), with Δφ correctly wrapped to [0, π].
+    """
+    if use_pseudorapidity:
+        d_y = a.eta() - b.eta()
+    else:
+        d_y = a.rap() - b.rap()
+
+    d_phi = a.delta_phi_to(b)
+
+    return math.sqrt(d_y ** 2 + d_phi ** 2)
+
+
+def encode_user_index(
+    pythia_event_index: int,
+    charged: bool
+) -> int:
+    """
+    Encode a particle's Pythia8 event-record index and charge into a single ``PseudoJet`` user index.
+
+    Parameters
+    ----------
+    pythia_event_index : int
+        Position of the particle in the Pythia8 event record. Must be strictly positive: index 0 is Pythia8's
+        reserved placeholder entry and never a real particle.
+    charged : bool
+        Whether the particle is electrically charged.
+
+    Returns
+    -------
+    user_index : int
+        The index itself if the particle is charged, its negation if it is neutral.
+
+    Raises
+    ------
+    ValueError
+        If ``pythia_event_index`` is not strictly positive, which would make the sign meaningless.
+    """
+    if pythia_event_index <= 0:
+        raise ValueError(
+            f"Expected a strictly positive Pythia8 event-record index, got {pythia_event_index}. Index 0 is "
+            f"Pythia8's reserved placeholder and should never belong to a real particle, and the sign of the "
+            f"user index is reserved for the particle's charge."
+        )
+
+    return pythia_event_index if charged else -pythia_event_index
+
+
+def is_charged(particle: fj.PseudoJet) -> bool:
+    """
+    Whether a particle is electrically charged, read back from its encoded user index.
+
+    Parameters
+    ----------
+    particle : fj.PseudoJet
+        A particle whose user index was written by ``encode_user_index``, i.e. one produced by
+        ``ParticleSelector``.
+
+    Returns
+    -------
+    charged : bool
+        True if the particle is charged.
+    """
+    return particle.user_index() > 0
+
+
+def pythia_index(particle: fj.PseudoJet) -> int:
+    """
+    The particle's position in the Pythia8 event record, read back from its encoded user index.
+
+    Parameters
+    ----------
+    particle : fj.PseudoJet
+        A particle whose user index was written by ``encode_user_index``, i.e. one produced by
+        ``ParticleSelector``.
+
+    Returns
+    -------
+    index : int
+        Position of the particle in the Pythia8 event record it was selected from.
+    """
+    return abs(particle.user_index())

jetgo/observables/__init__.py

@@ -0,0 +1,12 @@
+"""
+    @file:              __init__.py
+    @Author:            Maxence Larose
+
+    @Creation Date:     06/2026
+    @Last modification: 06/2026
+
+    @Description:       This file exposes the public observables of the ``observables`` package.
+"""
+
+from .double_differential_jet_cross_section import DoubleDifferentialJetCrossSection
+from .eec import EEC

jetgo/observables/base.py

@@ -0,0 +1,156 @@
+"""
+    @file:              base.py
+    @Author:            Maxence Larose
+
+    @Creation Date:     06/2026
+    @Last modification: 07/2026
+
+    @Description:       This file defines the Observable abstract base class. It provides the shared interface and
+                        common behavior that all physics observables inherit.
+
+                        Observables accumulate a histogram during the simulation, on bin edges fixed when the
+                        observable is constructed. Those edges are not a property of the physics being
+                        simulated -- they come from outside it, from whichever measurement or analysis the run
+                        will be read against -- so every observable takes them as a required argument rather
+                        than choosing any itself.
+
+                        Saving each contribution unbinned instead would leave that choice open until analysis
+                        time, which is the more flexible design and was the original one. It does not survive
+                        contact with the run sizes this project needs: an EEC stores one entry per constituent
+                        *pair*, which grows quadratically with jet multiplicity, and a million-event window
+                        costs gigabytes of RAM and JSON to keep values that are only ever histogrammed onto one
+                        known binning anyway. Binning at fill time makes the cost independent of the run
+                        length.
+
+                        What remains downstream (see ``finalize``) is the normalization, plus whatever
+                        re-binning each observable implements on top of the saved edges. Only the
+                        cross-section does: it sums whole bins onto coarser ones, and its edges are
+                        deliberately filled fine enough to be a refinement of every binning the analysis asks
+                        for. The EEC serves only the edges it was filled on.
+"""
+
+from abc import ABC, abstractmethod
+from typing import Dict, List
+
+import fastjet as fj
+import numpy as np
+
+from .identifier import ObservableIdentifier
+
+
+class Observable(ABC):
+    """
+    Abstract base class for physics observables.
+
+    Defines the shared interface and delegates observable-specific logic to abstract methods implemented by each
+    concrete subclass.
+    """
+
+    METADATA_KEY = "metadata"
+    VALUES_KEY = "values"
+
+    @property
+    @abstractmethod
+    def identifier(self) -> ObservableIdentifier:
+        """
+        Return the identifier of the observable. The identifier should match the observable identifier defined on the
+        HEP data website.
+
+        Returns
+        -------
+        identifier : ObservableIdentifier
+            Identifier of the observable.
+        """
+
+    @abstractmethod
+    def fill(
+            self,
+            jets: List[fj.PseudoJet],
+            cluster: fj.ClusterSequence
+    ) -> None:
+        """
+        Extract this observable's contributions from one event and accumulate them into its histogram. Called
+        once per generated event inside the simulation loop. Subclasses decide which quantities to extract and
+        which bin each one lands in.
+
+        Parameters
+        ----------
+        jets : List[fj.PseudoJet]
+            Reconstructed jets for this event, as returned by the jet finder.
+        cluster : fj.ClusterSequence
+            Cluster sequence for this event, as returned by the cluster finder.
+        """
+
+    @abstractmethod
+    def scale(
+            self,
+            scale_factor: float | int
+    ) -> None:
+        """
+        Record the event-normalization scale factor, applied later, at ``finalize`` time.
+
+        Parameters
+        ----------
+        scale_factor : float
+            Factor by which every bin will eventually be multiplied (``sigmaGen / weightSum`` from the generator).
+        """
+
+    @classmethod
+    @abstractmethod
+    def finalize(cls, values, bin_edges: np.ndarray, metadata: Dict) -> np.ndarray:
+        """
+        Normalize the histogram saved by ``to_dict`` onto ``bin_edges``.
+
+        Called downstream (by ``ObservableVisualizer`` or analysis code) -- never during the simulation itself.
+        ``bin_edges`` cannot reopen the binning decision, which was made when the observable was constructed:
+        it is the caller stating which edges it believes it is reading, and each subclass either serves it from
+        the edges it was filled on or raises. Implementations must never silently return values on edges other
+        than the ones asked for.
+
+        Parameters
+        ----------
+        values
+            The saved histogram, in whatever shape this observable writes it (see each subclass's ``to_dict``).
+        bin_edges : np.ndarray
+            The bin edges the result is wanted on.
+        metadata : Dict
+            This observable's saved metadata (see ``_get_complementary_metadata``).
+
+        Returns
+        -------
+        np.ndarray
+            Normalized values, one per target bin.
+        """
+
+    @abstractmethod
+    def _get_complementary_metadata(self) -> Dict:
+        """
+        Return complementary metadata for this observable.
+
+        Returns
+        -------
+        complementary_metadata : Dict
+            Additional metadata needed for that specific observable.
+        """
+
+    @abstractmethod
+    def _get_binned_values(self):
+        """
+        Return this observable's accumulated histogram (see ``fill``), one value per bin of the edges it was
+        constructed with.
+        """
+
+    def to_dict(self) -> Dict:
+        """
+        Return this observable's accumulated histogram and metadata as a JSON-serializable dict.
+
+        Returns
+        -------
+        Dict
+            Dictionary with keys ``VALUES_KEY`` (one value per bin) and ``METADATA_KEY`` (observable-specific
+            metadata, which always includes the bin edges the values sit on).
+        """
+        return {
+            self.VALUES_KEY: self._get_binned_values(),
+            self.METADATA_KEY: self._get_complementary_metadata(),
+        }

jetgo/observables/classes.py

@@ -0,0 +1,18 @@
+"""
+    @file:              classes.py
+    @Author:            Maxence Larose
+
+    @Creation Date:     08/2026
+    @Last modification: 08/2026
+
+    @Description:       Maps each ``ObservableIdentifier`` to the ``Observable`` subclass that implements it.
+"""
+
+from .double_differential_jet_cross_section import DoubleDifferentialJetCrossSection
+from .eec import EEC
+from .identifier import ObservableIdentifier
+
+OBSERVABLE_CLASSES = {
+    ObservableIdentifier.D2SIG_DPT_DYRAP: DoubleDifferentialJetCrossSection,
+    ObservableIdentifier.EEC: EEC,
+}

jetgo/observables/double_differential_jet_cross_section.py

@@ -0,0 +1,232 @@
+"""
+    @file:              double_differential_jet_cross_section.py
+    @Author:            Maxence Larose
+
+    @Creation Date:     06/2026
+    @Last modification: 07/2026
+
+    @Description:       This file defines the DoubleDifferentialJetCrossSection observable, the double
+                        differential inclusive jet cross-section d²σ/dpT dy. Jets are accumulated into a pT
+                        histogram on the bin edges the observable is constructed with (see ``Observable``);
+                        unlike the EEC, that grid can still be coarsened downstream, by summing whole bins
+                        (see ``finalize``).
+"""
+
+from typing import Dict, List, Optional
+
+import fastjet as fj
+import numpy as np
+
+from .base import Observable
+from .identifier import ObservableIdentifier
+
+
+class DoubleDifferentialJetCrossSection(Observable):
+    """
+    Double-differential inclusive jet cross-section d²σ/dpT dy (or d²σ/dpT dη, if jets were selected in
+    pseudorapidity instead -- see ``__init__``).
+    """
+
+    RAPIDITY_INTERVAL_KEY = "rapidity_interval"
+    SCALE_FACTOR_KEY = "scale_factor"
+    # The pT grid the values were accumulated on. Always written, and what ``finalize`` re-bins from -- a file
+    # without it predates forced binning and holds individual jet pT values instead.
+    BIN_EDGES_KEY = "bin_edges"
+
+    def __init__(
+        self,
+        bin_edges: np.ndarray,
+        max_abs_rapidity: Optional[float | int] = None,
+        max_abs_pseudorapidity: Optional[float | int] = None,
+    ) -> None:
+        """
+        Constructor for the JetCrossSection class.
+
+        Parameters
+        ----------
+        bin_edges : np.ndarray
+            The jet pT grid to accumulate the cross-section on. Required, so that the file size stays
+            independent of the run length and consistent with the EEC, which has no choice in the matter (see
+            ``Observable``).
+
+            The edges must be a *refinement* of every binning the analysis will later ask for: the extractors
+            finalize over the single bin ``[pt_min, pt_max]`` while ``PowerIndexExtractor`` uses five equal
+            sub-bins, so a grid dividing both reproduces them exactly. ``finalize`` re-bins by summing whole
+            bins and raises if the requested edges do not fall on this grid.
+        max_abs_rapidity : Optional[float | int]
+            Maximum absolute rapidity |y| used in the jet selection. Exactly one of ``max_abs_rapidity``/
+            ``max_abs_pseudorapidity`` must be given -- whichever cut ``JetFinder`` was actually configured
+            with -- since the acceptance interval width (Δy or Δη = 2 * the given value, assuming a symmetric
+            acceptance around mid-(pseudo)rapidity) is what the cross-section is normalized by.
+        max_abs_pseudorapidity : Optional[float | int]
+            Maximum absolute pseudorapidity |η| used in the jet selection, if that was the cut applied instead
+            of rapidity (e.g. matching a charged-particle-jet measurement).
+        """
+        if (max_abs_rapidity is None) == (max_abs_pseudorapidity is None):
+            raise ValueError(
+                "Exactly one of max_abs_rapidity or max_abs_pseudorapidity must be given, matching whichever "
+                "cut JetFinder was actually configured with."
+            )
+
+        acceptance = max_abs_rapidity if max_abs_rapidity is not None else max_abs_pseudorapidity
+        self._rapidity_interval = 2 * acceptance
+        self._scale_factor = 1.0
+        self._bin_edges = np.asarray(bin_edges, dtype=float)
+        self._counts = np.zeros(len(self._bin_edges) - 1, dtype=float)
+
+    @property
+    def identifier(self) -> ObservableIdentifier:
+        """
+        Identifier of the observable.
+
+        Returns
+        -------
+        identifier : ObservableIdentifier
+            Identifier of the observable.
+        """
+        return ObservableIdentifier.D2SIG_DPT_DYRAP
+
+    def scale(
+            self,
+            scale_factor: float | int
+    ) -> None:
+        """
+        Record the scale factor converting raw jet counts into an absolute cross-section in nb, applied later at
+        ``finalize`` time.
+
+        Parameters
+        ----------
+        scale_factor : float
+            Ratio sigmaGen / weightSum returned by the generator.
+        """
+        self._scale_factor = scale_factor
+
+    def fill(
+            self,
+            jets: List[fj.PseudoJet],
+            cluster: fj.ClusterSequence
+    ) -> None:
+        """
+        Accumulate every jet reconstructed in one event into its transverse momentum bin. Jets falling outside
+        the grid are dropped, exactly as they would be when histogramming stored pT values after the fact.
+
+        Parameters
+        ----------
+        jets : List[fj.PseudoJet]
+            Reconstructed jets for this event, as returned by the jet finder after kinematic selections.
+        cluster : fj.ClusterSequence
+            Cluster sequence for this event. Unused by this observable.
+        """
+        for jet in jets:
+            index = int(np.searchsorted(self._bin_edges, jet.pt(), side="right")) - 1
+            if 0 <= index < len(self._counts):
+                self._counts[index] += 1.0
+
+    @classmethod
+    def finalize(cls, values: List[float | int], bin_edges: np.ndarray, metadata: Dict) -> np.ndarray:
+        """
+        Sum the accumulated jet counts onto ``bin_edges`` and normalize into d²σ/dpT dy, by dividing by each
+        bin's width in pT and by the total rapidity interval Δy.
+
+        Parameters
+        ----------
+        values : List[float]
+            The jet count per bin of the grid the observable was filled on.
+        bin_edges : np.ndarray
+            The bin edges the result is wanted on. They must fall on that grid, so that each requested bin is a
+            whole number of grid bins.
+        metadata : Dict
+            This observable's saved metadata, must contain ``RAPIDITY_INTERVAL_KEY``, ``SCALE_FACTOR_KEY`` and
+            ``BIN_EDGES_KEY``.
+
+        Returns
+        -------
+        np.ndarray
+            d²σ/dpT dy in nb/GeV, one value per target bin.
+
+        Raises
+        ------
+        ValueError
+            If the requested edges do not lie on the filled grid, so the requested bins cannot be formed by
+            summing whole grid bins.
+        """
+        bin_edges = np.asarray(bin_edges, dtype=float)
+
+        if cls.BIN_EDGES_KEY not in metadata:
+            raise ValueError(
+                "This cross section holds individual jet pT values rather than a pT histogram, which no "
+                "version of this observable still produces. Regenerate the simulation, which will accumulate "
+                "it directly on a grid the requested edges can be summed from."
+            )
+
+        counts = cls._rebin(np.asarray(values, dtype=float),
+                            np.asarray(metadata[cls.BIN_EDGES_KEY], dtype=float), bin_edges)
+
+        bin_widths = np.diff(bin_edges)
+        rapidity_interval = metadata[cls.RAPIDITY_INTERVAL_KEY]
+        scale_factor = metadata[cls.SCALE_FACTOR_KEY]
+
+        return counts * scale_factor / (bin_widths * rapidity_interval)
+
+    @staticmethod
+    def _rebin(counts: np.ndarray, filled_edges: np.ndarray, requested_edges: np.ndarray) -> np.ndarray:
+        """
+        Sum a histogram's bins onto coarser edges that lie on its own grid.
+
+        Parameters
+        ----------
+        counts : np.ndarray
+            Counts on ``filled_edges``.
+        filled_edges : np.ndarray
+            The grid the observable was filled on.
+        requested_edges : np.ndarray
+            The coarser edges wanted, each of which must coincide with one of ``filled_edges``.
+
+        Returns
+        -------
+        np.ndarray
+            Counts on ``requested_edges``.
+
+        Raises
+        ------
+        ValueError
+            If any requested edge does not lie on the filled grid.
+        """
+        indices = []
+        for edge in requested_edges:
+            match = np.flatnonzero(np.isclose(filled_edges, edge, rtol=0.0, atol=1e-9))
+            if match.size == 0:
+                raise ValueError(
+                    f"Requested bin edge {edge:.6g} does not lie on the grid this cross section was filled on "
+                    f"({len(filled_edges) - 1} bins spanning [{filled_edges[0]:.6g}, {filled_edges[-1]:.6g}]). "
+                    f"The individual jet pT values are no longer there, so this bin cannot be formed."
+                )
+            indices.append(int(match[0]))
+
+        return np.array([counts[lo:hi].sum() for lo, hi in zip(indices[:-1], indices[1:])], dtype=float)
+
+    def _get_complementary_metadata(self) -> dict:
+        """
+        Return complementary metadata.
+
+        Returns
+        -------
+        complementary_metadata : dict
+            Additional metadata.
+        """
+        return {
+            self.RAPIDITY_INTERVAL_KEY: self._rapidity_interval,
+            self.SCALE_FACTOR_KEY: self._scale_factor,
+            self.BIN_EDGES_KEY: self._bin_edges.tolist(),
+        }
+
+    def _get_binned_values(self) -> List[float]:
+        """
+        Return the jet count accumulated in each pT bin so far.
+
+        Returns
+        -------
+        List[float]
+            The jet count per pT bin.
+        """
+        return self._counts.tolist()