[WIP] ops refactor

docs/guides/operators.md

@@ -20,7 +20,7 @@ import jax.numpy as jnp
 Galilean operators represent the basic transformations in classical mechanics:
 translations, rotations, and boosts.
 
-### GalileanSpatialTranslation
+### GalileanOp
 
 Translates position vectors by a fixed offset:
 
@@ -31,31 +31,31 @@ Translates position vectors by a fixed offset:
 
 ```{code-block} python
 >>> q = cxv.CartesianPos3D.from_([1, 2, 3], "kpc")
->>> op = cxo.GalileanSpatialTranslation.from_([10, 10, 10], "kpc")
+>>> op = cxo.GalileanOp.from_([10, 10, 10], "kpc")
 >>> op(q)
 CartesianPos3D(
     x=Quantity(11, unit='kpc'), y=Quantity(12, unit='kpc'), z=Quantity(13, unit='kpc')
 )
 ```
 
-### GalileanBoost
+### Galilean Boost
 
 Applies a velocity boost to a velocity vector:
 
 ```{code-block} python
->>> boost = cxo.GalileanBoost.from_([1, 1, 1], "km/s")
+>>> boost = cxo.Add.from_([1, 1, 1], "km/s")
 >>> boost(u.Quantity(1.0, "s"), q)[1]
 CartesianPos3D(
     x=Quantity(1., unit='kpc'), y=Quantity(2., unit='kpc'), z=Quantity(3., unit='kpc')
 )
 ```
 
-### GalileanRotation
+### Rotate
 
 Rotates vectors in space:
 
 ```{code-block} python
->>> rot = cxo.GalileanRotation.from_euler("z", u.Quantity(90, "deg"))
+>>> rot = cxo.Rotate.from_euler("z", u.Quantity(90, "deg"))
 >>> rot(q).round(2)
 CartesianPos3D(
     x=Quantity(-2., unit='kpc'),
@@ -72,8 +72,8 @@ Operators can be composed using the {class}`~coordinax.ops.Pipe` class or the
 `|` operator:
 
 ```{code-block} python
