1use crate::{builders::SourceBuilder, cu::Int, utilities::Mask};
26
27use super::{cu::Double, cu::Single, Centroiding, Cu, FromBuilder};
28use ffi::{bundle, dev2host, dev2host_int, source, vector};
29use serde::{Deserialize, Serialize};
30use skyangle::Conversion;
31
32use std::{
33 cell::UnsafeCell,
34 f32,
35 ffi::{CStr, CString},
36 fmt::Display,
37 usize,
38};
39
40pub(crate) const PHOTOMETRY: [&str; 10] = ["V", "Vs", "R", "I", "J", "H", "K", "Ks", "R+I", "VIS"];
41
42pub trait Propagation {
44 fn propagate(&mut self, src: &mut Source);
45 fn time_propagate(&mut self, secs: f64, src: &mut Source);
46}
47
48#[derive(Clone, Debug, PartialEq, Serialize, Deserialize)]
49#[serde(tag = "type")]
50pub enum PupilSampling {
51 SquareGrid {
52 size: Option<f64>,
53 resolution: usize,
54 },
55 UserSet(usize),
56}
57impl PupilSampling {
58 pub fn total(&self) -> usize {
59 match &self {
60 PupilSampling::SquareGrid { resolution, .. } => resolution * resolution,
61 PupilSampling::UserSet(n) => *n,
62 }
63 }
64 pub fn side(&self) -> usize {
65 match &self {
66 PupilSampling::SquareGrid { resolution, .. } => *resolution,
67 PupilSampling::UserSet(n) => *n,
68 }
69 }
70 pub fn size(&self) -> Option<f64> {
71 match &self {
72 PupilSampling::SquareGrid { size, .. } => *size,
73 _ => None,
74 }
75 }
76}
77
78pub struct Source {
80 pub(crate) _c_: source,
81 pub size: i32,
83 pub pupil_size: f64,
85 pub pupil_sampling: i32,
87 pub _wfe_rms: Vec<f32>,
88 pub _phase: Vec<f32>,
89 pub zenith: Vec<f32>,
90 pub azimuth: Vec<f32>,
91 pub magnitude: Vec<f32>,
92}
93impl Display for Source {
94 fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
95 writeln!(
96 f,
97 "Source(x{}) @ λ={:.0}nm",
98 self.size,
99 self.wavelength() * 1e9
100 )?;
101 writeln!(
102 f,
103 " zenith: {:.2?}arcsec",
104 self.zenith
105 .iter()
106 .map(|x| x.to_arcsec())
107 .collect::<Vec<_>>()
108 )?;
109 writeln!(
110 f,
111 " azimuth: {:.2?}deg",
112 self.azimuth
113 .iter()
114 .map(|x| x.to_degrees())
115 .collect::<Vec<_>>()
116 )?;
117 writeln!(f, " magnitude: {:?}", self.magnitude)?;
118 writeln!(
119 f,
120 " pupil: (size: {}m, sampling: {}px)",
121 self.pupil_size, self.pupil_sampling
122 )?;
123 Ok(())
124 }
125}
126
127impl PartialEq for Source {
128 fn eq(&self, other: &Self) -> bool {
129 Into::<SourceBuilder>::into(self) == Into::<SourceBuilder>::into(other)
130 }
131}
132impl FromBuilder for Source {
133 type ComponentBuilder = SourceBuilder;
134}
135impl Source {
136 pub fn empty() -> Source {
138 Source {
139 _c_: Default::default(),
140 size: 0,
141 pupil_size: 0.0,
142 pupil_sampling: 0,
143 _wfe_rms: vec![],
144 _phase: vec![],
145 zenith: vec![],
146 azimuth: vec![],
147 magnitude: vec![],
148 }
149 }
150 pub fn new(size: i32, pupil_size: f64, pupil_sampling: i32) -> Source {
155 Source {
156 _c_: Default::default(),
157 size,
158 pupil_size,
159 pupil_sampling,
160 _wfe_rms: vec![0.0; size as usize],
161 _phase: vec![0.0; (pupil_sampling * pupil_sampling * size) as usize],
162 zenith: vec![0.0; size as usize],
163 azimuth: vec![0.0; size as usize],
164 magnitude: vec![0.0; size as usize],
165 }
166 }
167 pub fn pupil_sampling(&self) -> usize {
168 self.pupil_sampling as usize
169 }
170 pub fn from(args: (i32, f64, i32)) -> Source {
171 Source::new(args.0, args.1, args.2)
172 }
173 pub fn build(
180 &mut self,
181 band: &str,
182 mut zenith: Vec<f32>,
183 mut azimuth: Vec<f32>,
184 mut magnitude: Vec<f32>,
185 ) -> &mut Self {
186 assert_eq!(zenith.len(), azimuth.len());
187 assert_eq!(zenith.len(), magnitude.len());
188 let band = CString::new(band).unwrap();
189 unsafe {
190 let origin = vector {
191 x: 0.0,
192 y: 0.0,
193 z: 25.0,
194 };
195 self.magnitude.copy_from_slice(magnitude.as_slice());
196 self._c_.setup7(
197 band.into_raw(),
198 magnitude.as_mut_ptr(),
199 zenith.as_mut_ptr(),
200 azimuth.as_mut_ptr(),
201 f32::INFINITY,
202 self.size,
203 self.pupil_size,
204 self.pupil_sampling,
205 origin,
206 );
207 }
208 self
209 }
210 pub fn as_raw_mut_ptr(&mut self) -> &mut source {
211 &mut self._c_
212 }
213 pub fn get_photometric_band(&self) -> String {
215 unsafe {
216 String::from(
217 CStr::from_ptr(self._c_.photometric_band)
218 .to_str()
219 .expect("CStr::to_str failed"),
220 )
221 }
222 }
223 pub fn xy(&self) -> (f64, f64) {
225 (self._c_.theta_x as f64, self._c_.theta_x as f64)
226 }
227 pub fn update(&mut self, mut zenith: Vec<f64>, mut azimuth: Vec<f64>) {
229 unsafe {
230 self._c_.update_directions(
231 zenith.as_mut_ptr(),
232 azimuth.as_mut_ptr(),
233 zenith.len() as i32,
234 );
235 }
236 }
237 pub fn wavelength(&self) -> f64 {
239 unsafe {
240 let src = UnsafeCell::new(self._c_);
241 (&mut *src.get()).wavelength() as f64
242 }
243 }
244 pub fn fwhm(&mut self, value: f64) {
246 self._c_.fwhm = value as f32;
247 }
248 pub fn rotate_rays(&mut self, angle: f64) {
250 self._c_.rays.rot_angle = angle;
251 }
252 pub fn xpupil(&mut self) -> &mut Self {
254 unsafe {
255 self._c_.wavefront.reset();
256 self._c_.opd2phase();
257 }
258 self
259 }
260 pub fn opd2phase(&mut self) -> &mut Self {
261 unsafe {
262 self._c_.opd2phase();
264 }
265 self
266 }
267 pub fn wfe_rms(&mut self) -> Vec<f64> {
269 unsafe {
270 self._c_.wavefront.rms(self._wfe_rms.as_mut_ptr());
271 }
272 self._wfe_rms
273 .clone()
274 .into_iter()
275 .map(|x| x as f64)
276 .collect()
277 }
278 pub fn wfe_rms_10e(&mut self, exp: i32) -> Vec<f64> {
280 unsafe {
281 self._c_.wavefront.rms(self._wfe_rms.as_mut_ptr());
282 }
283 self._wfe_rms
284 .iter()
285 .map(|x| *x as f64 * 10_f64.powi(-exp))
286 .collect()
287 }
288 pub fn gradients(&mut self) -> Vec<f64> {
289 let mut sxy: Vec<Vec<f32>> = vec![vec![0.; self.size as usize]; 2];
290 unsafe {
291 self._c_.wavefront.gradient_average1(
292 sxy[0].as_mut_ptr(),
293 sxy[1].as_mut_ptr(),
294 self._c_.rays.L as f32,
295 )
296 }
297 sxy.into_iter().flatten().map(|x| x as f64).collect()
298 }
299 pub fn segment_wfe(&mut self) -> Vec<(f64, f64)> {
301 let n_ray_total = self._c_.rays.N_RAY_TOTAL as usize;
302 let n_ray = n_ray_total / self.size as usize;
303 let mut mask = vec![0i32; n_ray_total];
304 unsafe {
305 dev2host_int(
306 mask.as_mut_ptr(),
307 self._c_.rays.d__piston_mask,
308 n_ray_total as i32,
309 );
310 }
311 self.phase();
312 let mut segment_wfe: Vec<(f64, f64)> = vec![];
313 for (mask, phase) in mask.chunks(n_ray).zip(self._phase.chunks(n_ray)) {
314 for k in 1..8 {
315 let segment_phase = mask
316 .iter()
317 .zip(phase)
318 .filter_map(|(&mask, &phase)| (mask == k).then_some(phase))
319 .collect::<Vec<f32>>();
320 let n = segment_phase.len() as f32;
321 if n > 0. {
322 let mean = segment_phase.iter().sum::<f32>() / n;
323 let var = segment_phase
324 .iter()
325 .map(|x| (x - mean).powi(2))
326 .sum::<f32>()
327 / n;
328 segment_wfe.push((mean as f64, var.sqrt() as f64));
329 }
330 }
331 }
332 segment_wfe
333 }
334 pub fn segment_wfe_10e(&mut self, exp: i32) -> Vec<(f64, f64)> {
335 self.segment_wfe()
336 .into_iter()
337 .map(|(p, s)| (p * 10_f64.powi(-exp), s * 10_f64.powi(-exp)))
338 .collect()
339 }
340 pub fn segment_dwfe_10e(&mut self, exp: i32) -> Vec<(f64, f64)> {
341 let data = self.segment_wfe();
342 let p7 = data[6].0;
343 data.into_iter()
344 .map(|(p, s)| ((p - p7) * 10_f64.powi(-exp), s * 10_f64.powi(-exp)))
345 .collect()
346 }
347 pub fn add_piston(&mut self, piston: &[f64]) -> &mut Self {
349 let n_ray_total = self._c_.rays.N_RAY_TOTAL as usize;
350 let n_ray = n_ray_total / self.size as usize;
351 let mask = self.segment_mask();
352 let mut opd = vec![0f32; n_ray_total];
353 for (mask, opd) in mask.chunks(n_ray).zip(opd.chunks_mut(n_ray)) {
354 for k in 1..8 {
355 mask.iter()
356 .zip(&mut *opd)
357 .filter_map(|(&mask, opd)| (mask == k).then_some(opd))
358 .for_each(|opd| *opd = piston[k as usize - 1] as f32)
359 }
360 }
361 self.add(&opd);
362 self
363 }
364 pub fn segment_wfe_rms(&mut self) -> Vec<f64> {
366 self.segment_wfe().into_iter().map(|(_, s)| s).collect()
367 }
368 pub fn segment_wfe_rms_10e(&mut self, exp: i32) -> Vec<f64> {
369 self.segment_wfe_rms()
370 .into_iter()
371 .map(|x| x * 10_f64.powi(-exp))
372 .collect()
373 }
374 pub fn segment_piston(&mut self) -> Vec<f64> {
376 self.segment_wfe().into_iter().map(|(p, _)| p).collect()
377 }
378 pub fn segment_piston_10e(&mut self, exp: i32) -> Vec<f64> {
379 self.segment_piston()
380 .iter()
381 .map(|x| x * 10_f64.powi(-exp))
382 .collect()
383 }
384 pub fn segment_mask(&mut self) -> Vec<i32> {
385 let mut mask = vec![0i32; self._c_.rays.N_RAY_TOTAL as usize];
386 unsafe {
387 dev2host_int(
388 mask.as_mut_ptr(),
389 self._c_.rays.d__piston_mask,
390 self._c_.rays.N_RAY_TOTAL,
391 );
392 }
393 mask
394 }
395 pub fn segment_gradients(&mut self) -> Vec<f64> {
397 let mut sxy: Vec<Vec<f32>> = vec![vec![0.; 7 * self.size as usize]; 2];
398 unsafe {
399 self._c_.wavefront.segments_gradient_averageFast(
400 sxy[0].as_mut_ptr(),
401 sxy[1].as_mut_ptr(),
402 self._c_.rays.L as f32,
403 self._c_.rays.d__piston_mask,
404 );
405 }
406 sxy.into_iter()
407 .flat_map(|x| {
408 x.into_iter()
409 .filter(|x| !x.is_nan())
410 .map(|x| x as f64)
411 .collect::<Vec<f64>>()
412 })
413 .collect()
414 }
415 pub fn lenslet_gradients(
417 &mut self,
418 n_lenslet: i32,
419 _lenslet_size: f64,
420 data: &mut Centroiding,
421 ) {
422 let lenslet_size = self.pupil_size / n_lenslet as f64;
423 unsafe {
424 if data.n_valid_lenslet_total() < data.n_lenslet_total {
425 self._c_.wavefront.finite_difference1(
426 data.__mut_ceo__().0.d__cx,
427 data.__mut_ceo__().0.d__cy,
428 n_lenslet,
429 lenslet_size as f32,
430 data.__mut_ceo__().1,
431 );
432 } else {
433 self._c_.wavefront.finite_difference(
434 data.__mut_ceo__().0.d__cx,
435 data.__mut_ceo__().0.d__cy,
436 n_lenslet,
437 lenslet_size as f32,
438 );
439 }
440 }
441 }
442 pub fn reset(&mut self) {
444 unsafe {
445 self._c_.wavefront.reset();
446 self._c_.reset_rays();
447 }
448 }
449 pub fn reset_phase(&mut self) -> &mut Self {
451 unsafe {
452 self._c_.wavefront.reset_phase();
453 }
454 self
455 }
456 pub fn add<'a, T>(&mut self, phase: &'a [T]) -> &mut Self
458 where
459 &'a [T]: Into<Cu<Single>>,
460 {
461 unsafe {
462 let mut cu: Cu<Single> = phase.into();
463 self._c_.wavefront.add_phase(1.0, cu.as_mut_ptr());
464 }
465 self
466 }
467 pub fn axpy<'a, T>(&mut self, alpha: f32, phase: &'a [T]) -> &mut Self
468 where
469 &'a [T]: Into<Cu<Single>>,
470 {
471 unsafe {
472 let mut cu: Cu<Single> = phase.into();
473 self._c_.wavefront.add_phase(alpha, cu.as_mut_ptr());
474 }
475 self
476 }
477 pub fn sub<'a, T>(&mut self, phase: &'a [T]) -> &mut Self
479 where
480 &'a [T]: Into<Cu<Single>>,
481 {
482 unsafe {
483 let mut cu: Cu<Single> = phase.into();
484 self._c_.wavefront.add_phase(-1.0, cu.as_mut_ptr());
485 }
486 self
487 }
488 pub fn add_same<'a, T>(&mut self, phase: &'a [T]) -> &mut Self
490 where
491 &'a [T]: Into<Cu<Single>>,
492 {
493 unsafe {
494 let mut cu: Cu<Single> = phase.into();
495 self._c_.wavefront.add_same_phase(1.0, cu.as_mut_ptr());
496 }
497 self
498 }
499 pub fn phase(&self) -> &Vec<f32> {
501 unsafe {
502 dev2host(
503 self._phase.as_ptr() as *mut _,
504 self._c_.wavefront.phase,
505 self._c_.wavefront.N_PX,
506 );
507 }
508 &self._phase
509 }
510 pub fn phase_as_ptr(&mut self) -> Cu<Single> {
511 let mut phase: Cu<Single> = Cu::vector(self._c_.wavefront.N_PX as usize);
512 phase.from_ptr(self._c_.wavefront.phase);
513 phase
514 }
515 pub fn phase_as_slice<'a, T>(&'a mut self) -> &'a [T]
516 where
517 Cu<Single>: Into<&'a [T]>,
518 {
519 let mut phase: Cu<Single> = Cu::vector(self._c_.wavefront.N_PX as usize);
520 phase.from_ptr(self._c_.wavefront.phase);
521 phase.into()
522 }
523 pub fn amplitude(&mut self) -> Vec<f32> {
525 let n = self._c_.wavefront.N_PX;
526 let mut a = vec![0f32; n as usize];
527 unsafe {
528 dev2host(a.as_mut_ptr(), self._c_.wavefront.amplitude, n);
529 }
530 a
531 }
532 pub fn rays_coordinates(&mut self) -> Vec<f64> {
536 let n = 3 * self._c_.rays.N_RAY_TOTAL as usize;
537 let mut d_xyz = Cu::<Double>::vector(n);
538 unsafe {
540 self._c_.rays.get_coordinates(d_xyz.malloc().as_mut_ptr());
541 }
542 d_xyz.into()
543 }
544 pub fn fluxlet(&mut self, n_let: usize) -> Vec<f32> {
546 let m = (self.pupil_sampling as usize - 1) / n_let;
547 assert_eq!(m * n_let + 1, self.pupil_sampling as usize);
548 let n = self.pupil_sampling as usize;
549 let a = self.amplitude();
550 let mut f = vec![0f32; (n_let * n_let) as usize];
551 for i_let in 0..n_let {
552 let ui = (m * i_let) as usize;
553 for j_let in 0..n_let {
554 let uj = (m * j_let) as usize;
555 let mut s = 0f32;
556 for i in 0..m as usize + 1 {
557 for j in 0..m as usize + 1 {
558 let k = ui + i + n * (uj + j);
559 s += a[k];
560 }
561 }
562 f[i_let + n_let * j_let] = s;
563 }
564 }
565 f
566 }
567 pub fn masklet(&mut self, n_let: usize, flux_threshold: f32) -> Vec<i8> {
569 let f = self.fluxlet(n_let);
570 let f_max = f.iter().cloned().fold(-f32::INFINITY, f32::max);
571 let t = flux_threshold * f_max;
572 f.iter().map(|x| if *x >= t { 1i8 } else { 0i8 }).collect()
573 }
574 pub fn through<T: Propagation>(&mut self, system: &mut T) -> &mut Self {
576 system.propagate(self);
577 self
578 }
579 pub fn n_photon(&mut self) -> Vec<f32> {
581 self.magnitude
582 .clone()
583 .iter()
584 .map(|m| unsafe { self._c_.n_photon1(*m) })
585 .collect()
586 }
587 pub fn light_collecting_area(&self) -> f32 {
589 self._c_.rays.V.area
590 }
591 pub fn rays(&self) -> Rays {
593 Rays {
594 _c_: UnsafeCell::new(self._c_.rays),
595 }
596 }
597}
598impl Drop for Source {
599 fn drop(&mut self) {
601 unsafe {
602 self._c_.cleanup();
603 }
604 }
605}
606impl Default for Source {
607 fn default() -> Self {
608 Self::empty()
609 }
610}
611
612pub struct Rays {
614 _c_: UnsafeCell<bundle>,
615}
616impl Rays {
617 pub fn coordinates(&self) -> Vec<f64> {
621 let rays = unsafe { &mut *self._c_.get() };
622 let n = 3 * rays.N_RAY_TOTAL as usize;
623 let mut d_xyz = Cu::<Double>::vector(n);
624 unsafe {
625 rays.get_coordinates(d_xyz.malloc().as_mut_ptr());
626 }
627 d_xyz.into()
628 }
629 pub fn directions(&self) -> Vec<f64> {
633 let rays = unsafe { &mut *self._c_.get() };
634 let n = 3 * rays.N_RAY_TOTAL as usize;
635 let mut d_klm = Cu::<Double>::vector(n);
636 unsafe {
637 rays.get_directions(d_klm.malloc().as_mut_ptr());
638 }
639 d_klm.into()
640 }
641 pub fn opd(&self) -> Vec<f64> {
643 let rays = unsafe { &mut *self._c_.get() };
644 let n = rays.N_RAY_TOTAL as usize;
645 let mut d_opd = Cu::<Double>::vector(n);
646 unsafe {
648 rays.get_optical_path_difference(d_opd.malloc().as_mut_ptr());
649 }
650 d_opd.into()
651 }
652 pub fn vignetting(&self) -> Vec<bool> {
654 let rays = unsafe { &mut *self._c_.get() };
655 let n = rays.N_RAY_TOTAL as usize;
656 let mut d_v = Cu::<Double>::vector(n);
657 unsafe {
658 rays.get_vignetting(d_v.malloc().as_mut_ptr());
659 }
660 Vec::<f64>::from(d_v)
661 .into_iter()
662 .map(|v| v.abs() > 0f64)
663 .collect()
664 }
665 pub fn mask(&self) -> Mask {
667 let rays = unsafe { *self._c_.get() };
668 Mask {
669 _c_: UnsafeCell::new(rays.V),
670 }
671 }
672 pub fn segment(&self) -> Vec<i32> {
677 let rays = unsafe { &mut *self._c_.get() };
678 let n = rays.N_RAY_TOTAL as usize;
679 let mut d_s = Cu::<Int>::vector(n);
680 d_s.from_ptr(unsafe { &mut *self._c_.get() }.d__piston_mask);
681 Vec::<i32>::from(d_s)
682 }
683}
684