1use std::iter::Flatten;
2use std::path::PathBuf;
3use std::slice::Iter;
4use serde::{Deserialize, Serialize};
5use uvex_fitrs::{Fits, FitsData, FitsDataArray, Hdu, HeaderValue};
6use astroimsim_geometry::coordinate_system::{CoordinateSystem, Coordinates};
7use astroimsim_geometry::grid2d::GRID2D;
8use astroimsim_geometry::points::Point;
9use astroimsim_geometry::grid2d::InterpolationData;
10use ndarray::Array2;
11use ndarray_conv::{get_fft_processor, ConvExt, ConvFFTExt, ConvMode, PaddingMode};
12use std::time::{Duration, Instant};
13
14use convolutions_rs::convolutions::*;
15use ndarray::*;
16use convolutions_rs::Padding;
17use convolve2d::{convolve2d, DynamicMatrix, Matrix};
18
19#[derive(Debug,Clone,Serialize)]
20pub struct PSF {
21 pub path: PathBuf,
22 pub data: Vec<Vec<f64>>,
23 pub x_pixels: usize,
24 pub y_pixels: usize,
25 pub center: Point,
26 pub size: (f64,f64),
27}
28
29pub struct DataFile {
30 pub description: String,
31 pub path: PathBuf,
32 pub x_pixels: usize,
33 pub y_pixels: usize,
34}
35
36
37pub enum Load{
38 FromKey(String),
39 FromValue(f64)
40}
41
42
43pub enum FITSType{
44 thirtytwo,
45 sixtyfour,
46}
47
48impl PSF {
49 pub fn load_file(file: PathBuf, center:(Load, Load), size:(Load, Load), x_num:usize, y_num:usize) -> PSF {
50 println!("Loading {:?} into a DataFrame ",file);
51 let fits = Fits::open(file.clone()).expect("Failed to open FITS file");
52 let primary_hdu= fits.iter().next().expect("Couldn't find primary HDU");
53 let (mut data,shape) = match primary_hdu.read_data() {
54 FitsData::FloatingPoint32(FitsDataArray { shape, data }) => (data.iter().map(|x| *x as f64).collect(),shape),
55 FitsData::FloatingPoint64(FitsDataArray { shape, data }) => (data,shape),
56 _ => panic!("Could not unpack PSF data")
57 }; let normalization:f64 = data.iter().sum();
60 let data:Vec<f64> = data.iter().map(|x|x/normalization).collect();
61 println!("PSF has been normalized to {:?}",data.iter().sum::<f64>());
62
63 assert_eq!(shape[0], x_num,"Diva down! Tried to load a file with data of the wrong x size"); assert_eq!(shape[1], y_num,"Diva down! Tried to load a file with data of the wrong y size");
65 let center_x:f64 = PSF::load(center.0, &primary_hdu);
68 let center_y:f64 = PSF::load(center.1, &primary_hdu);
69 println!("loaded center");
70 let size_x:f64 = PSF::load(size.0, &primary_hdu);
71 let size_y:f64 = PSF::load(size.1, &primary_hdu);
72 println!("loaded size");
73 let data = data.chunks(x_num).map(|i| i.to_vec()).collect();
74 PSF {
76 path: file,
77 data,
78 x_pixels: x_num,
79 y_pixels: y_num,
80 center: Point::new(center_x,center_y,Coordinates::ABSOLUTE),
81 size: (size_x,size_y),
82 }
83 }
84
85
86
87 pub fn write_file(&self, path:&str,center_keys:(String,String)){
88
89 let data:Vec<f64> = self.data.clone().iter().map(|x|x.to_owned()).flatten().collect();
90 let data= data.iter().map(|x| *x).collect();
91 println!("Trying to write to {path}, {:?}",data);
92 let mut primary_hdu = Hdu::new(&[64, 64], data);
93 println!("Done making hdu");
94 let (x,y) = self.center.to_absolute().values();
95 println!("x,y,{x} {y}");
96 primary_hdu.insert(center_keys.0,x.to_string().as_str() );
97 primary_hdu.insert(center_keys.1,y.to_string().as_str() );
98 println!("printed keys");
99 Fits::create(path, primary_hdu).expect("Failed to create");
102 println!("done??")
103
104
105 }
106
107 pub fn snap_to_grid(&self, grid: &GRID2D) -> usize{
108 let index = grid.snap(self.center.clone());
109 index
110 }
111
112
113 pub fn repack_data(flat_data: Vec<f64>) -> Vec<Vec<f64>>{
114 flat_data.chunks(64).map(|i| i.to_vec()).collect()
115 }
116
117 fn load(thingy:Load, header:&Hdu)-> f64{
118 match thingy{
119 Load::FromKey(Key) => {
120 match header.value(&Key).expect("failed to get key") {
121 HeaderValue::CharacterString(string) => { string.parse().expect("FAILED TO PARSE string") }
122 HeaderValue::RealFloatingNumber(i) => {*i}
123 _ => {panic!("failed to parse header value")}
124
125 }}
126 Load::FromValue(value) => value
127 }
128 }
129
130 pub fn convolve(&self, kernel:&Vec<Vec<f64>>) -> Vec<f64>{
131 let kernel:Vec<Vec<f64>> = kernel
132 .iter()
133 .rev()
134 .map(|v|v.iter().map(|x|*x).rev().collect()).collect();
135 let flat_data:Vec<f64> = self.data.iter().flatten().map(|x|*x).collect();
136 let flat_kernel:Vec<f64> = kernel.iter().flatten().map(|x|*x).collect();
137 let data = DynamicMatrix::new(self.x_pixels, self.y_pixels, flat_data).unwrap();
138 let kernel = DynamicMatrix::new(kernel[0].len(), kernel.len(), flat_kernel).unwrap();
139 let output = convolve2d(&data, &kernel);
140 output.get_data().to_vec()
141 }
142
143
144
145
146
147
148}
149
150