More typing fixes, mostly

Author: pchlenski

.github/workflows/test.yml

@@ -28,12 +28,9 @@ jobs:
       - name: Install dependencies
         run: |
           python -m pip install --upgrade pip
-          # Install base manify
           pip install -e .
-          # install manify-dev
           pip install -e ".[dev]"
-          # pip install -e ".[utils]"
-          # install all extras
+          pip install -e ".[utils]"
           pip install -e ".[all]" 
 
       - name: Smoke‑test pip install manify
@@ -45,31 +42,25 @@ jobs:
 
       # Code quality checks
       - name: Check code formatting with Black
-        run: |
-          black --check manify/ --line-length 120
+        run: black --check manify/ --line-length 120
         continue-on-error: true
 
       - name: Check import ordering with isort
-        run: |
-          isort --check-only --profile black manify/
+        run: isort --check-only --profile black manify/
         continue-on-error: true
 
       - name: Run pylint
-        run: |
-          pylint manify/
+        run: pylint manify/
         continue-on-error: true
 
       # Type checking
       - name: Check type annotations with MyPy
-        run: |
-          # only check core modules—skip untyped imports
-          mypy --ignore-missing-imports --no-warn-return-any --follow-imports=skip manify/
+        run: mypy manify/
         continue-on-error: true
 
       # Unit testing
       - name: Run unit tests & collect coverage
-        run: |
-          pytest tests --cov=manify --cov-report=xml:coverage.xml
+        run: pytest tests --cov=manify --cov-report=xml:coverage.xml
 
       # Code coverage
       - name: Upload coverage to Codecov

manify/__init__.py

@@ -1,3 +1,5 @@
+"""Manify: A Python Library for Learning Non-Euclidean Representations."""
+
 import manify.curvature_estimation
 import manify.embedders
 import manify.manifolds

manify/curvature_estimation/delta_hyperbolicity.py

@@ -2,9 +2,10 @@
 
 from __future__ import annotations
 
+from typing import Tuple
+
 import torch
 from jaxtyping import Float
