Skip to main content

nyx_space/dynamics/sequence/
python.rs

1use super::{AccelModels, Dynamics, ForceModels, PropagatorConfig, SpacecraftSequence, Thruster};
2#[cfg(feature = "python")]
3use crate::dynamics::{Drag, SolarPressure, SolidTides};
4use crate::propagators::{IntegratorMethod, IntegratorOptions};
5use crate::{dynamics::PointMasses, io::gravity::GravityFieldConfig};
6use pyo3::exceptions::PyException;
7use {
8    crate::Spacecraft,
9    anise::almanac::Almanac,
10    pyo3::prelude::*,
11    std::{collections::HashMap, sync::Arc},
12};
13
14#[cfg(feature = "python")]
15#[pymethods]
16impl SpacecraftSequence {
17    #[new]
18    fn py_new() -> Self {
19        SpacecraftSequence::default()
20    }
21
22    /// Load SpacecraftSequence from Dhall.
23    ///
24    /// :type dhall_str: str
25    /// :rtype: SpacecraftSequence
26    #[classmethod]
27    #[pyo3(name = "from_dhall")]
28    fn py_from_dhall(_cls: &Bound<'_, pyo3::types::PyType>, dhall_str: &str) -> PyResult<Self> {
29        serde_dhall::from_str(dhall_str)
30            .parse()
31            .map_err(|e| pyo3::exceptions::PyValueError::new_err(e.to_string()))
32    }
33
34    /// Load SpacecraftSequence from YAML.
35    ///
36    /// :type yaml_str: str
37    /// :rtype: SpacecraftSequence
38    #[classmethod]
39    #[pyo3(name = "from_yaml")]
40    fn py_from_yaml(_cls: &Bound<'_, pyo3::types::PyType>, yaml_str: &str) -> PyResult<Self> {
41        serde_yml::from_str(yaml_str)
42            .map_err(|e| pyo3::exceptions::PyValueError::new_err(e.to_string()))
43    }
44
45    /// Setup the sequence with the provided Almanac.
46    ///
47    /// :type almanac: Almanac
48    /// :rtype: None
49    #[pyo3(name = "setup")]
50    fn py_setup(&mut self, py: Python<'_>, almanac: Py<Almanac>) -> PyResult<()> {
51        let almanac_ref = almanac.borrow(py);
52        self.setup(Arc::new(almanac_ref.clone()))
53            .map_err(|e| pyo3::exceptions::PyValueError::new_err(e.to_string()))
54    }
55
56    /// Propagate the state through the sequence until a given phase.
57    ///
58    /// :type state: Spacecraft
59    /// :type until_phase: str | None
60    /// :type almanac: Almanac
61    /// :rtype: list[string, string]
62    #[pyo3(name = "propagate")]
63    fn py_propagate(
64        &self,
65        py: Python<'_>,
66        state: Spacecraft,
67        until_phase: Option<String>,
68        almanac: Py<Almanac>,
69    ) -> PyResult<Vec<(Option<String>, Vec<Spacecraft>)>> {
70        let almanac_ref = almanac.borrow(py);
71        let trajs = self
72            .propagate(state, until_phase, Arc::new(almanac_ref.clone()))
73            .map_err(|e| pyo3::exceptions::PyValueError::new_err(e.to_string()))?;
74
75        let result = trajs
76            .into_iter()
77            .map(|traj| (traj.name, traj.states))
78            .collect();
79        Ok(result)
80    }
81
82    #[getter]
83    fn get_thruster_sets(&self) -> HashMap<String, Thruster> {
84        self.thruster_sets.clone()
85    }
86
87    /// Insert a thruster with the given name into the thruster set.
88    ///
89    /// :type name: str
90    /// :type thruster: Thruster
91    /// :rtype: None
92    fn thruster_set_insert(&mut self, name: String, thruster: Thruster) {
93        self.thruster_sets.insert(name, thruster);
94    }
95
96    /// Remove a thruster with the given name from the thruster set.
97    ///
98    /// :type name: str
99    /// :rtype: None
100    fn thruster_set_remove(&mut self, name: String) -> PyResult<()> {
101        if self.thruster_sets.remove(&name).is_none() {
102            Err(PyException::new_err(format!("{name} not in thruster set")))
103        } else {
104            Ok(())
105        }
106    }
107}
108
109#[cfg(feature = "python")]
110#[cfg_attr(feature = "python", pymethods)]
111impl AccelModels {
112    #[pyo3(signature=(point_masses=None, gravity_field=None, solid_tides=None))]
113    #[new]
114    fn py_new(
115        point_masses: Option<PointMasses>,
116        gravity_field: Option<GravityFieldConfig>,
117        solid_tides: Option<SolidTides>,
118    ) -> Self {
119        Self {
120            point_masses,
121            gravity_field,
122            solid_tides,
123        }
124    }
125
126    fn __str__(&self) -> String {
127        format!("{self:?}")
128    }
129
130    fn __repr__(&self) -> String {
131        format!("{self:?} @ {self:p}")
132    }
133}
134
135#[cfg(feature = "python")]
136#[cfg_attr(feature = "python", pymethods)]
137impl ForceModels {
138    #[pyo3(signature=(solar_pressure=None, drag=None))]
139    #[new]
140    fn py_new(solar_pressure: Option<SolarPressure>, drag: Option<Drag>) -> Self {
141        Self {
142            solar_pressure,
143            drag,
144        }
145    }
146
147    fn __str__(&self) -> String {
148        format!("{self:?}")
149    }
150
151    fn __repr__(&self) -> String {
152        format!("{self:?} @ {self:p}")
153    }
154}
155
156#[cfg(feature = "python")]
157#[cfg_attr(feature = "python", pymethods)]
158impl Dynamics {
159    #[pyo3(signature=(accel_models=AccelModels::default(), force_models=ForceModels::default()))]
160    #[new]
161    fn py_new(accel_models: AccelModels, force_models: ForceModels) -> Self {
162        Self {
163            accel_models,
164            force_models,
165        }
166    }
167
168    fn __str__(&self) -> String {
169        format!("{self:?}")
170    }
171
172    fn __repr__(&self) -> String {
173        format!("{self:?} @ {self:p}")
174    }
175}
176
177#[cfg(feature = "python")]
178#[cfg_attr(feature = "python", pymethods)]
179impl PropagatorConfig {
180    #[new]
181    fn py_new(dynamics: Dynamics, method: IntegratorMethod, options: IntegratorOptions) -> Self {
182        Self {
183            dynamics,
184            method,
185            options,
186        }
187    }
188
189    fn __str__(&self) -> String {
190        format!("{self:?}")
191    }
192
193    fn __repr__(&self) -> String {
194        format!("{self:?} @ {self:p}")
195    }
196}