astroimsim_data/
psf_grid.rs1use std::fs;
2use astroimsim_geometry::coordinate_system::CoordinateSystem;
3use astroimsim_geometry::grid2d::{Corners, PlotPoint, GRID2D};
4use astroimsim_geometry::points::Point;
5use egui::emath::interpolation_factor;
6use plotpy::{Curve, Plot};
7use uvex_fitrs::{Fits, Hdu};
8use crate::psf::{DataFile, PSF, Load};
9
10#[derive(Debug,Clone)]
11pub struct PsfGrid {
12 label:String,
13 grid: GRID2D,
14 data:Vec<(usize,PSF)>,
15 valid: bool,
16 directory_path:String,
17 center_fits_keys:(String, String),
18
19}
20
21impl PsfGrid{
22 pub fn new(label:String,grid: GRID2D,directory_path:String,center_fits_keys:(String, String)) -> PsfGrid{
23 PsfGrid{
24 label,
25 data: vec![],
26 grid: grid,
27 valid:false,
28 directory_path,
29 center_fits_keys,
30 }
31 }
32 pub fn load_data_frames(&mut self, x_num:usize,y_num:usize){
33
34
35
36 println!("Loading data frames into grid. This overwrites any data previously loaded");
37 let mut data = vec![];
38 let paths = fs::read_dir(self.directory_path.clone()).unwrap();
39 let mut counter = 0;
40 for path in paths {
41 println!("loading {:?}",path);
42 counter += 1;
43 let path = path.unwrap().path();
44
45 let frame = PSF::load_file(
46 path,
47 (Load::FromKey(self.center_fits_keys.0.to_string()), Load::FromKey(self.center_fits_keys.1.to_string())),
48 (Load::FromValue(self.grid.x_size),Load::FromValue(self.grid.y_size)),
49 x_num,y_num,
50 );
51
52 let frame_index =frame.snap_to_grid(&self.grid);
53 data.push((frame_index,frame))
54 }
55 data.sort_by_key(|x|x.0);
56 self.data = data;
57 println!("Loaded {counter} files into a Grid Struct from {:?}",self.directory_path);
58 assert_eq!(counter,self.grid.num_points,"Loaded the wrong number of PSF files");
59
60
61 }
62
63
64
65 pub fn validate(&mut self) -> bool {
66 if self.data.len() != self.grid.num_points{
68 println!("Validation failed: expected {:?} data frames, have {:?}",self.grid.num_points,self.data.len());
69 self.valid = false;
70 return false
71 };
72 let mut missing = Vec::new();
74 let mut counter = 0;
75 for grid_number in 0..self.grid.num_points{
76 counter += 1;
77 if !self.data.iter().any(|(index,_)| *index==grid_number) {
78 missing.push(grid_number)
79 }
80 };
81 if missing.len() > 0{
82 println!("Validation failed! Missing data frames for the following grid positions: {:?} ",missing);
83 self.valid = false;
84 false
85 }else{
86 self.data.sort_by_key(|x|x.0);
87 self.valid = true;
88 true
89 }}
90
91
92 pub fn interpolated_psf(&self, point:&Point) -> Vec<Vec<f64>>{
93 let interpolation_data = self.grid.projected_interpolation_coefficients(&point.clone());
97
98 match interpolation_data.corners{
99 Corners::Four(Q12, Q22, Q21, Q11) => {
100 let q11 = self.data[Q11].clone();
101 let q12 = self.data[Q12].clone();
102 let q21 = self.data[Q21].clone();
103 let q22 = self.data[Q22].clone();
104
105 let c11 = interpolation_data.coefficients[0];
106 let c12 = interpolation_data.coefficients[1];
107 let c21 = interpolation_data.coefficients[2];
108 let c22 = interpolation_data.coefficients[3];
109
110 assert_eq!(q11.0,Q11);
111 assert_eq!(q12.0,Q12);
112 assert_eq!(q21.0,Q21);
113 assert_eq!(q22.0,Q22);
114
115 let interpolated_data:Vec<f64> =
116 q11.1.data.into_iter().flatten().zip(
117 q12.1.data.into_iter().flatten().zip(
118 q21.1.data.into_iter().flatten().zip(
119 q22.1.data.into_iter().flatten()))).map(
120 |(q11,(q12,(q21,q22)))| {
121 (q11*c11 + q12*c12 + q21*c21 + q22*c22 )/interpolation_data.normalization
122 }).collect();
123 PSF::repack_data(interpolated_data)
124
125 }
126 Corners::Two(Q1, Q2) => {
127 let q1 = self.data[Q1].clone();
128 let q2 = self.data[Q2].clone();
129
130 let c1 = interpolation_data.coefficients[0];
131 let c2 = interpolation_data.coefficients[1];
132
133
134 assert_eq!(q1.0,Q1);
135 assert_eq!(q2.0,Q2);
136
137
138 let interpolated_data:Vec<f64> =
139 q1.1.data.into_iter().flatten().zip(
140 q2.1.data.into_iter().flatten()).map(
141 |(q1,q2)| {
142 (q1*c1 + q2*c2 )/interpolation_data.normalization
143 }).collect();
144 PSF::repack_data(interpolated_data)
145
146 }
147 Corners::One(Q) => {
148 self.data[Q].clone().1.data
149 }
150 }
151 }
152
153
154
155
156 pub fn grid_psf(&self, index:usize)-> Vec<Vec<f64>>{
157 let (i, psf) = self.data[index].clone();
158 assert_eq!(index,i);
159 psf.data
160
161 }
162
163
164 pub fn gaussian_blur(&self, blurred_directory:String, std_in_pixels:f64){
165 let mut new_grid = PsfGrid::new("blurred grid".to_string(), self.grid.clone(), blurred_directory.clone(), self.center_fits_keys.clone());
166 let kernel = convolve2d::kernel::gaussian(std_in_pixels.ceil() as usize*10,std_in_pixels);
167 let (x,y,vec) = kernel.into_parts();
168 let square_kernel = vec.chunks_exact(x).map(|x|x.to_vec()).collect();
169 for (i,psf) in &self.data{
170 let output = psf.convolve(&square_kernel);
171 let path = format!("{blurred_directory}/{i}_blurred_{std_in_pixels}.fits");
172 let new_psf = PSF{
173 path: path.parse().unwrap(),
174 data: output.chunks_exact(psf.x_pixels).map(|x|x.to_vec()).collect(),
175 x_pixels: psf.x_pixels,
176 y_pixels: psf.y_pixels,
177 center: psf.center.clone(),
178 size: psf.size.clone()
179 };
180 new_grid.data.push((*i,new_psf.clone()));
181 new_psf.write_file(path.as_str(), self.center_fits_keys.clone());
182
183 }
184
185 }
186
187}
188
189
190