pyref_core/
loader.rs

1use astrors_fork::fits;
2use astrors_fork::io::hdulist::HDU;
3
4use polars::{lazy::prelude::*, prelude::*}; // Add the import statement for PolarsError
5use rayon::prelude::*;
6use std::fs;
7
8use crate::enums::HeaderValue;
9use crate::errors::FitsLoaderError;
10use crate::io::{add_calculated_domains, process_file_name, process_image, process_metadata};
11// Enum representing different types of experiments.
12
13// Polars Helper Function
14
15// workhorse functions for loading and processing CCD data.
16pub fn read_fits(
17    file_path: std::path::PathBuf,
18    header_items: &Vec<HeaderValue>,
19) -> Result<DataFrame, FitsLoaderError> {
20    if file_path.extension().and_then(|ext| ext.to_str()) != Some("fits") {
21        return Err(FitsLoaderError::NoData);
22    }
23
24    let hdul = fits::fromfile(
25        file_path
26            .to_str()
27            .ok_or_else(|| FitsLoaderError::InvalidFileName("Invalid path".into()))?,
28    )?;
29
30    let meta = match hdul.hdus.get(0) {
31        Some(HDU::Primary(hdu)) => process_metadata(hdu, header_items)?,
32        _ => return Err(FitsLoaderError::NoData),
33    };
34
35    let img_data = match hdul.hdus.get(2) {
36        Some(HDU::Image(hdu)) => process_image(hdu)?,
37        _ => return Err(FitsLoaderError::NoData),
38    };
39
40    let names = process_file_name(file_path);
41
42    let mut columns = meta;
43    columns.extend(img_data);
44    columns.extend(names);
45
46    DataFrame::new(columns).map_err(FitsLoaderError::PolarsError)
47}
48
49pub fn read_experiment(
50    dir: &str,
51    header_items: &Vec<HeaderValue>,
52) -> Result<DataFrame, FitsLoaderError> {
53    let dir_path = std::path::PathBuf::from(dir);
54
55    if !dir_path.exists() {
56        return Err(FitsLoaderError::NoData);
57    }
58
59    let entries: Vec<_> = fs::read_dir(dir)
60        .map_err(FitsLoaderError::IoError)?
61        .par_bridge()
62        .filter_map(|entry| entry.ok())
63        .filter(|entry| entry.path().extension().and_then(|ext| ext.to_str()) == Some("fits"))
64        .collect();
65
66    let dataframes: Result<Vec<DataFrame>, FitsLoaderError> = entries
67        .par_iter()
68        .map(|entry| read_fits(entry.path(), &header_items))
69        .collect();
70
71    let combined_df = dataframes?
72        .into_par_iter()
73        .reduce_with(|acc, df| acc.vstack(&df).unwrap_or(DataFrame::empty()))
74        .ok_or(FitsLoaderError::NoData)?;
75
76    Ok(add_calculated_domains(combined_df.lazy()))
77}
78
79pub fn _load() {
80    // let test_path = "/home/hduva/projects/pyref/test/stack/";
81
82    // let data = read_experiment(test_path.into(), &ExperimentType::Xrr.get_keys()).unwrap();
83    // println!("{:?}", data);
84}
pyref_core/loader.rs

pyref_core/
loader.rs
