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