1use rayon::prelude::*;
2use rustitude_core::{convert, prelude::*};
3use sphrs::{ComplexSH, SHEval};
4
5use crate::utils::{Decay, Frame, Sign, Wave};
6
7#[derive(Clone)]
8pub struct Ylm<F: Field> {
9 wave: Wave,
10 decay: Decay,
11 frame: Frame,
12 data: Vec<Complex<F>>,
13}
14impl<F: Field> Ylm<F> {
15 pub fn new(wave: Wave, decay: Decay, frame: Frame) -> Self {
16 Self {
17 wave,
18 decay,
19 frame,
20 data: Vec::default(),
21 }
22 }
23}
24impl<F: Field> Node<F> for Ylm<F> {
25 fn precalculate(&mut self, dataset: &Dataset<F>) -> Result<(), RustitudeError> {
26 self.data = dataset
27 .events
28 .par_iter()
29 .map(|event| {
30 let (_, _, _, p) =
31 self.frame
32 .coordinates(self.decay, self.decay.primary_p4(event), event);
33 ComplexSH::Spherical.eval(self.wave.l(), self.wave.m(), &p)
34 })
35 .collect();
36 Ok(())
37 }
38
39 fn calculate(&self, _parameters: &[F], event: &Event<F>) -> Result<Complex<F>, RustitudeError> {
40 Ok(self.data[event.index])
41 }
42}
43
44#[derive(Clone)]
45pub struct Zlm<F: Field> {
46 wave: Wave,
47 reflectivity: Sign,
48 decay: Decay,
49 frame: Frame,
50 data: Vec<Complex<F>>,
51}
52impl<F: Field> Zlm<F> {
53 pub fn new(wave: Wave, reflectivity: Sign, decay: Decay, frame: Frame) -> Self {
54 Self {
55 wave,
56 reflectivity,
57 decay,
58 frame,
59 data: Vec::default(),
60 }
61 }
62}
63impl<F: Field + num::Float> Node<F> for Zlm<F> {
64 fn precalculate(&mut self, dataset: &Dataset<F>) -> Result<(), RustitudeError> {
65 self.data = dataset
66 .events
67 .par_iter()
68 .map(|event| {
69 let (_, y, _, p) = self.decay.coordinates(self.frame, 0, event);
70 let ylm = ComplexSH::Spherical.eval(self.wave.l(), self.wave.m(), &p);
71 let big_phi = F::atan2(
72 y.dot(&event.eps),
73 event.beam_p4.direction().dot(&event.eps.cross(&y)),
74 );
75 let pgamma = event.eps_mag();
76 let phase = Complex::cis(-big_phi);
77 let zlm = ylm * phase;
78 match self.reflectivity {
79 Sign::Positive => Complex::new(
80 F::sqrt(F::one() + pgamma) * zlm.re,
81 F::sqrt(F::one() - pgamma) * zlm.im,
82 ),
83 Sign::Negative => Complex::new(
84 F::sqrt(F::one() - pgamma) * zlm.re,
85 F::sqrt(F::one() + pgamma) * zlm.im,
86 ),
87 }
88 })
89 .collect();
90 Ok(())
91 }
92 fn calculate(&self, _parameters: &[F], event: &Event<F>) -> Result<Complex<F>, RustitudeError> {
93 Ok(self.data[event.index])
94 }
95}
96
97#[derive(Clone)]
98pub struct OnePS<F: Field> {
99 reflectivity: Sign,
100 decay: Decay,
101 frame: Frame,
102 data: Vec<Complex<F>>,
103}
104impl<F: Field> OnePS<F> {
105 pub fn new(reflectivity: Sign, decay: Decay, frame: Frame) -> Self {
106 Self {
107 reflectivity,
108 decay,
109 frame,
110 data: Vec::default(),
111 }
112 }
113}
114impl<F: Field> Node<F> for OnePS<F> {
115 fn precalculate(&mut self, dataset: &Dataset<F>) -> Result<(), RustitudeError> {
116 self.data = dataset
117 .events
118 .par_iter()
119 .map(|event| {
120 let (_, y, _, _) = self.decay.coordinates(self.frame, 0, event);
121 let pol_angle = F::acos(event.eps[0]);
122 let big_phi = F::atan2(
123 y.dot(&event.eps),
124 event.beam_p4.direction().dot(&event.eps.cross(&y)),
125 );
126 let pgamma = event.eps_mag();
127 let phase = Complex::cis(-(pol_angle + big_phi));
128 match self.reflectivity {
129 Sign::Positive => Complex::new(
130 F::sqrt(F::one() + pgamma) * phase.re,
131 F::sqrt(F::one() - pgamma) * phase.im,
132 ),
133 Sign::Negative => Complex::new(
134 F::sqrt(F::one() - pgamma) * phase.re,
135 F::sqrt(F::one() + pgamma) * phase.im,
136 ),
137 }
138 })
139 .collect();
140 Ok(())
141 }
142
143 fn calculate(&self, _parameters: &[F], event: &Event<F>) -> Result<Complex<F>, RustitudeError> {
144 Ok(self.data[event.index])
145 }
146}
147
148#[derive(Clone)]
149pub struct TwoPS<F: Field> {
150 wave: Wave,
151 reflectivity: Sign,
152 decay: Decay,
153 frame: Frame,
154 data: Vec<Complex<F>>,
155}
156impl<F: Field> TwoPS<F> {
157 pub fn new(wave: Wave, reflectivity: Sign, decay: Decay, frame: Frame) -> Self {
158 Self {
159 wave,
160 reflectivity,
161 decay,
162 frame,
163 data: Vec::default(),
164 }
165 }
166}
167impl<F: Field> Node<F> for TwoPS<F> {
168 fn precalculate(&mut self, dataset: &Dataset<F>) -> Result<(), RustitudeError> {
169 self.data = dataset
170 .events
171 .par_iter()
172 .map(|event| {
173 let (_, _, _, p) = self.decay.coordinates(self.frame, 0, event);
174 let ylm_p = ComplexSH::Spherical
175 .eval(self.wave.l(), self.wave.m(), &p)
176 .conj();
177 let ylm_m = ComplexSH::Spherical
178 .eval(self.wave.l(), -self.wave.m(), &p)
179 .conj();
180 let m_refl = convert!(
181 if self.wave.m() % 2 == 0 {
182 self.reflectivity as isize
183 } else {
184 -(self.reflectivity as isize)
185 },
186 F
187 );
188 let big_theta = match self.wave.m().cmp(&0) {
189 std::cmp::Ordering::Less => F::zero(),
190 std::cmp::Ordering::Equal => convert!(0.5, F),
191 std::cmp::Ordering::Greater => F::FRAC_1_SQRT_2(),
192 };
193 let wigner_d_lm0_m = ylm_m.scale(F::sqrt(
194 convert!(4, F) * F::PI()
195 / (convert!(2, F) * convert!(self.wave.l(), F) + F::one()),
196 ));
197 ylm_p.scale(big_theta) - wigner_d_lm0_m.scale(m_refl)
198 })
199 .collect();
200 Ok(())
201 }
202
203 fn calculate(&self, _parameters: &[F], event: &Event<F>) -> Result<Complex<F>, RustitudeError> {
204 Ok(self.data[event.index])
205 }
206}