Skip to main content

laddu_core/variables/
mass.rs

1use std::fmt::Display;
2
3use serde::{Deserialize, Serialize};
4
5use super::{format_names, IntoP4Selection, P4Selection, Variable};
6use crate::{
7    data::{DatasetMetadata, EventLike},
8    reaction::Reaction,
9    vectors::Vec4,
10    LadduResult,
11};
12
13/// Source for a mass variable.
14#[derive(Clone, Debug, Serialize, Deserialize)]
15enum MassSource {
16    Selection(P4Selection),
17    Reaction {
18        reaction: Box<Reaction>,
19        particle: String,
20    },
21}
22
23/// A struct for obtaining the invariant mass of a selected or reaction-defined particle.
24#[derive(Clone, Debug, Serialize, Deserialize)]
25pub struct Mass {
26    source: MassSource,
27}
28
29impl Mass {
30    /// Create a new [`Mass`] from the sum of the four-momenta identified by `constituents` in the
31    /// [`EventData`](crate::data::EventData)'s `p4s` field.
32    pub fn new<C>(constituents: C) -> Self
33    where
34        C: IntoP4Selection,
35    {
36        Self {
37            source: MassSource::Selection(constituents.into_selection()),
38        }
39    }
40
41    /// Create a new [`Mass`] for a particle resolved through a [`Reaction`].
42    pub fn from_reaction(reaction: Reaction, particle: impl Into<String>) -> Self {
43        Self {
44            source: MassSource::Reaction {
45                reaction: Box::new(reaction),
46                particle: particle.into(),
47            },
48        }
49    }
50}
51
52impl Display for Mass {
53    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
54        match &self.source {
55            MassSource::Selection(constituents) => {
56                write!(
57                    f,
58                    "Mass(constituents=[{}])",
59                    format_names(constituents.names())
60                )
61            }
62            MassSource::Reaction { particle, .. } => write!(f, "Mass(particle={})", particle),
63        }
64    }
65}
66
67#[typetag::serde]
68impl Variable for Mass {
69    fn bind(&mut self, metadata: &DatasetMetadata) -> LadduResult<()> {
70        match &mut self.source {
71            MassSource::Selection(constituents) => constituents.bind(metadata),
72            MassSource::Reaction { .. } => Ok(()),
73        }
74    }
75
76    fn value(&self, event: &dyn EventLike) -> f64 {
77        match &self.source {
78            MassSource::Selection(constituents) => constituents
79                .indices()
80                .iter()
81                .map(|index| event.p4_at(*index))
82                .sum::<Vec4>()
83                .m(),
84            MassSource::Reaction { reaction, particle } => reaction
85                .p4(event, particle)
86                .unwrap_or_else(|err| panic!("failed to evaluate reaction mass: {err}"))
87                .m(),
88        }
89    }
90}