Large refactor that solidifies object relationships

Makefile

@@ -17,7 +17,7 @@ format:
 
 .PHONY: typecheck
 typecheck:
-	pyright --project pyproject.toml .
+	pyright --project pyproject.toml src/ tests/
 
 .PHONY: pre-commit
 pre-commit:

docs/examples/mnist.py

@@ -20,95 +20,67 @@
 from sklearn.datasets import fetch_openml
 from sklearn.decomposition import PCA
 
-from glass_box_umap.utils import get_accelerator
-
 app = typer.Typer(pretty_exceptions_enable=False)
 
 
-def load_and_preprocess_mnist(n_pcs: int, subset: int = 4000) -> tuple[np.ndarray, np.ndarray, PCA]:
+def load_and_preprocess_mnist(
+    n_pcs: int, subset: int = 4000
+) -> tuple[torch.Tensor, np.ndarray, np.ndarray, PCA]:
     """Fetches MNIST data and performs PCA preprocessing."""
     print("Fetching MNIST...")
     mnist = fetch_openml("mnist_784", version=1)
-    data_values = mnist.data.values[:subset, :]
-    # data_values = (data_values.T - data_values.mean(axis=1).T).T
+    data_raw = mnist.data.values[:subset, :].astype(np.float32)
     target_values = mnist.target.values[:subset]
 
+    data_centered = data_raw - data_raw.mean(axis=0)
+
     pca = PCA(n_components=n_pcs)
-    pca.fit(data_values)
+    pca.fit(data_raw)
+    data_pca = pca.transform(data_raw).astype(np.float32)
 
-    mnist_pca = pca.transform(data_values)
-    mnist_pca_centered = (mnist_pca.T - mnist_pca.mean(axis=1)).T
-    return mnist_pca_centered, target_values, pca
+    return torch.from_numpy(data_pca), target_values, data_centered, pca
 
 
 @app.command()
 def main(
     output_dir: Path = typer.Option(Path("output"), help="Directory for all output files"),
-    n_fits: int = typer.Option(1, help="Number of UMAP fits"),
-    n_pcs: int = typer.Option(100, help="Number of PCA components"),
-    epochs: int = typer.Option(50, help="Number of training epochs"),
+    n_pcs: int = typer.Option(25, help="Number of PCA components"),
+    epochs: int = typer.Option(100, help="Number of training epochs"),
     random_state: int = typer.Option(42, help="Random seed"),
-    hidden_size: int = typer.Option(512*2, help="Hidden layer size"),
+    hidden_size: int = typer.Option(512, help="Hidden layer size"),
     train: bool = typer.Option(True, help="Train model (False to load from disk)"),
 ) -> None:
     """Train GlassBoxUMAP on MNIST and save outputs."""
     output_dir.mkdir(parents=True, exist_ok=True)
-    model_path_pattern = str(output_dir / "models" / "umap_{i}.pth")
-    (output_dir / "models").mkdir(parents=True, exist_ok=True)
-
-    accelerator = get_accelerator()
-
-    if accelerator in ("cuda", "mps"):
-        subset = 20000
-        batch_size = 192
+    model_path = output_dir / "model.pt"
+
+    X_pca, y_target, X_centered, pca_model = load_and_preprocess_mnist(n_pcs)
+
+    if train:
+        print("Initializing GlassBoxUMAP...")
+        reducer = GlassBoxUMAP(
+            epochs=epochs,
+            random_state=random_state,
+            encoder_kwargs={"hidden_size": hidden_size},
+        )
+
+        print("Fitting GlassBoxUMAP...")
+        reducer.fit(X_pca)
+        reducer.save(model_path)
+        print(f"Saved model to {model_path}")
     else:
