Merge pull request #221 from nyx-space/relax-msr-generation

Relax tracking sim from Python

Cargo.toml

@@ -62,11 +62,11 @@ numpy = { version = "0.19", optional = true }
 indicatif = { version = "0.17", features = ["rayon"] }
 rstats = "1.2.50"
 thiserror = "1.0"
-parquet = { version = "45.0.0", default-features = false, features = [
+parquet = { version = "46.0.0", default-features = false, features = [
     "arrow",
-    "brotli",
+    "zstd",
 ] }
-arrow = "45.0.0"
+arrow = "46.0.0"
 shadow-rs = { version = "0.23.0", default-features = false }
 serde_yaml = "0.9.21"
 whoami = "1.3.0"
@@ -92,3 +92,8 @@ python = ["pyo3", "pyo3-log", "hifitime/python", "numpy", "pythonize"]
 [lib]
 crate-type = ["cdylib", "rlib"]
 name = "nyx_space"
+
+[target.x86_64-unknown-linux-gnu]
+# For flamegraph -- https://github.com/flamegraph-rs/flamegraph
+linker = "/usr/bin/clang"
+rustflags = ["-Clink-arg=-fuse-ld=lld", "-Clink-arg=-Wl,--no-rosegment"]

src/cosmic/cosm.rs

@@ -542,7 +542,7 @@ impl Cosm {
             let splt: Vec<_> = name.split(' ').collect();
             if splt[0] == "iau" {
                 // This is an IAU frame, so the orientation is specified first, and we don't capitalize the ephemeris name
-                vec![splt[0].to_string(), splt[1..splt.len()].join(" ")].join(" ")
+                [splt[0].to_string(), splt[1..splt.len()].join(" ")].join(" ")
             } else {
                 // Likely a default center and frame, so let's do some clever guessing and capitalize the words
                 let frame_name = capitalize(splt[splt.len() - 1]);
@@ -552,7 +552,7 @@ impl Cosm {
                     .collect::<Vec<_>>()
                     .join(" ");
 
-                vec![ephem_name, frame_name].join(" ")
+                [ephem_name, frame_name].join(" ")
             }
         }
     }

src/cosmic/spacecraft.rs

@@ -83,12 +83,6 @@ impl From<GuidanceMode> for f64 {
 #[derive(Clone, Copy, Debug, Serialize, Deserialize)]
 #[cfg_attr(feature = "python", pyclass)]
 #[cfg_attr(feature = "python", pyo3(module = "nyx_space.cosmic"))]
-#[cfg_attr(
-    feature = "python",
-    pyo3(
-        text_signature = "(orbit, dry_mass_kg, fuel_mass_kg, srp_area_m2, drag_area_m2, cr, cd, thruster, mode)"
-    )
-)]
 pub struct Spacecraft {
     /// Initial orbit the vehicle is in
     #[serde(deserialize_with = "orbit_from_str")]
@@ -127,7 +121,6 @@ impl Default for Spacecraft {
 }
 
 #[cfg_attr(feature = "python", pyclass)]
-#[cfg_attr(feature = "python", pyo3(text_signature = "(area_m2, cr=1.8)"))]
 #[cfg_attr(feature = "python", pyo3(module = "nyx_space.cosmic"))]
 #[derive(Copy, Clone, Debug, Serialize, Deserialize, PartialEq)]
 /// The Solar Radiation Pressure configuration for a spacecraft
@@ -158,7 +151,6 @@ impl Default for SrpConfig {
 }
 
 #[cfg_attr(feature = "python", pyclass)]
-#[cfg_attr(feature = "python", pyo3(text_signature = "(area_m2, cd=2.2)"))]
 #[cfg_attr(feature = "python", pyo3(module = "nyx_space.cosmic"))]
 #[derive(Copy, Clone, Debug, Serialize, Deserialize, PartialEq)]
 /// The drag configuration for a spacecraft

src/io/mod.rs

@@ -64,12 +64,6 @@ use pyo3::prelude::*;
 /// Configuration for exporting a trajectory to parquet.
 #[derive(Clone, Default, Serialize, Deserialize, TypedBuilder)]
 #[cfg_attr(feature = "python", pyclass)]
-#[cfg_attr(
-    feature = "python",
-    pyo3(
-        text_signature = "(timestamp=None, fields=None, start_epoch=None, step=None, end_epoch=None, metadata=None)"
-    )
-)]
 pub struct ExportCfg {
     /// Fields to export, if unset, defaults to all possible fields.
     #[builder(default, setter(strip_option))]
@@ -147,6 +141,9 @@ impl ExportCfg {
 #[pymethods]
 impl ExportCfg {
     #[new]
+    #[pyo3(
+        text_signature = "(timestamp=None, fields=None, start_epoch=None, step=None, end_epoch=None, metadata=None)"
+    )]
     fn py_new(
         timestamp: Option<bool>,
         fields: Option<Vec<StateParameter>>,

src/io/trajectory_data.rs

@@ -44,10 +44,6 @@ use crate::cosmic::Cosm;
 /// A dynamic trajectory allows loading a trajectory Parquet file and converting it
 /// to the concrete trajectory state type when desired.
 #[cfg_attr(feature = "python", pyclass)]
-#[cfg_attr(
-    feature = "python",
-    pyo3(text_signature = "(path, format='parquet', parquet_path=None, spacecraft_template=None)")
-)]
 #[cfg_attr(feature = "python", pyo3(module = "nyx_space.mission_design"))]
 #[derive(Clone, PartialEq)]
 pub struct TrajectoryLoader {
@@ -276,6 +272,9 @@ impl Display for TrajectoryLoader {
 impl TrajectoryLoader {
     /// Initializes a new dynamic trajectory from the provided file, and the format kind
     #[new]
+    #[pyo3(
+        text_signature = "(path, format='parquet', parquet_path=None, spacecraft_template=None)"
+    )]
     fn new(
         path: String,
         format: Option<String>,

src/io/watermark.rs

@@ -20,7 +20,7 @@ use std::collections::HashMap;
 
 use hifitime::Epoch;
 use parquet::{
-    basic::{BrotliLevel, Compression},
+    basic::{Compression, ZstdLevel},
     file::properties::WriterProperties,
     format::KeyValue,
 };
@@ -32,7 +32,7 @@ shadow!(build);
 /// The parquet writer properties
 pub(crate) fn pq_writer(metadata: Option<HashMap<String, String>>) -> Option<WriterProperties> {
     let bldr = WriterProperties::builder()
-        .set_compression(Compression::BROTLI(BrotliLevel::try_new(10).unwrap()));
+        .set_compression(Compression::ZSTD(ZstdLevel::try_new(10).unwrap()));
 
     let mut file_metadata = vec![
         KeyValue::new("Generated by".to_string(), prj_name_ver()),

src/md/events/mod.rs

@@ -51,12 +51,6 @@ where
 /// Defines a state parameter event finder
 #[derive(Clone, Debug)]
 #[cfg_attr(feature = "python", pyclass)]
-#[cfg_attr(
-    feature = "python",
-    pyo3(
-        text_signature = "(parameter, desired_value, epoch_precision=None, value_precision=None)"
-    )
-)]
 pub struct Event {
     /// The state parameter
     pub parameter: StateParameter,

src/md/trajectory/traj.rs

@@ -596,6 +596,25 @@ where
         Ok(traj)
     }
 
+    /// Rebuilds this trajectory with the provided epochs.
+    /// This may lead to aliasing due to the Nyquist–Shannon sampling theorem.
+    pub fn rebuild(&self, epochs: &[Epoch]) -> Result<Self, NyxError> {
+        if self.states.is_empty() {
+            return Err(NyxError::Trajectory(TrajError::CreationError(
+                "No trajectory to convert".to_string(),
+            )));
+        }
+
+        let mut traj = Self::new();
+        for epoch in epochs {
+            traj.states.push(self.at(*epoch)?);
+        }
+
+        traj.finalize();
+
+        Ok(traj)
+    }
+
     /// Export the difference in RIC from of this trajectory compare to the "other" trajectory in parquet format.
     ///
     /// # Notes

src/od/noise/gauss_markov.rs

@@ -61,7 +61,6 @@ use std::sync::Arc;
 /// This allows the users to model a white noise process without having to change the process type.
 #[derive(Copy, Clone, Debug, Serialize, Deserialize, PartialEq)]
 #[cfg_attr(feature = "python", pyclass)]
-#[cfg_attr(feature = "python", pyo3(text_signature = "(tau, sigma, state_state)"))]
 #[cfg_attr(feature = "python", pyo3(module = "nyx_space.orbit_determination"))]
 pub struct GaussMarkov {
     /// The time constant, tau gives the correlation time, or the time over which the intensity of the time correlation will fade to 1/e of its prior value. (This is sometimes incorrectly referred to as the "half-life" of the process.)
@@ -359,6 +358,7 @@ impl GaussMarkov {
 
     #[cfg(feature = "python")]
     #[new]
+    #[pyo3(text_signature = "(tau, sigma, state_state)")]
     fn py_new(
         tau: Option<Duration>,
         sigma: Option<f64>,

src/od/process/mod.rs

@@ -26,7 +26,6 @@ pub use crate::od::*;
 use crate::propagators::error_ctrl::ErrorCtrl;
 use crate::propagators::PropInstance;
 pub use crate::time::{Duration, Unit};
-use crate::State;
 mod conf;
 pub use conf::{IterationConf, SmoothingArc};
 mod trigger;

src/od/process/rejectcrit.rs

@@ -34,10 +34,6 @@ use std::sync::Arc;
 #[derive(Copy, Clone, Debug, Serialize, Deserialize, PartialEq)]
 #[cfg_attr(feature = "python", pyclass)]
 #[cfg_attr(feature = "python", pyo3(module = "nyx_space.orbit_determination"))]
-#[cfg_attr(
-    feature = "python",
-    pyo3(text_signature = "(min_accepted=None, num_sigmas=None)")
-)]
 pub struct FltResid {
     /// Minimum number of accepted measurements before applying the rejection criteria.
     pub min_accepted: usize,
@@ -49,6 +45,7 @@ pub struct FltResid {
 #[pymethods]
 impl FltResid {
     #[new]
+    #[pyo3(text_signature = "(min_accepted=None, num_sigmas=None)")]
     fn py_new(min_accepted: Option<usize>, num_sigmas: Option<f64>) -> Self {
         let mut me = Self::default();
         if let Some(min_accepted) = min_accepted {

src/od/simulator/arc.rs

@@ -156,6 +156,11 @@ where
         self.allow_overlap = false;
     }
 
+    /// Allows overlapping measurements
+    pub fn allow_overlap(&mut self) {
+        self.allow_overlap = true;
+    }
+
     /// Generates measurements from the simulated tracking arc.
     ///
     /// Notes:

src/od/simulator/trkconfig.rs

@@ -38,12 +38,6 @@ use super::Availability;
 #[derive(Clone, Debug, Serialize, Deserialize, PartialEq)]
 #[cfg_attr(feature = "python", pyclass)]
 #[cfg_attr(feature = "python", pyo3(module = "nyx_space.orbit_determination"))]
-#[cfg_attr(
-    feature = "python",
-    pyo3(
-        text_signature = "(start=None, end=None, schedule_on=None, schedule_off=None, sampling=None)"
-    )
-)]
 pub struct TrkConfig {
     /// Availability configuration to start the tracking arc
     #[serde(default)]

src/python/mission_design/events.rs

@@ -25,6 +25,9 @@ use hifitime::Unit;
 impl Event {
     /// Initializes a new event. Arguments are "parameter: StateParameter" and "desired_value: float".
     #[new]
+    #[pyo3(
+        text_signature = "(parameter, desired_value, epoch_precision=None, value_precision=None)"
+    )]
     fn py_new(
         parameter: StateParameter,
         desired_value: f64,

src/python/mission_design/orbit_trajectory.rs

@@ -69,6 +69,13 @@ impl OrbitTraj {
         Ok(Self { inner })
     }
 
+    /// Copies this object and rebuilds it with the provided epochs
+    fn rebuild(&self, epochs: Vec<Epoch>) -> Result<Self, NyxError> {
+        let inner = self.inner.rebuild(&epochs)?;
+
+        Ok(Self { inner })
+    }
+
     /// Finds a specific event in a trajectory.
     ///
     /// If a start or end epoch is provided (or both are provided), this function will return a list of a single event.

src/python/mission_design/sc_trajectory.rs

@@ -70,6 +70,13 @@ impl SpacecraftTraj {
         Ok(Self { inner })
     }
 
+    /// Copies this object and rebuilds it with the provided epochs
+    fn rebuild(&self, epochs: Vec<Epoch>) -> Result<Self, NyxError> {
+        let inner = self.inner.rebuild(&epochs)?;
+
+        Ok(Self { inner })
+    }
+
     /// Finds a specific event in a trajectory.
     ///
     /// If a start or end epoch is provided (or both are provided), this function will return a list of a single event.

src/python/mission_design/spacecraft.rs

@@ -35,6 +35,9 @@ use pythonize::{depythonize, pythonize};
 impl Spacecraft {
     /// Initialize a new Spacecraft with optional thruster, mode, SRP, and Drag parameters.
     #[new]
+    #[pyo3(
+        text_signature = "(orbit, dry_mass_kg, fuel_mass_kg, srp_area_m2, drag_area_m2, cr, cd, thruster, mode)"
+    )]
     pub fn py_new(
         orbit: Option<Orbit>,
         dry_mass_kg: Option<f64>,
@@ -149,6 +152,7 @@ impl Spacecraft {
 #[pymethods]
 impl SrpConfig {
     #[new]
+    #[pyo3(text_signature = "(area_m2, cr=1.8)")]
     pub fn py_new(area_m2: Option<f64>, cr: Option<f64>) -> Self {
         Self {
             area_m2: area_m2.unwrap_or(0.0),
@@ -207,6 +211,7 @@ impl SrpConfig {
 #[pymethods]
 impl DragConfig {
     #[new]
+    #[pyo3(text_signature = "(area_m2, cd=2.2)")]
     pub fn py_new(area_m2: Option<f64>, cd: Option<f64>) -> Self {
         Self {
             area_m2: area_m2.unwrap_or(0.0),

src/python/orbit_determination/arc.rs

@@ -16,22 +16,25 @@
     along with this program.  If not, see <https://www.gnu.org/licenses/>.
 */
 
-use std::collections::HashMap;
-
 use crate::cosmic::Cosm;
 use crate::io::trajectory_data::TrajectoryLoader;
 use crate::io::ExportCfg;
 use crate::od::msr::RangeDoppler;
 use crate::od::simulator::TrackingArcSim;
 pub use crate::od::simulator::TrkConfig;
 pub use crate::{io::ConfigError, od::prelude::GroundStation};
-use crate::{NyxError, Spacecraft};
+use crate::{NyxError, Orbit, Spacecraft};
+use either::Either;
 use pyo3::prelude::*;
+use std::collections::HashMap;
 
 #[derive(Clone)]
 #[pyclass]
 pub struct GroundTrackingArcSim {
-    inner: TrackingArcSim<Spacecraft, RangeDoppler, GroundStation>,
+    inner: Either<
+        TrackingArcSim<Spacecraft, RangeDoppler, GroundStation>,
+        TrackingArcSim<Orbit, RangeDoppler, GroundStation>,
+    >,
 }
 
 #[pymethods]
@@ -43,14 +46,37 @@ impl GroundTrackingArcSim {
         trajectory: TrajectoryLoader,
         configs: HashMap<String, TrkConfig>,
         seed: u64,
+        allow_overlap: Option<bool>,
     ) -> Result<Self, NyxError> {
-        // Try to convert the dynamic trajectory into an Orbit trajectory
-        let traj = trajectory
-            .to_traj()
-            .map_err(|e| NyxError::CustomError(e.to_string()))?;
+        // Try to convert the dynamic trajectory into a trajectory
+        let inner = if let Ok(sc_traj) = trajectory.to_traj::<Spacecraft>() {
+            let mut inner = TrackingArcSim::with_seed(devices, sc_traj, configs, seed)
+                .map_err(NyxError::ConfigError)?;
+
+            if let Some(allow_overlap) = allow_overlap {
+                if allow_overlap {
+                    inner.allow_overlap();
+                } else {
+                    inner.disallow_overlap();
+                }
+            }
+            Either::Left(inner)
+        } else if let Ok(traj) = trajectory.to_traj::<Orbit>() {
+            let mut inner = TrackingArcSim::with_seed(devices, traj, configs, seed)
+                .map_err(NyxError::ConfigError)?;
+
+            if let Some(allow_overlap) = allow_overlap {
+                if allow_overlap {
+                    inner.allow_overlap();
+                } else {
+                    inner.disallow_overlap();
+                }
+            }
 
-        let inner = TrackingArcSim::with_seed(devices, traj, configs, seed)
-            .map_err(NyxError::ConfigError)?;
+            Either::Right(inner)
+        } else {
+            return Err(NyxError::CustomError("Provided trajectory could neither be parsed as an orbit trajectory or a spacecraft trajectory".to_string()));
+        };
 
         Ok(Self { inner })
     }
@@ -64,7 +90,10 @@ impl GroundTrackingArcSim {
         export_cfg: ExportCfg,
     ) -> Result<String, NyxError> {
         let cosm = Cosm::de438();
-        let arc = self.inner.generate_measurements(cosm)?;
+        let arc = match &mut self.inner {
+            Either::Left(arc_sim) => arc_sim.generate_measurements(cosm)?,
+            Either::Right(arc_sim) => arc_sim.generate_measurements(cosm)?,
+        };
 
         // Save the tracking arc
         let maybe = arc.to_parquet(path, export_cfg);