Skip to content

Commit

Permalink
Add support for saving and loading of States from Parquet files (#146)
Browse files Browse the repository at this point in the history
* Add parquet support for SimultaneousStates
  • Loading branch information
dahlend authored Oct 29, 2024
1 parent c2391dc commit ce17749
Show file tree
Hide file tree
Showing 15 changed files with 249 additions and 29 deletions.
7 changes: 7 additions & 0 deletions CHANGELOG.md
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down
2 changes: 1 addition & 1 deletion Cargo.toml
Original file line number Diff line number Diff line change
Expand Up @@ -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"}
Expand Down
3 changes: 2 additions & 1 deletion src/kete/rust/lib.rs
Original file line number Diff line number Diff line change
Expand Up @@ -24,6 +24,7 @@
)]

use pyo3::prelude::*;
use state::PyState;

pub mod covariance;
pub mod elements;
Expand Down Expand Up @@ -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>()?;
Expand Down
24 changes: 24 additions & 0 deletions src/kete/rust/simult_states.rs
Original file line number Diff line number Diff line change
Expand Up @@ -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()
Expand Down
2 changes: 1 addition & 1 deletion src/kete/rust/spice/mod.rs
Original file line number Diff line number Diff line change
Expand Up @@ -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,
Expand Down
1 change: 1 addition & 0 deletions src/kete_core/Cargo.toml
Original file line number Diff line number Diff line change
Expand Up @@ -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"
Expand Down
32 changes: 31 additions & 1 deletion src/kete_core/src/frames/definitions.rs
Original file line number Diff line number Diff line change
Expand Up @@ -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),
}
}
}
Expand Down Expand Up @@ -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;
Expand Down Expand Up @@ -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)
}
Expand All @@ -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)
}
Expand All @@ -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)
}
Expand All @@ -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)
}
Expand All @@ -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)
}
Expand All @@ -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]])
Expand All @@ -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]);
Expand All @@ -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>,
Expand All @@ -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>,
Expand Down Expand Up @@ -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>,
Expand Down Expand Up @@ -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() {
Expand Down
Original file line number Diff line number Diff line change
@@ -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 {
Expand All @@ -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()))?;
Expand All @@ -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()))?;
Expand All @@ -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()))?;
Expand All @@ -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()))?;
Expand All @@ -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 {
Expand All @@ -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,
Expand All @@ -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))
}
4 changes: 3 additions & 1 deletion src/kete_core/src/io/mod.rs
Original file line number Diff line number Diff line change
@@ -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};
Expand Down
Loading

0 comments on commit ce17749

Please sign in to comment.