✨ feat(transforms): kinematic acts for Scale/Reflect/Shear without a base point

Only Boost and Rotate had TangentGeometry (velocity/acceleration) act paths; the
other linear transforms fell through to the generic jet prolongation, which
requires a base point `at` even for a constant Cartesian map where none is
needed (`act(Scale, None, velocity, cart3d, coord_vel)` raised TypeError).

A linear map's Jacobian is its (constant) matrix, so a Cartesian velocity or
acceleration transforms as `v -> M v` with no anchor. Register a shared
`pushforward` on `AbstractLinearTransform`: Cartesian charts apply `M` directly
(no `at`); non-Cartesian charts push through the chart Jacobian (which does need
`at`); time-dependent linear maps still route through the generic prolongation
for the `dot(M)` term. Rotate keeps its own (more specific) dispatch, so there
is no ambiguity. A clear TypeError now reports tangent/anchor components that do
not match the chart.

Verified: fast path equals the generic prolongation when `at` is supplied
(the suite's keystone property); Scale/Reflect/Shear velocity + acceleration;
non-Cartesian still requires `at`. New unit tests + doctest.

Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>

Author: nstarman

src/coordinax/transforms/_src/actions/linear.py

@@ -231,3 +231,118 @@ def _maybe(
         for i, (f, p) in enumerate(zip(chart.factors, parts, strict=True))
     )
     return chart.merge_components(mapped_parts)
+
+
+# ============================================================================
+# pushforward — tangent geometry (shared by every linear transform)
+
+
+def _linear_pushforward_cdict(
+    op: AbstractLinearTransform,
+    tau: Any,
+    x: CDict,
+    chart: cxc.AbstractChart,
+    rep: cxr.Representation,
+    /,
+    *,
+    at: CDict | None = None,
+    usys: OptUSys = None,
+) -> CDict:
+    r"""Frozen-$\tau$ pushforward of tangent data under a linear map: $v \mapsto M v$.
+
+    A linear map has a *constant* Jacobian equal to its matrix ``M``, so in the
+    canonical Cartesian chart the pushforward is simply ``M v`` and needs no
+    base point ``at``. For a non-Cartesian chart the tangent is pushed through
+    the chart Jacobian (which does require ``at``), ``M`` is applied in
+    Cartesian, then the result is pulled back to the original chart.
+
+    Examples
+    --------
+    Scale a Cartesian velocity vector (no base point needed):
+
+    >>> import quaxed.numpy as jnp
+    >>> import unxt as u
+    >>> import coordinax.charts as cxc
+    >>> import coordinax.representations as cxr
+    >>> import coordinax.transforms as cxfm
+
+    >>> op = cxfm.Scale.from_factors([2.0, 3.0, 1.0])
+    >>> v = {"x": u.Q(1.0, "m/s"), "y": u.Q(1.0, "m/s"), "z": u.Q(1.0, "m/s")}
+    >>> out = cxfm.act(op, None, v, cxc.cart3d, cxr.tangent_geom, cxr.coord_vel)
+    >>> jnp.stack([out[c].to_value("m/s") for c in ("x", "y", "z")]).round(3)
+    Array([2., 3., 1.], dtype=float64)
+
+    """
+    # The tangent must carry exactly the chart's components (a clear error
+    # instead of a raw KeyError when packing); an anchor, if given, must match.
+    expected = set(chart.components)
+    for name, d in (("tangent", x), *(() if at is None else (("base point", at),))):
+        if set(d) != expected:
+            missing = sorted(expected - set(d))
+            extra = sorted(set(d) - expected)
+            msg = (
+                f"pushforward({type(op).__name__}, ...): the {name} components "
+                f"do not match the chart's {sorted(expected)}"
+                + (f"; missing {missing}" if missing else "")
+                + (f"; unexpected {extra}" if extra else "")
+                + "."
+            )
+            raise TypeError(msg)
+
+    cart = chart.cartesian
+    matrix = op._matrix(cart, tau)
+
+    if chart is cart:
+        p_cart = x
+    else:
+        if at is None:
+            msg = (
+                f"pushforward({type(op).__name__}, ..., TangentGeometry) on a "
+                f"non-Cartesian chart ({chart!r}) requires 'at' (base point in "
+                "chart coords) so the Jacobian pushforward can be evaluated."
+            )
+            raise TypeError(msg)
+        at_cart = cxc.pt_map(at, chart, cart, usys=usys)
+        p_cart = cxr.tangent_map(x, chart, rep, cart, at=at, usys=usys)  # ty: ignore[missing-argument]
+
+    comps_cart = cart.components
+    v, unit = pack_uniform_unit(p_cart, keys=comps_cart)  # ty: ignore[no-matching-overload]
+    v_out = jnp.einsum("ij,...j->...i", matrix, v)
+    p_cart_out = cxc.cdict(v_out, unit, comps_cart)
+
+    if chart is cart:
+        return cast("CDict", p_cart_out)
+
+    # Map the base point forward (M @ at) to anchor the inverse Jacobian.
+    at_arr, at_unit = pack_uniform_unit(at_cart, keys=comps_cart)  # ty: ignore[no-matching-overload]
+    at_out_arr = jnp.einsum("ij,...j->...i", matrix, at_arr)
+    at_out = cxc.cdict(at_out_arr, at_unit, comps_cart)
+    return cast(
+        "CDict",
+        cxr.tangent_map(p_cart_out, cart, rep, chart, at=at_out, usys=usys),  # ty: ignore[missing-argument]
+    )
+
+
+@plum.dispatch
+def pushforward(
+    op: AbstractLinearTransform,
+    tau: Any,
+    v: CDict,
+    chart: cxc.AbstractChart,
+    rep: cxr.Representation,
+    /,
+    *,
+    at: CDict | None = None,
+    usys: OptUSys = None,
+) -> CDict:
+    r"""Frozen-$\tau$ pushforward of tangent data under a linear map.
+
+    A linear map's Jacobian is its (constant) matrix, so a Cartesian velocity
+    or acceleration transforms as ``v -> M v`` with no base point required —
+    matching the behavior already provided for `Rotate`. This is used by the
+    generic tangent ``act`` for displacement data and for every time-independent
+    linear transform (`Scale`, `Reflect`, `Shear`). Time-dependent linear maps
+    still route through the generic prolongation, which supplies the ``dot(M)``
+    term and requires the jet anchors.
+    """
+    return _linear_pushforward_cdict(op, tau, v, chart, rep, at=at, usys=usys)

tests/unit/transforms/test_prolongation.py

@@ -491,11 +491,24 @@ def test_prolong_additive_skips_intermediate_slots(self):
         out = cxfm.prolong(moving, u.Q(1.0, "s"), jet, cxc.cart3d)
         assert jnp.allclose(u.ustrip("km/s2", out[2]["x"]), 2.0)
 
-    def test_static_scale_vel_requires_at(self):
+    def test_static_scale_vel_no_at_needed(self):
+        # A static (time-independent) linear map has a constant Jacobian equal
+        # to its matrix, so a Cartesian velocity transforms as v -> M v with no
+        # base point required (matching Rotate).
         op = cxfm.Scale.from_factors([2.0, 3.0, 4.0])
         v = q3(1.0, 1.0, 1.0, "m/s")
-        with pytest.raises(TypeError, match="base point"):
-            cxfm.act(op, None, v, cxc.cart3d, cxr.coord_vel)
+        out = cxfm.act(op, None, v, cxc.cart3d, cxr.coord_vel)
+        assert jnp.allclose(u.ustrip("m/s", out["x"]), 2.0)
+        assert jnp.allclose(u.ustrip("m/s", out["y"]), 3.0)
+        assert jnp.allclose(u.ustrip("m/s", out["z"]), 4.0)
+
+    def test_static_scale_vel_noncartesian_still_requires_at(self):
+        # On a non-Cartesian chart the Jacobian varies with position, so the
+        # base point is still required.
+        op = cxfm.Scale.from_factors([2.0, 3.0, 4.0])
+        v = {"r": u.Q(1.0, "m/s"), "theta": u.Q(0.0, "rad/s"), "phi": u.Q(0.0, "rad/s")}
+        with pytest.raises(TypeError, match="'at'"):
+            cxfm.act(op, None, v, cxc.sph3d, cxr.coord_vel)
 
 
 # ============================================================================

tests/unit/transforms/test_spatial_linear_transforms.py

@@ -10,6 +10,8 @@
 import quaxed.numpy as jnp
 import unxt as u
 
+import coordinax.charts as cxc
+import coordinax.representations as cxr
 import coordinax.transforms as cxfm
 
 
@@ -73,3 +75,53 @@ def test_simplify_identity_scale_and_shear_to_identity() -> None:
 
     assert cxfm.simplify(s) is cxfm.identity
     assert cxfm.simplify(h) is cxfm.identity
+
+
+# ============================================================================
+# Kinematic (velocity / acceleration) acts — constant linear maps need no `at`
+
+
+def _vel(x, y, z):
+    return {"x": u.Q(x, "m/s"), "y": u.Q(y, "m/s"), "z": u.Q(z, "m/s")}
+
+
+@pytest.mark.parametrize(
+    ("op", "expected"),
+    [
+        (cxfm.Scale.from_factors([2.0, 3.0, 4.0]), (2.0, 3.0, 4.0)),
+        (cxfm.Reflect.from_normal([1.0, 0.0, 0.0]), (-1.0, 1.0, 1.0)),
+        (
+            cxfm.Shear(
+                jnp.asarray([[1.0, 1.0, 0.0], [0.0, 1.0, 0.0], [0.0, 0.0, 1.0]])
+            ),
+            (2.0, 1.0, 1.0),
+        ),
+    ],
+)
+def test_linear_transform_acts_on_velocity_without_at(op, expected) -> None:
+    """A constant linear map transforms a Cartesian velocity as v -> M v, no `at`."""
+    out = cxfm.act(op, None, _vel(1.0, 1.0, 1.0), cxc.cart3d, cxr.coord_vel)
+    got = tuple(round(float(u.ustrip("m/s", out[c])), 6) for c in ("x", "y", "z"))
+    assert got == expected
+
+
+def test_linear_transform_acts_on_acceleration_without_at() -> None:
+    """Acceleration also transforms as a -> M a for a constant linear map."""
+    op = cxfm.Scale.from_factors([2.0, 3.0, 4.0])
+    acc = {"x": u.Q(1.0, "m/s2"), "y": u.Q(0.0, "m/s2"), "z": u.Q(2.0, "m/s2")}
+    out = cxfm.act(op, None, acc, cxc.cart3d, cxr.coord_acc)
+    got = tuple(round(float(u.ustrip("m/s2", out[c])), 6) for c in ("x", "y", "z"))
+    assert got == (2.0, 0.0, 8.0)
+
+
+def test_linear_velocity_matches_generic_prolongation_with_at() -> None:
+    """Keystone: the no-`at` fast path equals the generic prolongation given `at`."""
+    op = cxfm.Scale.from_factors([2.0, 3.0, 4.0])
+    v = _vel(1.0, -2.0, 0.5)
+    at = {"x": u.Q(2.0, "m"), "y": u.Q(-1.0, "m"), "z": u.Q(3.0, "m")}
+    fast = cxfm.act(op, None, v, cxc.cart3d, cxr.coord_vel)
+    withat = cxfm.act(op, None, v, cxc.cart3d, cxr.coord_vel, at=at)
+    for c in ("x", "y", "z"):
+        np.testing.assert_allclose(
+            _to_np(fast[c], "m/s"), _to_np(withat[c], "m/s"), atol=1e-12
+        )