Skip to content

Commit 8b3cc1c

Browse files
authored
✨ feat(vectors): add Tangent and Coordinate vector bundle types (#522)
✨ feat(vectors): add Tangent and Coordinate vector bundle types with semantic transforms Add Tangent and Coordinate vector bundle types with explicit basis, semantic information, frame-aware transforms, and JAX/Quax integration. Highlights: - Add Tangent and Coordinate bundle types with chart/frame transforms via Jacobian pushforward - Support semantic transforms for displacement, velocity, and acceleration vectors - Register Quax/JAX primitives for arithmetic, broadcasting, dtype conversion, and equality - Improve Coordinate and Tangent validation, frame propagation, and representation consistency - Fix tangent rotation on non-Cartesian charts using Jacobian pushforward - Remove implicit chart alignment from Coordinate.__init__ - Enforce units on Tangent constructors and add missing overload dispatches - Preserve metadata in Quax broadcast/convert operations via _create_unchecked - Refine frame conversion, equality semantics, and add/sub mismatch diagnostics - Update specs, tutorials, doctests, and integration/unit tests - Reorganize configuration and tidy typing/JAX compatibility issues Also includes assorted fixes for transform composition, chart anchoring, Hypothesis compatibility, doctest formatting, and slow-schema test settings. Signed-off-by: nstarman <nstarman@users.noreply.github.com>
1 parent 2286ab9 commit 8b3cc1c

57 files changed

Lines changed: 5710 additions & 1029 deletions

File tree

Some content is hidden

Large Commits have some content hidden by default. Use the searchbox below for content that may be hidden.

docs/api/vectors.md

Lines changed: 26 additions & 16 deletions
Original file line numberDiff line numberDiff line change
@@ -13,21 +13,30 @@ For design philosophy, practical patterns, and worked examples, see [Working Wit
1313
```python
1414
import coordinax.main as cx
1515
import coordinax.charts as cxc
16-
17-
# Construct and convert
18-
v = cx.Point.from_([1, 2, 3], "m")
19-
v_sph = cx.cconvert(v, cxc.sph3d)
20-
21-
# Arithmetic. Technically points are NOT displacements,
22-
# but for convenience we allow subtraction to yield a new point.
23-
v2 = cx.Point.from_([4, 5, 6], "m")
24-
difference = v - v2
25-
26-
# Inspect converted components
27-
components = cx.cdict(v_sph)
16+
import coordinax.representations as cxr
17+
import unxt as u
18+
19+
# ── Point ──────────────────────────────────────────────────
20+
p = cx.Point.from_([1, 2, 3], "m")
21+
p_sph = cx.cconvert(p, cxc.sph3d)
22+
23+
# ── Tangent ────────────────────────────────────────────────
24+
# A velocity vector — transforms by Jacobian pushforward
25+
v = cx.Tangent.from_(
26+
{"x": u.Q(1.0, "m/s"), "y": u.Q(0.0, "m/s"), "z": u.Q(0.0, "m/s")},
27+
cxc.cart3d,
28+
cxr.coord_vel,
29+
)
30+
# Convert a Tangent — must supply the base Point via `at=`
31+
v_sph = v.cconvert(cxc.sph3d, at=p)
32+
33+
# ── Coordinate ─────────────────────────────────────────────
34+
# Bundle: base Point + named Tangent fibre fields
35+
pv = cx.Coordinate(point=p, velocity=v)
36+
pv_sph = pv.cconvert(cxc.sph3d) # converts point AND velocity together
2837
```
2938

30-
See [Working With Vectors](../guides/vectors.md#constructor-patterns) for all construction patterns and design rationale.
39+
See [Working With Vectors](../guides/vectors.md) and [Working With Tangent Vectors](../guides/tangents.md) for all construction patterns and design rationale.
3140

3241
## Functional API
3342

@@ -60,14 +69,15 @@ Additional utilities:
6069

6170
## Available Objects
6271

63-
- **`Point`**: the primary vector class storing data + chart + representation
72+
- **`Point`**: a geometric point storing data + chart + representation (always `PointGeometry`)
73+
- **`Tangent`**: a tangent-space vector with explicit basis and semantic kind (velocity, displacement, acceleration)
74+
- **`Coordinate`**: a vector bundle — a `Point` paired with named `Tangent` fibre fields anchored at that point
6475
- **`AbstractVector`**: base class defining the vector interface
65-
- **`Coordinate`**: a vector placed in a reference frame (see `coordinax.frames`)
6676
- **`ToUnitsOptions`**: configuration for unit conversion behavior
6777

6878
## Design & Integration
6979

70-
For design philosophy, architecture, and immutability details, see [Working With Vectors](../guides/vectors.md#architecture--design-philosophy). For JAX integration patterns (PyTree, scalar-first design, vmap/jit/grad), see [Working With Vectors](../guides/vectors.md#jax-integration--scaling).
80+
For design philosophy, architecture, and immutability details, see [Working With Vectors](../guides/vectors.md#architecture--design-philosophy). For tangent-vector patterns (basis, semantic kind, Jacobian pushforward), see [Working With Tangent Vectors](../guides/tangents.md). For JAX integration patterns (PyTree, scalar-first design, vmap/jit/grad), see [Working With Vectors](../guides/vectors.md#jax-integration--scaling).
7181

7282
```{automodule} coordinax.vectors
7383
:members:

docs/guides/representations.md

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -27,7 +27,7 @@ A representation is a triple $R = (K, B, S)$:
2727
2828
>>> rep = cxr.Representation(cxr.PointGeometry(), cxr.NoBasis(), cxr.Location())
2929
>>> rep
30-
Representation(geom_kind=PointGeometry(), basis=NoBasis(), semantic_kind=Location())
30+
point
3131
3232
>>> rep == cxr.point
3333
True

docs/guides/tangents.md

Lines changed: 251 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,251 @@
1+
# Working With Tangent Vectors
2+
3+
This guide explains `Tangent` — coordinax's type for **tangent-space quantities** (velocities, displacements, accelerations). For bundling a `Tangent` with a base `Point`, see the [Coordinate guide](./vectors.md) and the [Coordinate tutorial](../tutorials/coordinate_objects.md). For API reference see [the vector module reference](../api/vectors.md).
4+
5+
## Why Tangent Vectors?
6+
7+
Points and tangent vectors look similar — both are dictionaries of components attached to a chart — but they **transform differently** when you change coordinate system:
8+
9+
- **Point** `(r, θ, φ)`: transforms by the chart transition map $\varphi' \circ \varphi^{-1}$.
10+
- **Tangent** `(ṙ, θ̇, φ̇)`: transforms by the **Jacobian** $J^j{}_i = \partial \tilde{q}^j / \partial q^i$ evaluated at the base point.
11+
12+
Getting this wrong produces incorrect physical results. `Tangent` encodes the correct transformation law, so `cconvert` always does the right thing automatically.
13+
14+
Additionally, tangent vectors carry two extra pieces of metadata:
15+
16+
- **basis**: are the components expressed in the _coordinate basis_ ($\partial/\partial q^i$) or the _physical (orthonormal) basis_ ($\hat{e}_i$)?
17+
- **semantic**: what physical quantity do the components represent — `vel` (velocity), `dpl` (displacement), or `acc` (acceleration)?
18+
19+
## Basis And Semantic Kind
20+
21+
### Basis
22+
23+
| Object | Singleton | Meaning |
24+
| --- | --- | --- |
25+
| `CoordinateBasis` | `cxr.coord_basis` | components in coordinate (tangent) basis |
26+
| `PhysicalBasis` | `cxr.phys_basis` | components in orthonormal physical frame |
27+
28+
For **Cartesian charts** these are identical. For **spherical/cylindrical** charts the coordinate basis has components with units like `m/s`, `rad/s`, `rad/s`, while the physical basis has all components in `m/s` (scaled by the metric's scale factors).
29+
30+
### Semantic Kind
31+
32+
| Object | Singleton | Physical meaning |
33+
| -------------- | --------- | --------------------------------- |
34+
| `Displacement` | `cxr.dpl` | position displacement $\Delta q$ |
35+
| `Velocity` | `cxr.vel` | time derivative $\dot{q}$ |
36+
| `Acceleration` | `cxr.acc` | second time derivative $\ddot{q}$ |
37+
38+
### Pre-Built Representations
39+
40+
The most common combinations are available as ready-made singletons:
41+
42+
| Name | `geom_kind` | `basis` | `semantic` |
43+
| ---------------- | ----------------- | ------------- | ---------- |
44+
| `cxr.coord_disp` | `TangentGeometry` | `coord_basis` | `dpl` |
45+
| `cxr.coord_vel` | `TangentGeometry` | `coord_basis` | `vel` |
46+
| `cxr.coord_acc` | `TangentGeometry` | `coord_basis` | `acc` |
47+
| `cxr.phys_disp` | `TangentGeometry` | `phys_basis` | `dpl` |
48+
| `cxr.phys_vel` | `TangentGeometry` | `phys_basis` | `vel` |
49+
| `cxr.phys_acc` | `TangentGeometry` | `phys_basis` | `acc` |
50+
51+
## Constructing Tangent Vectors
52+
53+
::::{tab-set}
54+
55+
:::{tab-item} Dict + rep (most explicit)
56+
57+
Name every component and supply the representation directly:
58+
59+
```{code-block} python
60+
>>> import coordinax.main as cx
61+
>>> import coordinax.charts as cxc
62+
>>> import coordinax.representations as cxr
63+
>>> import unxt as u
64+
65+
>>> vel = cx.Tangent.from_(
66+
... {"x": u.Q(1, "m/s"), "y": u.Q(2, "m/s"), "z": u.Q(3, "m/s")},
67+
... cxc.cart3d, cxr.coord_vel)
68+
>>> print(vel)
69+
<Tangent: chart=Cart3D (x, y, z) [m / s]
70+
[1 2 3]>
71+
```
72+
73+
:::
74+
75+
:::{tab-item} Dict + basis + semantic
76+
77+
Specify basis and semantic separately (equivalent to the above):
78+
79+
```{code-block} python
80+
>>> vel = cx.Tangent.from_(
81+
... {"x": u.Q(1, "m/s"), "y": u.Q(2, "m/s"), "z": u.Q(3, "m/s")},
82+
... cxc.cart3d, cxr.coord_basis, cxr.vel)
83+
>>> print(vel)
84+
<Tangent: chart=Cart3D (x, y, z) [m / s]
85+
[1 2 3]>
86+
```
87+
88+
:::
89+
90+
:::{tab-item} Array + unit
91+
92+
Chart is inferred from the array length (3 → `cart3d`):
93+
94+
```{code-block} python
95+
>>> vel = cx.Tangent.from_([1, 2, 3], "m/s")
96+
>>> print(vel.chart) # Cart3D(M=Rn(3))
97+
Cart3D[('x', 'y', 'z'), ('length', 'length', 'length')](M=Rn(3))
98+
```
99+
100+
:::
101+
102+
:::{tab-item} Passthrough
103+
104+
Passing an existing `Tangent` returns it unchanged:
105+
106+
```{code-block} python
107+
>>> vel2 = cx.Tangent.from_(vel)
108+
>>> assert vel2 is vel
109+
```
110+
111+
:::
112+
113+
::::
114+
115+
## Accessing Data
116+
117+
```{code-block} python
118+
>>> vel = cx.Tangent.from_(
119+
... {"x": u.Q(1, "m/s"), "y": u.Q(2, "m/s"), "z": u.Q(3, "m/s")},
120+
... cxc.cart3d, cxr.coord_vel)
121+
122+
>>> # Component access by name
123+
>>> print(vel["x"]) # Q(1., 'm / s')
124+
Q['speed'](Array(1, dtype=int64, weak_type=True), unit='m / s')
125+
126+
# Metadata
127+
>>> print(vel.chart) # Cart3D(M=Rn(3))
128+
Cart3D[('x', 'y', 'z'), ('length', 'length', 'length')](M=Rn(3))
129+
130+
>>> print(vel.basis) # coord_basis
131+
CoordinateBasis()
132+
133+
>>> print(vel.semantic) # vel
134+
Velocity()
135+
136+
>>> print(vel.frame) # NoFrame()
137+
NoFrame()
138+
```
139+
140+
## Chart Conversion (Jacobian Pushforward)
141+
142+
Converting a `Tangent` to a new chart requires knowing the **base point** at which the Jacobian is evaluated. Pass it via the `at=` argument:
143+
144+
```{code-block} python
145+
>>> point = cx.Point.from_([1, 0, 0], "m") # Base point in Cartesian
146+
>>> vel_cart = cx.Tangent.from_(
147+
... {"x": u.Q(1, "m/s"), "y": u.Q(0, "m/s"), "z": u.Q(0, "m/s")},
148+
... cxc.cart3d, cxr.coord_vel)
149+
150+
>>> # Convert to spherical — Jacobian is evaluated at `point`
151+
>>> vel_sph = vel_cart.cconvert(cxc.sph3d, at=point)
152+
>>> print(vel_sph)
153+
<Tangent: chart=Spherical3D (r[m / s], theta[rad / s], phi[rad / s])
154+
[ 1. -0. 0.]>
155+
```
156+
157+
**Key point**: the `at=` base point must be in the **same chart** as the tangent vector. To convert a point and its tangent field together in one call, bundle them into a `Coordinate` — see the [Coordinate tutorial](../tutorials/coordinate_objects.md).
158+
159+
## Changing Basis
160+
161+
`change_basis` converts between coordinate and physical bases for `Tangent` objects directly. Pass the tangent vector, the target basis, and the base `Point` at which scale factors are evaluated:
162+
163+
```{code-block} python
164+
>>> # Velocity in coordinate basis at a spherical point
165+
>>> point_sph = cx.Point.from_(
166+
... {"r": u.Q(1, "m"), "theta": u.Q(0.5, "rad"), "phi": u.Q(0, "rad")}, cxc.sph3d
167+
... )
168+
>>> vel_sph_coord = cx.Tangent.from_(
169+
... {"r": u.Q(1, "m/s"), "theta": u.Q(0, "rad/s"), "phi": u.Q(0, "rad/s")},
170+
... cxc.sph3d, cxr.coord_vel)
171+
>>> # Convert to physical basis (all components in m/s)
172+
>>> vel_sph_phys = cxr.change_basis(vel_sph_coord, cxr.phys_basis, at=point_sph)
173+
>>> print(vel_sph_phys.rep) # phys_vel
174+
Representation(
175+
geom_kind=TangentGeometry(), basis=PhysicalBasis(), semantic_kind=Velocity()
176+
)
177+
```
178+
179+
## Promoting a Point to a Displacement
180+
181+
`change_basis` has a second role: it can **promote a `Point` to a `Tangent` with `Displacement` semantics**. This is useful when you have a position vector that you want to reinterpret as a tangent-space displacement — for example, when computing offsets or feeding a position into an operation that expects a displacement.
182+
183+
The component data is unchanged; only the geometric interpretation is recast from a manifold point (`PointGeometry`) to a tangent-space displacement (`TangentGeometry`, `Displacement`).
184+
185+
```{code-block} python
186+
>>> pt = cx.Point.from_([1, 2, 3], "m")
187+
188+
>>> # Promote to a coordinate-basis displacement
189+
>>> disp = cxr.change_basis(pt, cxr.coord_basis)
190+
>>> print(disp)
191+
<Tangent: chart=Cart3D (x, y, z) [m]
192+
[1 2 3]>
193+
```
194+
195+
You can also request the physical basis directly:
196+
197+
```{code-block} python
198+
>>> disp_phys = cxr.change_basis(pt, cxr.phys_basis)
199+
>>> print(disp_phys.rep) # phys_disp
200+
Representation(
201+
geom_kind=TangentGeometry(), basis=PhysicalBasis(), semantic_kind=Displacement()
202+
)
203+
```
204+
205+
The `at` and `usys` keyword arguments are accepted for API consistency but have no effect (no numerical transformation is performed).
206+
207+
### JIT
208+
209+
```{code-block} python
210+
>>> import jax
211+
>>> point = cx.Point.from_([1.0, 0.0, 0.0], "m")
212+
>>> vel = cx.Tangent.from_(
213+
... {"x": u.Q(1.0, "m/s"), "y": u.Q(0.0, "m/s"), "z": u.Q(0.0, "m/s")},
214+
... cxc.cart3d, cxr.coord_vel)
215+
216+
>>> convert_vel = jax.jit(lambda v, p: cx.cconvert(v, cxc.sph3d, at=p))
217+
>>> vel_sph = convert_vel(vel, point)
218+
>>> print(vel_sph.chart) # Spherical3D(M=Rn(3))
219+
Spherical3D[('r', 'theta', 'phi'), ('length', 'angle', 'angle')](M=Rn(3))
220+
```
221+
222+
### vmap
223+
224+
Use scalar-first design and batch with `vmap`:
225+
226+
```{code-block} python
227+
>>> import jax.numpy as jnp
228+
229+
>>> # Build a batched Tangent by stacking scalar values per component
230+
>>> batch_vel = cx.Tangent(
231+
... { "x": u.Q(jnp.array([1, 2, 3]), "m/s"),
232+
... "y": u.Q(jnp.zeros(3), "m/s"), "z": u.Q(jnp.zeros(3), "m/s") },
233+
... cxc.cart3d, cxr.coord_basis, cxr.vel)
234+
>>> batch_point = cx.Point(
235+
... { "x": u.Q(jnp.ones(3), "m"), "y": u.Q(jnp.zeros(3), "m"),
236+
... "z": u.Q(jnp.zeros(3), "m") }, cxc.cart3d)
237+
238+
>>> vec_fn = jax.vmap(lambda v, p: cx.cconvert(v, cxc.sph3d, at=p))
239+
>>> vels_sph = vec_fn(batch_vel, batch_point)
240+
>>> print(vels_sph.data) # Spherical3D(M=Rn(3))
241+
{'phi': Q([0., 0., 0.], 'rad / s'), 'r': Q([1., 2., 3.], 'm / s'), 'theta': Q([-0., -0., -0.], 'rad / s')}
242+
```
243+
244+
## When To Use Tangent
245+
246+
| You have | Use |
247+
| --- | --- |
248+
| Position only | `Point` — see [Point tutorial](../tutorials/point_objects.md) |
249+
| Velocity / displacement / acceleration only | `Tangent` (this guide) |
250+
| Position + tangent field(s) | `Coordinate` — see [Coordinate tutorial](../tutorials/coordinate_objects.md) |
251+
| Position to reinterpret as a displacement | `change_basis(pt, cxr.coord_basis)``Tangent[Displacement]` |

0 commit comments

Comments
 (0)