Skip to main content

astroimsim_data/
psf_grid.rs

1use 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        //check to make sure there are the same number of data frames as there are grid points
60        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        //check to make sure that every grid point has a data frame
66        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        //println!("Converting from {:?}", point);
87
88        //  println!("To {:?}",(x,y));
89        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