Add unit tests for ITRF93 transformations

moeyensj · moeyensj · commit e28438f51018 · 2025-01-24T13:42:42.000-08:00
diff --git a/src/adam_core/coordinates/tests/test_transforms_rotation.py b/src/adam_core/coordinates/tests/test_transforms_rotation.py
@@ -1,6 +1,11 @@
+import quivr as qv
+import spiceypy as sp
+
+from ...constants import KM_P_AU, S_P_DAY
 from ...time import Timestamp
+from ...utils.spice import get_perturber_state, setup_SPICE
 from ..cartesian import CartesianCoordinates
-from ..origin import Origin
+from ..origin import Origin, OriginCodes
 from ..transform import cartesian_to_frame, transform_coordinates
 from .test_transforms_translation import assert_coords_equal
 
@@ -121,3 +126,130 @@ def test_transform_coordinates_frame(orbital_elements, orbital_elements_equatori
     assert_coords_equal(
         cartesian_coordinates_equatorial_actual, cartesian_coordinates_equatorial
     )
+
+
+def test_transform_coordinates_to_itrf93():
+    """
+    Test that transform_coordinates correctly converts between ecliptic and
+    ITRF93 cartesian coordinates
+    """
+    setup_SPICE()
+
+    # Get perturber states for a few of the planets
+    times = Timestamp.from_mjd([59000, 59500, 60000], scale="tdb")
+    states = CartesianCoordinates.empty()
+    target_ids = []
+    for perturber in [
+        OriginCodes.MOON,
+        OriginCodes.VENUS,
+        OriginCodes.MARS_BARYCENTER,
+        OriginCodes.JUPITER_BARYCENTER,
+        OriginCodes.SATURN_BARYCENTER,
+        OriginCodes.URANUS_BARYCENTER,
+        OriginCodes.NEPTUNE_BARYCENTER,
+    ]:
+        states = qv.concatenate(
+            [
+                states,
+                get_perturber_state(
+                    perturber, times, origin=OriginCodes.SUN, frame="ecliptic"
+                ),
+            ]
+        )
+        target_ids.extend(perturber.name for _ in range(len(times)))
+
+    states_itrf93 = transform_coordinates(states, frame_out="itrf93")
+
+    # Repeat with SPICE
+    states_spice_itrf93 = CartesianCoordinates.empty()
+    for coord, target_id in zip(states_itrf93, target_ids):
+        # Rotate the coordinates to the ITRF93 frame using SPICE
+        state, lt = sp.spkezr(
+            target_id, coord.time.et()[0].as_py(), "ITRF93", "NONE", "SUN"
+        )
+        states_spice_itrf93 = qv.concatenate(
+            [
+                states_spice_itrf93,
+                CartesianCoordinates.from_kwargs(
+                    x=[state[0] / KM_P_AU],
+                    y=[state[1] / KM_P_AU],
+                    z=[state[2] / KM_P_AU],
+                    vx=[state[3] / KM_P_AU * S_P_DAY],
+                    vy=[state[4] / KM_P_AU * S_P_DAY],
+                    vz=[state[5] / KM_P_AU * S_P_DAY],
+                    time=coord.time,  # Add time from original coordinate
+                    frame="itrf93",  # Specify the frame
+                    origin=Origin.from_kwargs(code=["SUN"]),  # Specify the origin
+                ),
+            ]
+        )
+
+    # Test that the two coordinate sets are equal within tolerance
+    assert_coords_equal(
+        states_itrf93, states_spice_itrf93, position_tol_mm=10, velocity_tol_nm_s=500
+    )
+
+
+def test_transform_coordinates_from_itrf93():
+    """
+    Test that transform_coordinates correctly converts between ITRF93 and
+    ecliptic cartesian coordinates
+    """
+    setup_SPICE()
+
+    # Get perturber states for a few of the planets
+    times = Timestamp.from_mjd([59000, 59500, 60000], scale="tdb")
+    states = CartesianCoordinates.empty()
+    target_ids = []
+    for perturber in [
+        OriginCodes.MOON,
+        OriginCodes.VENUS,
+        OriginCodes.MARS_BARYCENTER,
+        OriginCodes.JUPITER_BARYCENTER,
+        OriginCodes.SATURN_BARYCENTER,
+        OriginCodes.URANUS_BARYCENTER,
+        OriginCodes.NEPTUNE_BARYCENTER,
+    ]:
+        states = qv.concatenate(
+            [
+                states,
+                get_perturber_state(
+                    perturber, times, origin=OriginCodes.SUN, frame="itrf93"
+                ),
+            ]
+        )
+        target_ids.extend(perturber.name for _ in range(len(times)))
+
+    states_ecliptic = transform_coordinates(states, frame_out="ecliptic")
+
+    # Repeat with SPICE
+    states_spice_ecliptic = CartesianCoordinates.empty()
+    for coord, target_id in zip(states_ecliptic, target_ids):
+        # Rotate the coordinates to the ecliptic frame using SPICE
+        state, lt = sp.spkezr(
+            target_id, coord.time.et()[0].as_py(), "ECLIPJ2000", "NONE", "SUN"
+        )
+        states_spice_ecliptic = qv.concatenate(
+            [
+                states_spice_ecliptic,
+                CartesianCoordinates.from_kwargs(
+                    x=[state[0] / KM_P_AU],
+                    y=[state[1] / KM_P_AU],
+                    z=[state[2] / KM_P_AU],
+                    vx=[state[3] / KM_P_AU * S_P_DAY],
+                    vy=[state[4] / KM_P_AU * S_P_DAY],
+                    vz=[state[5] / KM_P_AU * S_P_DAY],
+                    time=coord.time,  # Add time from original coordinate
+                    frame="ecliptic",  # Specify the frame
+                    origin=Origin.from_kwargs(code=["SUN"]),  # Specify the origin
+                ),
+            ]
+        )
+
+    # Test that the two coordinate sets are equal within tolerance
+    assert_coords_equal(
+        states_ecliptic,
+        states_spice_ecliptic,
+        position_tol_mm=10,
+        velocity_tol_nm_s=500,
+    )
diff --git a/src/adam_core/coordinates/tests/test_transforms_translation.py b/src/adam_core/coordinates/tests/test_transforms_translation.py
@@ -11,6 +11,8 @@
 def assert_coords_equal(
     have: CartesianCoordinates,
     want: CartesianCoordinates,
+    position_tol_mm: float = 10,
+    velocity_tol_nm_s: float = 10,
 ):
     diff = want.values - have.values
 
@@ -19,10 +21,10 @@ def assert_coords_equal(
     # Calculate offset in velocity in nm/s
     v_diff = np.linalg.norm(diff[:, 3:], axis=1) * (u.au / u.d).to(u.nm / u.s)
 
-    # Assert positions are to within 10 mm
-    np.testing.assert_array_less(r_diff, 10)
-    # Assert velocities are to within 10 nm/s
-    np.testing.assert_array_less(v_diff, 10)
+    # Assert positions are to within desired mm
+    np.testing.assert_array_less(r_diff, position_tol_mm)
+    # Assert velocities are to within desired nm/s
+    np.testing.assert_array_less(v_diff, velocity_tol_nm_s)
 
 
 def test_cartesian_to_origin(orbital_elements, orbital_elements_barycentric):
diff --git a/src/adam_core/coordinates/transform.py b/src/adam_core/coordinates/transform.py
@@ -1388,6 +1388,8 @@ def apply_time_varying_rotation(
     else:
         raise ValueError("Unsupported frame: {}".format(frame_in))
 
+    from ..constants import KM_P_AU, S_P_DAY
+
     # Loop over unique times and then rotate each coordinate
     coords_rotated = CartesianCoordinates.empty()
     indices = []
@@ -1405,24 +1407,35 @@ def apply_time_varying_rotation(
         # Store the indices so we can use to sort the coordinates later
         indices.extend(indices_time.to_pylist())
 
-        rotation_matrix_3x3 = sp.pxform(
+        # The units of the transformation matrix are km and km/s and while
+        # our states are in au and au/d. So we need to convert the transformation
+        # matrix to the correct units.
+        rotation_matrix_6x6_kms = sp.sxform(
             frame_spice_in, frame_spice_out, time.et().to_numpy(zero_copy_only=False)[0]
         )
-        # Create 6x6 rotation matrix for position and velocity
-        rotation_matrix_6x6 = np.zeros((6, 6))
-        rotation_matrix_6x6[:3, :3] = rotation_matrix_3x3
-        rotation_matrix_6x6[3:, 3:] = rotation_matrix_3x3
 
-        # Rotate the coordinates
-        coords_rotated_time = coords.apply_mask(time_mask).rotate(
-            rotation_matrix_6x6, frame_out
+        # Compute unit conversion matrix to convert from km to au and km/s to au/d
+        rotation_unit_conversion = np.zeros((6, 6))
+        rotation_unit_conversion[:3, :3] = np.identity(3) * KM_P_AU
+        rotation_unit_conversion[3:, 3:] = np.identity(3) * KM_P_AU / S_P_DAY
+
+        rotation_matrix_6x6_aud = (
+            np.linalg.inv(rotation_unit_conversion)
+            @ rotation_matrix_6x6_kms
+            @ rotation_unit_conversion
         )
 
+        # Rotate the coordinates
+        coords_time = coords.apply_mask(time_mask)
+        coords_rotated_time = coords_time.rotate(rotation_matrix_6x6_aud, frame_out)
+
         # Add the rotated coordinates to the total coordinates
         coords_rotated = qv.concatenate([coords_rotated, coords_rotated_time])
 
-    coords_rotated = coords_rotated.take(indices)
-
+    # Compute the indices that would sort the new coordinates
+    # to match the original coordinates
+    sorted_indices = np.argsort(indices)
+    coords_rotated = coords_rotated.take(sorted_indices)
     return coords_rotated
 
 
@@ -1506,7 +1519,7 @@ def transform_coordinates(
     if not isinstance(coords, CoordinatesClasses):
         raise TypeError("Unsupported coordinate type: {}".format(type(coords)))
 
-    if frame_out not in {None, "equatorial", "ecliptic"}:
+    if frame_out not in {None, "equatorial", "ecliptic", "itrf93"}:
         raise ValueError("Unsupported frame_out: {}".format(frame_out))
 
     if origin_out is not None and not isinstance(origin_out, OriginCodes):