->>> op1 = cxo.GalileanSpatialTranslation.from_([1, 0, 0], "kpc")
->>> op2 = cxo.GalileanRotation.from_euler("z", u.Quantity(90, "deg"))
+>>> op1 = cxo.GalileanOp.from_([1, 0, 0], "kpc")
+>>> op2 = cxo.Rotate.from_euler("z", u.Quantity(90, "deg"))
 >>> pipe = cxo.Pipe([op1, op2])
 >>> pipe(q).round(2)
 CartesianPos3D(
@@ -97,17 +97,14 @@ CartesianPos3D(
 
 - {class}`~coordinax.ops.Identity`: The do-nothing operator, useful for generic
   code.
-- {class}`~coordinax.ops.AbstractCompositeOperator`: Base for building custom
+- {class}`~coordinax.ops.Pipe`: Base for building custom
   operator pipelines.
 
 ---
 
 ## Utilities and Advanced Usage
 
-- {class}`~coordinax.ops.simplify_op`: Simplifies composed operators when
-  possible.
-- {class}`~coordinax.ops.convert_to_pipe_operators`: Utility to convert a list
-  of operators into a {class}`~coordinax.ops.Pipe`.
+- {class}`~coordinax.ops.simplify`: Simplifies composed operators when possible.
 
 ---
 

docs/index.md

@@ -315,7 +315,7 @@ set of vector operations that work seamlessly with all `coordinax` vector types.
 ```{code-block} python
 >>> import coordinax.ops as cxo
 
->>> op = cxo.GalileanSpatialTranslation.from_([10, 10, 10], "kpc")
+>>> op = cxo.GalileanOp.from_([10, 10, 10], "kpc")
 
 >>> print(op(q))
 <CartesianPos3D: (x, y, z) [kpc]

pyproject.toml

@@ -40,7 +40,7 @@
     "quax-blocks>=0.4.0",
     "quaxed>=0.10.2",
     "typing-extensions>=4.13.2",
-    "unxt>=1.7.3",
+    "unxt>=1.7.6",
     "wadler-lindig>=0.1.6",
     "xmmutablemap>=0.1",
 ]

src/coordinax/_coordinax_space_frames/frame_transforms.py

@@ -16,14 +16,7 @@
 from .galactocentric import Galactocentric
 from .icrs import ICRS
 from coordinax._src.distances import Distance
-from coordinax._src.operators import (
-    GalileanRotation,
-    GalileanSpatialTranslation,
-    Identity,
-    Pipe,
-    VelocityBoost,
-    simplify_op,
-)
+from coordinax._src.operators import Add, Identity, Pipe, Rotate, simplify
 
 ScalarAngle: TypeAlias = Shaped[u.Quantity["angle"] | u.Angle, ""]
 RotationMatrix: TypeAlias = Shaped[Array, "3 3"]
@@ -54,7 +47,7 @@ def frame_transform_op(
 
     >>> @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,15 +58,15 @@ def frame_transform_op(
     >>> op = cx.frames.frame_transform_op(my_frame, gcf_frame)
     >>> op
     Pipe((
-        GalileanRotation(rotation=f32[3,3]),
+        Rotate(rotation=f32[3,3]),
         ...
     ))
 
     """  # noqa: E501
     fromframe_to_icrs = frame_transform_op(from_frame, _icrs_frame)
     icrs_to_toframe = frame_transform_op(_icrs_frame, to_frame)
     pipe = fromframe_to_icrs | icrs_to_toframe
-    return simplify_op(pipe)
+    return simplify(pipe)
 
 
 # ---------------------------------------------------------------
@@ -116,22 +109,22 @@ def frame_transform_op(from_frame: Galactocentric, to_frame: Galactocentric, /)
     >>> 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]),
-        VelocityBoost(CartesianVel3D( ... ))
+        Add(CartesianVel3D( ... )),
+        Rotate(rotation=f32[3,3]),
+        Add(CartesianPos3D( ... )),
+        Rotate(rotation=f32[3,3]),
+        Rotate(rotation=f32[3,3]),
+        Add(CartesianPos3D( ... )),
+        Rotate(rotation=f32[3,3]),
+        Add(CartesianVel3D( ... ))
     ))
 
     """
     if from_frame == to_frame:
         return Pipe((Identity(),))
 
     # TODO: not go through ICRS for the self-transformation
-    return simplify_op(
+    return simplify(
         frame_transform_op(from_frame, ICRS()) | frame_transform_op(ICRS(), to_frame)
     )
 
@@ -193,10 +186,10 @@ def frame_transform_op(from_frame: ICRS, to_frame: Galactocentric, /) -> Pipe:
     >>> 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]),
-        VelocityBoost(CartesianVel3D( ... ))
+        Rotate(rotation=f32[3,3]),
+        Add(CartesianPos3D( ... )),
+        Rotate(rotation=f32[3,3]),
+        Add(CartesianVel3D( ... ))
     ))
 
     Apply the transformation:
@@ -238,25 +231,23 @@ def frame_transform_op(from_frame: ICRS, to_frame: Galactocentric, /) -> Pipe:
 
     """  # 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 = Add(-to_frame.galcen.distance * jnp.asarray([1, 0, 0]))
 
     # Post-rotation velocity offset
-    offset_v = VelocityBoost(to_frame.galcen_v_sun)
+    offset_v = Add(to_frame.galcen_v_sun)
 
     # Total Operator
     return R | offset_q | H | offset_v

src/coordinax/_src/frames/coordinate.py

@@ -276,7 +276,7 @@ class Coordinate(AbstractCoordinate):
     ...     data=space,
     ...     frame=cx.frames.TransformedReferenceFrame(
     ...         cx.frames.Galactocentric(),
-    ...         cx.ops.GalileanSpatialTranslation.from_([20, 0, 0], "kpc"),
+    ...         cx.ops.GalileanOp.from_([20, 0, 0], "kpc"),
     ...     ),
     ... )
 

src/coordinax/_src/frames/example.py

@@ -11,13 +11,7 @@
 
 from .base import AbstractReferenceFrame
 from coordinax._src.frames import api
-from coordinax._src.operators import (
-    GalileanBoost,
-    GalileanRotation,
-    GalileanSpatialTranslation,
-    Identity,
-    Pipe,
-)
+from coordinax._src.operators import Add, Identity, Pipe, Rotate
 
 
 @final
@@ -34,9 +28,9 @@ class Alice(AbstractReferenceFrame):
     >>> op = cxf.frame_transform_op(alice, bob)
     >>> print(op)
     Pipe((
-      GalileanSpatialTranslation(<CartesianPos3D: (x, y, z) [km]
+      Add(<CartesianPos3D: (x, y, z) [km]
           [100000  10000      0]>),
-      GalileanBoost(<CartesianVel3D: (x, y, z) [m / s]
+      Add(<CartesianVel3D: (x, y, z) [m / s]
           [2.698e+08 0.000e+00 0.000e+00]>)
     ))
 
@@ -62,9 +56,9 @@ class Bob(AbstractReferenceFrame):
     >>> op = cxf.frame_transform_op(alice, bob)
     >>> print(op)
     Pipe((
-      GalileanSpatialTranslation(<CartesianPos3D: (x, y, z) [km]
+      Add(<CartesianPos3D: (x, y, z) [km]
           [100000  10000      0]>),
-      GalileanBoost(<CartesianVel3D: (x, y, z) [m / s]
+      Add(<CartesianVel3D: (x, y, z) [m / s]
           [2.698e+08 0.000e+00 0.000e+00]>)
     ))
 
@@ -117,9 +111,9 @@ def frame_transform_op(from_frame: Alice, to_frame: FriendOfAlice, /) -> Pipe:
     >>> op = cx.frames.frame_transform_op(alice, friend)
     >>> print(op)
     Pipe((
-      GalileanSpatialTranslation(<CartesianPos3D: (x, y, z) [m]
+      Add(<CartesianPos3D: (x, y, z) [m]
           [10  0  0]>),
-      GalileanRotation([[ 0.         -0.99999994  0.        ]
+      Rotate([[ 0.         -0.99999994  0.        ]
                         [ 0.99999994  0.          0.        ]
                         [ 0.          0.          0.99999994]])
     ))
@@ -133,8 +127,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 = Add.from_([10, 0, 0], "m")
+    rotation = Rotate.from_euler("Z", u.Quantity(90, "deg"))
     return shift | rotation
 
 
@@ -154,9 +148,9 @@ def frame_transform_op(from_frame: Alice, to_frame: Bob, /) -> Pipe:
     >>> op = cxf.frame_transform_op(alice, bob)
     >>> print(op)
     Pipe((
-      GalileanSpatialTranslation(<CartesianPos3D: (x, y, z) [km]
+      Add(<CartesianPos3D: (x, y, z) [km]
           [100000  10000      0]>),
-      GalileanBoost(<CartesianVel3D: (x, y, z) [m / s]
+      Add(<CartesianVel3D: (x, y, z) [m / s]
           [2.698e+08 0.000e+00 0.000e+00]>)
     ))
 
@@ -175,8 +169,8 @@ def frame_transform_op(from_frame: Alice, to_frame: Bob, /) -> Pipe:
     ))
 
     """
-    shift = GalileanSpatialTranslation.from_([100_000, 10_000, 0], "km")
-    boost = GalileanBoost.from_([269_813_212.2, 0, 0], "m/s")
+    shift = Add.from_([100_000, 10_000, 0], "km")
+    boost = Add.from_([269_813_212.2, 0, 0], "m/s")
     return shift | boost
 
 
@@ -200,10 +194,10 @@ def frame_transform_op(
     >>> op = cx.frames.frame_transform_op(friend, alice)
     >>> print(op)
     Pipe((
-      GalileanRotation([[ 0.          0.99999994  0.        ]
+      Rotate([[ 0.          0.99999994  0.        ]
                         [-0.99999994  0.          0.        ]
                         [ 0.          0.          0.99999994]]),
-      GalileanSpatialTranslation(<CartesianPos3D: (x, y, z) [m]
+      Add(<CartesianPos3D: (x, y, z) [m]
           [-10   0   0]>)
     ))
 

src/coordinax/_src/frames/register_ops.py

@@ -3,14 +3,16 @@
 __all__: tuple[str, ...] = ()
 
 
+from plum import dispatch
+
 from dataclassish import replace
 
 from .coordinate import Coordinate
 from coordinax._src.operators import AbstractOperator
 
 
-@AbstractOperator.__call__.dispatch  # type: ignore[misc]
-def call(self: AbstractOperator, x: Coordinate, /) -> Coordinate:
+@dispatch
+def operate(self: AbstractOperator, obj: Coordinate, /) -> Coordinate:
     """Apply the operator to a coordinate.
 
     Examples
@@ -25,7 +27,7 @@ def call(self: AbstractOperator, x: Coordinate, /) -> Coordinate:
         frame=ICRS()
     )
 
-    >>> op = cx.ops.GalileanSpatialTranslation.from_([-1, -1, -1], "kpc")
+    >>> op = cx.ops.GalileanOp.from_([-1, -1, -1], "kpc")
 
     >>> new_coord = op(coord)
     >>> print(new_coord.data["length"])
@@ -34,4 +36,4 @@ def call(self: AbstractOperator, x: Coordinate, /) -> Coordinate:
 
     """
     # TODO: take the frame into account
-    return replace(x, data=self(x.data))
+    return replace(obj, data=self(obj.data))

src/coordinax/_src/frames/register_primitives.py

@@ -9,6 +9,7 @@
 from dataclassish import replace
 
 from .coordinate import Coordinate
+from coordinax._src.vectors.base_pos.core import AbstractPos
 
 
 @register(jax.lax.neg_p)
@@ -33,3 +34,26 @@ def neg_p_coord(x: Coordinate, /) -> Coordinate:
 
     """
     return replace(x, data=-x.data)
+
+
+@register(jax.lax.add_p)
+def add_p_coord_pos(x: Coordinate, y: AbstractPos, /) -> Coordinate:
+    r"""Add a position vector to a coordinate.
+
+    To understand this operation, let's consider a phase-space point $(x, v) \in
+    \mathbb{R}^3\times\mathbb{R}^3$ consisting of a position and a velocity. A
+    pure spatial translation is the map $T_{\Delta x} : (x,v) \mapsto (x+\Delta
+    x,\ v)$, i.e. only the position is shifted; velocity is unchanged.
+
+    """
+    # Get the Cartesian class for the coordinate's position
+    cart_cls = y.cartesian_type
+    # Convert the coordinate to that class. This changes the position, but also
+    # the other components, e.g. the velocity.
+    data = dict(x.data.vconvert(cart_cls))
+    # Now add the position vector to the position component only
+    data = replace(data, length=data["length"] + y)
+    # Transform back to the original vector types
+    # data.vconvert()  # TODO: all original types
+    # Reconstruct the Coordinate
+    return Coordinate(data, frame=x.frame)