[wip][clustering] visualization dashboard

.github/workflows/checks.yaml

@@ -40,6 +40,9 @@ jobs:
       - name: Print dependencies
         run: uv pip list
 
+      - name: Install Playwright
+        run: make install-playwright
+
       - name: Run basedpyright
         run: uv run basedpyright
 

.gitignore

@@ -3,6 +3,7 @@ docs/coverage/**
 artifacts/**
 docs/dep_graph/**
 tests/.temp/**
+*.prof
 
 **/out/
 neuronpedia_outputs/

Makefile

@@ -77,6 +77,50 @@ coverage:
 	uv run python -m coverage html --directory=$(COVERAGE_DIR)/html/
 
 
+BUNDLED_DASHBOARD_DIR=spd/clustering/dashboard/_bundled
+
+.PHONY: bundle-dashboard
+bundle-dashboard:
+	@mkdir -p $(BUNDLED_DASHBOARD_DIR)
+	uv run python -m muutils.web.bundle_html \
+		spd/clustering/dashboard/index.html \
+		--output $(BUNDLED_DASHBOARD_DIR)/index.html \
+		--source-dir spd/clustering/dashboard
+	uv run python -m muutils.web.bundle_html \
+		spd/clustering/dashboard/cluster.html \
+		--output $(BUNDLED_DASHBOARD_DIR)/cluster.html \
+		--source-dir spd/clustering/dashboard
+	@echo "Bundled HTML files to $(BUNDLED_DASHBOARD_DIR)/"
+
+.PHONY: clean-test-dashboard
+clean-test-dashboard:
+	rm -rf tests/.temp/dashboard-integration
+
+
+.PHONY: install-playwright
+install-playwright:
+	@echo "Install Playwright browsers, used for dashboard tests"
+	uv run playwright install chromium
+	uv run playwright install-deps
+
+.PHONY: test-dashboard
+test-dashboard: clean-test-dashboard bundle-dashboard
+	pytest tests/clustering/dashboard/test_dashboard_integration.py --runslow -v --durations 10
+
+
+.PHONY: clustering-dashboard
+clustering-dashboard: bundle-dashboard
+	uv run python spd/clustering/dashboard/run.py \
+		spd/clustering/dashboard/dashboard_config.yaml
+
+.PHONY: clustering-dashboard-profile
+clustering-dashboard-profile: bundle-dashboard
+	uv run python -m cProfile -o dashboard.prof spd/clustering/dashboard/run.py \
+		spd/clustering/dashboard/dashboard_config.yaml
+	@echo "\nProfile saved to dashboard.prof"
+	@echo "View with: python -m pstats dashboard.prof"
+	@echo "Or install snakeviz and run: snakeviz dashboard.prof"
+
 .PHONY: clean
 clean:
 	@echo "Cleaning Python cache and build artifacts..."

pyproject.toml

@@ -30,6 +30,7 @@ dependencies = [
     "simple_stories_train @ git+https://github.com/goodfire-ai/simple_stories_train.git@dev",
     "scipy>=1.14.1",
     "muutils",
+    "zanj", # for dashboard data saving/loading
     "fastapi",
     "uvicorn",
 ]
@@ -42,6 +43,7 @@ dev = [
     "ruff",
     "basedpyright<1.32.0", # pyright and wandb issues, see https://github.com/goodfire-ai/spd/pull/232
     "pre-commit",
+    "playwright", # for browser-based integration tests
 ]
 
 [project.scripts]

spd/clustering/activations.py

@@ -10,7 +10,7 @@
     ActivationsTensor,
     BoolActivationsTensor,
     ClusterCoactivationShaped,
-    ComponentLabels,
+    SubComponentKey,
 )
 from spd.clustering.util import ModuleFilterFunc
 from spd.models.component_model import ComponentModel, OutputWithCache
@@ -54,16 +54,16 @@ class FilteredActivations(NamedTuple):
     activations: ActivationsTensor
     "activations after filtering dead components"
 
-    labels: ComponentLabels
-    "list of length c with labels for each preserved component"
+    subcomponent_keys: list[SubComponentKey]
+    "list of length c with SubComponentInfo for each preserved component"
 
-    dead_components_labels: ComponentLabels | None
-    "list of labels for dead components, or None if no filtering was applied"
+    dead_subcomponent_keys: list[SubComponentKey] | None
+    "list of SubComponentInfo for dead components, or None if no filtering was applied"
 
     @property
     def n_alive(self) -> int:
         """Number of alive components after filtering."""
