Fixed basic implementaiton of FDEM simulation

Author: Tsuchijo

examples/test_fdem_functionality.py

@@ -38,11 +38,11 @@ def create_test_mesh():
     """Create a simple 3D test mesh"""
     print("Creating test mesh...")
 
-    # Create a 20x20x10 mesh with 5m cells
-    nx, ny, nz = 20, 20, 10
-    hx = torch.ones(nx) * 5.0  # 5m cell size
-    hy = torch.ones(ny) * 5.0
-    hz = torch.ones(nz) * 5.0
+    # Create a small 6x6x4 mesh with 10m cells for faster testing
+    nx, ny, nz = 6, 6, 4
+    hx = torch.ones(nx) * 10.0  # 10m cell size
+    hy = torch.ones(ny) * 10.0
+    hz = torch.ones(nz) * 10.0
 
     mesh = TensorMesh([hx, hy, hz])
     print(f"  Mesh shape: {mesh.shape_cells}")
@@ -64,9 +64,9 @@ def test_basic_magnetic_dipole():
     mesh = create_test_mesh()
 
     # Create receiver locations (line profile)
-    rx_x = torch.linspace(10, 90, 9)  # 9 receivers from 10m to 90m
-    rx_y = torch.ones_like(rx_x) * 50  # At y=50m
-    rx_z = torch.ones_like(rx_x) * 2.5  # At z=2.5m (just above surface)
+    rx_x = torch.linspace(15, 45, 4)  # 4 receivers from 15m to 45m
+    rx_y = torch.ones_like(rx_x) * 30  # At y=30m (center)
+    rx_z = torch.ones_like(rx_x) * 5.0  # At z=5m (just above surface)
     rx_locs = torch.stack([rx_x, rx_y, rx_z], dim=1)
 
     print(f"Receiver locations: {rx_locs.shape}")
@@ -81,7 +81,7 @@ def test_basic_magnetic_dipole():
     source = MagneticDipole(
         receiver_list=receivers,
         frequency=1000.0,  # 1 kHz
-        location=torch.tensor([50.0, 50.0, 0.0]),  # Center of domain at surface
+        location=torch.tensor([30.0, 30.0, 0.0]),  # Center of domain at surface
         moment=1.0,
         orientation="z",
     )
@@ -98,10 +98,10 @@ def test_basic_magnetic_dipole():
     # Add a conductive block
     cell_centers = mesh.cell_centers
     block_mask = (
-        (cell_centers[:, 0] > 40)
-        & (cell_centers[:, 0] < 60)
-        & (cell_centers[:, 1] > 40)
-        & (cell_centers[:, 1] < 60)
+        (cell_centers[:, 0] > 20)
+        & (cell_centers[:, 0] < 40)
+        & (cell_centers[:, 1] > 20)
+        & (cell_centers[:, 1] < 40)
         & (cell_centers[:, 2] > 10)
         & (cell_centers[:, 2] < 30)
     )
@@ -147,10 +147,10 @@ def test_multiple_frequencies():
     mesh = create_test_mesh()
 
     # Single receiver location
-    rx_loc = torch.tensor([[50.0, 50.0, 2.5]])
+    rx_loc = torch.tensor([[30.0, 30.0, 5.0]])
 
     sources = []
