Adapt forward modeling for multiple components

src/magali/_forward.py

@@ -0,0 +1,122 @@
+# Copyright (c) 2024 The Magali Developers.
+# Distributed under the terms of the BSD 3-Clause License.
+# SPDX-License-Identifier: BSD-3-Clause
+#
+# This code is part of the Fatiando a Terra project (https://www.fatiando.org)
+#
+"""
+Forward modeling code to calculate magnetic fields of finite sources.
+"""
+
+import harmonica as hm
+import verde as vd
+import xarray as xr
+
+from ._units import coordinates_micrometer_to_meter
+
+
+def dipole_magnetic(coordinates, dipole_coordinates, dipole_moments, field):
+    """
+    Compute the magnetic field produced by one or more dipoles.
+
+    Uses :func:`harmonica.dipole_magnetic` for the calculations and handles the
+    necessary unit versions from micrometers to meters.
+
+    Parameters
+    ----------
+    coordinates : tuple of float
+        Observation point coordinates in micrometers (μm).
+    dipole_coordinates : tuple of float
+        Dipole location coordinates in micrometers (μm).
+    dipole_moments : tuple of float
+        Dipole moment components (Am²).
+    field : str
+        Magnetic field that will be computed. The available fields are:
+
+        - The full magnetic vector: ``b``
+        - x component of the magnetic vector: ``b_x``
+        - y component of the magnetic vector: ``b_y``
+        - z (upward) component of the magnetic vector: ``b_z``
+
+    Returns
+    -------
+    magnetic_field : array or tuple of arrays
+        Computed magnetic field on every observation point in :math:`nT`.
+        If ``field`` is set to a single component, then a single array with the
+        computed magnetic field component will be returned.
+        If ``field`` is set to ``"b"``, then a tuple containing the three
+        components of the magnetic vector will be returned in the following
+        order: ``b_x``, ``b_y``, ``b_z``.
+    """
+    coordinates_m = coordinates_micrometer_to_meter(coordinates)
+    dipole_coordinates_m = coordinates_micrometer_to_meter(dipole_coordinates)
+    translator = {"b_x": "b_e", "b_y": "b_n", "b_z": "b_u", "b": "b"}
+    if field not in translator:
+        message = (
+            f"Invalid field argument '{field}'. Must be one of {translator.keys()}."
+        )
+        raise ValueError(message)
+    return hm.dipole_magnetic(
+        coordinates_m, dipole_coordinates_m, dipole_moments, field=translator[field]
+    )
+
+
+def dipole_magnetic_grid(
+    region, spacing, sensor_sample_distance, dipole_coordinates, dipole_moments, field
+):
+    """
+    Generate a grid of the magnetic field of one or more dipoles.
+
+    Uses :func:`magali.dipole_magnetic` to generate a regular grid of the
+    magnetic field produced by a dipole model. The grid is returned as an
+    :mod:`xarray` container type.
+
+    Parameters
+    ----------
+    region : tuple = (x_min, x_max, y_min, y_max)
+        The bounding box of the grid in micrometers (μm).
+    spacing : float
+        Grid spacing in micrometers (μm).
+    sensor_sample_distance : float
+        Distance of the sample from the grid in micrometers (μm).
+    dipole_coordinates : tuple = (x, y, z)
+        Arrays with the x, y, and z coordinates of the dipoles in micrometers
+        (μm). The arrays can have any shape as long as they all have the same
+        shape.
+    dipole_moments : tuple = (mx, my, mz)
+        Arrays with the x, y, and z components of the dipole moment of each
+        dipole in Am². The arrays should have a shape compatible with the
+        dipole coordinate arrays.
+    field : str
+        Magnetic field that will be computed. The available fields are:
+
+        - The full magnetic vector: ``b``
+        - x component of the magnetic vector: ``b_x``
+        - y component of the magnetic vector: ``b_y``
+        - z (upward) component of the magnetic vector: ``b_z``
+
+    Returns
+    -------
+    grid : :class:`xarray.DataArray` or :class:`xarray.Dataset`
+        Gridded magnetic field of the dipole model given. If *field* is
+        a single component, the grid is returned as
+        a :class:`xarray.DataArray`. If ``field ==  "b"``, the three components
+        of the magnetic field will be returned as a :class:`xarray.Dataset`.
+    """
+    coordinates = vd.grid_coordinates(
+        region=region,  # µm
+        spacing=spacing,  # µm
+        extra_coords=sensor_sample_distance,
+    )
+    data = dipole_magnetic(coordinates, dipole_coordinates, dipole_moments, field)
+    dims = ("y", "x")
+    coords = {
+        "x": coordinates[0][0, :],
+        "y": coordinates[1][:, 0],
+        "z": (dims, coordinates[2]),
+    }
+
+    grid.x.attrs = {"units": "µm"}
+    grid.y.attrs = {"units": "µm"}
+    grid.bz.attrs = {"long_name": "vertical magnetic field", "units": "nT"}
+    return data.bz

src/magali/_inversion.py

@@ -57,9 +57,9 @@ def fit(self, coordinates, data):
         Parameters
         ----------
         coordinates : tuple = (x, y, z)
-            Arrays with the x, y, and z coordinates of the observations points.
-            The arrays can have any shape as long as they all have the same
-            shape.
+            Arrays with the x, y, and z coordinates of the observations points
+            in micrometers (μm). The arrays can have any shape as long as they
+            all have the same shape.
         data : array
             Array with the observed Bz component of the magnetic field (in nT)
             at the locations provided in *coordinates*. Must have the same