-        n_alive: int = len(self.labels)
+        n_alive: int = len(self.subcomponent_keys)
         assert n_alive == self.activations.shape[1], (
             f"{n_alive = } != {self.activations.shape[1] = }"
         )
@@ -72,12 +72,12 @@ def n_alive(self) -> int:
     @property
     def n_dead(self) -> int:
         """Number of dead components after filtering."""
-        return len(self.dead_components_labels) if self.dead_components_labels else 0
+        return len(self.dead_subcomponent_keys) if self.dead_subcomponent_keys else 0
 
 
 def filter_dead_components(
     activations: ActivationsTensor,
-    labels: ComponentLabels,
+    subcomponent_keys: list[SubComponentKey],
     filter_dead_threshold: float = 0.01,
 ) -> FilteredActivations:
     """Filter out dead components based on a threshold
@@ -86,31 +86,29 @@ def filter_dead_components(
     activations and labels are returned as is, `dead_components_labels` is `None`.
 
     otherwise, components whose **maximum** activations across all samples is below the threshold
-    are considered dead and filtered out. The labels of these components are returned in `dead_components_labels`.
+    are considered dead and filtered out. The SubComponentInfo of these components are returned in `dead_components_labels`.
     `dead_components_labels` will also be `None` if no components were below the threshold.
     """
-    dead_components_lst: ComponentLabels | None = None
+    dead_components_lst: list[SubComponentKey] | None = None
     if filter_dead_threshold > 0:
-        dead_components_lst = ComponentLabels(list())
+        dead_components_lst = []
         max_act: Float[Tensor, " c"] = activations.max(dim=0).values
-        dead_components: Bool[Tensor, " c"] = max_act < filter_dead_threshold
+        dead_components_mask: Bool[Tensor, " c"] = max_act < filter_dead_threshold
 
