Start delta-hyperbolicity exploration

Author: pchlenski

manify/embedders/_losses.py

@@ -51,8 +51,8 @@ def distortion_loss(
         D_true = D_true.flatten()
         D_est = D_est.flatten()
 
-    # Mask out any infinite or nan values
-    mask = torch.isfinite(D_true) & ~torch.isnan(D_true)
+    # Mask out any infinite or nan values; also anywhere the true distance is zero
+    mask = torch.isfinite(D_true) & ~torch.isnan(D_true) & (D_true != 0)
     D_true = D_true[mask]
     D_est = D_est[mask]
 

manify/embedders/siamese.py

@@ -10,36 +10,49 @@
 
 from __future__ import annotations
 
-from typing import Optional
+import sys
+from typing import Dict, List, Optional, Tuple
 
+import numpy as np
 import torch
 from jaxtyping import Float
 
 from ..manifolds import ProductManifold
 from ._base import BaseEmbedder
+from ._losses import distortion_loss
 
+# TQDM: notebook or regular
+if "ipykernel" in sys.modules:
+    from tqdm.notebook import tqdm
+else:
+    from tqdm import tqdm
 
-class SiameseNetwork(torch.nn.Module):
+
+class SiameseNetwork(BaseEmbedder, torch.nn.Module):
     """Siamese network for embedding data into a product manifold space.
 
     A Siamese network consists of an encoder network that maps input data to a latent representation in a product
     manifold, and optionally a decoder network that maps the latent representation back to the original feature space.
 
     Attributes:
-        pm: The product manifold object defining the embedding space.
-        encoder: Neural network module that maps input data to the embedding space.
-        decoder: Optional neural network module for reconstructing input data from embeddings.
-        reconstruction_loss: Loss function for measuring reconstruction quality.
+        pm: Product manifold defining the structure of the latent space.
+        random_state: Random state for reproducibility.
+        encoder: Neural network that maps inputs to latent embeddings.
+        decoder: Neural network that reconstructs inputs from latent embeddings.
+        beta: Weight for the distortion term in the loss function.
+        device: Device for tensor computations.
+        reconstruction_loss: Type of reconstruction loss to use.
+        
 
     Args:
-        pm: Product manifold object defining the target embedding space.
-        encoder: Neural network module that maps inputs to the embedding space.
-        decoder: Optional neural network module that maps embeddings back to input space. If None, a no-op identity
-            module is used.
-        reconstruction_loss: Type of reconstruction loss to use. Currently only "mse" (mean squared error) is supported.
-
-    Raises:
-        ValueError: If an unsupported reconstruction_loss is specified.
+        pm: Product manifold defining the structure of the latent space.
+        encoder: Neural network module that maps inputs to the manifold's intrinsic dimension.
+            The output dimension should match the intrinsic dimension of the product manifold.
+        decoder: Neural network module that maps latent representations back to the input space.
+        random_state: Optional random state for reproducibility.
+        device: Optional device for tensor computations.
+        beta: Weight of the distortion term in the loss function.
+        reconstruction_loss: Type of reconstruction loss to use.
     """
 
     def __init__(

manify/embedders/vae.py

@@ -51,6 +51,9 @@ class ProductSpaceVAE(BaseEmbedder, torch.nn.Module):
         device: Device for tensor computations.
         n_samples: Number of samples for Monte Carlo estimation of KL divergence.
         reconstruction_loss: Type of reconstruction loss to use.
+        loss_history_: Dictionary to store the history of loss values during training.
+        is_fitted_: Boolean flag indicating whether the model has been fitted.
+
 
     Args:
         pm: Product manifold defining the structure of the latent space.