laddu_python/
variables.rs

1use std::fmt::{Debug, Display};
2
3use laddu_amplitudes::{DecayAmplitudeExt, ProductionAmplitudeExt};
4use laddu_core::{
5    data::{Dataset, DatasetMetadata, EventLike, OwnedEvent},
6    reaction::{Decay, Particle, Production, 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 two-to-two production view.
214    fn production(&self) -> PyResult<PyProduction> {
215        Ok(PyProduction(self.0.production()?))
216    }
217
218    /// Construct a Mandelstam variable.
219    fn mandelstam(&self, channel: &str) -> PyResult<PyMandelstam> {
220        Ok(PyMandelstam(self.0.mandelstam(channel.parse()?)?))
221    }
222
223    /// Construct a polarization-angle variable.
224    fn pol_angle(&self, angle_aux: String) -> PyPolAngle {
225        PyPolAngle(self.0.pol_angle(angle_aux))
226    }
227
228    /// Construct polarization variables.
229    #[pyo3(signature=(*, pol_magnitude, pol_angle))]
230    fn polarization(&self, pol_magnitude: String, pol_angle: String) -> PyResult<PyPolarization> {
231        if pol_magnitude == pol_angle {
232            return Err(PyValueError::new_err(
233                "`pol_magnitude` and `pol_angle` must reference distinct auxiliary columns",
234            ));
235        }
236        Ok(PyPolarization(
237            self.0.polarization(pol_magnitude, pol_angle),
238        ))
239    }
240
241    fn __repr__(&self) -> String {
242        format!("{:?}", self.0)
243    }
244
245    fn __str__(&self) -> String {
246        format!("{:?}", self.0)
247    }
248}
249
250/// A reaction-aware isobar decay view.
251#[pyclass(name = "Decay", module = "laddu", from_py_object)]
252#[derive(Clone, Serialize, Deserialize)]
253pub struct PyDecay(pub Decay);
254
255#[pymethods]
256impl PyDecay {
257    /// The enclosing reaction.
258    #[getter]
259    fn reaction(&self) -> PyReaction {
260        PyReaction(self.0.reaction().clone())
261    }
262
263    /// The parent particle.
264    #[getter]
265    fn parent(&self) -> String {
266        self.0.parent().to_string()
267    }
268
269    /// The first daughter particle identifier.
270    #[getter]
271    fn daughter_1(&self) -> String {
272        self.0.daughter_1().to_string()
273    }
274
275    /// The second daughter particle identifier.
276    #[getter]
277    fn daughter_2(&self) -> String {
278        self.0.daughter_2().to_string()
279    }
280
281    /// Ordered daughter particle identifiers.
282    fn daughters(&self) -> Vec<String> {
283        self.0.daughters().into_iter().map(str::to_string).collect()
284    }
285
286    /// Parent mass variable.
287    fn mass(&self) -> PyMass {
288        PyMass(self.0.mass())
289    }
290
291    /// Parent mass variable.
292    fn parent_mass(&self) -> PyMass {
293        PyMass(self.0.parent_mass())
294    }
295
296    /// First daughter mass variable.
297    fn daughter_1_mass(&self) -> PyMass {
298        PyMass(self.0.daughter_1_mass())
299    }
300
301    /// Second daughter mass variable.
302    fn daughter_2_mass(&self) -> PyMass {
303        PyMass(self.0.daughter_2_mass())
304    }
305
306    /// Mass variable for a selected daughter.
307    fn daughter_mass(&self, daughter: &str) -> PyResult<PyMass> {
308        Ok(PyMass(self.0.daughter_mass(daughter)?))
309    }
310
311    /// Decay costheta variable for the selected frame.
312    #[pyo3(signature=(daughter, frame="Helicity"))]
313    fn costheta(&self, daughter: &str, frame: &str) -> PyResult<PyCosTheta> {
314        Ok(PyCosTheta(self.0.costheta(daughter, frame.parse()?)?))
315    }
316
317    /// Decay phi variable for the selected frame.
318    #[pyo3(signature=(daughter, frame="Helicity"))]
319    fn phi(&self, daughter: &str, frame: &str) -> PyResult<PyPhi> {
320        Ok(PyPhi(self.0.phi(daughter, frame.parse()?)?))
321    }
322
323    /// Decay angle variables for the selected frame.
324    #[pyo3(signature=(daughter, frame="Helicity"))]
325    fn angles(&self, daughter: &str, frame: &str) -> PyResult<PyAngles> {
326        Ok(PyAngles(self.0.angles(daughter, frame.parse()?)?))
327    }
328
329    /// Construct the helicity-basis angular factor for one explicit helicity term.
330    #[pyo3(signature=(*tags, spin, projection, daughter, lambda_1, lambda_2, frame="Helicity"))]
331    #[allow(clippy::too_many_arguments)]
332    fn helicity_factor(
333        &self,
334        tags: &Bound<'_, PyTuple>,
335        spin: &Bound<'_, PyAny>,
336        projection: &Bound<'_, PyAny>,
337        daughter: &str,
338        lambda_1: &Bound<'_, PyAny>,
339        lambda_2: &Bound<'_, PyAny>,
340        frame: &str,
341    ) -> PyResult<PyExpression> {
342        Ok(PyExpression(self.0.helicity_factor(
343            py_tags(tags)?,
344            parse_angular_momentum(spin)?,
345            parse_projection(projection)?,
346            daughter,
347            parse_projection(lambda_1)?,
348            parse_projection(lambda_2)?,
349            frame.parse()?,
350        )?))
351    }
352
353    /// Construct the canonical-basis spin-angular factor for one explicit LS/helicity term.
354    #[pyo3(signature=(*tags, spin, projection, orbital_l, coupled_spin, daughter, daughter_1_spin, daughter_2_spin, lambda_1, lambda_2, frame="Helicity"))]
355    #[allow(clippy::too_many_arguments)]
356    fn canonical_factor(
357        &self,
358        tags: &Bound<'_, PyTuple>,
359        spin: &Bound<'_, PyAny>,
360        projection: &Bound<'_, PyAny>,
361        orbital_l: &Bound<'_, PyAny>,
362        coupled_spin: &Bound<'_, PyAny>,
363        daughter: &str,
364        daughter_1_spin: &Bound<'_, PyAny>,
365        daughter_2_spin: &Bound<'_, PyAny>,
366        lambda_1: &Bound<'_, PyAny>,
367        lambda_2: &Bound<'_, PyAny>,
368        frame: &str,
369    ) -> PyResult<PyExpression> {
370        Ok(PyExpression(self.0.canonical_factor(
371            py_tags(tags)?,
372            parse_angular_momentum(spin)?,
373            parse_projection(projection)?,
374            parse_orbital_angular_momentum(orbital_l)?,
375            parse_angular_momentum(coupled_spin)?,
376            daughter,
377            parse_angular_momentum(daughter_1_spin)?,
378            parse_angular_momentum(daughter_2_spin)?,
379            parse_projection(lambda_1)?,
380            parse_projection(lambda_2)?,
381            frame.parse()?,
382        )?))
383    }
384
385    fn __repr__(&self) -> String {
386        format!("{:?}", self.0)
387    }
388
389    fn __str__(&self) -> String {
390        format!("{:?}", self.0)
391    }
392}
393
394/// A reaction-aware two-to-two production view.
395#[pyclass(name = "Production", module = "laddu", from_py_object)]
396#[derive(Clone, Serialize, Deserialize)]
397pub struct PyProduction(pub Production);
398
399#[pymethods]
400impl PyProduction {
401    /// The enclosing reaction.
402    #[getter]
403    fn reaction(&self) -> PyReaction {
404        PyReaction(self.0.reaction().clone())
405    }
406
407    /// The produced-system particle.
408    #[getter]
409    fn produced(&self) -> String {
410        self.0.produced().to_string()
411    }
412
413    /// The recoil particle.
414    #[getter]
415    fn recoil(&self) -> String {
416        self.0.recoil().to_string()
417    }
418
419    /// Production costheta variable.
420    fn costheta(&self) -> PyResult<PyCosTheta> {
421        Ok(PyCosTheta(self.0.costheta()?))
422    }
423
424    /// Production phi variable.
425    fn phi(&self) -> PyResult<PyPhi> {
426        Ok(PyPhi(self.0.phi()?))
427    }
428
429    /// Production angle variables.
430    fn angles(&self) -> PyResult<PyAngles> {
431        Ok(PyAngles(self.0.angles()?))
432    }
433
434    /// Construct the helicity-basis production angular factor for one explicit helicity term.
435    #[pyo3(signature=(*tags, spin, projection, lambda_produced, lambda_recoil))]
436    #[allow(clippy::too_many_arguments)]
437    fn helicity_factor(
438        &self,
439        tags: &Bound<'_, PyTuple>,
440        spin: &Bound<'_, PyAny>,
441        projection: &Bound<'_, PyAny>,
442        lambda_produced: &Bound<'_, PyAny>,
443        lambda_recoil: &Bound<'_, PyAny>,
444    ) -> PyResult<PyExpression> {
445        Ok(PyExpression(self.0.helicity_factor(
446            py_tags(tags)?,
447            parse_angular_momentum(spin)?,
448            parse_projection(projection)?,
449            parse_projection(lambda_produced)?,
450            parse_projection(lambda_recoil)?,
451        )?))
452    }
453
454    /// Construct the canonical-basis production spin-angular factor for one LS/helicity term.
455    #[pyo3(signature=(*tags, spin, projection, orbital_l, coupled_spin, produced_spin, recoil_spin, lambda_produced, lambda_recoil))]
456    #[allow(clippy::too_many_arguments)]
457    fn canonical_factor(
458        &self,
459        tags: &Bound<'_, PyTuple>,
460        spin: &Bound<'_, PyAny>,
461        projection: &Bound<'_, PyAny>,
462        orbital_l: &Bound<'_, PyAny>,
463        coupled_spin: &Bound<'_, PyAny>,
464        produced_spin: &Bound<'_, PyAny>,
465        recoil_spin: &Bound<'_, PyAny>,
466        lambda_produced: &Bound<'_, PyAny>,
467        lambda_recoil: &Bound<'_, PyAny>,
468    ) -> PyResult<PyExpression> {
469        Ok(PyExpression(self.0.canonical_factor(
470            py_tags(tags)?,
471            parse_angular_momentum(spin)?,
472            parse_projection(projection)?,
473            parse_orbital_angular_momentum(orbital_l)?,
474            parse_angular_momentum(coupled_spin)?,
475            parse_angular_momentum(produced_spin)?,
476            parse_angular_momentum(recoil_spin)?,
477            parse_projection(lambda_produced)?,
478            parse_projection(lambda_recoil)?,
479        )?))
480    }
481
482    fn __repr__(&self) -> String {
483        format!("{:?}", self.0)
484    }
485
486    fn __str__(&self) -> String {
487        format!("{:?}", self.0)
488    }
489}
490
491/// The invariant mass of an arbitrary combination of constituent particles in an Event
492///
493/// This variable is calculated by summing up the 4-momenta of each particle listed by index in
494/// `constituents` and taking the invariant magnitude of the resulting 4-vector.
495///
496/// Parameters
497/// ----------
498/// constituents : str or list of str
499///     Particle names to combine when constructing the final four-momentum
500///
501/// See Also
502/// --------
503/// laddu.utils.vectors.Vec4.m
504///
505#[pyclass(name = "Mass", module = "laddu", from_py_object)]
506#[derive(Clone, Serialize, Deserialize)]
507pub struct PyMass(pub Mass);
508
509#[pymethods]
510impl PyMass {
511    #[new]
512    fn new(constituents: PyP4SelectionInput) -> Self {
513        Self(Mass::new(constituents.into_selection()))
514    }
515    /// The value of this Variable for the given Event
516    ///
517    /// Parameters
518    /// ----------
519    /// event : Event
520    ///     The Event upon which the Variable is calculated
521    ///
522    /// Returns
523    /// -------
524    /// value : float
525    ///     The value of the Variable for the given `event`
526    ///
527    fn value(&self, event: &PyEvent) -> PyResult<f64> {
528        let metadata = event
529            .metadata_opt()
530            .ok_or_else(|| PyValueError::new_err(
531                "This event is not associated with metadata; supply `p4_names`/`aux_names` when constructing it or evaluate via a Dataset.",
532            ))?;
533        let mut variable = self.0.clone();
534        variable.bind(metadata).map_err(PyErr::from)?;
535        Ok(variable.value(&event.event))
536    }
537    /// All values of this Variable on the given Dataset
538    ///
539    /// Parameters
540    /// ----------
541    /// dataset : Dataset
542    ///     The Dataset upon which the Variable is calculated
543    ///
544    /// Returns
545    /// -------
546    /// values : array_like
547    ///     The values of the Variable for each Event in the given `dataset`
548    ///
549    fn value_on<'py>(
550        &self,
551        py: Python<'py>,
552        dataset: &PyDataset,
553    ) -> PyResult<Bound<'py, PyArray1<f64>>> {
554        let values = self.0.value_on(&dataset.0).map_err(PyErr::from)?;
555        Ok(PyArray1::from_vec(py, values))
556    }
557    fn __eq__(&self, value: f64) -> PyVariableExpression {
558        PyVariableExpression(self.0.eq(value))
559    }
560    fn __lt__(&self, value: f64) -> PyVariableExpression {
561        PyVariableExpression(self.0.lt(value))
562    }
563    fn __gt__(&self, value: f64) -> PyVariableExpression {
564        PyVariableExpression(self.0.gt(value))
565    }
566    fn __le__(&self, value: f64) -> PyVariableExpression {
567        PyVariableExpression(self.0.le(value))
568    }
569    fn __ge__(&self, value: f64) -> PyVariableExpression {
570        PyVariableExpression(self.0.ge(value))
571    }
572    fn __repr__(&self) -> String {
573        format!("{:?}", self.0)
574    }
575    fn __str__(&self) -> String {
576        format!("{}", self.0)
577    }
578}
579
580/// The cosine of the polar decay angle in the rest frame of the given `resonance`
581///
582/// This Variable is calculated by forming the given frame (helicity or Gottfried-Jackson) and
583/// calculating the spherical angles according to one of the decaying `daughter` particles.
584///
585/// The helicity frame is defined in terms of the following Cartesian axes in the rest frame of
586/// the `resonance`:
587///
588/// .. math:: \hat{z} \propto -\vec{p}'_{\text{recoil}}
589/// .. math:: \hat{y} \propto \vec{p}_{\text{beam}} \times (-\vec{p}_{\text{recoil}})
590/// .. math:: \hat{x} = \hat{y} \times \hat{z}
591///
592/// where primed vectors are in the rest frame of the `resonance` and unprimed vectors are in
593/// the center-of-momentum frame.
594///
595/// The Gottfried-Jackson frame differs only in the definition of :math:`\hat{z}`:
596///
597/// .. math:: \hat{z} \propto \vec{p}'_{\text{beam}}
598///
599/// Parameters
600/// ----------
601/// reaction : laddu.Reaction
602///     Reaction describing the production kinematics and decay roots.
603/// daughter : list of str
604///     Names of particles which are combined to form one of the decay products of the
605///     resonance associated with the decay parent.
606/// frame : {'Helicity', 'HX', 'HEL', 'GottfriedJackson', 'Gottfried Jackson', 'GJ', 'Gottfried-Jackson'}
607///     The frame to use in the  calculation
608///
609/// Raises
610/// ------
611/// ValueError
612///     If `frame` is not one of the valid options
613///
614/// See Also
615/// --------
616/// laddu.utils.vectors.Vec3.costheta
617///
618#[pyclass(name = "CosTheta", module = "laddu", from_py_object)]
619#[derive(Clone, Serialize, Deserialize)]
620pub struct PyCosTheta(pub CosTheta);
621
622#[pymethods]
623impl PyCosTheta {
624    /// The value of this Variable for the given Event
625    ///
626    /// Parameters
627    /// ----------
628    /// event : Event
629    ///     The Event upon which the Variable is calculated
630    ///
631    /// Returns
632    /// -------
633    /// value : float
634    ///     The value of the Variable for the given `event`
635    ///
636    fn value(&self, event: &PyEvent) -> PyResult<f64> {
637        let metadata = event
638            .metadata_opt()
639            .ok_or_else(|| PyValueError::new_err(
640                "This event is not associated with metadata; supply `p4_names`/`aux_names` when constructing it or evaluate via a Dataset.",
641            ))?;
642        let mut variable = self.0.clone();
643        variable.bind(metadata).map_err(PyErr::from)?;
644        Ok(variable.value(&event.event))
645    }
646    /// All values of this Variable on the given Dataset
647    ///
648    /// Parameters
649    /// ----------
650    /// dataset : Dataset
651    ///     The Dataset upon which the Variable is calculated
652    ///
653    /// Returns
654    /// -------
655    /// values : array_like
656    ///     The values of the Variable for each Event in the given `dataset`
657    ///
658    fn value_on<'py>(
659        &self,
660        py: Python<'py>,
661        dataset: &PyDataset,
662    ) -> PyResult<Bound<'py, PyArray1<f64>>> {
663        let values = self.0.value_on(&dataset.0).map_err(PyErr::from)?;
664        Ok(PyArray1::from_vec(py, values))
665    }
666    fn __eq__(&self, value: f64) -> PyVariableExpression {
667        PyVariableExpression(self.0.eq(value))
668    }
669    fn __lt__(&self, value: f64) -> PyVariableExpression {
670        PyVariableExpression(self.0.lt(value))
671    }
672    fn __gt__(&self, value: f64) -> PyVariableExpression {
673        PyVariableExpression(self.0.gt(value))
674    }
675    fn __le__(&self, value: f64) -> PyVariableExpression {
676        PyVariableExpression(self.0.le(value))
677    }
678    fn __ge__(&self, value: f64) -> PyVariableExpression {
679        PyVariableExpression(self.0.ge(value))
680    }
681    fn __repr__(&self) -> String {
682        format!("{:?}", self.0)
683    }
684    fn __str__(&self) -> String {
685        format!("{}", self.0)
686    }
687}
688
689/// The aziumuthal decay angle in the rest frame of the given `resonance`
690///
691/// This Variable is calculated by forming the given frame (helicity or Gottfried-Jackson) and
692/// calculating the spherical angles according to one of the decaying `daughter` particles.
693///
694/// The helicity frame is defined in terms of the following Cartesian axes in the rest frame of
695/// the `resonance`:
696///
697/// .. math:: \hat{z} \propto -\vec{p}'_{\text{recoil}}
698/// .. math:: \hat{y} \propto \vec{p}_{\text{beam}} \times (-\vec{p}_{\text{recoil}})
699/// .. math:: \hat{x} = \hat{y} \times \hat{z}
700///
701/// where primed vectors are in the rest frame of the `resonance` and unprimed vectors are in
702/// the center-of-momentum frame.
703///
704/// The Gottfried-Jackson frame differs only in the definition of :math:`\hat{z}`:
705///
706/// .. math:: \hat{z} \propto \vec{p}'_{\text{beam}}
707///
708/// Parameters
709/// ----------
710/// reaction : laddu.Reaction
711///     Reaction describing the production kinematics and decay roots.
712/// daughter : list of str
713///     Names of particles which are combined to form one of the decay products of the
714///     resonance associated with the decay parent.
715/// frame : {'Helicity', 'HX', 'HEL', 'GottfriedJackson', 'Gottfried Jackson', 'GJ', 'Gottfried-Jackson'}
716///     The frame to use in the  calculation
717///
718/// Raises
719/// ------
720/// ValueError
721///     If `frame` is not one of the valid options
722///
723///
724/// See Also
725/// --------
726/// laddu.utils.vectors.Vec3.phi
727///
728#[pyclass(name = "Phi", module = "laddu", from_py_object)]
729#[derive(Clone, Serialize, Deserialize)]
730pub struct PyPhi(pub Phi);
731
732#[pymethods]
733impl PyPhi {
734    /// The value of this Variable for the given Event
735    ///
736    /// Parameters
737    /// ----------
738    /// event : Event
739    ///     The Event upon which the Variable is calculated
740    ///
741    /// Returns
742    /// -------
743    /// value : float
744    ///     The value of the Variable for the given `event`
745    ///
746    fn value(&self, event: &PyEvent) -> PyResult<f64> {
747        let metadata = event
748            .metadata_opt()
749            .ok_or_else(|| PyValueError::new_err(
750                "This event is not associated with metadata; supply `p4_names`/`aux_names` when constructing it or evaluate via a Dataset.",
751            ))?;
752        let mut variable = self.0.clone();
753        variable.bind(metadata).map_err(PyErr::from)?;
754        Ok(variable.value(&event.event))
755    }
756    /// All values of this Variable on the given Dataset
757    ///
758    /// Parameters
759    /// ----------
760    /// dataset : Dataset
761    ///     The Dataset upon which the Variable is calculated
762    ///
763    /// Returns
764    /// -------
765    /// values : array_like
766    ///     The values of the Variable for each Event in the given `dataset`
767    ///
768    fn value_on<'py>(
769        &self,
770        py: Python<'py>,
771        dataset: &PyDataset,
772    ) -> PyResult<Bound<'py, PyArray1<f64>>> {
773        let values = self.0.value_on(&dataset.0).map_err(PyErr::from)?;
774        Ok(PyArray1::from_vec(py, values))
775    }
776    fn __eq__(&self, value: f64) -> PyVariableExpression {
777        PyVariableExpression(self.0.eq(value))
778    }
779    fn __lt__(&self, value: f64) -> PyVariableExpression {
780        PyVariableExpression(self.0.lt(value))
781    }
782    fn __gt__(&self, value: f64) -> PyVariableExpression {
783        PyVariableExpression(self.0.gt(value))
784    }
785    fn __le__(&self, value: f64) -> PyVariableExpression {
786        PyVariableExpression(self.0.le(value))
787    }
788    fn __ge__(&self, value: f64) -> PyVariableExpression {
789        PyVariableExpression(self.0.ge(value))
790    }
791    fn __repr__(&self) -> String {
792        format!("{:?}", self.0)
793    }
794    fn __str__(&self) -> String {
795        format!("{}", self.0)
796    }
797}
798
799/// A Variable used to define both spherical decay angles in the given frame
800///
801/// This class combines ``laddu.CosTheta`` and ``laddu.Phi`` into a single
802/// object
803///
804/// Parameters
805/// ----------
806/// reaction : laddu.Reaction
807///     Reaction describing the production kinematics and decay roots.
808/// daughter : list of str
809///     Names of particles which are combined to form one of the decay products of the
810///     resonance associated with the decay parent.
811/// frame : {'Helicity', 'HX', 'HEL', 'GottfriedJackson', 'Gottfried Jackson', 'GJ', 'Gottfried-Jackson'}
812///     The frame to use in the  calculation
813///
814/// Raises
815/// ------
816/// ValueError
817///     If `frame` is not one of the valid options
818///
819/// See Also
820/// --------
821/// laddu.CosTheta
822/// laddu.Phi
823///
824#[pyclass(name = "Angles", module = "laddu", skip_from_py_object)]
825#[derive(Clone)]
826pub struct PyAngles(pub Angles);
827#[pymethods]
828impl PyAngles {
829    /// The Variable representing the cosine of the polar spherical decay angle
830    ///
831    /// Returns
832    /// -------
833    /// CosTheta
834    ///
835    #[getter]
836    fn costheta(&self) -> PyCosTheta {
837        PyCosTheta(self.0.costheta.clone())
838    }
839    // The Variable representing the polar azimuthal decay angle
840    //
841    // Returns
842    // -------
843    // Phi
844    //
845    #[getter]
846    fn phi(&self) -> PyPhi {
847        PyPhi(self.0.phi.clone())
848    }
849    fn __repr__(&self) -> String {
850        format!("{:?}", self.0)
851    }
852    fn __str__(&self) -> String {
853        format!("{}", self.0)
854    }
855}
856
857/// The polar angle of the given polarization vector with respect to the production plane
858///
859/// The `beam` and `recoil` particles define the plane of production, and this Variable
860/// describes the polar angle of the `beam` relative to this plane
861///
862/// Parameters
863/// ----------
864/// reaction : laddu.Reaction
865///     Reaction describing the production kinematics and decay roots.
866/// pol_angle : str
867///     Name of the auxiliary scalar column storing the polarization angle in radians
868///
869#[pyclass(name = "PolAngle", module = "laddu", from_py_object)]
870#[derive(Clone, Serialize, Deserialize)]
871pub struct PyPolAngle(pub PolAngle);
872
873#[pymethods]
874impl PyPolAngle {
875    #[new]
876    fn new(reaction: PyReaction, pol_angle: String) -> Self {
877        Self(PolAngle::new(reaction.0.clone(), pol_angle))
878    }
879    /// The value of this Variable for the given Event
880    ///
881    /// Parameters
882    /// ----------
883    /// event : Event
884    ///     The Event upon which the Variable is calculated
885    ///
886    /// Returns
887    /// -------
888    /// value : float
889    ///     The value of the Variable for the given `event`
890    ///
891    fn value(&self, event: &PyEvent) -> PyResult<f64> {
892        let metadata = event
893            .metadata_opt()
894            .ok_or_else(|| PyValueError::new_err(
895                "This event is not associated with metadata; supply `p4_names`/`aux_names` when constructing it or evaluate via a Dataset.",
896            ))?;
897        let mut variable = self.0.clone();
898        variable.bind(metadata).map_err(PyErr::from)?;
899        Ok(variable.value(&event.event))
900    }
901    /// All values of this Variable on the given Dataset
902    ///
903    /// Parameters
904    /// ----------
905    /// dataset : Dataset
906    ///     The Dataset upon which the Variable is calculated
907    ///
908    /// Returns
909    /// -------
910    /// values : array_like
911    ///     The values of the Variable for each Event in the given `dataset`
912    ///
913    fn value_on<'py>(
914        &self,
915        py: Python<'py>,
916        dataset: &PyDataset,
917    ) -> PyResult<Bound<'py, PyArray1<f64>>> {
918        let values = self.0.value_on(&dataset.0).map_err(PyErr::from)?;
919        Ok(PyArray1::from_vec(py, values))
920    }
921    fn __eq__(&self, value: f64) -> PyVariableExpression {
922        PyVariableExpression(self.0.eq(value))
923    }
924    fn __lt__(&self, value: f64) -> PyVariableExpression {
925        PyVariableExpression(self.0.lt(value))
926    }
927    fn __gt__(&self, value: f64) -> PyVariableExpression {
928        PyVariableExpression(self.0.gt(value))
929    }
930    fn __le__(&self, value: f64) -> PyVariableExpression {
931        PyVariableExpression(self.0.le(value))
932    }
933    fn __ge__(&self, value: f64) -> PyVariableExpression {
934        PyVariableExpression(self.0.ge(value))
935    }
936    fn __repr__(&self) -> String {
937        format!("{:?}", self.0)
938    }
939    fn __str__(&self) -> String {
940        format!("{}", self.0)
941    }
942}
943
944/// The magnitude of the given particle's polarization vector
945///
946/// This Variable simply represents the magnitude of the polarization vector of the particle
947/// with the index `beam`
948///
949/// Parameters
950/// ----------
951/// pol_magnitude : str
952///     Name of the auxiliary scalar column storing the magnitude of the polarization vector
953///
954/// See Also
955/// --------
956/// laddu.utils.vectors.Vec3.mag
957///
958#[pyclass(name = "PolMagnitude", module = "laddu", from_py_object)]
959#[derive(Clone, Serialize, Deserialize)]
960pub struct PyPolMagnitude(pub PolMagnitude);
961
962#[pymethods]
963impl PyPolMagnitude {
964    #[new]
965    fn new(pol_magnitude: String) -> Self {
966        Self(PolMagnitude::new(pol_magnitude))
967    }
968    /// The value of this Variable for the given Event
969    ///
970    /// Parameters
971    /// ----------
972    /// event : Event
973    ///     The Event upon which the Variable is calculated
974    ///
975    /// Returns
976    /// -------
977    /// value : float
978    ///     The value of the Variable for the given `event`
979    ///
980    fn value(&self, event: &PyEvent) -> PyResult<f64> {
981        let metadata = event
982            .metadata_opt()
983            .ok_or_else(|| PyValueError::new_err(
984                "This event is not associated with metadata; supply `p4_names`/`aux_names` when constructing it or evaluate via a Dataset.",
985            ))?;
986        let mut variable = self.0.clone();
987        variable.bind(metadata).map_err(PyErr::from)?;
988        Ok(variable.value(&event.event))
989    }
990    /// All values of this Variable on the given Dataset
991    ///
992    /// Parameters
993    /// ----------
994    /// dataset : Dataset
995    ///     The Dataset upon which the Variable is calculated
996    ///
997    /// Returns
998    /// -------
999    /// values : array_like
1000    ///     The values of the Variable for each Event in the given `dataset`
1001    ///
1002    fn value_on<'py>(
1003        &self,
1004        py: Python<'py>,
1005        dataset: &PyDataset,
1006    ) -> PyResult<Bound<'py, PyArray1<f64>>> {
1007        let values = self.0.value_on(&dataset.0).map_err(PyErr::from)?;
1008        Ok(PyArray1::from_vec(py, values))
1009    }
1010    fn __eq__(&self, value: f64) -> PyVariableExpression {
1011        PyVariableExpression(self.0.eq(value))
1012    }
1013    fn __lt__(&self, value: f64) -> PyVariableExpression {
1014        PyVariableExpression(self.0.lt(value))
1015    }
1016    fn __gt__(&self, value: f64) -> PyVariableExpression {
1017        PyVariableExpression(self.0.gt(value))
1018    }
1019    fn __le__(&self, value: f64) -> PyVariableExpression {
1020        PyVariableExpression(self.0.le(value))
1021    }
1022    fn __ge__(&self, value: f64) -> PyVariableExpression {
1023        PyVariableExpression(self.0.ge(value))
1024    }
1025    fn __repr__(&self) -> String {
1026        format!("{:?}", self.0)
1027    }
1028    fn __str__(&self) -> String {
1029        format!("{}", self.0)
1030    }
1031}
1032
1033/// A Variable used to define both the polarization angle and magnitude of the given particle``
1034///
1035/// This class combines ``laddu.PolAngle`` and ``laddu.PolMagnitude`` into a single
1036/// object
1037///
1038/// Parameters
1039/// ----------
1040/// reaction : laddu.Reaction
1041///     Reaction describing the production kinematics and decay roots.
1042/// pol_magnitude : str
1043///     Name of the auxiliary scalar storing the polarization magnitude
1044/// pol_angle : str
1045///     Name of the auxiliary scalar storing the polarization angle in radians
1046///
1047/// See Also
1048/// --------
1049/// laddu.PolAngle
1050/// laddu.PolMagnitude
1051///
1052#[pyclass(name = "Polarization", module = "laddu", skip_from_py_object)]
1053#[derive(Clone)]
1054pub struct PyPolarization(pub Polarization);
1055#[pymethods]
1056impl PyPolarization {
1057    #[new]
1058    #[pyo3(signature=(reaction, *, pol_magnitude, pol_angle))]
1059    fn new(reaction: PyReaction, pol_magnitude: String, pol_angle: String) -> PyResult<Self> {
1060        if pol_magnitude == pol_angle {
1061            return Err(PyValueError::new_err(
1062                "`pol_magnitude` and `pol_angle` must reference distinct auxiliary columns",
1063            ));
1064        }
1065        let polarization = Polarization::new(reaction.0.clone(), pol_magnitude, pol_angle);
1066        Ok(PyPolarization(polarization))
1067    }
1068    /// The Variable representing the magnitude of the polarization vector
1069    ///
1070    /// Returns
1071    /// -------
1072    /// PolMagnitude
1073    ///
1074    #[getter]
1075    fn pol_magnitude(&self) -> PyPolMagnitude {
1076        PyPolMagnitude(self.0.pol_magnitude.clone())
1077    }
1078    /// The Variable representing the polar angle of the polarization vector
1079    ///
1080    /// Returns
1081    /// -------
1082    /// PolAngle
1083    ///
1084    #[getter]
1085    fn pol_angle(&self) -> PyPolAngle {
1086        PyPolAngle(self.0.pol_angle.clone())
1087    }
1088    fn __repr__(&self) -> String {
1089        format!("{:?}", self.0)
1090    }
1091    fn __str__(&self) -> String {
1092        format!("{}", self.0)
1093    }
1094}
1095
1096/// Mandelstam variables s, t, and u
1097///
1098/// By convention, the metric is chosen to be :math:`(+---)` and the variables are defined as follows
1099/// (ignoring factors of :math:`c`):
1100///
1101/// .. math:: s = (p_1 + p_2)^2 = (p_3 + p_4)^2
1102///
1103/// .. math:: t = (p_1 - p_3)^2 = (p_4 - p_2)^2
1104///
1105/// .. math:: u = (p_1 - p_4)^2 = (p_3 - p_2)^2
1106///
1107/// Parameters
1108/// ----------
1109/// reaction : laddu.Reaction
1110///     Reaction describing the two-to-two kinematics whose Mandelstam channels should be evaluated.
1111/// channel: {'s', 't', 'u', 'S', 'T', 'U'}
1112///     The Mandelstam channel to calculate
1113///
1114/// Raises
1115/// ------
1116/// Exception
1117///     If more than one particle list is empty
1118/// ValueError
1119///     If `channel` is not one of the valid options
1120///
1121/// Notes
1122/// -----
1123/// ///
1124#[pyclass(name = "Mandelstam", module = "laddu", from_py_object)]
1125#[derive(Clone, Serialize, Deserialize)]
1126pub struct PyMandelstam(pub Mandelstam);
1127
1128#[pymethods]
1129impl PyMandelstam {
1130    #[new]
1131    fn new(reaction: PyReaction, channel: &str) -> PyResult<Self> {
1132        Ok(Self(reaction.0.mandelstam(channel.parse()?)?))
1133    }
1134    /// The value of this Variable for the given Event
1135    ///
1136    /// Parameters
1137    /// ----------
1138    /// event : Event
1139    ///     The Event upon which the Variable is calculated
1140    ///
1141    /// Returns
1142    /// -------
1143    /// value : float
1144    ///     The value of the Variable for the given `event`
1145    ///
1146    fn value(&self, event: &PyEvent) -> PyResult<f64> {
1147        let metadata = event
1148            .metadata_opt()
1149            .ok_or_else(|| PyValueError::new_err(
1150                "This event is not associated with metadata; supply `p4_names`/`aux_names` when constructing it or evaluate via a Dataset.",
1151            ))?;
1152        let mut variable = self.0.clone();
1153        variable.bind(metadata).map_err(PyErr::from)?;
1154        Ok(variable.value(&event.event))
1155    }
1156    /// All values of this Variable on the given Dataset
1157    ///
1158    /// Parameters
1159    /// ----------
1160    /// dataset : Dataset
1161    ///     The Dataset upon which the Variable is calculated
1162    ///
1163    /// Returns
1164    /// -------
1165    /// values : array_like
1166    ///     The values of the Variable for each Event in the given `dataset`
1167    ///
1168    fn value_on<'py>(
1169        &self,
1170        py: Python<'py>,
1171        dataset: &PyDataset,
1172    ) -> PyResult<Bound<'py, PyArray1<f64>>> {
1173        let values = self.0.value_on(&dataset.0).map_err(PyErr::from)?;
1174        Ok(PyArray1::from_vec(py, values))
1175    }
1176    fn __eq__(&self, value: f64) -> PyVariableExpression {
1177        PyVariableExpression(self.0.eq(value))
1178    }
1179    fn __lt__(&self, value: f64) -> PyVariableExpression {
1180        PyVariableExpression(self.0.lt(value))
1181    }
1182    fn __gt__(&self, value: f64) -> PyVariableExpression {
1183        PyVariableExpression(self.0.gt(value))
1184    }
1185    fn __le__(&self, value: f64) -> PyVariableExpression {
1186        PyVariableExpression(self.0.le(value))
1187    }
1188    fn __ge__(&self, value: f64) -> PyVariableExpression {
1189        PyVariableExpression(self.0.ge(value))
1190    }
1191    fn __repr__(&self) -> String {
1192        format!("{:?}", self.0)
1193    }
1194    fn __str__(&self) -> String {
1195        format!("{}", self.0)
1196    }
1197}
1198
1199#[typetag::serde]
1200impl Variable for PyVariable {
1201    fn bind(&mut self, metadata: &DatasetMetadata) -> LadduResult<()> {
1202        match self {
1203            PyVariable::Mass(mass) => mass.0.bind(metadata),
1204            PyVariable::CosTheta(cos_theta) => cos_theta.0.bind(metadata),
1205            PyVariable::Phi(phi) => phi.0.bind(metadata),
1206            PyVariable::PolAngle(pol_angle) => pol_angle.0.bind(metadata),
1207            PyVariable::PolMagnitude(pol_magnitude) => pol_magnitude.0.bind(metadata),
1208            PyVariable::Mandelstam(mandelstam) => mandelstam.0.bind(metadata),
1209        }
1210    }
1211
1212    fn value_on(&self, dataset: &Dataset) -> LadduResult<Vec<f64>> {
1213        match self {
1214            PyVariable::Mass(mass) => mass.0.value_on(dataset),
1215            PyVariable::CosTheta(cos_theta) => cos_theta.0.value_on(dataset),
1216            PyVariable::Phi(phi) => phi.0.value_on(dataset),
1217            PyVariable::PolAngle(pol_angle) => pol_angle.0.value_on(dataset),
1218            PyVariable::PolMagnitude(pol_magnitude) => pol_magnitude.0.value_on(dataset),
1219            PyVariable::Mandelstam(mandelstam) => mandelstam.0.value_on(dataset),
1220        }
1221    }
1222
1223    fn value(&self, event: &dyn EventLike) -> f64 {
1224        match self {
1225            PyVariable::Mass(mass) => mass.0.value(event),
1226            PyVariable::CosTheta(cos_theta) => cos_theta.0.value(event),
1227            PyVariable::Phi(phi) => phi.0.value(event),
1228            PyVariable::PolAngle(pol_angle) => pol_angle.0.value(event),
1229            PyVariable::PolMagnitude(pol_magnitude) => pol_magnitude.0.value(event),
1230            PyVariable::Mandelstam(mandelstam) => mandelstam.0.value(event),
1231        }
1232    }
1233}
laddu_python/variables.rs

laddu_python/
variables.rs
