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 for the selected frame.
420    #[pyo3(signature=(frame="Helicity"))]
421    fn costheta(&self, frame: &str) -> PyResult<PyCosTheta> {
422        Ok(PyCosTheta(self.0.costheta(frame.parse()?)?))
423    }
424
425    /// Production phi variable for the selected frame.
426    #[pyo3(signature=(frame="Helicity"))]
427    fn phi(&self, frame: &str) -> PyResult<PyPhi> {
428        Ok(PyPhi(self.0.phi(frame.parse()?)?))
429    }
430
431    /// Production angle variables for the selected frame.
432    #[pyo3(signature=(frame="Helicity"))]
433    fn angles(&self, frame: &str) -> PyResult<PyAngles> {
434        Ok(PyAngles(self.0.angles(frame.parse()?)?))
435    }
436
437    /// Construct the helicity-basis production angular factor for one explicit helicity term.
438    #[pyo3(signature=(*tags, spin, projection, lambda_produced, lambda_recoil, frame="Helicity"))]
439    #[allow(clippy::too_many_arguments)]
440    fn helicity_factor(
441        &self,
442        tags: &Bound<'_, PyTuple>,
443        spin: &Bound<'_, PyAny>,
444        projection: &Bound<'_, PyAny>,
445        lambda_produced: &Bound<'_, PyAny>,
446        lambda_recoil: &Bound<'_, PyAny>,
447        frame: &str,
448    ) -> PyResult<PyExpression> {
449        Ok(PyExpression(self.0.helicity_factor(
450            py_tags(tags)?,
451            parse_angular_momentum(spin)?,
452            parse_projection(projection)?,
453            parse_projection(lambda_produced)?,
454            parse_projection(lambda_recoil)?,
455            frame.parse()?,
456        )?))
457    }
458
459    /// Construct the canonical-basis production spin-angular factor for one LS/helicity term.
460    #[pyo3(signature=(*tags, spin, projection, orbital_l, coupled_spin, produced_spin, recoil_spin, lambda_produced, lambda_recoil, frame="Helicity"))]
461    #[allow(clippy::too_many_arguments)]
462    fn canonical_factor(
463        &self,
464        tags: &Bound<'_, PyTuple>,
465        spin: &Bound<'_, PyAny>,
466        projection: &Bound<'_, PyAny>,
467        orbital_l: &Bound<'_, PyAny>,
468        coupled_spin: &Bound<'_, PyAny>,
469        produced_spin: &Bound<'_, PyAny>,
470        recoil_spin: &Bound<'_, PyAny>,
471        lambda_produced: &Bound<'_, PyAny>,
472        lambda_recoil: &Bound<'_, PyAny>,
473        frame: &str,
474    ) -> PyResult<PyExpression> {
475        Ok(PyExpression(self.0.canonical_factor(
476            py_tags(tags)?,
477            parse_angular_momentum(spin)?,
478            parse_projection(projection)?,
479            parse_orbital_angular_momentum(orbital_l)?,
480            parse_angular_momentum(coupled_spin)?,
481            parse_angular_momentum(produced_spin)?,
482            parse_angular_momentum(recoil_spin)?,
483            parse_projection(lambda_produced)?,
484            parse_projection(lambda_recoil)?,
485            frame.parse()?,
486        )?))
487    }
488
489    fn __repr__(&self) -> String {
490        format!("{:?}", self.0)
491    }
492
493    fn __str__(&self) -> String {
494        format!("{:?}", self.0)
495    }
496}
497
498/// The invariant mass of an arbitrary combination of constituent particles in an Event
499///
500/// This variable is calculated by summing up the 4-momenta of each particle listed by index in
501/// `constituents` and taking the invariant magnitude of the resulting 4-vector.
502///
503/// Parameters
504/// ----------
505/// constituents : str or list of str
506///     Particle names to combine when constructing the final four-momentum
507///
508/// See Also
509/// --------
510/// laddu.utils.vectors.Vec4.m
511///
512#[pyclass(name = "Mass", module = "laddu", from_py_object)]
513#[derive(Clone, Serialize, Deserialize)]
514pub struct PyMass(pub Mass);
515
516#[pymethods]
517impl PyMass {
518    #[new]
519    fn new(constituents: PyP4SelectionInput) -> Self {
520        Self(Mass::new(constituents.into_selection()))
521    }
522    /// The value of this Variable for the given Event
523    ///
524    /// Parameters
525    /// ----------
526    /// event : Event
527    ///     The Event upon which the Variable is calculated
528    ///
529    /// Returns
530    /// -------
531    /// value : float
532    ///     The value of the Variable for the given `event`
533    ///
534    fn value(&self, event: &PyEvent) -> PyResult<f64> {
535        let metadata = event
536            .metadata_opt()
537            .ok_or_else(|| PyValueError::new_err(
538                "This event is not associated with metadata; supply `p4_names`/`aux_names` when constructing it or evaluate via a Dataset.",
539            ))?;
540        let mut variable = self.0.clone();
541        variable.bind(metadata).map_err(PyErr::from)?;
542        Ok(variable.value(&event.event))
543    }
544    /// All values of this Variable on the given Dataset
545    ///
546    /// Parameters
547    /// ----------
548    /// dataset : Dataset
549    ///     The Dataset upon which the Variable is calculated
550    ///
551    /// Returns
552    /// -------
553    /// values : array_like
554    ///     The values of the Variable for each Event in the given `dataset`
555    ///
556    fn value_on<'py>(
557        &self,
558        py: Python<'py>,
559        dataset: &PyDataset,
560    ) -> PyResult<Bound<'py, PyArray1<f64>>> {
561        let values = self.0.value_on(&dataset.0).map_err(PyErr::from)?;
562        Ok(PyArray1::from_vec(py, values))
563    }
564    fn __eq__(&self, value: f64) -> PyVariableExpression {
565        PyVariableExpression(self.0.eq(value))
566    }
567    fn __lt__(&self, value: f64) -> PyVariableExpression {
568        PyVariableExpression(self.0.lt(value))
569    }
570    fn __gt__(&self, value: f64) -> PyVariableExpression {
571        PyVariableExpression(self.0.gt(value))
572    }
573    fn __le__(&self, value: f64) -> PyVariableExpression {
574        PyVariableExpression(self.0.le(value))
575    }
576    fn __ge__(&self, value: f64) -> PyVariableExpression {
577        PyVariableExpression(self.0.ge(value))
578    }
579    fn __repr__(&self) -> String {
580        format!("{:?}", self.0)
581    }
582    fn __str__(&self) -> String {
583        format!("{}", self.0)
584    }
585}
586
587/// The cosine of the polar decay angle in the rest frame of the given `resonance`
588///
589/// This Variable is calculated by forming the given frame (helicity or Gottfried-Jackson) and
590/// calculating the spherical angles according to one of the decaying `daughter` particles.
591///
592/// The helicity frame is defined in terms of the following Cartesian axes in the rest frame of
593/// the `resonance`:
594///
595/// .. math:: \hat{z} \propto -\vec{p}'_{\text{recoil}}
596/// .. math:: \hat{y} \propto \vec{p}_{\text{beam}} \times (-\vec{p}_{\text{recoil}})
597/// .. math:: \hat{x} = \hat{y} \times \hat{z}
598///
599/// where primed vectors are in the rest frame of the `resonance` and unprimed vectors are in
600/// the center-of-momentum frame.
601///
602/// The Gottfried-Jackson frame differs only in the definition of :math:`\hat{z}`:
603///
604/// .. math:: \hat{z} \propto \vec{p}'_{\text{beam}}
605///
606/// Parameters
607/// ----------
608/// reaction : laddu.Reaction
609///     Reaction describing the production kinematics and decay roots.
610/// daughter : list of str
611///     Names of particles which are combined to form one of the decay products of the
612///     resonance associated with the decay parent.
613/// frame : {'Helicity', 'HX', 'HEL', 'GottfriedJackson', 'Gottfried Jackson', 'GJ', 'Gottfried-Jackson'}
614///     The frame to use in the  calculation
615///
616/// Raises
617/// ------
618/// ValueError
619///     If `frame` is not one of the valid options
620///
621/// See Also
622/// --------
623/// laddu.utils.vectors.Vec3.costheta
624///
625#[pyclass(name = "CosTheta", module = "laddu", from_py_object)]
626#[derive(Clone, Serialize, Deserialize)]
627pub struct PyCosTheta(pub CosTheta);
628
629#[pymethods]
630impl PyCosTheta {
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/// The aziumuthal decay angle in the rest frame of the given `resonance`
697///
698/// This Variable is calculated by forming the given frame (helicity or Gottfried-Jackson) and
699/// calculating the spherical angles according to one of the decaying `daughter` particles.
700///
701/// The helicity frame is defined in terms of the following Cartesian axes in the rest frame of
702/// the `resonance`:
703///
704/// .. math:: \hat{z} \propto -\vec{p}'_{\text{recoil}}
705/// .. math:: \hat{y} \propto \vec{p}_{\text{beam}} \times (-\vec{p}_{\text{recoil}})
706/// .. math:: \hat{x} = \hat{y} \times \hat{z}
707///
708/// where primed vectors are in the rest frame of the `resonance` and unprimed vectors are in
709/// the center-of-momentum frame.
710///
711/// The Gottfried-Jackson frame differs only in the definition of :math:`\hat{z}`:
712///
713/// .. math:: \hat{z} \propto \vec{p}'_{\text{beam}}
714///
715/// Parameters
716/// ----------
717/// reaction : laddu.Reaction
718///     Reaction describing the production kinematics and decay roots.
719/// daughter : list of str
720///     Names of particles which are combined to form one of the decay products of the
721///     resonance associated with the decay parent.
722/// frame : {'Helicity', 'HX', 'HEL', 'GottfriedJackson', 'Gottfried Jackson', 'GJ', 'Gottfried-Jackson'}
723///     The frame to use in the  calculation
724///
725/// Raises
726/// ------
727/// ValueError
728///     If `frame` is not one of the valid options
729///
730///
731/// See Also
732/// --------
733/// laddu.utils.vectors.Vec3.phi
734///
735#[pyclass(name = "Phi", module = "laddu", from_py_object)]
736#[derive(Clone, Serialize, Deserialize)]
737pub struct PyPhi(pub Phi);
738
739#[pymethods]
740impl PyPhi {
741    /// The value of this Variable for the given Event
742    ///
743    /// Parameters
744    /// ----------
745    /// event : Event
746    ///     The Event upon which the Variable is calculated
747    ///
748    /// Returns
749    /// -------
750    /// value : float
751    ///     The value of the Variable for the given `event`
752    ///
753    fn value(&self, event: &PyEvent) -> PyResult<f64> {
754        let metadata = event
755            .metadata_opt()
756            .ok_or_else(|| PyValueError::new_err(
757                "This event is not associated with metadata; supply `p4_names`/`aux_names` when constructing it or evaluate via a Dataset.",
758            ))?;
759        let mut variable = self.0.clone();
760        variable.bind(metadata).map_err(PyErr::from)?;
761        Ok(variable.value(&event.event))
762    }
763    /// All values of this Variable on the given Dataset
764    ///
765    /// Parameters
766    /// ----------
767    /// dataset : Dataset
768    ///     The Dataset upon which the Variable is calculated
769    ///
770    /// Returns
771    /// -------
772    /// values : array_like
773    ///     The values of the Variable for each Event in the given `dataset`
774    ///
775    fn value_on<'py>(
776        &self,
777        py: Python<'py>,
778        dataset: &PyDataset,
779    ) -> PyResult<Bound<'py, PyArray1<f64>>> {
780        let values = self.0.value_on(&dataset.0).map_err(PyErr::from)?;
781        Ok(PyArray1::from_vec(py, values))
782    }
783    fn __eq__(&self, value: f64) -> PyVariableExpression {
784        PyVariableExpression(self.0.eq(value))
785    }
786    fn __lt__(&self, value: f64) -> PyVariableExpression {
787        PyVariableExpression(self.0.lt(value))
788    }
789    fn __gt__(&self, value: f64) -> PyVariableExpression {
790        PyVariableExpression(self.0.gt(value))
791    }
792    fn __le__(&self, value: f64) -> PyVariableExpression {
793        PyVariableExpression(self.0.le(value))
794    }
795    fn __ge__(&self, value: f64) -> PyVariableExpression {
796        PyVariableExpression(self.0.ge(value))
797    }
798    fn __repr__(&self) -> String {
799        format!("{:?}", self.0)
800    }
801    fn __str__(&self) -> String {
802        format!("{}", self.0)
803    }
804}
805
806/// A Variable used to define both spherical decay angles in the given frame
807///
808/// This class combines ``laddu.CosTheta`` and ``laddu.Phi`` into a single
809/// object
810///
811/// Parameters
812/// ----------
813/// reaction : laddu.Reaction
814///     Reaction describing the production kinematics and decay roots.
815/// daughter : list of str
816///     Names of particles which are combined to form one of the decay products of the
817///     resonance associated with the decay parent.
818/// frame : {'Helicity', 'HX', 'HEL', 'GottfriedJackson', 'Gottfried Jackson', 'GJ', 'Gottfried-Jackson'}
819///     The frame to use in the  calculation
820///
821/// Raises
822/// ------
823/// ValueError
824///     If `frame` is not one of the valid options
825///
826/// See Also
827/// --------
828/// laddu.CosTheta
829/// laddu.Phi
830///
831#[pyclass(name = "Angles", module = "laddu", skip_from_py_object)]
832#[derive(Clone)]
833pub struct PyAngles(pub Angles);
834#[pymethods]
835impl PyAngles {
836    /// The Variable representing the cosine of the polar spherical decay angle
837    ///
838    /// Returns
839    /// -------
840    /// CosTheta
841    ///
842    #[getter]
843    fn costheta(&self) -> PyCosTheta {
844        PyCosTheta(self.0.costheta.clone())
845    }
846    // The Variable representing the polar azimuthal decay angle
847    //
848    // Returns
849    // -------
850    // Phi
851    //
852    #[getter]
853    fn phi(&self) -> PyPhi {
854        PyPhi(self.0.phi.clone())
855    }
856    fn __repr__(&self) -> String {
857        format!("{:?}", self.0)
858    }
859    fn __str__(&self) -> String {
860        format!("{}", self.0)
861    }
862}
863
864/// The polar angle of the given polarization vector with respect to the production plane
865///
866/// The `beam` and `recoil` particles define the plane of production, and this Variable
867/// describes the polar angle of the `beam` relative to this plane
868///
869/// Parameters
870/// ----------
871/// reaction : laddu.Reaction
872///     Reaction describing the production kinematics and decay roots.
873/// pol_angle : str
874///     Name of the auxiliary scalar column storing the polarization angle in radians
875///
876#[pyclass(name = "PolAngle", module = "laddu", from_py_object)]
877#[derive(Clone, Serialize, Deserialize)]
878pub struct PyPolAngle(pub PolAngle);
879
880#[pymethods]
881impl PyPolAngle {
882    #[new]
883    fn new(reaction: PyReaction, pol_angle: String) -> Self {
884        Self(PolAngle::new(reaction.0.clone(), pol_angle))
885    }
886    /// The value of this Variable for the given Event
887    ///
888    /// Parameters
889    /// ----------
890    /// event : Event
891    ///     The Event upon which the Variable is calculated
892    ///
893    /// Returns
894    /// -------
895    /// value : float
896    ///     The value of the Variable for the given `event`
897    ///
898    fn value(&self, event: &PyEvent) -> PyResult<f64> {
899        let metadata = event
900            .metadata_opt()
901            .ok_or_else(|| PyValueError::new_err(
902                "This event is not associated with metadata; supply `p4_names`/`aux_names` when constructing it or evaluate via a Dataset.",
903            ))?;
904        let mut variable = self.0.clone();
905        variable.bind(metadata).map_err(PyErr::from)?;
906        Ok(variable.value(&event.event))
907    }
908    /// All values of this Variable on the given Dataset
909    ///
910    /// Parameters
911    /// ----------
912    /// dataset : Dataset
913    ///     The Dataset upon which the Variable is calculated
914    ///
915    /// Returns
916    /// -------
917    /// values : array_like
918    ///     The values of the Variable for each Event in the given `dataset`
919    ///
920    fn value_on<'py>(
921        &self,
922        py: Python<'py>,
923        dataset: &PyDataset,
924    ) -> PyResult<Bound<'py, PyArray1<f64>>> {
925        let values = self.0.value_on(&dataset.0).map_err(PyErr::from)?;
926        Ok(PyArray1::from_vec(py, values))
927    }
928    fn __eq__(&self, value: f64) -> PyVariableExpression {
929        PyVariableExpression(self.0.eq(value))
930    }
931    fn __lt__(&self, value: f64) -> PyVariableExpression {
932        PyVariableExpression(self.0.lt(value))
933    }
934    fn __gt__(&self, value: f64) -> PyVariableExpression {
935        PyVariableExpression(self.0.gt(value))
936    }
937    fn __le__(&self, value: f64) -> PyVariableExpression {
938        PyVariableExpression(self.0.le(value))
939    }
940    fn __ge__(&self, value: f64) -> PyVariableExpression {
941        PyVariableExpression(self.0.ge(value))
942    }
943    fn __repr__(&self) -> String {
944        format!("{:?}", self.0)
945    }
946    fn __str__(&self) -> String {
947        format!("{}", self.0)
948    }
949}
950
951/// The magnitude of the given particle's polarization vector
952///
953/// This Variable simply represents the magnitude of the polarization vector of the particle
954/// with the index `beam`
955///
956/// Parameters
957/// ----------
958/// pol_magnitude : str
959///     Name of the auxiliary scalar column storing the magnitude of the polarization vector
960///
961/// See Also
962/// --------
963/// laddu.utils.vectors.Vec3.mag
964///
965#[pyclass(name = "PolMagnitude", module = "laddu", from_py_object)]
966#[derive(Clone, Serialize, Deserialize)]
967pub struct PyPolMagnitude(pub PolMagnitude);
968
969#[pymethods]
970impl PyPolMagnitude {
971    #[new]
972    fn new(pol_magnitude: String) -> Self {
973        Self(PolMagnitude::new(pol_magnitude))
974    }
975    /// The value of this Variable for the given Event
976    ///
977    /// Parameters
978    /// ----------
979    /// event : Event
980    ///     The Event upon which the Variable is calculated
981    ///
982    /// Returns
983    /// -------
984    /// value : float
985    ///     The value of the Variable for the given `event`
986    ///
987    fn value(&self, event: &PyEvent) -> PyResult<f64> {
988        let metadata = event
989            .metadata_opt()
990            .ok_or_else(|| PyValueError::new_err(
991                "This event is not associated with metadata; supply `p4_names`/`aux_names` when constructing it or evaluate via a Dataset.",
992            ))?;
993        let mut variable = self.0.clone();
994        variable.bind(metadata).map_err(PyErr::from)?;
995        Ok(variable.value(&event.event))
996    }
997    /// All values of this Variable on the given Dataset
998    ///
999    /// Parameters
1000    /// ----------
1001    /// dataset : Dataset
1002    ///     The Dataset upon which the Variable is calculated
1003    ///
1004    /// Returns
1005    /// -------
1006    /// values : array_like
1007    ///     The values of the Variable for each Event in the given `dataset`
1008    ///
1009    fn value_on<'py>(
1010        &self,
1011        py: Python<'py>,
1012        dataset: &PyDataset,
1013    ) -> PyResult<Bound<'py, PyArray1<f64>>> {
1014        let values = self.0.value_on(&dataset.0).map_err(PyErr::from)?;
1015        Ok(PyArray1::from_vec(py, values))
1016    }
1017    fn __eq__(&self, value: f64) -> PyVariableExpression {
1018        PyVariableExpression(self.0.eq(value))
1019    }
1020    fn __lt__(&self, value: f64) -> PyVariableExpression {
1021        PyVariableExpression(self.0.lt(value))
1022    }
1023    fn __gt__(&self, value: f64) -> PyVariableExpression {
1024        PyVariableExpression(self.0.gt(value))
1025    }
1026    fn __le__(&self, value: f64) -> PyVariableExpression {
1027        PyVariableExpression(self.0.le(value))
1028    }
1029    fn __ge__(&self, value: f64) -> PyVariableExpression {
1030        PyVariableExpression(self.0.ge(value))
1031    }
1032    fn __repr__(&self) -> String {
1033        format!("{:?}", self.0)
1034    }
1035    fn __str__(&self) -> String {
1036        format!("{}", self.0)
1037    }
1038}
1039
1040/// A Variable used to define both the polarization angle and magnitude of the given particle``
1041///
1042/// This class combines ``laddu.PolAngle`` and ``laddu.PolMagnitude`` into a single
1043/// object
1044///
1045/// Parameters
1046/// ----------
1047/// reaction : laddu.Reaction
1048///     Reaction describing the production kinematics and decay roots.
1049/// pol_magnitude : str
1050///     Name of the auxiliary scalar storing the polarization magnitude
1051/// pol_angle : str
1052///     Name of the auxiliary scalar storing the polarization angle in radians
1053///
1054/// See Also
1055/// --------
1056/// laddu.PolAngle
1057/// laddu.PolMagnitude
1058///
1059#[pyclass(name = "Polarization", module = "laddu", skip_from_py_object)]
1060#[derive(Clone)]
1061pub struct PyPolarization(pub Polarization);
1062#[pymethods]
1063impl PyPolarization {
1064    #[new]
1065    #[pyo3(signature=(reaction, *, pol_magnitude, pol_angle))]
1066    fn new(reaction: PyReaction, pol_magnitude: String, pol_angle: String) -> PyResult<Self> {
1067        if pol_magnitude == pol_angle {
1068            return Err(PyValueError::new_err(
1069                "`pol_magnitude` and `pol_angle` must reference distinct auxiliary columns",
1070            ));
1071        }
1072        let polarization = Polarization::new(reaction.0.clone(), pol_magnitude, pol_angle);
1073        Ok(PyPolarization(polarization))
1074    }
1075    /// The Variable representing the magnitude of the polarization vector
1076    ///
1077    /// Returns
1078    /// -------
1079    /// PolMagnitude
1080    ///
1081    #[getter]
1082    fn pol_magnitude(&self) -> PyPolMagnitude {
1083        PyPolMagnitude(self.0.pol_magnitude.clone())
1084    }
1085    /// The Variable representing the polar angle of the polarization vector
1086    ///
1087    /// Returns
1088    /// -------
1089    /// PolAngle
1090    ///
1091    #[getter]
1092    fn pol_angle(&self) -> PyPolAngle {
1093        PyPolAngle(self.0.pol_angle.clone())
1094    }
1095    fn __repr__(&self) -> String {
1096        format!("{:?}", self.0)
1097    }
1098    fn __str__(&self) -> String {
1099        format!("{}", self.0)
1100    }
1101}
1102
1103/// Mandelstam variables s, t, and u
1104///
1105/// By convention, the metric is chosen to be :math:`(+---)` and the variables are defined as follows
1106/// (ignoring factors of :math:`c`):
1107///
1108/// .. math:: s = (p_1 + p_2)^2 = (p_3 + p_4)^2
1109///
1110/// .. math:: t = (p_1 - p_3)^2 = (p_4 - p_2)^2
1111///
1112/// .. math:: u = (p_1 - p_4)^2 = (p_3 - p_2)^2
1113///
1114/// Parameters
1115/// ----------
1116/// reaction : laddu.Reaction
1117///     Reaction describing the two-to-two kinematics whose Mandelstam channels should be evaluated.
1118/// channel: {'s', 't', 'u', 'S', 'T', 'U'}
1119///     The Mandelstam channel to calculate
1120///
1121/// Raises
1122/// ------
1123/// Exception
1124///     If more than one particle list is empty
1125/// ValueError
1126///     If `channel` is not one of the valid options
1127///
1128/// Notes
1129/// -----
1130/// ///
1131#[pyclass(name = "Mandelstam", module = "laddu", from_py_object)]
1132#[derive(Clone, Serialize, Deserialize)]
1133pub struct PyMandelstam(pub Mandelstam);
1134
1135#[pymethods]
1136impl PyMandelstam {
1137    #[new]
1138    fn new(reaction: PyReaction, channel: &str) -> PyResult<Self> {
1139        Ok(Self(reaction.0.mandelstam(channel.parse()?)?))
1140    }
1141    /// The value of this Variable for the given Event
1142    ///
1143    /// Parameters
1144    /// ----------
1145    /// event : Event
1146    ///     The Event upon which the Variable is calculated
1147    ///
1148    /// Returns
1149    /// -------
1150    /// value : float
1151    ///     The value of the Variable for the given `event`
1152    ///
1153    fn value(&self, event: &PyEvent) -> PyResult<f64> {
1154        let metadata = event
1155            .metadata_opt()
1156            .ok_or_else(|| PyValueError::new_err(
1157                "This event is not associated with metadata; supply `p4_names`/`aux_names` when constructing it or evaluate via a Dataset.",
1158            ))?;
1159        let mut variable = self.0.clone();
1160        variable.bind(metadata).map_err(PyErr::from)?;
1161        Ok(variable.value(&event.event))
1162    }
1163    /// All values of this Variable on the given Dataset
1164    ///
1165    /// Parameters
1166    /// ----------
1167    /// dataset : Dataset
1168    ///     The Dataset upon which the Variable is calculated
1169    ///
1170    /// Returns
1171    /// -------
1172    /// values : array_like
1173    ///     The values of the Variable for each Event in the given `dataset`
1174    ///
1175    fn value_on<'py>(
1176        &self,
1177        py: Python<'py>,
1178        dataset: &PyDataset,
1179    ) -> PyResult<Bound<'py, PyArray1<f64>>> {
1180        let values = self.0.value_on(&dataset.0).map_err(PyErr::from)?;
1181        Ok(PyArray1::from_vec(py, values))
1182    }
1183    fn __eq__(&self, value: f64) -> PyVariableExpression {
1184        PyVariableExpression(self.0.eq(value))
1185    }
1186    fn __lt__(&self, value: f64) -> PyVariableExpression {
1187        PyVariableExpression(self.0.lt(value))
1188    }
1189    fn __gt__(&self, value: f64) -> PyVariableExpression {
1190        PyVariableExpression(self.0.gt(value))
1191    }
1192    fn __le__(&self, value: f64) -> PyVariableExpression {
1193        PyVariableExpression(self.0.le(value))
1194    }
1195    fn __ge__(&self, value: f64) -> PyVariableExpression {
1196        PyVariableExpression(self.0.ge(value))
1197    }
1198    fn __repr__(&self) -> String {
1199        format!("{:?}", self.0)
1200    }
1201    fn __str__(&self) -> String {
1202        format!("{}", self.0)
1203    }
1204}
1205
1206#[typetag::serde]
1207impl Variable for PyVariable {
1208    fn bind(&mut self, metadata: &DatasetMetadata) -> LadduResult<()> {
1209        match self {
1210            PyVariable::Mass(mass) => mass.0.bind(metadata),
1211            PyVariable::CosTheta(cos_theta) => cos_theta.0.bind(metadata),
1212            PyVariable::Phi(phi) => phi.0.bind(metadata),
1213            PyVariable::PolAngle(pol_angle) => pol_angle.0.bind(metadata),
1214            PyVariable::PolMagnitude(pol_magnitude) => pol_magnitude.0.bind(metadata),
1215            PyVariable::Mandelstam(mandelstam) => mandelstam.0.bind(metadata),
1216        }
1217    }
1218
1219    fn value_on(&self, dataset: &Dataset) -> LadduResult<Vec<f64>> {
1220        match self {
1221            PyVariable::Mass(mass) => mass.0.value_on(dataset),
1222            PyVariable::CosTheta(cos_theta) => cos_theta.0.value_on(dataset),
1223            PyVariable::Phi(phi) => phi.0.value_on(dataset),
1224            PyVariable::PolAngle(pol_angle) => pol_angle.0.value_on(dataset),
1225            PyVariable::PolMagnitude(pol_magnitude) => pol_magnitude.0.value_on(dataset),
1226            PyVariable::Mandelstam(mandelstam) => mandelstam.0.value_on(dataset),
1227        }
1228    }
1229
1230    fn value(&self, event: &dyn EventLike) -> f64 {
1231        match self {
1232            PyVariable::Mass(mass) => mass.0.value(event),
1233            PyVariable::CosTheta(cos_theta) => cos_theta.0.value(event),
1234            PyVariable::Phi(phi) => phi.0.value(event),
1235            PyVariable::PolAngle(pol_angle) => pol_angle.0.value(event),
1236            PyVariable::PolMagnitude(pol_magnitude) => pol_magnitude.0.value(event),
1237            PyVariable::Mandelstam(mandelstam) => mandelstam.0.value(event),
1238        }
1239    }
1240}
laddu_python/variables.rs

laddu_python/
variables.rs
