[REFACTOR] Restructure geophysics input classes for better magnetics and gravity handling (#41)

# Refactor Geophysics Input Structure and Improve Tensor Type Handling

## [FIX] Improve tensor type matching and error reporting in backend_tensor.py

- Fix inconsistent tensor type matching for PyTorch tensors
- Enhance error message to include unsupported tensor type details

## [ENH] Refactor magnetics input structure and update TMI computation

- Introduced `GravityInput` and `MagneticsInput` dataclasses to encapsulate related geophysical fields
- Updated `GeophysicsInput` to utilize `GravityInput` and `MagneticsInput` for improved modularity
- Modified `compute_tmi` to accept `MagneticsInput` instead of `GeophysicsInput` for better type specificity
- Adjusted `model_api.py` to align with new `MagneticsInput` structure

## [ENH] Refactor magnetics input handling and relocate `_direction_cosines`

- Simplified magnetics input condition check in `model_api.py` by replacing `mag_kernel` and `susceptibilities` with `magnetics_input`
- Moved `_direction_cosines` method within `magnetic_gradient.py` for better modularity and code organization

Author: Leguark

gempy_engine/API/model/model_api.py

@@ -59,10 +59,10 @@ def compute_model(interpolation_input: InterpolationInput, options: Interpolatio
 
             # Magnetics (optional)
             try:
-                if getattr(geophysics_input, 'mag_kernel', None) is not None and getattr(geophysics_input, 'susceptibilities', None) is not None:
+                if getattr(geophysics_input, 'magnetics_input', None) is not None:
                     magnetics = compute_tmi(
-                        geophysics_input=geophysics_input,
-                        root_ouput=first_level_last_field
+                        geophysics_input=geophysics_input.magnetics_input,
+                        root_output=first_level_last_field
                     )
                 else:
                     magnetics = None

gempy_engine/core/backend_tensor.py

@@ -199,9 +199,9 @@ def _concatenate(tensors, axis=0, dtype=None):
             match type(tensors[0]):
                 case _ if any(isinstance(t, numpy.ndarray) for t in tensors):
                     return numpy.concatenate(tensors, axis=axis)
-                case torch.Tensor:
+                case _ if isinstance(tensors[0], torch.Tensor):
                     return torch.cat(tensors, dim=axis)
-            raise TypeError("Unsupported tensor type")
+            raise TypeError(f"Unsupported tensor type:  {type(tensors[0])}")
 
         def _transpose(tensor, axes=None):
             return tensor.transpose(axes[0], axes[1])

gempy_engine/core/data/geophysics_input.py

@@ -1,21 +1,66 @@
-from dataclasses import dataclass
+import warnings
+from dataclasses import dataclass
 from typing import Annotated, Optional
 
 import numpy as np
 
 from .encoders.converters import numpy_array_short_validator
 
 
+@dataclass
+class GravityInput:
+    tz: Annotated[np.ndarray, numpy_array_short_validator]
+    densities: Annotated[np.ndarray, numpy_array_short_validator]
+
+
+@dataclass
+class MagneticsInput:
+    mag_kernel: np.ndarray
+    susceptibilities: np.ndarray
+    igrf_params: dict
+
+
 @dataclass
 class GeophysicsInput:
-    # 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
+    gravity_input: Optional[GravityInput] = None
+    magnetics_input: Optional[MagneticsInput] = None
+    
+    def __init__(self, gravity_input: Optional[GravityInput] = None,
+                 magnetics_input: Optional[MagneticsInput] = None,
+                 tz: Optional[Annotated[np.ndarray, numpy_array_short_validator]] = None,
+                 densities: Optional[Annotated[np.ndarray, numpy_array_short_validator]] = None):
+        if gravity_input is not None:
+            self.gravity_input = gravity_input
+        else:
+            warnings.warn("Using deprecated GeophysicsInput constructor. Use GravityInput instead.", DeprecationWarning)
+            self.gravity_input = GravityInput(tz=tz, densities=densities)
+        if magnetics_input is not None:
+            self.magnetics_input = magnetics_input
+
+    @property
+    def tz(self) -> Optional[Annotated[np.ndarray, numpy_array_short_validator]]:
+        if self.gravity_input is not None:
+            return self.gravity_input.tz
+        return None
+
+    @tz.setter
+    def tz(self, value: Optional[Annotated[np.ndarray, numpy_array_short_validator]]):
+        if value is not None:
+            if self.gravity_input is None:
+                self.gravity_input = GravityInput(tz=value, densities=None)
+            else:
+                self.gravity_input.tz = value
+
+    @property
+    def densities(self) -> Optional[Annotated[np.ndarray, numpy_array_short_validator]]:
+        if self.gravity_input is not None:
+            return self.gravity_input.densities
+        return None
+
+    @densities.setter
+    def densities(self, value: Optional[Annotated[np.ndarray, numpy_array_short_validator]]):
+        if value is not None:
+            if self.gravity_input is None:
+                self.gravity_input = GravityInput(tz=None, densities=value)
+            else:
+                self.gravity_input.densities = value

gempy_engine/modules/geophysics/fw_magnetic.py

@@ -1,7 +1,7 @@
 import numpy as np
 
 from .magnetic_gradient import _direction_cosines
-from ...core.data.geophysics_input import GeophysicsInput
+from ...core.data.geophysics_input import GeophysicsInput, MagneticsInput
 from ...core.data.interp_output import InterpOutput
 from ...core.backend_tensor import BackendTensor
 
@@ -14,7 +14,7 @@ def map_susceptibilities_to_ids_basic(ids_geophysics_grid, susceptibilities):
     return susceptibilities[ids_geophysics_grid - 1]
 
 
-def compute_tmi(geophysics_input: GeophysicsInput, root_output: InterpOutput) -> BackendTensor.t:
+def compute_tmi(geophysics_input: MagneticsInput, root_output: InterpOutput) -> BackendTensor.t:
     """
     Compute induced-only Total Magnetic Intensity (TMI) anomalies (nT) by combining
     precomputed per-voxel TMI kernel with voxel susceptibilities.

gempy_engine/modules/geophysics/magnetic_gradient.py

@@ -3,30 +3,6 @@
 from gempy_engine.core.data.centered_grid import CenteredGrid
 
 
-def _direction_cosines(inclination_deg: float, declination_deg: float) -> np.ndarray:
-    """Compute unit vector of Earth's field from inclination/declination.
-
-    Convention:
-    - Inclination I: positive downward from horizontal, in degrees [-90, 90]
-    - Declination D: clockwise from geographic north toward east, in degrees [-180, 180]
-
-    Returns unit vector l = [lx, ly, lz].
-    """
-    I = np.deg2rad(inclination_deg)
-    D = np.deg2rad(declination_deg)
-    cI = np.cos(I)
-    sI = np.sin(I)
-    cD = np.cos(D)
-    sD = np.sin(D)
-    # North (x), East (y), Down (z) convention
-    l = np.array([cI * cD, cI * sD, sI], dtype=float)
-    # Already unit length by construction, but normalize defensively
-    n = np.linalg.norm(l)
-    if n == 0:
-        return np.array([1.0, 0.0, 0.0], dtype=float)
-    return l / n
-
-
 def calculate_magnetic_gradient_components(centered_grid: CenteredGrid) -> np.ndarray:
     """
     Calculate the 6 independent magnetic gradient tensor components (V1-V6) for each voxel.
@@ -51,7 +27,7 @@ def calculate_magnetic_gradient_components(centered_grid: CenteredGrid) -> np.nd
         over rectangular prism voxels using the formulas from Blakely (1995).
         The sign convention follows Talwani (z-axis positive downwards).
     """
-    
+
     voxel_centers = centered_grid.kernel_grid_centers
     center_x, center_y, center_z = voxel_centers[:, 0], voxel_centers[:, 1], voxel_centers[:, 2]
 
@@ -163,3 +139,27 @@ def calculate_magnetic_gradient_tensor(
         result["V"] = V
 
     return result
+
+
+def _direction_cosines(inclination_deg: float, declination_deg: float) -> np.ndarray:
+    """Compute unit vector of Earth's field from inclination/declination.
+
+    Convention:
+    - Inclination I: positive downward from horizontal, in degrees [-90, 90]
+    - Declination D: clockwise from geographic north toward east, in degrees [-180, 180]
+
+    Returns unit vector l = [lx, ly, lz].
+    """
+    I = np.deg2rad(inclination_deg)
+    D = np.deg2rad(declination_deg)
+    cI = np.cos(I)
+    sI = np.sin(I)
+    cD = np.cos(D)
+    sD = np.sin(D)
+    # North (x), East (y), Down (z) convention
+    l = np.array([cI * cD, cI * sD, sI], dtype=float)
+    # Already unit length by construction, but normalize defensively
+    n = np.linalg.norm(l)
+    if n == 0:
+        return np.array([1.0, 0.0, 0.0], dtype=float)
+    return l / n