ADD: PINA-Topo Topological Diagnostics for FNOs

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/backends/__init__.py

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

pina/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/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/topology/monitor.py

@@ -0,0 +1,168 @@
+"""
+TopologyMonitor: PyTorch Lightning callback for topological health checking.
+Uses the TopologicalProfiler with lazy-loaded GUDHI backend.
+"""
+import torch
+import logging
+import warnings
+from pytorch_lightning import Callback, Trainer, LightningModule
+from typing import Optional, Union, List
+from .profiler import TopologicalProfiler
+from .backends import TopologyResult
+
+# Setup logger for this module
+logger = logging.getLogger(__name__)
+
+class TopologyMonitor(Callback):
+    """
+    Monitors topological health of model outputs during validation.
+
+    Args:
+        expected_beta_0: Expected number of connected components.
+        expected_beta_1: Expected number of holes.
+        monitor_freq: Run topology check every N validation epochs.
+        threshold_beta_0: Max allowed beta_0 before triggering alert.
+        threshold_beta_1: Max allowed beta_1 before triggering alert.
+        warmup_epochs: Number of epochs to wait before monitoring starts.
+        mode: 'warn' (log warning), 'stop' (graceful stop), 'log' (just log).
+        min_persistence: Minimum lifespan for features to be counted.
+        channel: Channel index to monitor (for multi-variate outputs).
+        collect_predictions: Function to extract predictions from the model.
+            If None, the user must override `on_validation_epoch_end`.
+    """
+    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
+
+        # Lazy-load the profiler (GUDHI error deferred until first use)
+        self.profiler = TopologicalProfiler(
+            min_persistence=min_persistence,
+            channel=channel
+        )
+
+    def on_validation_epoch_end(self, trainer: Trainer, pl_module: LightningModule):
+        """Runs topology check at the end of each validation epoch."""
+        # FIX 3: Use Lightning's state directly. No manual epoch counter.
+        if trainer.current_epoch < self.warmup_epochs:
+            return
+
+        # FIX 3: Sync frequency with actual training schedule
+        if (trainer.current_epoch - self.warmup_epochs) % self.monitor_freq != 0:
+            return
+
+        # Get predictions
+        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
+
+        # Compute topology
+        result = self.profiler.compute(predictions)
+        if not result.success:
+            warnings.warn(f"Topology computation failed: {result.error_msg}")
+            return
+
+        # FIX 2: Use Lightning's agnostic logging interface (pl_module.log)
+        self._log_results(pl_module, result)
+
+        # FIX 1: Check both beta_0 and beta_1 thresholds independently
+        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)
+
+            if self.mode == "stop":
+                # FIX 4: Graceful shutdown: log error, set flag, do NOT raise exception
+                logger.error(alert)
+                trainer.should_stop = True
+                # Let Lightning handle the shutdown gracefully
+            elif self.mode == "warn":
+                warnings.warn(alert)
+            # else 'log' mode: just log (already logged via `_log_results`)
+
+    def _get_predictions(self, trainer: Trainer, pl_module: LightningModule) -> Optional[torch.Tensor]:
+        """
+        Extracts predictions from the model.
+        If `collect_predictions` is provided, use it.
+        Otherwise, attempt to get the last batch from the validation dataloader.
+        """
+        if self.collect_predictions is not None:
+            return self.collect_predictions(trainer, pl_module)
+
+        # Fallback: warn user to implement collection
+        return None
+
+    def _log_results(self, pl_module: LightningModule, result: TopologyResult):
+        """
+        Logs topology metrics using PyTorch Lightning's agnostic interface.
+        FIX 2: Works with TensorBoard, WandB, CSV, Comet, etc.
+        """
+        # Log mean, std, max for beta_0
+        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)
+
+        # Log beta_1 only if available
+        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) -> str:
+        """
+        Generates an actionable alert message.
+        FIX 5: Diagnostics are based on the thresholds, not hardcoded offsets.
+        """
+        msg = (
+            f"\n{'='*60}\n"
+            f"[TOPOLOGY ALERT] Epoch: {self._current_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 += "\nDiagnostics:\n"
+        # FIX 5: Base diagnostic messages on thresholds, not expected+1
+        if result.beta_0_mean > self.threshold_beta_0:
+            msg += "  - Excessive connected components detected. Ghost islands may be present.\n"
+        if result.beta_1_mean is not None and result.beta_1_mean > self.threshold_beta_1:
+            msg += "  - Spurious holes detected. Check boundary conditions.\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