Add strict typing and linting

pyproject.toml

@@ -55,8 +55,10 @@ module = ["torch", "torch.*", "tensorflow", "tensorflow.*", "optuna", "matplotli
 ignore_missing_imports = true
 
 [tool.ruff]
-select = ["E", "F", "I"]
 src = ["src", "tests"]
 target-version = "py38"
 line-length = 120
 exclude = ["examples", "notebooks", "scripts", "experiments.py", "results", "data"]
+
+[tool.ruff.lint]
+select = ["E", "F", "I", "B", "UP", "ANN"]

src/psd/algorithms.py

@@ -11,7 +11,7 @@
 from __future__ import annotations
 
 import warnings
-from typing import Any, Callable, Optional, Tuple
+from collections.abc import Callable
 
 import numpy as np
 
@@ -25,7 +25,7 @@ def gradient_descent(
     step_size: float = 0.1,
     tol: float = 1e-6,
     max_iter: int = 10000,
-) -> Tuple[np.ndarray, int]:
+) -> tuple[np.ndarray, int]:
     """Perform basic gradient descent.
 
     Parameters
@@ -67,9 +67,9 @@ def psd(
     delta: float = 0.1,
     delta_f: float = 1.0,
     max_iter: int = 100000,
-    random_state: Optional[np.random.Generator] = None,
-    config: Optional[PSDConfig] = None,
-) -> Tuple[np.ndarray, int]:
+    random_state: np.random.Generator | None = None,
+    config: PSDConfig | None = None,
+) -> tuple[np.ndarray, int]:
     """Perturbed Saddle‑escape Descent (PSD).
 
     This implementation follows the algorithm described in the manuscript
@@ -189,8 +189,8 @@ def psgd(
     delta_fp: float = 0.1,
     delta: float = 0.1,
     delta_f: float = 1.0,
-    random_state: Optional[np.random.Generator] = None,
-) -> Tuple[np.ndarray, int]:
+    random_state: np.random.Generator | None = None,
+) -> tuple[np.ndarray, int]:
     """Perturbed stochastic gradient descent.
 
     This variant uses a noise proxy to choose the batch size and avoids
@@ -284,8 +284,8 @@ def psd_probe(
     rho: float,
     delta: float = 0.1,
     delta_f: float = 1.0,
-    random_state: Optional[np.random.Generator] = None,
-) -> Tuple[np.ndarray, int]:
+    random_state: np.random.Generator | None = None,
+) -> tuple[np.ndarray, int]:
     """Finite‑difference variant of PSD (PSD‑Probe).
 
     This implementation uses a probing step to detect negative curvature
