Deduplicate utilities and add kernel size guard

- L4: load_graphs/get_reference_dataset in 01-03 now delegate to
  utils/data.py instead of local copies
- L6: load_results extracted to utils/formatting.py, removed from
  4 format scripts
- P1: Warn when kernel matrix exceeds 10k samples in kernel_lr.py

Author: pjhartout

polygraph/metrics/base/kernel_lr.py

@@ -156,6 +156,14 @@ def fit(
         self.X_train_ = X
         self.y_train_ = y
 
+        n_samples = len(X) if isinstance(X, list) else int(X.shape[0])  # pyright: ignore[reportOptionalSubscript]
+        if n_samples > 10_000:
+            warnings.warn(
+                f"Kernel matrix will require ~{n_samples**2 * 8 / 1e9:.1f} GB "
+                f"of memory for {n_samples} samples. Consider reducing the "
+                f"dataset size."
+            )
+
         K = self._compute_kernel_matrix(X)
         if self.normalize_kernel:
             self.train_diag_ = np.diag(K).astype(np.float64)

reproducibility/01_subsampling/compute_mmd.py

@@ -14,15 +14,14 @@
 """
 
 import json
-import pickle
+import sys
 import time
 from pathlib import Path
 from typing import List, Literal, cast
 
 import hydra
 import networkx as nx
 import numpy as np
-import torch
 from loguru import logger
 from omegaconf import DictConfig
 from pyprojroot import here
@@ -41,6 +40,10 @@
 )
 from polygraph.utils.kernels import AdaptiveRBFKernel
 
+sys.path.insert(0, str(here() / "reproducibility"))
+from utils.data import get_reference_dataset
+from utils.data import load_graphs as _load
+
 REPO_ROOT = here()
 DATA_DIR = REPO_ROOT / "data"
 EXPERIMENT_RESULTS_DIR = (
@@ -64,49 +67,7 @@
 
 def load_graphs(model: str, dataset: str) -> List[nx.Graph]:
     """Load model-generated graphs from ``data/{model}/{dataset}.pkl``."""
-    pkl_path = DATA_DIR / model / f"{dataset}.pkl"
-    if not pkl_path.exists():
-        raise FileNotFoundError(f"{pkl_path} not found")
-    with open(pkl_path, "rb") as f:
-        graphs = pickle.load(f)
-
-    cleaned: List[nx.Graph] = []
-    for g in graphs:
-        if isinstance(g, nx.Graph):
-            simple = nx.Graph(g)
-        elif isinstance(g, (list, tuple)) and len(g) == 2:
-            try:
-                _node_feat, adj = g
-                if isinstance(adj, torch.Tensor):
-                    adj = adj.numpy()
-                simple = nx.from_numpy_array(adj)
-            except Exception as e:
-                logger.warning("Could not convert graph: {}", e)
-                continue
-        else:
-            logger.warning("Unknown graph format: {}", type(g))
-            continue
-        simple.remove_edges_from(nx.selfloop_edges(simple))
-        cleaned.append(simple)
-    return cleaned
-
-
-def get_reference_dataset(
-    dataset: str, split: str = "train", num_graphs: int = 4096
-) -> List[nx.Graph]:
-    """Get reference dataset from polygraph procedural generators."""
-    from polygraph.datasets.lobster import ProceduralLobsterGraphDataset
-    from polygraph.datasets.planar import ProceduralPlanarGraphDataset
-    from polygraph.datasets.sbm import ProceduralSBMGraphDataset
-
-    classes = {
-        "planar": ProceduralPlanarGraphDataset,
-        "lobster": ProceduralLobsterGraphDataset,
-        "sbm": ProceduralSBMGraphDataset,
-    }
-    if dataset not in classes:
-        raise ValueError(f"Unknown dataset: {dataset}")
-    return list(classes[dataset](split=split, num_graphs=num_graphs).to_nx())
+    return _load(DATA_DIR, model, dataset)
 
 
 def make_descriptor(name: str, reference_graphs: List[nx.Graph]):

reproducibility/01_subsampling/compute_pgd.py

@@ -14,7 +14,6 @@
 """
 
 import json
-import pickle
 import sys
 import time
 from importlib.metadata import version as pkg_version
@@ -23,7 +22,6 @@
 
 import hydra
 import networkx as nx
-import torch
 from loguru import logger
 from omegaconf import DictConfig
 from pyprojroot import here
@@ -34,6 +32,8 @@
 )
 
 sys.path.insert(0, str(here() / "reproducibility"))
