✨ feat: representations

Signed-off-by: nstarman <nstarman@users.noreply.github.com>

Author: nstarman

README.md

@@ -49,8 +49,8 @@ pip install coordinax
   components.
 - Metrics: a bilinear form _g_ on the tangent space defining inner products and
   norms (Euclidean, sphere intrinsic, Minkowski).
-- Transformations: convert between coordinate charts and roles, including
-  velocity and acceleration transformations
+- Transformations: convert between coordinate charts and representations,
+  including velocity and acceleration transformations
 - Operators: frame-aware vector operations (rotation, boost, translation, etc.)
 - `PointedVector`: combines a vector with a reference point or frame for context
 - Frames: reference systems (ICRS, Galactocentric, etc.) and their
@@ -62,7 +62,7 @@ pip install coordinax
 
 ```
 import coordinax.charts as cxc  # Charts
-import coordinax.roles as cxr  # Roles
+import coordinax.representations as cxr  # Representations
 import coordinax.metrics as cxm  # Metric
 import coordinax.tangents as cxt  # Transformations
 ```
@@ -126,7 +126,7 @@ frames.
 
 ```
 import jax.numpy as jnp
-import coordinax.roles as cxr
+import coordinax.representations as cxr
 from coordinax.objs import Vector
 
 q = {"x": u.Q(1.0, "km"), "y": u.Q(2.0, "km"), "z": u.Q(3.0, "km")}
@@ -139,7 +139,7 @@ bool(jnp.allclose(u.ustrip("km/s", v_back.data["x"]),
                   u.ustrip("km/s", vvec.data["x"])))
 ```
 
-Different vector roles transform via different mechanisms:
+Different vector representations transform via different mechanisms:
 
 | Role       | Transformation                      | Requires Base Point? |
 | ---------- | ----------------------------------- | -------------------- |
@@ -253,9 +253,9 @@ disp_from_origin = cx.as_disp(new_pos)
 
 ## Metrics and Representations
 
-Representations are coordinate charts; roles (PhysDisp, PhysVel, PhysAcc, ...)
-give vectors their physical meaning. A chart’s default metric defines how
-physical components are interpreted.
+Representations are coordinate charts; representations (PhysDisp, PhysVel,
+PhysAcc, ...) give vectors their physical meaning. A chart’s default metric
+defines how physical components are interpreted.
 
 ## Citation
 

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

@@ -0,0 +1,18 @@
+"""Representations."""
+
+__all__ = ("vconvert",)
+
+from typing import Any
+
+import plum
+
+
+@plum.dispatch.abstract
+def vconvert(*args: Any, **kwargs: Any) -> Any:
+    """Transform the current vector to the target chart.
+
+    This is an abstract API definition. See the main coordinax package for
+    concrete implementations and usage examples.
+
+    """
+    raise NotImplementedError  # pragma: no cover

pyproject.toml