-        subset = 4000
-        batch_size = 4096
-
-    X_pca_centered, y_target, pca_model = load_and_preprocess_mnist(n_pcs, subset=subset)
-
-    print("Initializing GlassBoxUMAP...")
-    reducer = GlassBoxUMAP(
-        n_fits=n_fits,
-        epochs=epochs if train else 0,
-        random_state=random_state,
-        input_size=n_pcs,
-        hidden_size=hidden_size,
-        min_dist=0.1,
-        lr=4e-4,
-        batch_size=batch_size
+        print(f"Loading model from {model_path}...")
+        reducer = GlassBoxUMAP.load(model_path)
+
+    print("Computing attributions...")
+    feature_contributions, jacobians = reducer.compute_attributions(
+        X=X_pca,
+        raw_features=X_centered,
+        projector=pca_model.components_.T,
     )
 
-    print("Fitting GlassBoxUMAP...")
-    reducer.fit(
-        X_pca_centered,
-        load_models=not train,
-        load_n_fits=n_fits,
-        save_models=train,
-        model_path_pattern=model_path_pattern,
-    )
-
-    print("Running manual Jacobian check...")
-    model = reducer._models[0]
-    encoder = model.encoder
-    device = reducer._device
-    encoder.eval().to(device)
-
-    sample_tensor = torch.tensor(X_pca_centered[:1, :].squeeze(), dtype=torch.float32).to(device)
-
-    jac_batch = torch.autograd.functional.jacobian(
-        encoder, sample_tensor, vectorize=True, strategy="reverse-mode"
-    )
-
-    jacobian_np = jac_batch.T.detach().cpu().numpy().T
-    reconstruction = jacobian_np @ X_pca_centered[0, :]
-    print("Jacobian Reconstruction:", reconstruction)
-
-    forward_pass = encoder(torch.tensor(X_pca_centered[0, :], dtype=torch.float32).to(device))
-    print("Forward Pass:", forward_pass)
-
     print("Plotting embedding...")
-    embedding = reducer._embeddings[0]
+    embedding = reducer.transform(X_pca)
 
     sns.set_style("whitegrid")
     ax = sns.scatterplot(x=embedding[:, 0], y=embedding[:, 1], hue=y_target, s=3)
@@ -119,12 +91,11 @@ def main(
     fig.savefig(embedding_path, bbox_inches="tight")
     print(f"Saved embedding plot to {embedding_path}")
 
-    print("Plotting linear operator...")
+    print("Plotting linear operator for first sample...")
     plt.figure()
 
-    linear_operator = (pca_model.components_.T @ jac_batch.T.detach().cpu().numpy())[:, 0].reshape(
-        [28, 28]
-    )
+    jacobian_pixel = jacobians[0].numpy() @ pca_model.components_
+    linear_operator = jacobian_pixel[0, :].reshape([28, 28])
 
     plt.imshow(linear_operator, cmap="RdBu")
     linear_op_path = output_dir / "linear_operator.png"

pyproject.toml

@@ -31,6 +31,7 @@ dependencies = [
     "umap-learn>=0.5.7",
     "pynndescent>=0.5.13",
     "scikit-learn>=1.6.0",
+    "tensorboard>=2.20.0",
 ]
 
 [[tool.uv.index]]
@@ -78,9 +79,9 @@ docs = [
 [tool.ruff]
 # The directories to consider when resolving first- vs. third-party imports.
 src = ["."]
+include = ["src/**/*.py", "tests/**/*.py", "**/*.ipynb"]
 line-length = 100
 indent-width = 4
-extend-include = ["*.ipynb"]
 
 [tool.ruff.lint.per-file-ignores]
 # Ignore star and unused imports.
@@ -109,10 +110,15 @@ line-ending = "auto"
 order-by-type = true
 no-lines-before = ["future", "standard-library"]
 
+[tool.pytest.ini_options]
+filterwarnings = [
+    "ignore:The 'train_dataloader' does not have many workers:UserWarning",
+]
+
 [tool.pyright]
 typeCheckingMode = "basic"
-exclude = ["./.venv", "./dist"]
+exclude = ["./.venv", "./dist", "./docs"]
 
 # Pyright reports a lot of unknown-member errors for some packages.
 # If this is a problem, set this to `false`.
-reportUnknownMemberType = true
+reportUnknownMemberType = false

scripts/visualize_mnist_embeddings.py

@@ -0,0 +1,56 @@
+from pathlib import Path
+
+import matplotlib.pyplot as plt
+import torch
+from glass_box_umap import GlassBoxUMAP, ParametricUMAP
+from umap import UMAP
+
+X = torch.load("./tests/fixtures/mnist_images.pt")
+y = torch.load("./tests/fixtures/mnist_labels.pt")
+
+color_dict = {
+    0: "red",
+    1: "green",
+    2: "blue",
+    3: "black",
+    4: "magenta",
+    5: "cyan",
+    6: "orange",
+    7: "salmon",
+    8: "violet",
+    9: "purple",
+}
+colors = [color_dict.get(x.item(), "grey") for x in y]
+
+reducer = UMAP()
+reducer.fit(X)
+output = reducer.transform(X)
+
+plt.scatter(output[:, 0], output[:, 1], c=colors, s=5)
+plt.savefig("umap_original.png")
+plt.close()
+
+for epoch in [1, 5, 10, 50, 100, 200, 500]:
+    reducer = ParametricUMAP(
+        epochs=epoch,
+        lr=1e-3,
+        batch_size=256,
+        repulsion_strength=1.0,
+        encoder_name="default",
+        checkpoint_dir=Path(f"tmp_{epoch}"),
+    )
+    # reducer = GlassBoxUMAP(
+    #    epochs=epoch,
+    #    lr=1e-3,
+    #    batch_size=256,
+    #    repulsion_strength=1.0,
+    #    encoder_kwargs={"hidden_size": 1024},
+    #    checkpoint_dir=Path(f"tmp_{epoch}"),
+    # )
+
+    reducer.fit(X)
+    output = reducer.transform(X)
+
+    plt.scatter(output[:, 0], output[:, 1], c=colors, s=5)
+    plt.savefig(f"umap_glassbox_{epoch}.png")
+    plt.close()

src/glass_box_umap/__init__.py

@@ -1,8 +1,7 @@
 from .core import GlassBoxUMAP
-from .parametric_umap import ParametricUMAP, load_pumap
+from .parametric_umap import ParametricUMAP
 
 __all__ = [
     "GlassBoxUMAP",
     "ParametricUMAP",
-    "load_pumap",
 ]

src/glass_box_umap/components.py

@@ -0,0 +1,60 @@
+from __future__ import annotations
+import math
+
+import torch
+from torch import nn
+
+
+class LayerNormDetached(nn.Module):
+    """A LayerNorm with detached variance calculation during evaluation."""
+
+    def __init__(self, emb_dim: int):
+        super().__init__()
+        self.scale = nn.Parameter(torch.ones(emb_dim))
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        mean = x.mean(dim=-1, keepdim=True)
+
+        # Detach variance calculation during evaluation
+        if not self.training:
+            var = x.clone().detach().var(dim=-1, keepdim=True, unbiased=False)
+        else:
+            var = x.var(dim=-1, keepdim=True, unbiased=False)
+
+        norm_x = (x - mean) / torch.sqrt(var + 1e-12)  # Added epsilon for stability
+        return self.scale * norm_x
+
+
+class DeepPReLUNet(nn.Module):
+    """A network with PReLU activation and LayerNormDetached."""
+
+    def __init__(
+        self,
+        input_dims: tuple[int, ...],
+        n_components: int = 2,
+        hidden_size: int = 256,
+        n_hidden_layers: int = 5,
+        dropout_rate: float = 0.1,
+    ):
+        super().__init__()
+
+        input_size = math.prod(input_dims)
+        self.flatten = nn.Flatten()
+
+        layers = []
+        for i in range(n_hidden_layers):
+            in_dim = input_size if i == 0 else hidden_size
+
+            layers.append(nn.Linear(in_dim, hidden_size, bias=False))
+            layers.append(nn.PReLU())
+
+            if i < n_hidden_layers - 1:
+                layers.append(LayerNormDetached(hidden_size))
+
+            layers.append(nn.Dropout(dropout_rate))
+
+        layers.append(nn.Linear(hidden_size, n_components, bias=False))
+        self.model = nn.Sequential(*layers)
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        return self.model(self.flatten(x))