future annotations + isort + black

Author: pchlenski

manify/__init__.py

@@ -1,6 +1,5 @@
 import manify.curvature_estimation
 import manify.embedders
+import manify.manifolds
 import manify.predictors
 import manify.utils
-
-import manify.manifolds

manify/curvature_estimation/delta_hyperbolicity.py

@@ -1,8 +1,10 @@
 """Compute delta-hyperbolicity of a metric space."""
 
+from __future__ import annotations
 
-from jaxtyping import Float
 import torch
+from jaxtyping import Float
+
 
 def sampled_delta_hyperbolicity(dismat, n_samples=1000, reference_idx=0):
     n = dismat.shape[0]
@@ -12,55 +14,62 @@ def sampled_delta_hyperbolicity(dismat, n_samples=1000, reference_idx=0):
     # Get gromov products
     # (j,k)_i = .5 (d(i,j) + d(i,k) - d(j,k))
 
-    x,y,z = indices.T
-    w = reference_idx # set reference point
+    x, y, z = indices.T
+    w = reference_idx  # set reference point
 
-    xy_w = .5 * (dismat[w,x] + dismat[w,y] - dismat[x,y])
-    xz_w = .5 * (dismat[w,x] + dismat[w,z] - dismat[x,z])
-    yz_w = .5 * (dismat[w,y] + dismat[w,z] - dismat[y,z])
+    xy_w = 0.5 * (dismat[w, x] + dismat[w, y] - dismat[x, y])
+    xz_w = 0.5 * (dismat[w, x] + dismat[w, z] - dismat[x, z])
+    yz_w = 0.5 * (dismat[w, y] + dismat[w, z] - dismat[y, z])
 
     # delta(x,y,z) = min((x,y)_w,(y-z)_w) - (x,z)_w
-    deltas = torch.minimum(xy_w,yz_w) - xz_w
+    deltas = torch.minimum(xy_w, yz_w) - xz_w
     diam = torch.max(dismat)
     rel_deltas = 2 * deltas / diam
 
     return rel_deltas, indices
 
+
 def iterative_delta_hyperbolicity(dismat):
     """delta(x,y,z) = min((x,y)_w,(y-z)_w) - (x,z)_w"""
     n = dismat.shape[0]
     w = 0
-    gromov_products = torch.zeros((n,n))
-    deltas = torch.zeros((n,n,n))
+    gromov_products = torch.zeros((n, n))
+    deltas = torch.zeros((n, n, n))
 
     # Get Gromov Products
     for x in range(n):
         for y in range(n):
-            gromov_products[x,y] = gromov_product(w,x,y,dismat)
+            gromov_products[x, y] = gromov_product(w, x, y, dismat)
 
     # Get Deltas
     for x in range(n):
         for y in range(n):
             for z in range(n):
-                xz_w = gromov_products[x,z]
-                xy_w = gromov_products[x,y]
-                yz_w = gromov_products[y,z]
-                deltas[x,y,z] = torch.minimum(xy_w,yz_w) - xz_w
-    
+                xz_w = gromov_products[x, z]
+                xy_w = gromov_products[x, y]
+                yz_w = gromov_products[y, z]
+                deltas[x, y, z] = torch.minimum(xy_w, yz_w) - xz_w
+
     diam = torch.max(dismat)
     rel_deltas = 2 * deltas / diam
 
     return rel_deltas, gromov_products
 
 
-def gromov_product(i,j,k,dismat):
+def gromov_product(i, j, k, dismat):
     """(j,k)_i = 0.5 (d(i,j) + d(i,k) - d(j,k))"""
-    d_ij = dismat[i,j]
-    d_ik = dismat[i,k]
-    d_jk = dismat[j,k]
+    d_ij = dismat[i, j]
+    d_ik = dismat[i, k]
+    d_jk = dismat[j, k]
     return 0.5 * (d_ij + d_ik - d_jk)
 
