laddu_python/
variables.rs

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

laddu_python/
variables.rs
