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};
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))]
85    #[new]
86    fn py_new(
87        point_masses: Option<PointMasses>,
88        gravity_field: Option<GravityFieldConfig>,
89    ) -> Self {
90        Self {
91            point_masses,
92            gravity_field,
93        }
94    }
95
96    fn __str__(&self) -> String {
97        format!("{self:?}")
98    }
99
100    fn __repr__(&self) -> String {
101        format!("{self:?} @ {self:p}")
102    }
103}
104
105#[cfg(feature = "python")]
106#[cfg_attr(feature = "python", pymethods)]
107impl ForceModels {
108    #[pyo3(signature=(solar_pressure=None, drag=None))]
109    #[new]
110    fn py_new(solar_pressure: Option<SolarPressure>, drag: Option<Drag>) -> Self {
111        Self {
112            solar_pressure,
113            drag,
114        }
115    }
116
117    fn __str__(&self) -> String {
118        format!("{self:?}")
119    }
120
121    fn __repr__(&self) -> String {
122        format!("{self:?} @ {self:p}")
123    }
124}
125
126#[cfg(feature = "python")]
127#[cfg_attr(feature = "python", pymethods)]
128impl Dynamics {
129    #[pyo3(signature=(accel_models=AccelModels::default(), force_models=ForceModels::default()))]
130    #[new]
131    fn py_new(accel_models: AccelModels, force_models: ForceModels) -> Self {
132        Self {
133            accel_models,
134            force_models,
135        }
136    }
137
138    fn __str__(&self) -> String {
139        format!("{self:?}")
140    }
141
142    fn __repr__(&self) -> String {
143        format!("{self:?} @ {self:p}")
144    }
145}
146
147#[cfg(feature = "python")]
148#[cfg_attr(feature = "python", pymethods)]
149impl PropagatorConfig {
150    #[new]
151    fn py_new(dynamics: Dynamics, method: IntegratorMethod, options: IntegratorOptions) -> Self {
152        Self {
153            dynamics,
154            method,
155            options,
156        }
157    }
158
159    fn __str__(&self) -> String {
160        format!("{self:?}")
161    }
162
163    fn __repr__(&self) -> String {
164        format!("{self:?} @ {self:p}")
165    }
166}