ADD: PINA-Topo Topological Diagnostics for FNOs

pina/_src/topology/__init__.py

@@ -0,0 +1,4 @@
+from .profiler import TopologicalProfiler
+from .monitor import TopologyMonitor
+
+__all__ = ["TopologicalProfiler", "TopologyMonitor"]

pina/_src/topology/backends/__init__.py

@@ -0,0 +1,4 @@
+from .base import TopologyBackend, TopologyResult
+from .gudhi_backend import GudhiBackend
+
+__all__ = ["TopologyBackend", "TopologyResult", "GudhiBackend"]

pina/_src/topology/backends/base.py

@@ -0,0 +1,29 @@
+from dataclasses import dataclass, field
+from abc import ABC, abstractmethod
+from typing import Dict, Any, Optional, List
+import torch
+
+@dataclass
+class TopologyResult:
+    beta_0_mean: float
+    beta_1_mean: Optional[float] = None  # Can be None if backend doesn't support β₁
+    beta_0_std: float = 0.0
+    beta_1_std: Optional[float] = None
+    beta_0_max: float = 0.0
+    beta_1_max: Optional[float] = None
+    per_sample_beta_0: List[int] = field(default_factory=list)
+    per_sample_beta_1: List[Optional[int]] = field(default_factory=list)
+    success: bool = True
+    error_msg: Optional[str] = None
+    backend_name: str = "unknown"
+    metadata: Dict[str, Any] = field(default_factory=dict)
+
+class TopologyBackend(ABC):
+    @abstractmethod
+    def compute(self, tensor: torch.Tensor, **kwargs) -> TopologyResult:
+        pass
+
+    @property
+    @abstractmethod
+    def name(self) -> str:
+        pass

pina/_src/topology/backends/gudhi_backend.py

@@ -0,0 +1,113 @@
+"""
+GUDHI-based backend using CubicalComplex.
+Implements defensive programming for production use.
+"""
+import numpy as np
+import torch
+import warnings
+from typing import Optional
+from .base import TopologyBackend, TopologyResult
+
+class GudhiBackend(TopologyBackend):
+    def __init__(self, min_persistence: float = 0.1):
+        self.min_persistence = min_persistence
+        self._gudhi_available = self._check_gudhi()
+
+    def _check_gudhi(self) -> bool:
+        try:
+            import gudhi
+            return True
+        except ImportError:
+            return False
+
+    @property
+    def name(self) -> str:
+        return "gudhi"
+
+    def _validate_and_prepare(self, tensor: torch.Tensor, channel: Optional[int] = None) -> np.ndarray:
+        """Fixes: GPU memory, dimensionality, channel selection."""
+        tensor = tensor.detach().cpu()
+
+        if tensor.ndim == 4:
+            if channel is None:
+                tensor = tensor[:, 0, :, :]
+            else:
+                if channel >= tensor.shape[1]:
+                    raise ValueError(f"Channel {channel} requested, but tensor has only {tensor.shape[1]} channels.")
+                tensor = tensor[:, channel, :, :]
+        elif tensor.ndim == 3:
+            pass
+        else:
+            raise ValueError(f"Unsupported tensor shape: {tensor.shape}. Expected 3D or 4D.")
+
+        if tensor.ndim == 2:
+            tensor = tensor.unsqueeze(0)
+
+        return tensor.numpy()
+
+    def _compute_single_betti(self, grid_2d: np.ndarray, min_persistence: float) -> tuple:
+        import gudhi
+        cc = gudhi.CubicalComplex(top_dimensional_cells=grid_2d.astype(np.float64))
+        cc.persistence()
+        intervals_0 = cc.persistence_intervals_in_dimension(0)
+        intervals_1 = cc.persistence_intervals_in_dimension(1)
+        beta_0 = sum(1 for (b, d) in intervals_0 if (d - b) > min_persistence)
+        beta_1 = sum(1 for (b, d) in intervals_1 if (d - b) > min_persistence)
+        return beta_0, beta_1
+
+    def compute(self, tensor: torch.Tensor, **kwargs) -> TopologyResult:
+        if not self._gudhi_available:
+            return TopologyResult(
+                success=False,
+                error_msg="GUDHI is not installed.",
+                backend_name=self.name
+            )
+
+        try:
+            channel = kwargs.get("channel", None)
+            negate = kwargs.get("negate", False)
+            min_pers = kwargs.get("min_persistence", self.min_persistence)
+
+            data = self._validate_and_prepare(tensor, channel=channel)
+            if negate:
+                data = -data  # convert to superlevel sets
+
+            batch_size = data.shape[0]
+
+            beta_0_list = []
+            beta_1_list = []
+
+            for i in range(batch_size):
+                b0, b1 = self._compute_single_betti(data[i], min_pers)
+                beta_0_list.append(b0)
+                beta_1_list.append(b1)
+
+            beta_0_arr = np.array(beta_0_list)
+            beta_1_arr = np.array(beta_1_list)
+
+            return TopologyResult(
+                beta_0_mean=float(beta_0_arr.mean()),
+                beta_1_mean=float(beta_1_arr.mean()),
+                beta_0_std=float(beta_0_arr.std()),
+                beta_1_std=float(beta_1_arr.std()),
+                beta_0_max=int(beta_0_arr.max()),
+                beta_1_max=int(beta_1_arr.max()),
+                per_sample_beta_0=beta_0_list,
+                per_sample_beta_1=beta_1_list,
+                success=True,
+                backend_name=self.name,
+                metadata={
+                    "batch_size": batch_size,
+                    "min_persistence": min_pers,
+                    "channel": channel,
+                    "negate": negate,
+                    "tensor_shape": tensor.shape
+                }
+            )
+        except Exception as e:
+            return TopologyResult(
+                success=False,
+                error_msg=f"GUDHI computation failed: {str(e)}",
+                backend_name=self.name,
+                metadata={"tensor_shape": tensor.shape}
+            )

