Updated type hints

Markus28 · Markus28 · commit 41c6d57b16a2 · 2025-08-06T15:03:30.000+02:00
diff --git a/polygraph/datasets/base/dataset.py b/polygraph/datasets/base/dataset.py
@@ -370,7 +370,7 @@ def url_for_split(self, split: str) -> str:
         ...
 
     @abstractmethod
-    def hash_for_split(self, split: str) -> str:
+    def hash_for_split(self, split: str) -> Optional[str]:
         """Gets the expected hash for a specific split's data.
 
         This hash is used to validate downloaded data.
diff --git a/polygraph/datasets/ego.py b/polygraph/datasets/ego.py
@@ -1,3 +1,4 @@
+from typing import Optional
 import networkx as nx
 
 from polygraph.datasets.base import OnlineGraphDataset
@@ -42,7 +43,7 @@ def url_for_split(self, split: str):
     def is_valid(graph: nx.Graph):
         return graph.number_of_nodes() > 0 and graph.number_of_edges() > 0
 
-    def hash_for_split(self, split: str) -> str:
+    def hash_for_split(self, split: str) -> Optional[str]:
         return self._HASH_FOR_SPLIT[split]
 
 
@@ -78,5 +79,5 @@ def url_for_split(self, split: str):
     def is_valid(graph: nx.Graph):
         return graph.number_of_nodes() > 0 and graph.number_of_edges() > 0
 
-    def hash_for_split(self, split: str) -> str:
+    def hash_for_split(self, split: str) -> Optional[str]:
         return self._HASH_FOR_SPLIT[split]
diff --git a/polygraph/datasets/lobster.py b/polygraph/datasets/lobster.py
@@ -16,15 +16,15 @@ def is_lobster_graph(graph: nx.Graph) -> bool:
     """Based on https://github.com/lrjconan/GRAN/blob/fc9c04a3f002c55acf892f864c03c6040947bc6b/utils/eval_helper.py#L426C3-L446C17"""
     graph = deepcopy(graph)
     if nx.is_tree(graph):
-        leaves = [n for n, d in graph.degree() if d == 1]
+        leaves = [n for n, d in graph.degree() if d == 1]   # pyright: ignore
         graph.remove_nodes_from(leaves)
 
-        leaves = [n for n, d in graph.degree() if d == 1]
+        leaves = [n for n, d in graph.degree() if d == 1]   # pyright: ignore
         graph.remove_nodes_from(leaves)
 
         num_nodes = len(graph.nodes())
