From 33dcba2f997e77200fc734e885bee4cd4a41f8e6 Mon Sep 17 00:00:00 2001 From: nstarman Date: Fri, 24 Jul 2026 23:05:22 -0400 Subject: [PATCH 1/9] =?UTF-8?q?=F0=9F=90=9B=20fix(charts):=20correct=20Poi?= =?UTF-8?q?ncarePolar6D=20as=20a=20symplectic=20phase-space=20chart?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit PoincarePolar6D declared its 5th coordinate `dt_pp_phi` with dimension `length/time**1.5` (a literal d/dt of `pp_phi` reading). But the authoritative definition — gala's `cartesian_to_poincare_polar` (Papaphilippou & Laskar 1996, A&A 307, 427) — makes it the *symplectic quadrature partner* `sqrt(2|Lz|)*sin(phi)`, with the SAME `length/time**0.5` dimension as `pp_phi`, not its time derivative. The `time**1.5` unit and `dt_` name were a placeholder inconsistency. * rename component `dt_pp_phi` -> `pp_phidot` (gala's `p_phi_dot`) * fix its dimension `length/time**1.5` -> `length/time**0.5` * document the chart as Poincaré *symplectic-polar* variables on phase space `T*R^3`, why it has no Riemannian metric (`M = no_manifold`: phase space is symplectic, `metric_matrix` is inapplicable) and no global Cartesian chart, and the forward-only / sign(Lz)-ambiguity caveat * drop the stale `# TODO: actual manifold` markers (no_manifold is correct) Updates the predefined-chart test and spec.md accordingly. (The gala forward transform from a Cartesian phase-space source chart is a follow-up: it needs a new 6D phase-space chart, which does not exist yet.) Co-Authored-By: Claude Opus 4.8 --- docs/spec.md | 4 +- src/coordinax/_src/charts/d6.py | 49 +++++++++++++++++++------ tests/unit/charts/test_predef_charts.py | 4 +- 3 files changed, 42 insertions(+), 15 deletions(-) diff --git a/docs/spec.md b/docs/spec.md index dee04e106..443c7e588 100644 --- a/docs/spec.md +++ b/docs/spec.md @@ -1701,9 +1701,9 @@ The `coordinax.charts` module provides the chart-facing API for representing poi - `PoincarePolar6D` is a final concrete chart type in the 6-D chart family. - Components (ordered): - `("rho", "pp_phi", "z", "dt_rho", "dt_pp_phi", "dt_z")`. + `("rho", "pp_phi", "z", "dt_rho", "pp_phidot", "dt_z")`. - Coordinate dimensions (ordered): - `("length", "length / time**0.5", "length", "speed", "length / time**1.5", "speed")`. + `("length", "length / time**0.5", "length", "speed", "length / time**0.5", "speed")`. - `poincarepolar6d` is the pre-defined `PoincarePolar6D()` instance. - Transition behavior currently registered in-core: identity transform only (`pt_map(p, PoincarePolar6D, PoincarePolar6D) -> p`). diff --git a/src/coordinax/_src/charts/d6.py b/src/coordinax/_src/charts/d6.py index 7b6cc9431..3dec1a95d 100644 --- a/src/coordinax/_src/charts/d6.py +++ b/src/coordinax/_src/charts/d6.py @@ -42,11 +42,11 @@ def __init_subclass__(cls, n: int | L["N"] | None = None, **kw: Any) -> None: PoincarePolarKeys = tuple[ - L["rho"], L["pp_phi"], L["z"], L["dt_rho"], L["dt_pp_phi"], L["dt_z"] + L["rho"], L["pp_phi"], L["z"], L["dt_rho"], L["pp_phidot"], L["dt_z"] ] PoincarePolarDims = tuple[ - Len, L["length / time**0.5"], Len, Spd, L["length / time**1.5"], Spd + Len, L["length / time**0.5"], Len, Spd, L["length / time**0.5"], Spd ] @@ -55,21 +55,46 @@ def __init_subclass__(cls, n: int | L["N"] | None = None, **kw: Any) -> None: class PoincarePolar6D( AbstractFixedComponentsChart[Any, PoincarePolarKeys, PoincarePolarDims], Abstract6D ): - r"""Six-dimensional Poincare-polar chart. + r"""Six-dimensional Poincaré symplectic-polar chart on phase space. - Components are ordered as - $(\rho,\;\mathrm{pp\_phi},\;z,\;\dot\rho,\;\dot{\mathrm{pp\_phi}},\;\dot z)$ + A chart on the 6-D phase space $T^*\mathbb{R}^3$ in the *Poincaré symplectic + polar* variables used in galactic dynamics. The azimuthal action-angle pair + $(\phi, L_z)$ is replaced by the Cartesian-like quadrature pair + + $\mathrm{pp\_phi} = \sqrt{2\,|L_z|}\,\cos\phi,$ + $\mathrm{pp\_phidot} = \sqrt{2\,|L_z|}\,\sin\phi,$ + + which removes the coordinate singularity of polar coordinates on the axis. + Following {mod}`gala` (``cartesian_to_poincare_polar``; Papaphilippou & + Laskar 1996, A&A 307, 427) the components are ordered as + $(\rho,\;\mathrm{pp\_phi},\;z,\;\dot\rho,\;\mathrm{pp\_phidot},\;\dot z)$ with dimensions - $(\mathrm{length},\;\mathrm{length}/\mathrm{time}^{1/2},\;\mathrm{length},\;\mathrm{speed},\;\mathrm{length}/\mathrm{time}^{3/2},\;\mathrm{speed})$. + $(\mathrm{length},\;\mathrm{length}/\mathrm{time}^{1/2},\;\mathrm{length},\;\mathrm{speed},\;\mathrm{length}/\mathrm{time}^{1/2},\;\mathrm{speed})$. + + ``pp_phi`` and ``pp_phidot`` are the two symplectic quadratures of the + azimuth (both $\sqrt{\text{action}}$, hence the ``length / time**0.5`` + dimension); ``pp_phidot`` is **not** the time derivative of ``pp_phi`` + despite the ``dot`` suffix (it is gala's ``p_phi_dot``). ``rho``, ``z`` are + cylindrical position and ``dt_rho``, ``dt_z`` their velocities. + + Notes + ----- + Phase space carries a *symplectic* form $\omega$, not a Riemannian metric, + so this chart has **no manifold metric**: ``M`` is ``no_manifold`` and there + is no ``metric_matrix``. It likewise has no global 6-D Cartesian chart + (``cartesian`` raises). The forward map discards $\mathrm{sign}(L_z)$ (via + $\sqrt{|L_z|}$), so it is not injective and the inverse is ambiguous in the + sign of angular momentum — matching gala, which provides only the forward + transform. Examples -------- >>> import coordinax.charts as cxc >>> cxc.poincarepolar6d.components - ('rho', 'pp_phi', 'z', 'dt_rho', 'dt_pp_phi', 'dt_z') + ('rho', 'pp_phi', 'z', 'dt_rho', 'pp_phidot', 'dt_z') >>> cxc.poincarepolar6d.coord_dimensions - ('length', 'length / time**0.5', 'length', 'speed', 'length / time**1.5', 'speed') + ('length', 'length / time**0.5', 'length', 'speed', 'length / time**0.5', 'speed') >>> isinstance(cxc.poincarepolar6d, cxc.PoincarePolar6D) True @@ -77,7 +102,9 @@ class PoincarePolar6D( """ _: dataclasses.KW_ONLY - M: AbstractManifold = no_manifold # TODO: actual manifold + # Phase space is symplectic, not Riemannian: there is no manifold metric, + # so `M` is `no_manifold` (a metric / `metric_matrix` is inapplicable). + M: AbstractManifold = no_manifold @override @property @@ -88,5 +115,5 @@ def cartesian(self) -> NoReturn: ) -poincarepolar6d: Final = PoincarePolar6D() # TODO: actual manifold -"""Six-dimensional Poincare-polar chart.""" +poincarepolar6d: Final = PoincarePolar6D() +"""Six-dimensional Poincaré symplectic-polar phase-space chart.""" diff --git a/tests/unit/charts/test_predef_charts.py b/tests/unit/charts/test_predef_charts.py index 63ee4f21c..e59a6b61b 100644 --- a/tests/unit/charts/test_predef_charts.py +++ b/tests/unit/charts/test_predef_charts.py @@ -111,13 +111,13 @@ "poincarepolar6d", cxc.poincarepolar6d, cxc.PoincarePolar6D, - ("rho", "pp_phi", "z", "dt_rho", "dt_pp_phi", "dt_z"), + ("rho", "pp_phi", "z", "dt_rho", "pp_phidot", "dt_z"), ( "length", "length / time**0.5", "length", "speed", - "length / time**1.5", + "length / time**0.5", "speed", ), 6, From 574435b86c9fe09d0e23cc6ac9429ffcc4f8c024 Mon Sep 17 00:00:00 2001 From: nstarman Date: Sat, 25 Jul 2026 10:39:45 -0400 Subject: [PATCH 2/9] =?UTF-8?q?=E2=9C=A8=20feat(charts):=20gala=20forward?= =?UTF-8?q?=20transform=20Cartesian=20phase=20space=20->=20PoincarePolar6D?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Register the pt_map from a two-factor Cartesian product chart (cart3d x cart3d = [position, velocity]) to PoincarePolar6D, implementing gala's cartesian_to_poincare_polar (Papaphilippou & Laskar 1996): rho = hypot(x, y); phi = atan2(x, y) [gala azimuth convention] Lz = x*vy - y*vx; s = sqrt(2|Lz|) pp_phi = s*cos(phi); pp_phidot = s*sin(phi) dt_rho = (x*vx + y*vy)/rho; dt_z = vz Reuses the existing product chart as the source (no new public chart). Forward only — sqrt(|Lz|) drops sign(Lz), so there is no inverse (matching gala). A non-(Cart3D, Cart3D) product source raises NotImplementedError. Known-value doctest covers the symplectic quadratures and the guard. Co-Authored-By: Claude Opus 4.8 --- src/coordinax/_src/charts/register_ptmap.py | 87 +++++++++++++++++++++ 1 file changed, 87 insertions(+) diff --git a/src/coordinax/_src/charts/register_ptmap.py b/src/coordinax/_src/charts/register_ptmap.py index b8a7be2a8..f21db72b8 100644 --- a/src/coordinax/_src/charts/register_ptmap.py +++ b/src/coordinax/_src/charts/register_ptmap.py @@ -35,6 +35,9 @@ from coordinax._src.base.manifold import AbstractManifold from coordinax._src.custom_types import CDict, OptUSys from coordinax._src.euclidean import RN, EuclideanManifold, Rn +from coordinax._src.null import NoManifold +from coordinax._src.product.chart import CartesianProductChart +from coordinax._src.product.manifold import CartesianProductManifold from coordinax._src.utils import uconvert_to_rad ##################################################################### @@ -1742,3 +1745,87 @@ def pt_map( p_out: Array = jnp.stack([p_to[comp] for comp in to_chart.components], axis=-1) return p_out + + +# =================================================================== +# Cartesian phase space (cart3d x cart3d) -> Poincaré symplectic polar + + +@plum.dispatch +def pt_map( + p: CDict, + from_M: CartesianProductManifold, + from_chart: CartesianProductChart, + to_M: NoManifold, + to_chart: PoincarePolar6D, + /, + *, + usys: OptUSys = None, +) -> CDict: + r"""Cartesian phase space ``cart3d x cart3d`` -> ``PoincarePolar6D`` (gala forward). + + The source is a two-factor Cartesian product chart: factor 0 is position + ``(x, y, z)``, factor 1 its velocity ``(vx, vy, vz)``. Implements gala's + ``cartesian_to_poincare_polar`` (Papaphilippou & Laskar 1996): + + ``rho = hypot(x, y)``, ``phi = atan2(x, y)`` (gala's azimuth convention), + ``dt_rho = (x*vx + y*vy) / rho``, ``Lz = x*vy - y*vx``, + ``pp_phi = sqrt(2|Lz|) cos(phi)``, ``pp_phidot = sqrt(2|Lz|) sin(phi)``, + ``dt_z = vz``. + + Forward only: ``sqrt(|Lz|)`` discards ``sign(Lz)``, so there is no inverse. + + >>> import coordinax.charts as cxc + >>> import unxt as u + >>> ps = cxc.CartesianProductChart((cxc.cart3d, cxc.cart3d), ("q", "p")) + >>> q = {"q.x": u.Q(3.0, "kpc"), "q.y": u.Q(4.0, "kpc"), "q.z": u.Q(5.0, "kpc"), + ... "p.x": u.Q(1.0, "kpc/Myr"), "p.y": u.Q(2.0, "kpc/Myr"), + ... "p.z": u.Q(0.5, "kpc/Myr")} + >>> out = cxc.pt_map(q, ps.M, ps, cxc.poincarepolar6d.M, cxc.poincarepolar6d) + >>> sorted(out) + ['dt_rho', 'dt_z', 'pp_phi', 'pp_phidot', 'rho', 'z'] + + Lz = x*vy - y*vx = 2 kpc^2/Myr, so sqrt(2|Lz|) = 2; phi = atan2(3, 4): + + >>> out["rho"], out["dt_rho"], out["dt_z"] + (Q(5., 'kpc'), Q(2.2, 'kpc / Myr'), Q(0.5, 'kpc / Myr')) + >>> out["pp_phi"].round(4), out["pp_phidot"].round(4) + (Q(1.6, 'kpc / Myr(1/2)'), Q(1.2, 'kpc / Myr(1/2)')) + + A non-Cartesian or wrong-arity product source is rejected: + + >>> bad = cxc.CartesianProductChart((cxc.cart3d, cxc.polar2d), ("q", "p")) + >>> try: + ... cxc.pt_map({}, bad.M, bad, cxc.poincarepolar6d.M, cxc.poincarepolar6d) + ... except NotImplementedError: + ... print("rejected") + rejected + + """ + del usys + if len(from_chart.factors) != 2 or not all( + isinstance(f, Cart3D) for f in from_chart.factors + ): + msg = ( + "pt_map to PoincarePolar6D requires a Cartesian phase-space source: a " + "two-factor CartesianProductChart of (Cart3D, Cart3D) [position, " + f"velocity]; got factors {from_chart.factors!r}." + ) + raise NotImplementedError(msg) + + pos, vel = from_chart.split_components(p) + x, y, z = pos["x"], pos["y"], pos["z"] + vx, vy, vz = vel["x"], vel["y"], vel["z"] + + rho = jnp.hypot(x, y) + phi = jnp.atan2(x, y) # gala convention: azimuth from +y + lz = x * vy - y * vx + s = jnp.sqrt(2 * jnp.abs(lz)) + return { + "rho": rho, + "pp_phi": s * jnp.cos(phi), + "z": z, + "dt_rho": (x * vx + y * vy) / rho, + "pp_phidot": s * jnp.sin(phi), + "dt_z": vz, + } From c8c2d030543fd99705f8d4ec9370b7f460096e6c Mon Sep 17 00:00:00 2001 From: nstarman Date: Sat, 25 Jul 2026 10:51:33 -0400 Subject: [PATCH 3/9] =?UTF-8?q?=E2=9C=A8=20feat(charts):=20partial=20inver?= =?UTF-8?q?se=20PoincarePolar6D=20->=20Cartesian=20phase=20space?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Adds the reverse pt_map (PoincarePolar6D -> cart3d x cart3d), derived algebraically from the forward gala map: s = hypot(pp_phi, pp_phidot); phi = atan2(pp_phidot, pp_phi); Lz = s**2/2 x = rho sin(phi); y = rho cos(phi); z = z vx = sin(phi) dt_rho - cos(phi) Lz/rho vy = cos(phi) dt_rho + sin(phi) Lz/rho; vz = dt_z Partial inverse: the forward's sqrt(2|Lz|) discards sign(Lz), so this assumes Lz >= 0 and is exact only for non-negative-Lz points (position always round-trips; velocity sign is lost otherwise). Guards a non-(Cart3D, Cart3D) target with NotImplementedError. Round-trip doctest included. Co-Authored-By: Claude Opus 4.8 --- src/coordinax/_src/charts/register_ptmap.py | 66 +++++++++++++++++++++ 1 file changed, 66 insertions(+) diff --git a/src/coordinax/_src/charts/register_ptmap.py b/src/coordinax/_src/charts/register_ptmap.py index f21db72b8..412be04ff 100644 --- a/src/coordinax/_src/charts/register_ptmap.py +++ b/src/coordinax/_src/charts/register_ptmap.py @@ -1829,3 +1829,69 @@ def pt_map( "pp_phidot": s * jnp.sin(phi), "dt_z": vz, } + + +@plum.dispatch +def pt_map( + p: CDict, + from_M: NoManifold, + from_chart: PoincarePolar6D, + to_M: CartesianProductManifold, + to_chart: CartesianProductChart, + /, + *, + usys: OptUSys = None, +) -> CDict: + r"""``PoincarePolar6D`` -> Cartesian phase space (partial inverse of gala map). + + Inverts the gala forward map. Because the forward uses ``sqrt(2|Lz|)`` the + sign of the angular momentum is not recoverable, so this assumes ``Lz >= 0`` + (the standard convention) and is a *partial* inverse — exact only when the + original point had non-negative ``Lz``: + + ``s = hypot(pp_phi, pp_phidot)``, ``phi = atan2(pp_phidot, pp_phi)``, + ``Lz = s**2 / 2``, ``x = rho sin(phi)``, ``y = rho cos(phi)``, + ``vx = sin(phi) dt_rho - cos(phi) Lz/rho``, + ``vy = cos(phi) dt_rho + sin(phi) Lz/rho``, ``vz = dt_z``. + + (Singular on the axis ``rho = 0``, inherent to the coordinates.) + + >>> import coordinax.charts as cxc + >>> import unxt as u + >>> ps = cxc.CartesianProductChart((cxc.cart3d, cxc.cart3d), ("q", "p")) + + Round-trips a forward result whose Lz >= 0: + + >>> q = {"q.x": u.Q(3.0, "kpc"), "q.y": u.Q(4.0, "kpc"), "q.z": u.Q(5.0, "kpc"), + ... "p.x": u.Q(1.0, "kpc/Myr"), "p.y": u.Q(2.0, "kpc/Myr"), + ... "p.z": u.Q(0.5, "kpc/Myr")} + >>> pp = cxc.pt_map(q, ps.M, ps, cxc.poincarepolar6d.M, cxc.poincarepolar6d) + >>> back = cxc.pt_map(pp, cxc.poincarepolar6d.M, cxc.poincarepolar6d, ps.M, ps) + >>> back["q.x"].round(6), back["q.y"].round(6), back["p.x"].round(6) + (Q(3., 'kpc'), Q(4., 'kpc'), Q(1., 'kpc / Myr')) + + """ + del usys + if len(to_chart.factors) != 2 or not all( + isinstance(f, Cart3D) for f in to_chart.factors + ): + msg = ( + "pt_map from PoincarePolar6D requires a Cartesian phase-space target: a " + "two-factor CartesianProductChart of (Cart3D, Cart3D) [position, " + f"velocity]; got factors {to_chart.factors!r}." + ) + raise NotImplementedError(msg) + + rho, z, dt_rho, dt_z = p["rho"], p["z"], p["dt_rho"], p["dt_z"] + s = jnp.hypot(p["pp_phi"], p["pp_phidot"]) + phi = jnp.atan2(p["pp_phidot"], p["pp_phi"]) + lz = s**2 / 2 # sign(Lz) is not recoverable from the forward map; take Lz >= 0 + sinp, cosp = jnp.sin(phi), jnp.cos(phi) + + pos = {"x": rho * sinp, "y": rho * cosp, "z": z} + vel = { + "x": sinp * dt_rho - cosp * lz / rho, + "y": cosp * dt_rho + sinp * lz / rho, + "z": dt_z, + } + return cast("CDict", to_chart.merge_components((pos, vel))) From 5c1721ac3d9d23c6cddc31299dd654a23c5484c8 Mon Sep 17 00:00:00 2001 From: nstarman Date: Mon, 27 Jul 2026 10:55:14 -0400 Subject: [PATCH 4/9] =?UTF-8?q?=F0=9F=90=9B=20fix(charts):=20address=20Poi?= =?UTF-8?q?ncarePolar6D=20pt=5Fmap=20review=20comments?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit - add manifold/chart consistency asserts to both new pt_map dispatches - guard dt_rho against 0/0 on the axis (rho == 0) -> dt_rho = 0 by convention - clarify forward docstring: no global inverse, but a partial (Lz >= 0) one exists - update docs/spec.md to list the new gala forward + partial-inverse transitions Co-Authored-By: Claude Opus 4.8 --- docs/spec.md | 7 +++++-- src/coordinax/_src/charts/register_ptmap.py | 12 ++++++++++-- 2 files changed, 15 insertions(+), 4 deletions(-) diff --git a/docs/spec.md b/docs/spec.md index 443c7e588..364a32c11 100644 --- a/docs/spec.md +++ b/docs/spec.md @@ -1706,8 +1706,11 @@ The `coordinax.charts` module provides the chart-facing API for representing poi `("length", "length / time**0.5", "length", "speed", "length / time**0.5", "speed")`. - `poincarepolar6d` is the pre-defined `PoincarePolar6D()` instance. - Transition behavior currently registered in-core: - identity transform only (`pt_map(p, PoincarePolar6D, PoincarePolar6D) -> p`). - - No dedicated Cartesian projection dispatch is currently defined for this chart family. + - identity transform (`pt_map(p, PoincarePolar6D, PoincarePolar6D) -> p`). + - gala forward map from a Cartesian phase-space product chart + (`CartesianProductChart((cart3d, cart3d))` = [position, velocity]) to + `PoincarePolar6D`, and its partial inverse back (exact only for + `Lz >= 0`, since `sqrt(|Lz|)` discards `sign(Lz)`).
diff --git a/src/coordinax/_src/charts/register_ptmap.py b/src/coordinax/_src/charts/register_ptmap.py index 412be04ff..32ab4075b 100644 --- a/src/coordinax/_src/charts/register_ptmap.py +++ b/src/coordinax/_src/charts/register_ptmap.py @@ -1773,7 +1773,8 @@ def pt_map( ``pp_phi = sqrt(2|Lz|) cos(phi)``, ``pp_phidot = sqrt(2|Lz|) sin(phi)``, ``dt_z = vz``. - Forward only: ``sqrt(|Lz|)`` discards ``sign(Lz)``, so there is no inverse. + ``sqrt(|Lz|)`` discards ``sign(Lz)``, so there is no *global* inverse. A + *partial* inverse (assuming ``Lz >= 0``) is registered below. >>> import coordinax.charts as cxc >>> import unxt as u @@ -1803,6 +1804,8 @@ def pt_map( """ del usys + assert from_M == from_chart.M # noqa: S101 + assert to_M == to_chart.M # noqa: S101 if len(from_chart.factors) != 2 or not all( isinstance(f, Cart3D) for f in from_chart.factors ): @@ -1821,11 +1824,14 @@ def pt_map( phi = jnp.atan2(x, y) # gala convention: azimuth from +y lz = x * vy - y * vx s = jnp.sqrt(2 * jnp.abs(lz)) + # On the axis (rho == 0) the numerator x*vx + y*vy is also 0; guard the + # denominator so dt_rho == 0 there by convention instead of 0/0 -> NaN. + dt_rho = (x * vx + y * vy) / jnp.where(rho == 0, jnp.ones_like(rho), rho) return { "rho": rho, "pp_phi": s * jnp.cos(phi), "z": z, - "dt_rho": (x * vx + y * vy) / rho, + "dt_rho": dt_rho, "pp_phidot": s * jnp.sin(phi), "dt_z": vz, } @@ -1872,6 +1878,8 @@ def pt_map( """ del usys + assert from_M == from_chart.M # noqa: S101 + assert to_M == to_chart.M # noqa: S101 if len(to_chart.factors) != 2 or not all( isinstance(f, Cart3D) for f in to_chart.factors ): From 08ecf93d6d9cc4a55f28f2e4960d736964fa010c Mon Sep 17 00:00:00 2001 From: nstarman Date: Mon, 27 Jul 2026 11:54:08 -0400 Subject: [PATCH 5/9] =?UTF-8?q?=F0=9F=93=9D=20docs(spec):=20reword=20Poinc?= =?UTF-8?q?arePolar6D=20transition=20to=20avoid=20invalid=20ctor=20call?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit `CartesianProductChart((cart3d, cart3d))` omitted the required factor_names arg; describe it as "a CartesianProductChart with factors (cart3d, cart3d)". Co-Authored-By: Claude Opus 4.8 --- docs/spec.md | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/docs/spec.md b/docs/spec.md index 364a32c11..2c0cf9767 100644 --- a/docs/spec.md +++ b/docs/spec.md @@ -1708,9 +1708,9 @@ The `coordinax.charts` module provides the chart-facing API for representing poi - Transition behavior currently registered in-core: - identity transform (`pt_map(p, PoincarePolar6D, PoincarePolar6D) -> p`). - gala forward map from a Cartesian phase-space product chart - (`CartesianProductChart((cart3d, cart3d))` = [position, velocity]) to - `PoincarePolar6D`, and its partial inverse back (exact only for - `Lz >= 0`, since `sqrt(|Lz|)` discards `sign(Lz)`). + (a `CartesianProductChart` with factors `(cart3d, cart3d)` = + [position, velocity]) to `PoincarePolar6D`, and its partial inverse + back (exact only for `Lz >= 0`, since `sqrt(|Lz|)` discards `sign(Lz)`).
From 0f1ac3ac8db878a9e319937c86de87cf3d9b85f1 Mon Sep 17 00:00:00 2001 From: Nathaniel Starkman Date: Mon, 27 Jul 2026 13:49:24 -0400 Subject: [PATCH 6/9] Potential fix for pull request finding Co-authored-by: Copilot Autofix powered by AI <175728472+Copilot@users.noreply.github.com> Signed-off-by: Nathaniel Starkman --- src/coordinax/_src/charts/register_ptmap.py | 5 +++-- 1 file changed, 3 insertions(+), 2 deletions(-) diff --git a/src/coordinax/_src/charts/register_ptmap.py b/src/coordinax/_src/charts/register_ptmap.py index 32ab4075b..a5c8b513d 100644 --- a/src/coordinax/_src/charts/register_ptmap.py +++ b/src/coordinax/_src/charts/register_ptmap.py @@ -1897,9 +1897,10 @@ def pt_map( sinp, cosp = jnp.sin(phi), jnp.cos(phi) pos = {"x": rho * sinp, "y": rho * cosp, "z": z} + lz_over_rho = jnp.where(rho == 0, jnp.zeros_like(dt_rho), lz / rho) vel = { - "x": sinp * dt_rho - cosp * lz / rho, - "y": cosp * dt_rho + sinp * lz / rho, + "x": sinp * dt_rho - cosp * lz_over_rho, + "y": cosp * dt_rho + sinp * lz_over_rho, "z": dt_z, } return cast("CDict", to_chart.merge_components((pos, vel))) From 8426719058b936c33f5a67a2a6f0dc10f6edd994 Mon Sep 17 00:00:00 2001 From: nstarman Date: Mon, 27 Jul 2026 16:01:02 -0400 Subject: [PATCH 7/9] =?UTF-8?q?=E2=99=BB=EF=B8=8F=20refactor(charts):=20gr?= =?UTF-8?q?ad-safe=20axis-singularity=20division=20via=20shared=20helper?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Consolidate the "divide by a coordinate that vanishes on an axis" pattern into `_ratio_zero_on_axis`, using the double-`where` idiom so both the value and its reverse-mode gradient stay finite on the axis. A plain `jnp.where(denom == 0, 0, num / denom)` returns the right value but leaks NaN into gradients (`0 * inf`). Applied to all three sites in register_ptmap.py: - LonCosLat -> Cart3D pole (`cos(lat) == 0`) - PoincarePolar6D forward `dt_rho` (already grad-safe; unified for clarity) - PoincarePolar6D partial inverse `lz / rho` (was the plain, grad-unsafe form) Adds a regression test asserting finite value + gradient at the singularity. Co-Authored-By: Claude Opus 4.8 --- src/coordinax/_src/charts/register_ptmap.py | 27 ++++++++--- tests/unit/charts/test_register_realize.py | 50 +++++++++++++++++++++ 2 files changed, 71 insertions(+), 6 deletions(-) diff --git a/src/coordinax/_src/charts/register_ptmap.py b/src/coordinax/_src/charts/register_ptmap.py index a5c8b513d..9c3e16fd5 100644 --- a/src/coordinax/_src/charts/register_ptmap.py +++ b/src/coordinax/_src/charts/register_ptmap.py @@ -40,6 +40,20 @@ from coordinax._src.product.manifold import CartesianProductManifold from coordinax._src.utils import uconvert_to_rad + +def _ratio_zero_on_axis(num: Array, denom: Array, /) -> Array: + """``num / denom``, defined as 0 where ``denom == 0`` (coordinate singularity). + + Uses the double-``where`` idiom so that both the value *and* its gradient are + finite on the axis. A plain ``jnp.where(denom == 0, 0, num / denom)`` still + evaluates ``num / denom`` in the unselected branch, leaking ``NaN`` into + reverse-mode gradients (``0 * inf``); guarding the denominator first avoids it. + """ + safe = jnp.where(denom == 0, jnp.ones_like(denom), denom) + ratio = num / safe + return jnp.where(denom == 0, jnp.zeros_like(ratio), ratio) + + ##################################################################### # Point transformations @@ -626,9 +640,9 @@ def pt_map( lon_coslat, r_ = p["lon_coslat"], p["distance"] lat = uconvert_to_rad(p["lat"], usys) - # Handle the poles where cos(lat) == 0 + # Longitude is undefined at the poles (cos(lat) == 0); set lon = 0 there. coslat = jnp.cos(lat) - lon = jnp.where(coslat == 0, 0, lon_coslat / coslat) + lon = _ratio_zero_on_axis(lon_coslat, coslat) lon = uconvert_to_rad(lon, usys) # Convert to Cartesian x = r_ * jnp.cos(lat) * jnp.cos(lon) @@ -1824,9 +1838,9 @@ def pt_map( phi = jnp.atan2(x, y) # gala convention: azimuth from +y lz = x * vy - y * vx s = jnp.sqrt(2 * jnp.abs(lz)) - # On the axis (rho == 0) the numerator x*vx + y*vy is also 0; guard the - # denominator so dt_rho == 0 there by convention instead of 0/0 -> NaN. - dt_rho = (x * vx + y * vy) / jnp.where(rho == 0, jnp.ones_like(rho), rho) + # On the axis (rho == 0) the numerator x*vx + y*vy is also 0; define + # dt_rho == 0 there by convention instead of 0/0 -> NaN. + dt_rho = _ratio_zero_on_axis(x * vx + y * vy, rho) return { "rho": rho, "pp_phi": s * jnp.cos(phi), @@ -1897,7 +1911,8 @@ def pt_map( sinp, cosp = jnp.sin(phi), jnp.cos(phi) pos = {"x": rho * sinp, "y": rho * cosp, "z": z} - lz_over_rho = jnp.where(rho == 0, jnp.zeros_like(dt_rho), lz / rho) + # On the axis (rho == 0, where lz == 0 too) define lz/rho == 0 by convention. + lz_over_rho = _ratio_zero_on_axis(lz, rho) vel = { "x": sinp * dt_rho - cosp * lz_over_rho, "y": cosp * dt_rho + sinp * lz_over_rho, diff --git a/tests/unit/charts/test_register_realize.py b/tests/unit/charts/test_register_realize.py index d7ad2872d..6f34b6358 100644 --- a/tests/unit/charts/test_register_realize.py +++ b/tests/unit/charts/test_register_realize.py @@ -3,6 +3,7 @@ cartesian_chart, pt_map. """ +import jax import jax.numpy as jnp import plum import pytest @@ -11,7 +12,9 @@ import unxt as u import coordinax.charts as cxc +import coordinax.manifolds as cxm import coordinaxs.hypothesis.main as cxst +from coordinax._src.charts.register_ptmap import _ratio_zero_on_axis # ============================================================================= # cartesian_chart @@ -122,3 +125,50 @@ def test_spherical_to_prolate(self) -> None: p = {"r": u.Q(3.0, "m"), "theta": u.Q(0.6, "rad"), "phi": u.Q(0.4, "rad")} out = cxc.pt_map(p, cxc.sph3d, prolate) assert set(out) == {"mu", "nu", "phi"} + + +# ============================================================================= +# Coordinate-singularity division is grad-safe (double-where idiom) +# ============================================================================= + + +class TestAxisSingularityGradSafe: + """``pt_map`` divisions at a coordinate singularity stay finite in value and grad. + + A plain ``jnp.where(denom == 0, 0, num / denom)`` returns the right value but + leaks ``NaN`` into reverse-mode gradients; these guard against regressing to it. + """ + + def test_ratio_helper_value_and_grad(self): + def f(d): + return _ratio_zero_on_axis(2.0 * d, d) + + d0 = jnp.asarray(0.0) + assert float(f(d0)) == 0.0 + assert bool(jnp.isfinite(jax.grad(f)(d0))) + + def test_loncoslat_pole_is_finite(self): + """LonCosLat -> Cart3D at the pole (cos(lat) == 0) is finite.""" + p = { + "lon_coslat": u.Q(0.4, "rad"), + "lat": u.Q(90.0, "deg"), + "distance": u.Q(2.0, "m"), + } + out = cxc.pt_map(p, cxm.R3, cxc.loncoslat_sph3d, cxm.R3, cxc.cart3d) + assert all(bool(jnp.isfinite(v.value)) for v in out.values()) + + def test_poincarepolar6d_axis_roundtrip_is_finite(self): + """Forward + inverse through the rho == 0 axis stays finite (dt_rho, lz/rho).""" + ps = cxc.CartesianProductChart((cxc.cart3d, cxc.cart3d), ("q", "p")) + q = { + "q.x": u.Q(0.0, "kpc"), + "q.y": u.Q(0.0, "kpc"), + "q.z": u.Q(5.0, "kpc"), + "p.x": u.Q(0.0, "kpc/Myr"), + "p.y": u.Q(0.0, "kpc/Myr"), + "p.z": u.Q(0.5, "kpc/Myr"), + } + pp = cxc.pt_map(q, ps.M, ps, cxc.poincarepolar6d.M, cxc.poincarepolar6d) + back = cxc.pt_map(pp, cxc.poincarepolar6d.M, cxc.poincarepolar6d, ps.M, ps) + assert all(bool(jnp.isfinite(v.value)) for v in pp.values()) + assert all(bool(jnp.isfinite(v.value)) for v in back.values()) From 04f708cbd63c66976a370c47dd69463ab9d264ae Mon Sep 17 00:00:00 2001 From: nstarman Date: Mon, 27 Jul 2026 17:24:24 -0400 Subject: [PATCH 8/9] =?UTF-8?q?=F0=9F=93=9D=20docs(charts):=20clarify=20Lo?= =?UTF-8?q?nCosLat=20pole=20guard=20fires=20only=20on=20exact=20cos(lat)?= =?UTF-8?q?=3D=3D0?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit The _ratio_zero_on_axis guard triggers only when cos(lat) is exactly 0; a floating-point near-pole value (cos(pi/2) ~= 6e-17) still divides but the cos(lat) factor multiplies the result back, keeping x, y finite. Co-Authored-By: Claude Opus 4.8 --- src/coordinax/_src/charts/register_ptmap.py | 5 ++++- 1 file changed, 4 insertions(+), 1 deletion(-) diff --git a/src/coordinax/_src/charts/register_ptmap.py b/src/coordinax/_src/charts/register_ptmap.py index 9c3e16fd5..eb0da26ba 100644 --- a/src/coordinax/_src/charts/register_ptmap.py +++ b/src/coordinax/_src/charts/register_ptmap.py @@ -640,7 +640,10 @@ def pt_map( lon_coslat, r_ = p["lon_coslat"], p["distance"] lat = uconvert_to_rad(p["lat"], usys) - # Longitude is undefined at the poles (cos(lat) == 0); set lon = 0 there. + # Longitude is undefined at the poles. The guard fires only when cos(lat) is + # *exactly* 0 (giving lon = 0); a floating-point near-pole value (e.g. + # cos(pi/2) ~= 6e-17) still divides, but the cos(lat) factor below multiplies + # the result back so x, y stay finite. coslat = jnp.cos(lat) lon = _ratio_zero_on_axis(lon_coslat, coslat) lon = uconvert_to_rad(lon, usys) From f93b212b3abec1a7a666f3b532f61387b7bf36ca Mon Sep 17 00:00:00 2001 From: nstarman Date: Mon, 27 Jul 2026 21:56:31 -0400 Subject: [PATCH 9/9] =?UTF-8?q?=E2=99=BB=EF=B8=8F=20refactor(charts):=20de?= =?UTF-8?q?dupe=20PoincarePolar6D=20phase-space=20validation;=20skip=20sqr?= =?UTF-8?q?t?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit - Extract `_require_cart3d_phase_space(chart, *, direction)` used by both the forward and inverse PoincarePolar6D dispatches (single source for the two-factor Cart3D check and error wording). - Compute the inverse-map `lz = (pp_phi**2 + pp_phidot**2) / 2` directly instead of `hypot(...)**2 / 2`, avoiding an unnecessary sqrt. Co-Authored-By: Claude Opus 4.8 --- src/coordinax/_src/charts/register_ptmap.py | 43 +++++++++++---------- 1 file changed, 23 insertions(+), 20 deletions(-) diff --git a/src/coordinax/_src/charts/register_ptmap.py b/src/coordinax/_src/charts/register_ptmap.py index eb0da26ba..c9c5393c9 100644 --- a/src/coordinax/_src/charts/register_ptmap.py +++ b/src/coordinax/_src/charts/register_ptmap.py @@ -54,6 +54,24 @@ def _ratio_zero_on_axis(num: Array, denom: Array, /) -> Array: return jnp.where(denom == 0, jnp.zeros_like(ratio), ratio) +def _require_cart3d_phase_space(chart: Any, /, *, direction: str) -> None: + """Validate a two-factor ``Cart3D`` Cartesian phase-space product chart. + + ``direction`` is ``"to"`` (``chart`` is the source of ``pt_map`` *to* + ``PoincarePolar6D``) or ``"from"`` (``chart`` is the target of ``pt_map`` + *from* it); it only selects the error wording. Raises ``NotImplementedError`` + unless ``chart`` has exactly two ``Cart3D`` factors [position, velocity]. + """ + if len(chart.factors) != 2 or not all(isinstance(f, Cart3D) for f in chart.factors): + role = "source" if direction == "to" else "target" + msg = ( + f"pt_map {direction} PoincarePolar6D requires a Cartesian phase-space " + f"{role}: a two-factor CartesianProductChart of (Cart3D, Cart3D) " + f"[position, velocity]; got factors {chart.factors!r}." + ) + raise NotImplementedError(msg) + + ##################################################################### # Point transformations @@ -1823,15 +1841,7 @@ def pt_map( del usys assert from_M == from_chart.M # noqa: S101 assert to_M == to_chart.M # noqa: S101 - if len(from_chart.factors) != 2 or not all( - isinstance(f, Cart3D) for f in from_chart.factors - ): - msg = ( - "pt_map to PoincarePolar6D requires a Cartesian phase-space source: a " - "two-factor CartesianProductChart of (Cart3D, Cart3D) [position, " - f"velocity]; got factors {from_chart.factors!r}." - ) - raise NotImplementedError(msg) + _require_cart3d_phase_space(from_chart, direction="to") pos, vel = from_chart.split_components(p) x, y, z = pos["x"], pos["y"], pos["z"] @@ -1897,20 +1907,13 @@ def pt_map( del usys assert from_M == from_chart.M # noqa: S101 assert to_M == to_chart.M # noqa: S101 - if len(to_chart.factors) != 2 or not all( - isinstance(f, Cart3D) for f in to_chart.factors - ): - msg = ( - "pt_map from PoincarePolar6D requires a Cartesian phase-space target: a " - "two-factor CartesianProductChart of (Cart3D, Cart3D) [position, " - f"velocity]; got factors {to_chart.factors!r}." - ) - raise NotImplementedError(msg) + _require_cart3d_phase_space(to_chart, direction="from") rho, z, dt_rho, dt_z = p["rho"], p["z"], p["dt_rho"], p["dt_z"] - s = jnp.hypot(p["pp_phi"], p["pp_phidot"]) phi = jnp.atan2(p["pp_phidot"], p["pp_phi"]) - lz = s**2 / 2 # sign(Lz) is not recoverable from the forward map; take Lz >= 0 + # Lz = s²/2 with s = hypot(pp_phi, pp_phidot); compute directly to skip the + # sqrt (sign(Lz) is not recoverable from the forward map, so take Lz >= 0). + lz = (p["pp_phi"] ** 2 + p["pp_phidot"] ** 2) / 2 sinp, cosp = jnp.sin(phi), jnp.cos(phi) pos = {"x": rho * sinp, "y": rho * cosp, "z": z}