1use std::fs;
2use astroimsim_geometry::coordinate_system::CoordinateSystem;
3use astroimsim_geometry::grid2d::{Corners, GRID2D};
4use astroimsim_geometry::points::Point;
5use crate::psf::{DataFile, PSF, Load};
6
7
8pub struct PsfGrid {
9 data: Vec<(usize,PSF)>,
10 pub(crate) grid: GRID2D,
11 valid: bool
12}
13
14impl PsfGrid{
15 pub fn new(grid: GRID2D) -> PsfGrid{
16 PsfGrid{
17 data: vec![],
18 grid: grid,
19 valid:false,
20
21 }
22 }
23 pub fn load_data_frames(&mut self, directory_path:&str, center_fits_keys:(&str, &str), pixels:(usize, usize), size:(f64, f64)){
24 println!("Loading data frames into grid. This overwrites any data previously loaded");
25 let mut data = vec![];
26 let paths = fs::read_dir(directory_path).unwrap();
27 let mut counter = 0;
28 for path in paths {
29 println!("loading {:?}",path);
30 counter += 1;
31 let path = path.unwrap().path();
32 let (x_pixels,y_pixels) = pixels;
33 let data_file = DataFile{
34 description: counter.to_string(),
35 path,
36 x_pixels,
37 y_pixels,
38 };
39 let frame = PSF::load_file(data_file,
40 (Load::FromKey(center_fits_keys.0.to_string()), Load::FromKey(center_fits_keys.1.to_string())),
41 (Load::FromValue(size.0),Load::FromValue(size.0)));
42 let frame_index = frame.snap_to_grid(&self.grid);
43 data.push((frame_index,frame))
44 }
45 data.sort_by_key(|x|x.0);
46 self.data = data;
47 println!("Loaded {counter} files into a Grid Struct from {directory_path}");
48 assert_eq!(counter,self.grid.num_points,"Loaded the wrong number of PSF files");
49 }
50
51
52
53 pub fn validate(&mut self) -> bool {
54 if self.data.len() != self.grid.num_points{
56 println!("Validation failed: expected {:?} data frames, have {:?}",self.grid.num_points,self.data.len());
57 self.valid = false;
58 return false
59 };
60 let mut missing = Vec::new();
62 let mut counter = 0;
63 for grid_number in 0..self.grid.num_points{
64 counter += 1;
65 if !self.data.iter().any(|(index,_)| *index==grid_number) {
66 missing.push(grid_number)
67 }
68 };
69 if missing.len() > 0{
70 println!("Validation failed! Missing data frames for the following grid positions: {:?} ",missing);
71 self.valid = false;
72 false
73 }else{
74 self.data.sort_by_key(|x|x.0);
75 self.valid = true;
76 true
77 }}
78
79
80 pub fn interpolated_psf(&self, point:&Point) -> Vec<Vec<f32>>{
81 let ((Q12, Q22, Q21, Q11),(c11,c12,c21,c22),normalization) = self.interpolation_coefficients(point.clone());
86
87 let q11 = self.data[Q11].clone();
88 let q12 = self.data[Q12].clone();
89 let q21 = self.data[Q21].clone();
90 let q22 = self.data[Q22].clone();
91
92 assert_eq!(q11.0,Q11);
93 assert_eq!(q12.0,Q12);
94 assert_eq!(q21.0,Q21);
95 assert_eq!(q22.0,Q22);
96
97 let interpolated_data:Vec<f32> =
98 q11.1.data.into_iter().flatten().zip(
99 q12.1.data.into_iter().flatten().zip(
100 q21.1.data.into_iter().flatten().zip(
101 q22.1.data.into_iter().flatten()))).map(
102 |(q11,(q12,(q21,q22)))| {
103 (q11*c11 as f32 + q12*c12 as f32 + q21*c21 as f32 + q22*c22 as f32)/normalization as f32
104 }).collect();
105 PSF::repack_data(interpolated_data)
106
107
108
109 }
110
111
112
113 pub fn interpolation_coefficients(&self, point:Point) -> ((usize,usize,usize,usize),(f64,f64,f64,f64),f64){
114 if self.valid == false{
115 panic!("Must validate Grid before attempting to interpolate")
116 };
117
118 match self.grid.find_corners(point.clone()){
119 Corners::Four(Q12, Q22, Q21, Q11) => { let Q11point = self.grid.locate(Q11);
122 let Q22point = self.grid.locate(Q22);
123
124
125 let (x1,y1) = (Q11point.x,Q11point.y);
126 let (x2,y2) = (Q22point.x,Q22point.y);
127 let (x,y) = point.convert(&self.grid.coordinates).values();
128 let c11 = (x2-x)*(y2-y);
129 let c12 = (x2-x)*(y-y1);
130 let c21 = (x-x1)*(y2-y);
131 let c22 = (x-x1)*(y-y1);
132 println!("The coefficients are {:?}",(c11,c12,c21,c22));
133 let normalization = (x2-x1)*(y2-y1);
134 ((Q12, Q22, Q21, Q11),(c11,c12,c21,c22),normalization)
135
136 }
137 Corners::Two(Q1, Q2) => {((Q1, Q2, Q1, Q2),(1.0,1.0,1.0,1.0),1.0)} Corners::One(Q1) => {((Q1, Q1, Q1, Q1),(1.0,1.0,1.0,1.0),1.0)}
139 }
140
141 }
142
143 pub fn grid_psf(&self, index:usize)-> Vec<Vec<f32>>{
144 let (i, psf) = self.data[index].clone();
145 assert_eq!(index,i);
146 psf.data
147
148 }
149
150}