@@ -341,7 +341,7 @@ def psd_probe(
         return x, grad_evals
 
 
-def deprecated_psd(*args: Any, **kwargs: Any) -> Tuple[np.ndarray, int]:
+def deprecated_psd(*args: object, **kwargs: object) -> tuple[np.ndarray, int]:
     """Deprecated alias for :func:`psd`.
 
     This function will be removed in a future release.  Use :func:`psd`

src/psd/feature_flags.py

@@ -10,7 +10,7 @@ class FeatureFlags:
     new_escape_condition: bool = False
 
 
-FLAGS = FeatureFlags()
+FLAGS: FeatureFlags = FeatureFlags()
 
 
 def enable(flag: str) -> None:

src/psd/framework_optimizers.py

@@ -48,7 +48,7 @@
 # mypy: ignore-errors
 import logging
 import math
-from typing import Callable, Iterable, List, Optional, Tuple
+from collections.abc import Callable, Iterable
 
 try:  # Optional import for environments without PyTorch
     import torch
@@ -61,7 +61,7 @@
     tf = None  # type: ignore[assignment]
 
 
-logger = logging.getLogger(__name__)
+logger: logging.Logger = logging.getLogger(__name__)
 
 
 class PSDTorch(torch.optim.Optimizer):  # type: ignore[misc]
@@ -98,7 +98,7 @@ def __init__(
         g_thres: float = 1e-3,
         t_thres: int = 10,
         r: float = 1e-3,
-        max_grad_norm: Optional[float] = 1.0,
+        max_grad_norm: float | None = 1.0,
         **kwargs: object,
     ) -> None:
         if torch is None:  # pragma: no cover - ensures clear error if torch missing
@@ -114,7 +114,7 @@ def __init__(
         super().__init__(params, defaults, **kwargs)
 
     @torch.no_grad()  # type: ignore[misc]
-    def step(self, closure: Optional[Callable[[], float]] = None) -> Optional[float]:
+    def step(self, closure: Callable[[], float] | None = None) -> float | None:
         """Perform a single optimization step.
 
         Parameters
@@ -130,7 +130,7 @@ def step(self, closure: Optional[Callable[[], float]] = None) -> Optional[float]
                 loss = closure()
 
         for group in self.param_groups:
-            params: List[torch.nn.Parameter] = [p for p in group["params"] if p.grad is not None]
+            params: list[torch.nn.Parameter] = [p for p in group["params"] if p.grad is not None]
             if not params:
                 continue
 
@@ -217,7 +217,7 @@ def __init__(
         t_thres: int = 10,
         r: float = 1e-3,
         name: str = "PSDTensorFlow",
-        **kwargs,
+        **kwargs: object,
     ) -> None:
         if tf is None:  # pragma: no cover - ensures clear error if TF missing
             raise ImportError("TensorFlow is required to use PSDTensorFlow")
@@ -227,7 +227,7 @@ def __init__(
         self._set_hyper("t_thres", float(t_thres))
         self._set_hyper("r", r)
 
-    def _create_slots(self, var_list: List[tf.Variable]) -> None:  # pragma: no cover - TF specific
+    def _create_slots(self, var_list: list[tf.Variable]) -> None:  # pragma: no cover - TF specific
         for var in var_list:
             # ``t_noise`` stores the iteration of the last perturbation.  It is
             # initialised to a very negative value so that a perturbation is
@@ -237,10 +237,10 @@ def _create_slots(self, var_list: List[tf.Variable]) -> None:  # pragma: no cove
     @tf.function
     def apply_gradients(
         self,
-        grads_and_vars: Iterable[Tuple[tf.Tensor, tf.Variable]],
-        name: Optional[str] = None,
-        **kwargs,
-    ):  # pragma: no cover - TF specific
+        grads_and_vars: Iterable[tuple[tf.Tensor, tf.Variable]],
+        name: str | None = None,
+        **kwargs: object,
+    ) -> None:  # pragma: no cover - TF specific
         grads_and_vars = [(g, v) for g, v in grads_and_vars if g is not None]
         if not grads_and_vars:
             return

src/psd/functions.py

@@ -9,8 +9,8 @@
 
 from __future__ import annotations
 
+from collections.abc import Callable
 from dataclasses import dataclass
-from typing import Callable, Dict
 
 import numpy as np
 
@@ -322,7 +322,7 @@ def random_quadratic_hess(A: np.ndarray) -> np.ndarray:
     hess=rosenbrock_hess,
 )
 
