Merge branch 'main' of https://github.com/eta077/astro-rs

eta077 · Jun 1, 2022 · d65c286 · d65c286
2 parents ffc2681 + b83e98a
commit d65c286
Show file tree

Hide file tree

Showing 5 changed files with 114 additions and 24 deletions.
diff --git a/src/fits/hdu_types.rs b/src/fits/hdu_types.rs
@@ -210,7 +210,7 @@ pub mod binary_table_hdu {
             }
         }
         if let Some(tform) = tform {
-            return Some(*tform.get_values(hdu.get_data(), column_start, row_len, num_rows));
+            return Some(*tform.create_column(hdu.data_raw(), column_start, row_len, num_rows));
         }
         None
     }
@@ -254,7 +254,7 @@ pub mod binary_table_hdu {
             }
         }
         if let Some(tform) = tform {
-            return Some(*tform.get_values(hdu.get_data(), column_start, row_len, num_rows));
+            return Some(*tform.create_column(hdu.data_raw(), column_start, row_len, num_rows));
         }
         None
     }

diff --git a/src/fits/header_value/tform.rs b/src/fits/header_value/tform.rs
@@ -75,8 +75,8 @@ impl TForm {
         self.r * type_bytes
     }
 
-    /// Gets the function to transform raw bytes to the given type.
-    pub fn get_values<T>(
+    /// Creates a column of values from the given data.
+    pub fn create_column<T>(
         &self,
         data: &[u8],
         column_start: usize,
@@ -177,17 +177,6 @@ impl TForm {
 }
 
 mod tform_macros {
-    /// Creates a box of the given value and casts it to an implicit return type.
-    #[macro_export]
-    macro_rules! return_box {
-        ($result: expr) => {{
-            let b = Box::new($result);
-            let ptr = Box::into_raw(b);
-            let new_ptr = ptr.cast();
-            Box::from_raw(new_ptr)
-        }};
-    }
-
     /// Creates a boxed vector deserialized with the given function, or a default function if one is not given.
     #[macro_export]
     macro_rules! deserialize_column {

diff --git a/src/fits/mod.rs b/src/fits/mod.rs
@@ -308,12 +308,9 @@ impl<R: Read> HduList<R> {
                 .and_then(|card| card.get_value::<Bitpix>().ok())
             {
                 let mut data_len = 1;
+                let mut naxisx_keyword = FitsHeaderKeyword::from(NAXIS_KEYWORD);
                 for x in 1..=naxis {
-                    let mut naxisx_keyword = NAXIS_KEYWORD;
-                    let x_bytes = x.to_string().into_bytes();
-                    for (index, i) in x_bytes.iter().enumerate() {
-                        naxisx_keyword[index + 5] = *i;
-                    }
+                    naxisx_keyword.append_number(x);
 
                     let naxisx = *header
                         .get_card(naxisx_keyword)
@@ -367,16 +364,94 @@ impl Hdu {
     /// ```
     pub fn get_name(&mut self) -> String {
         self.header
-            .cards
-            .iter_mut()
-            .find(|card| *card.keyword() == EXTNAME_KEYWORD)
+            .get_card(EXTNAME_KEYWORD)
             .and_then(|card| card.get_value::<String>().ok())
             .map(|name| name.trim().to_owned())
             .unwrap_or_default()
     }
 
     /// Gets the data section of the HDU.
-    pub fn get_data(&mut self) -> &Vec<u8> {
+    pub fn data_raw(&self) -> &Vec<u8> {
         &self.data_raw
     }
+
+    /// Creates a Vec containing the data section of the HDU as the type defined by the BITPIX keyword.
+    pub fn get_data<T>(&mut self) -> Vec<T> {
+        *self
+            .header
+            .get_card(BITPIX_KEYWORD)
+            .and_then(|card| card.get_value::<Bitpix>().ok())
+            .map(|bitpix| unsafe {
+                match *bitpix {
+                    Bitpix::U8 => crate::return_box!(self.data_raw.to_owned()),
+                    Bitpix::I16 => {
+                        let mut data = Vec::with_capacity(self.data_raw.len() / 2);
+                        for chunk in self.data_raw.chunks_exact(2) {
+                            data.push(i16::from_be_bytes(chunk.try_into().unwrap()));
+                        }
+                        crate::return_box!(data)
+                    }
+                    Bitpix::I32 => {
+                        let mut data = Vec::with_capacity(self.data_raw.len() / 4);
+                        for chunk in self.data_raw.chunks_exact(4) {
+                            data.push(i32::from_be_bytes(chunk.try_into().unwrap()));
+                        }
+                        crate::return_box!(data)
+                    }
+                    Bitpix::F32 => {
+                        let mut data = Vec::with_capacity(self.data_raw.len() / 4);
+                        for chunk in self.data_raw.chunks_exact(4) {
+                            data.push(f32::from_be_bytes(chunk.try_into().unwrap()));
+                        }
+                        crate::return_box!(data)
+                    }
+                    Bitpix::F64 => {
+                        let mut data = Vec::with_capacity(self.data_raw.len() / 8);
+                        for chunk in self.data_raw.chunks_exact(8) {
+                            data.push(f64::from_be_bytes(chunk.try_into().unwrap()));
+                        }
+                        crate::return_box!(data)
+                    }
+                }
+            })
+            .unwrap_or_default()
+    }
+
+    /// Creates a Vec containing the dimensions of the data section of the HDU as defined by the NAXIS keywords.
+    pub fn get_dimensions(&mut self) -> Vec<usize> {
+        let naxis = *self
+            .header
+            .get_card(NAXIS_KEYWORD)
+            .and_then(|card| card.get_value::<u16>().ok())
+            .unwrap_or_default();
+        if naxis == 0 {
+            return Vec::new();
+        }
+        let mut result = Vec::with_capacity(naxis as usize);
+        let mut naxisx_keyword = FitsHeaderKeyword::from(NAXIS_KEYWORD);
+        for x in 1..=naxis {
+            naxisx_keyword.append_number(x);
+
+            let naxisx = *self
+                .header
+                .get_card(naxisx_keyword)
+                .and_then(|card| card.get_value::<u32>().ok())
+                .unwrap_or_default() as usize;
+            result.push(naxisx);
+        }
+        result
+    }
+}
+
+pub(crate) mod hdu_macros {
+    /// Creates a box of the given value and casts it to an implicit return type.
+    #[macro_export]
+    macro_rules! return_box {
+        ($result: expr) => {{
+            let b = Box::new($result);
+            let ptr = Box::into_raw(b);
+            let new_ptr = ptr.cast();
+            Box::from_raw(new_ptr)
+        }};
+    }
 }
diff --git a/src/lib.rs b/src/lib.rs
@@ -1,4 +1,5 @@
 #![deny(missing_docs)]
+#![deny(clippy::all)]
 //! This library provides utilities to interact with astronomical data.
 
 pub mod coordinates;

diff --git a/tests/fits.rs b/tests/fits.rs
@@ -64,3 +64,28 @@ fn test_hdu_table_column() -> Result<(), Box<dyn Error>> {
 
     Ok(())
 }
+
+#[test]
+fn test_hdu_image() -> Result<(), Box<dyn Error>> {
+    let fits_file = File::open("assets/eagle_nebula/502nmos.fits")?;
+    let fits_file_reader = BufReader::new(fits_file);
+    let mut hdu_list = HduList::new(fits_file_reader);
+    let primary_hdu = hdu_list.first_mut().unwrap();
+
+    let dimensions = primary_hdu.get_dimensions();
+    assert_eq!(dimensions, vec![1600, 1600]);
+
+    let data = primary_hdu.get_data::<f32>();
+    let mut data_min = 0.0;
+    let mut data_max = 0.0;
+    for value in &data {
+        if *value > data_max {
+            data_max = *value;
+        } else if *value < data_min {
+            data_min = *value;
+        }
+    }
+    assert_eq!(data_max, 2925.8718);
+    assert_eq!(data_min, -12.439324);
+    Ok(())
+}