+from utils.data import get_reference_dataset
+from utils.data import load_graphs as _load
 from utils.data import make_tabpfn_classifier
 
 
@@ -46,49 +46,7 @@
 
 def load_graphs(model: str, dataset: str) -> List[nx.Graph]:
     """Load model-generated graphs from ``data/{model}/{dataset}.pkl``."""
-    pkl_path = DATA_DIR / model / f"{dataset}.pkl"
-    if not pkl_path.exists():
-        raise FileNotFoundError(f"{pkl_path} not found")
-    with open(pkl_path, "rb") as f:
-        graphs = pickle.load(f)
-
-    cleaned: List[nx.Graph] = []
-    for g in graphs:
-        if isinstance(g, nx.Graph):
-            simple = nx.Graph(g)
-        elif isinstance(g, (list, tuple)) and len(g) == 2:
-            try:
-                _node_feat, adj = g
-                if isinstance(adj, torch.Tensor):
-                    adj = adj.numpy()
-                simple = nx.from_numpy_array(adj)
-            except Exception as e:
-                logger.warning("Could not convert graph: {}", e)
-                continue
-        else:
-            logger.warning("Unknown graph format: {}", type(g))
-            continue
-        simple.remove_edges_from(nx.selfloop_edges(simple))
-        cleaned.append(simple)
-    return cleaned
-
-
-def get_reference_dataset(
-    dataset: str, split: str = "train", num_graphs: int = 4096
-) -> List[nx.Graph]:
-    """Get reference dataset from polygraph procedural generators."""
-    from polygraph.datasets.lobster import ProceduralLobsterGraphDataset
-    from polygraph.datasets.planar import ProceduralPlanarGraphDataset
-    from polygraph.datasets.sbm import ProceduralSBMGraphDataset
-
-    classes = {
-        "planar": ProceduralPlanarGraphDataset,
-        "lobster": ProceduralLobsterGraphDataset,
-        "sbm": ProceduralSBMGraphDataset,
-    }
-    if dataset not in classes:
-        raise ValueError(f"Unknown dataset: {dataset}")
-    return list(classes[dataset](split=split, num_graphs=num_graphs).to_nx())
+    return _load(DATA_DIR, model, dataset)
 
 
 @hydra.main(

reproducibility/03_model_quality/compute.py

@@ -12,11 +12,10 @@
 """
 
 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
+from typing import Any, List, Literal, Tuple, cast
 
 import hydra
 import networkx as nx
@@ -30,6 +29,7 @@
 )
 
 sys.path.insert(0, str(here() / "reproducibility"))
+from utils.data import load_graphs as _load
 from utils.data import make_tabpfn_classifier
 
 REPO_ROOT = here()
@@ -40,46 +40,37 @@
 
 
 def load_graphs(path: Path) -> List[nx.Graph]:
-    """Load graphs from pickle file and convert to networkx."""
-    if not path.exists():
-        logger.warning("{} not found", path)
-        return []
-    with open(path, "rb") as f:
-        data = pickle.load(f)
-    graphs = []
-    for item in data:
-        if isinstance(item, nx.Graph):
-            graphs.append(item)
-        elif isinstance(item, (tuple, list)) and len(item) >= 2:
-            adj = item[1]
-            if hasattr(adj, "numpy"):
-                adj = adj.numpy()
-            graphs.append(nx.from_numpy_array(adj))
-        else:
-            graphs.append(nx.from_numpy_array(np.array(item)))
-    return graphs
+    """Load graphs from a single pickle file.
+
+    Delegates to ``utils.data.load_graphs`` by extracting the parent
+    directory and stem so that the caller-facing ``(path)`` signature
+    is preserved.
+    """
+    return _load(path.parent, "", path.stem)
 
 
 def get_reference_dataset(
     dataset: str,
     split: Literal["train", "val", "test"] = "train",
     num_graphs: int = 2048,
-):
-    """Get reference dataset from polygraph library."""
-    if dataset == "planar":
-        from polygraph.datasets.planar import ProceduralPlanarGraphDataset
-
-        ds = ProceduralPlanarGraphDataset(split=split, num_graphs=num_graphs)
-    elif dataset == "sbm":
-        from polygraph.datasets.sbm import ProceduralSBMGraphDataset
-
-        ds = ProceduralSBMGraphDataset(split=split, num_graphs=num_graphs)
-    elif dataset == "lobster":
-        from polygraph.datasets.lobster import ProceduralLobsterGraphDataset
-
-        ds = ProceduralLobsterGraphDataset(split=split, num_graphs=num_graphs)
-    else:
+) -> Tuple[Any, List[nx.Graph]]:
+    """Get reference dataset from polygraph library.
+
+    Returns ``(dataset_object, graphs)`` so callers can also call
+    ``dataset_object.is_valid()``.
+    """
+    from polygraph.datasets.lobster import ProceduralLobsterGraphDataset
+    from polygraph.datasets.planar import ProceduralPlanarGraphDataset
+    from polygraph.datasets.sbm import ProceduralSBMGraphDataset
+
+    procedural = {
+        "planar": ProceduralPlanarGraphDataset,
+        "lobster": ProceduralLobsterGraphDataset,
+        "sbm": ProceduralSBMGraphDataset,
+    }
+    if dataset not in procedural:
         raise ValueError(f"Unknown dataset: {dataset}")
+    ds = procedural[dataset](split=split, num_graphs=num_graphs)
     return ds, list(ds.to_nx())
 
 

reproducibility/05_benchmark/format.py

@@ -7,10 +7,8 @@
     python format.py
 """
 
-import json
 import sys
-from pathlib import Path
-from typing import Dict, List
+from typing import Dict
 
 import pandas as pd
 import typer
@@ -25,6 +23,7 @@
     MODELS,
     best_two,
     fmt_pgs,
+    load_results,
 )
 
 app = typer.Typer()
