[WIP] Galilean operators refactor

src/coordinax/_coordinax_space_frames/frame_transforms.py

@@ -17,10 +17,10 @@
 from .icrs import ICRS
 from coordinax._src.distances import Distance
 from coordinax._src.operators import (
-    GalileanRotation,
-    GalileanSpatialTranslation,
     Identity,
     Pipe,
+    Rotate,
+    Translate,
     VelocityBoost,
     simplify_op,
 )
@@ -54,7 +54,7 @@
 
     >>> @dispatch
     ... def frame_transform_op(from_frame: MySpaceFrame, to_frame: ICRS, /) -> cx.ops.AbstractOperator:
-    ...     return cx.ops.GalileanRotation.from_euler("z", u.Quantity(10, "deg"))
+    ...     return cx.ops.Rotate.from_euler("z", u.Quantity(10, "deg"))
 
     We can transform from `MySpaceFrame` to a Galacocentric frame, even though
     we don't have a direct transformation defined:
@@ -65,7 +65,7 @@
     >>> op = cx.frames.frame_transform_op(my_frame, gcf_frame)
     >>> op
     Pipe((
-        GalileanRotation(rotation=f32[3,3]),
+        Rotate(rotation=f32[3,3]),
         ...
     ))
 
@@ -115,14 +115,14 @@
     >>> gcf_frame2 = cxf.Galactocentric(roll=u.Quantity(10, "deg"))
     >>> frame_op2 = cxf.frame_transform_op(gcf_frame, gcf_frame2)
     >>> frame_op2
     Pipe((
         VelocityBoost(CartesianVel3D( ... )),
-        GalileanRotation(rotation=f32[3,3]),
-        GalileanSpatialTranslation(CartesianPos3D( ... )),
-        GalileanRotation(rotation=f32[3,3]),
-        GalileanRotation(rotation=f32[3,3]),
-        GalileanSpatialTranslation(CartesianPos3D( ... )),
-        GalileanRotation(rotation=f32[3,3]),
+        Rotate(rotation=f32[3,3]),
+        Translate(CartesianPos3D( ... )),
+        Rotate(rotation=f32[3,3]),
+        Rotate(rotation=f32[3,3]),
+        Translate(CartesianPos3D( ... )),
+        Rotate(rotation=f32[3,3]),
         VelocityBoost(CartesianVel3D( ... ))
     ))
 
@@ -192,23 +192,23 @@
 
     >>> frame_op = cx.frames.frame_transform_op(icrs_frame, gcf_frame)
     >>> frame_op
     Pipe((
-        GalileanRotation(rotation=f32[3,3]),
-        GalileanSpatialTranslation(CartesianPos3D( ... )),
-        GalileanRotation(rotation=f32[3,3]),
+        Rotate(rotation=f32[3,3]),
+        Translate(CartesianPos3D( ... )),
+        Rotate(rotation=f32[3,3]),
         VelocityBoost(CartesianVel3D( ... ))
     ))
 
     Apply the transformation:
 
     >>> vega_gcf_q, vega_gcf_p = frame_op(vega_q, vega_p)
     >>> vega_gcf_q = vega_gcf_q.vconvert(cx.vecs.CartesianPos3D)
     >>> vega_gcf_p = vega_gcf_p.vconvert(cx.vecs.CartesianVel3D, vega_gcf_q)
     >>> print(vega_gcf_q)
     <CartesianPos3D: (x, y, z) [pc]
         [-8112.898    21.799    29.01...]>
     >>> print(vega_gcf_p.uconvert({u.dimension("speed"): "km/s"}))
     <CartesianVel3D: (x, y, z) [km / s]
         [ 34.067 234.616 -28.76 ]>
 
     It matches!
@@ -222,15 +222,15 @@
     >>> p = u.Quantity([0., 0, 0], "km/s")
 
     >>> newq, newp = frame_op(q, p)
     >>> newq, newp
     (Quantity(Array([-8121.972,     0.   ,    20.8  ], dtype=float32), unit='pc'),
      Quantity(Array([ 12.9 , 245.6 ,   7.78], dtype=float32), unit='km / s'))
 
     >>> q = cx.CartesianPos3D.from_([0, 0, 0], "pc")
     >>> p = cx.CartesianVel3D.from_([0, 0, 0], "km/s")
 
     >>> newq, newp = frame_op(q, p)
     >>> print(newq, newp, sep="\n")
     <CartesianPos3D: (x, y, z) [pc]
         [-8121.972     0.       20.8  ]>
     <CartesianVel3D: (x, y, z) [km / s]
