laddu_python/utils/
angular_momentum.rs

1use laddu_core::{
2    allowed_projections, helicity_combinations, AngularMomentum, AngularMomentumProjection,
3    LadduError, LadduResult,
4};
5use num::rational::Ratio;
6use pyo3::{
7    exceptions::PyValueError,
8    prelude::*,
9    types::{PyAny, PyBool, PyModule},
10    IntoPyObjectExt,
11};
12
13type PyQuantumNumber = Py<PyAny>;
14type PyHelicityCombination = (PyQuantumNumber, PyQuantumNumber, PyQuantumNumber);
15
16fn parse_angular_momentum(input: &Bound<'_, PyAny>) -> PyResult<AngularMomentum> {
17    parse_ratio_like(input)
18        .and_then(AngularMomentum::from_ratio)
19        .map_err(py_value_error)
20}
21
22fn parse_ratio_like(input: &Bound<'_, PyAny>) -> LadduResult<Ratio<i32>> {
23    if input.is_instance_of::<PyBool>() {
24        return Err(LadduError::Custom(
25            "quantum number cannot be a bool".to_string(),
26        ));
27    }
28    if let Ok(value) = input.extract::<i32>() {
29        return Ok(Ratio::from_integer(value));
30    }
31    if let Ok(value) = input.extract::<f64>() {
32        let twice = AngularMomentumProjection::from_f64(value)?.value();
33        return Ok(Ratio::new(twice, 2));
34    }
35    let numerator = input
36        .getattr("numerator")
37        .and_then(|value| value.extract::<i32>());
38    let denominator = input
39        .getattr("denominator")
40        .and_then(|value| value.extract::<i32>());
41    if let (Ok(numerator), Ok(denominator)) = (numerator, denominator) {
42        if denominator == 0 {
43            return Err(LadduError::Custom(
44                "quantum number denominator cannot be zero".to_string(),
45            ));
46        }
47        return Ok(Ratio::new(numerator, denominator));
48    }
49    Err(LadduError::Custom(
50        "quantum number must be an int, float, or fractions.Fraction".to_string(),
51    ))
52}
53
54fn py_value_error(err: LadduError) -> PyErr {
55    PyValueError::new_err(err.to_string())
56}
57
58fn projection_to_python(
59    py: Python<'_>,
60    projection: AngularMomentumProjection,
61) -> PyResult<PyQuantumNumber> {
62    let twice = projection.value();
63    if twice % 2 == 0 {
64        Ok((twice / 2).into_bound_py_any(py)?.unbind())
65    } else {
66        let fractions = PyModule::import(py, "fractions")?;
67        let fraction = fractions.getattr("Fraction")?;
68        Ok(fraction.call1((twice, 2))?.unbind())
69    }
70}
71
72/// Enumerate allowed spin projections.
73#[pyfunction(name = "allowed_projections")]
74pub fn py_allowed_projections(
75    py: Python<'_>,
76    spin: &Bound<'_, PyAny>,
77) -> PyResult<Vec<PyQuantumNumber>> {
78    allowed_projections(parse_angular_momentum(spin)?)
79        .into_iter()
80        .map(|projection| projection_to_python(py, projection))
81        .collect()
82}
83
84/// Enumerate daughter helicity combinations.
85#[pyfunction(name = "helicity_combinations")]
86pub fn py_helicity_combinations(
87    py: Python<'_>,
88    spin_1: &Bound<'_, PyAny>,
89    spin_2: &Bound<'_, PyAny>,
90) -> PyResult<Vec<PyHelicityCombination>> {
91    helicity_combinations(
92        parse_angular_momentum(spin_1)?,
93        parse_angular_momentum(spin_2)?,
94    )
95    .into_iter()
96    .map(|combination| {
97        Ok((
98            projection_to_python(py, combination.lambda_1())?,
99            projection_to_python(py, combination.lambda_2())?,
100            projection_to_python(py, combination.helicity())?,
101        ))
102    })
103    .collect()
104}
laddu_python/utils/angular_momentum.rs

laddu_python/utils/
angular_momentum.rs
