Add support for saving and loading of States from Parquet files

CHANGELOG.md

@@ -6,6 +6,13 @@ The format is based on [Keep a Changelog](https://keepachangelog.com/en/1.0.0/),
 and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0.html).
 
 
+## [Unreleased]
+
+### Added
+
+- Added support for saving and loading `SimultaneousStates` as Parquet files.
+
+
 ## [v1.0.3]
 
 ### Added

Cargo.toml

@@ -8,7 +8,7 @@ license = "BSD-3-Clause"
 keywords = ["physics", "simulation", "astronomy"]
 
 [dependencies]
-kete_core = { version = "*", path = "src/kete_core", features=["pyo3"]}
+kete_core = { version = "*", path = "src/kete_core", features=["pyo3", "polars"]}
 pyo3 = { version =  "^0.22.1", features = ["extension-module"] }
 serde = { version = "^1.0.203", features = ["derive"] }
 nalgebra = {version = "^0.33.0"}

src/kete/rust/lib.rs

@@ -24,6 +24,7 @@
 )]
 
 use pyo3::prelude::*;
+use state::PyState;
 
 pub mod covariance;
 pub mod elements;
@@ -56,7 +57,7 @@ pub mod vector;
 #[pymodule]
 fn _core(_py: Python, m: &Bound<'_, PyModule>) -> PyResult<()> {
     m.add_class::<frame::PyFrames>()?;
-    m.add_class::<state::PyState>()?;
+    m.add_class::<PyState>()?;
     m.add_class::<vector::Vector>()?;
     m.add_class::<elements::PyCometElements>()?;
     m.add_class::<simult_states::PySimultaneousStates>()?;

src/kete/rust/simult_states.rs

@@ -119,6 +119,30 @@ impl PySimultaneousStates {
         Ok(())
     }
 
+    /// Save states as a parquet file.
+    pub fn save_parquet(&self, filename: String) -> PyResult<()> {
+        if self.0.fov.is_some() {
+            Err(Error::IOError(
+                "Cannot save a SimultaneousStates object which has a FOV as parquet. \
+                Parquet can only support a basic table format and saving metadata such \
+                as a field of view is not feasible. Consider using the binary saving \
+                method `SimultaneousStates.save`."
+                    .into(),
+            ))?;
+        }
+        kete_core::io::parquet::write_states_parquet(&self.0.states, &filename)?;
+        Ok(())
+    }
+
+    /// Load states from a parquet file.
+    #[staticmethod]
+    pub fn load_parquet(filename: String) -> PyResult<Self> {
+        let states = kete_core::io::parquet::read_states_parquet(&filename)?;
+        Ok(PySimultaneousStates(Box::new(
+            SimultaneousStates::new_exact(states, None)?,
+        )))
+    }
+
     /// Length of states
     pub fn __len__(&self) -> usize {
         self.0.states.len()

src/kete/rust/spice/mod.rs

@@ -15,7 +15,7 @@ use pyo3::pyfunction;
 #[pyo3(name = "observatory_codes")]
 pub fn obs_codes() -> Vec<(f64, f64, f64, String, String)> {
     let mut codes = Vec::new();
-    for row in kete_core::io::obs_codes::OBS_CODES.iter() {
+    for row in kete_core::spice::OBS_CODES.iter() {
         codes.push((
             row.lat,
             row.lon,

src/kete_core/Cargo.toml

@@ -18,6 +18,7 @@ itertools = "^0.13.0"
 kdtree = "^0.7.0"
 lazy_static = "^1.5.0"
 nalgebra = {version = "^0.33.0"}
+polars = {version = "0.43.1", optional=true, features=["parquet", "polars-io"]}
 pathfinding = "^4.10.0"
 pyo3 = { version =  "^0.22.1", optional=true}
 rayon = "^1.10.0"

src/kete_core/src/frames/definitions.rs

@@ -47,12 +47,24 @@ pub enum Frame {
 impl From<Frame> for i32 {
     fn from(value: Frame) -> Self {
         match value {
-            Frame::Unknown(_) => 0,
             Frame::Equatorial => 1,
             Frame::Ecliptic => 2,
             Frame::FK4 => 3,
             Frame::Galactic => 4,
             Frame::EclipticNonInertial(..) => 5,
+            Frame::Unknown(_) => 0,
+        }
+    }
+}
+
+impl From<i32> for Frame {
+    fn from(value: i32) -> Self {
+        match value {
+            1 => Frame::Equatorial,
+            2 => Frame::Ecliptic,
+            3 => Frame::FK4,
+            4 => Frame::Galactic,
+            i => Frame::Unknown(i as usize),
         }
     }
 }
@@ -128,6 +140,7 @@ const OBLIQUITY: f64 = 0.40909280422232897;
 ///
 /// The equation here is from the 2010 Astronomical Almanac.
 ///
+#[inline(always)]
 pub fn calc_obliquity(jd: f64) -> f64 {
     // centuries from j2000
     let c = (jd - Time::j2000().jd) / 365.25 / 100.0;
@@ -166,6 +179,7 @@ lazy_static! {
 /// # Arguments
 ///
 /// * `vec` - Vector, arbitrary units.
+#[inline(always)]
 pub fn fk4_to_ecliptic(vec: &Vector3<f64>) -> Vector3<f64> {
     FK4_ECLIPTIC_ROT.transform_vector(vec)
 }
@@ -175,6 +189,7 @@ pub fn fk4_to_ecliptic(vec: &Vector3<f64>) -> Vector3<f64> {
 /// # Arguments
 ///
 /// * `vec` - Vector, arbitrary units.
+#[inline(always)]
 pub fn ecliptic_to_fk4(vec: &Vector3<f64>) -> Vector3<f64> {
     FK4_ECLIPTIC_ROT.inverse().transform_vector(vec)
 }
@@ -184,6 +199,7 @@ pub fn ecliptic_to_fk4(vec: &Vector3<f64>) -> Vector3<f64> {
 /// # Arguments
 ///
 /// * `vec` - Vector, arbitrary units.
+#[inline(always)]
 pub fn galactic_to_ecliptic(vec: &Vector3<f64>) -> Vector3<f64> {
     GALACTIC_ECLIPTIC_ROT.transform_vector(vec)
 }
@@ -193,6 +209,7 @@ pub fn galactic_to_ecliptic(vec: &Vector3<f64>) -> Vector3<f64> {
 /// # Arguments
 ///
 /// * `vec` - Vector, arbitrary units.
+#[inline(always)]
 pub fn ecliptic_to_galactic(vec: &Vector3<f64>) -> Vector3<f64> {
     GALACTIC_ECLIPTIC_ROT.inverse().transform_vector(vec)
 }
@@ -203,6 +220,7 @@ pub fn ecliptic_to_galactic(vec: &Vector3<f64>) -> Vector3<f64> {
 /// # Arguments
 ///
 /// * `vec` - Vector, arbitrary units.
+#[inline(always)]
 pub fn ecliptic_to_equatorial(vec: &Vector3<f64>) -> Vector3<f64> {
     ECLIPTIC_EQUATORIAL_ROT.transform_vector(vec)
 }
@@ -213,17 +231,20 @@ pub fn ecliptic_to_equatorial(vec: &Vector3<f64>) -> Vector3<f64> {
 /// # Arguments
 ///
 /// * `vec` - Vector, arbitrary units.
+#[inline(always)]
 pub fn equatorial_to_ecliptic(vec: &Vector3<f64>) -> Vector3<f64> {
     ECLIPTIC_EQUATORIAL_ROT.inverse().transform_vector(vec)
 }
 
 /// Derivative of the z rotation matrix with respect to the rotation angle.
+#[inline(always)]
 fn rot_z_der(angle: f64) -> Matrix3<f64> {
     let (sin_a, cos_a) = angle.sin_cos();
     Matrix3::<f64>::from([[-sin_a, cos_a, 0.0], [-cos_a, -sin_a, 0.0], [0.0, 0.0, 0.0]])
 }
 
 /// Derivative of the x rotation matrix with respect to the rotation angle.
+#[inline(always)]
 fn rot_x_der(angle: f64) -> Matrix3<f64> {
     let (sin_a, cos_a) = angle.sin_cos();
     Matrix3::<f64>::from([[0.0, 0.0, 0.0], [0.0, -sin_a, cos_a], [0.0, -cos_a, -sin_a]])
@@ -237,6 +258,7 @@ fn rot_x_der(angle: f64) -> Matrix3<f64> {
 /// second is the derivative of the 3x3 matrix with respect to time. These two matrices
 /// may be used to compute the new position and velocities when moving from one frame
 /// to another.
+#[inline(always)]
 pub fn noninertial_rotation(frame_angles: &[f64; 6]) -> (Matrix3<f64>, Matrix3<f64>) {
     let r_z1 = Rotation3::from_axis_angle(&Vector3::z_axis(), frame_angles[0]);
     let r_x = Rotation3::from_axis_angle(&Vector3::x_axis(), frame_angles[1]);
@@ -262,6 +284,7 @@ pub fn noninertial_rotation(frame_angles: &[f64; 6]) -> (Matrix3<f64>, Matrix3<f
 /// defined by the provided angles. The first 3 angles here define the rotation ZXZ, the
 /// second three values define the derivative of the 3 angles. These angles define the
 /// rotation from the inertial to the non-inertial frame.
+#[inline(always)]
 pub fn noninertial_to_inertial(
     frame_angles: &[f64; 6],
     pos: &Vector3<f64>,
@@ -279,6 +302,7 @@ pub fn noninertial_to_inertial(
 /// defined by the provided angles. The first 3 angles here define the rotation ZXZ, the
 /// second three values define the derivative of the 3 angles. These angles define the
 /// rotation from the inertial to the non-inertial frame.
+#[inline(always)]
 pub fn inertial_to_noninertial(
     frame_angles: &[f64; 6],
     pos: &Vector3<f64>,
@@ -328,6 +352,7 @@ pub fn to_lat_lon(x: f64, y: f64, z: f64) -> (f64, f64) {
 /// * `rotation_vec` - The single vector around which to rotate the vectors.
 /// * `angle` - The angle in radians to rotate the vectors.
 ///
+#[inline(always)]
 pub fn rotate_around(
     vector: &Vector3<f64>,
     rotation_vec: Vector3<f64>,
@@ -383,6 +408,11 @@ mod tests {
         assert!((0.2 - vel_return[1]).abs() <= 10.0 * f64::EPSILON);
         assert!((0.3 - vel_return[2]).abs() <= 10.0 * f64::EPSILON);
     }
+    #[test]
+    fn test_frame_conversion() {
+        let f: Frame = 2i32.into();
+        assert!(f == Frame::Ecliptic);
+    }
 
     #[test]
     fn test_lat_lon() {

src/kete_core/src/spice/binary.rs → src/kete_core/src/io/bytes.rs

@@ -1,8 +1,9 @@
+//! Converting to and from bytes
 use crate::errors::{Error, KeteResult};
 use std::io::Read;
 
 /// Read the exact number of specified bytes from the file.
-pub fn read_bytes_exact<T: Read>(buffer: T, n_bytes: usize) -> KeteResult<Box<[u8]>> {
+pub(crate) fn read_bytes_exact<T: Read>(buffer: T, n_bytes: usize) -> KeteResult<Box<[u8]>> {
     let mut bytes = Vec::with_capacity(n_bytes);
     let n_read = buffer.take(n_bytes as u64).read_to_end(&mut bytes)?;
     if n_read != n_bytes {
@@ -13,7 +14,7 @@ pub fn read_bytes_exact<T: Read>(buffer: T, n_bytes: usize) -> KeteResult<Box<[u
 }
 
 /// Change a collection of bytes into a f64.
-pub fn bytes_to_f64(bytes: &[u8], little_endian: bool) -> KeteResult<f64> {
+pub(crate) fn bytes_to_f64(bytes: &[u8], little_endian: bool) -> KeteResult<f64> {
     let bytes: [u8; 8] = bytes
         .try_into()
         .map_err(|_| Error::IOError("File is not correctly formatted".into()))?;
@@ -25,7 +26,7 @@ pub fn bytes_to_f64(bytes: &[u8], little_endian: bool) -> KeteResult<f64> {
 }
 
 /// Change a collection of bytes into a vector of f64s.
-pub fn bytes_to_f64_vec(bytes: &[u8], little_endian: bool) -> KeteResult<Box<[f64]>> {
+pub(crate) fn bytes_to_f64_vec(bytes: &[u8], little_endian: bool) -> KeteResult<Box<[f64]>> {
     let byte_len = bytes.len();
     if byte_len % 8 != 0 {
         Err(Error::IOError("File is not correctly formatted".into()))?;
@@ -37,7 +38,7 @@ pub fn bytes_to_f64_vec(bytes: &[u8], little_endian: bool) -> KeteResult<Box<[f6
 }
 
 /// Change a collection of bytes into a vector of i32s.
-pub fn bytes_to_i32_vec(bytes: &[u8], little_endian: bool) -> KeteResult<Box<[i32]>> {
+pub(crate) fn bytes_to_i32_vec(bytes: &[u8], little_endian: bool) -> KeteResult<Box<[i32]>> {
     let byte_len = bytes.len();
     if byte_len % 4 != 0 {
         Err(Error::IOError("File is not correctly formatted".into()))?;
@@ -49,7 +50,7 @@ pub fn bytes_to_i32_vec(bytes: &[u8], little_endian: bool) -> KeteResult<Box<[i3
 }
 
 /// Change a collection of bytes into a i32.
-pub fn bytes_to_i32(bytes: &[u8], little_endian: bool) -> KeteResult<i32> {
+pub(crate) fn bytes_to_i32(bytes: &[u8], little_endian: bool) -> KeteResult<i32> {
     let bytes: [u8; 4] = bytes
         .try_into()
         .map_err(|_| Error::IOError("File is not correctly formatted".into()))?;
@@ -61,7 +62,7 @@ pub fn bytes_to_i32(bytes: &[u8], little_endian: bool) -> KeteResult<i32> {
 }
 
 /// Change a collection of bytes into a String.
-pub fn bytes_to_string(bytes: &[u8]) -> String {
+pub(crate) fn bytes_to_string(bytes: &[u8]) -> String {
     let mut bytes = bytes.to_vec();
     bytes.iter_mut().for_each(|x| {
         if x == &0x00 {
@@ -72,7 +73,7 @@ pub fn bytes_to_string(bytes: &[u8]) -> String {
 }
 
 /// Read a multiple contiguous f64s from the file.
-pub fn read_f64_vec<T: Read>(
+pub(crate) fn read_f64_vec<T: Read>(
     buffer: T,
     n_floats: usize,
     little_endian: bool,
@@ -84,7 +85,7 @@ pub fn read_f64_vec<T: Read>(
 /// Read a string of the specified length from the file.
 /// 0x00 are replaced with new lines, and new lines are stripped from the end of the
 /// string.
-pub fn read_str<T: Read>(buffer: T, length: usize) -> KeteResult<String> {
+pub(crate) fn read_str<T: Read>(buffer: T, length: usize) -> KeteResult<String> {
     let bytes = read_bytes_exact(buffer, length)?;
     Ok(bytes_to_string(&bytes))
 }

src/kete_core/src/io/mod.rs

@@ -1,6 +1,8 @@
 //! File IO related tools
 
-pub mod obs_codes;
+pub mod bytes;
+#[cfg(feature = "polars")]
+pub mod parquet;
 pub mod serde_const_arr;
 
 use crate::prelude::{Error, KeteResult};