Repository: fatiando/harmonica (python)

Issue:
Update magnetic forward functions for dipoles
**Description of the desired feature:**

After solving #446 and #447 with #448, we should apply similar changes to the dipole magnetic forward functions.

In a nutshell:
- merge the `dipole_magnetic` and `dipole_magnetic_component` functions into a single one,
- change the shape and type of the expected `magnetic_moment` argument: it should be a tuple of three arrays: `mag_moment_e`, `mag_moment_n`, `mag_moment_u`.


**Are you willing to help implement and maintain this feature?** 

Sure!

Pull request:
Merge magnetic forward functions for dipoles
Merge the two forward functions for magnetic fields of dipoles into a single one named `dipole_magnetic` that includes a new `field` argument. Make the `dipole_magnetic` function to ask magnetic moment vectors as a tuple of three arrays, where each one of these arrays. Update tests accordingly. Update the user guide page about forward modelling of prisms.


**Relevant issues/PRs:**

Fixes #450

Implements the same interface as #448


Patch:
diff --git a/doc/api/index.rst b/doc/api/index.rst
index 15cb6b8..8a522a2 100644
--- a/doc/api/index.rst
+++ b/doc/api/index.rst
@@ -80,7 +80,6 @@ Magnetic fields:
     :toctree: generated/
 
     dipole_magnetic
-    dipole_magnetic_component
     prism_magnetic
 
 Layers and meshes:
diff --git a/doc/user_guide/forward_modelling/dipole.rst b/doc/user_guide/forward_modelling/dipole.rst
index e4f0dee..e7a5e9a 100644
--- a/doc/user_guide/forward_modelling/dipole.rst
+++ b/doc/user_guide/forward_modelling/dipole.rst
@@ -4,11 +4,7 @@ Magnetic dipoles
 ================
 
 We can compute the magnetic field generated by dipole sources through
