laddu_core/variables/
mass.rs1use 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#[derive(Clone, Debug, Serialize, Deserialize)]
15enum MassSource {
16 Selection(P4Selection),
17 Reaction {
18 reaction: Box<Reaction>,
19 particle: String,
20 },
21}
22
23#[derive(Clone, Debug, Serialize, Deserialize)]
25pub struct Mass {
26 source: MassSource,
27}
28
29impl Mass {
30 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 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}