@@ -87,9 +87,9 @@ def jacobian(self, coordinates):
         Parameters
         ----------
         coordinates : tuple = (x, y, z)
-            Arrays with the x, y, and z coordinates of the observations points.
-            The arrays can have any shape as long as they all have the same
-            shape.
+            Arrays with the x, y, and z coordinates of the observations points
+            in micrometers (μm). The arrays can have any shape as long as they
+            all have the same shape.
 
         Returns
         -------
@@ -185,9 +185,10 @@ def predict(self, coordinates):
 
         Parameters
         ----------
-        coordinates : tuple of array-like
-            Arrays with the x, y,    z coordinates of the observation points,
-            in µm. The arrays must have the same shape.
+        coordinates : tuple = (x, y, z)
+            Arrays with the x, y, and z coordinates of the computation points
+            in micrometers (μm). The arrays can have any shape as long as they
+            all have the same shape.
 
         Returns
         -------
@@ -221,9 +222,9 @@ def jacobian(self, coordinates, location, moment, jacobian):
         Parameters
         ----------
         coordinates : tuple = (x, y, z)
-            Arrays with the x, y, and z coordinates of the observations points.
-            The arrays can have any shape as long as they all have the same
-            shape.
+            Arrays with the x, y, and z coordinates of the observation points
+            in micrometers (μm). The arrays can have any shape as long as they
+            all have the same shape.
         location : array-like
             Dipole location (x, y, z), in µm.
         moment : array-like
@@ -286,9 +287,10 @@ def fit(self, coordinates, data):
 
         Parameters
         ----------
-        coordinates : tuple of array-like
-            Arrays with the x, y, z coordinates of the observation points.
-            The arrays can have any shape as long as they all have the same shape.
+        coordinates : tuple = (x, y, z)
+            Arrays with the x, y, and z coordinates of the observation points
+            in micrometers (μm). The arrays can have any shape as long as they
+            all have the same shape.
         data : array-like
             Observed Bz component of the magnetic field (in nT) at the observation
             points. Must have the same shape as the coordinate arrays.

src/magali/_synthetic.py

@@ -10,9 +10,6 @@
 
 import harmonica as hm
 import numpy as np
-import verde as vd
-
-from ._units import coordinates_micrometer_to_meter
 
 
 def random_directions(
@@ -145,72 +142,3 @@ def _rotate_vector(x, y, z, inclination, declination):
     y_r = np.sin(phi) * (x * np.cos(theta) + z * np.sin(theta)) + y * np.cos(phi)
     z_r = -x * np.sin(theta) + z * np.cos(theta)
     return x_r, y_r, z_r
-
-
-def dipole_bz(coordinates, dipole_coordinates, dipole_moments):
-    """
-    Compute the vertical component of the magnetic field (Bz) produced by a magnetic dipole.
-
-    Parameters
-    ----------
-    coordinates : tuple of float
-        Observation point coordinates in micrometers (μm).
-    dipole_coordinates : tuple of float
-        Dipole location coordinates in micrometers (μm).
-    dipole_moments : tuple of float
-        Dipole moment components (Am²).
-
-    Returns
-    -------
-    bz : float
-        The vertical component of the magnetic field (Bz) at the given observation point.
-    """
-    coordinates_m = coordinates_micrometer_to_meter(coordinates)
-    dipole_coordinates_m = coordinates_micrometer_to_meter(dipole_coordinates)
-    return hm.dipole_magnetic(
-        coordinates_m, dipole_coordinates_m, dipole_moments, field="b_u"
-    )
-
-
-def dipole_bz_grid(
-    region, spacing, sensor_sample_distance, dipole_coordinates, dipole_moments
-):
-    """
-    Generate a grid of the vertical component of the magnetic field of a dipole.
-
-    Parameters
-    ----------
-    region : tuple of float
-        The spatial region for the grid in micrometers (μm), defined as
-        (x_min, x_max, y_min, y_max).
-    spacing : float
-        Grid spacing in micrometers (μm).
-    sensor_sample_distance : float
-        Distance of the sensor from the grid in micrometers (μm).
-    dipole_coordinates : tuple of float
-        Dipole location coordinates in micrometers (μm).
-    dipole_moments : tuple of float
-        Dipole moment components (Am²).
-
-    Returns
-    -------
-    data : xarray.DataArray
-        Gridded dataset containing the vertical component of the magnetic
-        field (Bz).
-        The dataset includes:
-        - "bz" : vertical magnetic field (nT)
-        - "x" and "y" coordinates with units in micrometers (μm)
-    """
-    coordinates = vd.grid_coordinates(
-        region=region,  # µm
-        spacing=spacing,  # µm
-        extra_coords=sensor_sample_distance,
-    )
-    bz = dipole_bz(coordinates, dipole_coordinates, dipole_moments)
-    data = vd.make_xarray_grid(
-        coordinates, bz, data_names=["bz"], dims=("y", "x"), extra_coords_names="z"
-    )
-    data.x.attrs = {"units": "µm"}
-    data.y.attrs = {"units": "µm"}
-    data.bz.attrs = {"long_name": "vertical magnetic field", "units": "nT"}
-    return data.bz