@@ -238,22 +238,20 @@
 
     """  # noqa: E501
     # rotation matrix to align x(ICRS) with the vector to the Galactic center
-    rot_lat = GalileanRotation.from_euler("y", to_frame.galcen.lat)
-    rot_lon = GalileanRotation.from_euler("z", -to_frame.galcen.lon)
+    rot_lat = Rotate.from_euler("y", to_frame.galcen.lat)
+    rot_lon = Rotate.from_euler("z", -to_frame.galcen.lon)
     # extra roll away from the Galactic x-z plane
-    roll = GalileanRotation.from_euler("x", to_frame.roll - to_frame.roll0)
+    roll = Rotate.from_euler("x", to_frame.roll - to_frame.roll0)
     # construct transformation matrix
     R = (roll @ rot_lat @ rot_lon).simplify()
 
     # Translation by Sun-Galactic center distance around x' and rotate about y'
     # to account for tilt due to Sun's height above the plane
     z_d = u.ustrip("", to_frame.z_sun / to_frame.galcen.distance)  # [radian]
-    H = GalileanRotation.from_euler("y", u.Quantity(jnp.asin(z_d), "rad"))
+    H = Rotate.from_euler("y", u.Quantity(jnp.asin(z_d), "rad"))
 
     # Post-rotation spatial offset to Galactic center.
-    offset_q = GalileanSpatialTranslation(
-        -to_frame.galcen.distance * jnp.asarray([1, 0, 0])
-    )
+    offset_q = Translate(-to_frame.galcen.distance * jnp.asarray([1, 0, 0]))
 
     # Post-rotation velocity offset
     offset_v = VelocityBoost(to_frame.galcen_v_sun)

src/coordinax/_src/frames/example.py

@@ -13,10 +13,10 @@
 from coordinax._src.frames import api
 from coordinax._src.operators import (
     GalileanBoost,
-    GalileanRotation,
-    GalileanSpatialTranslation,
     Identity,
     Pipe,
+    Rotate,
+    Translate,
 )
 
 
@@ -133,8 +133,8 @@ def frame_transform_op(from_frame: Alice, to_frame: FriendOfAlice, /) -> Pipe:
      ))
 
     """
-    shift = GalileanSpatialTranslation.from_([10, 0, 0], "m")
-    rotation = GalileanRotation.from_euler("Z", u.Quantity(90, "deg"))
+    shift = Translate.from_([10, 0, 0], "m")
+    rotation = Rotate.from_euler("Z", u.Quantity(90, "deg"))
     return shift | rotation
 
 
@@ -175,7 +175,7 @@ def frame_transform_op(from_frame: Alice, to_frame: Bob, /) -> Pipe:
     ))
 
     """
-    shift = GalileanSpatialTranslation.from_([100_000, 10_000, 0], "km")
+    shift = Translate.from_([100_000, 10_000, 0], "km")
     boost = GalileanBoost.from_([269_813_212.2, 0, 0], "m/s")
     return shift | boost
 

src/coordinax/_src/operators/__init__.py

@@ -1,5 +1,4 @@
 """Coordinax Operator package."""
-# ruff: noqa: F401
 
 __all__ = [
     "simplify_op",
@@ -12,10 +11,9 @@
     "AbstractGalileanOperator",
     "GalileanBoost",
     "GalileanOperator",
-    "GalileanRotation",
-    "GalileanSpatialTranslation",
-    "GalileanTranslation",
     # Misc
+    "Rotate",
+    "Translate",
     "VelocityBoost",
     # Utils
     "convert_to_pipe_operators",
@@ -28,11 +26,10 @@
 from .galilean.base import AbstractGalileanOperator
 from .galilean.boost import GalileanBoost
 from .galilean.composite import GalileanOperator
-from .galilean.rotation import GalileanRotation
-from .galilean.spatial_translation import GalileanSpatialTranslation
-from .galilean.translation import GalileanTranslation
 from .identity import Identity
 from .pipe import Pipe, convert_to_pipe_operators
+from .rotate import Rotate
+from .translate import Translate
 
 # isort: split
 from . import compat, register_simplify

src/coordinax/_src/operators/base.py

@@ -98,11 +98,10 @@ def simplify(self) -> "AbstractOperator":
         >>> op.simplify()
         Identity()
 
-        >>> op = cxo.GalileanOperator(translation=u.Quantity([0., 2., 3., 4.], "km"))
+        >>> op = cxo.GalileanOperator(translation=u.Quantity([1., 2., 3.], "km"))
         >>> op.simplify()
-        GalileanTranslation(
-            delta_t=Quantity(f32[], unit='s'),
-            delta_q=CartesianPos3D( ... )
+        Translate(
+            delta=CartesianPos3D( ... )
         )
 
         """