@@ -297,6 +297,7 @@ ignore = [
   "ANN401",  # Dynamically typed expressions (typing.Any) are disallowed
   "ARG001",  # Unused function argument  # TODO: resolve
   "B008",    # Do not perform function calls in argument defaults
+  "B024",    # X is an ABC, but it has no abstract methods or properties
   "COM812",  # <- for ruff.format
   "D103",    # Missing docstring in public function  # TODO: resolve
   "D105",    # Missing docstring in magic method
@@ -339,6 +340,7 @@ ignore = [
 "coordinax.core"            = "cx"
 "coordinax.distances"       = "cxd"
 "coordinax.hypothesis.core" = "cxst"
+"coordinax.representations" = "cxr"
 equinox                     = "eqx"
 functools                   = "ft"
 unxt                        = "u"

src/coordinax/representations/__init__.py

@@ -0,0 +1,41 @@
+"""`coordinax.representations` Module."""
+
+__all__ = (
+    "vconvert",
+    # Representations
+    "Representation",
+    "point",
+    # Geometric Kinds
+    "AbstractGeometry",
+    "PointGeometry",
+    "point_geom",
+    # Coordinate Bases
+    "AbstractBasis",
+    "NoBasis",
+    "nobasis",
+    # Semantic Kinds
+    "AbstractSemanticKind",
+    "Location",
+    "location",
+)
+
+from ._setup_package import install_import_hook
+
+with install_import_hook("coordinax.representations"):
+    from ._src import (
+        AbstractBasis,
+        AbstractGeometry,
+        AbstractSemanticKind,
+        Location,
+        NoBasis,
+        PointGeometry,
+        Representation,
+        location,
+        nobasis,
+        point,
+        point_geom,
+    )
+    from coordinax.api.representations import vconvert
+
+
+del install_import_hook

src/coordinax/representations/_setup_package.py

@@ -0,0 +1,61 @@
+"""Package setup information.
+
+Note that this module is NOT public API nor are any of its contents.
+Stability is NOT guaranteed.
+This module exposes package setup information for the `unxt` package.
+
+"""
+
+__all__: tuple[str, ...] = ("RUNTIME_TYPECHECKER", "install_import_hook")
+
+import contextlib
+import os
+
+from collections.abc import Sequence
+from jaxtyping import install_import_hook as _install_import_hook
+from typing import Any, Final, Literal
+
+_RUNTIME_TYPECHECKER: str | None | Literal[False]
+match os.getenv("COORDINAX_ENABLE_RUNTIME_TYPECHECKING", "False"):
+    case "False":
+        _RUNTIME_TYPECHECKER = False
+    case "None":
+        _RUNTIME_TYPECHECKER = None
+    case str() as _name:
+        _RUNTIME_TYPECHECKER = _name
+
+RUNTIME_TYPECHECKER: Final[str | None | Literal[False]] = _RUNTIME_TYPECHECKER
+"""Runtime type checking variable "COORDINAX_ENABLE_RUNTIME_TYPECHECKING".
+
+Set to "False" to disable runtime typechecking (default).
+Set to "None" to only enable typechecking for `@jaxtyped`-decorated functions.
+Set to "beartype.beartype" to enable runtime typechecking.
+
+See https://docs.kidger.site/jaxtyping/api/runtime-type-checking for more
+information on options.
+
+
+"""
+
+
+def install_import_hook(
+    modules: str | Sequence[str], /
+) -> contextlib.AbstractContextManager[Any, None]:
+    """Install the jaxtyping import hook for the given modules.
+
+    Parameters
+    ----------
+    modules
+        Module name or sequence of module names to install the import hook for.
+
+    Returns
+    -------
+    contextlib.AbstractContextManager
+        Context manager that installs the import hook on entry and removes it on exit.
+
+    """
+    return (
+        _install_import_hook(modules, RUNTIME_TYPECHECKER)
+        if RUNTIME_TYPECHECKER is not False
+        else contextlib.nullcontext()
+    )

src/coordinax/representations/_src/__init__.py

@@ -0,0 +1,7 @@
+"""Vector."""
+
+from .basis import *
+from .core import *
+from .geom import *
+from .representations import *
+from .semantics import *

src/coordinax/representations/_src/basis.py

@@ -0,0 +1,31 @@
+"""Vector."""
+
+__all__ = ("AbstractBasis", "NoBasis", "nobasis")
+
+
+import abc
+import dataclasses
+
+from typing import final
+
+import jax.tree_util as jtu
+
+
+@jtu.register_static
+class AbstractBasis(metaclass=abc.ABCMeta):
+    pass
+
+
+@jtu.register_static
+@final
+@dataclasses.dataclass(frozen=True, slots=True)
+class NoBasis(AbstractBasis):
+    """No basis.
+
+    Points, and other geometric objects that do not have a vector space
+    structure do not have a basis.
+
+    """
+
+
+nobasis = NoBasis()

src/coordinax/representations/_src/core.py

@@ -0,0 +1,59 @@
+"""Vector Conversion."""
+
+__all__ = ("vconvert",)
+
+import plum
+
+import coordinax.api.representations as api
+import coordinax.charts as cxc
+from .basis import AbstractBasis, NoBasis
+from .geom import PointGeometry
+from .representations import Representation
+from .semantics import AbstractSemanticKind, Location
+from coordinax.internal.custom_types import CDict
+
+
+@plum.dispatch
+def vconvert(
+    rep: Representation,
+    to_chart: cxc.AbstractChart,
+    from_chart: cxc.AbstractChart,
+    x: CDict,
+) -> CDict:
+    # redispatch on the combination of GeometryKind, Basis, and SemanticKind
+    return api.vconvert(
+        rep.geom_kind, rep.basis, rep.semantic_kind, to_chart, from_chart, x
+    )
+
+
+# =======================================================================
+
+
+@plum.dispatch
+def vconvert(
+    geom_kind: PointGeometry,
+    basis: AbstractBasis,
+    semantic_kind: AbstractSemanticKind,
+    to_chart: cxc.AbstractChart,
+    from_chart: cxc.AbstractChart,
+    x: CDict,
+) -> CDict:
+    msg = (
+        "For point-role representations, vconvert is only implemented for the "
+        "combination of PointGeometry, NoBasis, and Location semantic kind. "
+    )
+    raise TypeError(msg)
+
+
+@plum.dispatch
+def vconvert(
+    geom_kind: PointGeometry,
+    basis: NoBasis,
+    semantic_kind: Location,
+    to_chart: cxc.AbstractChart,
+    from_chart: cxc.AbstractChart,
+    x: CDict,
+) -> CDict:
+    """Convert a point-role representation from one chart to another."""
+    del geom_kind, basis, semantic_kind  # only used for dispatching to here.
+    return cxc.coord_map(to_chart, from_chart, x)

src/coordinax/representations/_src/geom.py

@@ -0,0 +1,41 @@
+"""Vector."""
+
+__all__ = ("AbstractGeometry", "PointGeometry", "point_geom")
+
+
+import abc
+import dataclasses
+
+from typing import final
+
+import jax.tree_util as jtu
+
+
+@jtu.register_static
+class AbstractGeometry(metaclass=abc.ABCMeta):
+    """Flag for geometric role (position, velocity, acceleration, etc.)."""
+
+
+@jtu.register_static
+@final
+@dataclasses.dataclass(frozen=True, slots=True)
+class PointGeometry(AbstractGeometry):
+    r"""Point geometric kind.
+
+    Mathematical Definition:
+
+    A **point** is an element of a manifold or affine space $M$. Points do not
+    form a vector space in general; e.g., on a curved manifold you cannot
+    meaningfully add two points.
+
+    - In Euclidean space $\mathbb{R}^n$, points can be identified with vectors
+      from the origin, but this identification is basis-dependent.
+    - On manifolds (e.g., a sphere), points are elements of the manifold and
+      have no additive structure.
+    - Point coordinates may have mixed dimensions (e.g., spherical: length +
+      angles).
+
+    """
+
+
+point_geom = PointGeometry()

src/coordinax/representations/_src/representations.py

@@ -0,0 +1,26 @@
+"""Vector."""
+
+__all__ = ("AbstractGeometry",)
+
+
+import dataclasses
+
+from typing import final
+
+import jax.tree_util as jtu
+
+from .basis import AbstractBasis, nobasis
+from .geom import AbstractGeometry, point_geom
+from .semantics import AbstractSemanticKind, location
+
+
+@jtu.register_static
+@final
+@dataclasses.dataclass(frozen=True, slots=True)
+class Representation:
+    geom_kind: AbstractGeometry
+    basis: AbstractBasis
+    semantic_kind: AbstractSemanticKind
+
+
+point = Representation(point_geom, nobasis, location)