laddu_python/utils/
variables.rs

1use crate::{
2    amplitudes::PyExpression,
3    data::{PyDataset, PyEvent},
4    utils::vectors::PyVec4,
5};
6use laddu_core::{
7    data::{Dataset, DatasetMetadata, Event, NamedEventView},
8    traits::Variable,
9    utils::angular_momentum::{AngularMomentum, AngularMomentumProjection, OrbitalAngularMomentum},
10    utils::reaction::{Decay, Particle, Reaction},
11    utils::variables::{
12        Angles, CosTheta, IntoP4Selection, Mandelstam, Mass, P4Selection, Phi, PolAngle,
13        PolMagnitude, Polarization, VariableExpression,
14    },
15    LadduError, LadduResult,
16};
17use num::rational::Ratio;
18use numpy::PyArray1;
19use pyo3::{exceptions::PyValueError, prelude::*, types::PyBool};
20use serde::{Deserialize, Serialize};
21use std::fmt::{Debug, Display};
22
23#[derive(FromPyObject, Clone, Serialize, Deserialize)]
24pub enum PyVariable {
25    #[pyo3(transparent)]
26    Mass(PyMass),
27    #[pyo3(transparent)]
28    CosTheta(PyCosTheta),
29    #[pyo3(transparent)]
30    Phi(PyPhi),
31    #[pyo3(transparent)]
32    PolAngle(PyPolAngle),
33    #[pyo3(transparent)]
34    PolMagnitude(PyPolMagnitude),
35    #[pyo3(transparent)]
36    Mandelstam(PyMandelstam),
37}
38
39impl Debug for PyVariable {
40    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
41        match self {
42            Self::Mass(v) => write!(f, "{:?}", v.0),
43            Self::CosTheta(v) => write!(f, "{:?}", v.0),
44            Self::Phi(v) => write!(f, "{:?}", v.0),
45            Self::PolAngle(v) => write!(f, "{:?}", v.0),
46            Self::PolMagnitude(v) => write!(f, "{:?}", v.0),
47            Self::Mandelstam(v) => write!(f, "{:?}", v.0),
48        }
49    }
50}
51impl Display for PyVariable {
52    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
53        match self {
54            Self::Mass(v) => write!(f, "{}", v.0),
55            Self::CosTheta(v) => write!(f, "{}", v.0),
56            Self::Phi(v) => write!(f, "{}", v.0),
57            Self::PolAngle(v) => write!(f, "{}", v.0),
58            Self::PolMagnitude(v) => write!(f, "{}", v.0),
59            Self::Mandelstam(v) => write!(f, "{}", v.0),
60        }
61    }
62}
63
64impl PyVariable {
65    pub(crate) fn bind_in_place(&mut self, metadata: &DatasetMetadata) -> PyResult<()> {
66        match self {
67            Self::Mass(mass) => mass.0.bind(metadata).map_err(PyErr::from),
68            Self::CosTheta(cos_theta) => cos_theta.0.bind(metadata).map_err(PyErr::from),
69            Self::Phi(phi) => phi.0.bind(metadata).map_err(PyErr::from),
70            Self::PolAngle(pol_angle) => pol_angle.0.bind(metadata).map_err(PyErr::from),
71            Self::PolMagnitude(pol_magnitude) => {
72                pol_magnitude.0.bind(metadata).map_err(PyErr::from)
73            }
74            Self::Mandelstam(mandelstam) => mandelstam.0.bind(metadata).map_err(PyErr::from),
75        }
76    }
77
78    pub(crate) fn bound(&self, metadata: &DatasetMetadata) -> PyResult<Self> {
79        let mut cloned = self.clone();
80        cloned.bind_in_place(metadata)?;
81        Ok(cloned)
82    }
83
84    pub(crate) fn evaluate_event(&self, event: &Event) -> PyResult<f64> {
85        let dataset = Dataset::new_with_metadata(vec![event.data_arc()], event.metadata_arc());
86        let event_view = dataset.event_view(0);
87        Ok(self.value(&event_view))
88    }
89}
90
91#[pyclass(name = "VariableExpression", module = "laddu")]
92pub struct PyVariableExpression(pub VariableExpression);
93
94#[pymethods]
95impl PyVariableExpression {
96    fn __and__(&self, rhs: &PyVariableExpression) -> PyVariableExpression {
97        PyVariableExpression(self.0.clone() & rhs.0.clone())
98    }
99    fn __or__(&self, rhs: &PyVariableExpression) -> PyVariableExpression {
100        PyVariableExpression(self.0.clone() | rhs.0.clone())
101    }
102    fn __invert__(&self) -> PyVariableExpression {
103        PyVariableExpression(!self.0.clone())
104    }
105    fn __str__(&self) -> String {
106        format!("{}", self.0)
107    }
108}
109
110#[derive(Clone, FromPyObject)]
111pub enum PyP4SelectionInput {
112    #[pyo3(transparent)]
113    Name(String),
114    #[pyo3(transparent)]
115    Names(Vec<String>),
116}
117
118impl PyP4SelectionInput {
119    fn into_selection(self) -> P4Selection {
120        match self {
121            PyP4SelectionInput::Name(name) => name.into_selection(),
122            PyP4SelectionInput::Names(names) => names.into_selection(),
123        }
124    }
125}
126
127/// A kinematic particle used to define reaction-aware variables.
128#[pyclass(name = "Particle", module = "laddu", from_py_object)]
129#[derive(Clone, Serialize, Deserialize)]
130pub struct PyParticle(pub Particle);
131
132#[pymethods]
133impl PyParticle {
134    /// Construct a measured particle from one or more p4 column names.
135    #[staticmethod]
136    fn measured(label: &str, p4: PyP4SelectionInput) -> Self {
137        Self(Particle::measured(label, p4.into_selection()))
138    }
139
140    /// Construct a particle with fixed event-independent four-momentum.
141    #[staticmethod]
142    fn fixed(label: &str, p4: &PyVec4) -> Self {
143        Self(Particle::fixed(label, p4.0))
144    }
145
146    /// Construct a missing particle solved by the reaction topology.
147    #[staticmethod]
148    fn missing(label: &str) -> Self {
149        Self(Particle::missing(label))
150    }
151
152    /// Construct a composite particle from daughter particles.
153    #[staticmethod]
154    fn composite(label: &str, daughters: Vec<PyParticle>) -> PyResult<Self> {
155        Ok(Self(Particle::composite(
156            label,
157            daughters.iter().map(|daughter| &daughter.0),
158        )?))
159    }
160
161    /// The particle label.
162    #[getter]
163    fn label(&self) -> String {
164        self.0.label().to_string()
165    }
166
167    fn __repr__(&self) -> String {
168        format!("{:?}", self.0)
169    }
170
171    fn __str__(&self) -> String {
172        self.0.to_string()
173    }
174}
175
176/// A reaction topology with direct particle definitions.
177#[pyclass(name = "Reaction", module = "laddu", from_py_object)]
178#[derive(Clone, Serialize, Deserialize)]
179pub struct PyReaction(pub Reaction);
180
181#[pymethods]
182impl PyReaction {
183    /// Construct a two-to-two reaction from `p1, p2, p3, p4`.
184    #[staticmethod]
185    fn two_to_two(
186        p1: &PyParticle,
187        p2: &PyParticle,
188        p3: &PyParticle,
189        p4: &PyParticle,
190    ) -> PyResult<Self> {
191        Ok(Self(Reaction::two_to_two(&p1.0, &p2.0, &p3.0, &p4.0)?))
192    }
193
194    /// Construct a particle mass variable.
195    fn mass(&self, particle: &PyParticle) -> PyMass {
196        PyMass(self.0.mass(&particle.0))
197    }
198
199    /// Construct an isobar decay view.
200    fn decay(&self, parent: &PyParticle) -> PyResult<PyDecay> {
201        Ok(PyDecay(self.0.decay(&parent.0)?))
202    }
203
204    /// Construct a Mandelstam variable.
205    fn mandelstam(&self, channel: &str) -> PyResult<PyMandelstam> {
206        Ok(PyMandelstam(self.0.mandelstam(channel.parse()?)))
207    }
208
209    /// Construct a polarization-angle variable.
210    fn pol_angle(&self, angle_aux: String) -> PyPolAngle {
211        PyPolAngle(self.0.pol_angle(angle_aux))
212    }
213
214    /// Construct polarization variables.
215    fn polarization(&self, pol_magnitude: String, pol_angle: String) -> PyResult<PyPolarization> {
216        if pol_magnitude == pol_angle {
217            return Err(PyValueError::new_err(
218                "`pol_magnitude` and `pol_angle` must reference distinct auxiliary columns",
219            ));
220        }
221        Ok(PyPolarization(
222            self.0.polarization(pol_magnitude, pol_angle),
223        ))
224    }
225
226    fn __repr__(&self) -> String {
227        format!("{:?}", self.0)
228    }
229}
230
231/// A reaction-aware isobar decay view.
232#[pyclass(name = "Decay", module = "laddu", from_py_object)]
233#[derive(Clone, Serialize, Deserialize)]
234pub struct PyDecay(pub Decay);
235
236#[pymethods]
237impl PyDecay {
238    /// The parent particle.
239    #[getter]
240    fn parent(&self) -> PyParticle {
241        PyParticle(self.0.parent().clone())
242    }
243
244    /// The first daughter particle.
245    #[getter]
246    fn daughter_1(&self) -> PyParticle {
247        PyParticle(self.0.daughter_1().clone())
248    }
249
250    /// The second daughter particle.
251    #[getter]
252    fn daughter_2(&self) -> PyParticle {
253        PyParticle(self.0.daughter_2().clone())
254    }
255
256    /// Ordered daughter particles.
257    fn daughters(&self) -> Vec<PyParticle> {
258        self.0
259            .daughters()
260            .into_iter()
261            .map(|daughter| PyParticle(daughter.clone()))
262            .collect()
263    }
264
265    /// Parent mass variable.
266    fn mass(&self) -> PyMass {
267        PyMass(self.0.mass())
268    }
269
270    /// Parent mass variable.
271    fn parent_mass(&self) -> PyMass {
272        PyMass(self.0.parent_mass())
273    }
274
275    /// First daughter mass variable.
276    fn daughter_1_mass(&self) -> PyMass {
277        PyMass(self.0.daughter_1_mass())
278    }
279
280    /// Second daughter mass variable.
281    fn daughter_2_mass(&self) -> PyMass {
282        PyMass(self.0.daughter_2_mass())
283    }
284
285    /// Mass variable for a selected daughter.
286    fn daughter_mass(&self, daughter: &PyParticle) -> PyResult<PyMass> {
287        Ok(PyMass(self.0.daughter_mass(&daughter.0)?))
288    }
289
290    /// Decay costheta variable for the selected frame.
291    #[pyo3(signature=(daughter, frame="Helicity"))]
292    fn costheta(&self, daughter: &PyParticle, frame: &str) -> PyResult<PyCosTheta> {
293        Ok(PyCosTheta(self.0.costheta(&daughter.0, frame.parse()?)?))
294    }
295
296    /// Decay phi variable for the selected frame.
297    #[pyo3(signature=(daughter, frame="Helicity"))]
298    fn phi(&self, daughter: &PyParticle, frame: &str) -> PyResult<PyPhi> {
299        Ok(PyPhi(self.0.phi(&daughter.0, frame.parse()?)?))
300    }
301
302    /// Decay angle variables for the selected frame.
303    #[pyo3(signature=(daughter, frame="Helicity"))]
304    fn angles(&self, daughter: &PyParticle, frame: &str) -> PyResult<PyAngles> {
305        Ok(PyAngles(self.0.angles(&daughter.0, frame.parse()?)?))
306    }
307
308    /// Construct the helicity-basis angular factor for one explicit helicity term.
309    #[pyo3(signature=(name, spin, projection, daughter, lambda_1, lambda_2, frame="Helicity"))]
310    #[allow(clippy::too_many_arguments)]
311    fn helicity_factor(
312        &self,
313        name: &str,
314        spin: &Bound<'_, PyAny>,
315        projection: &Bound<'_, PyAny>,
316        daughter: &PyParticle,
317        lambda_1: &Bound<'_, PyAny>,
318        lambda_2: &Bound<'_, PyAny>,
319        frame: &str,
320    ) -> PyResult<PyExpression> {
321        Ok(PyExpression(self.0.helicity_factor(
322            name,
323            parse_angular_momentum(spin)?,
324            parse_projection(projection)?,
325            &daughter.0,
326            parse_projection(lambda_1)?,
327            parse_projection(lambda_2)?,
328            frame.parse()?,
329        )?))
330    }
331
332    /// Construct the canonical-basis spin-angular factor for one explicit LS/helicity term.
333    #[pyo3(signature=(name, spin, projection, orbital_l, coupled_spin, daughter, daughter_1_spin, daughter_2_spin, lambda_1, lambda_2, frame="Helicity"))]
334    #[allow(clippy::too_many_arguments)]
335    fn canonical_factor(
336        &self,
337        name: &str,
338        spin: &Bound<'_, PyAny>,
339        projection: &Bound<'_, PyAny>,
340        orbital_l: &Bound<'_, PyAny>,
341        coupled_spin: &Bound<'_, PyAny>,
342        daughter: &PyParticle,
343        daughter_1_spin: &Bound<'_, PyAny>,
344        daughter_2_spin: &Bound<'_, PyAny>,
345        lambda_1: &Bound<'_, PyAny>,
346        lambda_2: &Bound<'_, PyAny>,
347        frame: &str,
348    ) -> PyResult<PyExpression> {
349        Ok(PyExpression(self.0.canonical_factor(
350            name,
351            parse_angular_momentum(spin)?,
352            parse_projection(projection)?,
353            parse_orbital_angular_momentum(orbital_l)?,
354            parse_angular_momentum(coupled_spin)?,
355            &daughter.0,
356            parse_angular_momentum(daughter_1_spin)?,
357            parse_angular_momentum(daughter_2_spin)?,
358            parse_projection(lambda_1)?,
359            parse_projection(lambda_2)?,
360            frame.parse()?,
361        )?))
362    }
363
364    fn __repr__(&self) -> String {
365        format!("{:?}", self.0)
366    }
367}
368
369fn parse_angular_momentum(input: &Bound<'_, PyAny>) -> PyResult<AngularMomentum> {
370    parse_ratio_like(input)
371        .and_then(AngularMomentum::from_ratio)
372        .map_err(py_value_error)
373}
374
375fn parse_projection(input: &Bound<'_, PyAny>) -> PyResult<AngularMomentumProjection> {
376    parse_ratio_like(input)
377        .and_then(AngularMomentumProjection::from_ratio)
378        .map_err(py_value_error)
379}
380
381fn parse_orbital_angular_momentum(input: &Bound<'_, PyAny>) -> PyResult<OrbitalAngularMomentum> {
382    parse_ratio_like(input)
383        .and_then(OrbitalAngularMomentum::from_ratio)
384        .map_err(py_value_error)
385}
386
387fn parse_ratio_like(input: &Bound<'_, PyAny>) -> LadduResult<Ratio<i32>> {
388    if input.is_instance_of::<PyBool>() {
389        return Err(LadduError::Custom(
390            "quantum number cannot be a bool".to_string(),
391        ));
392    }
393    if let Ok(value) = input.extract::<i32>() {
394        return Ok(Ratio::from_integer(value));
395    }
396    if let Ok(value) = input.extract::<f64>() {
397        let twice = AngularMomentumProjection::from_f64(value)?.value();
398        return Ok(Ratio::new(twice, 2));
399    }
400    let numerator = input
401        .getattr("numerator")
402        .and_then(|value| value.extract::<i32>());
403    let denominator = input
404        .getattr("denominator")
405        .and_then(|value| value.extract::<i32>());
406    if let (Ok(numerator), Ok(denominator)) = (numerator, denominator) {
407        if denominator == 0 {
408            return Err(LadduError::Custom(
409                "quantum number denominator cannot be zero".to_string(),
410            ));
411        }
412        return Ok(Ratio::new(numerator, denominator));
413    }
414    Err(LadduError::Custom(
415        "quantum number must be an int, float, or fractions.Fraction".to_string(),
416    ))
417}
418
419fn py_value_error(err: LadduError) -> PyErr {
420    PyValueError::new_err(err.to_string())
421}
422/// The invariant mass of an arbitrary combination of constituent particles in an Event
423///
424/// This variable is calculated by summing up the 4-momenta of each particle listed by index in
425/// `constituents` and taking the invariant magnitude of the resulting 4-vector.
426///
427/// Parameters
428/// ----------
429/// constituents : str or list of str
430///     Particle names to combine when constructing the final four-momentum
431///
432/// See Also
433/// --------
434/// laddu.utils.vectors.Vec4.m
435///
436#[pyclass(name = "Mass", module = "laddu", from_py_object)]
437#[derive(Clone, Serialize, Deserialize)]
438pub struct PyMass(pub Mass);
439
440#[pymethods]
441impl PyMass {
442    #[new]
443    fn new(constituents: PyP4SelectionInput) -> Self {
444        Self(Mass::new(constituents.into_selection()))
445    }
446    /// The value of this Variable for the given Event
447    ///
448    /// Parameters
449    /// ----------
450    /// event : Event
451    ///     The Event upon which the Variable is calculated
452    ///
453    /// Returns
454    /// -------
455    /// value : float
456    ///     The value of the Variable for the given `event`
457    ///
458    fn value(&self, event: &PyEvent) -> PyResult<f64> {
459        let metadata = event
460            .metadata_opt()
461            .ok_or_else(|| PyValueError::new_err(
462                "This event is not associated with metadata; supply `p4_names`/`aux_names` when constructing it or evaluate via a Dataset.",
463            ))?;
464        let mut variable = self.0.clone();
465        variable.bind(metadata).map_err(PyErr::from)?;
466        let dataset =
467            Dataset::new_with_metadata(vec![event.event.data_arc()], event.event.metadata_arc());
468        let event_view = dataset.event_view(0);
469        Ok(variable.value(&event_view))
470    }
471    /// All values of this Variable on the given Dataset
472    ///
473    /// Parameters
474    /// ----------
475    /// dataset : Dataset
476    ///     The Dataset upon which the Variable is calculated
477    ///
478    /// Returns
479    /// -------
480    /// values : array_like
481    ///     The values of the Variable for each Event in the given `dataset`
482    ///
483    fn value_on<'py>(
484        &self,
485        py: Python<'py>,
486        dataset: &PyDataset,
487    ) -> PyResult<Bound<'py, PyArray1<f64>>> {
488        let values = self.0.value_on(&dataset.0).map_err(PyErr::from)?;
489        Ok(PyArray1::from_vec(py, values))
490    }
491    fn __eq__(&self, value: f64) -> PyVariableExpression {
492        PyVariableExpression(self.0.eq(value))
493    }
494    fn __lt__(&self, value: f64) -> PyVariableExpression {
495        PyVariableExpression(self.0.lt(value))
496    }
497    fn __gt__(&self, value: f64) -> PyVariableExpression {
498        PyVariableExpression(self.0.gt(value))
499    }
500    fn __le__(&self, value: f64) -> PyVariableExpression {
501        PyVariableExpression(self.0.le(value))
502    }
503    fn __ge__(&self, value: f64) -> PyVariableExpression {
504        PyVariableExpression(self.0.ge(value))
505    }
506    fn __repr__(&self) -> String {
507        format!("{:?}", self.0)
508    }
509    fn __str__(&self) -> String {
510        format!("{}", self.0)
511    }
512}
513
514/// The cosine of the polar decay angle in the rest frame of the given `resonance`
515///
516/// This Variable is calculated by forming the given frame (helicity or Gottfried-Jackson) and
517/// calculating the spherical angles according to one of the decaying `daughter` particles.
518///
519/// The helicity frame is defined in terms of the following Cartesian axes in the rest frame of
520/// the `resonance`:
521///
522/// .. math:: \hat{z} \propto -\vec{p}'_{\text{recoil}}
523/// .. math:: \hat{y} \propto \vec{p}_{\text{beam}} \times (-\vec{p}_{\text{recoil}})
524/// .. math:: \hat{x} = \hat{y} \times \hat{z}
525///
526/// where primed vectors are in the rest frame of the `resonance` and unprimed vectors are in
527/// the center-of-momentum frame.
528///
529/// The Gottfried-Jackson frame differs only in the definition of :math:`\hat{z}`:
530///
531/// .. math:: \hat{z} \propto \vec{p}'_{\text{beam}}
532///
533/// Parameters
534/// ----------
535/// reaction : laddu.Reaction
536///     Reaction describing the production kinematics and decay roots.
537/// daughter : list of str
538///     Names of particles which are combined to form one of the decay products of the
539///     resonance associated with the decay parent.
540/// frame : {'Helicity', 'HX', 'HEL', 'GottfriedJackson', 'Gottfried Jackson', 'GJ', 'Gottfried-Jackson'}
541///     The frame to use in the  calculation
542///
543/// Raises
544/// ------
545/// ValueError
546///     If `frame` is not one of the valid options
547///
548/// See Also
549/// --------
550/// laddu.utils.vectors.Vec3.costheta
551///
552#[pyclass(name = "CosTheta", module = "laddu", from_py_object)]
553#[derive(Clone, Serialize, Deserialize)]
554pub struct PyCosTheta(pub CosTheta);
555
556#[pymethods]
557impl PyCosTheta {
558    /// The value of this Variable for the given Event
559    ///
560    /// Parameters
561    /// ----------
562    /// event : Event
563    ///     The Event upon which the Variable is calculated
564    ///
565    /// Returns
566    /// -------
567    /// value : float
568    ///     The value of the Variable for the given `event`
569    ///
570    fn value(&self, event: &PyEvent) -> PyResult<f64> {
571        let metadata = event
572            .metadata_opt()
573            .ok_or_else(|| PyValueError::new_err(
574                "This event is not associated with metadata; supply `p4_names`/`aux_names` when constructing it or evaluate via a Dataset.",
575            ))?;
576        let mut variable = self.0.clone();
577        variable.bind(metadata).map_err(PyErr::from)?;
578        let dataset =
579            Dataset::new_with_metadata(vec![event.event.data_arc()], event.event.metadata_arc());
580        let event_view = dataset.event_view(0);
581        Ok(variable.value(&event_view))
582    }
583    /// All values of this Variable on the given Dataset
584    ///
585    /// Parameters
586    /// ----------
587    /// dataset : Dataset
588    ///     The Dataset upon which the Variable is calculated
589    ///
590    /// Returns
591    /// -------
592    /// values : array_like
593    ///     The values of the Variable for each Event in the given `dataset`
594    ///
595    fn value_on<'py>(
596        &self,
597        py: Python<'py>,
598        dataset: &PyDataset,
599    ) -> PyResult<Bound<'py, PyArray1<f64>>> {
600        let values = self.0.value_on(&dataset.0).map_err(PyErr::from)?;
601        Ok(PyArray1::from_vec(py, values))
602    }
603    fn __eq__(&self, value: f64) -> PyVariableExpression {
604        PyVariableExpression(self.0.eq(value))
605    }
606    fn __lt__(&self, value: f64) -> PyVariableExpression {
607        PyVariableExpression(self.0.lt(value))
608    }
609    fn __gt__(&self, value: f64) -> PyVariableExpression {
610        PyVariableExpression(self.0.gt(value))
611    }
612    fn __le__(&self, value: f64) -> PyVariableExpression {
613        PyVariableExpression(self.0.le(value))
614    }
615    fn __ge__(&self, value: f64) -> PyVariableExpression {
616        PyVariableExpression(self.0.ge(value))
617    }
618    fn __repr__(&self) -> String {
619        format!("{:?}", self.0)
620    }
621    fn __str__(&self) -> String {
622        format!("{}", self.0)
623    }
624}
625
626/// The aziumuthal decay angle in the rest frame of the given `resonance`
627///
628/// This Variable is calculated by forming the given frame (helicity or Gottfried-Jackson) and
629/// calculating the spherical angles according to one of the decaying `daughter` particles.
630///
631/// The helicity frame is defined in terms of the following Cartesian axes in the rest frame of
632/// the `resonance`:
633///
634/// .. math:: \hat{z} \propto -\vec{p}'_{\text{recoil}}
635/// .. math:: \hat{y} \propto \vec{p}_{\text{beam}} \times (-\vec{p}_{\text{recoil}})
636/// .. math:: \hat{x} = \hat{y} \times \hat{z}
637///
638/// where primed vectors are in the rest frame of the `resonance` and unprimed vectors are in
639/// the center-of-momentum frame.
640///
641/// The Gottfried-Jackson frame differs only in the definition of :math:`\hat{z}`:
642///
643/// .. math:: \hat{z} \propto \vec{p}'_{\text{beam}}
644///
645/// Parameters
646/// ----------
647/// reaction : laddu.Reaction
648///     Reaction describing the production kinematics and decay roots.
649/// daughter : list of str
650///     Names of particles which are combined to form one of the decay products of the
651///     resonance associated with the decay parent.
652/// frame : {'Helicity', 'HX', 'HEL', 'GottfriedJackson', 'Gottfried Jackson', 'GJ', 'Gottfried-Jackson'}
653///     The frame to use in the  calculation
654///
655/// Raises
656/// ------
657/// ValueError
658///     If `frame` is not one of the valid options
659///
660///
661/// See Also
662/// --------
663/// laddu.utils.vectors.Vec3.phi
664///
665#[pyclass(name = "Phi", module = "laddu", from_py_object)]
666#[derive(Clone, Serialize, Deserialize)]
667pub struct PyPhi(pub Phi);
668
669#[pymethods]
670impl PyPhi {
671    /// The value of this Variable for the given Event
672    ///
673    /// Parameters
674    /// ----------
675    /// event : Event
676    ///     The Event upon which the Variable is calculated
677    ///
678    /// Returns
679    /// -------
680    /// value : float
681    ///     The value of the Variable for the given `event`
682    ///
683    fn value(&self, event: &PyEvent) -> PyResult<f64> {
684        let metadata = event
685            .metadata_opt()
686            .ok_or_else(|| PyValueError::new_err(
687                "This event is not associated with metadata; supply `p4_names`/`aux_names` when constructing it or evaluate via a Dataset.",
688            ))?;
689        let mut variable = self.0.clone();
690        variable.bind(metadata).map_err(PyErr::from)?;
691        let dataset =
692            Dataset::new_with_metadata(vec![event.event.data_arc()], event.event.metadata_arc());
693        let event_view = dataset.event_view(0);
694        Ok(variable.value(&event_view))
695    }
696    /// All values of this Variable on the given Dataset
697    ///
698    /// Parameters
699    /// ----------
700    /// dataset : Dataset
701    ///     The Dataset upon which the Variable is calculated
702    ///
703    /// Returns
704    /// -------
705    /// values : array_like
706    ///     The values of the Variable for each Event in the given `dataset`
707    ///
708    fn value_on<'py>(
709        &self,
710        py: Python<'py>,
711        dataset: &PyDataset,
712    ) -> PyResult<Bound<'py, PyArray1<f64>>> {
713        let values = self.0.value_on(&dataset.0).map_err(PyErr::from)?;
714        Ok(PyArray1::from_vec(py, values))
715    }
716    fn __eq__(&self, value: f64) -> PyVariableExpression {
717        PyVariableExpression(self.0.eq(value))
718    }
719    fn __lt__(&self, value: f64) -> PyVariableExpression {
720        PyVariableExpression(self.0.lt(value))
721    }
722    fn __gt__(&self, value: f64) -> PyVariableExpression {
723        PyVariableExpression(self.0.gt(value))
724    }
725    fn __le__(&self, value: f64) -> PyVariableExpression {
726        PyVariableExpression(self.0.le(value))
727    }
728    fn __ge__(&self, value: f64) -> PyVariableExpression {
729        PyVariableExpression(self.0.ge(value))
730    }
731    fn __repr__(&self) -> String {
732        format!("{:?}", self.0)
733    }
734    fn __str__(&self) -> String {
735        format!("{}", self.0)
736    }
737}
738
739/// A Variable used to define both spherical decay angles in the given frame
740///
741/// This class combines ``laddu.CosTheta`` and ``laddu.Phi`` into a single
742/// object
743///
744/// Parameters
745/// ----------
746/// reaction : laddu.Reaction
747///     Reaction describing the production kinematics and decay roots.
748/// daughter : list of str
749///     Names of particles which are combined to form one of the decay products of the
750///     resonance associated with the decay parent.
751/// frame : {'Helicity', 'HX', 'HEL', 'GottfriedJackson', 'Gottfried Jackson', 'GJ', 'Gottfried-Jackson'}
752///     The frame to use in the  calculation
753///
754/// Raises
755/// ------
756/// ValueError
757///     If `frame` is not one of the valid options
758///
759/// See Also
760/// --------
761/// laddu.CosTheta
762/// laddu.Phi
763///
764#[pyclass(name = "Angles", module = "laddu", skip_from_py_object)]
765#[derive(Clone)]
766pub struct PyAngles(pub Angles);
767#[pymethods]
768impl PyAngles {
769    /// The Variable representing the cosine of the polar spherical decay angle
770    ///
771    /// Returns
772    /// -------
773    /// CosTheta
774    ///
775    #[getter]
776    fn costheta(&self) -> PyCosTheta {
777        PyCosTheta(self.0.costheta.clone())
778    }
779    // The Variable representing the polar azimuthal decay angle
780    //
781    // Returns
782    // -------
783    // Phi
784    //
785    #[getter]
786    fn phi(&self) -> PyPhi {
787        PyPhi(self.0.phi.clone())
788    }
789    fn __repr__(&self) -> String {
790        format!("{:?}", self.0)
791    }
792    fn __str__(&self) -> String {
793        format!("{}", self.0)
794    }
795}
796
797/// The polar angle of the given polarization vector with respect to the production plane
798///
799/// The `beam` and `recoil` particles define the plane of production, and this Variable
800/// describes the polar angle of the `beam` relative to this plane
801///
802/// Parameters
803/// ----------
804/// reaction : laddu.Reaction
805///     Reaction describing the production kinematics and decay roots.
806/// pol_angle : str
807///     Name of the auxiliary scalar column storing the polarization angle in radians
808///
809#[pyclass(name = "PolAngle", module = "laddu", from_py_object)]
810#[derive(Clone, Serialize, Deserialize)]
811pub struct PyPolAngle(pub PolAngle);
812
813#[pymethods]
814impl PyPolAngle {
815    #[new]
816    fn new(reaction: PyReaction, pol_angle: String) -> Self {
817        Self(PolAngle::new(reaction.0.clone(), pol_angle))
818    }
819    /// The value of this Variable for the given Event
820    ///
821    /// Parameters
822    /// ----------
823    /// event : Event
824    ///     The Event upon which the Variable is calculated
825    ///
826    /// Returns
827    /// -------
828    /// value : float
829    ///     The value of the Variable for the given `event`
830    ///
831    fn value(&self, event: &PyEvent) -> PyResult<f64> {
832        let metadata = event
833            .metadata_opt()
834            .ok_or_else(|| PyValueError::new_err(
835                "This event is not associated with metadata; supply `p4_names`/`aux_names` when constructing it or evaluate via a Dataset.",
836            ))?;
837        let mut variable = self.0.clone();
838        variable.bind(metadata).map_err(PyErr::from)?;
839        let dataset =
840            Dataset::new_with_metadata(vec![event.event.data_arc()], event.event.metadata_arc());
841        let event_view = dataset.event_view(0);
842        Ok(variable.value(&event_view))
843    }
844    /// All values of this Variable on the given Dataset
845    ///
846    /// Parameters
847    /// ----------
848    /// dataset : Dataset
849    ///     The Dataset upon which the Variable is calculated
850    ///
851    /// Returns
852    /// -------
853    /// values : array_like
854    ///     The values of the Variable for each Event in the given `dataset`
855    ///
856    fn value_on<'py>(
857        &self,
858        py: Python<'py>,
859        dataset: &PyDataset,
860    ) -> PyResult<Bound<'py, PyArray1<f64>>> {
861        let values = self.0.value_on(&dataset.0).map_err(PyErr::from)?;
862        Ok(PyArray1::from_vec(py, values))
863    }
864    fn __eq__(&self, value: f64) -> PyVariableExpression {
865        PyVariableExpression(self.0.eq(value))
866    }
867    fn __lt__(&self, value: f64) -> PyVariableExpression {
868        PyVariableExpression(self.0.lt(value))
869    }
870    fn __gt__(&self, value: f64) -> PyVariableExpression {
871        PyVariableExpression(self.0.gt(value))
872    }
873    fn __le__(&self, value: f64) -> PyVariableExpression {
874        PyVariableExpression(self.0.le(value))
875    }
876    fn __ge__(&self, value: f64) -> PyVariableExpression {
877        PyVariableExpression(self.0.ge(value))
878    }
879    fn __repr__(&self) -> String {
880        format!("{:?}", self.0)
881    }
882    fn __str__(&self) -> String {
883        format!("{}", self.0)
884    }
885}
886
887/// The magnitude of the given particle's polarization vector
888///
889/// This Variable simply represents the magnitude of the polarization vector of the particle
890/// with the index `beam`
891///
892/// Parameters
893/// ----------
894/// pol_magnitude : str
895///     Name of the auxiliary scalar column storing the magnitude of the polarization vector
896///
897/// See Also
898/// --------
899/// laddu.utils.vectors.Vec3.mag
900///
901#[pyclass(name = "PolMagnitude", module = "laddu", from_py_object)]
902#[derive(Clone, Serialize, Deserialize)]
903pub struct PyPolMagnitude(pub PolMagnitude);
904
905#[pymethods]
906impl PyPolMagnitude {
907    #[new]
908    fn new(pol_magnitude: String) -> Self {
909        Self(PolMagnitude::new(pol_magnitude))
910    }
911    /// The value of this Variable for the given Event
912    ///
913    /// Parameters
914    /// ----------
915    /// event : Event
916    ///     The Event upon which the Variable is calculated
917    ///
918    /// Returns
919    /// -------
920    /// value : float
921    ///     The value of the Variable for the given `event`
922    ///
923    fn value(&self, event: &PyEvent) -> PyResult<f64> {
924        let metadata = event
925            .metadata_opt()
926            .ok_or_else(|| PyValueError::new_err(
927                "This event is not associated with metadata; supply `p4_names`/`aux_names` when constructing it or evaluate via a Dataset.",
928            ))?;
929        let mut variable = self.0.clone();
930        variable.bind(metadata).map_err(PyErr::from)?;
931        let dataset =
932            Dataset::new_with_metadata(vec![event.event.data_arc()], event.event.metadata_arc());
933        let event_view = dataset.event_view(0);
934        Ok(variable.value(&event_view))
935    }
936    /// All values of this Variable on the given Dataset
937    ///
938    /// Parameters
939    /// ----------
940    /// dataset : Dataset
941    ///     The Dataset upon which the Variable is calculated
942    ///
943    /// Returns
944    /// -------
945    /// values : array_like
946    ///     The values of the Variable for each Event in the given `dataset`
947    ///
948    fn value_on<'py>(
949        &self,
950        py: Python<'py>,
951        dataset: &PyDataset,
952    ) -> PyResult<Bound<'py, PyArray1<f64>>> {
953        let values = self.0.value_on(&dataset.0).map_err(PyErr::from)?;
954        Ok(PyArray1::from_vec(py, values))
955    }
956    fn __eq__(&self, value: f64) -> PyVariableExpression {
957        PyVariableExpression(self.0.eq(value))
958    }
959    fn __lt__(&self, value: f64) -> PyVariableExpression {
960        PyVariableExpression(self.0.lt(value))
961    }
962    fn __gt__(&self, value: f64) -> PyVariableExpression {
963        PyVariableExpression(self.0.gt(value))
964    }
965    fn __le__(&self, value: f64) -> PyVariableExpression {
966        PyVariableExpression(self.0.le(value))
967    }
968    fn __ge__(&self, value: f64) -> PyVariableExpression {
969        PyVariableExpression(self.0.ge(value))
970    }
971    fn __repr__(&self) -> String {
972        format!("{:?}", self.0)
973    }
974    fn __str__(&self) -> String {
975        format!("{}", self.0)
976    }
977}
978
979/// A Variable used to define both the polarization angle and magnitude of the given particle``
980///
981/// This class combines ``laddu.PolAngle`` and ``laddu.PolMagnitude`` into a single
982/// object
983///
984/// Parameters
985/// ----------
986/// reaction : laddu.Reaction
987///     Reaction describing the production kinematics and decay roots.
988/// pol_magnitude : str
989///     Name of the auxiliary scalar storing the polarization magnitude
990/// pol_angle : str
991///     Name of the auxiliary scalar storing the polarization angle in radians
992///
993/// See Also
994/// --------
995/// laddu.PolAngle
996/// laddu.PolMagnitude
997///
998#[pyclass(name = "Polarization", module = "laddu", skip_from_py_object)]
999#[derive(Clone)]
1000pub struct PyPolarization(pub Polarization);
1001#[pymethods]
1002impl PyPolarization {
1003    #[new]
1004    #[pyo3(signature=(reaction, *, pol_magnitude, pol_angle))]
1005    fn new(reaction: PyReaction, pol_magnitude: String, pol_angle: String) -> PyResult<Self> {
1006        if pol_magnitude == pol_angle {
1007            return Err(PyValueError::new_err(
1008                "`pol_magnitude` and `pol_angle` must reference distinct auxiliary columns",
1009            ));
1010        }
1011        let polarization = Polarization::new(reaction.0.clone(), pol_magnitude, pol_angle);
1012        Ok(PyPolarization(polarization))
1013    }
1014    /// The Variable representing the magnitude of the polarization vector
1015    ///
1016    /// Returns
1017    /// -------
1018    /// PolMagnitude
1019    ///
1020    #[getter]
1021    fn pol_magnitude(&self) -> PyPolMagnitude {
1022        PyPolMagnitude(self.0.pol_magnitude.clone())
1023    }
1024    /// The Variable representing the polar angle of the polarization vector
1025    ///
1026    /// Returns
1027    /// -------
1028    /// PolAngle
1029    ///
1030    #[getter]
1031    fn pol_angle(&self) -> PyPolAngle {
1032        PyPolAngle(self.0.pol_angle.clone())
1033    }
1034    fn __repr__(&self) -> String {
1035        format!("{:?}", self.0)
1036    }
1037    fn __str__(&self) -> String {
1038        format!("{}", self.0)
1039    }
1040}
1041
1042/// Mandelstam variables s, t, and u
1043///
1044/// By convention, the metric is chosen to be :math:`(+---)` and the variables are defined as follows
1045/// (ignoring factors of :math:`c`):
1046///
1047/// .. math:: s = (p_1 + p_2)^2 = (p_3 + p_4)^2
1048///
1049/// .. math:: t = (p_1 - p_3)^2 = (p_4 - p_2)^2
1050///
1051/// .. math:: u = (p_1 - p_4)^2 = (p_3 - p_2)^2
1052///
1053/// Parameters
1054/// ----------
1055/// reaction : laddu.Reaction
1056///     Reaction describing the two-to-two kinematics whose Mandelstam channels should be evaluated.
1057/// channel: {'s', 't', 'u', 'S', 'T', 'U'}
1058///     The Mandelstam channel to calculate
1059///
1060/// Raises
1061/// ------
1062/// Exception
1063///     If more than one particle list is empty
1064/// ValueError
1065///     If `channel` is not one of the valid options
1066///
1067/// Notes
1068/// -----
1069/// ///
1070#[pyclass(name = "Mandelstam", module = "laddu", from_py_object)]
1071#[derive(Clone, Serialize, Deserialize)]
1072pub struct PyMandelstam(pub Mandelstam);
1073
1074#[pymethods]
1075impl PyMandelstam {
1076    #[new]
1077    fn new(reaction: PyReaction, channel: &str) -> PyResult<Self> {
1078        Ok(Self(Mandelstam::new(reaction.0.clone(), channel.parse()?)))
1079    }
1080    /// The value of this Variable for the given Event
1081    ///
1082    /// Parameters
1083    /// ----------
1084    /// event : Event
1085    ///     The Event upon which the Variable is calculated
1086    ///
1087    /// Returns
1088    /// -------
1089    /// value : float
1090    ///     The value of the Variable for the given `event`
1091    ///
1092    fn value(&self, event: &PyEvent) -> PyResult<f64> {
1093        let metadata = event
1094            .metadata_opt()
1095            .ok_or_else(|| PyValueError::new_err(
1096                "This event is not associated with metadata; supply `p4_names`/`aux_names` when constructing it or evaluate via a Dataset.",
1097            ))?;
1098        let mut variable = self.0.clone();
1099        variable.bind(metadata).map_err(PyErr::from)?;
1100        let dataset =
1101            Dataset::new_with_metadata(vec![event.event.data_arc()], event.event.metadata_arc());
1102        let event_view = dataset.event_view(0);
1103        Ok(variable.value(&event_view))
1104    }
1105    /// All values of this Variable on the given Dataset
1106    ///
1107    /// Parameters
1108    /// ----------
1109    /// dataset : Dataset
1110    ///     The Dataset upon which the Variable is calculated
1111    ///
1112    /// Returns
1113    /// -------
1114    /// values : array_like
1115    ///     The values of the Variable for each Event in the given `dataset`
1116    ///
1117    fn value_on<'py>(
1118        &self,
1119        py: Python<'py>,
1120        dataset: &PyDataset,
1121    ) -> PyResult<Bound<'py, PyArray1<f64>>> {
1122        let values = self.0.value_on(&dataset.0).map_err(PyErr::from)?;
1123        Ok(PyArray1::from_vec(py, values))
1124    }
1125    fn __eq__(&self, value: f64) -> PyVariableExpression {
1126        PyVariableExpression(self.0.eq(value))
1127    }
1128    fn __lt__(&self, value: f64) -> PyVariableExpression {
1129        PyVariableExpression(self.0.lt(value))
1130    }
1131    fn __gt__(&self, value: f64) -> PyVariableExpression {
1132        PyVariableExpression(self.0.gt(value))
1133    }
1134    fn __le__(&self, value: f64) -> PyVariableExpression {
1135        PyVariableExpression(self.0.le(value))
1136    }
1137    fn __ge__(&self, value: f64) -> PyVariableExpression {
1138        PyVariableExpression(self.0.ge(value))
1139    }
1140    fn __repr__(&self) -> String {
1141        format!("{:?}", self.0)
1142    }
1143    fn __str__(&self) -> String {
1144        format!("{}", self.0)
1145    }
1146}
1147
1148#[typetag::serde]
1149impl Variable for PyVariable {
1150    fn bind(&mut self, metadata: &DatasetMetadata) -> LadduResult<()> {
1151        match self {
1152            PyVariable::Mass(mass) => mass.0.bind(metadata),
1153            PyVariable::CosTheta(cos_theta) => cos_theta.0.bind(metadata),
1154            PyVariable::Phi(phi) => phi.0.bind(metadata),
1155            PyVariable::PolAngle(pol_angle) => pol_angle.0.bind(metadata),
1156            PyVariable::PolMagnitude(pol_magnitude) => pol_magnitude.0.bind(metadata),
1157            PyVariable::Mandelstam(mandelstam) => mandelstam.0.bind(metadata),
1158        }
1159    }
1160
1161    fn value_on(&self, dataset: &Dataset) -> LadduResult<Vec<f64>> {
1162        match self {
1163            PyVariable::Mass(mass) => mass.0.value_on(dataset),
1164            PyVariable::CosTheta(cos_theta) => cos_theta.0.value_on(dataset),
1165            PyVariable::Phi(phi) => phi.0.value_on(dataset),
1166            PyVariable::PolAngle(pol_angle) => pol_angle.0.value_on(dataset),
1167            PyVariable::PolMagnitude(pol_magnitude) => pol_magnitude.0.value_on(dataset),
1168            PyVariable::Mandelstam(mandelstam) => mandelstam.0.value_on(dataset),
1169        }
1170    }
1171
1172    fn value(&self, event: &NamedEventView<'_>) -> f64 {
1173        match self {
1174            PyVariable::Mass(mass) => mass.0.value(event),
1175            PyVariable::CosTheta(cos_theta) => cos_theta.0.value(event),
1176            PyVariable::Phi(phi) => phi.0.value(event),
1177            PyVariable::PolAngle(pol_angle) => pol_angle.0.value(event),
1178            PyVariable::PolMagnitude(pol_magnitude) => pol_magnitude.0.value(event),
1179            PyVariable::Mandelstam(mandelstam) => mandelstam.0.value(event),
1180        }
1181    }
1182}
laddu_python/utils/variables.rs

laddu_python/utils/
variables.rs