-def delta_hyperbolicity(dismat: Float[torch.Tensor, "n_points n_points"], relative=True, device='cpu', full=False) -> Float[torch.Tensor, ""]:
+
+def delta_hyperbolicity(
+    dismat: Float[torch.Tensor, "n_points n_points"],
+    relative=True,
+    device="cpu",
+    full=False,
+) -> Float[torch.Tensor, ""]:
     """
     Compute the delta-hyperbolicity of a metric space.
 
@@ -69,7 +78,7 @@ def delta_hyperbolicity(dismat: Float[torch.Tensor, "n_points n_points"], relati
         relative: Whether to return the relative delta-hyperbolicity.
         device: Device to run the computation on.
         full: Whether to return the full delta tensor or just the maximum delta.
-    
+
     Returns:
         delta: Delta-hyperbolicity of the metric space.
     """
@@ -95,7 +104,5 @@ def delta_hyperbolicity(dismat: Float[torch.Tensor, "n_points n_points"], relati
     if relative:
         diam = torch.max(dismat).item()
         delta = 2 * delta / diam
-    
-    return delta
 
-    
+    return delta

manify/curvature_estimation/greedy_method.py

@@ -1,5 +1,7 @@
 """Greedy selection of signatures, as described in Tabaghi et al. at https://arxiv.org/pdf/2102.10204"""
 
+from __future__ import annotations
+
 from typing import Tuple
 
 import torch
@@ -10,7 +12,11 @@
 def greedy_curvature_method(
     pm: ProductManifold,
     dists: torch.Tensor,
-    candidate_components: Tuple[Tuple[float, int], ...] = ((-1.0, 2), (0.0, 2), (1.0, 2)),
+    candidate_components: Tuple[Tuple[float, int], ...] = (
+        (-1.0, 2),
+        (0.0, 2),
+        (1.0, 2),
+    ),
     max_components: int = 3,
 ):
     """The greedy curvature estimation method from Tabaghi et al. at https://arxiv.org/pdf/2102.10204"""

manify/curvature_estimation/sectional_curvature.py

@@ -1,8 +1,9 @@
+from __future__ import annotations
+
 import random
 
 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
@@ -60,7 +61,7 @@ def sample(D, size, n_samples=100):
 def estimate(D, size, n_samples):
     samples = sample(D, size, n_samples)
     m1 = np.mean(samples)
-    m2 = np.mean(samples**2)
+    m2 = np.mean(samples ** 2)
     return samples
 
 

manify/embedders/coordinate_learning.py

@@ -1,15 +1,17 @@
 """Implementation for coordinate training and optimization"""
 
+from __future__ import annotations
+
 import sys
-from typing import Tuple, List, Dict
-from jaxtyping import Float, Int
+from typing import Dict, List, Tuple
 
+import geoopt
 import numpy as np
 import torch
-import geoopt
+from jaxtyping import Float, Int
 
-from .losses import distortion_loss, d_avg
 from ..manifolds import ProductManifold
+from .losses import d_avg, distortion_loss
 
 # TQDM: notebook or regular
 if "ipykernel" in sys.modules:

manify/embedders/losses.py

@@ -1,10 +1,12 @@
 """Implementation of different measurement metrics"""
 
+from __future__ import annotations
+
 from typing import List
-from jaxtyping import Float
 
-import torch
 import networkx as nx
+import torch
+from jaxtyping import Float
 
 from ..manifolds import ProductManifold
 
@@ -24,7 +26,7 @@ def distortion_loss(
     Returns:
         float: A float indicating the distortion loss, calculated as the sum of the squared relative
          errors between the estimated and true squared distances.
-        
+
     See also: square_loss in HazyResearch hyperbolics repo:
     https://github.com/HazyResearch/hyperbolics/blob/master/pytorch/hyperbolic_models.py#L178
     """

manify/embedders/siamese.py

@@ -1,8 +1,11 @@
 """Siamese network embedder"""
 
-from typing import List, Optional, Literal
-from jaxtyping import Float
+from __future__ import annotations
+
+from typing import Optional
+
 import torch
+from jaxtyping import Float
 
 from ..manifolds import ProductManifold
 

manify/embedders/vae.py

@@ -1,9 +1,11 @@
 """Product space variational autoencoder implementation"""
 
+from __future__ import annotations
+
 from typing import List, Tuple
-from jaxtyping import Float
 
 import torch
+from jaxtyping import Float
 
 from ..manifolds import ProductManifold
 
@@ -40,7 +42,7 @@ def __init__(
 
     def encode(
         self, x: Float[torch.Tensor, "batch_size n_features"]
-    ) -> Tuple[Float[torch.Tensor, "batch_size n_latent"], Float[torch.Tensor, "batch_size n_latent"]]:
+    ) -> Tuple[Float[torch.Tensor, "batch_size n_latent"], Float[torch.Tensor, "batch_size n_latent"],]:
         """Must return z_mean, z_logvar"""
         z_mean_tangent, z_logvar = self.encoder(x)
         z_mean_ambient = z_mean_tangent @ self.pm.projection_matrix  # Adds zeros in the right places
@@ -51,7 +53,9 @@ def decode(self, z: Float[torch.Tensor, "batch_size n_latent"]) -> Float[torch.T
         """Decoding in product space VAE"""
         return self.decoder(z)
 
-    def forward(self, x: Float[torch.Tensor, "batch_size n_features"]) -> Tuple[
+    def forward(
+        self, x: Float[torch.Tensor, "batch_size n_features"]
+    ) -> Tuple[
         Float[torch.Tensor, "batch_size n_features"],
         Float[torch.Tensor, "batch_size n_latent"],
         List[Float[torch.Tensor, "n_latent n_latent"]],

manify/manifolds.py

@@ -8,8 +8,10 @@
 includes functions for different key geometric operations.
 """
 
+from __future__ import annotations
+
 import warnings
-from typing import Callable, List, Literal, Optional, Tuple, Union
+from typing import TYPE_CHECKING, Callable, List, Literal, Optional, Tuple, Union
 
 import geoopt
 import torch
@@ -31,7 +33,13 @@ class Manifold:
     stereographic: (bool) Whether to use stereographic coordinates for the manifold.
     """
 
-    def __init__(self, curvature: float, dim: int, device: str = "cpu", stereographic: bool = False):
+    def __init__(
+        self,
+        curvature: float,
+        dim: int,
+        device: str = "cpu",
+        stereographic: bool = False,
+    ):
         # Device management
         self.device = device
 
@@ -95,7 +103,9 @@ def to(self, device: str) -> "Manifold":
         return self
 
     def inner(
-        self, X: Float[torch.Tensor, "n_points1 n_dim"], Y: Float[torch.Tensor, "n_points2 n_dim"]
+        self,
+        X: Float[torch.Tensor, "n_points1 n_dim"],
+        Y: Float[torch.Tensor, "n_points2 n_dim"],
     ) -> Float[torch.Tensor, "n_points1 n_points2"]:
         """
         Compute the inner product of manifolds.
@@ -116,7 +126,9 @@ def inner(
         return X_fixed @ Y.T * scaler
 
     def dist(
-        self, X: Float[torch.Tensor, "n_points1 n_dim"], Y: Float[torch.Tensor, "n_points2 n_dim"]
+        self,
+        X: Float[torch.Tensor, "n_points1 n_dim"],
+        Y: Float[torch.Tensor, "n_points2 n_dim"],
     ) -> Float[torch.Tensor, "n_points1 n_points2"]:
         """
         Inherit distance function from the geoopt manifold.
@@ -131,7 +143,9 @@ def dist(
         return self.manifold.dist(X[:, None], Y[None, :])
 
     def dist2(
-        self, X: Float[torch.Tensor, "n_points1 n_dim"], Y: Float[torch.Tensor, "n_points2 n_dim"]
+        self,
+        X: Float[torch.Tensor, "n_points1 n_dim"],
+        Y: Float[torch.Tensor, "n_points2 n_dim"],
     ) -> Float[torch.Tensor, "n_points1 n_points2"]:
         """
         Inherit squared distance function from the geoopt manifold.
@@ -194,7 +208,10 @@ def sample(
         sigma: Optional[Float[torch.Tensor, "n_points n_dim n_dim"]] = None,
     ) -> Union[
         Float[torch.Tensor, "n_points n_ambient_dim"],
-        Tuple[Float[torch.Tensor, "n_points n_ambient_dim"], Float[torch.Tensor, "n_points n_dim"]],
+        Tuple[
+            Float[torch.Tensor, "n_points n_ambient_dim"],
+            Float[torch.Tensor, "n_points n_dim"],
+        ],
     ]:
         """
         Sample from the variational distribution.
@@ -304,7 +321,9 @@ def log_likelihood(
             return ll - (n - 1) * torch.log(R * torch.abs(sin_M(u_norm / R) / u_norm) + 1e-8)
 
     def logmap(
-        self, x: Float[torch.Tensor, "n_points n_dim"], base: Optional[Float[torch.Tensor, "n_points n_dim"]] = None
+        self,
+        x: Float[torch.Tensor, "n_points n_dim"],
+        base: Optional[Float[torch.Tensor, "n_points n_dim"]] = None,
     ) -> Float[torch.Tensor, "n_points n_dim"]:
         """
         Logarithmic map of point on manifold x at base point.
@@ -322,7 +341,9 @@ def logmap(
         return self.manifold.logmap(x=base, y=x)
 
     def expmap(
-        self, u: Float[torch.Tensor, "n_points n_dim"], base: Optional[Float[torch.Tensor, "n_points n_dim"]] = None
+        self,
+        u: Float[torch.Tensor, "n_points n_dim"],
+        base: Optional[Float[torch.Tensor, "n_points n_dim"]] = None,
     ) -> Float[torch.Tensor, "n_points n_dim"]:
         """
         Exponential map of tangent vector u at base point.
@@ -394,7 +415,7 @@ def inverse_stereographic(self, *points: Float[torch.Tensor, "n_points n_dim_ste
             return orig_manifold, *points  # type: ignore
 
         # Inverse projection for points
-        norm_squared = [(Y**2).sum(dim=1, keepdim=True) for Y in points]
+        norm_squared = [(Y ** 2).sum(dim=1, keepdim=True) for Y in points]
         sign = torch.sign(self.curvature)  # type: ignore
 
         X0 = (1 + sign * norm_squared) / (1 - sign * norm_squared)
@@ -436,7 +457,12 @@ class ProductManifold(Manifold):
     stereographic: (bool) Whether to use stereographic coordinates for the manifold.
     """
 
-    def __init__(self, signature: List[Tuple[float, int]], device: str = "cpu", stereographic: bool = False):
+    def __init__(
+        self,
+        signature: List[Tuple[float, int]],
+        device: str = "cpu",
+        stereographic: bool = False,
+    ):
         # Device management
         self.device = device
 
@@ -459,7 +485,12 @@ def __init__(self, signature: List[Tuple[float, int]], device: str = "cpu", ster
 
         # Manifold <-> Dimension mapping
         self.ambient_dim, self.n_manifolds, self.dim = 0, 0, 0
-        self.dim2man, self.man2dim, self.man2intrinsic, self.intrinsic2man = {}, {}, {}, {}
+        self.dim2man, self.man2dim, self.man2intrinsic, self.intrinsic2man = (
+            {},
+            {},
+            {},
+            {},
+        )
 
         for M in self.P:
             for d in range(self.ambient_dim, self.ambient_dim + M.ambient_dim):
@@ -518,7 +549,10 @@ def sample(
         sigma_factorized: Optional[List[Float[torch.Tensor, "n_points n_dim_manifold n_dim_manifold"]]] = None,
     ) -> Union[
         Float[torch.Tensor, "n_points n_ambient_dim"],
-        Tuple[Float[torch.Tensor, "n_points n_ambient_dim"], Float[torch.Tensor, "n_points n_dim"]],
+        Tuple[
+            Float[torch.Tensor, "n_points n_ambient_dim"],
+            Float[torch.Tensor, "n_points n_dim"],
+        ],
     ]:
         """
         Sample from the variational distribution.