1use std::f64::consts::PI;
4
5use crate::{imag, real};
6use ndarray::Array;
7use num_complex::Complex64;
8use statrs::function::erf::erf;
9
10use serde::{Deserialize, Serialize};
11
12const J: Complex64 = Complex64::new(0f64, 1f64);
13
14pub fn apply_phase_and_detuning(
16 iq_values: &mut [Complex64],
17 phase: f64,
18 detuning: f64,
19 sample_rate: f64,
20) {
21 for (index, value) in iq_values.iter_mut().enumerate() {
22 *value *= (2.0 * PI * J * (detuning * (index as f64) / sample_rate + phase)).exp();
23 }
24}
25
26fn ceiling_with_epsilon(value: f64) -> f64 {
33 let truncated = value - (value * 10.0 * f64::EPSILON);
34 truncated.ceil()
35}
36
37fn polar_to_rectangular(magnitude: f64, angle: crate::units::Radians<f64>) -> Complex64 {
39 magnitude * imag!(angle.0).exp()
40}
41
42pub trait WaveformTemplate {
43 fn into_iq_values(self) -> Vec<Complex64>;
44}
45
46#[derive(Clone, Copy, Debug, PartialEq, Serialize, Deserialize)]
47pub struct BoxcarKernel {
48 pub phase: crate::units::Cycles<f64>,
49 pub scale: f64,
50 pub sample_count: u64,
51}
52
53impl BoxcarKernel {
54 pub fn into_iq_value(self) -> Complex64 {
55 polar_to_rectangular(self.scale / self.sample_count as f64, self.phase.into())
56 }
57}
58
59#[derive(Clone, Copy, Debug, PartialEq, Serialize, Deserialize)]
61pub struct ErfSquare {
62 pub duration: f64,
64 pub risetime: f64,
66 pub sample_rate: f64,
68 pub pad_left: f64,
70 pub pad_right: f64,
72 pub positive_polarity: bool,
74 pub scale: f64,
76 pub phase: f64,
78 pub detuning: f64,
80}
81
82impl WaveformTemplate for ErfSquare {
83 fn into_iq_values(self) -> Vec<Complex64> {
84 let length = ceiling_with_epsilon(self.duration * self.sample_rate);
85 let mut time_steps = Array::<f64, _>::range(0f64, length, 1f64) / self.sample_rate;
86
87 let fwhm = 0.5 * self.risetime;
88 let t1 = fwhm;
89 let t2 = self.duration - fwhm;
90 let sigma = 0.5 * fwhm / (2f64 * 2f64.ln()).sqrt();
91 time_steps.mapv_inplace(|el| 0.5 * (erf((el - t1) / sigma) - erf((el - t2) / sigma)));
92 if !self.positive_polarity {
93 time_steps *= -1f64;
94 }
95
96 let left_pad_length = (self.pad_left * self.sample_rate).ceil() as usize;
97 let mut waveform = vec![real!(0f64); left_pad_length];
98 waveform.extend(time_steps.into_iter().map(|el| self.scale * real!(el)));
99
100 let right_pad_length = (self.pad_right * self.sample_rate).ceil() as usize;
101 waveform.extend_from_slice(&vec![real!(0f64); right_pad_length]);
102 apply_phase_and_detuning(&mut waveform, self.phase, self.detuning, self.sample_rate);
103 waveform
104 }
105}
106
107#[derive(Clone, Copy, Debug, PartialEq, Serialize, Deserialize)]
109pub struct Gaussian {
110 pub duration: f64,
112 pub fwhm: f64,
114 pub t0: f64,
116 pub sample_rate: f64,
118 pub scale: f64,
120 pub phase: f64,
122 pub detuning: f64,
124}
125
126impl WaveformTemplate for Gaussian {
127 fn into_iq_values(self) -> Vec<Complex64> {
128 let length = ceiling_with_epsilon(self.duration * self.sample_rate);
129 let mut time_steps = Array::<f64, _>::range(0f64, length, 1f64) / self.sample_rate;
130
131 let sigma = 0.5 * self.fwhm / (2f64 * 2f64.ln()).sqrt();
132 time_steps.mapv_inplace(|el| (-0.5 * (el - self.t0).powf(2f64) / sigma.powf(2f64)).exp());
133
134 let mut waveform = time_steps
135 .into_iter()
136 .map(|el| self.scale * real!(el))
137 .collect::<Vec<_>>();
138 apply_phase_and_detuning(&mut waveform, self.phase, self.detuning, self.sample_rate);
139 waveform
140 }
141}
142
143#[derive(Clone, Copy, Debug, PartialEq, Serialize, Deserialize)]
149pub struct DragGaussian {
150 pub duration: f64,
152 pub fwhm: f64,
154 pub t0: f64,
156 pub anh: f64,
158 pub alpha: f64,
160 pub sample_rate: f64,
162 pub scale: f64,
164 pub phase: f64,
166 pub detuning: f64,
168}
169
170impl WaveformTemplate for DragGaussian {
171 fn into_iq_values(self) -> Vec<Complex64> {
172 let length = ceiling_with_epsilon(self.duration * self.sample_rate);
173 let time_steps = Array::<f64, _>::range(0f64, length, 1f64) / self.sample_rate;
174
175 let sigma = 0.5 * self.fwhm / (2f64 * 2f64.ln()).sqrt();
176
177 let mut waveform = vec![Complex64::new(0f64, 0f64,); length as usize];
178 for i in 0..length as usize {
179 let env = (-0.5 * (time_steps[i] - self.t0).powf(2f64) / sigma.powf(2f64)).exp();
181 let env_mod = (self.alpha * (1f64 / (2f64 * PI * self.anh * sigma.powf(2f64))))
183 * (time_steps[i] - self.t0)
184 * env;
185 waveform[i] = self.scale * Complex64::new(env, env_mod);
186 }
187 apply_phase_and_detuning(&mut waveform, self.phase, self.detuning, self.sample_rate);
188
189 waveform
190 }
191}
192
193#[derive(Clone, Copy, Debug, PartialEq, Serialize, Deserialize)]
203pub struct HermiteGaussian {
204 pub duration: f64,
206 pub fwhm: f64,
208 pub t0: f64,
210 pub anh: f64,
212 pub alpha: f64,
214 pub sample_rate: f64,
216 pub second_order_hrm_coeff: f64,
218 pub scale: f64,
220 pub phase: f64,
222 pub detuning: f64,
224}
225
226impl WaveformTemplate for HermiteGaussian {
227 fn into_iq_values(self) -> Vec<Complex64> {
228 let length = ceiling_with_epsilon(self.duration * self.sample_rate);
229 let time_steps = Array::<f64, _>::range(0f64, length, 1f64) / self.sample_rate;
230
231 let sigma = 0.5 * self.fwhm / (2f64 * 2f64.ln()).sqrt();
232 let deriv_prefactor = -self.alpha / (2f64 * PI * self.anh);
233
234 let mut waveform = vec![Complex64::new(0f64, 0f64,); length as usize];
235 for i in 0..length as usize {
236 let exp_t = 0.5 * (time_steps[i] - self.t0).powf(2f64) / sigma.powf(2f64);
237 let g = (-exp_t).exp();
238 let env = (1f64 - self.second_order_hrm_coeff * exp_t) * g;
239 let env_derived = deriv_prefactor * (time_steps[i] - self.t0) / sigma.powf(2f64)
240 * g
241 * (self.second_order_hrm_coeff * (exp_t - 1f64) - 1f64);
242 waveform[i] = self.scale * Complex64::new(env, env_derived);
243 }
244 apply_phase_and_detuning(&mut waveform, self.phase, self.detuning, self.sample_rate);
245
246 waveform
247 }
248}
249
250#[cfg(test)]
251mod tests {
252 use super::*;
253
254 fn format_snapshot_name(item: impl std::fmt::Debug) -> String {
257 format!("{item:?}")
258 .replace(['{', '}', ':', ','], "")
259 .replace([' ', '.'], "_")
260 }
261
262 fn assert_almost_eq(left: Complex64, right: Complex64, epsilon: f64) {
263 assert!(
264 (left - right).norm() < epsilon,
265 "Expected {} to be almost equal to {} with epsilon {}",
266 left,
267 right,
268 epsilon
269 );
270 }
271
272 #[rstest::rstest]
273 #[case(BoxcarKernel {
274 phase: crate::units::Cycles(0.0),
275 scale: 1.0,
276 sample_count: 10,
277 }, Complex64::new(0.1, 0.0))]
278 #[case(BoxcarKernel {
279 phase: crate::units::Cycles(0.5),
280 scale: 1.0,
281 sample_count: 10,
282 }, Complex64::new(-0.1, 0.0))]
283 #[case(BoxcarKernel {
284 phase: crate::units::Cycles(0.0),
285 scale: -1.0,
286 sample_count: 10,
287 }, Complex64::new(-0.1, 0.0))]
288 #[case(BoxcarKernel {
289 phase: crate::units::Cycles(0.0),
290 scale: 0.0,
291 sample_count: 10,
292 }, Complex64::new(0.0, 0.0))]
293 fn boxcar_kernel(#[case] kernel: BoxcarKernel, #[case] expected: Complex64) {
294 let iq_value = kernel.into_iq_value();
295 assert_almost_eq(iq_value, expected, 1e-10);
296 }
297
298 #[rstest::rstest]
299 #[case(0.0, 0.0)]
300 #[case(0.0-f64::EPSILON, 0.0)]
301 #[case(0.0+f64::EPSILON, 1.0)]
302 #[case(8.800_000_000_000_001e-8 * 1.0e9, 88.0)]
304 fn ceiling_with_epsilon(#[case] value: f64, #[case] expected: f64) {
305 let result = super::ceiling_with_epsilon(value);
306 assert_eq!(result, expected);
307 }
308
309 #[rstest::rstest]
324 #[case(ErfSquare { duration: 1e-4, risetime: 1e-5, sample_rate: 1e6, pad_left: 0.0, pad_right: 0.0, positive_polarity: true, scale: 1.0, phase: 0.0, detuning: 0.0 })]
325 #[case(ErfSquare { duration: 1e-4, risetime: 1e-5, sample_rate: 1e6, pad_left: 1e-5, pad_right: 2e-5, positive_polarity: false, scale: 1.0, phase: 0.0, detuning: 0.0 })]
326 #[case(ErfSquare { duration: 1e-4, risetime: 1e-4, sample_rate: 1e6, pad_left: 0.0, pad_right: 0.0, positive_polarity: false, scale: 1.0, phase: 0.0, detuning: 1e6 })]
327 #[case(ErfSquare { duration: 1e-4, risetime: 1e-5, sample_rate: 1e6, pad_left: 0.0, pad_right: 0.0, positive_polarity: true, scale: 1.0, phase: 0.5, detuning: 0.0 })]
328 #[case(ErfSquare { duration: 1e-4, risetime: 1e-5, sample_rate: 1e6, pad_left: 0.0, pad_right: 0.0, positive_polarity: true, scale: -1.0, phase: 0.0, detuning: 0.0 })]
329 #[case(ErfSquare { duration: 1e-4, risetime: 1e-5, sample_rate: 1e6, pad_left: 0.0, pad_right: 0.0, positive_polarity: true, scale: 0.0, phase: 0.0, detuning: 0.0 })]
330 #[case(ErfSquare { duration: 1e-4, risetime: 1e-5, sample_rate: 1e6, pad_left: 0.0, pad_right: 0.0, positive_polarity: false, scale: 0.5, phase: 0.4, detuning: -1e6 })]
331 #[case(Gaussian { duration: 1e-4, fwhm: 1e-5, t0: 0.0, sample_rate: 1e6, scale: 1.0, phase: 0.0, detuning: 0.0 })]
332 #[case(Gaussian { duration: 1e-4, fwhm: 1e-5, t0: 5e-5, sample_rate: 1e6, scale: 1.0, phase: 0.0, detuning: 1e6 })]
333 #[case(Gaussian { duration: 1e-4, fwhm: 2e-5, t0: 5e-5, sample_rate: 1e6, scale: 0.5, phase: 0.0, detuning: 0.0 })]
334 #[case(Gaussian { duration: 1e-4, fwhm: 4e-5, t0: 5e-5, sample_rate: 1e6, scale: 0.5, phase: 0.5, detuning: 0.0 })]
335 #[case(Gaussian { duration: 1e-4, fwhm: 4e-5, t0: 5e-5, sample_rate: 1e6, scale: -1.0, phase: 0.0, detuning: 0.0 })]
336 #[case(DragGaussian { duration: 1e-4, fwhm: 1e-5, t0: 0.0, anh: 1e6, alpha: 1.0, sample_rate: 1e6, scale: 1.0, phase: 0.0, detuning: 0.0 })]
337 #[case(DragGaussian { duration: 1e-4, fwhm: 1e-5, t0: 0.0, anh: 1e6, alpha: 1.0, sample_rate: 1e6, scale: 1.0, phase: 0.0, detuning: 1e6 })]
338 #[case(HermiteGaussian { duration: 1e-4, fwhm: 1e-5, t0: 0.0, anh: 1e6, alpha: 1.0, sample_rate: 1e6, second_order_hrm_coeff: 0.1, scale: 1.0, phase: 0.0, detuning: 0.0 })]
339 #[case(HermiteGaussian { duration: 1e-4, fwhm: 1e-5, t0: 0.0, anh: 1e6, alpha: 1.0, sample_rate: 1e6, second_order_hrm_coeff: 0.1, scale: 1.0, phase: 0.0, detuning: 1e6 })]
340 fn into_iq_values(#[case] parameters: impl WaveformTemplate + Copy + std::fmt::Debug) {
341 let iq_values = parameters.into_iq_values();
342 let count = iq_values.len();
343
344 let all_values_zero = iq_values.iter().all(|el| el == &Complex64::new(0.0, 0.0));
345
346 if count <= 200 && !all_values_zero {
349 let split = iq_values.clone().into_iter().fold(
350 (vec![], vec![]),
351 |(mut reals, mut imags), el| {
352 reals.push(el.re);
353 imags.push(el.im);
354 (reals, imags)
355 },
356 );
357 let split = vec![split.0, split.1];
358
359 let res = rasciigraph::plot_many(
360 split,
361 rasciigraph::Config::default()
362 .with_width(count as u32 + 10)
363 .with_height(20),
364 );
365
366 insta::assert_snapshot!(format!("{}__plot", format_snapshot_name(parameters)), res);
369 }
370
371 let neat_iq_values = iq_values
372 .iter()
373 .map(|el| format!("{:+.5e}, {:+.5e}", el.re, el.im))
374 .collect::<Vec<_>>()
375 .join("\n");
376
377 insta::assert_snapshot!(
378 format!("{}__data", format_snapshot_name(parameters)),
379 neat_iq_values
380 )
381 }
382}