@@ -147,27 +146,24 @@ def __str__(self) -> str:
         >>> import jax.numpy as jnp
         >>> import coordinax as cx
 
-        >>> op = cx.ops.GalileanRotation([[1, 0, 0], [0, 1, 0], [0, 0, 1]])
+        >>> op = cx.ops.Rotate([[1, 0, 0], [0, 1, 0], [0, 0, 1]])
         >>> print(op)
-        GalileanRotation([[1 0 0]
+        Rotate([[1 0 0]
                           [0 1 0]
                           [0 0 1]])
 
         >>> op = cx.ops.GalileanOperator(
-        ...     translation=u.Quantity([0., 2, 3, 4], "km"),
+        ...     translation=u.Quantity([2, 3, 4], "km"),
         ...     velocity=u.Quantity([1., 2, 3], "km/s"),
         ...     rotation=jnp.eye(3).at[0, 2].set(1),
         ... )
         >>> print(op)
         GalileanOperator(
-            rotation=GalileanRotation([[1. 0. 1.]
+            rotation=Rotate([[1. 0. 1.]
                                        [0. 1. 0.]
                                        [0. 0. 1.]]),
-            translation=GalileanTranslation(
-                delta_t=Quantity(0., unit='s'),
-                delta_q=<CartesianPos3D: (x, y, z) [km]
-                    [2. 3. 4.]>
-            ),
+            translation=Translate(<CartesianPos3D: (x, y, z) [km]
+                                    [2 3 4]>),
             velocity=GalileanBoost(<CartesianVel3D: (x, y, z) [km / s]
                     [1. 2. 3.]>)
         )
@@ -187,15 +183,15 @@ def __or__(self, other: "AbstractOperator") -> "coordinax.ops.Pipe":
         --------
         >>> import coordinax.ops as cxo
 
-        >>> op1 = cxo.GalileanRotation([[1, 0, 0], [0, 1, 0], [0, 0, 1]])
+        >>> op1 = cxo.Rotate([[1, 0, 0], [0, 1, 0], [0, 0, 1]])
         >>> op2 = cxo.Identity()
         >>> op3 = op1 | op2
         >>> op3
-        Pipe((GalileanRotation(rotation=i32[3,3]), Identity()))
+        Pipe((Rotate(rotation=i32[3,3]), Identity()))
 
         >>> op4 = cxo.Identity() | op3
         >>> op4
-        Pipe((Identity(), GalileanRotation(rotation=i32[3,3]), Identity()))
+        Pipe((Identity(), Rotate(rotation=i32[3,3]), Identity()))
 
         """
         from .pipe import Pipe
@@ -251,18 +247,15 @@ def from_(
     --------
     >>> import coordinax.ops as cxo
 
-    >>> op = cxo.GalileanSpatialTranslation.from_([1, 1, 1], "km")
+    >>> op = cxo.Translate.from_([1, 1, 1], "km")
     >>> print(op)
-    GalileanSpatialTranslation(<CartesianPos3D: (x, y, z) [km]
+    Translate(<CartesianPos3D: (x, y, z) [km]
         [1 1 1]>)
 
-    >>> op = cxo.GalileanTranslation.from_([3e5, 1, 1, 1], "km")
+    >>> op = cxo.Translate.from_([1, 1, 1], "km")
     >>> print(op)
-    GalileanTranslation(
-      delta_t=Quantity(1.0006922, unit='s'),
-      delta_q=<CartesianPos3D: (x, y, z) [km]
-          [1. 1. 1.]>
-    )
+    Translate(<CartesianPos3D: (x, y, z) [km]
+                [1 1 1]>)
 
     >>> op = cxo.GalileanBoost.from_([1, 1, 1], "km/s")
     >>> print(op)