laddu_python/utils/
variables.rs

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

laddu_python/utils/
variables.rs
