WIP

Signed-off-by: nstarman <[email protected]>

Author: nstarman

docs/conf.py

@@ -6,6 +6,7 @@
 """
 
 from datetime import datetime
+
 from typing import Any
 
 import pytz

pyproject.toml

@@ -290,7 +290,8 @@ COORDINAX_ENABLE_RUNTIME_TYPECHECKING = "beartype.beartype"
 
     [tool.ruff.lint.isort]
       combine-as-imports = true
-      extra-standard-library = ["typing_extensions"]
+      sections = { typing = ["typing", "typing_extensions", "jaxtyping"] }
+      section-order = ["future", "standard-library", "typing", "third-party", "first-party", "local-folder"]
       known-first-party = ["dataclassish", "optional_dependencies", "quaxed", "unxt", "xmmutablemap"]
       known-local-folder = ["coordinax"]
 
@@ -301,3 +302,6 @@ constraint-dependencies = [
     "jax<0.7",
     "jaxlib<0.7",
 ]
+
+[tool.uv.sources]
+unxt = { path = "../unxt", editable = true }

src/coordinax/__init__.py

@@ -5,6 +5,7 @@
     # modules
     "angle",
     "distance",
+    "r",
     "vecs",
     "ops",
     "frames",
@@ -13,6 +14,7 @@
     # common vecs objects
     "vector",
     "vconvert",
+    "Vector",
     "CartesianPos3D",
     "CartesianVel3D",
     "SphericalPos",
@@ -26,20 +28,11 @@
 from .setup_package import install_import_hook
 
 with install_import_hook("coordinax"):
-    from . import angle, distance, frames, ops, vecs
+    from . import angle, distance, frames, ops, r, vecs
     from ._version import version as __version__  # noqa: F401
     from .distance import Distance
     from .frames import Coordinate
-    from .vecs import (
-        CartesianPos3D,
-        CartesianVel3D,
-        FourVector,
-        KinematicSpace,
-        SphericalPos,
-        SphericalVel,
-        vconvert,
-        vector,
-    )
+    from .vecs import KinematicSpace, Vector, vconvert, vector
 
     # isort: split
     # Interoperability

src/coordinax/_coordinax_space_frames/frame_transforms.py

@@ -3,9 +3,9 @@
 __all__: tuple[str, ...] = ()
 
 
+from jaxtyping import Array, Shaped
 from typing import TypeAlias
 
-from jaxtyping import Array, Shaped
 from plum import dispatch
 
 import quaxed.numpy as jnp

src/coordinax/_coordinax_space_frames/galactocentric.py

@@ -3,10 +3,10 @@
 __all__ = ("Galactocentric",)
 
 
+from jaxtyping import Array, Shaped
 from typing import ClassVar, TypeAlias, final
 
 import equinox as eqx
-from jaxtyping import Array, Shaped
 
 import unxt as u
 from dataclassish.converters import Unless

src/coordinax/_coordinax_space_frames/icrs.py

@@ -3,9 +3,8 @@
 __all__ = ("ICRS",)
 
 
-from typing import TypeAlias, final
-
 from jaxtyping import Array, Shaped
+from typing import TypeAlias, final
 
 import unxt as u
 

src/coordinax/_interop/coordinax_interop_astropy/converters.py

@@ -5,78 +5,75 @@
 __all__: tuple[str, ...] = ()
 
 
+from jaxtyping import Shaped
 from typing import cast
 
 import astropy.coordinates as apyc
 import astropy.units as apyu
-from jaxtyping import Shaped
 from plum import conversion_method, convert
 
 import unxt as u
 
 import coordinax as cx
+import coordinax.r as cxr
+import coordinax.vecs as cxv
 
 #####################################################################
 
 # =====================================
 # Quantity
 
 
-@conversion_method(cx.vecs.AbstractPos3D, apyu.Quantity)  # type: ignore[arg-type]
-def vec_to_q(obj: cx.vecs.AbstractPos3D, /) -> Shaped[apyu.Quantity, "*batch 3"]:
+@conversion_method(cxv.Vector, apyu.Quantity)  # type: ignore[arg-type]
+def vec_to_q(obj: cxv.Vector, /) -> Shaped[apyu.Quantity, "*batch 3"]:
     """`coordinax.AbstractPos3D` -> `astropy.units.Quantity`.
 
     Examples
     --------
-    >>> import coordinax as cx
+    >>> import coordinax.vecs as cxv
     >>> from plum import convert
     >>> import astropy.units as apyu
 
