API reference

The NumPy entry points work with the base installation. Tensor, history, and data helpers need the relevant optional dependencies.

Clustering

flashjet.cluster(p4, mask=None, R=0.4, algorithm='antikt', p=None, backend='auto', validate=True, decode=True)

Cluster particles with a generalized-kt sequential recombination.

Parameters:

• p4 – (n, 4) numpy array for a single event (returns the NumPy reference ClusterSequenceRef), or a (B, N, 4) torch tensor (px, py, pz, E) for a padded batch.

• mask – (B, N) bool tensor for the batched path (default: all true).

• R – jet radius.

• algorithm – ‘antikt’ | ‘kt’ | ‘cambridge’ (ignored if p given).

• p – generalized-kt exponent overriding algorithm.

• backend – ‘auto’ | ‘triton’ | ‘triton-large’ | ‘cpu’ | ‘torch’. ‘auto’ picks the fused register kernel (CUDA, N <= 16, where it measures fastest), then the scratch NN-array kernel (CUDA, N <= 16384), then the compiled C++ CPU kernel for CPU tensors, then torch. Set FLASHJET_TUNE=1 to let the triton-large backend autotune its launch params once per GPU model (see flashjet/tune.py; persisted, reproducible afterwards), and FLASHJET_COMPILE_DECODE=1 to torch.compile its history decode (see history.py; helps large direct calls at small N).

• validate – reject non-finite four-momenta in masked-in slots up front (a single inf row can otherwise abort the CUDA context in the decode’s scatter). Costs one device sync per call; pass False in hot loops with trusted inputs.

• decode – triton-large only – when False, skip the per-particle jet_idx pointer-jump decode (ClusterOutput.jet_idx is then None) for callers that only need the merge history / substructure features (ClusterOutput.splitting_scales). Other backends compute jet_idx in-kernel and ignore this flag.

Returns:

ClusterSequenceRef (single event) or ClusterOutput (batch).

class flashjet.ClusterOutput(jet_idx, n_jets, hist_p1, hist_p2, hist_child, hist_d, mask=None)

Batched clustering result (all tensors live on the input device).

jet_idx

(B, N) int64, jet index per input particle in beam-merge order, -1 for padding (or particles below ptmin filters applied later).

Type:

object

n_jets

(B,) int64.

Type:

object

hist_p1/p2/child

(B, N) int64 merge tree (pseudojet ids; initial particles are 0..n-1 in mask order, -1 = beam, -2 = pad).

hist_d

(B, N) float, d_min of each merge step.

Type:

object

mask

(B, N) bool, True for real particles (needed by the exclusive-jet / substructure decoders that map slots -> initial pseudojet ids). None if the output was built without it (older callers); the exclusive/grooming helpers then require an explicit mask=.

Type:

object

exclusive_jets(n_jets=None, d_cut=None, mask=None)

Exclusive-jet particle assignment (undo the sequence’s last merges).

Pass exactly one of n_jets (leave this many jets) or d_cut (undo every pair-merge with d >= d_cut, the exclusive-y_cut form on the stored d = min(w_i, w_j) * dR^2 / R^2). Returns (excl_idx (B, N) int64, n_excl (B,)); excl_idx feeds jets_p4() directly. This is the classic kt exclusive mode; on kt (p=1) d is monotonic so both cuts are exact.

groomed_jets(p4, R, z_cut=0.1, beta=0.0, mu=None, n_jets_max=None, mask=None)

Soft-drop / mass-drop grooming of each jet (declustering tagger).

Walks each jet down the harder branch, dropping soft wide-angle radiation until a split satisfies z > z_cut*(dR/R)**beta (beta=0 is mMDT); pass mu= to also require max(m_i,m_j) < mu*m (Mass-Drop Tagger). Returns a dict with groomed_p4 (B, J, 4), tagged (B, J) bool, and the passing split’s z / dR / mu_split / n_drop; jets in beam-merge order aligned with jets_p4(). See history.groom_from_history.

jets_p4(p4, n_jets_max=None)

Differentiable jet four-momenta via scatter_add of constituents.

Parameters:

p4 – the (B, N, 4) input tensor (may require grad).

Returns:

(B, J, 4) tensor, J = n_jets.max() (or n_jets_max), zero-padded, jets in beam-merge order (use sort_jets_by_pt to reorder). Constituents of jets >= n_jets_max are dropped, not folded into the last slot.

Note

Recovery uses scatter_add, whose CUDA kernel accumulates in a nondeterministic order: unlike the clustering itself (the kernels are bitwise-deterministic), the returned momenta can vary run to run at the ulp level. Enable torch.use_deterministic_algorithms(True) if you need bitwise-stable jet four-momenta.

lund_coordinates(p4, R, n_jets_max=None, mask=None)

Per-jet, per-split Lund-plane coordinates (B, J, S, 6).

Channels: (z, dR, kt, ln 1/dR, ln kt, d), one row per de-clustering split in the same order as splitting_scales() (slot 0 = widest split), aligned with jets_p4(). Extends splitting_scales with the full Lund inputs; see history.lund_coordinates_from_history.

