laddu_core/utils/
variables.rs

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