-        num_degree_one = [d for n, d in graph.degree() if d == 1]
-        num_degree_two = [d for n, d in graph.degree() if d == 2]
+        num_degree_one = [d for n, d in graph.degree() if d == 1]   # pyright: ignore
+        num_degree_two = [d for n, d in graph.degree() if d == 2]   # pyright: ignore
 
         if sum(num_degree_one) == 2 and sum(num_degree_two) == 2 * (
             num_nodes - 2
@@ -166,5 +166,5 @@ def is_valid(graph: nx.Graph) -> bool:
         """Check if a graph is a valid lobster graph."""
         return is_lobster_graph(graph)
 
-    def hash_for_split(self, split: str) -> str:
+    def hash_for_split(self, split: str) -> Optional[str]:
         return self._HASH_FOR_SPLIT[split]
diff --git a/polygraph/datasets/modelnet.py b/polygraph/datasets/modelnet.py
@@ -1,3 +1,4 @@
+from typing import Optional
 import networkx as nx
 
 from polygraph.datasets.base import OnlineGraphDataset
@@ -44,5 +45,5 @@ def url_for_split(self, split: str):
     def is_valid(graph: nx.Graph):
         return graph.number_of_nodes() > 0 and graph.number_of_edges() > 0
 
-    def hash_for_split(self, split: str) -> str:
+    def hash_for_split(self, split: str) -> Optional[str]:
         return self._HASH_FOR_SPLIT[split]
diff --git a/polygraph/datasets/planar.py b/polygraph/datasets/planar.py
@@ -1,4 +1,4 @@
-from typing import Literal
+from typing import Literal, Optional
 
 import joblib
 import networkx as nx
@@ -37,7 +37,7 @@ def __init__(
         memmap: bool = False,
         show_generation_progress: bool = False,
     ):
-        config_hash = joblib.hash(
+        config_hash: str = joblib.hash(
             (num_graphs, n_nodes, seed, split), hash_name="md5"
         )
         self._rng = np.random.default_rng(
@@ -132,5 +132,5 @@ def is_valid(graph: nx.Graph) -> bool:
         """
         return is_planar_graph(graph)
 
-    def hash_for_split(self, split: str) -> str:
+    def hash_for_split(self, split: str) -> Optional[str]:
         return self._HASH_FOR_SPLIT[split]
diff --git a/polygraph/datasets/point_clouds.py b/polygraph/datasets/point_clouds.py
@@ -1,3 +1,4 @@
+from typing import Optional
 import networkx as nx
 
 from polygraph.datasets.base import OnlineGraphDataset
@@ -43,5 +44,5 @@ def url_for_split(self, split: str):
     def is_valid(graph: nx.Graph):
         return graph.number_of_nodes() > 0 and graph.number_of_edges() > 0
 
-    def hash_for_split(self, split: str) -> str:
+    def hash_for_split(self, split: str) -> Optional[str]:
         return self._HASH_FOR_SPLIT[split]
diff --git a/polygraph/datasets/proteins.py b/polygraph/datasets/proteins.py
@@ -1,3 +1,4 @@
+from typing import Optional
 import networkx as nx
 
 from polygraph.datasets.base import OnlineGraphDataset
@@ -55,5 +56,5 @@ def url_for_split(self, split: str):
     def is_valid(graph: nx.Graph):
         return graph.number_of_nodes() > 0 and graph.number_of_edges() > 0
 
-    def hash_for_split(self, split: str) -> str:
+    def hash_for_split(self, split: str) -> Optional[str]:
         return self._HASH_FOR_SPLIT[split]
diff --git a/polygraph/datasets/sbm.py b/polygraph/datasets/sbm.py
@@ -1,4 +1,4 @@
-from typing import Literal, Tuple
+from typing import Literal, Tuple, Optional
 
 import joblib
 import networkx as nx
@@ -220,7 +220,7 @@ def __init__(
         memmap: bool = False,
         show_generation_progress: bool = False,
     ):
-        config_hash = joblib.hash(
+        config_hash: str = joblib.hash(
             (
                 num_graphs,
                 intra_p,
@@ -358,5 +358,5 @@ def is_valid(graph: nx.Graph) -> bool:
     def is_valid_alt(graph: nx.Graph) -> bool:
         return is_sbm_graph_alt(graph)
 
-    def hash_for_split(self, split: str) -> str:
+    def hash_for_split(self, split: str) -> Optional[str]:
         return self._HASH_FOR_SPLIT[split]
diff --git a/polygraph/metrics/base/metric_interval.py b/polygraph/metrics/base/metric_interval.py
@@ -47,7 +47,7 @@ def from_samples(
 
         return cls(mean=mean, std=std, low=low, high=high, coverage=coverage)
 
-    def __getitem__(self, key: str) -> float:
+    def __getitem__(self, key: str) -> Optional[float]:
         if key == "mean":
             return self.mean
         elif key == "std":
diff --git a/polygraph/metrics/base/mmd.py b/polygraph/metrics/base/mmd.py
@@ -39,8 +39,7 @@
 """
 
 from abc import ABC, abstractmethod
-from collections import namedtuple
-from typing import Collection, Dict, Literal, Tuple, Union
+from typing import Any, Collection, Literal, Union
 
 import networkx as nx
 import numpy as np
@@ -67,6 +66,8 @@ class DescriptorMMD2(GenerationMetric):
         variant: Which MMD estimator to use ('biased', 'umve', or 'ustat')
     """
 
+    _variant: Literal["biased", "umve", "ustat"]
+
     def __init__(
         self,
         reference_graphs: Collection[nx.Graph],
@@ -136,13 +137,16 @@ def compute(self, generated_graphs: Collection[nx.Graph]) -> float:
             Maximum MMD² value across kernel parameters
         """
         multi_kernel_result = super().compute(generated_graphs)
+        assert isinstance(multi_kernel_result, np.ndarray)
         idx = int(np.argmax(multi_kernel_result))
         return multi_kernel_result[idx]
 
 
 class _DescriptorMMD2Interval(ABC):
     """Base class for computing MMD² confidence intervals through subsampling."""
 
+    _variant: Literal["biased", "umve", "ustat"]
+
     def __init__(
         self,
         reference_graphs: Collection[nx.Graph],
diff --git a/polygraph/metrics/base/polygraphscore.py b/polygraph/metrics/base/polygraphscore.py
@@ -62,51 +62,26 @@ def _scores_to_jsd(ref_scores, gen_scores, eps: float = 1e-10):
     return np.sqrt(np.clip(divergence, 0, 1))
 
 
-def _scores_to_informedness_and_threshold(ref_scores: np.ndarray, gen_scores: np.ndarray):
+def _scores_to_informedness_and_threshold(ref_scores: np.ndarray, gen_scores: np.ndarray) -> Tuple[float, float]:
     ground_truth = np.concatenate(
         [np.ones(len(ref_scores)), np.zeros(len(gen_scores))]
     )
-    if ref_scores.ndim == 2:
-        assert (
-            gen_scores.ndim == 2 and ref_scores.shape[1] == gen_scores.shape[1]
-        )
-        all_rocs = [
-            roc_curve(
-                ground_truth,
-                np.concatenate([ref_scores[:, i], gen_scores[:, i]]),
-            )
-            for i in range(gen_scores.shape[1])
-        ]
-        all_j_statistics = [tpr - fpr for fpr, tpr, _ in all_rocs]
-        all_thresholds = [thresholds for _, _, thresholds in all_rocs]
-        optimal_idxs = [
-            np.argmax(j_statistic) for j_statistic in all_j_statistics
-        ]
-        assert len(all_thresholds) == len(optimal_idxs)
-        optimal_threshold = [
-            thresholds[idx]
-            for thresholds, idx in zip(all_thresholds, optimal_idxs)
-        ]
-        assert len(all_j_statistics) == len(optimal_threshold)
-        j_statistic = np.array(
-            [
-                j_statistic[idx]
-                for j_statistic, idx in zip(all_j_statistics, optimal_idxs)
-            ]
-        )
-    else:
-        assert ref_scores.ndim == 1 and gen_scores.ndim == 1
-        fpr, tpr, thresholds = roc_curve(
-            ground_truth, np.concatenate([ref_scores, gen_scores])
-        )
-        j_statistic = tpr - fpr
-        optimal_idx = np.argmax(j_statistic)
-        optimal_threshold = thresholds[optimal_idx]
-        j_statistic = j_statistic[optimal_idx]
+    if ref_scores.ndim != 1:
+        raise RuntimeError("ref_scores must be 1-dimensional, got shape {ref_scores.shape}. This should not happen, please file a bug report.")
+
+    assert ref_scores.ndim == 1 and gen_scores.ndim == 1
+    fpr, tpr, thresholds = roc_curve(
+        ground_truth, np.concatenate([ref_scores, gen_scores])
+    )
+    j_statistic = tpr - fpr
+    optimal_idx = np.argmax(j_statistic)
+    optimal_threshold = thresholds[optimal_idx]
+    j_statistic = j_statistic[optimal_idx]
     return j_statistic, optimal_threshold
 
 
-def _scores_and_threshold_to_informedness(ref_scores: np.ndarray, gen_scores: np.ndarray, threshold: float):
+def _scores_and_threshold_to_informedness(ref_scores: np.ndarray, gen_scores: np.ndarray, threshold: float) -> float:
+    assert ref_scores.ndim == 1 and gen_scores.ndim == 1
     ref_pred = (ref_scores >= threshold).astype(int)
     gen_pred = (gen_scores >= threshold).astype(int)
     tpr = np.mean(ref_pred, axis=0)
@@ -212,6 +187,7 @@ def _descriptions_to_classifier_metric(
     rng = np.random.default_rng(0) if rng is None else rng
 
     if isinstance(ref_descriptions, csr_array):
+        assert isinstance(gen_descriptions, csr_array)
         # Convert to dense array
         num_features = (
             max(
@@ -262,17 +238,17 @@ def _descriptions_to_classifier_metric(
     train_gen_descriptions = scaler.transform(train_gen_descriptions)
 
     if classifier == "logistic":
-        classifier = LogisticRegression(penalty="l2", max_iter=1000)
+        clf = LogisticRegression(penalty="l2", max_iter=1000)
     elif classifier == "tabpfn":
-        classifier = TabPFNClassifier(
+        clf = TabPFNClassifier(
             device="cuda" if torch.cuda.is_available() else "cpu"
         )
     else:
         raise ValueError(f"Invalid classifier: {classifier}")
 
     # Use custom cross-validation function
     scores = _classifier_cross_validation(
-        classifier,
+        clf,
         train_ref_descriptions,
         train_gen_descriptions,
         variant,
@@ -299,8 +275,8 @@ def _descriptions_to_classifier_metric(
         )
         predict_proba = lambda x: np.ones((x.shape[0], 2)) * 0.5
     else:
-        classifier.fit(train_all_descriptions, train_labels)
-        predict_proba = classifier.predict_proba
+        clf.fit(train_all_descriptions, train_labels)
+        predict_proba = clf.predict_proba
 
     ref_test_pred = predict_proba(test_ref_descriptions)[:, 1]
     gen_test_pred = predict_proba(test_gen_descriptions)[:, 1]
@@ -331,6 +307,9 @@ class ClassifierMetric(GenerationMetric):
         variant: Classifier metric to compute. To estimate the Jensen-Shannon distance, use "jsd". To estimate total variation distance, use "informedness".
         classifier: Binary classifier to fit
     """
+    _variant: Literal["informedness", "jsd"]
+    _classifier: Literal["logistic", "tabpfn"]
+
     def __init__(
         self,
         reference_graphs: Collection[nx.Graph],
@@ -366,6 +345,9 @@ def compute(
 
 
 class _ClassifierMetricSamples:
+    _variant: Literal["informedness", "jsd"]
+    _classifier: Literal["logistic", "tabpfn"]
+
     def __init__(
         self,
         reference_graphs: Collection[nx.Graph],
@@ -420,6 +402,9 @@ class PolyGraphScore(GenerationMetric):
         variant: Classifier metric to compute. To estimate the Jensen-Shannon distance, use "jsd". To estimate total variation distance, use "informedness".
         classifier: Binary classifier to fit
     """
+    _variant: Literal["informedness", "jsd"]
+    _classifier: Literal["logistic", "tabpfn"]
+
     def __init__(
         self,
         reference_graphs: Collection[nx.Graph],
@@ -463,10 +448,13 @@ def compute(self, generated_graphs: Collection[nx.Graph]) -> PolyGraphScoreResul
                 name: metric[1] for name, metric in all_metrics.items()
             },
         }
-        return result
+        return PolyGraphScoreResult(**result)
 
 
 class PolyGraphScoreInterval(GenerationMetricInterval):
+    _variant: Literal["informedness", "jsd"]
+    _classifier: Literal["logistic", "tabpfn"]
+
     def __init__(
         self,
         reference_graphs: Collection[nx.Graph],
@@ -522,4 +510,4 @@ def compute(
                 for key in self._sub_metrics.keys()
             },
         }
-        return result
+        return PolyGraphScoreIntervalResult(**result)
diff --git a/polygraph/metrics/base/vun.py b/polygraph/metrics/base/vun.py
@@ -10,7 +10,7 @@
 """
 
 from collections import defaultdict, namedtuple
-from typing import Any, Callable, DefaultDict, Dict, Iterable, List, Optional
+from typing import Any, Callable, Collection, DefaultDict, Dict, Iterable, List, Optional, Union
 
 import networkx as nx
 from scipy.stats import binomtest
@@ -133,10 +133,10 @@ def __init__(
 
     def compute(
         self,
-        generated_graphs: Iterable[nx.Graph],
+        generated_graphs: Collection[nx.Graph],
         uncertainty: bool = True,
         confidence_level: float = 0.95,
-    ) -> Dict[str, BinomConfidenceInterval]:
+    ) -> Union[Dict[str, BinomConfidenceInterval], Dict[str, float]]:
         """Computes VUN metrics for a collection of generated graphs.
 
         Args:
diff --git a/polygraph/two_sample_tests/permutation_mmd.py b/polygraph/two_sample_tests/permutation_mmd.py
@@ -14,6 +14,8 @@
 
 
 class _BootStrapTestBase:
+    _variant: Literal["biased", "umve", "ustat"]
+
     def __init__(
         self,
         reference_graphs: Collection[nx.Graph],
@@ -117,6 +119,7 @@ def compute(
         realized_mmd, mmd_samples = self._get_realized_and_samples(
             generated_graphs, num_samples
         )
+        assert isinstance(realized_mmd, np.ndarray) and isinstance(mmd_samples, np.ndarray)
         assert realized_mmd.ndim == 1 and mmd_samples.ndim == 2
         realized_mmd = np.max(realized_mmd)
         mmd_samples = np.max(mmd_samples, axis=1)
diff --git a/polygraph/utils/graph_descriptors.py b/polygraph/utils/graph_descriptors.py
diff --git a/polygraph/utils/kernels.py b/polygraph/utils/kernels.py
