✨ feat: manifolds

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

Author: nstarman

packages/coordinax-api/src/coordinax_api/__init__.py

@@ -16,6 +16,8 @@
     "project_tangent",
     # metrics
     "metric_of",
+    # manifolds
+    "manifold_of",
     # frames
     "frame_of",
     "frame_transform_op",
@@ -55,6 +57,7 @@
 )
 from ._embeddings import embed_point, embed_tangent, project_point, project_tangent
 from ._frames import frame_of, frame_transform_op
+from ._manifolds import manifold_of
 from ._metrics import metric_of
 from ._objs import vconvert
 from ._operators import apply_op, simplify

packages/coordinax-api/src/coordinax_api/_manifolds.py

@@ -0,0 +1,13 @@
+"""Manifold API for coordinax."""
+
+__all__ = ("manifold_of",)
+
+from typing import Any
+
+import plum
+
+
+@plum.dispatch.abstract
+def manifold_of(*args: Any) -> Any:
+    """Return the default manifold associated with the input."""
+    raise NotImplementedError  # pragma: no cover

src/coordinax/__init__.py

@@ -8,6 +8,7 @@
     "distances",
     "embeddings",
     "frames",
+    "manifolds",
     "metrics",
     "objs",
     "ops",
@@ -36,6 +37,7 @@
         distances,
         embeddings,
         frames,
+        manifolds,
         metrics,
         objs,
         ops,

src/coordinax/manifolds/__init__.py

@@ -0,0 +1,18 @@
+"""`coordinax.manifolds` module."""
+
+__all__ = (
+    "AbstractManifold",
+    "Atlas",
+    "Euclidean",
+    "EuclideanAtlas",
+    "manifold_of",
+)
+
+from coordinax import setup_package
+
+with setup_package.install_import_hook("coordinax.manifolds"):
+    from ._src import AbstractManifold, Atlas, Euclidean, EuclideanAtlas
+    from coordinax.api import manifold_of
+
+
+del setup_package

src/coordinax/manifolds/_src/__init__.py

@@ -0,0 +1,4 @@
+"""Manifolds in coordinax."""
+
+from .base import *
+from .euclidean import *

src/coordinax/manifolds/_src/base.py

@@ -0,0 +1,47 @@
+"""Manifold definitions and manifold inference helpers."""
+
+__all__ = (
+    "AbstractManifold",
+    "Atlas",
+)
+
+import abc
+
+from typing import Any, Protocol
+
+import coordinax.charts as cxc
+import coordinax.metrics as cxm
+
+
+class Atlas(Protocol):
+    """Atlas protocol for manifolds."""
+
+    def default_chart(self) -> cxc.AbstractChart[Any, Any]:
+        """Return a default chart from the atlas."""
+        ...
+
+    # def supports(self, chart: cxc.AbstractChart[Any, Any]) -> bool:
+    #     """Return whether the atlas supports the given chart."""
+    #     ...
+
+
+class AbstractManifold(metaclass=abc.ABCMeta):
+    """Abstract manifold interface."""
+
+    dim: int
+    """Intrinsic dimension of the manifold."""
+
+    metric: cxm.AbstractMetric
+    """(Pseudo-)Riemannian metric on the manifold."""
+
+    atlas: Atlas
+    """Charts compatible with this manifold."""
+
+    @property
+    def default_chart(self) -> cxc.AbstractChart[Any, Any]:
+        """Return a default chart from the atlas."""
+        return self.atlas.default_chart()
+
+    # def has_chart(self, chart: cxc.AbstractChart[Any, Any], /) -> bool:
+    #     """Return whether ``chart`` belongs to this manifold atlas."""
+    #     return self.atlas.supports(chart)

src/coordinax/manifolds/_src/euclidean.py

@@ -0,0 +1,55 @@
+"""Euclidean manifolds."""
+
+__all__ = ("Euclidean", "EuclideanAtlas")
+
+from dataclasses import dataclass
+
+from typing import Any, final
+
+import coordinax.charts as cxc
+import coordinax.metrics as cxm
+from .base import AbstractManifold, Atlas
+
+
+@dataclass(frozen=True, slots=True)
+class EuclideanAtlas(Atlas):
+    """Atlas for Euclidean manifolds."""
+
+    dim: int
+    """Dimension of the Euclidean manifold."""
+
+    def default_chart(self) -> cxc.AbstractChart[Any, Any]:
+        # TODO: make this autodetect
+        chart: cxc.AbstractChart[Any, Any]
+        match self.dim:
+            case 0:
+                chart = cxc.cart0d
+            case 1:
+                chart = cxc.cart1d
+            case 2:
+                chart = cxc.cart2d
+            case 3:
+                chart = cxc.cart3d
+            case 6:
+                chart = cxc.poincarepolar6d
+            case _:
+                msg = f"Euclidean({self.dim}) is unsupported for now."
+                raise ValueError(msg)
+        return chart
+
+    # def supports(self, chart: cxc.AbstractChart[Any, Any]) -> bool:
+    #     return any(c == chart for c in self._charts)
+
+
+@final
+@dataclass(frozen=True, slots=True)
+class Euclidean(AbstractManifold):
+    """Euclidean manifold with identity metric."""
+
+    dim: int
+    """Intrinsic dimension of the manifold."""
+
+    def __init__(self, dim: int, /) -> None:
+        object.__setattr__(self, "dim", dim)
+        object.__setattr__(self, "metric", cxm.EuclideanMetric(dim))
+        object.__setattr__(self, "atlas", EuclideanAtlas(self.dim))