Deduplicate TabPFN factory and VUN logic, add sparse eigenvalue path

- H6: 05_benchmark/compute_vun.py now imports compute_vun_parallel
  from utils.vun instead of duplicating ~120 lines of VUN logic
- L5: Extracted make_tabpfn_classifier to utils/data.py, replaced
  6 local copies across compute scripts
- P4: EigenvalueHistogram uses scipy.sparse.linalg.eigsh for graphs
  with >500 nodes, avoiding dense conversion of large Laplacians

Author: pjhartout

polygraph/utils/descriptors/generic_descriptors.py

@@ -18,6 +18,7 @@
 import orbit_count
 import torch
 from scipy.sparse import csgraph, csr_array
+from scipy.sparse.linalg import eigsh
 from sklearn.decomposition import TruncatedSVD
 from sklearn.preprocessing import StandardScaler
 from torch_geometric.data import Batch
@@ -225,14 +226,24 @@ def __init__(self, n_bins: int = 200, sparse: bool = False):
         else:
             self._bins = None
 
+    _SPARSE_THRESHOLD = 500
+
     def __call__(
         self, graphs: Iterable[nx.Graph]
     ) -> Union[np.ndarray, csr_array]:
         all_eigs = []
         for g in graphs:
-            eigs = np.linalg.eigvalsh(
-                nx.normalized_laplacian_matrix(g).todense()
-            )
+            n = g.number_of_nodes()
+            laplacian = nx.normalized_laplacian_matrix(g)
+            if n > self._SPARSE_THRESHOLD:
+                k = min(n - 2, self._n_bins)
+                eigs = eigsh(
+                    laplacian.astype(np.float64),
+                    k=k,
+                    return_eigenvectors=False,
+                )
+            else:
+                eigs = np.linalg.eigvalsh(laplacian.todense())
             all_eigs.append(eigs)
 
         if self._sparse:

reproducibility/01_subsampling/compute_pgd.py

@@ -15,6 +15,7 @@
 
 import json
 import pickle
+import sys
 import time
 from importlib.metadata import version as pkg_version
 from pathlib import Path
@@ -32,25 +33,8 @@
     maybe_append_jsonl,
 )
 
-
-def _make_tabpfn_classifier(weights_version: str):
-    """Create a TabPFN classifier for the given weights version."""
-    from tabpfn import TabPFNClassifier
-    from tabpfn.classifier import ModelVersion
-
-    version_map = {
-        "v2": ModelVersion.V2,
-        "v2.5": ModelVersion.V2_5,
-    }
-    if weights_version not in version_map:
-        raise ValueError(
-            f"Unknown weights_version: {weights_version!r}. Must be one of {list(version_map)}"
-        )
-    return TabPFNClassifier.create_default_for_version(
-        version_map[weights_version],
-        device="auto",
-        n_estimators=4,
-    )
+sys.path.insert(0, str(here() / "reproducibility"))
+from utils.data import make_tabpfn_classifier
 
 
 REPO_ROOT = here()
@@ -122,7 +106,7 @@ def main(cfg: DictConfig) -> None:
     model: str = cfg.model
     subsample_size: int = cfg.subsample_size
     num_bootstrap: int = 3 if cfg.subset else cfg.num_bootstrap
-    classifier = _make_tabpfn_classifier(tabpfn_weights_version)
+    classifier = make_tabpfn_classifier(tabpfn_weights_version)
 
     logger.info(
         "PGD subsampling: dataset={}, model={}, n={}, bootstraps={}",

reproducibility/02_perturbation/compute.py

@@ -26,6 +26,7 @@
 import gc
 import json
 import random
+import sys
 from importlib.metadata import version as pkg_version
 from itertools import product
 from typing import Any, Callable, Dict, List, Literal, Optional, Tuple, cast
@@ -43,6 +44,9 @@
     maybe_append_jsonl,
 )
 