pina/_src/topology/monitor.py

@@ -0,0 +1,115 @@
+"""
+TopologyMonitor: PyTorch Lightning callback for topological health checking.
+"""
+import torch
+import logging
+import warnings
+from lightning.pytorch import Callback, Trainer, LightningModule
+from typing import Optional
+from pina._src.topology.profiler import TopologicalProfiler
+from pina._src.topology.backends import TopologyResult
+
+logger = logging.getLogger(__name__)
+
+class TopologyMonitor(Callback):
+    def __init__(
+        self,
+        expected_beta_0: int = 1,
+        expected_beta_1: int = 0,
+        monitor_freq: int = 10,
+        threshold_beta_0: Optional[float] = None,
+        threshold_beta_1: Optional[float] = None,
+        warmup_epochs: int = 10,
+        mode: str = "warn",
+        min_persistence: float = 0.1,
+        channel: Optional[int] = None,
+        collect_predictions: Optional[callable] = None,
+    ):
+        super().__init__()
+        self.expected_beta_0 = expected_beta_0
+        self.expected_beta_1 = expected_beta_1
+        self.monitor_freq = monitor_freq
+        self.threshold_beta_0 = threshold_beta_0 if threshold_beta_0 is not None else expected_beta_0 + 0.5
+        self.threshold_beta_1 = threshold_beta_1 if threshold_beta_1 is not None else expected_beta_1 + 0.5
+        self.warmup_epochs = warmup_epochs
+        self.mode = mode
+        self.collect_predictions = collect_predictions
+
+        self.profiler = TopologicalProfiler(
+            min_persistence=min_persistence,
+            channel=channel
+        )
+
+    def on_validation_epoch_end(self, trainer: Trainer, pl_module: LightningModule):
+        if trainer.current_epoch < self.warmup_epochs:
+            return
+
+        if (trainer.current_epoch - self.warmup_epochs) % self.monitor_freq != 0:
+            return
+
+        predictions = self._get_predictions(trainer, pl_module)
+        if predictions is None:
+            warnings.warn(
+                "No predictions available for TopologyMonitor. "
+                "Override `_get_predictions` or provide `collect_predictions`."
+            )
+            return
+
+        result = self.profiler.compute(predictions)
+        if not result.success:
+            warnings.warn(f"Topology computation failed: {result.error_msg}")
+            return
+
+        self._log_results(pl_module, result)
+
+        trigger_beta_0 = result.beta_0_mean > self.threshold_beta_0
+        trigger_beta_1 = (result.beta_1_mean is not None and
+                          result.beta_1_mean > self.threshold_beta_1)
+
+        if trigger_beta_0 or trigger_beta_1:
+            alert = self._generate_alert(result, trainer.current_epoch)
+
+            if self.mode == "stop":
+                logger.error(alert)
+                trainer.should_stop = True
+            elif self.mode == "warn":
+                warnings.warn(alert)
+
+    def _get_predictions(self, trainer: Trainer, pl_module: LightningModule) -> Optional[torch.Tensor]:
+        if self.collect_predictions is not None:
+            return self.collect_predictions(trainer, pl_module)
+        return None
+
+    def _log_results(self, pl_module: LightningModule, result: TopologyResult):
+        pl_module.log("topology/beta_0_mean", result.beta_0_mean, sync_dist=True)
+        pl_module.log("topology/beta_0_std", result.beta_0_std, sync_dist=True)
+        pl_module.log("topology/beta_0_max", float(result.beta_0_max), sync_dist=True)
+
+        if result.beta_1_mean is not None:
+            pl_module.log("topology/beta_1_mean", result.beta_1_mean, sync_dist=True)
+            pl_module.log("topology/beta_1_std", result.beta_1_std, sync_dist=True)
+            if result.beta_1_max is not None:
+                pl_module.log("topology/beta_1_max", float(result.beta_1_max), sync_dist=True)
+
+    def _generate_alert(self, result: TopologyResult, epoch: int) -> str:
+        msg = (
+            f"\n{'='*60}\n"
+            f"[TOPOLOGY ALERT] Epoch: {epoch}\n"
+        )
+
+        if result.beta_0_mean > self.threshold_beta_0:
+            msg += f"β₀ = {result.beta_0_mean:.2f} (threshold: {self.threshold_beta_0})\n"
+        if result.beta_1_mean is not None and result.beta_1_mean > self.threshold_beta_1:
+            msg += f"β₁ = {result.beta_1_mean:.2f} (threshold: {self.threshold_beta_1})\n"
+
+        msg += "\nSuggested actions:\n"
+        suggestions = []
+        if result.beta_0_mean > self.threshold_beta_0:
+            suggestions.append("Increase `n_modes` to capture high-frequency edges.")
+        if result.beta_1_mean is not None and result.beta_1_mean > self.threshold_beta_1:
+            suggestions.append("Check boundary conditions; spurious loops may arise from BC mismatches.")
+        if not suggestions:
+            suggestions.append("Reduce learning_rate to stabilize spectral weight convergence.")
+        msg += "  - " + "\n  - ".join(suggestions)
+        msg += f"\n{'='*60}"
+        return msg