mass_drop(p4, R, mu=0.67, y_cut=0.09, n_jets_max=None, mask=None)

Original Mass-Drop Tagger (Butterworth-Davison-Rubin-Salam).

Convenience wrapper over groomed_jets with the mass-drop mu and a z_cut derived from y_cut (z_cut = y_cut, beta = 0); returns the same dict. Use groomed_jets(…, beta, mu) for the general soft-drop form.

sort_jets_by_pt(jets_p4)

Return (sorted_jets, order) with jets sorted pt-descending per event.

splitting_scales(n_jets_max=None)

Per-jet sequential-recombination splitting scales (B, J, S) float.

out[:, j, 0] is jet j’s last merge / first de-clustering split (the d_12 scale), [:, j, 1] the next (d_23), …, zero-padded past each jet’s merge count; jets are in beam-merge order, ALIGNED with jets_p4 (so per-jet features concatenate – if you sort_jets_by_pt the jets, gather these with the same order). For the kt algorithm these are the exclusive d_12 >= d_23 >= … scales (kt-splitting / Lund inputs); for C/A and anti-kt the entries are the de-clustering sequence but not value-sorted. See history.splitting_scales_from_history for details.

Reads only the merge history, so it works whether or not the per-particle decode ran (cluster(…, decode=False) is fine).

NumPy reference

flashjet.cluster_event(p4, R=0.4, p=-1.0)

Cluster a single event; p4 is (n, 4) with columns px, py, pz, E.

class flashjet.ClusterSequenceRef(p4, history=<factory>, beam_jets=<factory>, n_initial=0)

Result of clustering one event.

constituents(pseudojet_id)

Indices of the initial particles contained in a pseudojet.

inclusive_jets(ptmin=0.0)

Jet 4-momenta (pt-sorted, descending) with pt > ptmin.

jet_constituents(ptmin=0.0)

Constituent index lists, ordered like inclusive_jets(ptmin).

History helpers

Decode per-particle jet assignments from the merge history.

The clustering kernels record, per step, the two parent pseudojet ids and the child id (or a beam merge). That tree determines the particle->jet mapping, so the kernel does not need to maintain a destination array in its hot loop (which costs global-memory traffic every step at large N). Instead we rebuild the mapping afterwards with pointer jumping: O(log N) batched gathers.

On CUDA the whole decode runs as a single Triton launch (one program per event, ping-pong scratch so each jumping round is synchronous like the eager gathers); it is bitwise-identical to the eager torch-op loop, which remains the spec and the CPU path.

Ids: initial particles are 0..n-1 in mask order, children continue from n.

flashjet.history.exclusive_jets_from_history(hist_p1, hist_p2, hist_child, hist_d, mask, n_jets=None, d_cut=None)

Exclusive-jet particle assignment: undo the last merges of the sequence.

Exactly one of n_jets (stop when this many exclusive jets remain) or d_cut (undo every pair-merge with d >= d_cut) must be given. This is the kt-family exclusive jet definition (FastJet ClusterSequence:: exclusive_jets / exclusive_jets_ycut) read off the recorded sequence: a pair-merge is kept only if it happened before the stopping point, and each particle is assigned to the highest pseudojet it reaches through kept merges. d_cut compares against the SAME d stored in hist_d (min(w_i,w_j) * dR^2 / R^2); for a p_t-style y_cut pass R**2 * d.

Returns (excl_idx (B, N) int64, n_excl (B,) int64): jet index per particle in ascending-root order, -1 for padding. For the trivial cut (n_jets == n_inclusive, or d_cut past the largest d) this reproduces the inclusive jet_idx partition (same particles grouped), though jet numbering follows root order, not beam-merge order.

flashjet.history.groom_from_history(hist_p1, hist_p2, hist_child, hist_d, mask, p4, R, z_cut=0.1, beta=0.0, mu=None, n_jets_max=None)

Soft-drop / mass-drop grooming by declustering each jet’s tree.

Walks every jet from its root down the HARDER (higher-pt) branch, undoing the widest split first (this is the C/A declustering picture, and is exact for any recorded tree since we follow the stored merge structure). At each node with parents i, j it tests the soft-drop condition:

z > z_cut * (dR / R)**beta ,   z = min(pt_i,pt_j)/(pt_i+pt_j)

(beta=0 is the modified Mass-Drop Tagger / mMDT). If it passes, that node is the groomed jet and the walk stops; otherwise the softer parent is dropped and the walk continues into the harder parent. A jet that declusters to a single particle without ever passing is untagged.

If mu is given, the additional mass-drop requirement max(m_i,m_j) < mu*m (m = mass of the current node) must also hold for a node to pass – the original Mass-Drop Tagger.

Returns dict of tensors (all B x J, beam-merge order, aligned with jets_p4 / splitting_scales):

groomed_p4 (B, J, 4): 4-momentum of the groomed subjet (0 if untagged)
tagged     (B, J) bool: whether a split passed the condition
z, dR, mu_split (B, J): the passing split's z, dR, and mass ratio
                        max(m_i,m_j)/m (0 where untagged)
