Skip to content

Commit cb44f2f

Browse files
committed
✨ feat: tangent representations
Signed-off-by: nstarman <nstarman@users.noreply.github.com>
1 parent 58e18ad commit cb44f2f

14 files changed

Lines changed: 1308 additions & 126 deletions

File tree

docs/spec.md

Lines changed: 39 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -946,7 +946,7 @@ A non-exhaustive table of exported objects are:
946946
| `coordinax.angles` | `AbstractAngle`, `Angle`, `wrap_to` |
947947
| `coordinax.distances` | `AbstractDistance`, `Distance` |
948948
| `coordinax.charts` | `CartesianProductChart`, </br> `cartesian_chart`, `guess_chart`, `cdict`, `pt_map`, `jac_pt_map`, </br> `cart0d`, </br> `cart1d`, `radial1d`, `time1d`, </br> `cart2d`, `polar2d`, </br> `cart3d`, `cyl3d`, `sph3d`, `lonlat_sph3d`, `loncoslat_sph3d`, `math_sph3d`, </br> `cartnd`, </br> `spacetimect` |
949-
| `coordinax.representations` | `cconvert`, `change_basis`, </br> `Representation`, `point`, `coord_disp`, `coord_vel`, `coord_acc`, `phys_disp`, `phys_vel`, `phys_acc`, </br> `PointGeometry`, `point_geom`, `TangentGeometry`, `tangent_geom`, </br> `NoBasis`, `no_basis`, `CoordinateBasis`, `coord_basis`, `PhysicalBasis`, `phys_basis`, </br> `Location`, `loc`, `Displacement`, `dpl`, `Velocity`, `vel`, `Acceleration`, `acc`, </br> `guess_geometry_kind`, `guess_semantic_kind`, `guess_rep` |
949+
| `coordinax.representations` | `cconvert`, `change_basis`, `tangent_map`, </br> `Representation`, `point`, `coord_disp`, `coord_vel`, `coord_acc`, `phys_disp`, `phys_vel`, `phys_acc`, </br> `PointGeometry`, `point_geom`, `TangentGeometry`, `tangent_geom`, </br> `NoBasis`, `no_basis`, `CoordinateBasis`, `coord_basis`, `PhysicalBasis`, `phys_basis`, </br> `Location`, `loc`, `Displacement`, `dpl`, `Velocity`, `vel`, `Acceleration`, `acc`, </br> `guess_geometry_kind`, `guess_semantic_kind`, `guess_rep` |
950950
| `coordinax.vectors` | `Point`, `ToUnitsOptions` |
951951
| `coordinax.manifolds` | `guess_manifold`, `scale_factors`, `angle_between`, </br> `EuclideanManifold`, `EuclideanMetric`, `euclidean3d`, </br> `EmbeddedManifold`, `EmbeddedChart` </br> `twosphere`, `embedded_twosphere`, </br> `CustomManifold`,`CustomAtlas`, |
952952
| `coordinax.transforms` | `act`, `simplify`, `compose`, `materialize_transform`, </br> `AbstractTransform`, `Identity`, `Composed`, `Translate`, `Rotate`, `Reflect`, `Scale`, `Shear`, `identity`, </br> `AbstractTransformGroup`, `IdentityGroup`, `DiffeomorphismGroup`, `AffineGroup`, `EuclideanGroup`, `OrthogonalGroup`, `SpecialOrthogonalGroup`, `PoincareGroup`, `LorentzGroup`, `ProperOrthochronousLorentzGroup` |
@@ -1728,6 +1728,44 @@ A representation is therefore **not** the same thing as a chart: the chart deter
17281728

17291729
- Inherits all failure semantics from `guess_geometry_kind`.
17301730

