laddu_amplitudes/
ylm.rs

1use serde::{Deserialize, Serialize};
2
3use laddu_core::{
4    amplitudes::{Amplitude, AmplitudeID},
5    data::Event,
6    resources::{Cache, ComplexScalarID, Parameters, Resources},
7    utils::{
8        functions::spherical_harmonic,
9        variables::{Angles, Variable},
10    },
11    Complex, DVector, Float, LadduError,
12};
13
14#[cfg(feature = "python")]
15use laddu_python::{amplitudes::PyAmplitude, utils::variables::PyAngles};
16#[cfg(feature = "python")]
17use pyo3::prelude::*;
18
19/// An [`Amplitude`] for the spherical harmonic function $`Y_\ell^m(\theta, \phi)`$.
20#[derive(Clone, Serialize, Deserialize)]
21pub struct Ylm {
22    name: String,
23    l: usize,
24    m: isize,
25    angles: Angles,
26    csid: ComplexScalarID,
27}
28
29impl Ylm {
30    /// Construct a new [`Ylm`] with the given name, angular momentum (`l`) and moment (`m`) over
31    /// the given set of [`Angles`].
32    pub fn new(name: &str, l: usize, m: isize, angles: &Angles) -> Box<Self> {
33        Self {
34            name: name.to_string(),
35            l,
36            m,
37            angles: angles.clone(),
38            csid: ComplexScalarID::default(),
39        }
40        .into()
41    }
42}
43
44#[typetag::serde]
45impl Amplitude for Ylm {
46    fn register(&mut self, resources: &mut Resources) -> Result<AmplitudeID, LadduError> {
47        self.csid = resources.register_complex_scalar(None);
48        resources.register_amplitude(&self.name)
49    }
50
51    fn precompute(&self, event: &Event, cache: &mut Cache) {
52        cache.store_complex_scalar(
53            self.csid,
54            spherical_harmonic(
55                self.l,
56                self.m,
57                self.angles.costheta.value(event),
58                self.angles.phi.value(event),
59            ),
60        );
61    }
62
63    fn compute(&self, _parameters: &Parameters, _event: &Event, cache: &Cache) -> Complex<Float> {
64        cache.get_complex_scalar(self.csid)
65    }
66
67    fn compute_gradient(
68        &self,
69        _parameters: &Parameters,
70        _event: &Event,
71        _cache: &Cache,
72        _gradient: &mut DVector<Complex<Float>>,
73    ) {
74        // This amplitude is independent of free parameters
75    }
76}
77
78/// An spherical harmonic Amplitude
79///
80/// Computes a spherical harmonic (:math:`Y_{\ell}^m(\theta, \varphi)`)
81///
82/// Parameters
83/// ----------
84/// name : str
85///     The Amplitude name
86/// l : int
87///     The total orbital momentum (:math:`l \geq 0`)
88/// m : int
89///     The orbital moment (:math:`-l \leq m \leq l`)
90/// angles : laddu.Angles
91///     The spherical angles to use in the calculation
92///     
93/// Returns
94/// -------
95/// laddu.Amplitude
96///     An Amplitude which can be registered by a laddu.Manager
97///
98/// See Also
99/// --------
100/// laddu.Manager
101///
102#[cfg(feature = "python")]
103#[pyfunction(name = "Ylm")]
104pub fn py_ylm(name: &str, l: usize, m: isize, angles: &PyAngles) -> PyAmplitude {
105    PyAmplitude(Ylm::new(name, l, m, &angles.0))
106}
107
108#[cfg(test)]
109mod tests {
110    use std::sync::Arc;
111
112    use super::*;
113    use approx::assert_relative_eq;
114    use laddu_core::{data::test_dataset, Frame, Manager};
115
116    #[test]
117    fn test_ylm_evaluation() {
118        let mut manager = Manager::default();
119        let angles = Angles::new(0, [1], [2], [2, 3], Frame::Helicity);
120        let amp = Ylm::new("ylm", 1, 1, &angles);
121        let aid = manager.register(amp).unwrap();
122
123        let dataset = Arc::new(test_dataset());
124        let expr = aid.into();
125        let model = manager.model(&expr);
126        let evaluator = model.load(&dataset);
127
128        let result = evaluator.evaluate(&[]);
129
130        assert_relative_eq!(result[0].re, 0.27133944, epsilon = Float::EPSILON.sqrt());
131        assert_relative_eq!(result[0].im, 0.14268971, epsilon = Float::EPSILON.sqrt());
132    }
133
134    #[test]
135    fn test_ylm_gradient() {
136        let mut manager = Manager::default();
137        let angles = Angles::new(0, [1], [2], [2, 3], Frame::Helicity);
138        let amp = Ylm::new("ylm", 1, 1, &angles);
139        let aid = manager.register(amp).unwrap();
140
141        let dataset = Arc::new(test_dataset());
142        let expr = aid.into();
143        let model = manager.model(&expr);
144        let evaluator = model.load(&dataset);
145
146        let result = evaluator.evaluate_gradient(&[]);
147        assert_eq!(result[0].len(), 0); // amplitude has no parameters
148    }
149}
laddu_amplitudes/ylm.rs

laddu_amplitudes/
ylm.rs
