Skip to main content

astroimsim_data/
psf_grid.rs

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