minor updates and ensuring bias free models, working with pretty mnist and gene_exp scripts

Author: james-golden-arcadia

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

@@ -168,16 +168,37 @@ def fit(self, X: NDArray[np.floating] | Tensor) -> Self:
 
             if isinstance(X, torch.Tensor):                        
                 X = X.detach().cpu().squeeze()
-
-            graph = get_umap_graph(
-                X,
-                n_neighbors=self.n_neighbors,
-                metric=self.metric,
-                random_state=self.random_state,
-            )
+            
+            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, graph),
@@ -226,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/graph.py

@@ -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

@@ -44,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(),
         )
 
@@ -56,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)
 
@@ -160,17 +160,19 @@ def __init__(
         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.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, a=a, nonlinearity=nonlin)
-        init.kaiming_uniform_(self.fc2.weight, a=a, nonlinearity=nonlin)
+        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)