-:func:`harmonica.dipole_magnetic` and
-:func:`harmonica.dipole_magnetic_component` functions.
-The former computes the three components of the magnetic field generated by
-a set of dipoles on every observation point, while the latter computes only
-a single component.
+the :func:`harmonica.dipole_magnetic` function.
 
 Each dipole can be defined as a tuple containing its coordinates in the
 following order: *easting*, *northing*, *upward* (in Cartesian coordinates and
@@ -27,7 +23,7 @@ Define the dipole and its magnetic moment vector:
 .. jupyter-execute::
 
    dipole = (20, 40, -50)
-   magnetic_moment = [100, 100, 100]
+   magnetic_moment = (100, 100, 100)
 
 Define the observation point above it:
 
@@ -35,20 +31,21 @@ Define the observation point above it:
 
    coordinates = (20, 40, 10)
 
-And compute the three components of the magnetic field the dipole generates:
+And compute the three components of the magnetic field the dipole generates by
+choosing ``field="b"``:
 
 .. jupyter-execute::
 
-   b_e, b_n, b_u = hm.dipole_magnetic(coordinates, dipole, magnetic_moment)
+   b_e, b_n, b_u = hm.dipole_magnetic(coordinates, dipole, magnetic_moment, field="b")
    print(b_e, b_n, b_u)
 
-We can use :func:`harmonica.dipole_magnetic_component` to compute just a single
-component:
+We can compute just a single component with ``field`` set to ``"b_e"``,
+``"b_n"`` or ``"b_u"``:
 
 .. jupyter-execute::
 
-   b_u = hm.dipole_magnetic_component(
-      coordinates, dipole, magnetic_moment, component="upward"
+   b_u = hm.dipole_magnetic(
+      coordinates, dipole, magnetic_moment, field="b_u"
    )
    print(b_u)
 
@@ -56,7 +53,7 @@ component:
 Multiple dipoles
 ----------------
 
-These two functions also allow us to compute the magnetic fields of multiple
+We can also use this function to compute the magnetic fields of multiple
 dipoles on multiple observation points.
 
 Let's define a regular grid of observation points:
@@ -83,19 +80,17 @@ And a set of dipoles with their magnetic moments:
    upward = [-200, -100, -300, -150]
    dipoles = (easting, northing, upward)
 
-   magnetic_moments = np.array([
-      [1e3, 1e3, 1e3],
-      [2e3, 2e3, 2e3],
-      [500, 500, 500],
-      [2e3, 2e3, 2e3],
-   ])
+   mag_e = [1e3, 2e3, 500, 2e3]
+   mag_n = [1e3, 2e3, 500, 2e3]
+   mag_u = [1e3, 2e3, 500, 2e3]
+   magnetic_moments = (mag_e, mag_n, mag_u)
 
 Now, let's compute the magnetic field components that the dipoles generate on
 every observation point:
 
 .. jupyter-execute::
 
-   b_e, b_n, b_u = hm.dipole_magnetic(coordinates, dipoles, magnetic_moments)
+   b_e, b_n, b_u = hm.dipole_magnetic(coordinates, dipoles, magnetic_moments, field="b")
 
 .. jupyter-execute::
 
diff --git a/harmonica/__init__.py b/harmonica/__init__.py
index 2195e75..4902d8f 100644
--- a/harmonica/__init__.py
+++ b/harmonica/__init__.py
@@ -9,7 +9,7 @@
 from ._equivalent_sources.cartesian import EquivalentSources
 from ._equivalent_sources.gradient_boosted import EquivalentSourcesGB
 from ._equivalent_sources.spherical import EquivalentSourcesSph
-from ._forward.dipole import dipole_magnetic, dipole_magnetic_component
+from ._forward.dipole import dipole_magnetic
 from ._forward.point import point_gravity
 from ._forward.prism_gravity import prism_gravity
 from ._forward.prism_layer import DatasetAccessorPrismLayer, prism_layer
diff --git a/harmonica/_forward/dipole.py b/harmonica/_forward/dipole.py
index b3567e5..62aeec8 100644
--- a/harmonica/_forward/dipole.py
+++ b/harmonica/_forward/dipole.py
@@ -7,37 +7,65 @@
 """
 Forward modelling for magnetic fields of dipoles
 """
+from math import prod
+
 import numpy as np
 from choclo.dipole import magnetic_e, magnetic_field, magnetic_n, magnetic_u
 from numba import jit, prange
 
+from .utils import initialize_progressbar
+
+VALID_FIELDS = ("b", "b_e", "b_n", "b_u")
+FORWARD_FUNCTIONS = {
+    "b_e": magnetic_e,
+    "b_n": magnetic_n,
+    "b_u": magnetic_u,
+}
+
 
 def dipole_magnetic(
     coordinates,
     dipoles,
     magnetic_moments,
+    field,
     parallel=True,
     dtype="float64",
+    progressbar=False,
     disable_checks=False,
 ):
     """
     Magnetic field of dipoles in Cartesian coordinates
 
+    Compute the component(s) of the magnetic field vector generated by
+    a collection of dipoles on a set of observation points.
+
+    .. important::
+
+        The component(s) of the magnetic field are returned in :math:`nT`.
+
     Parameters
     ----------
     coordinates : list of arrays
         List of arrays containing the ``easting``, ``northing`` and ``upward``
         coordinates of the computation points defined on a Cartesian coordinate
         system. All coordinates should be in meters.
-    dipoles : list of arrays
-        List of arrays containing the ``easting``, ``northing`` and ``upward``
+    dipoles : tuple of arrays
+        Tuple of arrays containing the ``easting``, ``northing`` and ``upward``
         locations of the dipoles defined on a Cartesian coordinate system. All
         coordinates should be in meters.
-    magnetic_moments : 2d-array
-        Two dimensional array containing the magnetic moments of each dipole in
-        :math:`Am^2`. Each row contains the components of the magnetic moment
-        for each dipole in the following order: ``mag_moment_easting``,
+    magnetic_moments : tuple of arrays
+        Tuple containing the three arrays corresponding to the magnetic moment
+        components of each dipole in :math:`Am^2`. These arrays should
+        be provided in the following order: ``mag_moment_easting``,
         ``mag_moment_northing``, ``mag_moment_upward``.
+    field : str
+        Magnetic field that will be computed. The available fields are:
+
+        - The full magnetic vector: ``b``
+        - Easting component of the magnetic vector: ``b_e``
+        - Northing component of the magnetic vector: ``b_n``
+        - Upward component of the magnetic vector: ``b_u``
+
     parallel : bool (optional)
         If True the computations will run in parallel using Numba built-in
         parallelization. If False, the forward model will run on a single core.
@@ -46,6 +74,10 @@ def dipole_magnetic(
     dtype : data-type (optional)
         Data type assigned to the resulting magnetic field. Default to
         ``np.float64``.
+    progressbar : bool (optional)
+        If True, a progress bar of the computation will be printed to standard
+        error (stderr). Requires :mod:`numba_progress` to be installed.
+        Default to ``False``.
     disable_checks : bool (optional)
         Flag that controls whether to perform a sanity check on the model.
         Should be set to ``True`` only when it is certain that the input model
@@ -54,137 +86,207 @@ def dipole_magnetic(
 
     Returns
     -------
-    magnetic_field : tuple of arrays
-        Tuple containing each component of the magnetic field generated by the
-        dipoles as arrays. The three components are returned in the following
-        order: ``b_e``, ``b_n``, ``b_u``. Magnetic field components are in
-        :math:`nT`.
+    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_e``, ``b_n``, ``b_u``.
     """
-    # Figure out the shape and size of the output arrays
+    # Check if field is valid
+    if field not in VALID_FIELDS:
+        raise ValueError(
+            f"Invalid field '{field}'. "
+            f"Please choose one of '{', '.join(VALID_FIELDS)}'."
+        )
+    # Figure out the shape and size of the output array(s)
     cast = np.broadcast(*coordinates[:3])
-    b_e, b_n, b_u = tuple(np.zeros(cast.size, dtype=dtype) for _ in range(3))
     # Convert coordinates, dipoles and magnetic moments to arrays
     coordinates = tuple(np.atleast_1d(i).ravel() for i in coordinates[:3])
     dipoles = tuple(np.atleast_1d(i).ravel() for i in dipoles[:3])
-    magnetic_moments = np.atleast_2d(magnetic_moments)
+    magnetic_moments = tuple(np.atleast_1d(m).ravel() for m in magnetic_moments)
     # Sanity checks
     if not disable_checks:
         _check_dipoles_and_magnetic_moments(dipoles, magnetic_moments)
-    # Compute the magnetic fields of the dipoles
-    if parallel:
-        _jit_dipole_magnetic_field_cartesian_parallel(
-            coordinates, dipoles, magnetic_moments, b_e, b_n, b_u
+    # Run computations
+    if field == "b":
+        result = _dipole_magnetic_vector(
+            coordinates,
+            dipoles,
+            magnetic_moments,
+            cast.shape,
+            dtype,
+            parallel,
+            progressbar,
         )
     else:
-        _jit_dipole_magnetic_field_cartesian_serial(
-            coordinates, dipoles, magnetic_moments, b_e, b_n, b_u
+        forward_func = FORWARD_FUNCTIONS[field]
+        result = _dipole_magnetic_component(
+            coordinates,
+            dipoles,
+            magnetic_moments,
+            forward_func,
+            cast.shape,
+            dtype,
+            parallel,
+            progressbar,
         )
+    return result
+
+
+def _dipole_magnetic_vector(
+    coordinates,
+    dipoles,
+    magnetic_moments,
+    shape,
+    dtype,
+    parallel,
+    progressbar,
+):
+    """
+    Forward model the three components of the magnetic vector
+
+    Parameters
+    ----------
+    coordinates : tuple of arrays
+        Tuple containing ``easting``, ``northing`` and ``upward`` of the
+        computation points as arrays, all defined on a Cartesian coordinate
+        system and in meters.
+    dipoles : tuple of arrays
+        Tuple of arrays containing the ``easting``, ``northing`` and ``upward``
+        locations of the dipoles defined on a Cartesian coordinate system. All
+        coordinates should be in meters.
+    magnetic_moments : tuple of arrays
+        Tuple containing the three arrays corresponding to the magnetic moment
+        components of each dipole in :math:`Am^2`. These arrays should
+        be provided in the following order: ``mag_moment_easting``,
+        ``mag_moment_northing``, ``mag_moment_upward``.
+    shape : tuple of int
+        Shape of the expected output arrays.
+    dtype : np.dtype
+        Data type of the expected output arrays.
+    parallel : bool
+        If True, the forward modelling will be run in parallel. If False, it
+        will be run in a single thread.
+    progressbar : bool
+        If True, a progress bar of the computation will be printed to standard
+        error (stderr). Requires :mod:`numba_progress` to be installed.
+
+    Returns
+    -------
+    magnetic_components : tuple of arrays
+        Tuple containing the three components of the magnetic vector in
+        :math:`nT`: ``b_e``, ``b_n``, ``b_u``.
+    """
+    # Decide which function should be used
+    if parallel:
+        jit_func = _jit_dipole_magnetic_field_cartesian_parallel
+    else:
+        jit_func = _jit_dipole_magnetic_field_cartesian_serial
+    # Run forward model
+    size = prod(shape)
+    b_e, b_n, b_u = tuple(np.zeros(size, dtype=dtype) for _ in range(3))
+    with initialize_progressbar(coordinates[0].size, progressbar) as progress_proxy:
+        jit_func(coordinates, dipoles, magnetic_moments, b_e, b_n, b_u, progress_proxy)
     # Convert to nT
     b_e *= 1e9
     b_n *= 1e9
     b_u *= 1e9
-    return b_e.reshape(cast.shape), b_n.reshape(cast.shape), b_u.reshape(cast.shape)
+    # Cast shape and form the tuple
+    result = tuple(component.reshape(shape) for component in (b_e, b_n, b_u))
+    return result
 
 
-def dipole_magnetic_component(
+def _dipole_magnetic_component(
     coordinates,
     dipoles,
     magnetic_moments,
-    component,
-    parallel=True,
-    dtype="float64",
-    disable_checks=False,
+    forward_func,
+    shape,
+    dtype,
+    parallel,
+    progressbar,
 ):
     """
-    Single magnetic field component of dipoles in Cartesian coordinates
-
-    .. important::
-
-        Use this function only if you need to compute a single component of the
-        magnetic field. Use :func:`harmonica.dipole_magnetic` to compute the
-        three components more efficiently.
+    Forward model a single component of the magnetic vector
 
     Parameters
     ----------
-    coordinates : list of arrays
-        List of arrays containing the ``easting``, ``northing`` and ``upward``
-        coordinates of the computation points defined on a Cartesian coordinate
-        system. All coordinates should be in meters.
-    dipoles : list of arrays
-        List of arrays containing the ``easting``, ``northing`` and ``upward``
+    coordinates : tuple of arrays
+        Tuple containing ``easting``, ``northing`` and ``upward`` of the
+        computation points as arrays, all defined on a Cartesian coordinate
+        system and in meters.
+    dipoles : tuple of arrays
+        Tuple of arrays containing the ``easting``, ``northing`` and ``upward``
         locations of the dipoles defined on a Cartesian coordinate system. All
         coordinates should be in meters.
-    magnetic_moments : 2d-array
-        Two dimensional array containing the magnetic moments of each dipole in
-        :math:`Am^2`. Each row contains the components of the magnetic moment
-        for each dipole in the following order: ``mag_moment_easting``,
+    magnetic_moments : tuple of arrays
+        Tuple containing the three arrays corresponding to the magnetic moment
+        components of each dipole in :math:`Am^2`. These arrays should
+        be provided in the following order: ``mag_moment_easting``,
         ``mag_moment_northing``, ``mag_moment_upward``.
-    component : str
-        Computed that will be computed. Available options are: ``"easting"``,
-        ``"northing"`` or ``"upward"``.
-    parallel : bool (optional)
-        If True the computations will run in parallel using Numba built-in
-        parallelization. If False, the forward model will run on a single core.
-        Might be useful to disable parallelization if the forward model is run
-        by an already parallelized workflow. Default to True.
-    dtype : data-type (optional)
-        Data type assigned to the resulting magnetic field. Default to
-        ``np.float64``.
-    disable_checks : bool (optional)
-        Flag that controls whether to perform a sanity check on the model.
-        Should be set to ``True`` only when it is certain that the input model
-        is valid and it does not need to be checked.
-        Default to ``False``.
+    forward_func : callable
+        Forward function to be used to compute the desired component of the
+        magnetic field. Choose one of :func:`choclo.dipole.magnetic_easting`,
+        :func:`choclo.dipole.magnetic_northing` or
+        :func:`choclo.dipole.magnetic_upward`.
+    shape : tuple of int
+        Shape of the expected output array.
+    dtype : np.dtype
+        Data type of the expected output array.
+    parallel : bool
+        If True, the forward modelling will be run in parallel. If False, it
+        will be run in a single thread.
+    progressbar : bool
+        If True, a progress bar of the computation will be printed to standard
+        error (stderr). Requires :mod:`numba_progress` to be installed.
 
     Returns
     -------
-    magnetic_field :  array
-        Array containing the magnetic field component in :math:`nT`.
+    magnetic_component : arrays
+        Array containing the desired magnetic component in :math:`nT`.
     """
-    # Figure out the shape and size of the output arrays
-    cast = np.broadcast(*coordinates[:3])
-    result = np.zeros(cast.size, dtype=dtype)
-    # Convert coordinates, dipoles and magnetic moments to arrays
-    coordinates = tuple(np.atleast_1d(i).ravel() for i in coordinates[:3])
-    dipoles = tuple(np.atleast_1d(i).ravel() for i in dipoles[:3])
-    magnetic_moments = np.atleast_2d(magnetic_moments)
-    # Choose forward modelling function based on the chosen component
-    forward_function = _get_magnetic_forward_function(component)
-    # Sanity checks
-    if not disable_checks:
-        _check_dipoles_and_magnetic_moments(dipoles, magnetic_moments)
-    # Compute the magnetic fields of the dipoles
+    # Decide which function should be used
     if parallel:
-        _jit_dipole_magnetic_component_cartesian_parallel(
-            coordinates, dipoles, magnetic_moments, result, forward_function
-        )
+        jit_func = _jit_dipole_magnetic_component_cartesian_parallel
     else:
-        _jit_dipole_magnetic_component_cartesian_serial(
-            coordinates, dipoles, magnetic_moments, result, forward_function
+        jit_func = _jit_dipole_magnetic_component_cartesian_serial
+    # Run computations
+    size = prod(shape)
+    result = np.zeros(size, dtype=dtype)
+    with initialize_progressbar(coordinates[0].size, progressbar) as progress_proxy:
+        jit_func(
+            coordinates,
+            dipoles,
+            magnetic_moments,
+            result,
+            forward_func,
+            progress_proxy,
         )
     # Convert to nT
     result *= 1e9
-    return result.reshape(cast.shape)
+    return result.reshape(shape)
 
 
 def _check_dipoles_and_magnetic_moments(dipoles, magnetic_moments):
     """
     Check if dipoles and magnetic moments have valid shape and size
     """
-    if magnetic_moments.shape[1] != 3:
+    if (size := len(magnetic_moments)) != 3:
         raise ValueError(
-            f"Invalid magnetic moments with '{magnetic_moments.shape[1]}' elements."
+            f"Invalid magnetic moments with '{size}' elements."
             " Magnetic moments vectors should have 3 components."
         )
-    if magnetic_moments.shape[0] != dipoles[0].size:
+    if magnetic_moments[0].size != dipoles[0].size:
         raise ValueError(
-            f"Number of elements in magnetic_moments ({magnetic_moments.shape[0]})"
+            f"Number of elements in magnetic_moments ({magnetic_moments[0].size})"
             f" mismatch the number of dipoles ({dipoles[0].size})."
         )
 
 
 def _jit_dipole_magnetic_field_cartesian(
-    coordinates, dipoles, magnetic_moments, b_e, b_n, b_u
+    coordinates, dipoles, magnetic_moments, b_e, b_n, b_u, progress_proxy
 ):
     """
     Compute the magnetic field components of dipoles in Cartesian coordinates
@@ -208,9 +310,18 @@ def _jit_dipole_magnetic_field_cartesian(
     b_u : 1d-array
         Array where the resulting values of the upward component of the
         magnetic field will be stored.
+    progress_proxy : :class:`numba_progress.ProgressBar` or None
+        Instance of :class:`numba_progress.ProgressBar` that gets updated after
+        each iteration on the observation points. Use None if no progress bar
+        is should be used.
     """
+    # Check if we need to update the progressbar on each iteration
+    update_progressbar = progress_proxy is not None
+    # Unpack coordinates and magnetic_moments
     easting, northing, upward = coordinates
     easting_p, northing_p, upward_p = dipoles
+    mag_e, mag_n, mag_u = magnetic_moments
+    # Iterate over computation points and dipoles
     for l in prange(easting.size):
         for m in range(easting_p.size):
             easting_comp, northing_comp, upward_comp = magnetic_field(
@@ -220,17 +331,20 @@ def _jit_dipole_magnetic_field_cartesian(
                 easting_p[m],
                 northing_p[m],
                 upward_p[m],
-                magnetic_moments[m, 0],
-                magnetic_moments[m, 1],
-                magnetic_moments[m, 2],
+                mag_e[m],
+                mag_n[m],
+                mag_u[m],
             )
             b_e[l] += easting_comp
             b_n[l] += northing_comp
             b_u[l] += upward_comp
+        # Update progress bar if called
+        if update_progressbar:
+            progress_proxy.update(1)
 
 
 def _jit_dipole_magnetic_component_cartesian(
-    coordinates, dipoles, magnetic_moments, result, forward_func
+    coordinates, dipoles, magnetic_moments, result, forward_func, progress_proxy
 ):
     """
     Compute a single magnetic component of dipoles in Cartesian coordinates
@@ -253,9 +367,18 @@ def _jit_dipole_magnetic_component_cartesian(
         magnetic field. Choose one of :func:`choclo.dipole.magnetic_easting`,
         :func:`choclo.dipole.magnetic_northing` or
         :func:`choclo.dipole.magnetic_upward`.
+    progress_proxy : :class:`numba_progress.ProgressBar` or None
+        Instance of :class:`numba_progress.ProgressBar` that gets updated after
+        each iteration on the observation points. Use None if no progress bar
+        is should be used.
     """
+    # Check if we need to update the progressbar on each iteration
+    update_progressbar = progress_proxy is not None
+    # Unpack coordinates and magnetic_moments
     easting, northing, upward = coordinates
     easting_p, northing_p, upward_p = dipoles
+    mag_e, mag_n, mag_u = magnetic_moments
+    # Iterate over computation points and dipoles
     for l in prange(easting.size):
         for m in range(easting_p.size):
             result[l] += forward_func(
@@ -265,34 +388,13 @@ def _jit_dipole_magnetic_component_cartesian(
                 easting_p[m],
                 northing_p[m],
                 upward_p[m],
-                magnetic_moments[m, 0],
-                magnetic_moments[m, 1],
-                magnetic_moments[m, 2],
+                mag_e[m],
+                mag_n[m],
+                mag_u[m],
             )
-
-
-def _get_magnetic_forward_function(component):
-    """
-    Returns the Choclo magnetic forward modelling function for the desired
-    component
-
-    Parameters
-    ----------
-    component : str
-        Magnetic field component.
-
-    Returns
-    -------
-    forward_function : callable
-        Forward modelling function for the desired component.
-    """
-    if component not in ("easting", "northing", "upward"):
-        raise ValueError(
-            f"Invalid component '{component}'. "
-            "It must be either 'easting', 'northing' or 'upward'."
-        )
-    functions = {"easting": magnetic_e, "northing": magnetic_n, "upward": magnetic_u}
-    return functions[component]
+        # Update progress bar if called
+        if update_progressbar:
+            progress_proxy.update(1)
 
 
 _jit_dipole_magnetic_field_cartesian_serial = jit(nopython=True)(


Test patch:
diff --git a/harmonica/tests/test_dipole.py b/harmonica/tests/test_dipole.py
index 8549518..7486adc 100644
--- a/harmonica/tests/test_dipole.py
+++ b/harmonica/tests/test_dipole.py
@@ -14,89 +14,93 @@ import pytest
 import verde as vd
 from choclo.dipole import magnetic_field
 
-from .. import dipole_magnetic, dipole_magnetic_component
+try:
+    from numba_progress import ProgressBar
+except ImportError:
+    ProgressBar = None
+
+from .. import dipole_magnetic
+from .._forward.dipole import VALID_FIELDS
 from .utils import run_only_with_numba
 
 
-def test_invalid_component():
-    "Check if passing an invalid component raises an error"
+def test_invalid_field():
+    "Check if passing an invalid field raises an error"
     coordinates = [0, 0, 0]
-    dipoles = [[-100, 100], [-100, 100], [-200, -100]]
-    magnetic_moments = [[1, 1, 2], [-1, -2, 3]]
-    with pytest.raises(ValueError, match="Invalid component"):
-        dipole_magnetic_component(
-            coordinates, dipoles, magnetic_moments, component="Not a valid field"
+    dipoles = ([-100, 100], [-100, 100], [-200, -100])
+    magnetic_moments = ([1, -1], [1, -2], [2, 3])
+    with pytest.raises(ValueError, match="Invalid field"):
+        dipole_magnetic(
+            coordinates, dipoles, magnetic_moments, field="Not a valid field"
         )
 
 
+@pytest.mark.parametrize("field", VALID_FIELDS)
+def test_disable_checks(field):
+    """Test if disabling checks works as expected"""
+    coordinates = [0, 0, 0]
+    dipoles = ([1, 1, 2], [-1, 0, 2], [-10, -2, -4])
+    magnetic_moments = (
+        [1, 4, 7],
+        [2, 5, 8],
+        [3, 6, 9],
+    )
+    dipole_magnetic(
+        coordinates, dipoles, magnetic_moments, disable_checks=True, field=field
+    )
+
+
+@pytest.mark.use_numba
+@pytest.mark.skipif(ProgressBar is None, reason="requires numba_progress")
+@pytest.mark.parametrize("field", VALID_FIELDS)
+def test_progress_bar(field):
+    """
+    Check if forward modelling results with and without progress bar match
+    """
+    dipoles = ([1, 1], [-1, 0], [-10, -2])
+    magnetic_moments = ([1.0, 1.0], [1.0, -1.0], [1.0, 5.0])
+    coordinates = vd.grid_coordinates(
+        region=(-100, 100, -100, 100), spacing=20, extra_coords=10
+    )
+    result_progress_true = dipole_magnetic(
+        coordinates, dipoles, magnetic_moments, field=field, progressbar=True
+    )
+    result_progress_false = dipole_magnetic(
+        coordinates, dipoles, magnetic_moments, field=field, progressbar=False
+    )
+    npt.assert_allclose(result_progress_true, result_progress_false)
+
+
 class TestInvalidMagneticMoments:
-    @pytest.mark.parametrize("component", ("all", "easting", "northing", "upward"))
-    def test_magnetic_moments_and_dipoles(self, component):
+    @pytest.mark.parametrize("field", ("b", "b_e", "b_n", "b_u"))
+    def test_magnetic_moments_and_dipoles(self, field):
         """
         Test error when dipoles and magnetic moments mismatch
         """
         coordinates = [0, 0, 0]
-        dipoles = [[1, 1], [-1, 0], [-10, -2]]
+        dipoles = ([1, 1], [-1, 0], [-10, -2])
         # Define invalid magnetic moments (more moments than dipoles)
-        magnetic_moments = [
-            [1, 2, 3],
-            [4, 5, 6],
-            [7, 8, 9],
-        ]
+        magnetic_moments = (
+            [1, 4, 7],
+            [2, 5, 8],
+            [3, 6, 9],
+        )
         msg = "Number of elements in magnetic_moments"
-        if component == "all":
-            with pytest.raises(ValueError, match=msg):
-                dipole_magnetic(coordinates, dipoles, magnetic_moments)
-        else:
-            with pytest.raises(ValueError, match=msg):
-                dipole_magnetic_component(
-                    coordinates, dipoles, magnetic_moments, component
-                )
+        with pytest.raises(ValueError, match=msg):
+            dipole_magnetic(coordinates, dipoles, magnetic_moments, field=field)
 
-    @pytest.mark.parametrize("component", ("all", "easting", "northing", "upward"))
-    def test_magnetic_moments_components(self, component):
+    @pytest.mark.parametrize("field", VALID_FIELDS)
+    def test_magnetic_moments_field(self, field):
         """
-        Test if error is raised when magnetic moments have != 3 components
+        Test if error is raised when magnetic moments have != 3 elements
         """
         coordinates = [0, 0, 0]
-        dipoles = [[1, 1], [-1, 0], [-10, -2]]
+        dipoles = ([1, 1], [-1, 0], [-10, -2])
         # Define invalid magnetic moments (more than three components)
-        magnetic_moments = [
-            [1, 2, 3, 5],
-            [4, 5, 6, 7],
-        ]
+        magnetic_moments = ([1, 4], [2, 5], [3, 6], [5, 7])
         msg = "Invalid magnetic moments with '4' elements."
-        if component == "all":
-            with pytest.raises(ValueError, match=msg):
-                dipole_magnetic(coordinates, dipoles, magnetic_moments)
-        else:
-            with pytest.raises(ValueError, match=msg):
-                dipole_magnetic_component(
-                    coordinates, dipoles, magnetic_moments, component
-                )
-
-    @pytest.mark.parametrize("component", ("all", "easting", "northing", "upward"))
-    def test_disable_checks(self, component):
-        """
-        Test disable checks on invalid magnetic moments
-        """
-        coordinates = [0, 0, 0]
-        dipoles = [[1, 1], [-1, 0], [-10, -2]]
-        # Define invalid magnetic moments (more than three components)
-        magnetic_moments = [
-            [1, 2, 3, 5],
-            [4, 5, 6, 7],
-        ]
-        if component == "all":
-            dipole_magnetic(coordinates, dipoles, magnetic_moments, disable_checks=True)
-        else:
-            dipole_magnetic_component(
-                coordinates,
-                dipoles,
-                magnetic_moments,
-                component,
-                disable_checks=True,
-            )
+        with pytest.raises(ValueError, match=msg):
+            dipole_magnetic(coordinates, dipoles, magnetic_moments, field=field)
 
 
 class TestSerialVsParallel:
@@ -104,37 +108,30 @@ class TestSerialVsParallel:
     Test serial vs parallel
     """
 
-    @pytest.mark.parametrize("component", ("all", "easting", "northing", "upward"))
-    def test_dipoles_parallel_vs_serial_no_numba(self, component):
+    @pytest.mark.parametrize("field", VALID_FIELDS)
+    def test_dipoles_parallel_vs_serial_no_numba(self, field):
         """
         Check results of parallelized and serials runs
         Run a small problem with Numba disable to count for test coverage.
         """
         coordinates = ([0, 10], [0, 10], [0, 10])
-        dipoles = [[-100, 0], [0, 100], [-20, -50]]
-        magnetic_moments = [
-            [1.0, 1.0, 1.0],
-            [1.0, -1.0, 5.0],
-        ]
-        if component == "all":
-            parallel = dipole_magnetic(
-                coordinates, dipoles, magnetic_moments, parallel=True
-            )
-            serial = dipole_magnetic(
-                coordinates, dipoles, magnetic_moments, parallel=False
-            )
-        else:
-            parallel = dipole_magnetic_component(
-                coordinates, dipoles, magnetic_moments, component, parallel=True
-            )
-            serial = dipole_magnetic_component(
-                coordinates, dipoles, magnetic_moments, component, parallel=False
-            )
+        dipoles = ([-100, 0], [0, 100], [-20, -50])
+        magnetic_moments = (
+            [1.0, 1.0],
+            [1.0, -1.0],
+            [1.0, 5.0],
+        )
+        parallel = dipole_magnetic(
+            coordinates, dipoles, magnetic_moments, field=field, parallel=True
+        )
+        serial = dipole_magnetic(
+            coordinates, dipoles, magnetic_moments, field=field, parallel=False
+        )
         npt.assert_allclose(parallel, serial)
 
     @run_only_with_numba
-    @pytest.mark.parametrize("component", ("all", "easting", "northing", "upward"))
-    def test_dipoles_parallel_vs_serial(self, component):
+    @pytest.mark.parametrize("field", VALID_FIELDS)
+    def test_dipoles_parallel_vs_serial(self, field):
         """
         Check results of parallelized and serials runs
         Run a large problem only with Numba enabled.
@@ -142,25 +139,18 @@ class TestSerialVsParallel:
         coordinates = vd.grid_coordinates(
             region=(-100, 100, -100, 100), spacing=20, extra_coords=10
         )
-        dipoles = [[-100, 0], [0, 100], [-20, -50]]
+        dipoles = ([-100, 0], [0, 100], [-20, -50])
         magnetic_moments = [
-            [1.0, 1.0, 1.0],
-            [1.0, -1.0, 5.0],
+            [1.0, 1.0],
+            [1.0, -1.0],
+            [1.0, 5.0],
         ]
-        if component == "all":
-            parallel = dipole_magnetic(
-                coordinates, dipoles, magnetic_moments, parallel=True
-            )
-            serial = dipole_magnetic(
-                coordinates, dipoles, magnetic_moments, parallel=False
-            )
-        else:
-            parallel = dipole_magnetic_component(
-                coordinates, dipoles, magnetic_moments, component, parallel=True
-            )
-            serial = dipole_magnetic_component(
-                coordinates, dipoles, magnetic_moments, component, parallel=False
-            )
+        parallel = dipole_magnetic(
+            coordinates, dipoles, magnetic_moments, field=field, parallel=True
+        )
+        serial = dipole_magnetic(
+            coordinates, dipoles, magnetic_moments, field=field, parallel=False
+        )
         npt.assert_allclose(parallel, serial)
 
 
@@ -186,14 +176,10 @@ class TestAgainstChoclo:
         """
         Return sample magnetic moment vectors for the dipoles
         """
-        magnetic_moments = np.array(
-            [
-                [1.0, 1.0, 1],
-                [-1.0, -0.5, 2.0],
-                [3, -1.0, -3.0],
-                [-1.5, 2.0, -2.0],
-            ],
-            dtype=np.float64,
+        magnetic_moments = (
+            np.array([1.0, -1.0, 3.0, -1.5]),
+            np.array([1.0, -0.5, -1.0, 2.0]),
+            np.array([1.0, 2.0, -3.0, -2.0]),
         )
         return magnetic_moments
 
@@ -230,7 +216,9 @@ class TestAgainstChoclo:
                     easting_p[j],
                     northing_p[j],
                     upward_p[j],
-                    *sample_magnetic_moments[j, :],
+                    sample_magnetic_moments[0][j],
+                    sample_magnetic_moments[1][j],
+                    sample_magnetic_moments[2][j],
                 )
                 expected_magnetic_e[i] += b_e
                 expected_magnetic_n[i] += b_n
@@ -241,15 +229,15 @@ class TestAgainstChoclo:
         expected_magnetic_u *= 1e9
         # Compare with harmonica results
         b_e, b_n, b_u = dipole_magnetic(
-            sample_coordinates, sample_dipoles, sample_magnetic_moments
+            sample_coordinates, sample_dipoles, sample_magnetic_moments, field="b"
         )
         npt.assert_allclose(b_e, expected_magnetic_e)
         npt.assert_allclose(b_n, expected_magnetic_n)
         npt.assert_allclose(b_u, expected_magnetic_u)
 
-    @pytest.mark.parametrize("component", ("easting", "northing", "upward"))
+    @pytest.mark.parametrize("field", ("b_e", "b_n", "b_u"))
     def test_component_against_choclo(
-        self, sample_coordinates, sample_dipoles, sample_magnetic_moments, component
+        self, sample_coordinates, sample_dipoles, sample_magnetic_moments, field
     ):
         """
         Test dipole_magnetic_component against raw Choclo runs
@@ -260,7 +248,7 @@ class TestAgainstChoclo:
         n_coords = easting.size
         n_dipoles = sample_dipoles[0].size
         expected_result = np.zeros(n_coords, dtype=np.float64)
-        forward_func = getattr(choclo.dipole, f"magnetic_{component[0]}")
+        forward_func = getattr(choclo.dipole, f"magnetic_{field[-1]}")
         for i in range(n_coords):
             for j in range(n_dipoles):
                 expected_result[i] += forward_func(
@@ -270,15 +258,17 @@ class TestAgainstChoclo:
                     easting_p[j],
                     northing_p[j],
                     upward_p[j],
-                    *sample_magnetic_moments[j, :],
+                    sample_magnetic_moments[0][j],
+                    sample_magnetic_moments[1][j],
+                    sample_magnetic_moments[2][j],
                 )
         # Convert to nT
         expected_result *= 1e9
         # Compare with harmonica results
-        result = dipole_magnetic_component(
+        result = dipole_magnetic(
             sample_coordinates,
             sample_dipoles,
             sample_magnetic_moments,
-            component,
+            field=field,
         )
         npt.assert_allclose(result, expected_result)