[ENH] Add magnetic forward modeling for Total Magnetic Intensity (TMI) calculation (#39)

# Add magnetics (TMI) forward modeling to geophysics workflows

Introduce functionality for induced-only TMI anomaly computation using precomputed voxel-based TMI kernels and geologic unit susceptibilities. Updates include:

- `compute_tmi`: Core TMI calculation logic based on magnetics kernel and voxel susceptibilities in `fw_magnetic.py`.
- Extended `GeophysicsInput` to include magnetics-specific fields such as `mag_kernel`, `susceptibilities`, and `igrf_params`.
- Modified `model_api.py` to integrate magnetics into geophysics workflows alongside gravity computations.
- Updated `solutions.py` to include magnetics in forward modeling outputs.

Ensures magnetics functionality remains optional without impacting gravity workflows.

## Additional changes

- Updated `model_api.py` to make gravity computations conditional, enabling workflows without gravity inputs.
- Modified `GeophysicsInput` to allow `tz` and `densities` as optional fields for magnetics-only workflows.
- Added `tmi` alias property in `solutions.py` for Total Magnetic Intensity outputs.
- Improved formatting in geophysics README for better readability.
- Added analytical and symmetry benchmark tests for magnetics implementation validation.
- Included placeholder magnetics gradient computation module with core physics for initial testing.

Author: Leguark

gempy_engine/API/model/model_api.py

@@ -7,6 +7,7 @@
 from ...core.data.interp_output import InterpOutput
 from ...core.data.geophysics_input import GeophysicsInput
 from ...modules.geophysics.fw_gravity import compute_gravity
+from ...modules.geophysics.fw_magnetic import compute_tmi
 from ...core.data.dual_contouring_mesh import DualContouringMesh
 from ..dual_contouring.multi_scalar_dual_contouring import dual_contouring_multi_scalar
 from ..interp_single.interp_features import interpolate_n_octree_levels
@@ -46,12 +47,31 @@ def compute_model(interpolation_input: InterpolationInput, options: Interpolatio
 
         if geophysics_input is not None:
             first_level_last_field: InterpOutput = output[0].outputs_centers[-1]
-            gravity = compute_gravity(
-                geophysics_input=geophysics_input,
-                root_ouput=first_level_last_field
-            )
+
+            # Gravity (optional)
+            if getattr(geophysics_input, 'tz', None) is not None and getattr(geophysics_input, 'densities', None) is not None:
+                gravity = compute_gravity(
+                    geophysics_input=geophysics_input,
+                    root_ouput=first_level_last_field
+                )
+            else:
+                gravity = None
+
+            # Magnetics (optional)
+            try:
+                if getattr(geophysics_input, 'mag_kernel', None) is not None and getattr(geophysics_input, 'susceptibilities', None) is not None:
+                    magnetics = compute_tmi(
+                        geophysics_input=geophysics_input,
+                        root_ouput=first_level_last_field
+                    )
+                else:
+                    magnetics = None
+            except Exception:
+                # Keep gravity working even if magnetics paths are incomplete
+                magnetics = None
         else:
             gravity = None
+            magnetics = None
 
         # endregion
 
@@ -71,6 +91,7 @@ def compute_model(interpolation_input: InterpolationInput, options: Interpolatio
             octrees_output=output,
             dc_meshes=meshes,
             fw_gravity=gravity,
+            fw_magnetics=magnetics,
             block_solution_type=options.block_solutions_type
         )
 

gempy_engine/core/data/centered_grid.py

@@ -91,3 +91,49 @@ def create_irregular_grid_kernel(grid_resolution, scaling_factor, base_spacing=0
         flattened_right_voxel_edges = np.vstack(tuple(map(np.ravel, right_voxel_edges))).T.astype("float64")
 
         return flattened_grid_centers, flattened_left_voxel_edges, flattened_right_voxel_edges
+
+    def get_number_of_voxels_per_device(self) -> int:
+        """
+        Calculate the number of voxels in the kernel grid for a single device.
+
+        Returns:
+            int: Number of voxels per observation device
+
+        Notes:
+            - X and Y axes use symmetric grids: (resolution + 1) points each
+            - Z axis uses asymmetric grid (downward only): (resolution + 1) points
+            - Total voxels = (rx + 1) × (ry + 1) × (rz + 1)
+
+        Example:
+            >>> grid = CenteredGrid(centers=[[500, 500, 600]], 
+            ...                      resolution=[10, 10, 10], 
+            ...                      radius=[100, 100, 100])
+            >>> grid.get_number_of_voxels_per_device()
+            1331  # = 11 × 11 × 11
+        """
+        resolution = np.atleast_1d(self.resolution)
+
+        # Calculate points per axis following the create_irregular_grid_kernel logic
+        n_x = int(resolution[0] // 2) * 2 + 1  # Symmetric: 2 * (res//2) + 1
+        n_y = int(resolution[1] // 2) * 2 + 1  # Symmetric: 2 * (res//2) + 1
+        n_z = int(resolution[2] // 1) + 1  # Asymmetric: res + 1
+
+        return n_x * n_y * n_z
+
+    def get_total_number_of_voxels(self) -> int:
+        """
+        Calculate the total number of voxels across all observation devices.
+
+        Returns:
+            int: Total number of voxels (n_devices × voxels_per_device)
+
+        Example:
+            >>> grid = CenteredGrid(centers=[[500, 500, 600], [600, 500, 600]], 
+            ...                      resolution=[10, 10, 10], 
+            ...                      radius=[100, 100, 100])
+            >>> grid.get_total_number_of_voxels()
+            2662  # = 2 devices × 1331 voxels
+        """
+        n_devices = self.centers.shape[0]
+        voxels_per_device = self.get_number_of_voxels_per_device()
+        return n_devices * voxels_per_device

gempy_engine/core/data/geophysics_input.py

@@ -1,5 +1,5 @@
 from dataclasses import dataclass
-from typing import Annotated
+from typing import Annotated, Optional
 
 import numpy as np
 
@@ -8,5 +8,14 @@
 
 @dataclass
 class GeophysicsInput:
-    tz: Annotated[np.ndarray,  numpy_array_short_validator]
-    densities: Annotated[np.ndarray,  numpy_array_short_validator]
+    # Gravity inputs (optional to allow magnetics-only workflows)
+    tz: Optional[Annotated[np.ndarray, numpy_array_short_validator]] = None
+    densities: Optional[Annotated[np.ndarray, numpy_array_short_validator]] = None
+
+    # Magnetics inputs (optional for Phase 1)
+    # Pre-projected TMI kernel per voxel (per device geometry), shape: (n_voxels_per_device,)
+    mag_kernel: Optional[np.ndarray] = None
+    # Susceptibilities per geologic unit (dimensionless SI), shape: (n_units,)
+    susceptibilities: Optional[np.ndarray] = None
+    # IGRF parameters metadata used to build kernel (inclination, declination, intensity (nT))
+    igrf_params: Optional[dict] = None

gempy_engine/core/data/solutions.py

@@ -25,10 +25,12 @@ class Solutions:
     block_solution_type: RawArraysSolution.BlockSolutionType 
 
     def __init__(self, octrees_output: List[OctreeLevel], dc_meshes: List[DualContouringMesh] = None, fw_gravity: np.ndarray = None,
+                 fw_magnetics: np.ndarray = None,
                  block_solution_type: RawArraysSolution.BlockSolutionType = RawArraysSolution.BlockSolutionType.OCTREE):
         self.octrees_output = octrees_output
         self.dc_meshes = dc_meshes
         self.gravity = fw_gravity
+        self.magnetics = fw_magnetics
         self.block_solution_type = block_solution_type
 
         self._set_scalar_field_at_surface_points_and_elements_order(octrees_output)
@@ -124,3 +126,9 @@ def _set_scalar_field_at_surface_points_and_elements_order(self, octrees_output:
 
             # Order self_scalar_field_at_surface_points
             self._ordered_elements.append(np.argsort(scalar_field_at_surface_points_data)[::-1])
+
+
+    @property
+    def tmi(self) -> np.ndarray:
+        """Alias for magnetics Total Magnetic Intensity (nT)."""
+        return self.magnetics

gempy_engine/modules/geophysics/README.md

@@ -257,14 +257,14 @@ Magnetic field modeling is more complex than gravity because it involves **vecto
 
 ### Key Differences from Gravity
 
-| Aspect | Gravity | Magnetics |
-|--------|---------|-----------|
-| **Field type** | Scalar (gz only) | Vector (3 components) |
-| **Physical property** | Density (ρ) | Susceptibility (χ) + Remanence (Mr) |
-| **Source** | Mass distribution | Magnetic dipoles |
-| **Ambient field** | None (constant g) | Earth's field (varies by location/time) |
-| **Measurement** | Vertical acceleration | Total field intensity |
-| **Kernel complexity** | Single component (tz) | Full tensor (9 components) |
+| Aspect                | Gravity               | Magnetics                               |
+|-----------------------|-----------------------|-----------------------------------------|
+| **Field type**        | Scalar (gz only)      | Vector (3 components)                   |
+| **Physical property** | Density (ρ)           | Susceptibility (χ) + Remanence (Mr)     |
+| **Source**            | Mass distribution     | Magnetic dipoles                        |
+| **Ambient field**     | None (constant g)     | Earth's field (varies by location/time) |
+| **Measurement**       | Vertical acceleration | Total field intensity                   |
+| **Kernel complexity** | Single component (tz) | Full tensor (9 components)              |
 
 ### Mathematical Framework
 

gempy_engine/modules/geophysics/fw_magnetic.py

@@ -0,0 +1,141 @@
+import numpy as np
+
+from .magnetic_gradient import _direction_cosines
+from ...core.data.geophysics_input import GeophysicsInput
+from ...core.data.interp_output import InterpOutput
+from ...core.backend_tensor import BackendTensor
+
+
+def map_susceptibilities_to_ids_basic(ids_geophysics_grid, susceptibilities):
+    """
+    Map per-unit susceptibilities to voxel centers using 1-based geological IDs.
+    Matches gravity's basic mapper behavior.
+    """
+    return susceptibilities[ids_geophysics_grid - 1]
+
+
+def compute_tmi(geophysics_input: GeophysicsInput, root_output: InterpOutput) -> BackendTensor.t:
+    """
+    Compute induced-only Total Magnetic Intensity (TMI) anomalies (nT) by combining
+    precomputed per-voxel TMI kernel with voxel susceptibilities.
+
+    Expectations for Phase 1 (MVP):
+    - geophysics_input.mag_kernel is a scalar kernel per voxel (for one device geometry)
+      that already includes projection along the IGRF field direction and unit handling
+      (nT per unit susceptibility).
+    - geophysics_input.susceptibilities is an array of susceptibilities per geologic unit (SI).
+    - root_output.ids_geophysics_grid are 1-based IDs mapping each voxel to a geologic unit.
+
+    Returns:
+        BackendTensor.t array of shape (n_devices,) with TMI anomaly in nT.
+        
+    Notes:
+        This function works with pre-projected TMI kernels from 
+        calculate_magnetic_gradient_tensor(..., compute_tmi=True).
+        For more flexibility (e.g., remanent magnetization in Phase 2), 
+        consider using the raw V components and compute_magnetic_forward().
+    """
+    if geophysics_input.mag_kernel is None:
+        raise ValueError("GeophysicsInput.mag_kernel is required for magnetic forward modeling.")
+    if geophysics_input.susceptibilities is None:
+        raise ValueError("GeophysicsInput.susceptibilities is required for magnetic forward modeling.")
+
+    # Kernel for one device geometry
+    mag_kernel = BackendTensor.t.array(geophysics_input.mag_kernel)
+
+    # Map susceptibilities to voxel centers according to IDs (1-based indexing)
+    chi = map_susceptibilities_to_ids_basic(
+        ids_geophysics_grid=root_output.ids_geophysics_grid,
+        susceptibilities=BackendTensor.t.array(geophysics_input.susceptibilities)
+    )
+
+    # Determine how many devices are present by comparing lengths
+    n_devices = chi.shape[0] // mag_kernel.shape[0]
+
+    # Reshape for batch multiply-sum across devices
+    mag_kernel = mag_kernel.reshape(1, -1)
+    chi = chi.reshape(n_devices, -1)
+
+    # Weighted sum per device
+    tmi = BackendTensor.t.sum(chi * mag_kernel, axis=1)
+    return tmi
+
+
+def compute_magnetic_forward(
+        V: np.ndarray,
+        susceptibilities: np.ndarray,
+        igrf_params: dict,
+        n_devices: int = 1,
+        units_nT: bool = True
+) -> np.ndarray:
+    """
+    Compute Total Magnetic Intensity (TMI) anomalies from precomputed tensor components.
+
+    This follows the legacy implementation workflow: combine geometry-dependent V components
+    with susceptibility and IGRF field parameters to compute TMI.
+
+    Args:
+        V: np.ndarray of shape (6, n_voxels_per_device) from calculate_magnetic_gradient_components()
+        susceptibilities: np.ndarray of shape (n_total_voxels,) - susceptibility per voxel for all devices
+        igrf_params: dict with keys {"inclination", "declination", "intensity"} (intensity in nT)
+        n_devices: Number of observation devices (default 1)
+        units_nT: If True, output in nT; if False, output in Tesla
+
+    Returns:
+        dT: np.ndarray of shape (n_devices,) - TMI anomaly at each observation point
+
+    Notes:
+        Implements the formula from legacy code:
+        - Compute induced magnetization: J = chi * B_ext
+        - Compute directional components: Jx, Jy, Jz
+        - Apply gradient tensor: Tx, Ty, Tz using V components
+        - Project onto field direction: dT = Tx*dir_x + Ty*dir_y + Tz*dir_z
+    """
+    # Extract IGRF parameters
+    incl = float(igrf_params.get("inclination", 0.0))
+    decl = float(igrf_params.get("declination", 0.0))
+    B_ext = float(igrf_params.get("intensity", 50000.0))  # in nT
+
+    # Convert to Tesla for internal computation
+    B_ext_tesla = B_ext * 1e-9
+
+    # Get field direction cosines
+    dir_x, dir_y, dir_z = _direction_cosines(incl, decl)
+
+    # Compute induced magnetization [Tesla]
+    # J = chi * B_ext (susceptibility is dimensionless)
+    J = susceptibilities * B_ext_tesla
+
+    # Compute magnetization components along field direction
+    Jx = dir_x * J
+    Jy = dir_y * J
+    Jz = dir_z * J
+
+    # Tile V for multiple devices (repeat the kernel for each device)
+    V_tiled = np.tile(V, (1, n_devices))
+
+    # Compute directional magnetic effect on each voxel using V components
+    # This is equation (3.19) from the theory
+    # Bx = Jx*V1 + Jy*V2 + Jz*V3
+    # By = Jx*V2 + Jy*V4 + Jz*V5  (V2 = V[1] because tensor is symmetric)
+    # Bz = Jx*V3 + Jy*V5 + Jz*V6
+    Tx = (Jx * V_tiled[0, :] + Jy * V_tiled[1, :] + Jz * V_tiled[2, :]) / (4 * np.pi)
+    Ty = (Jx * V_tiled[1, :] + Jy * V_tiled[3, :] + Jz * V_tiled[4, :]) / (4 * np.pi)
+    Tz = (Jx * V_tiled[2, :] + Jy * V_tiled[4, :] + Jz * V_tiled[5, :]) / (4 * np.pi)
+
+    # Sum over voxels for each device
+    n_voxels_per_device = V.shape[1]
+    Tx = np.sum(Tx.reshape((n_devices, n_voxels_per_device)), axis=1)
+    Ty = np.sum(Ty.reshape((n_devices, n_voxels_per_device)), axis=1)
+    Tz = np.sum(Tz.reshape((n_devices, n_voxels_per_device)), axis=1)
+
+    # Project onto field direction to get TMI
+    # "Total field magnetometers can measure only that part of the anomalous field 
+    # which is in the direction of the Earth's main field" (SimPEG documentation)
+    dT = Tx * dir_x + Ty * dir_y + Tz * dir_z
+
+    if units_nT:
+        # Convert to nT
+        dT = dT * 1e9
+
+    return dT