@@ -43,26 +42,6 @@
 ]
 
 
-def load_results(results_dir: Path) -> List[Dict]:
-    results = []
-    for f in sorted(results_dir.glob("*.json")):
-        with open(f) as fh:
-            results.append(json.load(fh))
-    return results
-
-
-def _reshape(result_list: List[Dict]) -> Dict[str, Dict]:
-    all_results: Dict[str, Dict] = {}
-    for r in result_list:
-        r = r.copy()
-        ds = r.pop("dataset", None)
-        model = r.pop("model", None)
-        r.pop("error", None)
-        if ds and model:
-            all_results.setdefault(ds, {})[model] = r
-    return all_results
-
-
 def generate_benchmark_table(all_results: Dict) -> str:
     lines = []
     lines.append(
@@ -167,14 +146,12 @@ def generate_benchmark_table(all_results: Dict) -> str:
 def main():
     """Generate LaTeX tables from pre-computed JSON results."""
     results_dir = OUTPUT_DIR / "results" / "benchmark"
-    result_list = load_results(results_dir)
-    if not result_list:
+    all_results = load_results(results_dir)
+    if not all_results:
         logger.error(
             "No results found in {}. Run compute.py first.", results_dir
         )
         return
-
-    all_results = _reshape(result_list)
     OUTPUT_DIR.mkdir(parents=True, exist_ok=True)
 
     table = generate_benchmark_table(all_results)

reproducibility/06_mmd/format.py

@@ -10,9 +10,7 @@
     python format.py
 """
 
-import json
 import sys
-from pathlib import Path
 from typing import Dict, List
 
 import pandas as pd
@@ -28,6 +26,7 @@
     MODELS,
     best_two,
     fmt_sci,
+    load_results,
 )
 
 app = typer.Typer()
@@ -38,19 +37,6 @@
 BENCHMARK_RESULTS_DIR = OUTPUT_DIR / "results" / "benchmark"
 
 
-def load_results(results_dir: Path) -> Dict[str, Dict]:
-    all_r: Dict[str, Dict] = {}
-    if not results_dir.exists():
-        return all_r
-    for f in sorted(results_dir.glob("*.json")):
-        with open(f) as fh:
-            r = json.load(fh)
-        ds, model = r.get("dataset"), r.get("model")
-        if ds and model:
-            all_r.setdefault(ds, {})[model] = r
-    return all_r
-
-
 def _fmt_pgs(mean: float, std: float, is_best=False, is_second=False) -> str:
     """MMD tables use 3 decimal places for PGD scores (not 1 like benchmark)."""
     if pd.isna(mean):