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 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        //check to make sure there are the same number of data frames as there are grid points
55        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        //check to make sure that every grid point has a data frame
61        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        //println!("Converting from {:?}", point);
82
83        //  println!("To {:?}",(x,y));
84
85        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) => { // using the wikipedia convention https://en.wikipedia.org/wiki/Bilinear_interpolation
120
121                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)} //TODO !!!!
138            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}