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    #[classmethod]
23    #[pyo3(name = "from_dhall")]
24    fn py_from_dhall(_cls: &Bound<'_, pyo3::types::PyType>, dhall_str: &str) -> PyResult<Self> {
25        serde_dhall::from_str(dhall_str)
26            .parse()
27            .map_err(|e| pyo3::exceptions::PyValueError::new_err(e.to_string()))
28    }
29
30    #[classmethod]
31    #[pyo3(name = "from_yaml")]
32    fn py_from_yaml(_cls: &Bound<'_, pyo3::types::PyType>, yaml_str: &str) -> PyResult<Self> {
33        serde_yml::from_str(yaml_str)
34            .map_err(|e| pyo3::exceptions::PyValueError::new_err(e.to_string()))
35    }
36
37    #[pyo3(name = "setup")]
38    fn py_setup(&mut self, py: Python<'_>, almanac: Py<Almanac>) -> PyResult<()> {
39        let almanac_ref = almanac.borrow(py);
40        self.setup(Arc::new(almanac_ref.clone()))
41            .map_err(|e| pyo3::exceptions::PyValueError::new_err(e.to_string()))
42    }
43
44    #[pyo3(name = "propagate")]
45    fn py_propagate(
46        &self,
47        py: Python<'_>,
48        state: Spacecraft,
49        until_phase: Option<String>,
50        almanac: Py<Almanac>,
51    ) -> PyResult<Vec<(Option<String>, Vec<Spacecraft>)>> {
52        let almanac_ref = almanac.borrow(py);
53        let trajs = self
54            .propagate(state, until_phase, Arc::new(almanac_ref.clone()))
55            .map_err(|e| pyo3::exceptions::PyValueError::new_err(e.to_string()))?;
56
57        let result = trajs
58            .into_iter()
59            .map(|traj| (traj.name, traj.states))
60            .collect();
61        Ok(result)
62    }
63
64    #[getter]
65    fn get_thruster_sets(&self) -> HashMap<String, Thruster> {
66        self.thruster_sets.clone()
67    }
68
69    fn thruster_set_insert(&mut self, name: String, thruster: Thruster) {
70        self.thruster_sets.insert(name, thruster);
71    }
72    fn thruster_set_remove(&mut self, name: String) -> PyResult<()> {
73        if self.thruster_sets.remove(&name).is_none() {
74            Err(PyException::new_err(format!("{name} not in thruster set")))
75        } else {
76            Ok(())
77        }
78    }
79}
80
81#[cfg(feature = "python")]
82#[cfg_attr(feature = "python", pymethods)]
83impl AccelModels {
84    #[pyo3(signature=(point_masses=None, gravity_field=None, solid_tides=None))]
85    #[new]
86    fn py_new(
87        point_masses: Option<PointMasses>,
88        gravity_field: Option<GravityFieldConfig>,
89        solid_tides: Option<SolidTides>,
90    ) -> Self {
91        Self {
92            point_masses,
93            gravity_field,
94            solid_tides,
95        }
96    }
97
98    fn __str__(&self) -> String {
99        format!("{self:?}")
100    }
101
102    fn __repr__(&self) -> String {
103        format!("{self:?} @ {self:p}")
104    }
105}
106
107#[cfg(feature = "python")]
108#[cfg_attr(feature = "python", pymethods)]
109impl ForceModels {
110    #[pyo3(signature=(solar_pressure=None, drag=None))]
111    #[new]
112    fn py_new(solar_pressure: Option<SolarPressure>, drag: Option<Drag>) -> Self {
113        Self {
114            solar_pressure,
115            drag,
116        }
117    }
118
119    fn __str__(&self) -> String {
120        format!("{self:?}")
121    }
122
123    fn __repr__(&self) -> String {
124        format!("{self:?} @ {self:p}")
125    }
126}
127
128#[cfg(feature = "python")]
129#[cfg_attr(feature = "python", pymethods)]
130impl Dynamics {
131    #[pyo3(signature=(accel_models=AccelModels::default(), force_models=ForceModels::default()))]
132    #[new]
133    fn py_new(accel_models: AccelModels, force_models: ForceModels) -> Self {
134        Self {
135            accel_models,
136            force_models,
137        }
138    }
139
140    fn __str__(&self) -> String {
141        format!("{self:?}")
142    }
143
144    fn __repr__(&self) -> String {
145        format!("{self:?} @ {self:p}")
146    }
147}
148
149#[cfg(feature = "python")]
150#[cfg_attr(feature = "python", pymethods)]
151impl PropagatorConfig {
152    #[new]
153    fn py_new(dynamics: Dynamics, method: IntegratorMethod, options: IntegratorOptions) -> Self {
154        Self {
155            dynamics,
156            method,
157            options,
158        }
159    }
160
161    fn __str__(&self) -> String {
162        format!("{self:?}")
163    }
164
165    fn __repr__(&self) -> String {
166        format!("{self:?} @ {self:p}")
167    }
168}