pyref_core/
io.rs

1use astrors_fork::io::hdus::{
2    image::{imagehdu::ImageHDU, ImageData},
3    primaryhdu::PrimaryHDU,
4};
5use ndarray::{ArrayBase, Axis, Dim, IxDynImpl, OwnedRepr};
6use polars::export::chrono::NaiveDateTime;
7use polars::prelude::*;
8use std::ops::Mul;
9
10use crate::enums::HeaderValue;
11use crate::errors::FitsLoaderError;
12
13// Find the theta offset between theta and the ccd theta or 2theta
14pub fn theta_offset(theta: f64, ccd_theta: f64) -> f64 {
15    // 2theta = -ccd_theta / 2 - theta rounded to 3 decimal places
16    let ccd_theta = ccd_theta / 2.0;
17    ccd_theta - theta
18}
19
20pub fn q(lam: f64, theta: f64, angle_offset: f64) -> f64 {
21    let theta = theta - angle_offset;
22    match 4.0 * std::f64::consts::PI * theta.to_radians().sin() / lam {
23        q if q < 0.0 => 0.0,
24        q => q,
25    }
26}
27
28pub fn col_from_array(
29    name: PlSmallStr,
30    array: ArrayBase<OwnedRepr<i64>, Dim<IxDynImpl>>,
31) -> Result<Column, PolarsError> {
32    let size0 = array.len_of(Axis(0));
33    let size1 = array.len_of(Axis(1));
34
35    let mut s = Column::new_empty(
36        name,
37        &DataType::List(Box::new(DataType::List(Box::new(DataType::Int64)))),
38    );
39
40    let mut chunked_builder = ListPrimitiveChunkedBuilder::<Int64Type>::new(
41        PlSmallStr::EMPTY,
42        array.len_of(Axis(0)),
43        array.len_of(Axis(1)) * array.len_of(Axis(0)),
44        DataType::Int64,
45    );
46    for row in array.axis_iter(Axis(0)) {
47        match row.as_slice() {
48            Some(row) => chunked_builder.append_slice(row),
49            None => chunked_builder.append_slice(&row.to_owned().into_raw_vec()),
50        }
51    }
52    let new_series = chunked_builder
53        .finish()
54        .into_series()
55        .implode()
56        .unwrap()
57        .into_column();
58    let _ = s.extend(&new_series);
59    let s = s.cast(&DataType::Array(
60        Box::new(DataType::Array(Box::new(DataType::Int32), size1)),
61        size0,
62    ));
63    s
64}
65// ================== CCD Raw Data Processing ============
66pub fn add_calculated_domains(lzf: LazyFrame) -> DataFrame {
67    let h = physical_constants::PLANCK_CONSTANT_IN_EV_PER_HZ;
68    let c = physical_constants::SPEED_OF_LIGHT_IN_VACUUM * 1e10;
69    // Calculate lambda and q values in angstrom
70    let lz = lzf
71        .sort(["Sample Name"], Default::default())
72        .sort(["Scan ID"], Default::default())
73        .sort(["Frame Number"], Default::default())
74        .with_columns(&[
75            col("EXPOSURE [s]").round(3).alias("EXPOSURE [s]"),
76            col("Higher Order Suppressor [mm]")
77                .round(2)
78                .alias("Higher Order Suppressor [mm]"),
79            col("Horizontal Exit Slit Size [um]")
80                .round(1)
81                .alias("Horizontal Exit Slit Size [um]"),
82            col("Beamline Energy [eV]")
83                .round(1)
84                .alias("Beamline Energy [eV]"),
85        ])
86        .with_column(
87            col("Beamline Energy [eV]")
88                .pow(-1)
89                .mul(lit(h * c))
90                .alias("Lambda [Å]"),
91        );
92
93    let angle_offset = lz
94        .clone()
95        .filter(col("Sample Theta [deg]").eq(0.0))
96        .last()
97        .select(&[col("CCD Theta [deg]"), col("Sample Theta [deg]")])
98        .with_column(
99            as_struct(vec![col("Sample Theta [deg]"), col("CCD Theta [deg]")])
100                .map(
101                    move |s| {
102                        let struc = s.struct_()?;
103                        let th_series = struc.field_by_name("Sample Theta [deg]")?;
104                        let theta = th_series.f64()?;
105                        let ccd_th_series = struc.field_by_name("CCD Theta [deg]")?;
106                        let ccd_theta = ccd_th_series.f64()?;
107                        let out: Float64Chunked = theta
108                            .into_iter()
109                            .zip(ccd_theta.iter())
110                            .map(|(theta, ccd_theta)| match (theta, ccd_theta) {
111                                (Some(theta), Some(ccd_theta)) => {
112                                    Some(theta_offset(theta, ccd_theta))
113                                }
114                                _ => Some(0.0),
115                            })
116                            .collect();
117                        Ok(Some(out.into_column()))
118                    },
119                    GetOutput::from_type(DataType::Float64),
120                )
121                .alias("Theta Offset [deg]"),
122        )
123        .select(&[col("Theta Offset [deg]")])
124        .collect()
125        .unwrap()
126        .get(0)
127        .unwrap()[0]
128        .try_extract::<f64>()
129        .unwrap();
130    // get the row cor
131    let lz = lz
132        .with_column(lit(angle_offset).alias("Theta Offset [deg]"))
133        .with_column(
134            as_struct(vec![col("Sample Theta [deg]"), col("Lambda [Å]")])
135                .map(
136                    move |s| {
137                        let struc = s.struct_()?;
138                        let th_series = struc.field_by_name("Sample Theta [deg]")?;
139                        let theta = th_series.f64()?;
140                        let lam_series = struc.field_by_name("Lambda [Å]")?;
141                        let lam = lam_series.f64()?;
142
143                        let out: Float64Chunked = theta
144                            .into_iter()
145                            .zip(lam.iter())
146                            .map(|(theta, lam)| match (theta, lam) {
147                                (Some(theta), Some(lam)) => Some(q(lam, theta, angle_offset)),
148                                _ => None,
149                            })
150                            .collect();
151
152                        Ok(Some(out.into_column()))
153                    },
154                    GetOutput::from_type(DataType::Float64),
155                )
156                .alias("Q [Å⁻¹]"),
157        );
158    lz.collect().unwrap()
159}
160
161//
162// ================== CCD Data Loader ==================
163pub fn process_image(img: &ImageHDU) -> Result<Vec<Column>, FitsLoaderError> {
164    let bzero = img
165        .header
166        .get_card("BZERO")
167        .ok_or_else(|| FitsLoaderError::MissingHeaderKey("BZERO".into()))?
168        .value
169        .as_int()
170        .ok_or_else(|| FitsLoaderError::FitsError("BZERO not an integer".into()))?;
171
172    match &img.data {
173        ImageData::I16(image) => {
174            let data = image.map(|&x| i64::from(x as i64 + bzero));
175            Ok(vec![col_from_array("Raw".into(), data.clone()).unwrap()])
176        }
177        _ => Err(FitsLoaderError::UnsupportedImageData),
178    }
179}
180
181pub fn process_metadata(
182    hdu: &PrimaryHDU,
183    keys: &Vec<HeaderValue>,
184) -> Result<Vec<Column>, FitsLoaderError> {
185    if keys.is_empty() {
186        Ok(hdu
187            .header
188            .iter()
189            .map(|card| {
190                let name = card.keyword.as_str();
191                let value = card.value.as_float().unwrap_or(0.0);
192                Column::new(name.into(), &[value])
193            })
194            .collect())
195    } else {
196        keys.iter()
197            .map(|key| {
198                let val = hdu
199                    .header
200                    .get_card(key.hdu())
201                    .ok_or_else(|| FitsLoaderError::MissingHeaderKey(key.hdu().to_string()))?
202                    .value
203                    .as_float()
204                    .ok_or_else(|| {
205                        FitsLoaderError::Other(format!("Value for {} is not a float", key.hdu()))
206                    })?;
207                Ok(Column::new(key.name().into(), &[val]))
208            })
209            .collect()
210    }
211}
212
213pub fn process_file_name(path: std::path::PathBuf) -> Vec<Column> {
214    let ts = path
215        .metadata()
216        .unwrap()
217        .modified()
218        .unwrap()
219        .duration_since(std::time::UNIX_EPOCH)
220        .unwrap()
221        .as_secs();
222
223    let ts = NaiveDateTime::from_timestamp(ts as i64, 0);
224
225    let file_name = path.file_stem().unwrap().to_str().unwrap_or("");
226
227    let filename_segments = file_name.split('-');
228
229    let frame = filename_segments
230        .clone()
231        .last()
232        .unwrap()
233        .parse::<u32>()
234        .unwrap_or(0);
235
236    let remaining = filename_segments
237        .clone()
238        .take(filename_segments.clone().count() - 1)
239        .collect::<String>();
240
241    if remaining.is_empty() {
242        return vec![
243            Column::new("Scan ID".into(), vec![0]),
244            Column::new("Sample Name".into(), vec![""]),
245            Column::new("Frame Number".into(), vec![0]),
246            Column::new("Timestamp".into(), vec![ts]),
247        ];
248    }
249
250    let scan_id = remaining
251        .chars()
252        .rev()
253        .take(5)
254        .collect::<String>()
255        .chars()
256        .rev()
257        .collect::<String>();
258
259    let sample_name = remaining
260        .chars()
261        .take(remaining.len() - 5)
262        .collect::<String>()
263        .trim_end_matches(|c| c == '-' || c == '_')
264        .to_string();
265
266    vec![
267        Column::new("File Name".into(), vec![remaining]),
268        Column::new("Scan ID".into(), vec![scan_id]),
269        Column::new("Sample Name".into(), vec![sample_name]),
270        Column::new("Frame Number".into(), vec![frame]),
271        Column::new("Timestamp".into(), vec![ts]),
272    ]
273}
pyref_core/io.rs

pyref_core/
io.rs
