1use rayon::prelude::*;
2use rustitude_core::prelude::*;
3use sphrs::SHCoordinates;
4
5use crate::utils::{Decay, Frame};
6
7#[derive(Clone)]
8pub struct TwoPiSDME<F: Field> {
9 decay: Decay,
10 frame: Frame,
11 data: Vec<(F, F, F, F, F, F)>,
12}
13
14impl<F: Field> TwoPiSDME<F> {
15 pub fn new(decay: Decay, frame: Frame) -> Self {
16 match decay {
17 Decay::TwoBodyDecay(_) => Self {
18 decay,
19 frame,
20 data: Vec::default(),
21 },
22
23 _ => unimplemented!(),
24 }
25 }
26}
27
28impl<F: Field> Node<F> for TwoPiSDME<F> {
29 fn precalculate(&mut self, dataset: &Dataset<F>) -> Result<(), RustitudeError> {
30 self.data = dataset
31 .events
32 .par_iter()
33 .map(|event| {
34 let (_, y, _, p) = self.decay.coordinates(self.frame, 0, event);
35 let big_phi = F::atan2(
36 y.dot(&event.eps),
37 event.beam_p4.direction().dot(&event.eps.cross(&y)),
38 );
39 let pgamma = event.eps_mag();
40 (
41 p.theta_cos().powi(2),
42 F::sin(p.theta()).powi(2),
43 F::sin(convert!(2, F) * p.theta()),
44 p.phi(),
45 big_phi,
46 pgamma,
47 )
48 })
49 .collect();
50 Ok(())
51 }
52
53 fn calculate(&self, parameters: &[F], event: &Event<F>) -> Result<Complex<F>, RustitudeError> {
54 let (cossqtheta, sinsqtheta, sin2theta, phi, big_phi, pgamma) = self.data[event.index];
55 let rho_000 = parameters[0];
56 let rho_100 = parameters[1];
57 let rho_1n10 = parameters[2];
58 let rho_111 = parameters[3];
59 let rho_001 = parameters[4];
60 let rho_101 = parameters[5];
61 let rho_1n11 = parameters[6];
62 let rho_102 = parameters[7];
63 let rho_1n12 = parameters[8];
64
65 Ok(Complex::from(F::sqrt(F::abs(
66 (convert!(3, F) / (convert!(4, F) * F::PI()))
67 * (convert!(0.5, F) * (F::one() - rho_000)
68 + convert!(0.5, F) * (convert!(3, F) * rho_000 - F::one()) * cossqtheta
69 - F::SQRT_2() * rho_100 * sin2theta * F::cos(phi)
70 - rho_1n10 * sinsqtheta * F::cos(convert!(2, F) * phi))
71 - pgamma
72 * F::cos(convert!(2, F) * big_phi)
73 * (rho_111 * sinsqtheta + rho_001 * cossqtheta
74 - F::SQRT_2() * rho_101 * sin2theta * F::cos(phi)
75 - rho_1n11 * sinsqtheta * F::cos(convert!(2, F) * phi))
76 - pgamma
77 * F::sin(convert!(2, F) * big_phi)
78 * (F::SQRT_2() * rho_102 * sin2theta * F::sin(phi)
79 + rho_1n12 * sinsqtheta * F::sin(convert!(2, F) * phi)),
80 ))))
81 }
82
83 fn parameters(&self) -> Vec<String> {
84 vec![
85 "rho_000".to_string(),
86 "rho_100".to_string(),
87 "rho_1n10".to_string(),
88 "rho_111".to_string(),
89 "rho_001".to_string(),
90 "rho_101".to_string(),
91 "rho_1n11".to_string(),
92 "rho_102".to_string(),
93 "rho_1n12".to_string(),
94 ]
95 }
96}
97
98#[derive(Clone)]
99pub struct ThreePiSDME<F: Field> {
100 decay: Decay,
101 frame: Frame,
102 data: Vec<(F, F, F, F, F, F)>,
103}
104
105impl<F: Field> ThreePiSDME<F> {
106 pub fn new(decay: Decay, frame: Frame) -> Self {
107 match decay {
108 Decay::ThreeBodyDecay(_) => Self {
109 decay,
110 frame,
111 data: Vec::default(),
112 },
113
114 _ => unimplemented!(),
115 }
116 }
117}
118
119impl<F: Field> Node<F> for ThreePiSDME<F> {
120 fn precalculate(&mut self, dataset: &Dataset<F>) -> Result<(), RustitudeError> {
121 self.data = dataset
122 .events
123 .par_iter()
124 .map(|event| {
125 let res_p4 = self.decay.resonance_p4(event);
126 let p1_res_p4 = self.decay.primary_p4(event).boost_along(&res_p4);
127 let p2_res_p4 = self.decay.primary_p4(event).boost_along(&res_p4);
128 let norm = p1_res_p4.momentum().cross(&p2_res_p4.momentum()).unit();
129 let (_, y, _, p) = self
130 .frame
131 .coordinates_from_boosted_vec(self.decay, &norm, event);
132 let big_phi = F::atan2(
133 y.dot(&event.eps),
134 event.beam_p4.direction().dot(&event.eps.cross(&y)),
135 );
136 let pgamma = event.eps_mag();
137 (
138 p.theta_cos().powi(2),
139 F::sin(p.theta()).powi(2),
140 F::sin(convert!(2, F) * p.theta()),
141 p.phi(),
142 big_phi,
143 pgamma,
144 )
145 })
146 .collect();
147 Ok(())
148 }
149
150 fn calculate(&self, parameters: &[F], event: &Event<F>) -> Result<Complex<F>, RustitudeError> {
151 let (cossqtheta, sinsqtheta, sin2theta, phi, big_phi, pgamma) = self.data[event.index];
152 let rho_000 = parameters[0];
153 let rho_100 = parameters[1];
154 let rho_1n10 = parameters[2];
155 let rho_111 = parameters[3];
156 let rho_001 = parameters[4];
157 let rho_101 = parameters[5];
158 let rho_1n11 = parameters[6];
159 let rho_102 = parameters[7];
160 let rho_1n12 = parameters[8];
161
162 Ok(Complex::from(F::sqrt(F::abs(
163 (convert!(3, F) / (convert!(4, F) * F::PI()))
164 * (convert!(0.5, F) * (F::one() - rho_000)
165 + convert!(0.5, F) * (convert!(3, F) * rho_000 - F::one()) * cossqtheta
166 - F::SQRT_2() * rho_100 * sin2theta * F::cos(phi)
167 - rho_1n10 * sinsqtheta * F::cos(convert!(2, F) * phi))
168 - pgamma
169 * F::cos(convert!(2, F) * big_phi)
170 * (rho_111 * sinsqtheta + rho_001 * cossqtheta
171 - F::SQRT_2() * rho_101 * sin2theta * F::cos(phi)
172 - rho_1n11 * sinsqtheta * F::cos(convert!(2, F) * phi))
173 - pgamma
174 * F::sin(convert!(2, F) * big_phi)
175 * (F::SQRT_2() * rho_102 * sin2theta * F::sin(phi)
176 + rho_1n12 * sinsqtheta * F::sin(convert!(2, F) * phi)),
177 ))))
178 }
179
180 fn parameters(&self) -> Vec<String> {
181 vec![
182 "rho_000".to_string(),
183 "rho_100".to_string(),
184 "rho_1n10".to_string(),
185 "rho_111".to_string(),
186 "rho_001".to_string(),
187 "rho_101".to_string(),
188 "rho_1n11".to_string(),
189 "rho_102".to_string(),
190 "rho_1n12".to_string(),
191 ]
192 }
193}
194
195#[derive(Clone)]
196pub struct VecRadiativeSDME<F: Field> {
197 decay: Decay,
198 frame: Frame,
199 data: Vec<(F, F, F, F, F, F)>,
200}
201
202impl<F: Field> VecRadiativeSDME<F> {
203 pub fn new(decay: Decay, frame: Frame) -> Self {
204 match decay {
205 Decay::TwoBodyDecay(_) => Self {
206 decay,
207 frame,
208 data: Vec::default(),
209 },
210
211 _ => unimplemented!(),
212 }
213 }
214}
215
216impl<F: Field> Node<F> for VecRadiativeSDME<F> {
217 fn precalculate(&mut self, dataset: &Dataset<F>) -> Result<(), RustitudeError> {
218 self.data = dataset
219 .events
220 .par_iter()
221 .map(|event| {
222 let (_, y, _, p) = self.decay.coordinates(self.frame, 0, event);
223 let big_phi = F::atan2(
224 y.dot(&event.eps),
225 event.beam_p4.direction().dot(&event.eps.cross(&y)),
226 );
227 let pgamma = event.eps_mag();
228 (
229 p.theta_cos().powi(2),
230 F::sin(p.theta()).powi(2),
231 F::sin(convert!(2, F) * p.theta()),
232 p.phi(),
233 big_phi,
234 pgamma,
235 )
236 })
237 .collect();
238 Ok(())
239 }
240
241 fn calculate(&self, parameters: &[F], event: &Event<F>) -> Result<Complex<F>, RustitudeError> {
242 let (cossqtheta, sinsqtheta, sin2theta, phi, big_phi, pgamma) = self.data[event.index];
243 let rho_000 = parameters[0];
244 let rho_100 = parameters[1];
245 let rho_1n10 = parameters[2];
246 let rho_111 = parameters[3];
247 let rho_001 = parameters[4];
248 let rho_101 = parameters[5];
249 let rho_1n11 = parameters[6];
250 let rho_102 = parameters[7];
251 let rho_1n12 = parameters[8];
252
253 Ok(Complex::from(F::sqrt(F::abs(
254 (convert!(3, F) / (convert!(8, F) * F::PI()))
255 * (F::one()
256 - sinsqtheta * convert!(0.5, F) * (F::one() - rho_000)
257 - rho_000 * cossqtheta
258 + rho_1n10 * sinsqtheta * F::cos(convert!(2, F) * phi)
259 + F::SQRT_2() * rho_100 * sin2theta * F::cos(phi)
260 - pgamma
261 * F::cos(convert!(2, F) * big_phi)
262 * (convert!(2, F) * rho_111
263 + (rho_001 - rho_111) * sinsqtheta
264 + rho_1n11 * sinsqtheta * F::cos(convert!(2, F) * phi)
265 + F::SQRT_2() * rho_101 * sin2theta * F::cos(phi))
266 + pgamma
267 * F::sin(convert!(2, F) * big_phi)
268 * (rho_1n12 * sinsqtheta * F::sin(convert!(2, F) * phi)
269 + F::SQRT_2() * rho_102 * sin2theta * F::sin(phi))),
270 ))))
271 }
272
273 fn parameters(&self) -> Vec<String> {
274 vec![
275 "rho_000".to_string(),
276 "rho_100".to_string(),
277 "rho_1n10".to_string(),
278 "rho_111".to_string(),
279 "rho_001".to_string(),
280 "rho_101".to_string(),
281 "rho_1n11".to_string(),
282 "rho_102".to_string(),
283 "rho_1n12".to_string(),
284 ]
285 }
286}