jg/pretty minimal

scripts/visualize_mnist_embeddings.py

@@ -30,7 +30,7 @@
 plt.savefig("umap_original.png")
 plt.close()
 
-for epoch in [1, 5, 10, 50, 100, 200, 500]:
+for epoch in [4]:#1, 5, 10, 50, 100, 200, 500]:
     reducer = ParametricUMAP(
         epochs=epoch,
         lr=1e-3,
@@ -52,5 +52,5 @@
     output = reducer.transform(X)
 
     plt.scatter(output[:, 0], output[:, 1], c=colors, s=5)
-    plt.savefig(f"umap_glassbox_{epoch}.png")
+    plt.savefig(f"feb_26_pr_umap_glassbox_{epoch}.png")
     plt.close()

src/glass_box_umap/parametric_umap/core.py

@@ -75,6 +75,7 @@ class ParametricUMAP:
     repulsion_strength: float = 3.0
     num_workers: int = 0
     checkpoint_dir: Path | None = None
+    best_ckpt: bool = False
 
     _model: UMAPLightningModule | None = field(init=False, default=None)
     _pca: PCA | None = field(init=False, default=None)
@@ -164,22 +165,52 @@ def fit(self, X: NDArray[np.floating] | Tensor) -> Self:
             input_dims = tuple(X.shape[1:])
             self._model = self._build_model(input_dims)
 
-            graph = get_umap_graph(
-                X.detach().cpu().numpy(),
-                n_neighbors=self.n_neighbors,
-                metric=self.metric,
-                random_state=self.random_state,
-            )
+
+            if isinstance(X, torch.Tensor):                        
+                X = X.detach().cpu().squeeze()
+
+            if len(X.shape)>2:
+                conv_flag=True
+                # X = X.reshape([-1,784])
+                graph = get_umap_graph(
+                    X.reshape([-1,784]),
+                    n_neighbors=self.n_neighbors,
+                    metric=self.metric,
+                    random_state=self.random_state,
+                )
+            else:
+                conv_flag=False
+                graph = get_umap_graph(
+                    X,
+                    n_neighbors=self.n_neighbors,
+                    metric=self.metric,
+                    random_state=self.random_state,
+                )
+
+            # graph = get_umap_graph(
+            #     X,
+            #     n_neighbors=self.n_neighbors,
+            #     metric=self.metric,
+            #     random_state=self.random_state,
+            # )
+
+            if isinstance(X, torch.Tensor):                        
+                X = X.detach().cpu().numpy()
+            if conv_flag:
+                X = torch.tensor(X.squeeze()).unsqueeze(1).detach().cpu().numpy() #reshape([-1,28,28]).unsqueeze(1)
+                print("Shape: ", X.shape)
             datamodule = UMAPDataModule(
-                UMAPDataset(X.detach().cpu().numpy(), graph),
+                # UMAPDataset(X.detach().cpu().numpy(), graph),
+                UMAPDataset(X, graph),
                 self.batch_size,
                 self.num_workers,
             )
 
             trainer.fit(model=self._model, datamodule=datamodule)
-
-            best_ckpt = torch.load(best_checkpoint.best_model_path, map_location="cpu")
-            self._model.load_state_dict(best_ckpt["state_dict"])
+
+            if self.best_ckpt:
+                best_ckpt = torch.load(best_checkpoint.best_model_path, map_location="cpu")
+                self._model.load_state_dict(best_ckpt["state_dict"])
 
         self._model.to(self._device)
 
@@ -216,7 +247,7 @@ def transform(
 
         return torch.cat(results).numpy()
 
-    def fit_transform(self, X: NDArray[np.floating] | Tensor) -> NDArray[np.floating]:
+    def fit_transform(self, X: Tensor) -> NDArray[np.floating]:
         self.fit(X)
         return self.transform(X)
 

src/glass_box_umap/parametric_umap/data.py

@@ -20,24 +20,36 @@ class UMAPDataset(Dataset[tuple[Tensor, Tensor]]):
         n_epochs: Number of training epochs for computing edge sampling frequency.
     """
 
+    # def __init__(
+    #     self,
+    #     data: NDArray[np.floating],
+    #     graph: csr_matrix,
+    #     n_epochs: int = 200,
+    # ) -> None:
+    #     _, epochs_per_sample, head, tail, _, _ = get_graph_elements(graph, n_epochs)
+
+    #     edges_to_exp = np.repeat(head, epochs_per_sample.astype(np.intp))
+    #     edges_from_exp = np.repeat(tail, epochs_per_sample.astype(np.intp))
+
     def __init__(
         self,
         data: NDArray[np.floating],
         graph: csr_matrix,
-        n_epochs: int = 200,
+        n_epochs: int = 40,#200,#
     ) -> None:
         _, epochs_per_sample, head, tail, _, _ = get_graph_elements(graph, n_epochs)
-
-        edges_to_exp = np.repeat(head, epochs_per_sample.astype(np.intp))
-        edges_from_exp = np.repeat(tail, epochs_per_sample.astype(np.intp))
+        print("n_epochs: ", n_epochs)
+        print("Graph len: ", len(head),len(tail))
+        edges_to_exp = np.repeat(head,1)# epochs_per_sample.astype(np.intp))
+        edges_from_exp = np.repeat(tail, 1)#epochs_per_sample.astype(np.intp))
 
         shuffle_mask = np.random.permutation(np.arange(len(edges_to_exp)))
         self.edges_to_exp = edges_to_exp[shuffle_mask].astype(np.int64)
         self.edges_from_exp = edges_from_exp[shuffle_mask].astype(np.int64)
         self.data = torch.as_tensor(data, dtype=torch.float32)
 
     def __len__(self) -> int:
-        return self.data.shape[0]
+        return self.edges_to_exp.shape[0]
 
     def __getitem__(self, index: int) -> tuple[Tensor, Tensor]:
         edges_to_exp = self.data[self.edges_to_exp[index]]

src/glass_box_umap/parametric_umap/graph.py

@@ -21,7 +21,7 @@
 def get_umap_graph(
     X: NDArray[np.floating],
     n_neighbors: int = 10,
-    metric: str = "cosine",
+    metric: str = "euclidean",
     random_state: RandomState | int | None = None,
 ) -> csr_matrix:
     """Build a UMAP graph from input data using nearest neighbor descent.
@@ -97,6 +97,12 @@ def get_graph_elements(graph: csr_matrix, n_epochs: int) -> GraphElements:
     graph_coo.data[graph_coo.data < (graph_coo.data.max() / float(n_epochs))] = 0.0
     graph_coo.eliminate_zeros()
 
+    # print("Graph len 0: ", len(graph_coo.data))
+    # graph_cutoff_sorted = np.percentile(graph_coo.data, 90)
+    # graph_coo.data[graph_coo.data < graph_cutoff_sorted]=0
+    # graph_coo.eliminate_zeros()
+    # print("Graph len 1: ", len(graph_coo.data))
+
     epochs_per_sample = (n_epochs * graph_coo.data).astype(np.float32)
     head = graph_coo.row
     tail = graph_coo.col

src/glass_box_umap/parametric_umap/models.py

@@ -3,6 +3,10 @@
 import torch
 from torch import Tensor, nn
 
+import torch.nn.init as init
+
+from glass_box_umap.components import LayerNormDetached
+
 DEFAULT_HIDDEN_DIMS = [100, 100, 100]
 
 
@@ -40,10 +44,10 @@ def __init__(
             hidden_dims = [512, 512]
 
         self.conv_layers = nn.Sequential(
-            nn.Conv2d(in_channels=in_channels, out_channels=64, kernel_size=3, stride=2, padding=1),
-            nn.ReLU(),
-            nn.Conv2d(in_channels=64, out_channels=128, kernel_size=3, stride=2, padding=1),
-            nn.ReLU(),
+            nn.Conv2d(in_channels=in_channels, out_channels=64, kernel_size=3, stride=2, padding=1,bias=False),
+            nn.PReLU(),
+            nn.Conv2d(in_channels=64, out_channels=64, kernel_size=3, stride=2, padding=1,bias=False),
+            nn.PReLU(),
             nn.Flatten(),
         )
 
@@ -52,9 +56,9 @@ def __init__(
         mlp_layers: list[nn.Module] = []
         prev_dim = flattened_size
         for dim in hidden_dims:
-            mlp_layers.extend([nn.Linear(prev_dim, dim), nn.ReLU()])
+            mlp_layers.extend([nn.Linear(prev_dim, dim, bias=False), nn.PReLU()])
             prev_dim = dim
-        mlp_layers.append(nn.Linear(prev_dim, n_components))
+        mlp_layers.append(nn.Linear(prev_dim, n_components, bias=False))
 
         self.mlp = nn.Sequential(*mlp_layers)
 
@@ -63,35 +67,35 @@ def forward(self, x: Tensor) -> Tensor:
         return self.mlp(x)
 
 
-class DefaultEncoder(nn.Module):
-    """Default MLP encoder for flattened data.
+# class DefaultEncoder(nn.Module):
+#     """Default MLP encoder for flattened data.
 
-    A simple feedforward encoder that flattens input and projects through
-    hidden layers to the embedding space.
+#     A simple feedforward encoder that flattens input and projects through
+#     hidden layers to the embedding space.
 
-    Args:
-        input_dims: Shape of input data (excluding batch dimension).
-        n_components: Dimensionality of the output embedding space.
-        hidden_dims: Sizes of hidden layers.
-    """
+#     Args:
+#         input_dims: Shape of input data (excluding batch dimension).
+#         n_components: Dimensionality of the output embedding space.
+#         hidden_dims: Sizes of hidden layers.
+#     """
 
-    def __init__(
-        self,
-        input_dims: tuple[int, ...],
-        n_components: int = 2,
-        hidden_dims: list[int] = DEFAULT_HIDDEN_DIMS,
-    ) -> None:
-        super().__init__()
-        layers: list[nn.Module] = [nn.Flatten()]
-        prev_dim = math.prod(input_dims)
-        for dim in hidden_dims:
-            layers.extend([nn.Linear(prev_dim, dim), nn.ReLU()])
-            prev_dim = dim
-        layers.append(nn.Linear(prev_dim, n_components))
-        self.encoder = nn.Sequential(*layers)
+#     def __init__(
+#         self,
+#         input_dims: tuple[int, ...],
+#         n_components: int = 2,
+#         hidden_dims: list[int] = DEFAULT_HIDDEN_DIMS,
+#     ) -> None:
+#         super().__init__()
+#         layers: list[nn.Module] = [nn.Flatten()]
+#         prev_dim = math.prod(input_dims)
+#         for dim in hidden_dims:
+#             layers.extend([nn.Linear(prev_dim, dim), nn.ReLU()])
+#             prev_dim = dim
+#         layers.append(nn.Linear(prev_dim, n_components))
+#         self.encoder = nn.Sequential(*layers)
 
-    def forward(self, x: Tensor) -> Tensor:
-        return self.encoder(x)
+#     def forward(self, x: Tensor) -> Tensor:
+#         return self.encoder(x)
 
 
 class DefaultDecoder(nn.Module):
@@ -128,3 +132,117 @@ def __init__(
 
     def forward(self, x: Tensor) -> Tensor:
         return self.decoder(x).view(x.shape[0], *self.dims)
+
+
+class ResidualMLPBlock(nn.Module):
+    """
+    Residual MLP block: x -> Linear -> (Norm) -> Act -> Linear -> (Norm) -> +skip
+    Keeps everything bias-free and piecewise-linear (if Norm has no affine).
+    """
+    def __init__(
+        self,
+        dim: int,
+        hidden_dim: int | None = None,
+        activation: str = "relu",   # or "leaky_relu"
+        negative_slope: float = 0.01,
+        use_norm: bool = True,
+    ) -> None:
+        super().__init__()
+        hidden_dim = hidden_dim or dim
+
+        self.fc1 = nn.Linear(dim, hidden_dim, bias=False)
+        self.fc2 = nn.Linear(hidden_dim, dim, bias=False)
+
+        # Norm that does NOT introduce an additive learned offset
+        # self.norm1 = nn.LayerNorm(hidden_dim, elementwise_affine=False) if use_norm else nn.Identity()
+        # self.norm2 = nn.LayerNorm(dim, elementwise_affine=False) if use_norm else nn.Identity()
+        self.norm1 = LayerNormDetached(hidden_dim) if use_norm else nn.Identity()
+        self.norm2 = LayerNormDetached(dim) if use_norm else nn.Identity()
+
+        if activation == "leaky_relu":
+            # self.act = nn.LeakyReLU(negative_slope=negative_slope)
+            # a = negative_slope
+
+            self.act = nn.PReLU()#negative_slope=negative_slope)
+            nonlin = "leaky_relu"
+        else:
+            self.act = nn.ReLU()
+            a = 0.0
+            nonlin = "relu"
+
+        # He/Kaiming init for ReLU-family activations
+        init.kaiming_uniform_(self.fc1.weight, nonlinearity=nonlin)
+        init.kaiming_uniform_(self.fc2.weight, nonlinearity=nonlin)
+
+    def forward(self, x: Tensor) -> Tensor:
+        h = self.fc1(x)
+        h = self.norm1(h)
+        h = self.act(h)
+        h = self.fc2(h)
+        h = self.norm2(h)
+        return x + h
+
+
+class DefaultEncoder(nn.Module):
+    """
+    Bias-free residual MLP encoder for tabular data.
+
+    - Mean-centering/scaling done outside the module (recommended).
+    - All Linear layers use bias=False to preserve your local Jacobian property.
+    - Residual blocks greatly improve trainability for depth.
+    """
+    def __init__(
+        self,
+        input_dims: tuple[int, ...],
+        n_components: int = 2,
+        width: int = 128,
+        depth: int = 3,              # number of residual blocks
+        mlp_ratio: float = 2.0,      # inner expansion in each block
+        activation: str = "leaky_relu",    # or "leaky_relu"
+        negative_slope: float = 0.01,
+        use_norm: bool = True,
+    ) -> None:
+        super().__init__()
+        in_dim = math.prod(input_dims)
+        hidden_dim = int(round(width * mlp_ratio))
+
+        self.flatten = nn.Flatten()
+
+        # Input projection (bias-free)
+        self.in_proj = nn.Linear(in_dim, width, bias=False)
+
+        # Init
+        if activation == "leaky_relu":
+            a = negative_slope
+            nonlin = "leaky_relu"
+        else:
+            a = 0.0
+            nonlin = "relu"
+        init.kaiming_uniform_(self.in_proj.weight, a=a, nonlinearity=nonlin)
+
+        # Residual trunk
+        self.blocks = nn.Sequential(*[
+            ResidualMLPBlock(
+                dim=width,
+                hidden_dim=hidden_dim,
+                activation=activation,
+                negative_slope=negative_slope,
+                use_norm=use_norm,
+            )
+            for _ in range(depth)
+        ])
+
+        # Optional final activation (keeps piecewise linearity)
+        self.out_act = nn.Identity()  # or nn.ReLU() / nn.LeakyReLU(...) if you want
+
+        # Output projection (bias-free!)
+        self.out_proj = nn.Linear(width, n_components, bias=False)
+        init.kaiming_uniform_(self.out_proj.weight, a=a, nonlinearity=nonlin)
+
+    def forward(self, x: Tensor) -> Tensor:
+        x = self.flatten(x)
+        x = self.in_proj(x)
+        x = self.blocks(x)
+        x = self.out_act(x)
+        x = self.out_proj(x)
+        return x