-        if dead_components.any():
-            activations = activations[:, ~dead_components]
-            alive_labels: list[tuple[str, bool]] = [
-                (lbl, bool(keep.item()))
-                for lbl, keep in zip(labels, ~dead_components, strict=False)
+        if dead_components_mask.any():
+            activations = activations[:, ~dead_components_mask]
+            alive_labels: list[tuple[SubComponentKey, bool]] = [
+                (comp, bool(keep.item()))
+                for comp, keep in zip(subcomponent_keys, ~dead_components_mask, strict=False)
             ]
             # re-assign labels only if we are filtering
-            labels = ComponentLabels([label for label, keep in alive_labels if keep])
-            dead_components_lst = ComponentLabels(
-                [label for label, keep in alive_labels if not keep]
-            )
+            subcomponent_keys = [comp for comp, keep in alive_labels if keep]
+            dead_components_lst = [comp for comp, keep in alive_labels if not keep]
 
     return FilteredActivations(
         activations=activations,
-        labels=labels,
-        dead_components_labels=dead_components_lst if dead_components_lst else None,
+        subcomponent_keys=subcomponent_keys,
+        dead_subcomponent_keys=dead_components_lst if dead_components_lst else None,
     )
 
 
@@ -124,11 +122,11 @@ class ProcessedActivations:
     activations: ActivationsTensor
     "activations after filtering and concatenation"
 
-    labels: ComponentLabels
-    "list of length c with labels for each preserved component, format `{module_name}:{component_index}`"
+    subcomponent_keys: list[SubComponentKey]
+    "list of length c with SubComponentInfo for each preserved component"
 
-    dead_components_lst: ComponentLabels | None
-    "list of labels for dead components, or None if no filtering was applied"
+    dead_subcomponent_keys: list[SubComponentKey] | None
+    "list of SubComponentInfo for dead components, or None if no filtering was applied"
 
     def validate(self) -> None:
         """Validate the processed activations"""
@@ -143,7 +141,7 @@ def n_components_original(self) -> int:
     @property
     def n_components_alive(self) -> int:
         """Number of alive components after filtering. equal to the length of `labels`"""
-        n_alive: int = len(self.labels)
+        n_alive: int = len(self.subcomponent_keys)
         assert n_alive + self.n_components_dead == self.n_components_original, (
             f"({n_alive = }) + ({self.n_components_dead = }) != ({self.n_components_original = })"
         )
@@ -156,26 +154,26 @@ def n_components_alive(self) -> int:
     @property
     def n_components_dead(self) -> int:
         """Number of dead components after filtering. equal to the length of `dead_components_lst` if it is not None, or 0 otherwise"""
-        return len(self.dead_components_lst) if self.dead_components_lst else 0
+        return len(self.dead_subcomponent_keys) if self.dead_subcomponent_keys else 0
 
     @cached_property
     def label_index(self) -> dict[str, int | None]:
-        """Create a mapping from label to alive index (`None` if dead)"""
+        """Create a mapping from label string to alive index (`None` if dead)"""
         return {
-            **{label: i for i, label in enumerate(self.labels)},
+            **{comp.label: i for i, comp in enumerate(self.subcomponent_keys)},
             **(
-                {label: None for label in self.dead_components_lst}
-                if self.dead_components_lst
+                {comp.label: None for comp in self.dead_subcomponent_keys}
+                if self.dead_subcomponent_keys
                 else {}
             ),
         }
 
     def get_label_index(self, label: str) -> int | None:
-        """Get the index of a label in the activations, or None if it is dead"""
+        """Get the index of a label string in the activations, or None if it is dead"""
         return self.label_index[label]
 
     def get_label_index_alive(self, label: str) -> int:
-        """Get the index of a label in the activations, or raise if it is dead"""
+        """Get the index of a label string in the activations, or raise if it is dead"""
         idx: int | None = self.get_label_index(label)
         if idx is None:
             raise ValueError(f"Label '{label}' is dead and has no index in the activations.")
@@ -239,10 +237,10 @@ def process_activations(
 
     # compute the labels and total component count
     total_c: int = 0
-    labels: ComponentLabels = ComponentLabels(list())
+    labels: list[SubComponentKey] = []
     for key, act in activations_.items():
         c: int = act.shape[-1]
-        labels.extend([f"{key}:{i}" for i in range(c)])
+        labels.extend([SubComponentKey(module=key, index=i) for i in range(c)])
         total_c += c
 
     # concat the activations
@@ -251,7 +249,7 @@ def process_activations(
     # filter dead components
     filtered_components: FilteredActivations = filter_dead_components(
         activations=act_concat,
-        labels=labels,
+        subcomponent_keys=labels,
         filter_dead_threshold=filter_dead_threshold,
     )
 
@@ -262,6 +260,6 @@ def process_activations(
     return ProcessedActivations(
         activations_raw=activations_,
         activations=filtered_components.activations,
-        labels=filtered_components.labels,
-        dead_components_lst=filtered_components.dead_components_labels,
+        subcomponent_keys=filtered_components.subcomponent_keys,
+        dead_subcomponent_keys=filtered_components.dead_subcomponent_keys,
     )

spd/clustering/consts.py

@@ -1,8 +1,10 @@
 """Constants and shared abstractions for clustering pipeline."""
 
+import hashlib
 from abc import ABC, abstractmethod
+from dataclasses import dataclass
 from pathlib import Path
-from typing import Literal, NewType
+from typing import Literal, NewType, override
 
 import numpy as np
 from jaxtyping import Bool, Float, Int
@@ -15,8 +17,39 @@
 DistancesArray = Float[np.ndarray, "n_iters n_ens n_ens"]
 
 # Component and label types (NewType for stronger type safety)
-ComponentLabel = NewType("ComponentLabel", str)  # Format: "module_name:component_index"
-ComponentLabels = NewType("ComponentLabels", list[str])
+SubComponentLabel = NewType("SubComponentLabel", str)  # Format: "module_name:component_index"
+
+
+@dataclass(frozen=True, slots=True, kw_only=True)
+class SubComponentKey:
+    """unique identifier of a subcomponent. indices can refer to dead components"""
+
+    module: str
+    index: int
+
+    @property
+    def label(self) -> SubComponentLabel:
+        """Component label as 'module:index'."""
+        return SubComponentLabel(f"{self.module}:{self.index}")
+
+    @classmethod
+    def from_label(cls, label: SubComponentLabel) -> "SubComponentKey":
+        """Create SubComponentInfo from a component label."""
+        assert label.count(":") == 1, (
+            "Invalid component label format, expected '{{module}}:{{index}}'"
+        )
+        module, index_str = label.rsplit(":", 1)
+        return cls(module=module, index=int(index_str))
+
+    @override
+    def __str__(self) -> str:
+        return self.label
+
+    @override
+    def __hash__(self) -> int:
+        return int(hashlib.md5(str(self).encode()).hexdigest(), 16)
+
+
 BatchId = NewType("BatchId", str)
 
 # Path types