laddu_core/utils/
variables.rs

1use dyn_clone::DynClone;
2use serde::{Deserialize, Serialize};
3use std::{
4    fmt::{Debug, Display},
5    sync::Arc,
6};
7
8#[cfg(feature = "rayon")]
9use rayon::prelude::*;
10
11#[cfg(feature = "mpi")]
12use crate::mpi::LadduMPI;
13use crate::{
14    data::{Dataset, Event},
15    utils::{
16        enums::{Channel, Frame},
17        vectors::Vec3,
18    },
19    Float, LadduError,
20};
21#[cfg(feature = "mpi")]
22use mpi::{datatype::PartitionMut, topology::SimpleCommunicator, traits::*};
23
24/// Standard methods for extracting some value out of an [`Event`].
25#[typetag::serde(tag = "type")]
26pub trait Variable: DynClone + Send + Sync + Debug + Display {
27    /// This method takes an [`Event`] and extracts a single value (like the mass of a particle).
28    fn value(&self, event: &Event) -> Float;
29
30    /// This method distributes the [`Variable::value`] method over each [`Event`] in a
31    /// [`Dataset`] (non-MPI version).
32    ///
33    /// # Notes
34    ///
35    /// This method is not intended to be called in analyses but rather in writing methods
36    /// that have `mpi`-feature-gated versions. Most users should just call [`Variable::value_on`] instead.
37    fn value_on_local(&self, dataset: &Arc<Dataset>) -> Vec<Float> {
38        #[cfg(feature = "rayon")]
39        let local_values: Vec<Float> = dataset.events.par_iter().map(|e| self.value(e)).collect();
40        #[cfg(not(feature = "rayon"))]
41        let local_values: Vec<Float> = dataset.events.iter().map(|e| self.value(e)).collect();
42        local_values
43    }
44
45    /// This method distributes the [`Variable::value`] method over each [`Event`] in a
46    /// [`Dataset`] (MPI-compatible version).
47    ///
48    /// # Notes
49    ///
50    /// This method is not intended to be called in analyses but rather in writing methods
51    /// that have `mpi`-feature-gated versions. Most users should just call [`Variable::value_on`] instead.
52    #[cfg(feature = "mpi")]
53    fn value_on_mpi(&self, dataset: &Arc<Dataset>, world: &SimpleCommunicator) -> Vec<Float> {
54        let local_weights = self.value_on_local(dataset);
55        let n_events = dataset.n_events();
56        let mut buffer: Vec<Float> = vec![0.0; n_events];
57        let (counts, displs) = world.get_counts_displs(n_events);
58        {
59            let mut partitioned_buffer = PartitionMut::new(&mut buffer, counts, displs);
60            world.all_gather_varcount_into(&local_weights, &mut partitioned_buffer);
61        }
62        buffer
63    }
64
65    /// This method distributes the [`Variable::value`] method over each [`Event`] in a
66    /// [`Dataset`].
67    fn value_on(&self, dataset: &Arc<Dataset>) -> Vec<Float> {
68        #[cfg(feature = "mpi")]
69        {
70            if let Some(world) = crate::mpi::get_world() {
71                return self.value_on_mpi(dataset, &world);
72            }
73        }
74        self.value_on_local(dataset)
75    }
76}
77dyn_clone::clone_trait_object!(Variable);
78
79fn sort_indices<T: AsRef<[usize]>>(indices: T) -> Vec<usize> {
80    let mut indices = indices.as_ref().to_vec();
81    indices.sort();
82    indices
83}
84
85fn indices_to_string<T: AsRef<[usize]>>(indices: T) -> String {
86    indices
87        .as_ref()
88        .iter()
89        .map(|n| n.to_string())
90        .collect::<Vec<_>>()
91        .join(", ")
92}
93
94/// A struct for obtaining the mass of a particle by indexing the four-momenta of an event, adding
95/// together multiple four-momenta if more than one index is given.
96#[derive(Clone, Debug, Serialize, Deserialize)]
97pub struct Mass(Vec<usize>);
98impl Mass {
99    /// Create a new [`Mass`] from the sum of the four-momenta at the given indices in the
100    /// [`Event`]'s `p4s` field.
101    pub fn new<T: AsRef<[usize]>>(constituents: T) -> Self {
102        Self(sort_indices(constituents))
103    }
104}
105impl Display for Mass {
106    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
107        write!(f, "Mass(constituents=[{}])", indices_to_string(&self.0))
108    }
109}
110#[typetag::serde]
111impl Variable for Mass {
112    fn value(&self, event: &Event) -> Float {
113        event.get_p4_sum(&self.0).m()
114    }
115}
116
117/// A struct for obtaining the $`\cos\theta`$ (cosine of the polar angle) of a decay product in
118/// a given reference frame of its parent resonance.
119#[derive(Clone, Debug, Serialize, Deserialize)]
120pub struct CosTheta {
121    beam: usize,
122    recoil: Vec<usize>,
123    daughter: Vec<usize>,
124    resonance: Vec<usize>,
125    frame: Frame,
126}
127impl Display for CosTheta {
128    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
129        write!(
130            f,
131            "CosTheta(beam={}, recoil=[{}], daughter=[{}], resonance=[{}], frame={})",
132            self.beam,
133            indices_to_string(&self.recoil),
134            indices_to_string(&self.daughter),
135            indices_to_string(&self.resonance),
136            self.frame
137        )
138    }
139}
140impl CosTheta {
141    /// Construct the angle given the four-momentum indices for each specified particle. Fields
142    /// which can take lists of more than one index will add the relevant four-momenta to make a
143    /// new particle from the constituents. See [`Frame`] for options regarding the reference
144    /// frame.
145    pub fn new<T: AsRef<[usize]>, U: AsRef<[usize]>, V: AsRef<[usize]>>(
146        beam: usize,
147        recoil: T,
148        daughter: U,
149        resonance: V,
150        frame: Frame,
151    ) -> Self {
152        Self {
153            beam,
154            recoil: recoil.as_ref().into(),
155            daughter: daughter.as_ref().into(),
156            resonance: resonance.as_ref().into(),
157            frame,
158        }
159    }
160}
161impl Default for CosTheta {
162    fn default() -> Self {
163        Self {
164            beam: 0,
165            recoil: vec![1],
166            daughter: vec![2],
167            resonance: vec![2, 3],
168            frame: Frame::Helicity,
169        }
170    }
171}
172#[typetag::serde]
173impl Variable for CosTheta {
174    fn value(&self, event: &Event) -> Float {
175        let beam = event.p4s[self.beam];
176        let recoil = event.get_p4_sum(&self.recoil);
177        let daughter = event.get_p4_sum(&self.daughter);
178        let resonance = event.get_p4_sum(&self.resonance);
179        let daughter_res = daughter.boost(&-resonance.beta());
180        match self.frame {
181            Frame::Helicity => {
182                let recoil_res = recoil.boost(&-resonance.beta());
183                let z = -recoil_res.vec3().unit();
184                let y = beam.vec3().cross(&-recoil.vec3()).unit();
185                let x = y.cross(&z);
186                let angles = Vec3::new(
187                    daughter_res.vec3().dot(&x),
188                    daughter_res.vec3().dot(&y),
189                    daughter_res.vec3().dot(&z),
190                );
191                angles.costheta()
192            }
193            Frame::GottfriedJackson => {
194                let beam_res = beam.boost(&-resonance.beta());
195                let z = beam_res.vec3().unit();
196                let y = beam.vec3().cross(&-recoil.vec3()).unit();
197                let x = y.cross(&z);
198                let angles = Vec3::new(
199                    daughter_res.vec3().dot(&x),
200                    daughter_res.vec3().dot(&y),
201                    daughter_res.vec3().dot(&z),
202                );
203                angles.costheta()
204            }
205        }
206    }
207}
208
209/// A struct for obtaining the $`\phi`$ angle (azimuthal angle) of a decay product in a given
210/// reference frame of its parent resonance.
211#[derive(Clone, Debug, Serialize, Deserialize)]
212pub struct Phi {
213    beam: usize,
214    recoil: Vec<usize>,
215    daughter: Vec<usize>,
216    resonance: Vec<usize>,
217    frame: Frame,
218}
219impl Display for Phi {
220    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
221        write!(
222            f,
223            "Phi(beam={}, recoil=[{}], daughter=[{}], resonance=[{}], frame={})",
224            self.beam,
225            indices_to_string(&self.recoil),
226            indices_to_string(&self.daughter),
227            indices_to_string(&self.resonance),
228            self.frame
229        )
230    }
231}
232impl Phi {
233    /// Construct the angle given the four-momentum indices for each specified particle. Fields
234    /// which can take lists of more than one index will add the relevant four-momenta to make a
235    /// new particle from the constituents. See [`Frame`] for options regarding the reference
236    /// frame.
237    pub fn new<T: AsRef<[usize]>, U: AsRef<[usize]>, V: AsRef<[usize]>>(
238        beam: usize,
239        recoil: T,
240        daughter: U,
241        resonance: V,
242        frame: Frame,
243    ) -> Self {
244        Self {
245            beam,
246            recoil: recoil.as_ref().into(),
247            daughter: daughter.as_ref().into(),
248            resonance: resonance.as_ref().into(),
249            frame,
250        }
251    }
252}
253impl Default for Phi {
254    fn default() -> Self {
255        Self {
256            beam: 0,
257            recoil: vec![1],
258            daughter: vec![2],
259            resonance: vec![2, 3],
260            frame: Frame::Helicity,
261        }
262    }
263}
264#[typetag::serde]
265impl Variable for Phi {
266    fn value(&self, event: &Event) -> Float {
267        let beam = event.p4s[self.beam];
268        let recoil = event.get_p4_sum(&self.recoil);
269        let daughter = event.get_p4_sum(&self.daughter);
270        let resonance = event.get_p4_sum(&self.resonance);
271        let daughter_res = daughter.boost(&-resonance.beta());
272        match self.frame {
273            Frame::Helicity => {
274                let recoil_res = recoil.boost(&-resonance.beta());
275                let z = -recoil_res.vec3().unit();
276                let y = beam.vec3().cross(&-recoil.vec3()).unit();
277                let x = y.cross(&z);
278                let angles = Vec3::new(
279                    daughter_res.vec3().dot(&x),
280                    daughter_res.vec3().dot(&y),
281                    daughter_res.vec3().dot(&z),
282                );
283                angles.phi()
284            }
285            Frame::GottfriedJackson => {
286                let beam_res = beam.boost(&-resonance.beta());
287                let z = beam_res.vec3().unit();
288                let y = beam.vec3().cross(&-recoil.vec3()).unit();
289                let x = y.cross(&z);
290                let angles = Vec3::new(
291                    daughter_res.vec3().dot(&x),
292                    daughter_res.vec3().dot(&y),
293                    daughter_res.vec3().dot(&z),
294                );
295                angles.phi()
296            }
297        }
298    }
299}
300
301/// A struct for obtaining both spherical angles at the same time.
302#[derive(Clone, Debug, Serialize, Deserialize)]
303pub struct Angles {
304    /// See [`CosTheta`].
305    pub costheta: CosTheta,
306    /// See [`Phi`].
307    pub phi: Phi,
308}
309
310impl Display for Angles {
311    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
312        write!(
313            f,
314            "Angles(beam={}, recoil=[{}], daughter=[{}], resonance=[{}], frame={})",
315            self.costheta.beam,
316            indices_to_string(&self.costheta.recoil),
317            indices_to_string(&self.costheta.daughter),
318            indices_to_string(&self.costheta.resonance),
319            self.costheta.frame
320        )
321    }
322}
323impl Angles {
324    /// Construct the angles given the four-momentum indices for each specified particle. Fields
325    /// which can take lists of more than one index will add the relevant four-momenta to make a
326    /// new particle from the constituents. See [`Frame`] for options regarding the reference
327    /// frame.
328    pub fn new<T: AsRef<[usize]>, U: AsRef<[usize]>, V: AsRef<[usize]>>(
329        beam: usize,
330        recoil: T,
331        daughter: U,
332        resonance: V,
333        frame: Frame,
334    ) -> Self {
335        Self {
336            costheta: CosTheta::new(beam, &recoil, &daughter, &resonance, frame),
337            phi: Phi {
338                beam,
339                recoil: recoil.as_ref().into(),
340                daughter: daughter.as_ref().into(),
341                resonance: resonance.as_ref().into(),
342                frame,
343            },
344        }
345    }
346}
347
348/// A struct defining the polarization angle for a beam relative to the production plane.
349#[derive(Clone, Debug, Serialize, Deserialize)]
350pub struct PolAngle {
351    beam: usize,
352    recoil: Vec<usize>,
353    beam_polarization: usize,
354}
355impl Display for PolAngle {
356    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
357        write!(
358            f,
359            "PolAngle(beam={}, recoil=[{}], beam_polarization={})",
360            self.beam,
361            indices_to_string(&self.recoil),
362            self.beam_polarization,
363        )
364    }
365}
366impl PolAngle {
367    /// Constructs the polarization angle given the four-momentum indices for each specified
368    /// particle. Fields which can take lists of more than one index will add the relevant
369    /// four-momenta to make a new particle from the constituents.
370    pub fn new<T: AsRef<[usize]>>(beam: usize, recoil: T, beam_polarization: usize) -> Self {
371        Self {
372            beam,
373            recoil: recoil.as_ref().into(),
374            beam_polarization,
375        }
376    }
377}
378#[typetag::serde]
379impl Variable for PolAngle {
380    fn value(&self, event: &Event) -> Float {
381        let beam = event.p4s[self.beam];
382        let recoil = event.get_p4_sum(&self.recoil);
383        let y = beam.vec3().cross(&-recoil.vec3()).unit();
384        Float::atan2(
385            y.dot(&event.aux[self.beam_polarization]),
386            beam.vec3()
387                .unit()
388                .dot(&event.aux[self.beam_polarization].cross(&y)),
389        )
390    }
391}
392
393/// A struct defining the polarization magnitude for a beam relative to the production plane.
394#[derive(Copy, Clone, Default, Debug, Serialize, Deserialize)]
395pub struct PolMagnitude {
396    beam_polarization: usize,
397}
398impl Display for PolMagnitude {
399    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
400        write!(
401            f,
402            "PolMagnitude(beam_polarization={})",
403            self.beam_polarization,
404        )
405    }
406}
407impl PolMagnitude {
408    /// Constructs the polarization magnitude given the four-momentum index for the beam.
409    pub fn new(beam_polarization: usize) -> Self {
410        Self { beam_polarization }
411    }
412}
413#[typetag::serde]
414impl Variable for PolMagnitude {
415    fn value(&self, event: &Event) -> Float {
416        event.aux[self.beam_polarization].mag()
417    }
418}
419
420/// A struct for obtaining both the polarization angle and magnitude at the same time.
421#[derive(Clone, Debug, Serialize, Deserialize)]
422pub struct Polarization {
423    /// See [`PolMagnitude`].
424    pub pol_magnitude: PolMagnitude,
425    /// See [`PolAngle`].
426    pub pol_angle: PolAngle,
427}
428impl Display for Polarization {
429    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
430        write!(
431            f,
432            "Polarization(beam={}, recoil=[{}], beam_polarization={})",
433            self.pol_angle.beam,
434            indices_to_string(&self.pol_angle.recoil),
435            self.pol_angle.beam_polarization,
436        )
437    }
438}
439impl Polarization {
440    /// Constructs the polarization angle and magnitude given the four-momentum indices for
441    /// the beam and target (recoil) particle. Fields which can take lists of more than one index will add
442    /// the relevant four-momenta to make a new particle from the constituents.
443    pub fn new<T: AsRef<[usize]>>(beam: usize, recoil: T, beam_polarization: usize) -> Self {
444        Self {
445            pol_magnitude: PolMagnitude::new(beam_polarization),
446            pol_angle: PolAngle::new(beam, recoil, beam_polarization),
447        }
448    }
449}
450
451/// A struct used to calculate Mandelstam variables ($`s`$, $`t`$, or $`u`$).
452///
453/// By convention, the metric is chosen to be $`(+---)`$ and the variables are defined as follows
454/// (ignoring factors of $`c`$):
455///
456/// $`s = (p_1 + p_2)^2 = (p_3 + p_4)^2`$
457///
458/// $`t = (p_1 - p_3)^2 = (p_4 - p_2)^2`$
459///
460/// $`u = (p_1 - p_4)^2 = (p_3 - p_2)^2`$
461#[derive(Clone, Debug, Serialize, Deserialize)]
462pub struct Mandelstam {
463    p1: Vec<usize>,
464    p2: Vec<usize>,
465    p3: Vec<usize>,
466    p4: Vec<usize>,
467    missing: Option<u8>,
468    channel: Channel,
469}
470impl Display for Mandelstam {
471    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
472        write!(
473            f,
474            "Mandelstam(p1=[{}], p2=[{}], p3=[{}], p4=[{}], channel={})",
475            indices_to_string(&self.p1),
476            indices_to_string(&self.p2),
477            indices_to_string(&self.p3),
478            indices_to_string(&self.p4),
479            self.channel,
480        )
481    }
482}
483impl Mandelstam {
484    /// Constructs the Mandelstam variable for the given `channel` and particles.
485    /// Fields which can take lists of more than one index will add
486    /// the relevant four-momenta to make a new particle from the constituents.
487    pub fn new<T, U, V, W>(p1: T, p2: U, p3: V, p4: W, channel: Channel) -> Result<Self, LadduError>
488    where
489        T: AsRef<[usize]>,
490        U: AsRef<[usize]>,
491        V: AsRef<[usize]>,
492        W: AsRef<[usize]>,
493    {
494        let mut missing = None;
495        if p1.as_ref().is_empty() {
496            missing = Some(1)
497        }
498        if p2.as_ref().is_empty() {
499            if missing.is_none() {
500                missing = Some(2)
501            } else {
502                return Err(LadduError::Custom("A maximum of one particle may be ommitted while constructing a Mandelstam variable!".to_string()));
503            }
504        }
505        if p3.as_ref().is_empty() {
506            if missing.is_none() {
507                missing = Some(3)
508            } else {
509                return Err(LadduError::Custom("A maximum of one particle may be ommitted while constructing a Mandelstam variable!".to_string()));
510            }
511        }
512        if p4.as_ref().is_empty() {
513            if missing.is_none() {
514                missing = Some(4)
515            } else {
516                return Err(LadduError::Custom("A maximum of one particle may be ommitted while constructing a Mandelstam variable!".to_string()));
517            }
518        }
519        Ok(Self {
520            p1: p1.as_ref().into(),
521            p2: p2.as_ref().into(),
522            p3: p3.as_ref().into(),
523            p4: p4.as_ref().into(),
524            missing,
525            channel,
526        })
527    }
528}
529
530#[typetag::serde]
531impl Variable for Mandelstam {
532    fn value(&self, event: &Event) -> Float {
533        match self.channel {
534            Channel::S => match self.missing {
535                None | Some(3) | Some(4) => {
536                    let p1 = event.get_p4_sum(&self.p1);
537                    let p2 = event.get_p4_sum(&self.p2);
538                    (p1 + p2).mag2()
539                }
540                Some(1) | Some(2) => {
541                    let p3 = event.get_p4_sum(&self.p3);
542                    let p4 = event.get_p4_sum(&self.p4);
543                    (p3 + p4).mag2()
544                }
545                _ => unreachable!(),
546            },
547            Channel::T => match self.missing {
548                None | Some(2) | Some(4) => {
549                    let p1 = event.get_p4_sum(&self.p1);
550                    let p3 = event.get_p4_sum(&self.p3);
551                    (p1 - p3).mag2()
552                }
553                Some(1) | Some(3) => {
554                    let p2 = event.get_p4_sum(&self.p2);
555                    let p4 = event.get_p4_sum(&self.p4);
556                    (p4 - p2).mag2()
557                }
558                _ => unreachable!(),
559            },
560            Channel::U => match self.missing {
561                None | Some(2) | Some(3) => {
562                    let p1 = event.get_p4_sum(&self.p1);
563                    let p4 = event.get_p4_sum(&self.p4);
564                    (p1 - p4).mag2()
565                }
566                Some(1) | Some(4) => {
567                    let p2 = event.get_p4_sum(&self.p2);
568                    let p3 = event.get_p4_sum(&self.p3);
569                    (p3 - p2).mag2()
570                }
571                _ => unreachable!(),
572            },
573        }
574    }
575}
576
577#[cfg(test)]
578mod tests {
579    use approx::assert_relative_eq;
580
581    use crate::data::{test_dataset, test_event};
582
583    use super::*;
584    #[test]
585    fn test_mass_single_particle() {
586        let event = test_event();
587        let mass = Mass::new([1]);
588        assert_relative_eq!(mass.value(&event), 1.007);
589    }
590
591    #[test]
592    fn test_mass_multiple_particles() {
593        let event = test_event();
594        let mass = Mass::new([2, 3]);
595        assert_relative_eq!(
596            mass.value(&event),
597            1.37437863,
598            epsilon = Float::EPSILON.sqrt()
599        );
600    }
601
602    #[test]
603    fn test_mass_display() {
604        let mass = Mass::new([2, 3]);
605        assert_eq!(mass.to_string(), "Mass(constituents=[2, 3])");
606    }
607
608    #[test]
609    fn test_costheta_helicity() {
610        let event = test_event();
611        let costheta = CosTheta::new(0, [1], [2], [2, 3], Frame::Helicity);
612        assert_relative_eq!(
613            costheta.value(&event),
614            -0.4611175,
615            epsilon = Float::EPSILON.sqrt()
616        );
617    }
618
619    #[test]
620    fn test_costheta_display() {
621        let costheta = CosTheta::new(0, [1], [2], [2, 3], Frame::Helicity);
622        assert_eq!(
623            costheta.to_string(),
624            "CosTheta(beam=0, recoil=[1], daughter=[2], resonance=[2, 3], frame=Helicity)"
625        );
626    }
627
628    #[test]
629    fn test_phi_helicity() {
630        let event = test_event();
631        let phi = Phi::new(0, [1], [2], [2, 3], Frame::Helicity);
632        assert_relative_eq!(
633            phi.value(&event),
634            -2.65746258,
635            epsilon = Float::EPSILON.sqrt()
636        );
637    }
638
639    #[test]
640    fn test_phi_display() {
641        let phi = Phi::new(0, [1], [2], [2, 3], Frame::Helicity);
642        assert_eq!(
643            phi.to_string(),
644            "Phi(beam=0, recoil=[1], daughter=[2], resonance=[2, 3], frame=Helicity)"
645        );
646    }
647
648    #[test]
649    fn test_costheta_gottfried_jackson() {
650        let event = test_event();
651        let costheta = CosTheta::new(0, [1], [2], [2, 3], Frame::GottfriedJackson);
652        assert_relative_eq!(
653            costheta.value(&event),
654            0.09198832,
655            epsilon = Float::EPSILON.sqrt()
656        );
657    }
658
659    #[test]
660    fn test_phi_gottfried_jackson() {
661        let event = test_event();
662        let phi = Phi::new(0, [1], [2], [2, 3], Frame::GottfriedJackson);
663        assert_relative_eq!(
664            phi.value(&event),
665            -2.71391319,
666            epsilon = Float::EPSILON.sqrt()
667        );
668    }
669
670    #[test]
671    fn test_angles() {
672        let event = test_event();
673        let angles = Angles::new(0, [1], [2], [2, 3], Frame::Helicity);
674        assert_relative_eq!(
675            angles.costheta.value(&event),
676            -0.4611175,
677            epsilon = Float::EPSILON.sqrt()
678        );
679        assert_relative_eq!(
680            angles.phi.value(&event),
681            -2.65746258,
682            epsilon = Float::EPSILON.sqrt()
683        );
684    }
685
686    #[test]
687    fn test_angles_display() {
688        let angles = Angles::new(0, [1], [2], [2, 3], Frame::Helicity);
689        assert_eq!(
690            angles.to_string(),
691            "Angles(beam=0, recoil=[1], daughter=[2], resonance=[2, 3], frame=Helicity)"
692        );
693    }
694
695    #[test]
696    fn test_pol_angle() {
697        let event = test_event();
698        let pol_angle = PolAngle::new(0, vec![1], 0);
699        assert_relative_eq!(
700            pol_angle.value(&event),
701            1.93592989,
702            epsilon = Float::EPSILON.sqrt()
703        );
704    }
705
706    #[test]
707    fn test_pol_angle_display() {
708        let pol_angle = PolAngle::new(0, vec![1], 0);
709        assert_eq!(
710            pol_angle.to_string(),
711            "PolAngle(beam=0, recoil=[1], beam_polarization=0)"
712        );
713    }
714
715    #[test]
716    fn test_pol_magnitude() {
717        let event = test_event();
718        let pol_magnitude = PolMagnitude::new(0);
719        assert_relative_eq!(
720            pol_magnitude.value(&event),
721            0.38562805,
722            epsilon = Float::EPSILON.sqrt()
723        );
724    }
725
726    #[test]
727    fn test_pol_magnitude_display() {
728        let pol_magnitude = PolMagnitude::new(0);
729        assert_eq!(
730            pol_magnitude.to_string(),
731            "PolMagnitude(beam_polarization=0)"
732        );
733    }
734
735    #[test]
736    fn test_polarization() {
737        let event = test_event();
738        let polarization = Polarization::new(0, vec![1], 0);
739        assert_relative_eq!(
740            polarization.pol_angle.value(&event),
741            1.93592989,
742            epsilon = Float::EPSILON.sqrt()
743        );
744        assert_relative_eq!(
745            polarization.pol_magnitude.value(&event),
746            0.38562805,
747            epsilon = Float::EPSILON.sqrt()
748        );
749    }
750
751    #[test]
752    fn test_polarization_display() {
753        let polarization = Polarization::new(0, vec![1], 0);
754        assert_eq!(
755            polarization.to_string(),
756            "Polarization(beam=0, recoil=[1], beam_polarization=0)"
757        );
758    }
759
760    #[test]
761    fn test_mandelstam() {
762        let event = test_event();
763        let s = Mandelstam::new([0], [], [2, 3], [1], Channel::S).unwrap();
764        let t = Mandelstam::new([0], [], [2, 3], [1], Channel::T).unwrap();
765        let u = Mandelstam::new([0], [], [2, 3], [1], Channel::U).unwrap();
766        let sp = Mandelstam::new([], [0], [1], [2, 3], Channel::S).unwrap();
767        let tp = Mandelstam::new([], [0], [1], [2, 3], Channel::T).unwrap();
768        let up = Mandelstam::new([], [0], [1], [2, 3], Channel::U).unwrap();
769        assert_relative_eq!(
770            s.value(&event),
771            18.50401105,
772            epsilon = Float::EPSILON.sqrt()
773        );
774        assert_relative_eq!(s.value(&event), sp.value(&event),);
775        assert_relative_eq!(
776            t.value(&event),
777            -0.19222859,
778            epsilon = Float::EPSILON.sqrt()
779        );
780        assert_relative_eq!(t.value(&event), tp.value(&event),);
781        assert_relative_eq!(
782            u.value(&event),
783            -14.40419893,
784            epsilon = Float::EPSILON.sqrt()
785        );
786        assert_relative_eq!(u.value(&event), up.value(&event),);
787        let m2_beam = test_event().get_p4_sum([0]).m2();
788        let m2_recoil = test_event().get_p4_sum([1]).m2();
789        let m2_res = test_event().get_p4_sum([2, 3]).m2();
790        assert_relative_eq!(
791            s.value(&event) + t.value(&event) + u.value(&event) - m2_beam - m2_recoil - m2_res,
792            1.00,
793            epsilon = 1e-2
794        );
795        // Note: not very accurate, but considering the values in test_event only go to about 3
796        // decimal places, this is probably okay
797    }
798
799    #[test]
800    fn test_mandelstam_display() {
801        let s = Mandelstam::new([0], [], [2, 3], [1], Channel::S).unwrap();
802        assert_eq!(
803            s.to_string(),
804            "Mandelstam(p1=[0], p2=[], p3=[2, 3], p4=[1], channel=s)"
805        );
806    }
807
808    #[test]
809    fn test_variable_value_on() {
810        let dataset = Arc::new(test_dataset());
811        let mass = Mass::new(vec![2, 3]);
812
813        let values = mass.value_on(&dataset);
814        assert_eq!(values.len(), 1);
815        assert_relative_eq!(values[0], 1.37437863, epsilon = Float::EPSILON.sqrt());
816    }
817}