1731+
!!! info `tangent_map`
1732+
1733+
Transform a tangent vector from one chart to another.
1734+
1735+
**Signature:**
1736+
1737+
```text
1738+
tangent_map(v, from_chart, from_geom, from_rep, to_chart, to_geom, to_rep, /, *, at) -> CDict
1739+
```
1740+
1741+
A 4-argument shorthand form is also supported:
1742+
1743+
```text
1744+
tangent_map(v, from_chart, from_rep, to_chart, /, *, at) -> CDict
1745+
```
1746+
1747+
**Arguments:**
1748+
1749+
- `v`: `CDict` — tangent vector components in `from_chart` with basis `from_rep.basis`.
1750+
- `from_chart`: source chart.
1751+
- `from_geom`: source geometry (e.g. `TangentGeometry`).
1752+
- `from_rep`: source `Representation` (must have `TangentGeometry`).
1753+
- `to_chart`: target chart.
1754+
- `to_geom`: target geometry.
1755+
- `to_rep`: target `Representation`.
1756+
- `at`: `CDict` — base point in `from_chart` coordinates at which the tangent space is attached. Required for non-Cartesian charts (since the Jacobian depends on the base point).
1757+
1758+
**Semantics by basis:**
1759+
1760+
- **`CoordinateBasis`**: delegates to `jac_pt_map(at, from_chart, to_chart)` to obtain the Jacobian $J^j{}_i = \partial\tilde{q}^j/\partial q^i$, then applies $\tilde{v}^j = J^j{}_i v^i$.
1761+
- **`PhysicalBasis`**: fetch the orthonormal frame matrices $B_{\rm from}$ (columns = physical basis vectors in Cartesian) and $B_{\rm to}$ via `frame_cart`, compute $R = B_{\rm to}^T B_{\rm from}$, apply $\hat{v}' = R \hat{v}$.
1762+
1763+
**Same-chart optimisation:** when `from_chart is to_chart`, returns `v` unchanged.
1764+
1765+
**`cconvert` integration:** `cconvert` dispatches to `tangent_map` when the source representation has `TangentGeometry`, passing `at` through the `at` keyword argument.
1766+
1767+
**Same-chart basis conversion:** `change_basis(v, chart, from_basis, to_basis, /, *, at)` changes tangent component conventions without changing charts. In v1 it is defined only for Cartesian charts and `CoordinateBasis` $
1768+
17311769
</br>
17321770

17331771
### Geometric Kind

packages/coordinax.api/src/coordinax/api/representations.py

Lines changed: 17 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -32,6 +32,23 @@ def change_basis(*args: Any, **kwargs: Any) -> Any:
3232
raise NotImplementedError # pragma: no cover
3333

3434