-from typing import Tuple
 
 
 def sampled_delta_hyperbolicity(

manify/curvature_estimation/greedy_method.py

@@ -2,7 +2,7 @@
 
 from __future__ import annotations
 
-from typing import Tuple, Any
+from typing import Any, Tuple
 
 import torch
 

manify/curvature_estimation/sectional_curvature.py

@@ -2,76 +2,81 @@
 
 import random
 
+import networkx as nx
 import numpy as np
 
 # The next couple functions are taken from this repo:
 # https://github.com/HazyResearch/hyperbolics
 # Paper: https://openreview.net/pdf?id=HJxeWnCcF7
 
 
-def Ka(D, m, b, c, a):
-    if a == m:
-        return 0.0
-    k = D[a][m] ** 2 + D[b][c] ** 2 / 4.0 - (D[a][b] ** 2 + D[a][c] ** 2) / 2.0
-    k /= 2 * D[a][m]
-    return k
-
-
-def K(D, n, m, b, c):
-    ks = [Ka(D, m, b, c, a) for a in range(n)]
-    return np.mean(ks)
-
-
-def ref(D, size, n, m, b, c):
-    ks = []
-    for i in range(n):
-        a = random.randint(0, size - 1)
-        if a == b or a == c:
-            continue
-        else:
-            ks.append(Ka(D, m, b, c, a))
-    return np.mean(ks)
-
-
-def estimate_curvature(G, D, n):
-    for m in range(n):
-        ks = []
-        edges = list(G.edges(m))
-        for i in range(len(edges)):
-            for j in range(b, len(edges)):
-                b = edges[i]
-                c = edges[j]
-                ks.append(K(D, n, b, c))
-    return None
-
-
-def sample(D, size, n_samples=100):
-    samples = []
-    _cnt = 0
-    while _cnt < n_samples:
-        a, b, c, m = random.sample(range(0, size), 4)
-        k = Ka(D, m, b, c, a)
-        samples.append(k)
-
-        _cnt += 1
-
-    return np.array(samples)
-
-
-def estimate(D, size, n_samples):
-    samples = sample(D, size, n_samples)
-    m1 = np.mean(samples)
-    m2 = np.mean(samples**2)
-    return samples
-
-
-def estimate_diff(D, size, n_sample, num):
-    samples = []
-    _cnt = 0
-    while _cnt < n_sample:
-        b, c, m = random.sample(range(0, size), 3)
-        k = ref(D, size, num, m, b, c)
-        # k=K(D, n, m, b, c)
-        samples.append(k)
-        _cnt += 1
-    return np.array(samples)
+def sectional_curvature(G: nx.Graph) -> np.ndarray:
+    raise NotImplementedError
+
+
+# def Ka(D: np.ndarray, m: int, b: int, c: int, a: int) -> float:
+#     if a == m:
+#         return 0.0
+#     k = D[a][m] ** 2 + D[b][c] ** 2 / 4.0 - (D[a][b] ** 2 + D[a][c] ** 2) / 2.0
+#     k /= 2 * D[a][m]
+#     return float(k)
+
+
+# def K(D: np.ndarray, n: int, m: int, b: int, c: int) -> float:
+#     ks = [Ka(D, m, b, c, a) for a in range(n)]
+#     return float(np.mean(ks))
+
+
+# def ref(D: np.ndarray, size: int, n: int, m: int, b: int, c: int) -> float:
+#     ks = []
+#     for i in range(n):
+#         a = random.randint(0, size - 1)
+#         if a == b or a == c:
+#             continue
+#         else:
+#             ks.append(Ka(D, m, b, c, a))
+#     return float(np.mean(ks))
+
+
+# def estimate_curvature(G: nx.Graph, D: np.ndarray, n: int) -> None:
+#     for m in range(n):
+#         ks = []
+#         edges = list(G.edges(m))
+#         for i in range(len(edges)):
+#             for j in range(i + 1, len(edges)):
+#                 b = edges[i]
+#                 c = edges[j]
+#                 ks.append(K(D, n, m, b, c))
+#     return None
+
+
+# def sample(D: np.ndarray, size: int, n_samples: int = 100) -> np.ndarray:
+#     samples = []
+#     _cnt = 0
+#     while _cnt < n_samples:
+#         a, b, c, m = random.sample(range(0, size), 4)
+#         k = Ka(D, m, b, c, a)
+#         samples.append(k)
+
+#         _cnt += 1
+
+#     return np.array(samples)
+
+
+# def estimate(D: np.ndarray, size: int, n_samples: int) -> np.ndarray:
+#     samples = sample(D, size, n_samples)
+#     m1 = np.mean(samples)
+#     m2 = np.mean(samples**2)
+#     return samples
+
+
+# def estimate_diff(D: np.ndarray, size: int, n_sample: int, num: int) -> np.ndarray:
+#     samples = []
+#     _cnt = 0
+#     while _cnt < n_sample:
+#         b, c, m = random.sample(range(0, size), 3)
+#         k = ref(D, size, num, m, b, c)
+#         # k=K(D, n, m, b, c)
+#         samples.append(k)
+#         _cnt += 1
+#     return np.array(samples)

manify/embedders/losses.py manify/embedders/_losses.pymanify/embedders/losses.py renamed to manify/embedders/_losses.py

@@ -11,7 +11,7 @@
 from jaxtyping import Float, Int
 
 from ..manifolds import ProductManifold
-from .losses import d_avg, distortion_loss
+from ._losses import d_avg, distortion_loss
 
 # TQDM: notebook or regular
 if "ipykernel" in sys.modules:

manify/embedders/coordinate_learning.py

@@ -1,11 +1,10 @@
-"""
-Tools for generating Riemannian manifolds and product manifolds.
+"""Tools for generating Riemannian manifolds and product manifolds.
 
-The module consists of two classes: Manifold and ProductManifold .The Manifold class
-represents hyperbolic, Euclidean, or spherical manifolds based on curvature.
-The ProductManifold class supports products of multiple manifolds,
-combining their geometric properties to create mixed-curvature. Both classes
-includes functions for different key geometric operations.
+The module consists of two classes: `Manifold` and `ProductManifold`. The `Manifold` class represents hyperbolic,
+Euclidean, or spherical manifolds of constant Gaussian curvature. The `ProductManifold` class supports Cartesian
+products of multiple manifolds, combining their geometric properties to create mixed-curvature. Both classes
+include methods for different key geometric operations, and are built on top of their corresponding `geoopt` classes
+(`Lorentz`, `Euclidean`, `Sphere`, `Scaled` and `ProductManifold`)
 """
 
 from __future__ import annotations
@@ -21,8 +20,7 @@
 
 
 class Manifold:
-    """
-    Tools for generating Riemannian manifolds.
+    """Tools for generating Riemannian manifolds.
 
     Parameters
     ----------