Merge pull request #454 from nyx-space/feat/python-parallel-exec

Enable Rust multithreading speeds for Python: search for the `many` keyword in the functions

Author: ChristopherRabotin

Cargo.toml

@@ -44,6 +44,7 @@ pyo3 = { version = "0.25", features = ["multiple-pymethods"] }
 pyo3-log = "0.12"
 numpy = "0.25"
 ndarray = ">= 0.15, < 0.17"
+rayon = "1.10.0"
 
 anise = { version = "0.6.1", path = "anise", default-features = false }
 

anise-py/README.md

@@ -1,6 +1,6 @@
 # ANISE (Attitude, Navigation, Instrument, Spacecraft, Ephemeris)
 
-ANISE is a rewrite of the core functionalities of the NAIF SPICE toolkit with enhanced performance, and ease of use, while leveraging Rust's safety and speed.
+ANISE is a modern rewrite of the core functionalities of the NAIF SPICE toolkit with enhanced features, and ease of use, while leveraging Rust's safety and speed.
 
 [**Please fill out our user survey**](https://7ug5imdtt8v.typeform.com/to/qYDB14Hj)
 
@@ -10,6 +10,7 @@ In the realm of space exploration, navigation, and astrophysics, precise and eff
 
 + Loading SPK, BPC, PCK, FK, and TPC files.
 + High-precision translations, rotations, and their combination (rigid body transformations).
++ Querying SPICE files in parallel at incredible speeds (~ 125,000 queries per second), search for the functions with the `many` keyword
 + Comprehensive time system conversions using the hifitime library (including TT, TAI, ET, TDB, UTC, GPS time, and more).
 
 ANISE stands validated against the traditional SPICE toolkit, ensuring accuracy and reliability, with translations achieving machine precision (2e-16) and rotations presenting minimal error (less than two arcseconds in the pointing of the rotation axis and less than one arcsecond in the angle about this rotation axis).

anise-py/anise.pyi

@@ -92,6 +92,15 @@ The obstructing body _should_ be a tri-axial ellipsoid body, e.g. IAU_MOON_FRAME
 6. Compute the elevation, and ensure it is between +/- 180 degrees.
 7. Compute the azimuth with a quadrant check, and ensure it is between 0 and 360 degrees."""
 
+    def azimuth_elevation_range_sez_many(self, rx_tx_states: List[Orbit], obstructing_body: Frame=None, ab_corr: Aberration=None) -> List[AzElRange]:
+        """Computes the azimuth (in degrees), elevation (in degrees), and range (in kilometers) of the
+receiver states (first item in tuple) seen from the transmitter state (second item in states tuple), once converted into the SEZ frame of the transmitter.
+
+Note: if any computation fails, the error will be printed to the stderr.
+Note: the output AER will be chronologically sorted, regardless of transmitter.
+
+Refer to [azimuth_elevation_range_sez] for details."""
+
     def beta_angle_deg(self, state: Orbit, ab_corr: Aberration=None) -> float:
         """Computes the Beta angle (β) for a given orbital state, in degrees. A Beta angle of 0° indicates that the orbit plane is edge-on to the Sun, leading to maximum eclipse time. Conversely, a Beta angle of +90° or -90° means the orbit plane is face-on to the Sun, resulting in continuous sunlight exposure and no eclipses.
 
@@ -190,6 +199,14 @@ This function performs a recursion of no more than twice the MAX_TREE_DEPTH."""
 This function calls `occultation` where the back object is the Sun in the J2000 frame, and the front object
 is the provided eclipsing frame."""
 
+    def solar_eclipsing_many(self, eclipsing_frame: Frame, observers: List[Orbit], ab_corr: Aberration=None) -> List[Occultation]:
+        """Computes the solar eclipsing of all the observers due to the eclipsing_frame, computed in parallel under the hood.
+
+Note: if any computation fails, the error will be printed to the stderr.
+Note: the output AER will be chronologically sorted, regardless of transmitter.
+
+Refer to [solar_eclipsing] for details."""
+
     def spk_domain(self, id: int) -> typing.Tuple:
         """Returns the applicable domain of the request id, i.e. start and end epoch that the provided id has loaded data."""
 
@@ -269,10 +286,21 @@ will return exactly the same data as the spkerz SPICE call.
 # Note
 The units will be those of the underlying ephemeris data (typically km and km/s)"""
 
-    def transform_to(self, state: Orbit, observer_frame: Frame, ab_corr: Aberration=None) -> Orbit:
-        """Translates a state with its origin (`to_frame`) and given its units (distance_unit, time_unit), returns that state with respect to the requested frame
+    def transform_many(self, target_frame: Orbit, observer_frame: Frame, time_series: TimeSeries, ab_corr: Aberration=None) -> List[Orbit]:
+        """Returns a chronologically sorted list of the Cartesian states that transform the `from_frame` to the `to_frame` for each epoch of the time series, computed in parallel under the hood.
+Note: if any transformation fails, the error will be printed to the stderr.
+
+Refer to [transform] for details."""
 
-**WARNING:** This function only performs the translation and no rotation _whatsoever_. Use the `transform_state_to` function instead to include rotations."""
+    def transform_many_to(self, states: List[Orbit], observer_frame: Frame, ab_corr: Aberration=None) -> List[Orbit]:
+        """Returns a chronologically sorted list of the provided states as seen from the observer frame, given the aberration.
+Note: if any transformation fails, the error will be printed to the stderr.
+Note: the input ordering is lost: the output states will not be in the same order as the input states if these are not chronologically sorted!
+
+Refer to [transform_to] for details."""
+
+    def transform_to(self, state: Orbit, observer_frame: Frame, ab_corr: Aberration=None) -> Orbit:
+        """Returns the provided state as seen from the observer frame, given the aberration."""
 
     def translate(self, target_frame: Orbit, observer_frame: Frame, epoch: Epoch, ab_corr: Aberration=None) -> Orbit:
         """Returns the Cartesian state of the target frame as seen from the observer frame at the provided epoch, and optionally given the aberration correction.
@@ -2409,26 +2437,26 @@ class rotation:
         def transpose(self) -> DCM:
             """Returns the transpose of this DCM"""
 
-        def __eq__(self, value: typing.Any) -> bool:
-            """Return self==value."""
+    def __eq__(self, value: typing.Any) -> bool:
+        """Return self==value."""
 
-        def __ge__(self, value: typing.Any) -> bool:
-            """Return self>=value."""
+    def __ge__(self, value: typing.Any) -> bool:
+        """Return self>=value."""
 
-        def __gt__(self, value: typing.Any) -> bool:
-            """Return self>value."""
+    def __gt__(self, value: typing.Any) -> bool:
+        """Return self>value."""
 
-        def __le__(self, value: typing.Any) -> bool:
-            """Return self<=value."""
+    def __le__(self, value: typing.Any) -> bool:
+        """Return self<=value."""
 
-        def __lt__(self, value: typing.Any) -> bool:
-            """Return self<value."""
+    def __lt__(self, value: typing.Any) -> bool:
+        """Return self<value."""
 
-        def __ne__(self, value: typing.Any) -> bool:
-            """Return self!=value."""
+    def __ne__(self, value: typing.Any) -> bool:
+        """Return self!=value."""
 
-        def __repr__(self) -> str:
-            """Return repr(self)."""
+    def __repr__(self) -> str:
+        """Return repr(self)."""
 
-        def __str__(self) -> str:
-            """Return str(self)."""
+    def __str__(self) -> str:
+        """Return str(self)."""

anise-py/tests/test_almanac.py

@@ -6,7 +6,7 @@
 from anise.astro import *
 from anise.astro.constants import Frames
 from anise.rotation import DCM
-from anise.time import Epoch
+from anise.time import Duration, Epoch, TimeSeries, Unit
 from anise.utils import convert_tpc
 
 from os import environ
@@ -27,7 +27,7 @@ def test_state_transformation():
         print(meta)
         # Process the files to be loaded
         try:
-            ctx = meta.process()
+            almanac = meta.process()
         except Exception as e:
             if "lfs" in str(e):
                 # Must be some LFS error in the CI again
@@ -36,14 +36,14 @@ def test_state_transformation():
     else:
         data_path = Path(__file__).parent.joinpath("..", "..", "data")
         # Must ensure that the path is a string
-        ctx = Almanac(str(data_path.joinpath("de440s.bsp")))
+        almanac = Almanac(str(data_path.joinpath("de440s.bsp")))
         # Let's add another file here -- note that the Almanac will load into a NEW variable, so we must overwrite it!
         # This prevents memory leaks (yes, I promise)
-        ctx = ctx.load(str(data_path.joinpath("pck08.pca"))).load(
+        almanac = almanac.load(str(data_path.joinpath("pck08.pca"))).load(
             str(data_path.joinpath("earth_latest_high_prec.bpc"))
         )
 
-    eme2k = ctx.frame_info(Frames.EME2000)
+    eme2k = almanac.frame_info(Frames.EME2000)
     assert eme2k.mu_km3_s2() == 398600.435436096
     assert eme2k.shape.polar_radius_km == 6356.75
     assert abs(eme2k.shape.flattening() - 0.0033536422844278) < 2e-16
@@ -94,7 +94,7 @@ def test_state_transformation():
     # In Python, we can set the aberration to None
     aberration = None
 
-    state_itrf93 = ctx.transform_to(orig_state, Frames.EARTH_ITRF93, aberration)
+    state_itrf93 = almanac.transform_to(orig_state, Frames.EARTH_ITRF93, aberration)
 
     print(orig_state)
     print(state_itrf93)
@@ -104,7 +104,7 @@ def test_state_transformation():
     assert abs(state_itrf93.height_km() - 1814.503598063825) < 1e-10
 
     # Convert back
-    from_state_itrf93_to_eme2k = ctx.transform_to(state_itrf93, Frames.EARTH_J2000)
+    from_state_itrf93_to_eme2k = almanac.transform_to(state_itrf93, Frames.EARTH_J2000)
 
     print(from_state_itrf93_to_eme2k)
 
@@ -113,7 +113,7 @@ def test_state_transformation():
     # Demo creation of a ground station
     mean_earth_angular_velocity_deg_s = 0.004178079012116429
     # Grab the loaded frame info
-    itrf93 = ctx.frame_info(Frames.EARTH_ITRF93)
+    itrf93 = almanac.frame_info(Frames.EARTH_ITRF93)
     paris = Orbit.from_latlongalt(
         48.8566,
         2.3522,
@@ -147,7 +147,30 @@ def test_state_transformation():
         "line_of_sight_obstructed",
         "azimuth_elevation_range_sez",
     ]:
-        assert hasattr(ctx, fname)
+        assert hasattr(almanac, fname)
+
+    # Test the parallel function calls
+    start = Epoch("2021-10-29 12:34:56 TDB")
+    stop = Epoch("2022-10-29 12:34:56 TDB")
+    time_series = TimeSeries(
+        start,
+        stop,
+        Duration("1 min"),
+        False,
+    )
+
+    tick = Epoch.system_now()
+
+    states = almanac.transform_many(
+        Frames.EARTH_J2000,
+        Frames.SUN_J2000,
+        time_series,
+        None,
+    )
+
+    clock_time = Epoch.system_now().timedelta(tick)
+    print(f"Queried {len(states)} states in {clock_time}")
+    assert len(states) == int(stop.timedelta(start).to_unit(Unit.Minute))
 
 
 def test_convert_tpc():
@@ -213,5 +236,5 @@ def test_frame_defs():
     # test_meta_load()
     # test_exports()
     # test_frame_defs()
-    test_convert_tpc()
+    # test_convert_tpc()
     test_state_transformation()

anise/Cargo.toml

@@ -43,6 +43,7 @@ rust-embed = { version = "8.4.0", features = [
     "include-exclude",
 ], optional = true }
 regex = { version = "1.10.5", optional = true }
+rayon = { workspace = true, optional = true } # Only used when building with Python
 
 [dev-dependencies]
 rust-spice = "0.7.6"
@@ -66,7 +67,7 @@ ureq = { version = "3.0.10", default-features = false, optional = true, features
 
 [features]
 default = ["metaload"]
-python = ["pyo3", "pyo3-log", "numpy", "ndarray"]
+python = ["pyo3", "pyo3-log", "numpy", "ndarray", "rayon"]
 metaload = ["url", "ureq", "platform-dirs", "regex", "serde_dhall"]
 embed_ephem = ["rust-embed", "ureq"]
 # Enabling this flag significantly increases compilation times due to Arrow and Polars.