diff --git a/pina/_src/topology/__init__.py b/pina/_src/topology/__init__.py new file mode 100644 index 000000000..a142050d3 --- /dev/null +++ b/pina/_src/topology/__init__.py @@ -0,0 +1,4 @@ +from .profiler import TopologicalProfiler +from .monitor import TopologyMonitor + +__all__ = ["TopologicalProfiler", "TopologyMonitor"] diff --git a/pina/_src/topology/backends/__init__.py b/pina/_src/topology/backends/__init__.py new file mode 100644 index 000000000..c25b42974 --- /dev/null +++ b/pina/_src/topology/backends/__init__.py @@ -0,0 +1,4 @@ +from .base import TopologyBackend, TopologyResult +from .gudhi_backend import GudhiBackend + +__all__ = ["TopologyBackend", "TopologyResult", "GudhiBackend"] \ No newline at end of file diff --git a/pina/_src/topology/backends/base.py b/pina/_src/topology/backends/base.py new file mode 100644 index 000000000..c4da4692a --- /dev/null +++ b/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 \ No newline at end of file diff --git a/pina/_src/topology/backends/gudhi_backend.py b/pina/_src/topology/backends/gudhi_backend.py new file mode 100644 index 000000000..467e27144 --- /dev/null +++ b/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} + ) \ No newline at end of file diff --git a/pina/_src/topology/monitor.py b/pina/_src/topology/monitor.py new file mode 100644 index 000000000..a5c037b28 --- /dev/null +++ b/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 diff --git a/pina/_src/topology/profiler.py b/pina/_src/topology/profiler.py new file mode 100644 index 000000000..b27752c9d --- /dev/null +++ b/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 diff --git a/pina/callbacks/__init__.py b/pina/callbacks/__init__.py new file mode 100644 index 000000000..cf8f98064 --- /dev/null +++ b/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) \ No newline at end of file diff --git a/pina/topology/__init__.py b/pina/topology/__init__.py new file mode 100644 index 000000000..8e0ff73f2 --- /dev/null +++ b/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"] diff --git a/test_monitor.py b/test_monitor.py new file mode 100644 index 000000000..d9c7a43e5 --- /dev/null +++ b/test_monitor.py @@ -0,0 +1,156 @@ +""" +Minimal test for TopologyMonitor. +""" +import torch +import numpy as np +from scipy import io +from torch.utils.data import TensorDataset, DataLoader +from pina.model import FNO, FeedForward +from pina.solver import SupervisedSingleModelSolver +from pina.problem.zoo import SupervisedProblem +from pina import Trainer, LabelTensor +from pina.topology import TopologyMonitor, TopologicalProfiler, GudhiBackend + +print("=" * 60) +print("Testing TopologyMonitor with Darcy dataset") +print("=" * 60) + +# 1. Load data (Channels-Last for FNO) +print("\n[1] Loading Darcy data...") +data = io.loadmat("Data_Darcy.mat") +k_train = torch.tensor(data["k_train"], dtype=torch.float) +u_train = torch.tensor(data["u_train"], dtype=torch.float) + +k_train = k_train.unsqueeze(-1) +u_train = u_train.unsqueeze(-1) + +print(f" Total samples: {k_train.shape[0]}") +print(f" Grid size: {k_train.shape[1]}x{k_train.shape[2]}") +print(f" Input shape: {k_train.shape}") + +# 2. Split data +n_samples = 100 +train_split = int(0.8 * n_samples) +k_train_subset = k_train[:n_samples] +u_train_subset = u_train[:n_samples] + +train_input = k_train_subset[:train_split] +train_output = u_train_subset[:train_split] +val_input = k_train_subset[train_split:] +val_output = u_train_subset[train_split:] + +print(f" Training samples: {train_split}") +print(f" Validation samples: {n_samples - train_split}") + +# 3. Create problem with explicit variable names +problem = SupervisedProblem(input_=train_input, output_=train_output) +# CRITICAL: Set the variable names for LabelTensor extraction +problem.input_variables = ['x'] +problem.output_variables = ['u'] + +# 4. Create FNO model +lifting_net = FeedForward(input_dimensions=1, output_dimensions=20, layers=[30, 30]) +projecting_net = FeedForward(input_dimensions=20, output_dimensions=1, layers=[30, 30]) + +fno = FNO( + lifting_net=lifting_net, + projecting_net=projecting_net, + n_modes=8, + dimensions=2, + n_layers=2, + padding=4, +) + +# 5. Architecture Sanity Check +print("\n[2] Running architecture sanity check...") +with torch.no_grad(): + test_input = k_train_subset[:1] + test_out = fno(test_input) + print(f" FNO Output Shape: {test_out.shape}") + assert test_out.shape == (1, k_train.shape[1], k_train.shape[2], 1), \ + f"Shape mismatch: expected (1, {k_train.shape[1]}, {k_train.shape[2]}, 1), got {test_out.shape}" + print(" ✅ Shape check passed!") + +# 6. Create solver +solver = SupervisedSingleModelSolver( + problem=problem, + model=fno +) + +# 7. collect_predictions hook (Handles LabelTensor dict) +def collect_predictions(trainer, pl_module): + """Extract predictions from validation batch and convert to channels-first.""" + val_loader = trainer.val_dataloaders + if val_loader is None: + return None + + with torch.no_grad(): + for batch in val_loader: + if isinstance(batch, list) and len(batch) > 0: + _, data_dict = batch[0] + input_label = data_dict['input'] + # Convert LabelTensor to raw tensor + inputs = input_label.as_subclass(torch.Tensor).to(pl_module.device) + pred = pl_module.model(inputs) + # Convert Channels-Last to Channels-First + return pred.permute(0, 3, 1, 2) + else: + try: + input_batch = batch[0].to(pl_module.device) + pred = pl_module.model(input_batch) + return pred.permute(0, 3, 1, 2) + except: + continue + return None + +print("\n[3] Creating TopologyMonitor with collect_predictions hook...") +monitor = TopologyMonitor( + expected_beta_0=1, + expected_beta_1=0, + monitor_freq=2, + warmup_epochs=1, + mode="warn", + min_persistence=0.5, + channel=0, + collect_predictions=collect_predictions, +) + +# 8. Create Trainer +trainer = Trainer( + solver=solver, + max_epochs=5, + callbacks=[monitor], + accelerator="cpu", +) + +# 9. PINA-compliant collate function with LabelTensor and problem variables +print("\n[4] Setting up DataLoader streams...") + +def pina_collate(batch): + """Wraps raw PyTorch tensors into PINA's expected format with LabelTensor.""" + inputs = torch.stack([item[0] for item in batch]) + outputs = torch.stack([item[1] for item in batch]) + condition_name = list(solver.problem.conditions.keys())[0] + # Use the variable names we set on the problem + input_vars = solver.problem.input_variables # ['x'] + output_vars = solver.problem.output_variables # ['u'] + data_dict = { + 'input': LabelTensor(inputs, input_vars), + 'target': LabelTensor(outputs, output_vars) + } + return [(condition_name, data_dict)] + +train_dataset = TensorDataset(train_input, train_output) +train_loader = DataLoader(train_dataset, batch_size=32, shuffle=True, collate_fn=pina_collate) + +val_dataset = TensorDataset(val_input, val_output) +val_loader = DataLoader(val_dataset, batch_size=32, shuffle=False, collate_fn=pina_collate) + +# 10. Train with explicit data streams +print("\n[5] Starting training...") +trainer.fit(solver, train_dataloaders=train_loader, val_dataloaders=val_loader) + +print("\n" + "=" * 60) +print("✅ Test completed successfully!") +print("✅ TopologyMonitor executed with GUDHI!") +print("=" * 60) \ No newline at end of file diff --git a/tests/test_gudhi_backend.py b/tests/test_gudhi_backend.py new file mode 100644 index 000000000..5e0d2ad48 --- /dev/null +++ b/tests/test_gudhi_backend.py @@ -0,0 +1,54 @@ +""" +Smoke test for the MorphologicalBackend (Kornia fallback). +""" +import torch +from pina.topology.backends.morphological_backend import MorphologicalBackend + +def test_morphological_backend(): + print("=" * 50) + print("Testing MorphologicalBackend (Kornia Fallback)") + print("=" * 50) + + data = torch.zeros(3, 32, 32) + val = 1.0 + + data[0, 5:10, 5:10] = val + data[0, 20:25, 20:25] = val + + data[1, 5:25, 5:25] = val + + data[2, 4:28, 4:28] = val + data[2, 10:22, 10:22] = 0.0 + + backend = MorphologicalBackend(low_threshold=0.3, high_threshold=0.7) + result = backend.compute(data) + + print(f"\nBackend: {result.backend_name}") + print(f"Success: {result.success}") + if not result.success: + print(f"Error: {result.error_msg}") + return + + print("\n--- Per-Sample Results ---") + for i in range(3): + print(f"Sample {i}: β₀ = {result.per_sample_beta_0[i]}, β₁ = {result.per_sample_beta_1[i]}") + + print("\n--- Metadata ---") + print(f"Batch Size: {result.metadata.get('batch_size')}") + print(f"low_threshold: {result.metadata.get('low_threshold')}") + print(f"high_threshold: {result.metadata.get('high_threshold')}") + print(f"β₁ Note: {result.metadata.get('note')}") + + expected_beta_0 = [2, 1, 1] + expected_beta_1 = [None, None, None] + + print("\n--- Checking Results ---") + assert result.per_sample_beta_0 == expected_beta_0, \ + f"Beta_0 mismatch: got {result.per_sample_beta_0}, expected {expected_beta_0}" + assert result.per_sample_beta_1 == expected_beta_1, \ + f"Beta_1 mismatch: got {result.per_sample_beta_1}, expected {expected_beta_1}" + + print("\n✅ MorphologicalBackend test passed!") + +if __name__ == "__main__": + test_morphological_backend() \ No newline at end of file diff --git a/tutorials/tutorial_topology.ipynb b/tutorials/tutorial_topology.ipynb new file mode 100644 index 000000000..c0f71f373 --- /dev/null +++ b/tutorials/tutorial_topology.ipynb @@ -0,0 +1,588 @@ +{ + "cells": [ + { + "cell_type": "markdown", + "id": "1702c35d", + "metadata": {}, + "source": [ + "# Monitoring Topological Invariants in Neural Operators using PINA\n", + "\n", + "This tutorial demonstrates how to use PINA's `TopologyMonitor` alongside a **Fourier Neural Operator (FNO)** to track topological properties (Betti numbers) during training on the Darcy Flow dataset.\n", + "\n", + "## Why Topology Matters\n", + "FNOs operate in the frequency domain. By truncating high frequencies (`n_modes`), they often **smooth out sharp boundaries**, creating non-physical \"ghost islands\" (disconnected blobs). \n", + "\n", + "Standard MSE loss ignores these structural errors. The `TopologyMonitor` uses GUDHI's persistent homology to track:\n", + "- **β₀ (Connected Components):** Should be `1` (one single blob).\n", + "- **β₁ (Holes):** Should be `0` (no artificial loops).\n", + "\n", + "If β₀ jumps to `2` or `3`, the model is hallucinating a ghost island." + ] + }, + { + "cell_type": "code", + "execution_count": 7, + "id": "22836a8b", + "metadata": {}, + "outputs": [ + { + "name": "stdout", + "output_type": "stream", + "text": [ + "✅ All imports loaded successfully.\n" + ] + } + ], + "source": [ + "import torch\n", + "import numpy as np\n", + "import matplotlib.pyplot as plt\n", + "from scipy import io\n", + "from torch.utils.data import DataLoader, TensorDataset\n", + "\n", + "from pina import LabelTensor\n", + "from pina.model import FNO, FeedForward\n", + "from pina.solver import SupervisedSingleModelSolver\n", + "from pina.problem.zoo import SupervisedProblem\n", + "from pina import Trainer\n", + "from pina.topology.monitor import TopologyMonitor\n", + "\n", + "print(\"✅ All imports loaded successfully.\")" + ] + }, + { + "cell_type": "code", + "execution_count": 8, + "id": "dc2f1654", + "metadata": {}, + "outputs": [ + { + "name": "stdout", + "output_type": "stream", + "text": [ + "Loading Darcy data...\n", + "Train: 80, Validation: 20\n", + "Grid size: 20x20\n" + ] + } + ], + "source": [ + "print(\"Loading Darcy data...\")\n", + "data = io.loadmat(\"Data_Darcy.mat\")\n", + "\n", + "# FNO expects Channels-Last: (Batch, Height, Width, Channels)\n", + "k_train = torch.tensor(data[\"k_train\"], dtype=torch.float).unsqueeze(-1)\n", + "u_train = torch.tensor(data[\"u_train\"], dtype=torch.float).unsqueeze(-1)\n", + "\n", + "# Use 100 samples for speed\n", + "n_samples = 100\n", + "train_split = 80\n", + "\n", + "k_sub = k_train[:n_samples]\n", + "u_sub = u_train[:n_samples]\n", + "\n", + "train_input = k_sub[:train_split]\n", + "train_output = u_sub[:train_split]\n", + "val_input = k_sub[train_split:]\n", + "val_output = u_sub[train_split:]\n", + "\n", + "print(f\"Train: {len(train_input)}, Validation: {len(val_input)}\")\n", + "print(f\"Grid size: {k_train.shape[1]}x{k_train.shape[2]}\")\n" + ] + }, + { + "cell_type": "code", + "execution_count": 9, + "id": "f8ab8e97", + "metadata": {}, + "outputs": [ + { + "name": "stdout", + "output_type": "stream", + "text": [ + "✅ Problem, model, and solver created.\n" + ] + } + ], + "source": [ + "# Define the supervised problem and explicitly set variable names\n", + "problem = SupervisedProblem(input_=train_input, output_=train_output)\n", + "problem.input_variables = ['x']\n", + "problem.output_variables = ['u']\n", + "\n", + "# Build the FNO architecture\n", + "lifting_net = FeedForward(input_dimensions=1, output_dimensions=20, layers=[30, 30])\n", + "projecting_net = FeedForward(input_dimensions=20, output_dimensions=1, layers=[30, 30])\n", + "\n", + "fno = FNO(\n", + " lifting_net=lifting_net,\n", + " projecting_net=projecting_net,\n", + " n_modes=8,\n", + " dimensions=2,\n", + " n_layers=2,\n", + " padding=4,\n", + ")\n", + "\n", + "# Create the solver\n", + "solver = SupervisedSingleModelSolver(problem=problem, model=fno)\n", + "\n", + "print(\"✅ Problem, model, and solver created.\")" + ] + }, + { + "cell_type": "markdown", + "id": "b22fa3cb", + "metadata": {}, + "source": [ + "## The Data Pipeline Bridge (Engineering Deep-Dive)\n", + "\n", + "**The Problem:** PyTorch `DataLoader`s return raw tuples `(input, output)`. However, PINA's internal `ConditionAggregatorMixin` expects the batch to be a list of `(condition_name, data_dict)` pairs. If we pass raw tuples, it crashes with `ValueError: too many values to unpack`.\n", + "\n", + "**The Solution:** We implement a custom `collate_fn` that dynamically reads the condition name from the solver and wraps the tensors into the exact `{'input': inputs, 'target': outputs}` dictionary that PINA requires.\n", + "\n", + "**Why this is Senior-level:** We read `list(solver.problem.conditions.keys())[0]` dynamically. If PINA changes the internal naming convention in a future update, our code survives." + ] + }, + { + "cell_type": "code", + "execution_count": 10, + "id": "244e999d", + "metadata": {}, + "outputs": [ + { + "name": "stdout", + "output_type": "stream", + "text": [ + "✅ DataLoaders are PINA-compliant!\n" + ] + } + ], + "source": [ + "def pina_collate(batch):\n", + " \"\"\"Bridges PyTorch tuples to PINA's condition-based format.\"\"\"\n", + " inputs = torch.stack([item[0] for item in batch])\n", + " outputs = torch.stack([item[1] for item in batch])\n", + " \n", + " condition_name = list(solver.problem.conditions.keys())[0]\n", + " \n", + " # Keys: 'input' and 'target' (required by InputTargetCondition)\n", + " # Labels: ['x'] and ['u'] (required by forward decorator)\n", + " data_dict = {\n", + " 'input': LabelTensor(inputs, ['x']),\n", + " 'target': LabelTensor(outputs, ['u'])\n", + " }\n", + " \n", + " return [(condition_name, data_dict)]\n", + "\n", + "# Create the DataLoaders with the fix\n", + "train_loader = DataLoader(\n", + " TensorDataset(train_input, train_output),\n", + " batch_size=32,\n", + " shuffle=True,\n", + " collate_fn=pina_collate\n", + ")\n", + "\n", + "val_loader = DataLoader(\n", + " TensorDataset(val_input, val_output),\n", + " batch_size=32,\n", + " shuffle=False,\n", + " collate_fn=pina_collate\n", + ")\n", + "\n", + "print(\"✅ DataLoaders are PINA-compliant!\")" + ] + }, + { + "cell_type": "code", + "execution_count": 11, + "id": "8759464a", + "metadata": {}, + "outputs": [ + { + "name": "stdout", + "output_type": "stream", + "text": [ + "✅ TopologyMonitor is ready!\n" + ] + } + ], + "source": [ + "def collect_predictions(trainer, pl_module):\n", + " \"\"\"Extract predictions and permute to Channels-First for GUDHI.\"\"\"\n", + " val_loader = trainer.val_dataloaders\n", + " if val_loader is None:\n", + " return None\n", + "\n", + " with torch.no_grad():\n", + " for batch in val_loader:\n", + " # batch is [(condition, {'input': LabelTensor, 'target': LabelTensor})]\n", + " _, data_dict = batch[0]\n", + " # Unpack using 'input' key\n", + " inputs = data_dict['input'].as_subclass(torch.Tensor).to(pl_module.device)\n", + " pred = pl_module.model(inputs) # Shape: (B, H, W, C)\n", + " # Permute to (B, C, H, W) because GUDHI expects Channels-First\n", + " return pred.permute(0, 3, 1, 2)\n", + " return None\n", + "\n", + "# Instantiate the TopologyMonitor\n", + "monitor = TopologyMonitor(\n", + " expected_beta_0=1, # Expect one connected component\n", + " expected_beta_1=0, # Expect no holes\n", + " monitor_freq=2,\n", + " warmup_epochs=1,\n", + " mode=\"warn\",\n", + " min_persistence=0.5,\n", + " collect_predictions=collect_predictions,\n", + ")\n", + "\n", + "print(\"✅ TopologyMonitor is ready!\")" + ] + }, + { + "cell_type": "markdown", + "id": "8230041c", + "metadata": {}, + "source": [ + "## The Prediction Collector (Topological Translation)\n", + "\n", + "FNO outputs are **Channels-Last**: `(Batch, Height, Width, Channels)`. \n", + "However, GUDHI's CubicalComplex (used by the monitor) expects **Channels-First**: `(Batch, Channels, Height, Width)`. \n", + "\n", + "The `collect_predictions` hook intercepts the FNO output, permutes the dimensions, and passes the correct grid to the `TopologyMonitor` without crashing the GPU training loop." + ] + }, + { + "cell_type": "code", + "execution_count": 12, + "id": "a2c0f177", + "metadata": {}, + "outputs": [ + { + "name": "stderr", + "output_type": "stream", + "text": [ + "GPU available: False, used: False\n", + "TPU available: False, using: 0 TPU cores\n", + "💡 Tip: For seamless cloud logging and experiment tracking, try installing [litlogger](https://pypi.org/project/litlogger/) to enable LitLogger, which logs metrics and artifacts automatically to the Lightning Experiments platform.\n", + "💡 Tip: For seamless cloud uploads and versioning, try installing [litmodels](https://pypi.org/project/litmodels/) to enable LitModelCheckpoint, which syncs automatically with the Lightning model registry.\n" + ] + }, + { + "name": "stdout", + "output_type": "stream", + "text": [ + "Starting training...\n" + ] + }, + { + "name": "stderr", + "output_type": "stream", + "text": [ + "\n", + " | Name | Type | Params | Mode | FLOPs\n", + "------------------------------------------------------------\n", + "0 | _pina_models | ModuleList | 106 K | train | 0 \n", + "1 | _loss_fn | MSELoss | 0 | train | 0 \n", + "------------------------------------------------------------\n", + "106 K Trainable params\n", + "0 Non-trainable params\n", + "106 K Total params\n", + "0.426 Total estimated model params size (MB)\n", + "27 Modules in train mode\n", + "0 Modules in eval mode\n", + "0 Total Flops\n" + ] + }, + { + "data": { + "application/vnd.jupyter.widget-view+json": { + "model_id": "a01300399ce24c3d9b2890dbf904821d", + "version_major": 2, + "version_minor": 0 + }, + "text/plain": [ + "Sanity Checking: | | 0/? [00:00" + ] + }, + "metadata": {}, + "output_type": "display_data" + }, + { + "name": "stdout", + "output_type": "stream", + "text": [ + "Topological Signature: β₀ = 1.00, β₁ = 0.00\n" + ] + } + ], + "source": [ + "from pina.topology.profiler import TopologicalProfiler\n", + "\n", + "# Get a sample from validation\n", + "sample_input = val_input[0:1]\n", + "sample_gt = val_output[0]\n", + "\n", + "# Forward pass\n", + "with torch.no_grad():\n", + " pred = fno(sample_input) # Shape: (1, H, W, 1)\n", + "\n", + "# Convert to Channels-First for the profiler\n", + "pred_channels_first = pred.permute(0, 3, 1, 2)\n", + "\n", + "# Compute Betti numbers\n", + "profiler = TopologicalProfiler(min_persistence=0.5)\n", + "result = profiler.compute(pred_channels_first)\n", + "\n", + "# Plotting\n", + "fig, ax = plt.subplots(1, 2, figsize=(10, 4))\n", + "\n", + "ax[0].imshow(sample_gt.squeeze().numpy(), cmap='jet')\n", + "ax[0].set_title(\"Ground Truth Pressure\")\n", + "ax[0].axis('off')\n", + "\n", + "ax[1].imshow(pred[0, ..., 0].numpy(), cmap='jet')\n", + "ax[1].set_title(f\"FNO Prediction\\nβ₀ = {result.beta_0_mean:.0f}, β₁ = {result.beta_1_mean:.0f}\")\n", + "ax[1].axis('off')\n", + "\n", + "plt.tight_layout()\n", + "plt.show()\n", + "\n", + "print(f\"Topological Signature: β₀ = {result.beta_0_mean:.2f}, β₁ = {result.beta_1_mean:.2f}\")" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "a136a61a", + "metadata": {}, + "outputs": [], + "source": [] + } + ], + "metadata": { + "kernelspec": { + "display_name": "Python (pina-dev)", + "language": "python", + "name": "pina-dev" + }, + "language_info": { + "codemirror_mode": { + "name": "ipython", + "version": 3 + }, + "file_extension": ".py", + "mimetype": "text/x-python", + "name": "python", + "nbconvert_exporter": "python", + "pygments_lexer": "ipython3", + "version": "3.10.20" + } + }, + "nbformat": 4, + "nbformat_minor": 5 +}