Updated DCR tests to use new formulation

Author: Tsuchijo

examples/dc_torch_example.py

@@ -0,0 +1,76 @@
+import torch
+from simpegtorch.discretize import TensorMesh
+from simpegtorch.simulation.resistivity import (
+    DC3DCellCentered,
+    SrcDipole,
+    RxDipole,
+    Survey,
+)
+
+from simpegtorch.simulation.base import DirectSolver, mappings
+
+from simpegtorch.discretize.utils import (
+    ndgrid,
+)
+
+# Create a tensor mesh
+hx = torch.ones(10) * 25
+hy = torch.ones(10) * 25
+hz = torch.ones(10) * 25
+
+# 250m x 250m x 250m mesh
+
+# Origin at (-125, -125, -250) to center the mesh
+origin = torch.tensor([-125.0, -125.0, -250.0])
+
+mesh = TensorMesh(
+    [hx, hy, hz],
+    origin=origin,
+)
+
+sigma = torch.ones(mesh.nC) * 1e-2  # Uniform conductivity
+sigma_map = mappings.BaseMapping(sigma)
+
+# Set up survey parameters for numeric solution
+x = mesh.cell_centers_x[(mesh.cell_centers_x > -75.0) & (mesh.cell_centers_x < 75.0)]
+y = mesh.cell_centers_y[(mesh.cell_centers_y > -75.0) & (mesh.cell_centers_y < 75.0)]
+
+M = ndgrid(x - 25.0, y, [0.0])
+N = ndgrid(x + 25.0, y, [0.0])
+
+# create a dipole dipole survey
+rx = RxDipole(
+    locations_m=M,
+    locations_n=N,
+)
+
+loc_a = torch.tensor([-25.0, 0.0, 0.0])
+loc_b = torch.tensor([25.0, 0.0, 0.0])
+
+src = SrcDipole(
+    [rx],
+    loc_a,  # location of A
+    loc_b,  # location of B
+    current=1.0,  # current in Amperes
+)
+
+survey = Survey([src])
+
+## Setup the Problem as a Cell Centered PDE problem
+problem = DC3DCellCentered(
+    mesh,
+    survey,
+    sigma_map,
+    bc_type="Dirichlet",
+)
+
+# Create the solver we will use
+solver = DirectSolver(problem)
+
+# Solve the forward problem
+Data = solver.forward()
+
+# Print results
+print("Forward simulation completed successfully!")
+print(f"Data values: {Data}")
+print(f"Data min: {Data.min():.6e}, max: {Data.max():.6e}")

simpegtorch/__init__.py

@@ -1,3 +1,2 @@
 from . import discretize
 from . import utils
-from . import maps

simpegtorch/simulation/base/__init__.py

@@ -1,12 +1,2 @@
-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",
-]
+from .basePDE import BasePDE
+from .direct_solver import DirectSolver

simpegtorch/simulation/base/basePDE.py

@@ -1,7 +1,7 @@
 from abc import ABC, abstractmethod
 import torch
 from simpegtorch.discretize.base import BaseMesh
-from mappings import BaseMapping
+from .mappings import BaseMapping
 
 
 class BasePDE(ABC):

simpegtorch/simulation/resistivity/__init__.py

@@ -25,6 +25,7 @@
 """
 
 from .simulation import Simulation3DCellCentered, Simulation3DNodal
+from .dc_pde import DC3DCellCentered, DC3DNodal
 from .sources import BaseSrc, Pole as SrcPole, Dipole as SrcDipole, Multipole
 from .receivers import BaseRx, Pole as RxPole, Dipole as RxDipole
 from .survey import Survey
@@ -44,6 +45,9 @@
     # 3D simulations
     "Simulation3DCellCentered",
     "Simulation3DNodal",
+    # PDE problems
+    "DC3DCellCentered",
+    "DC3DNodal",
     # Sources
     "BaseSrc",
     "Src",

simpegtorch/simulation/resistivity/dc_pde.py

@@ -1,5 +1,5 @@
 import torch
-from ..basePDE import BasePDE
+from simpegtorch.simulation.base.basePDE import BasePDE
 from simpegtorch.discretize import TensorMesh
 from simpegtorch.discretize.utils import sdiag
 
@@ -34,7 +34,7 @@ def __init__(
         survey : Survey
             Survey object with sources and receivers
         mapping : BaseMapping
-            Parameter generating function
+            Parameter generating function, by default this is conductivity
         """
         super().__init__(mesh, mapping)
         self.survey = survey
