Added base classes for new architecture, updated mat solver for smart batching

Author: Tsuchijo

simpegtorch/simulation/base/__init__.py

@@ -0,0 +1,12 @@
+from .universal_simulation import UniversalSimulation
+from .dc_pde import DCResistivityPDE
+from .fdem_pde import FDEMPDE
+from ..basePDE import BasePDE, BaseMapping
+
+__all__ = [
+    "UniversalSimulation",
+    "DCResistivityPDE",
+    "FDEMPDE",
+    "BasePDE",
+    "BaseMapping",
+]

simpegtorch/simulation/base/basePDE.py

@@ -0,0 +1,77 @@
+from abc import ABC, abstractmethod
+import torch
+from simpegtorch.discretize.base import BaseMesh
+from mappings import BaseMapping
+
+
+class BasePDE(ABC):
+    """
+    Minimal PyTorch-native PDE interface for forward simulations.
+
+    This class defines the essential interface for any PDE formulation,
+    focusing on the three core operations needed for simulation:
+    1. System matrix construction
+    2. RHS vector construction
+    3. Field-to-data projection
+    4. A vec product (optional, for experimental indirect solver)
+
+    All gradient computation is handled automatically by PyTorch autograd.
+    """
+
+    def __init__(self, mesh: BaseMesh, mapping: BaseMapping):
+        self.mesh = mesh
+        self.mapping = mapping
+
+    @abstractmethod
+    def get_system_matrices(self) -> torch.Tensor:
+        """
+        Construct system matrices for the PDE.
+
+        Parameters
+        ----------
+        model_params : torch.Tensor
+            Physical model parameters (conductivity, permeability, etc.)
+
+        Returns
+        -------
+        Union[torch.Tensor, List[torch.Tensor]]
+            Single system matrix for problems like DC resistivity,
+            or list of matrices for frequency-domain problems.
+            Shape depends on problem type:
+            - DC: [1, n_cells, n_cells]
+            - FDEM: [n_frequencies, n_cells, n_cells]
+        """
+        pass
+
+    @abstractmethod
+    def get_rhs_tensors(self) -> torch.Tensor:
+        """
+        Construct right-hand side vectors for all sources.
+
+        Parameters
+
+        Returns
+        -------
+        torch.Tensor
+            Batched RHS tensors. Shape depends on problem type:
+            - DC: [1, n_sources, n_cells]
+            - FDEM: [n_frequencies, n_sources, n_cells]
+        """
+        pass
+
+    @abstractmethod
+    def fields_to_data(self, fields: torch.Tensor) -> torch.Tensor:
+        """
+        Project solution fields to predicted data.
+
+        Parameters
+        ----------
+        fields : torch.Tensor
+            Solution fields from PDE solve
+
+        Returns
+        -------
+        torch.Tensor
+            Predicted data vector [n_data]
+        """
+        pass

simpegtorch/simulation/base/direct_solver.py

@@ -0,0 +1,59 @@
+import torch
+from torch import nn
+from .basePDE import BasePDE
+from simpegtorch.torchmatsolver import batched_mumps_solve, batched_sparse_solve
+
+
+class DirectSolver(nn.Module):
+    """
+    PyTorch-native simulation class for a direct solver on any PDE formualtion.
+
+    This class provides a minimal, clean interface that works with any PDE. All gradient computation is handled
+    automatically by PyTorch autograd.
+
+    The entire simulation is a single forward pass which gives predicted data
+    """
+
+    def __init__(self, pde: BasePDE):
+        """
+        Initialize universal simulation.
+
+        Parameters
+        ----------
+        pde : BasePDE
+            PDE formulation defining the physics
+        survey : Survey
+            Survey object with sources and receivers
+        solver_method : str
+            Solver method ("Direct" or "Iterative")
+        """
+        super().__init__()
+        self.pde = pde
+
+    def forward(self) -> torch.Tensor:
+        """
+        Forward simulation: model parameters → predicted data.
+
+        Solves the PDE using a direct solver through matrixn inversion.
+        Total process:
+        1. Gets a RHS tensor of shape 1xkxn for n parameters problem and k different sources
+        2. Gets a A matrix tensor of fx1xnxn for n parameterrs and f different problems (e.g. different frequencies)
+        3. Runs a direct solver over this to solve for a fxkxn field solution tensor
+        4. for j different receivers projects the fields into fxj different measurements
+        """
+        # 1. Apply mapping if provided
+
+        # 2. Get system matrices from PDE
+        system_matrices = self.pde.get_system_matrices()
+
+        # 3. Get RHS vectors from PDE
+        rhs_tensors = self.pde.get_rhs_tensors()
+
+        try:
+            fields = batched_mumps_solve(system_matrices, rhs_tensors)
+        except ImportError:
+            print("MUMPS not installed, falling back")
+            fields = batched_sparse_solve(system_matrices, rhs_tensors)
+
+        # 5. Project to data and return
+        return self.pde.fields_to_data(fields)