n_drop     (B, J) int64: number of soft branches dropped before passing

flashjet.history.lund_coordinates_from_history(hist_p1, hist_p2, hist_child, hist_d, mask, p4, R, n_jets_max=None)

Per-jet, per-split Lund-plane coordinates from the merge history.

Returns (B, J, S, C) float in the SAME per-jet de-clustering order as splitting_scales_from_history (slot 0 = the jet’s first / widest split), zero-padded past each jet’s split count. Channels C = 6:

0 z      = min(pt_i, pt_j) / (pt_i + pt_j)          in (0, 0.5]
1 dR     = sqrt(dy^2 + dphi^2) of the two parents
2 kt     = min(pt_i, pt_j) * dR                     (the Lund kt)
3 ln(1/dR)
4 ln(kt)
5 d      = hist_d for that split (== splitting_scales entry; sanity tie)

Jets are aligned with jet_idx / jets_p4 (beam-merge order); a caller that sort_jets_by_pt’s the jets applies the same permutation here. pt is the transverse momentum of each parent pseudojet (E-scheme sum), recovered from the tree, so this needs the input p4 and mask (unlike splitting_scales).

flashjet.history.splitting_scales_from_history(hist_p1, hist_p2, hist_child, hist_d, n_jets_max=None)

Per-jet sequential-recombination splitting scales, in de-clustering order.

Returns (B, J, S) float (J = n_jets.max() or n_jets_max, S = the largest per-jet merge count): out[b, j, 0] is the LAST merge that formed jet j (the widest / d_12 splitting), out[b, j, 1] the next (d_23), …, zero- padded past each jet’s merge count. Jets are numbered in beam-merge order, ALIGNED with jet_idx and ClusterOutput.jets_p4 (so per-jet features concatenate); a caller that sort_jets_by_pt’s the jets must apply the same permutation here.

The value is flashjet’s merge distance d = min(w_i, w_j) * dR^2 / R^2 with w = kt^(2p). Entry 0 is the LAST merge that built the jet – equivalently the jet’s FIRST de-clustering split, the d_12 scale – entry 1 the next, etc. For the kt algorithm (p=1) d_min is monotonic, so the sequence is additionally value-sorted (d_12 >= d_23 >= …): the exclusive kt scales. For Cambridge/Aachen and anti-kt d_min is NOT monotonic (recombination can lower it), so the entries are the de-clustering sequence in merge order but not value-sorted. Recover conventional sqrt(d_ij) as (out * R**2).sqrt().

Data helpers

Pipeline I/O: ship ragged CPU events to the GPU for clustering.

Real datasets are ragged (variable particles per event) and live on the CPU (awkward arrays from uproot/parquet). flashjet.cluster() wants padded (B, N, 4) tensors + mask on the GPU. This module closes that gap without per-event Python loops or synchronous copies:

• collate(): vectorized ragged -> padded+mask (one fancy-indexing scatter), optionally truncating to the hardest-pt particles, optionally writing into preallocated pinned buffers.

• to_gpu_batches(): a generator that double-buffers pinned staging memory and issues async H2D copies on a side stream, so collation of batch k+1 and the copy overlap with GPU compute on batch k.

The transfer itself is cheap (a 6000-particle event is ~100 kB ~ 8 us over PCIe vs ~3 ms to cluster); what this module avoids is the Python-loop padding and blocking copies that would otherwise dominate.

flashjet.data.collate(events, n_max=None, truncate='pt', out=None)

Vectorized ragged -> (p4 (B, N, 4) float32, mask (B, N) bool) tensors.

Parameters:

• events – awkward Array (px/py/pz/E fields) or sequence of (n_i, 4) arrays (px, py, pz, E columns).

• n_max – pad/truncate width (default: longest event in the batch).

• truncate – ‘pt’ keeps the hardest particles of oversized events, ‘first’ keeps the leading slice, ‘error’ raises.

• out – optional (p4, mask) preallocated tensors (e.g. pinned) to fill; must be at least (B, n_max, 4) / (B, n_max).

Returns (p4, mask) torch tensors (views into out when given).

flashjet.data.gpu_batch_ready(batch)

Make the current stream wait for a batch yielded by to_gpu_batches and return (p4, mask). Call this right before using the tensors.

Holding several un-readied batches (or dropping one) is data-safe — the yielded event object is re-recorded two batches later, but only ever to a LATER point on the same copy stream, so the cost is over-synchronization, never stale tensors.

flashjet.data.to_gpu_batches(events, batch_size, n_max=None, device='cuda', truncate='pt')

Yield (p4, mask) GPU batches from a ragged dataset, overlapping collation and H2D copies with downstream GPU compute.

The ragged dataset is flattened ONCE up front (the only awkward-array work); each batch is then a pure numpy scatter from precomputed offsets into a ring of two pinned staging buffers, copied on a dedicated stream. While the caller clusters batch k on the default stream, batch k+1 is collated and copied asynchronously.