Add a modicum of support for loading orbit estimates and spacecraft d…

…ynamics

src/dynamics/spacecraft.rs

@@ -37,10 +37,14 @@ use crate::cosmic::Cosm;
 #[cfg(feature = "python")]
 use crate::io::ConfigRepr;
 #[cfg(feature = "python")]
+use pyo3::class::basic::CompareOp;
+#[cfg(feature = "python")]
 use pyo3::prelude::*;
 #[cfg(feature = "python")]
 use pyo3::types::PyType;
 #[cfg(feature = "python")]
+use pythonize::depythonize;
+#[cfg(feature = "python")]
 use std::collections::HashMap;
 
 const NORM_ERR: f64 = 1e-4;
@@ -50,6 +54,7 @@ const NORM_ERR: f64 = 1e-4;
 /// Note: if the spacecraft runs out of fuel, the propagation segment will return an error.
 #[derive(Clone)]
 #[cfg_attr(feature = "python", pyclass)]
+#[cfg_attr(feature = "python", pyo3(module = "nyx_space.mission_design"))]
 pub struct SpacecraftDynamics {
     pub orbital_dyn: OrbitalDynamics,
     pub force_models: Vec<Arc<dyn ForceModel>>,
@@ -212,6 +217,25 @@ impl SpacecraftDynamics {
     fn __repr__(&self) -> String {
         format!("{self}")
     }
+
+    #[cfg(feature = "python")]
+    fn __richcmp__(&self, other: &Self, op: CompareOp) -> Result<bool, NyxError> {
+        match op {
+            CompareOp::Eq => Ok(self.__repr__() == other.__repr__()),
+            CompareOp::Ne => Ok(self.__repr__() != other.__repr__()),
+            _ => Err(NyxError::CustomError(format!("{op:?} not available"))),
+        }
+    }
+
+    #[cfg(feature = "python")]
+    #[classmethod]
+    /// Loads the SpacecraftDynamics from its YAML representation
+    fn loads(_cls: &PyType, state: &PyAny) -> Result<Self, ConfigError> {
+        <Self as Configurable>::from_config(
+            depythonize(state).map_err(|e| ConfigError::InvalidConfig(e.to_string()))?,
+            Cosm::de438(),
+        )
+    }
 }
 
 impl fmt::Display for SpacecraftDynamics {

src/od/noise/gauss_markov.rs

@@ -455,7 +455,7 @@ impl GaussMarkov {
     }
 
     #[cfg(feature = "python")]
-    /// Tries to load a GroundStation from the provided Python data
+    /// Loads the SpacecraftDynamics from its YAML representation
     #[classmethod]
     fn loads(_cls: &PyType, data: &PyAny) -> Result<Vec<Self>, ConfigError> {
         if let Ok(as_list) = data.downcast::<PyList>() {

src/python/mission_design/spacecraft.rs

@@ -103,6 +103,12 @@ impl Spacecraft {
         self.dumps(py)
     }
 
+    #[classmethod]
+    /// Loads the SpacecraftDynamics from its YAML representation
+    fn loads(_cls: &PyType, state: &PyAny) -> Result<Self, ConfigError> {
+        depythonize(state).map_err(|e| ConfigError::InvalidConfig(e.to_string()))
+    }
+
     fn __setstate__(&mut self, state: &PyAny) -> Result<(), ConfigError> {
         *self = depythonize(state).map_err(|e| ConfigError::InvalidConfig(e.to_string()))?;
         Ok(())

src/python/orbit_determination/estimate.rs

@@ -26,8 +26,10 @@ use crate::{
 };
 use nalgebra::Matrix6;
 use numpy::PyReadonlyArrayDyn;
+use pyo3::class::basic::CompareOp;
 use pyo3::prelude::*;
 use pyo3::types::PyType;
+use pythonize::depythonize;
 
 use super::ConfigError;
 
@@ -122,6 +124,23 @@ impl OrbitEstimate {
         Ok(selves)
     }
 
+    fn __richcmp__(&self, other: &Self, op: CompareOp) -> Result<bool, NyxError> {
+        match op {
+            CompareOp::Eq => Ok(self.__repr__() == other.__repr__()),
+            CompareOp::Ne => Ok(self.__repr__() != other.__repr__()),
+            _ => Err(NyxError::CustomError(format!("{op:?} not available"))),
+        }
+    }
+
+    #[classmethod]
+    /// Loads the SpacecraftDynamics from its YAML representation
+    fn loads(_cls: &PyType, state: &PyAny) -> Result<Self, ConfigError> {
+        <Self as Configurable>::from_config(
+            depythonize(state).map_err(|e| ConfigError::InvalidConfig(e.to_string()))?,
+            Cosm::de438(),
+        )
+    }
+
     // Manual getter/setters -- waiting on https://github.com/PyO3/pyo3/pull/2786
 
     #[getter]

src/python/orbit_determination/ground_station.rs

@@ -115,8 +115,8 @@ impl GroundStation {
         <Self as ConfigRepr>::load_named(path)
     }
 
-    /// Tries to load a GroundStation from the provided Python data
     #[classmethod]
+    /// Loads the SpacecraftDynamics from its YAML representation
     fn loads(_cls: &PyType, data: &PyAny) -> Result<HashMap<String, Self>, ConfigError> {
         if let Ok(as_list) = data.downcast::<PyList>() {
             let mut as_map = HashMap::new();

tests/python/test_cosmic.py

@@ -79,6 +79,8 @@ def test_define_spacecraft():
     # Check that we can pickle and dump the config of this spacecraft
     pkld = pickle.dumps(sc)
     print(sc.dumps())
+    sc_loaded = Spacecraft.loads(sc.dumps())
+    assert sc_loaded == sc
     unpkld = pickle.loads(pkld)
     assert unpkld == sc
 

tests/python/test_mission_design.py

@@ -5,6 +5,7 @@
 from timeit import timeit
 
 import pandas as pd
+import yaml
 from nyx_space.cosmic import Cosm, Orbit, Spacecraft, SrpConfig
 from nyx_space.mission_design import (
     Event,
@@ -283,6 +284,13 @@ def test_merge_traj():
         str(config_path.joinpath("dynamics.yaml"))
     )["lofi"]
 
+    # Check loading from the YAML read from Python
+    with open(config_path.joinpath("dynamics.yaml")) as fh:
+        data = yaml.safe_load(fh)
+
+    loaded = SpacecraftDynamics.loads(data["lofi"])
+    assert loaded == dynamics
+
     sc2, traj1 = propagate(sc1, dynamics, Unit.Day * 5)
     # And propagate again
     sc3, traj2 = propagate(sc2, dynamics, Unit.Day * 5)
@@ -303,4 +311,5 @@ def test_merge_traj():
 
 
 if __name__ == "__main__":
-    test_propagate()
+    # test_propagate()
+    test_merge_traj()

tests/python/test_orbit_determination.py

@@ -1,31 +1,32 @@
 import logging
-from pathlib import Path
 import pickle
-import numpy as np
-import pandas as pd
 import sys
+from pathlib import Path
 
+import numpy as np
+import pandas as pd
+import yaml
+from nyx_space.analysis import diff_traj_parquet
+from nyx_space.cosmic import Cosm, Spacecraft
+from nyx_space.mission_design import SpacecraftDynamics, TrajectoryLoader, propagate
 from nyx_space.orbit_determination import (
+    DynamicTrackingArc,
+    ExportCfg,
     GroundStation,
     GroundTrackingArcSim,
-    TrkConfig,
     OrbitEstimate,
-    process_tracking_arc,
+    TrkConfig,
     predictor,
-    DynamicTrackingArc,
-    ExportCfg,
+    process_tracking_arc,
 )
-from nyx_space.mission_design import TrajectoryLoader, SpacecraftDynamics, propagate
-from nyx_space.cosmic import Spacecraft, Cosm
-from nyx_space.time import Unit, TimeSeries
 from nyx_space.plots.od import (
     plot_covar,
     plot_estimates,
-    plot_residuals,
     plot_residual_histogram,
+    plot_residuals,
 )
 from nyx_space.plots.traj import plot_orbit_elements
-from nyx_space.analysis import diff_traj_parquet
+from nyx_space.time import TimeSeries, Unit
 
 
 def test_filter_arc():
@@ -101,6 +102,13 @@ def test_filter_arc():
         sc.orbit, covar=np.diag([100.0, 100.0, 100.0, 1.0, 1.0, 1.0])
     )
 
+    # Check loading from the YAML read from Python
+    with open(config_path.joinpath("orbit_estimates.yaml")) as fh:
+        data = yaml.safe_load(fh)
+
+    loaded = OrbitEstimate.loads(data["example 1"])
+    print(loaded)
+
     msr_noise = [1e-3, 0, 0, 1e-6]
     # Switch from sequential to EKF after 100 measurements
     ekf_num_msr_trig = 100