diff --git a/.github/workflows/python.yml b/.github/workflows/python.yml index f9ece15d..d78dbc80 100644 --- a/.github/workflows/python.yml +++ b/.github/workflows/python.yml @@ -123,7 +123,7 @@ jobs: - uses: actions/checkout@v4 - uses: actions/setup-python@v5 with: - python-version: 3.x + python-version: 3.12 - name: Build wheels uses: PyO3/maturin-action@v1 with: diff --git a/crates/lox-orbits/src/python.rs b/crates/lox-orbits/src/python.rs index 75f543a4..08e39d03 100644 --- a/crates/lox-orbits/src/python.rs +++ b/crates/lox-orbits/src/python.rs @@ -999,13 +999,40 @@ impl From for PyErr { } } -#[pyclass(name = "ElevationMask", module = "lox_space", frozen)] +#[pyclass(name = "ElevationMask", module = "lox_space", frozen, eq)] +#[derive(Debug, Clone, PartialEq)] pub struct PyElevationMask(pub ElevationMask); #[pymethods] impl PyElevationMask { #[new] + #[pyo3(signature = (azimuth=None, elevation=None, min_elevation=None))] fn new( + azimuth: Option<&Bound<'_, PyArray1>>, + elevation: Option<&Bound<'_, PyArray1>>, + min_elevation: Option, + ) -> PyResult { + if let Some(min_elevation) = min_elevation { + return Ok(PyElevationMask(ElevationMask::with_fixed_elevation( + min_elevation, + ))); + } + if let (Some(azimuth), Some(elevation)) = (azimuth, elevation) { + let azimuth = azimuth.to_vec()?; + let elevation = elevation.to_vec()?; + return Ok(PyElevationMask(ElevationMask::new(azimuth, elevation)?)); + } + Err(PyValueError::new_err("invalid argument combination, either `min_elevation` or `azimuth` and `elevation` arrays need to be present")) + } + + #[classmethod] + fn fixed(_cls: &Bound<'_, PyType>, min_elevation: f64) -> Self { + PyElevationMask(ElevationMask::with_fixed_elevation(min_elevation)) + } + + #[classmethod] + fn variable( + _cls: &Bound<'_, PyType>, azimuth: &Bound<'_, PyArray1>, elevation: &Bound<'_, PyArray1>, ) -> PyResult { @@ -1014,9 +1041,29 @@ impl PyElevationMask { Ok(PyElevationMask(ElevationMask::new(azimuth, elevation)?)) } - #[classmethod] - fn fixed(_cls: &Bound<'_, PyType>, min_elevation: f64) -> Self { - PyElevationMask(ElevationMask::with_fixed_elevation(min_elevation)) + fn __getnewargs__(&self) -> (Option>, Option>, Option) { + (self.azimuth(), self.elevation(), self.min_elevation()) + } + + fn azimuth(&self) -> Option> { + match &self.0 { + ElevationMask::Fixed(_) => None, + ElevationMask::Variable(series) => Some(series.x().to_vec()), + } + } + + fn elevation(&self) -> Option> { + match &self.0 { + ElevationMask::Fixed(_) => None, + ElevationMask::Variable(series) => Some(series.y().to_vec()), + } + } + + fn min_elevation(&self) -> Option { + match &self.0 { + ElevationMask::Fixed(min_elevation) => Some(*min_elevation), + ElevationMask::Variable(_) => None, + } } } diff --git a/crates/lox-space/src/lib.rs b/crates/lox-space/src/lib.rs index afd2eb35..c56fbddc 100644 --- a/crates/lox-space/src/lib.rs +++ b/crates/lox-space/src/lib.rs @@ -9,9 +9,9 @@ use lox_bodies::python::{PyBarycenter, PyMinorBody, PyPlanet, PySatellite, PySun}; use lox_ephem::python::PySpk; use lox_orbits::python::{ - elevation, find_events, find_windows, visibility, PyEvent, PyFrame, PyGroundLocation, - PyGroundPropagator, PyKeplerian, PyObservables, PySgp4, PyState, PyTopocentric, PyTrajectory, - PyVallado, PyWindow, + elevation, find_events, find_windows, visibility, PyElevationMask, PyEvent, PyFrame, + PyGroundLocation, PyGroundPropagator, PyKeplerian, PyObservables, PySgp4, PyState, + PyTopocentric, PyTrajectory, PyVallado, PyWindow, }; use pyo3::prelude::*; @@ -50,5 +50,6 @@ fn lox_space(m: &Bound<'_, PyModule>) -> PyResult<()> { m.add_class::()?; m.add_class::()?; m.add_class::()?; + m.add_class::()?; Ok(()) } diff --git a/crates/lox-space/tests/test_pickle.py b/crates/lox-space/tests/test_pickle.py index 03a9cb36..b90f33e6 100644 --- a/crates/lox-space/tests/test_pickle.py +++ b/crates/lox-space/tests/test_pickle.py @@ -16,6 +16,7 @@ lox.Planet("Earth"), lox.Satellite("Moon"), lox.MinorBody("Ceres"), + lox.ElevationMask.fixed(0.0), ]) def test_pickle(obj): pickled = pickle.dumps(obj)