-TEST_FUNCTIONS: Dict[str, TestFunction] = {
+TEST_FUNCTIONS: dict[str, TestFunction] = {
     "separable_quartic": SEPARABLE_QUARTIC,
     "coupled_quartic": COUPLED_QUARTIC,
     "rosenbrock": ROSENBROCK,

src/psd/graph.py

@@ -11,21 +11,21 @@
 
 import logging
 from collections import deque
+from collections.abc import Hashable
 from math import isfinite
 from time import perf_counter
-from typing import Any, Dict, List, Tuple
 
-logger = logging.getLogger(__name__)
+logger: logging.Logger = logging.getLogger(__name__)
 
 
-Graph = Dict[Any, Dict[Any, float]]
+Graph = dict[Hashable, dict[Hashable, float]]
 """Type alias for an adjacency list representing a weighted directed graph."""
 
 MAX_PATH_WEIGHT = 1e12
 """Maximum allowable weight for any path to guard against overflow."""
 
 
-def _validate_graph(graph: Graph, start: Any, end: Any) -> None:
+def _validate_graph(graph: Graph, start: Hashable, end: Hashable) -> None:
     """Validate graph structure and ensure non-negative, finite edge weights.
 
     Parameters
@@ -49,41 +49,41 @@ def _validate_graph(graph: Graph, start: Any, end: Any) -> None:
     if start not in graph or end not in graph:
         raise ValueError("Start or end node not present in graph.")
 
-    for node, neighbours in graph.items():
+    for _node, neighbours in graph.items():
         if not isinstance(neighbours, dict):
             raise ValueError("Graph adjacency lists must be dictionaries.")
-        for neighbour, weight in neighbours.items():
+        for _neighbour, weight in neighbours.items():
             if weight < 0:
                 raise ValueError("Graph contains negative edge weights.")
             if not isfinite(weight) or weight > MAX_PATH_WEIGHT:
                 raise OverflowError("Edge weight exceeds safe maximum.")
 
 
-def _initialize_state(graph: Graph, start: Any) -> Tuple[Dict[Any, float], Dict[Any, Any]]:
+def _initialize_state(graph: Graph, start: Hashable) -> tuple[dict[Hashable, float], dict[Hashable, Hashable]]:
     """Initialise distance estimates and predecessor map."""
 
-    distances: Dict[Any, float] = {node: float("inf") for node in graph}
-    previous: Dict[Any, Any] = {}
+    distances: dict[Hashable, float] = {node: float("inf") for node in graph}
+    previous: dict[Hashable, Hashable] = {}
     distances[start] = 0.0
     return distances, previous
 
 
-def _topological_sort(graph: Graph) -> List[Any]:
+def _topological_sort(graph: Graph) -> list[Hashable]:
     """Return a topological ordering of ``graph`` or raise ``ValueError``.
 
     The function performs Kahn's algorithm while ensuring that all nodes are
     included in the ordering.  A ``ValueError`` is raised if the graph contains
     a cycle which would prevent such an ordering.
     """
 
-    indegree: Dict[Any, int] = {node: 0 for node in graph}
-    for node, neighbours in graph.items():
+    indegree: dict[Hashable, int] = {node: 0 for node in graph}
+    for _node, neighbours in graph.items():
         for neighbour in neighbours:
             indegree.setdefault(neighbour, 0)
             indegree[neighbour] += 1
 
-    queue: deque[Any] = deque([n for n, d in indegree.items() if d == 0])
-    order: List[Any] = []
+    queue: deque[Hashable] = deque([n for n, d in indegree.items() if d == 0])
+    order: list[Hashable] = []
     while queue:
         node = queue.popleft()
         order.append(node)
@@ -98,7 +98,7 @@ def _topological_sort(graph: Graph) -> List[Any]:
     return order
 
 
-def _reconstruct_path(previous: Dict[Any, Any], start: Any, end: Any) -> List[Any]:
+def _reconstruct_path(previous: dict[Hashable, Hashable], start: Hashable, end: Hashable) -> list[Hashable]:
     """Rebuild the path from ``start`` to ``end`` using ``previous`` map.
 
     Prior to this change the function blindly followed the predecessor map.
@@ -109,8 +109,8 @@ def _reconstruct_path(previous: Dict[Any, Any], start: Any, end: Any) -> List[An
     encountered.
     """
 
-    path: List[Any] = [end]
-    visited = set()
+    path: list[Hashable] = [end]
+    visited: set[Hashable] = set()
     while path[-1] != start:
         node = path[-1]
         if node in visited:
@@ -125,7 +125,7 @@ def _reconstruct_path(previous: Dict[Any, Any], start: Any, end: Any) -> List[An
     return path
 
 
-def find_optimal_path(graph: Graph, start: Any, end: Any) -> List[Any]:
+def find_optimal_path(graph: Graph, start: Hashable, end: Hashable) -> list[Hashable]:
     """Find the shortest path from ``start`` to ``end`` in a DAG.
 
     The graph is represented as an adjacency list mapping each node to a

src/psd/optimizer_tf.py

@@ -10,7 +10,7 @@
 from __future__ import annotations
 
 # mypy: ignore-errors
-from typing import Dict, Iterable, Optional, Tuple
+from collections.abc import Iterable
 
 import tensorflow as tf
 
@@ -59,7 +59,7 @@ def _resource_apply_dense(
         self,
         grad: tf.Tensor,
         var: tf.Variable,
-        apply_state: Optional[Dict[str, object]] = None,
+        apply_state: dict[str, object] | None = None,
     ) -> None:
         lr_t = tf.cast(self._decayed_lr(var.dtype), var.dtype)
         var.assign_sub(grad * lr_t)
@@ -69,15 +69,15 @@ def _resource_apply_sparse(
         grad: tf.Tensor,
         var: tf.Variable,
         indices: tf.Tensor,
-        apply_state: Optional[Dict[str, object]] = None,
+        apply_state: dict[str, object] | None = None,
     ) -> None:
         lr_t = tf.cast(self._decayed_lr(var.dtype), var.dtype)
         var.scatter_sub(tf.IndexedSlices(grad * lr_t, indices))
 
     def apply_gradients(
         self,
-        grads_and_vars: Iterable[Tuple[tf.Tensor, tf.Variable]],
-        name: Optional[str] = None,
+        grads_and_vars: Iterable[tuple[tf.Tensor, tf.Variable]],
+        name: str | None = None,
         **kwargs: object,
     ) -> None:
         grads_and_vars = [(g, v) for g, v in grads_and_vars if g is not None]
@@ -103,7 +103,7 @@ def apply_gradients(
         step = tf.identity(self.iterations)
         need_perturb = tf.logical_and(global_norm <= g_thres, (step - self.t_noise) > t_thres)
 
-        def perturb_vars():
+        def perturb_vars() -> None:
             for v in vars:
                 v_dtype = v.dtype.base_dtype
                 noise = tf.random.normal(tf.shape(v), dtype=v_dtype)
@@ -118,8 +118,8 @@ def perturb_vars():
         super().apply_gradients(grads_and_vars, name, **kwargs)
         return None
 
-    def get_config(self) -> Dict[str, object]:
-        config: Dict[str, object] = super().get_config()
+    def get_config(self) -> dict[str, object]:
+        config: dict[str, object] = super().get_config()
         config.update(
             {
                 "learning_rate": self._serialize_hyperparameter("learning_rate"),

src/psd_optimizer/optimizer.py

@@ -1,7 +1,7 @@
 from __future__ import annotations
 
 import logging
-from typing import Callable, Iterable, Optional
+from collections.abc import Callable, Iterable
 
 import torch
 from torch import Tensor
@@ -31,15 +31,15 @@ def __init__(
         epsilon: float = 1e-3,
         r: float = 1e-3,
         T: int = 10,
-        max_grad_norm: Optional[float] = 1.0,
+        max_grad_norm: float | None = 1.0,
         **kwargs: object,
     ) -> None:
         if lr <= 0:
             raise ValueError("Invalid learning rate")
         defaults = dict(lr=lr, epsilon=epsilon, r=r, T=T, max_grad_norm=max_grad_norm)
         super().__init__(params, defaults, **kwargs)
 
-    def step(self, closure: Optional[Callable[[], Tensor]] = None) -> Tensor:
+    def step(self, closure: Callable[[], Tensor] | None = None) -> Tensor:
         """Perform a single optimization step.
 
         Parameters