pina/_src/topology/profiler.py

@@ -0,0 +1,46 @@
+"""
+Topological Profiler: Lazy-loads GUDHI backend with state consistency.
+"""
+import torch
+from typing import Optional
+from pina._src.topology.backends import TopologyResult
+
+class TopologicalProfiler:
+    def __init__(
+        self,
+        min_persistence: float = 0.1,
+        channel: Optional[int] = None,
+    ):
+        self.min_persistence = min_persistence
+        self.channel = channel
+        self._backend = None
+
+    def _get_backend(self):
+        if self._backend is None:
+            try:
+                from pina._src.topology.backends import GudhiBackend
+                self._backend = GudhiBackend(min_persistence=self.min_persistence)
+            except ImportError as e:
+                raise ImportError(
+                    "\n" + "=" * 60 + "\n"
+                    "GUDHI BACKEND REQUIRED FOR TOPOLOGICAL PROFILING.\n"
+                    "Please install GUDHI:\n"
+                    "    pip install gudhi\n"
+                    "=" * 60
+                ) from e
+        return self._backend
+
+    def compute(self, tensor: torch.Tensor, **kwargs) -> TopologyResult:
+        channel = kwargs.pop("channel", self.channel)
+        min_pers = kwargs.pop("min_persistence", self.min_persistence)
+
+        backend = self._get_backend()
+
+        if backend.min_persistence != min_pers:
+            backend.min_persistence = min_pers
+
+        return backend.compute(tensor, channel=channel, min_persistence=min_pers, **kwargs)
+
+    @property
+    def backend_name(self) -> str:
+        return self._get_backend().name

pina/callbacks/__init__.py

@@ -0,0 +1,7 @@
+"""
+PINA Callbacks module with lazy loading for topology monitor.
+"""
+def TopologyMonitor(*args, **kwargs):
+    """Lazy loader for TopologyMonitor to avoid import overhead."""
+    from pina.topology.monitor import TopologyMonitor as _TopologyMonitor
+    return _TopologyMonitor(*args, **kwargs)

pina/topology/__init__.py

@@ -0,0 +1,5 @@
+from pina._src.topology.monitor import TopologyMonitor
+from pina._src.topology.profiler import TopologicalProfiler
+from pina._src.topology.backends import GudhiBackend
+
+__all__ = ["TopologyMonitor", "TopologicalProfiler", "GudhiBackend"]