35+
@plum.dispatch.abstract
36+
def tangent_map(*args: Any, **kwargs: Any) -> Any:
37+
"""Compute the tangent map (Jacobian) of a chart transition.
38+
39+
This is an abstract API definition. See the main coordinax package for
40+
concrete implementations.
41+
42+
Examples
43+
--------
44+
>>> import jax.numpy as jnp
45+
>>> import coordinax.charts as cxc
46+
>>> import coordinax.representations as cxr
47+
48+
"""
49+
raise NotImplementedError # pragma: no cover
50+
51+
3552
@plum.dispatch.abstract
3653
def cconvert(*args: Any, **kwargs: Any) -> Any:
3754
"""Transform the current vector to the target chart.

src/coordinax/main/__init__.py

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -78,6 +78,7 @@
7878
"phys_disp",
7979
"phys_vel",
8080
"phys_acc",
81+
"tangent_map",
8182
"change_basis",
8283
# vectors
8384
"Point",
@@ -151,6 +152,7 @@
151152
point_geom,
152153
subtract,
153154
tangent_geom,
155+
tangent_map,
154156
vel,
155157
)
156158
from coordinax.transforms import (

src/coordinax/representations/__init__.py

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -96,6 +96,7 @@
9696
"guess_rep",
9797
"guess_semantic_kind",
9898
"subtract",
99+
"tangent_map",
99100
"change_basis",
100101
# Representations
101102
"Representation",
@@ -171,6 +172,7 @@
171172
point,
172173
point_geom,
173174
tangent_geom,
175+
tangent_map,
174176
vel,
175177
)
176178
from coordinax.api.representations import (

src/coordinax/representations/_src/__init__.py

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -8,3 +8,4 @@
88
from .register_cx import *
99
from .rep import *
1010
from .semantics import *
11+
from .tangent_map import *

src/coordinax/representations/_src/basis_change.py

Lines changed: 3 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -3,7 +3,7 @@
33
__all__ = ("change_basis",)
44

55
from jaxtyping import ArrayLike
6-
from typing import Any
6+
from typing import Any, TypeVar
77

88
import jax
99
import jax.scipy.linalg
@@ -28,6 +28,8 @@
2828
from coordinax.internal import QuantityMatrix, UnitsMatrix
2929
from coordinax.internal.custom_types import CDict, OptUSys
3030

31+
T = TypeVar("T", bound=u.Quantity)
32+
3133
_RAD = u.unit("rad")
3234

3335

src/coordinax/representations/_src/core.py

Lines changed: 9 additions & 27 deletions
Original file line numberDiff line numberDiff line change
@@ -330,19 +330,14 @@ def cconvert(
330330
at: CDict | None = None,
331331
usys: OptUSys = None,
332332
) -> Any:
333-
r"""Convert tangent data between basis conventions in the same chart.
334-
335-
Tangent conversions are basis changes when source and target charts are
336-
identical. In this case, `cconvert` redispatches to `change_basis`.
333+
r"""Convert tangent data between charts via Jacobian pushforward.
337334
338335
Examples
339336
--------
340-
Convert tangent data between coordinate and physical basis in the same
341-
chart:
342-
343337
>>> import jax.numpy as jnp
344338
>>> import coordinax.charts as cxc
345339
>>> import coordinax.representations as cxr
340+
346341
>>> v = {"r": jnp.array(5.0), "theta": jnp.array(1.0), "phi": jnp.array(2.0)}
347342
>>> at = {"r": jnp.array(3.0), "theta": jnp.array(0.5), "phi": jnp.array(0.0)}
348343
>>> cxr.cconvert(v, cxc.sph3d, cxr.tangent_geom, cxr.coord_disp,
@@ -351,28 +346,15 @@ def cconvert(
351346
'theta': Array(3., dtype=float64, ...),
352347
'phi': Array(..., dtype=float64, ...)}
353348
354-
Tangent conversion across different charts is not implemented by this
355-
dispatch:
356-
357-
>>> cxr.cconvert(v, cxc.sph3d, cxr.tangent_geom, cxr.coord_disp,
358-
... cxc.cart3d, cxr.tangent_geom, cxr.coord_disp, at=at)
359-
Traceback (most recent call last):
360-
...
361-
NotImplementedError: Tangent cconvert between different charts is not implemented;
362-
use the same chart for basis changes.
349+
>>> v = {"x": jnp.array(1.0), "y": jnp.array(0.0)}
350+
>>> at = {"x": jnp.array(1.0), "y": jnp.array(0.0)}
351+
>>> cxr.cconvert(v, cxc.cart2d, cxr.coord_disp, cxc.polar2d, cxr.coord_disp, at=at)
352+
{'r': Array(1., ...), 'theta': Array(0., ...)}
363353
364354
"""
365-
del from_geom, to_geom # represented by dispatch signature
366-
367-
if from_chart != to_chart:
368-
msg = (
369-
"Tangent cconvert between different charts is not implemented; "
370-
"use the same chart for basis changes."
371-
)
372-
raise NotImplementedError(msg)
373-
374-
return cxrapi.change_basis(
375-
x, from_chart, from_rep.basis, to_rep.basis, at=at, usys=usys
355+
del from_geom, to_geom
356+
return cxrapi.tangent_map(
357+
x, from_chart, from_rep, to_chart, to_rep, at=at, usys=usys
376358
)
377359

378360

0 commit comments

Comments
 (0)