-    frequencies = [100.0, 1000.0, 10000.0]  # 100 Hz, 1 kHz, 10 kHz
+    frequencies = [1000.0, 10000.0]  # 1 kHz, 10 kHz (reduced to 2 frequencies)
 
     for freq in frequencies:
         receivers = [
@@ -161,7 +161,7 @@ def test_multiple_frequencies():
         source = MagneticDipole(
             receiver_list=receivers,
             frequency=freq,
-            location=torch.tensor([50.0, 50.0, 0.0]),
+            location=torch.tensor([30.0, 30.0, 0.0]),
             moment=1.0,
             orientation="z",
         )
@@ -202,7 +202,7 @@ def test_different_sources():
     print("=" * 60)
 
     mesh = create_test_mesh()
-    rx_loc = torch.tensor([[60.0, 50.0, 2.5]])
+    rx_loc = torch.tensor([[40.0, 30.0, 5.0]])
 
     # Test each source type
     source_tests = []
@@ -213,7 +213,7 @@ def test_different_sources():
         mag_dipole = MagneticDipole(
             receiver_list=receivers,
             frequency=1000.0,
-            location=torch.tensor([50.0, 50.0, 0.0]),
+            location=torch.tensor([30.0, 30.0, 0.0]),
             moment=1.0,
             orientation="z",
         )
@@ -234,7 +234,7 @@ def test_different_sources():
         elec_dipole = ElectricDipole(
             receiver_list=receivers,
             frequency=1000.0,
-            location=torch.tensor([50.0, 50.0, 0.0]),
+            location=torch.tensor([30.0, 30.0, 0.0]),
             current=1.0,
             length=1.0,
             orientation="z",
@@ -256,7 +256,7 @@ def test_different_sources():
         loop_source = LoopSource(
             receiver_list=receivers,
             frequency=1000.0,
-            location=torch.tensor([50.0, 50.0, 0.0]),
+            location=torch.tensor([30.0, 30.0, 0.0]),
             radius=10.0,
             current=1.0,
             orientation="z",
@@ -283,7 +283,7 @@ def test_receivers():
     print("=" * 60)
 
     mesh = create_test_mesh()
-    rx_loc = torch.tensor([[60.0, 50.0, 2.5]])
+    rx_loc = torch.tensor([[40.0, 30.0, 5.0]])
 
     receiver_tests = []
 
@@ -317,7 +317,7 @@ def test_receivers():
             source = MagneticDipole(
                 receiver_list=[receiver],
                 frequency=1000.0,
-                location=torch.tensor([50.0, 50.0, 0.0]),
+                location=torch.tensor([30.0, 30.0, 0.0]),
                 moment=1.0,
                 orientation="z",
             )

fdem_test_results.png

@@ -10,8 +10,8 @@ def __init__(self, locations, orientation="z", component="real"):
         ----------
         locations : array_like
             Receiver locations (n_receivers, 3)
-        orientation : str
-            Field component orientation ('x', 'y', 'z')
+        orientation : str or array_like
+            Field component orientation. Can be 'x', 'y', 'z' or a 3-vector like [0, 0, 1]
         component : str
             Complex component ('real', 'imag')
         """
@@ -21,7 +21,16 @@ def __init__(self, locations, orientation="z", component="real"):
             self.locations = torch.tensor(locations, dtype=torch.float64)
         if self.locations.ndim == 1:
             self.locations = self.locations.reshape(1, -1)
-        self.orientation = orientation
+
+        # Convert orientation to vector format
+        if isinstance(orientation, str):
+            orient_map = {"x": [1, 0, 0], "y": [0, 1, 0], "z": [0, 0, 1]}
+            self.orientation = torch.tensor(
+                orient_map[orientation], dtype=torch.float64
+            )
+        else:
+            self.orientation = torch.tensor(orientation, dtype=torch.float64)
+
         self.component = component
 
     @property
@@ -33,6 +42,36 @@ def evaluate(self, fields, simulation):
         """Extract data from fields at receiver locations"""
         raise NotImplementedError("evaluate must be implemented in derived classes")
 
+    def _get_projection_matrix(self, mesh, projected_grid):
+        """
+        Get projection matrix from mesh to receivers (following SimPEG pattern).
+
+        Parameters
+        ----------
+        mesh : TensorMesh
+            The mesh
+        projected_grid : str
+
+        Returns
+        -------
+        torch.Tensor
+            Projection matrix
+        """
+        # Initialize zero matrix
+        n_grid = getattr(mesh, f"n_{projected_grid[:-1]}")
+        P = torch.zeros(self.locations.shape[0], n_grid, dtype=torch.complex128)
+
+        # Component-wise interpolation like SimPEG
+        for strength, comp in zip(self.orientation, ["x", "y", "z"]):
+            if strength != 0.0:
+                location_type = projected_grid + comp  # e.g., "faces_x"
+                P_comp = mesh.get_interpolation_matrix(
+                    self.locations, location_type=location_type
+                )
+                P = P + strength * P_comp.to(dtype=torch.complex128)
+
+        return P
+
 
 class PointMagneticFluxDensity(BaseFDEMReceiver):
     """Point receiver for magnetic flux density measurements."""
@@ -53,29 +92,17 @@ def evaluate(self, fields, simulation):
         torch.Tensor
             Data at receiver locations
         """
-        # Get interpolation matrix from faces to receiver locations
-        P = simulation.mesh.get_interpolation_matrix(self.locations, location_type="F")
+        # Get projection matrix using component-wise interpolation (like SimPEG)
+        P = self._get_projection_matrix(simulation.mesh, "faces_")
 
         # Interpolate fields to receiver locations
         b_interp = P @ fields
 
-        # For 3D, need to extract the correct component
-        if simulation.mesh.dim == 3:
-            # Reshape to (n_receivers, 3) for vector field
-            b_vector = b_interp.reshape(-1, 3)
-
-            # Extract component
-            comp_idx = {"x": 0, "y": 1, "z": 2}[self.orientation]
-            b_comp = b_vector[:, comp_idx]
-        else:
-            # For 2D, assume z-component
-            b_comp = b_interp
-
         # Return real or imaginary part
         if self.component == "real":
-            return b_comp.real
+            return b_interp.real
         else:
-            return b_comp.imag
+            return b_interp.imag
 
 
 class PointMagneticFluxDensitySecondary(PointMagneticFluxDensity):
@@ -105,7 +132,7 @@ def evaluate(self, fields, simulation):
         E = -iωμ⁻¹∇×B
         """
         # Get curl matrix
-        C = simulation.mesh.edge_curl
+        C = simulation.mesh.edge_curl.to(dtype=torch.complex128)
 
         # Compute electric field from magnetic flux density
         # E = -iω * μ⁻¹ * C.T * B
@@ -118,20 +145,12 @@ def evaluate(self, fields, simulation):
         e_field = -1j * omega / mu0 * (C.T @ fields)
 
         # Interpolate to receiver locations
-        P = simulation.mesh.get_interpolation_matrix(self.locations, location_type="E")
+        P = self._get_projection_matrix(simulation.mesh, "edges_")
 
         e_interp = P @ e_field
 
-        # Extract component
-        if simulation.mesh.dim == 3:
-            e_vector = e_interp.reshape(-1, 3)
-            comp_idx = {"x": 0, "y": 1, "z": 2}[self.orientation]
-            e_comp = e_vector[:, comp_idx]
-        else:
-            e_comp = e_interp
-
         # Return real or imaginary part
         if self.component == "real":
-            return e_comp.real
+            return e_interp.real
         else:
-            return e_comp.imag
+            return e_interp.imag