@@ -44,14 +44,10 @@ def __init__(
     def setBC(self):
         mesh = self.mesh
         # Standard cell-centered finite volume discretization
-        # Volume scaling is needed for proper conservation
         V = sdiag(mesh.cell_volumes)
         self.Div = V @ mesh.face_divergence
         self.Grad = self.Div.T
 
-        # Initialize MfRhoI here, it will be updated in getA
-        self.MfRhoI = None
-
         if self.bc_type == "Dirichlet":
             print("Homogeneous Dirichlet is the natural BC for this CC discretization")
             return
@@ -108,12 +104,12 @@ def get_system_matrices(self) -> torch.Tensor:
         D = self.Div
         G = self.Grad
 
-        # Use face inner product with inverse resistivity (conductivity)
-        self.MfRhoI = self.mesh.get_face_inner_product(
+        # Use face inner product with resistivity from mapping (inverted to get conductivity)
+        MfRhoI = self.mesh.get_face_inner_product(
             self.mapping.forward(), invert_matrix=True
         )
 
-        A = D @ self.MfRhoI @ G
+        A = D @ MfRhoI @ G
 
         if self.bc_type.lower() == "neumann":
             A = self.condition_matrix(A)
@@ -156,7 +152,162 @@ def fields_to_data(self, fields: torch.Tensor) -> torch.Tensor:
 
             # Get receiver projection tensor and project to data
             rx_tensor = src.build_receiver_tensor(self.mesh, "CC")
-            rx_data = torch.mv(rx_tensor, src_field)
+
+            # Handle sparse tensor matrix multiplication
+            # rx_tensor shape: [1, n_receivers, n_mesh_cells]
+            # src_field shape: [n_mesh_cells]
+            if rx_tensor.is_sparse:
+                # For sparse tensors, we need to handle the multiplication differently
+                # Select the first batch (index 0) and multiply with each receiver projection
+                rx_tensor_2d = rx_tensor[0]  # [n_receivers, n_mesh_cells]
+                rx_data = torch.sparse.mm(
+                    rx_tensor_2d, src_field.unsqueeze(-1)
+                ).squeeze(-1)
+            else:
+                # Dense tensor multiplication
+                rx_data = rx_tensor.squeeze(0) @ src_field
+
+            data_list.append(rx_data)
+
+        return torch.cat(data_list, dim=0)
+
+
+class DC3DNodal(BasePDE):
+    """
+    DC Resistivity PDE for Node centered formulation
+    """
+
+    def __init__(
+        self,
+        mesh: TensorMesh,
+        survey,
+        mapping,
+        bc_type: str = "Dirichlet",
+    ):
+        """
+        Initialize DC Resistivity PDE.
+
+        Parameters
+        ----------
+        mesh : TensorMesh
+            Discretization mesh
+        survey : Survey
+            Survey object with sources and receivers
+        mapping : BaseMapping
+            Parameter generating function
+        """
+        super().__init__(mesh, mapping)
+        self.survey = survey
+        self.bc_type = bc_type
+        self.setBC()
+
+    def setBC(self):
+        valid_bc_types = ["neumann"]
+        if self.bc_type.lower() not in valid_bc_types:
+            raise ValueError(f"Unsupported boundary condition type: {self.bc_type}")
+
+    def get_system_matrices(self) -> torch.Tensor:
+        """
+        Construct system matrix for the nodal DC resistivity PDE.
+        Returns
+        -------
+        torch.Tensor
+            Single system matrix for nodal DC resistivity.
+            Shape [1, nN, nN]
+        """
+        # Use edge inner product with conductivity (1/resistivity)
+        MeSigma = self.mesh.get_edge_inner_product(1.0 / self.mapping.forward())
+        Grad = self.mesh.nodal_gradient
+        A = Grad.T @ MeSigma @ Grad
+
+        if self.bc_type.lower() == "neumann":
+            A = self.condition_matrix(A)
+
+        return torch.stack([A], dim=0)  # Shape [1, nN, nN]
+
+    def condition_matrix(self, A: torch.Tensor) -> torch.Tensor:
+        """
+        Condition the system matrix for numerical stability.
+        Used for Neumann boundary conditions.
+        """
+        # Create a sparse tensor that represents the modification
+        # It will have 1 at (0,0) and 0 elsewhere
+        mod_indices = torch.tensor([[0], [0]], dtype=torch.long, device=A.device)
+        mod_values = torch.tensor([1.0], dtype=A.dtype, device=A.device)
+        modification_matrix = torch.sparse_coo_tensor(
+            mod_indices, mod_values, A.shape
+        ).coalesce()
+
+        # Create a mask for the first row of A
+        first_row_mask = A.indices()[0] == 0
+
+        # Get the values of the first row of A
+        first_row_values = A.values()[first_row_mask]
+        first_row_cols = A.indices()[1, first_row_mask]
+
+        # Create a sparse tensor representing the negative of the first row of A
+        neg_first_row_indices = torch.stack(
+            [torch.zeros_like(first_row_cols), first_row_cols]
+        )
+        neg_first_row_values = -first_row_values
+        neg_first_row_matrix = torch.sparse_coo_tensor(
+            neg_first_row_indices, neg_first_row_values, A.shape
+        ).coalesce()
+
+        # Add the modification matrix and the negative of the first row to A
+        A = A + neg_first_row_matrix + modification_matrix
+        return A
+
+    def get_rhs_tensors(self) -> torch.Tensor:
+        """
+        Construct RHS vectors for all sources.
+        Returns
+        -------
+        torch.Tensor
+            Batched RHS tensors. Shape [n_sources, nC]
+        """
+
+        return self.survey.get_source_tensor(self.mesh, projected_grid="N")
+
+    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, shape [1, n_sources, nC]
+        Returns
+        -------
+        torch.Tensor
+            Predicted data vector [n_data]
+        """
+        # Remove frequency dimension if present (shape: [n_sources, nC])
+        if fields.dim() == 3:
+            fields = fields.squeeze(0)
+
+        sources = self.survey.source_list
+        data_list = []
+
+        for i, src in enumerate(sources):
+            # Get field for this source
+            src_field = fields[i]
+
+            # Get receiver projection tensor and project to data
+            rx_tensor = src.build_receiver_tensor(self.mesh, "N")
+
+            # Handle sparse tensor matrix multiplication
+            # rx_tensor shape: [1, n_receivers, n_mesh_nodes]
+            # src_field shape: [n_mesh_nodes]
+            if rx_tensor.is_sparse:
+                # For sparse tensors, we need to handle the multiplication differently
+                # Select the first batch (index 0) and multiply with each receiver projection
+                rx_tensor_2d = rx_tensor[0]  # [n_receivers, n_mesh_nodes]
+                rx_data = torch.sparse.mm(
+                    rx_tensor_2d, src_field.unsqueeze(-1)
+                ).squeeze(-1)
+            else:
+                # Dense tensor multiplication
+                rx_data = rx_tensor.squeeze(0) @ src_field
 
             data_list.append(rx_data)
 

simpegtorch/simulation/resistivity/sources.py

@@ -127,7 +127,16 @@ def evaluate(self, mesh, projected_grid: str = "CC"):
         """
 
         indices = mesh.closest_points_index(self.location, grid_loc=projected_grid)
-        q = torch.zeros(mesh.nC, dtype=self.current.dtype, device=self.current.device)
+
+        # Choose the right vector size based on formulation
+        if projected_grid in ["N", "nodes"]:
+            vector_size = mesh.nN  # Number of nodes for nodal formulation
+        else:
+            vector_size = mesh.nC  # Number of cells for cell-centered formulation
+
+        q = torch.zeros(
+            vector_size, dtype=self.current.dtype, device=self.current.device
+        )
         q[indices] = self.current
         return q