-    >>> vec = cx.CartesianPos3D.from_([1, 2, 3], "km")
+    >>> vec = cxv.CartesianPos3D.from_([1, 2, 3], "km")
     >>> convert(vec, apyu.Quantity)
     <Quantity [1., 2., 3.] km>
 
-    >>> vec = cx.SphericalPos(r=apyu.Quantity(1, unit="km"),
+    >>> vec = cxv.SphericalPos(r=apyu.Quantity(1, unit="km"),
     ...                       theta=apyu.Quantity(2, unit="deg"),
     ...                       phi=apyu.Quantity(3, unit="deg"))
     >>> convert(vec, apyu.Quantity)
     <Quantity [0.03485167, 0.0018265 , 0.99939084] km>
 
-    >>> vec = cx.vecs.CylindricalPos(rho=apyu.Quantity(1, unit="km"),
+    >>> vec = cxv.CylindricalPos(rho=apyu.Quantity(1, unit="km"),
     ...                              phi=apyu.Quantity(2, unit="deg"),
     ...                              z=apyu.Quantity(3, unit="m"))
     >>> convert(vec, apyu.Quantity)
     <Quantity [0.99939084, 0.0348995 , 0.003     ] km>
 
-    """
-    return convert(convert(obj, u.Quantity), apyu.Quantity)
-
-
-@conversion_method(cx.vecs.CartesianAcc3D, apyu.Quantity)  # type: ignore[arg-type]
-@conversion_method(cx.CartesianVel3D, apyu.Quantity)  # type: ignore[arg-type]
-def vec_diff_to_q(
-    obj: cx.CartesianVel3D | cx.vecs.CartesianAcc3D, /
-) -> Shaped[apyu.Quantity, "*batch 3"]:
-    """`coordinax.CartesianVel3D` -> `astropy.units.Quantity`.
-
-    Examples
-    --------
-    >>> import coordinax as cx
-    >>> from plum import convert
-    >>> from astropy.units import Quantity as AstropyQuantity
-
-    >>> dif = cx.CartesianVel3D.from_([1, 2, 3], "km/s")
+    >>> dif = cxv.CartesianVel3D.from_([1, 2, 3], "km/s")
     >>> convert(dif, AstropyQuantity)
     <Quantity [1., 2., 3.] km / s>
 
-    >>> dif2 = cx.vecs.CartesianAcc3D.from_([1, 2, 3], "km/s2")
+    >>> dif2 = cxv.CartesianAcc3D.from_([1, 2, 3], "km/s2")
     >>> convert(dif2, AstropyQuantity)
     <Quantity [1., 2., 3.] km / s2>
 
     """
     return convert(convert(obj, u.Quantity), apyu.Quantity)
 
 
+# =====================================
+
+
+@conversion_method(cxv.Vector, apyc.CartesianRepresentation)
+def convert_vector_to_astropy(obj: cxv.Vector, /) -> apyc.CartesianRepresentation:
+    return apyc.CartesianRepresentation(
+        x=convert(obj["x"], apyu.Quantity),
+        y=convert(obj["y"], apyu.Quantity),
+        z=convert(obj["z"], apyu.Quantity),
+    )
+
+
 # =====================================
 # CartesianPos3D
 

src/coordinax/_src/custom_types.py

@@ -2,14 +2,14 @@
 
 __all__: tuple[str, ...] = ()
 
+from jaxtyping import Real, Shaped
 from typing import TypeAlias
 
 from astropy.units import (
     CompositeUnit as AstropyCompositeUnit,
     Unit as AstropyUnit,
     UnitBase as AstropyUnitBase,
 )
-from jaxtyping import Real, Shaped
 
 import unxt as u
 

src/coordinax/_src/distances/funcs.py

@@ -7,9 +7,9 @@
 )
 
 
+from jaxtyping import ArrayLike
 from typing import Any
 
-from jaxtyping import ArrayLike
 from plum import dispatch
 
 import quaxed.numpy as jnp

src/coordinax/_src/distances/measures.py

@@ -3,12 +3,13 @@
 __all__ = ("Distance", "DistanceModulus", "Parallax")
 
 from dataclasses import KW_ONLY
+
+from jaxtyping import Array, Shaped
 from typing import Any, final
 
 import equinox as eqx
 import jax.numpy as jnp
 import wadler_lindig as wl
-from jaxtyping import Array, Shaped
 
 import quaxed.numpy as jnp
 import unxt as u