+sys.path.insert(0, str(here() / "reproducibility"))
+from utils.data import make_tabpfn_classifier
+
 from polygraph.datasets.ego import EgoGraphDataset
 from polygraph.datasets.lobster import ProceduralLobsterGraphDataset
 from polygraph.datasets.planar import ProceduralPlanarGraphDataset
@@ -328,22 +332,7 @@ def load_dataset(
 def _make_classifier(name: str, tabpfn_weights_version: str = "v2.5"):
     """Build a classifier by name. For TabPFN, respects weights version."""
     if name == "tabpfn":
-        from tabpfn import TabPFNClassifier
-        from tabpfn.classifier import ModelVersion
-
-        version_map = {
-            "v2": ModelVersion.V2,
-            "v2.5": ModelVersion.V2_5,
-        }
-        if tabpfn_weights_version not in version_map:
-            raise ValueError(
-                f"Unknown tabpfn_weights_version: {tabpfn_weights_version!r}. Must be one of {list(version_map)}"
-            )
-        return TabPFNClassifier.create_default_for_version(
-            version_map[tabpfn_weights_version],
-            device="auto",
-            n_estimators=4,
-        )
+        return make_tabpfn_classifier(tabpfn_weights_version)
     elif name == "lr":
         return LogisticRegression(max_iter=1000)
     else:

reproducibility/03_model_quality/compute.py

@@ -13,6 +13,7 @@
 
 import json
 import pickle
+import sys
 from importlib.metadata import version as pkg_version
 from pathlib import Path
 from typing import Any, List, Literal, cast
@@ -28,6 +29,9 @@
     maybe_append_jsonl,
 )
 
+sys.path.insert(0, str(here() / "reproducibility"))
+from utils.data import make_tabpfn_classifier
+
 REPO_ROOT = here()
 DATA_DIR = REPO_ROOT / "data"
 _RESULTS_DIR_BASE = (
@@ -192,22 +196,7 @@ def _parse_steps(p: Path) -> int:
     if subset:
         ref = ref[:30]
 
-    from tabpfn import TabPFNClassifier
-    from tabpfn.classifier import ModelVersion
-
-    version_map = {
-        "v2": ModelVersion.V2,
-        "v2.5": ModelVersion.V2_5,
-    }
-    if tabpfn_weights_version not in version_map:
-        raise ValueError(
-            f"Unknown tabpfn_weights_version: {tabpfn_weights_version!r}. Must be one of {list(version_map)}"
-        )
-    classifier = TabPFNClassifier.create_default_for_version(
-        version_map[tabpfn_weights_version],
-        device="auto",
-        n_estimators=4,
-    )
+    classifier = make_tabpfn_classifier(tabpfn_weights_version)
 
     pgd_metric = StandardPGD(
         reference_graphs=ref,

reproducibility/05_benchmark/compute.py

@@ -25,6 +25,7 @@
 sys.path.insert(0, str(here() / "reproducibility"))
 from utils.data import get_reference_dataset as _get_ref
 from utils.data import load_graphs as _load
+from utils.data import make_tabpfn_classifier
 
 REPO_ROOT = here()
 DATA_DIR = REPO_ROOT / "data"
@@ -129,22 +130,7 @@ def main(cfg: DictConfig) -> None:
     subset = cfg.subset
     skip_vun = cfg.get("skip_vun", False)
 
-    from tabpfn import TabPFNClassifier
-    from tabpfn.classifier import ModelVersion
-
-    version_map = {
-        "v2": ModelVersion.V2,
-        "v2.5": ModelVersion.V2_5,
-    }
-    if tabpfn_weights_version not in version_map:
-        raise ValueError(
-            f"Unknown tabpfn_weights_version: {tabpfn_weights_version!r}. Must be one of {list(version_map)}"
-        )
-    classifier = TabPFNClassifier.create_default_for_version(
-        version_map[tabpfn_weights_version],
-        device="auto",
-        n_estimators=4,
-    )
+    classifier = make_tabpfn_classifier(tabpfn_weights_version)
 
     logger.info("Computing benchmark for {}/{}", model, dataset)