gammalooprs/observables/clustering/
mod.rs

1mod genkt;
2mod types;
3
4use super::events::GenericEvent;
5use crate::utils::FloatLike;
6use bincode_trait_derive::{Decode, Encode};
7use schemars::JsonSchema;
8use serde::{Deserialize, Serialize};
9use smallvec::SmallVec;
10
11pub use types::{ClusteringResult, Jet};
12
13#[derive(
14    Debug, Clone, Copy, Default, Serialize, Deserialize, Encode, Decode, PartialEq, Eq, JsonSchema,
15)]
16#[serde(rename_all = "snake_case")]
17#[cfg_attr(feature = "python_api", pyo3::pyclass(from_py_object))]
18pub enum JetAlgorithm {
19    Kt,
20    CambridgeAachen,
21    #[default]
22    AntiKt,
23}
24
25#[derive(Debug, Clone, Default)]
26pub struct JetClustering {
27    algorithm: JetAlgorithm,
28    clustered_pdgs: Vec<isize>,
29    d_r: f64,
30    min_jpt: f64,
31}
32
33impl JetClustering {
34    pub fn new(
35        algorithm: JetAlgorithm,
36        d_r: f64,
37        min_jpt: f64,
38        clustered_pdgs: Vec<isize>,
39    ) -> Self {
40        Self {
41            algorithm,
42            clustered_pdgs,
43            d_r,
44            min_jpt,
45        }
46    }
47
48    pub fn process_event<T: FloatLike>(&self, event: &GenericEvent<T>) -> ClusteringResult<T> {
49        let candidates = event
50            .cut_info
51            .particle_pdgs
52            .1
53            .iter()
54            .copied()
55            .zip(event.kinematic_configuration.1.iter().cloned())
56            .enumerate()
57            .filter(|(_, (pdg, _))| is_qcd_like_jet_pdg(&self.clustered_pdgs, *pdg))
58            .map(|(index, (_pdg, momentum))| types::PseudoJet::from_momentum(index, momentum))
59            .collect::<SmallVec<[_; 8]>>();
60        self.cluster_candidates(candidates)
61    }
62
63    pub fn cluster_momenta<T: FloatLike>(
64        &self,
65        momenta: &[crate::momentum::FourMomentum<crate::utils::F<T>>],
66    ) -> ClusteringResult<T> {
67        let candidates = momenta
68            .iter()
69            .cloned()
70            .enumerate()
71            .map(|(index, momentum)| types::PseudoJet::from_momentum(index, momentum))
72            .collect::<SmallVec<[_; 8]>>();
73        self.cluster_candidates(candidates)
74    }
75
76    fn cluster_candidates<T: FloatLike>(
77        &self,
78        candidates: SmallVec<[types::PseudoJet<T>; 8]>,
79    ) -> ClusteringResult<T> {
80        genkt::cluster_candidates(self.algorithm, self.d_r, self.min_jpt, candidates)
81    }
82}
83
84fn is_qcd_like_jet_pdg(clustered_pdgs: &[isize], pdg: isize) -> bool {
85    clustered_pdgs.binary_search(&pdg).is_ok()
86}
gammalooprs/observables/clustering/mod.rs

gammalooprs/observables/clustering/
mod.rs
