Repository: simpeg/simpeg
Pull request #1527 (2024-08-28 21:52:17)

Implement gravity equivalent sources with Choclo as engine
#### Summary
<!-- Add a summary of this Pull Request -->

Implement gravity equivalent sources using Choclo as the engine. Allow gravity equivalent sources to take `engine="choclo"`. Implement new Numba-based forward and sensitivity matrix functions that use Choclo's forward modelling functions for a prism. These new functions can take a 2D mesh and the arrays for the top and bottom boundaries for each cell and compute the forward or build the sensitivity matrix by iterating over each prism in the layer. Add extended tests for the new implementation that also extend the tests for the geoana-based implementation.

#### PR Checklist
* [ ] If this is a work in progress PR, set as a Draft PR
* [ ] Linted my code according to the [style guides](https://docs.simpeg.xyz/content/getting_started/contributing/code-style.html).
* [ ] Added [tests](https://docs.simpeg.xyz/content/getting_started/practices.html#testing) to verify changes to the code.
* [ ] Added necessary documentation to any new functions/classes following the
      expect [style](https://docs.simpeg.xyz/content/getting_started/practices.html#documentation).
* [ ] Marked as ready for review (if this is was a draft PR), and converted
      to a Pull Request
* [ ] Tagged ``@simpeg/simpeg-developers`` when ready for review.

#### Reference issue
<!--Example: write "Closes #NNNN" to automatically close that issue on merge.-->

Part of solution to #1373

#### What does this implement/fix?
<!--Please explain your changes.-->

#### Additional information
<!--Any additional information you think is important.-->


**TODO:**

- [x] Make magnetic equivalent sources to raise an error if `engine="choclo"`. We can implement that one in a separate PR.


codecov[bot]: ## [Codecov](https://app.codecov.io/gh/simpeg/simpeg/pull/1527?dropdown=coverage&src=pr&el=h1&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=simpeg) Report
Attention: Patch coverage is `0%` with `162 lines` in your changes missing coverage. Please review.
> Project coverage is 23.51%. Comparing base [(`5e4b583`)](https://app.codecov.io/gh/simpeg/simpeg/commit/5e4b583be3e5dafa3622dcff59894a2a02e5a304?dropdown=coverage&el=desc&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=simpeg) to head [(`8941796`)](https://app.codecov.io/gh/simpeg/simpeg/commit/894179605f8c1e645ac724a9c11eecb7e28ebafe?dropdown=coverage&el=desc&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=simpeg).

| [Files with missing lines](https://app.codecov.io/gh/simpeg/simpeg/pull/1527?dropdown=coverage&src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=simpeg) | Patch % | Lines |
|---|---|---|
| [tests/pf/test\_equivalent\_sources.py](https://app.codecov.io/gh/simpeg/simpeg/pull/1527?src=pr&el=tree&filepath=tests%2Fpf%2Ftest_equivalent_sources.py&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=simpeg#diff-dGVzdHMvcGYvdGVzdF9lcXVpdmFsZW50X3NvdXJjZXMucHk=) | 0.00% | [77 Missing :warning: ](https://app.codecov.io/gh/simpeg/simpeg/pull/1527?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=simpeg) |
| [simpeg/potential\_fields/gravity/simulation.py](https://app.codecov.io/gh/simpeg/simpeg/pull/1527?src=pr&el=tree&filepath=simpeg%2Fpotential_fields%2Fgravity%2Fsimulation.py&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=simpeg#diff-c2ltcGVnL3BvdGVudGlhbF9maWVsZHMvZ3Jhdml0eS9zaW11bGF0aW9uLnB5) | 0.00% | [59 Missing :warning: ](https://app.codecov.io/gh/simpeg/simpeg/pull/1527?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=simpeg) |
| [...impeg/potential\_fields/gravity/\_numba\_functions.py](https://app.codecov.io/gh/simpeg/simpeg/pull/1527?src=pr&el=tree&filepath=simpeg%2Fpotential_fields%2Fgravity%2F_numba_functions.py&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=simpeg#diff-c2ltcGVnL3BvdGVudGlhbF9maWVsZHMvZ3Jhdml0eS9fbnVtYmFfZnVuY3Rpb25zLnB5) | 0.00% | [16 Missing :warning: ](https://app.codecov.io/gh/simpeg/simpeg/pull/1527?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=simpeg) |
| [simpeg/potential\_fields/base.py](https://app.codecov.io/gh/simpeg/simpeg/pull/1527?src=pr&el=tree&filepath=simpeg%2Fpotential_fields%2Fbase.py&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=simpeg#diff-c2ltcGVnL3BvdGVudGlhbF9maWVsZHMvYmFzZS5weQ==) | 0.00% | [10 Missing :warning: ](https://app.codecov.io/gh/simpeg/simpeg/pull/1527?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=simpeg) |

> :exclamation:  There is a different number of reports uploaded between BASE (5e4b583) and HEAD (8941796). Click for more details.
> 
> <details><summary>HEAD has 19 uploads less than BASE</summary>
>
>| Flag | BASE (5e4b583) | HEAD (8941796) |
>|------|------|------|
>||20|1|
></details>

<details><summary>Additional details and impacted files</summary>


```diff
@@             Coverage Diff             @@
##             main    #1527       +/-   ##
===========================================
- Coverage   86.23%   23.51%   -62.72%     
===========================================
  Files         396      298       -98     
  Lines       50366    39802    -10564     
  Branches     5309     4860      -449     
===========================================
- Hits        43431     9361    -34070     
- Misses       5483    29752    +24269     
+ Partials     1452      689      -763     
```

</details>

[:umbrella: View full report in Codecov by Sentry](https://app.codecov.io/gh/simpeg/simpeg/pull/1527?dropdown=coverage&src=pr&el=continue&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=simpeg).   
:loudspeaker: Have feedback on the report? [Share it here](https://about.codecov.io/codecov-pr-comment-feedback/?utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=simpeg).

santisoler: @thibaut-kobold I think this implementation of the gravity equivalent sources that can use choclo as the engine is quite ready. I wrote extensive tests for it and I quite like it. Would you like to give it a try and provide feedback?

Patch:
diff --git a/simpeg/potential_fields/base.py b/simpeg/potential_fields/base.py
index 33bc55694..9f4e4e3f0 100644
--- a/simpeg/potential_fields/base.py
+++ b/simpeg/potential_fields/base.py
@@ -395,7 +395,7 @@ class BasePFSimulation(LinearSimulation):
             nodes = self.mesh.nodes
         else:
             raise TypeError(f"Invalid mesh of type {self.mesh.__class__.__name__}.")
-        # Get original cell_nodes but only for active cells
+        # Get original cell_nodes
         cell_nodes = self.mesh.cell_nodes
         # If all cells in the mesh are active, return nodes and cell_nodes
         if self.nC == self.mesh.n_cells:
@@ -437,6 +437,7 @@ class BaseEquivalentSourceLayerSimulation(BasePFSimulation):
     """
 
     def __init__(self, mesh, cell_z_top, cell_z_bottom, **kwargs):
+
         if mesh.dim != 2:
             raise AttributeError("Mesh to equivalent source layer must be 2D.")
 
@@ -454,6 +455,8 @@ class BaseEquivalentSourceLayerSimulation(BasePFSimulation):
                 "cells, and match the number of active cells.",
             )
 
+        self._cell_z_top, self._cell_z_bottom = cell_z_top, cell_z_bottom
+
         all_nodes = self._nodes[self._unique_inv]
         all_nodes = [
             np.c_[all_nodes[0], cell_z_bottom],
@@ -468,6 +471,32 @@ class BaseEquivalentSourceLayerSimulation(BasePFSimulation):
         self._nodes = np.stack(all_nodes, axis=0)
         self._unique_inv = None
 
+    @property
+    def cell_z_top(self) -> np.ndarray:
+        """
+        Elevations for the top face of all cells in the layer.
+        """
+        return self._cell_z_top
+
+    @property
+    def cell_z_bottom(self) -> np.ndarray:
+        """
+        Elevations for the bottom face of all cells in the layer.
+        """
+        return self._cell_z_bottom
+
+    def _check_engine_and_mesh_dimensions(self):
+        """
+        Check dimensions of the mesh
+
+        Overwrite the parent's method: the equivalent sources class needs 2D
+        meshes, while the potential field simulations work only with 3D meshes.
+
+        This check will run for any given engine.
+        """
+        if self.mesh.dim != 2:
+            raise AttributeError("Mesh to equivalent source layer must be 2D.")
+
 
 def progress(iteration, prog, final):
     """Progress (% complete) for constructing sensitivity matrix.
diff --git a/simpeg/potential_fields/gravity/_numba_functions.py b/simpeg/potential_fields/gravity/_numba_functions.py
index 3eb6ae9bf..fe2e69e20 100644
--- a/simpeg/potential_fields/gravity/_numba_functions.py
+++ b/simpeg/potential_fields/gravity/_numba_functions.py
@@ -192,8 +192,161 @@ def _evaluate_kernel(
     return kernel_func(dx, dy, dz, distance)
 
 
+def _forward_gravity_2d_mesh(
+    receivers,
+    cells_bounds,
+    top,
+    bottom,
+    densities,
+    fields,
+    forward_func,
+    constant_factor,
+):
+    """
+    Forward gravity fields of 2D meshes.
+
+    This function is designed to be used with equivalent sources, where the
+    mesh is a 2D mesh (prism layer). The top and bottom boundaries of each cell
+    are passed through the ``top`` and ``bottom`` arrays.
+
+    This function should be used with a `numba.jit` decorator, for example:
+
+    .. code::
+
+        from numba import jit
+
+        jit_function = jit(nopython=True, parallel=True)(_forward_gravity_2d_mesh)
+
+    Parameters
+    ----------
+    receivers : (n_receivers, 3) numpy.ndarray
+        Array with the locations of the receivers
+    cells_bounds : (n_active_cells, 4) numpy.ndarray
+        Array with the bounds of each active cell in the 2D mesh. For each row, the
+        bounds should be passed in the following order: ``x_min``, ``x_max``,
+        ``y_min``, ``y_max``.
+    top : (n_active_cells) np.ndarray
+        Array with the top boundaries of each active cell in the 2D mesh.
+    bottom : (n_active_cells) np.ndarray
+        Array with the bottom boundaries of each active cell in the 2D mesh.
+    densities : (n_active_cells) numpy.ndarray
+        Array with densities of each active cell in the mesh.
+    fields : (n_receivers) numpy.ndarray
+        Array full of zeros where the gravity fields on each receiver will be
+        stored. This could be a preallocated array or a slice of it.
+    forward_func : callable
+        Forward function that will be evaluated on each node of the mesh. Choose
+        one of the forward functions in ``choclo.prism``.
+    constant_factor : float
+        Constant factor that will be used to multiply each element of the
+        ``fields`` array.
+
+    Notes
+    -----
+    The constant factor is applied here to each element of fields because
+    it's more efficient than doing it afterwards: it would require to
+    index the elements that corresponds to each component.
+    """
+    n_receivers = receivers.shape[0]
+    n_cells = cells_bounds.shape[0]
+    # Forward model the gravity field of each cell on each receiver location
+    for i in prange(n_receivers):
+        for j in range(n_cells):
+            fields[i] += constant_factor * forward_func(
+                receivers[i, 0],
+                receivers[i, 1],
+                receivers[i, 2],
+                cells_bounds[j, 0],
+                cells_bounds[j, 1],
+                cells_bounds[j, 2],
+                cells_bounds[j, 3],
+                bottom[j],
+                top[j],
+                densities[j],
+            )
+
+
+def _sensitivity_gravity_2d_mesh(
+    receivers,
+    cells_bounds,
+    top,
+    bottom,
+    sensitivity_matrix,
+    forward_func,
+    constant_factor,
+):
+    """
+    Fill the sensitivity matrix
+
+    This function should be used with a `numba.jit` decorator, for example:
+
+    .. code::
+
+        from numba import jit
+
+        jit_function = jit(nopython=True, parallel=True)(_sensitivity_gravity_2d_mesh)
+
+    Parameters
+    ----------
+    receivers : (n_receivers, 3) numpy.ndarray
+        Array with the locations of the receivers
+    cells_bounds : (n_active_cells, 4) numpy.ndarray
+        Array with the bounds of each active cell in the 2D mesh. For each row, the
+        bounds should be passed in the following order: ``x_min``, ``x_max``,
+        ``y_min``, ``y_max``.
+    top : (n_active_cells) np.ndarray
+        Array with the top boundaries of each active cell in the 2D mesh.
+    bottom : (n_active_cells) np.ndarray
+        Array with the bottom boundaries of each active cell in the 2D mesh.
+    sensitivity_matrix : (n_receivers, n_active_nodes) array
+        Empty 2d array where the sensitivity matrix elements will be filled.
+        This could be a preallocated empty array or a slice of it.
+    forward_func : callable
+        Forward function that will be evaluated on each node of the mesh. Choose
+        one of the forward functions in ``choclo.prism``.
+    constant_factor : float
+        Constant factor that will be used to multiply each element of the
+        sensitivity matrix.
+
+    Notes
+    -----
+    The constant factor is applied here to each row of the sensitivity matrix
+    because it's more efficient than doing it afterwards: it would require to
+    index the rows that corresponds to each component.
+    """
+    n_receivers = receivers.shape[0]
+    n_cells = cells_bounds.shape[0]
+    # Evaluate kernel function on each node, for each receiver location
+    for i in prange(n_receivers):
+        for j in range(n_cells):
+            sensitivity_matrix[i, j] = constant_factor * forward_func(
+                receivers[i, 0],
+                receivers[i, 1],
+                receivers[i, 2],
+                cells_bounds[j, 0],
+                cells_bounds[j, 1],
+                cells_bounds[j, 2],
+                cells_bounds[j, 3],
+                bottom[j],
+                top[j],
+                1.0,  # use unitary density to get sensitivities
+            )
+
+
 # Define decorated versions of these functions
 _sensitivity_gravity_parallel = jit(nopython=True, parallel=True)(_sensitivity_gravity)
 _sensitivity_gravity_serial = jit(nopython=True, parallel=False)(_sensitivity_gravity)
 _forward_gravity_parallel = jit(nopython=True, parallel=True)(_forward_gravity)
 _forward_gravity_serial = jit(nopython=True, parallel=False)(_forward_gravity)
+_forward_gravity_2d_mesh_serial = jit(nopython=True, parallel=False)(
+    _forward_gravity_2d_mesh
+)
+_forward_gravity_2d_mesh_parallel = jit(nopython=True, parallel=True)(
+    _forward_gravity_2d_mesh
+)
+_sensitivity_gravity_2d_mesh_serial = jit(nopython=True, parallel=False)(
+    _sensitivity_gravity_2d_mesh
+)
+_sensitivity_gravity_2d_mesh_parallel = jit(nopython=True, parallel=True)(
+    _sensitivity_gravity_2d_mesh
+)
diff --git a/simpeg/potential_fields/gravity/simulation.py b/simpeg/potential_fields/gravity/simulation.py
index 1d1f667e2..2d68c7c0c 100644
--- a/simpeg/potential_fields/gravity/simulation.py
+++ b/simpeg/potential_fields/gravity/simulation.py
@@ -1,6 +1,8 @@
+from __future__ import annotations
 import warnings
 import numpy as np
 import scipy.constants as constants
+from discretize import TensorMesh, TreeMesh
 from geoana.kernels import prism_fz, prism_fzx, prism_fzy, prism_fzz
 from scipy.constants import G as NewtG
 
@@ -16,6 +18,10 @@ from ._numba_functions import (
     _sensitivity_gravity_parallel,
     _forward_gravity_serial,
     _forward_gravity_parallel,
+    _forward_gravity_2d_mesh_serial,
+    _forward_gravity_2d_mesh_parallel,
+    _sensitivity_gravity_2d_mesh_serial,
+    _sensitivity_gravity_2d_mesh_parallel,
 )
 
 if choclo is not None:
@@ -30,6 +36,48 @@ if choclo is not None:
         )
         return result
 
+    @jit(nopython=True)
+    def gravity_uv(
+        easting,
+        northing,
+        upward,
+        prism_west,
+        prism_east,
+        prism_south,
+        prism_north,
+        prism_bottom,
+        prism_top,
+        density,
+    ):
+        """Forward model the Guv gradiometry component."""
+        result = 0.5 * (
+            choclo.prism.gravity_nn(
+                easting,
+                northing,
+                upward,
+                prism_west,
+                prism_east,
+                prism_south,
+                prism_north,
+                prism_bottom,
+                prism_top,
+                density,
+            )
+            - choclo.prism.gravity_ee(
+                easting,
+                northing,
+                upward,
+                prism_west,
+                prism_east,
+                prism_south,
+                prism_north,
+                prism_bottom,
+                prism_top,
+                density,
+            )
+        )
+        return result
+
     CHOCLO_KERNELS = {
         "gx": choclo.prism.kernel_e,
         "gy": choclo.prism.kernel_n,
@@ -43,6 +91,19 @@ if choclo is not None:
         "guv": kernel_uv,
     }
 
+    CHOCLO_FORWARD_FUNCS = {
+        "gx": choclo.prism.gravity_e,
+        "gy": choclo.prism.gravity_n,
+        "gz": choclo.prism.gravity_u,
+        "gxx": choclo.prism.gravity_ee,
+        "gyy": choclo.prism.gravity_nn,
+        "gzz": choclo.prism.gravity_uu,
+        "gxy": choclo.prism.gravity_en,
+        "gxz": choclo.prism.gravity_eu,
+        "gyz": choclo.prism.gravity_nu,
+        "guv": gravity_uv,
+    }
+
 
 def _get_conversion_factor(component):
     """
@@ -57,6 +118,42 @@ def _get_conversion_factor(component):
     return conversion_factor
 
 
+def _get_cell_bounds(mesh: TensorMesh | TreeMesh):
+    """
+    Bounds of each cell in a 2D TensorMesh or TreeMesh.
+
+    The bounds are defined as ``x_min``, ``x_max``, ``y_min``, ``y_max``.
+
+    ..note:
+
+        This private function could be replaced by calling some `cell_bounds`
+        method of the meshes directly from discretize.
+
+    Parameters
+    ----------
+    mesh : discretize.TensorMesh or discretize.TreeMesh
+        A 2D mesh.
+
+    Returns
+    -------
+    bounds : (n_cells, 4) array
+        Array with the bounds of each cell in the mesh.
+    """
+    if not isinstance(mesh, (TensorMesh, TreeMesh)):
+        raise TypeError(f"Invalid mesh of type {mesh.__class__}.")
+    if mesh.dim != 2:
+        raise TypeError(
+            f"Invalid mesh with '{mesh.dim}' dimensions. Only 2D meshes can be passed."
+        )
+    centers, widths = mesh.cell_centers, mesh.h_gridded
+    xmin = centers[:, 0] - widths[:, 0] / 2
+    xmax = centers[:, 0] + widths[:, 0] / 2
+    ymin = centers[:, 1] - widths[:, 1] / 2
+    ymax = centers[:, 1] + widths[:, 1] / 2
+    bounds = np.hstack(tuple(v[:, np.newaxis] for v in (xmin, xmax, ymin, ymax)))
+    return bounds
+
+
 class Simulation3DIntegral(BasePFSimulation):
     """
     Gravity simulation in integral form.
@@ -432,13 +529,137 @@ class SimulationEquivalentSourceLayer(
     mesh : discretize.BaseMesh
         A 2D tensor or tree mesh defining discretization along the x and y directions
     cell_z_top : numpy.ndarray or float
-        Define the elevations for the top face of all cells in the layer. If an array it should be the same size as
-        the active cell set.
+        Define the elevations for the top face of all cells in the layer.
+        If an array it should be the same size as the active cell set.
     cell_z_bottom : numpy.ndarray or float
-        Define the elevations for the bottom face of all cells in the layer. If an array it should be the same size as
-        the active cell set.
+        Define the elevations for the bottom face of all cells in the layer.
+        If an array it should be the same size as the active cell set.
+    engine : {"geoana", "choclo"}, optional
+        Choose which engine should be used to run the forward model.
+    numba_parallel : bool, optional
+        If True, the simulation will run in parallel. If False, it will
+        run in serial. If ``engine`` is not ``"choclo"`` this argument will be
+        ignored.
     """
 
+    def __init__(
+        self,
+        mesh,
+        cell_z_top,
+        cell_z_bottom,
+        engine="geoana",
+        numba_parallel=True,
+        **kwargs,
+    ):
+        super().__init__(
+            mesh,
+            cell_z_top,
+            cell_z_bottom,
+            engine=engine,
+            numba_parallel=numba_parallel,
+            **kwargs,
+        )
+
+        if self.engine == "choclo":
+            if self.numba_parallel:
+                self._sensitivity_gravity = _sensitivity_gravity_2d_mesh_parallel
+                self._forward_gravity = _forward_gravity_2d_mesh_parallel
+            else:
+                self._sensitivity_gravity = _sensitivity_gravity_2d_mesh_serial
+                self._forward_gravity = _forward_gravity_2d_mesh_serial
+
+    def _forward(self, densities):
+        """
+        Forward model the fields of active cells in the mesh on receivers.
+
+        Parameters
+        ----------
+        densities : (n_active_cells) numpy.ndarray
+            Array containing the densities of the active cells in the mesh, in
+            g/cc.
+
+        Returns
+        -------
+        (nD,) numpy.ndarray
+            Always return a ``np.float64`` array.
+        """
+        # Get cells in the 2D mesh
+        cells_bounds = _get_cell_bounds(self.mesh)
+        # Keep only active cells
+        cells_bounds_active = cells_bounds[self.active_cells]
+        # Allocate fields array
+        fields = np.zeros(self.survey.nD, dtype=self.sensitivity_dtype)
+        # Compute fields
+        index_offset = 0
+        for components, receivers in self._get_components_and_receivers():
+            n_components = len(components)
+            n_elements = n_components * receivers.shape[0]
+            for i, component in enumerate(components):
+                forward_func = CHOCLO_FORWARD_FUNCS[component]
+                conversion_factor = _get_conversion_factor(component)
+                vector_slice = slice(
+                    index_offset + i, index_offset + n_elements, n_components
+                )
+                self._forward_gravity(
+                    receivers,
+                    cells_bounds_active,
+                    self.cell_z_top,
+                    self.cell_z_bottom,
+                    densities,
+                    fields[vector_slice],
+                    forward_func,
+                    conversion_factor,
+                )
+            index_offset += n_elements
+        return fields
+
+    def _sensitivity_matrix(self):
+        """
+        Compute the sensitivity matrix G
+
+        Returns
+        -------
+        (nD, n_active_cells) numpy.ndarray
+        """
+        # Get cells in the 2D mesh
+        cells_bounds = _get_cell_bounds(self.mesh)
+        # Keep only active cells
+        cells_bounds_active = cells_bounds[self.active_cells]
+        # Allocate sensitivity matrix
+        shape = (self.survey.nD, self.nC)
+        if self.store_sensitivities == "disk":
+            sensitivity_matrix = np.memmap(
+                self.sensitivity_path,
+                shape=shape,
+                dtype=self.sensitivity_dtype,
+                order="C",  # it's more efficient to write in row major
+                mode="w+",
+            )
+        else:
+            sensitivity_matrix = np.empty(shape, dtype=self.sensitivity_dtype)
+        # Start filling the sensitivity matrix
+        index_offset = 0
+        for components, receivers in self._get_components_and_receivers():
+            n_components = len(components)
+            n_rows = n_components * receivers.shape[0]
+            for i, component in enumerate(components):
+                forward_func = CHOCLO_FORWARD_FUNCS[component]
+                conversion_factor = _get_conversion_factor(component)
+                matrix_slice = slice(
+                    index_offset + i, index_offset + n_rows, n_components
+                )
+                self._sensitivity_gravity(
+                    receivers,
+                    cells_bounds_active,
+                    self.cell_z_top,
+                    self.cell_z_bottom,
+                    sensitivity_matrix[matrix_slice, :],
+                    forward_func,
+                    conversion_factor,
+                )
+            index_offset += n_rows
+        return sensitivity_matrix
+
 
 class Simulation3DDifferential(BasePDESimulation):
     r"""Finite volume simulation class for gravity.
diff --git a/simpeg/potential_fields/magnetics/simulation.py b/simpeg/potential_fields/magnetics/simulation.py
index 91ec72a7e..64cf9eba1 100644
--- a/simpeg/potential_fields/magnetics/simulation.py
+++ b/simpeg/potential_fields/magnetics/simulation.py
@@ -792,6 +792,29 @@ class SimulationEquivalentSourceLayer(
 
     """
 
+    def __init__(
+        self,
+        mesh,
+        cell_z_top,
+        cell_z_bottom,
+        engine="geoana",
+        numba_parallel=True,
+        **kwargs,
+    ):
+        if engine == "choclo":
+            raise NotImplementedError(
+                "Magnetic equivalent sources with choclo as engine has not been"
+                " implemented yet. Use 'geoana' instead."
+            )
+        super().__init__(
+            mesh,
+            cell_z_top,
+            cell_z_bottom,
+            engine=engine,
+            numba_parallel=numba_parallel,
+            **kwargs,
+        )
+
 
 class Simulation3DDifferential(BaseMagneticPDESimulation):
     """


Test patch:
diff --git a/tests/pf/test_equivalent_sources.py b/tests/pf/test_equivalent_sources.py
new file mode 100644
index 000000000..141476866
--- /dev/null
+++ b/tests/pf/test_equivalent_sources.py
@@ -0,0 +1,579 @@
+import pytest
+
+import numpy as np
+from discretize import TensorMesh
+from discretize.utils import mesh_builder_xyz, mkvc
+
+import simpeg
+from simpeg.optimization import ProjectedGNCG
+from simpeg.potential_fields import gravity, magnetics, base
+
+COMPONENTS = ["gx", "gy", "gz", "gxx", "gyy", "gzz", "gxy", "gxz", "gyz", "guv"]
+
+
+def create_grid(x_range, y_range, size):
+    """Create a 2D horizontal coordinates grid."""
+    x_start, x_end = x_range
+    y_start, y_end = y_range
+    x, y = np.meshgrid(
+        np.linspace(x_start, x_end, size), np.linspace(y_start, y_end, size)
+    )
+    return x, y
+
+
+@pytest.fixture()
+def mesh_params():
+    """Parameters for building the sample meshes."""
+    h = [5, 5]
+    padding_distances = np.ones((2, 2)) * 50
+    return h, padding_distances
+
+
+@pytest.fixture()
+def tensor_mesh(mesh_params, coordinates):
+    """Sample 2D tensor mesh to use with equivalent sources."""
+    mesh_type = "tensor"
+    h, padding_distance = mesh_params
+    mesh = mesh_builder_xyz(
+        coordinates[:, :2], h, padding_distance=padding_distance, mesh_type=mesh_type
+    )
+    return mesh
+
+
+@pytest.fixture()
+def tree_mesh(mesh_params, coordinates):
+    """Sample 2D tree mesh to use with equivalent sources."""
+    mesh_type = "tree"
+    h, padding_distance = mesh_params
+    mesh = mesh_builder_xyz(
+        coordinates[:, :2], h, padding_distance=padding_distance, mesh_type=mesh_type
+    )
+    mesh.refine_points(coordinates[:, :2], padding_cells_by_level=[2, 4])
+    return mesh
+
+
+@pytest.fixture(params=["tensor", "tree"])
+def mesh(tensor_mesh, tree_mesh, request):
+    """Sample 2D mesh to use with equivalent sources."""
+    mesh_type = request.param
+    if mesh_type == "tree":
+        return tree_mesh
+    elif mesh_type == "tensor":
+        return tensor_mesh
+    else:
+        raise ValueError(f"Invalid mesh type: '{mesh_type}'")
+
+
+@pytest.fixture
+def mesh_top():
+    """Top boundary of the mesh."""
+    return -20.0
+
+
+@pytest.fixture
+def mesh_bottom():
+    """Bottom boundary of the mesh."""
+    return -50.0
+
+
+@pytest.fixture
+def coordinates():
+    """Synthetic observation points grid."""
+    x, y = create_grid(x_range=(-50, 50), y_range=(-50, 50), size=11)
+    z = np.full_like(x, fill_value=5.0)
+    return np.c_[mkvc(x), mkvc(y), mkvc(z)]
+
+
+def get_block_model(mesh, phys_property: float):
+    """Build a block model."""
+    model = simpeg.utils.model_builder.add_block(
+        mesh.cell_centers,
+        np.zeros(mesh.n_cells),
+        np.r_[-20, -20],
+        np.r_[20, 20],
+        phys_property,
+    )
+    return model
+
+
+def get_mapping(mesh):
+    """Get an identity map for the given mesh."""
+    return simpeg.maps.IdentityMap(nP=mesh.n_cells)
+
+
+@pytest.fixture
+def gravity_survey(coordinates):
+    """
+    Sample survey for the gravity equivalent sources.
+    """
+    return build_gravity_survey(coordinates, components="gz")
+
+
+def build_gravity_survey(coordinates, components):
+    """
+    Build a gravity survey.
+    """
+    receivers = gravity.Point(coordinates, components=components)
+    source_field = gravity.SourceField([receivers])
+    survey = gravity.Survey(source_field)
+    return survey
+
+
+class Test3DMeshError:
+    """
+    Test if error is raised after passing a 3D mesh to equivalent sources.
+    """
+
+    @pytest.fixture
+    def mesh_3d(self):
+        mesh = TensorMesh([2, 3, 4])
+        return mesh
+
+    @pytest.mark.parametrize("engine", ("geoana", "choclo"))
+    def test_error_on_gravity(self, mesh_3d, engine):
+        """
+        Test error is raised after passing a 3D mesh to gravity eq source class.
+        """
+        msg = "Mesh to equivalent source layer must be 2D."
+        with pytest.raises(AttributeError, match=msg):
+            gravity.SimulationEquivalentSourceLayer(
+                mesh=mesh_3d, cell_z_top=0.0, cell_z_bottom=-2.0, engine=engine
+            )
+
+    @pytest.mark.parametrize("engine", ("geoana",))
+    def test_error_on_mag(self, mesh_3d, engine):
+        """
+        Test error is raised after passing a 3D mesh to magnetic eq source class.
+        """
+        msg = "Mesh to equivalent source layer must be 2D."
+        with pytest.raises(AttributeError, match=msg):
+            magnetics.SimulationEquivalentSourceLayer(
+                mesh=mesh_3d, cell_z_top=0.0, cell_z_bottom=-2.0, engine=engine
+            )
+
+    def test_error_on_base_class(self, mesh_3d):
+        """
+        Test error is raised after passing a 3D mesh to the eq source base class.
+        """
+        msg = "Mesh to equivalent source layer must be 2D."
+        with pytest.raises(AttributeError, match=msg):
+            base.BaseEquivalentSourceLayerSimulation(
+                mesh=mesh_3d, cell_z_top=0.0, cell_z_bottom=-2.0
+            )
+
+    @pytest.mark.parametrize(
+        "eq_sources_class",
+        (
+            gravity.SimulationEquivalentSourceLayer,
+            magnetics.SimulationEquivalentSourceLayer,
+        ),
+    )
+    def test_overridden_method(
+        self, mesh_3d, tensor_mesh, mesh_top, mesh_bottom, eq_sources_class
+    ):
+        """
+        Test for the overridden _check_engine_and_mesh_dimensions method.
+
+        This method is rarely going to trigger an error because if a 3D mesh is
+        passed to the constructor, it'll catch it and raise an error.
+        This test is added to extend coverage.
+        """
+        # Initialize instance with a 2D mesh
+        eq_sources = eq_sources_class(
+            mesh=tensor_mesh, cell_z_top=mesh_top, cell_z_bottom=mesh_bottom
+        )
+        # Set the mesh to the 3D mesh
+        eq_sources.mesh = mesh_3d
+        # Check error in _check_engine_and_mesh_dimensions method
+        msg = "Mesh to equivalent source layer must be 2D."
+        with pytest.raises(AttributeError, match=msg):
+            eq_sources._check_engine_and_mesh_dimensions()
+
+
+class TestGravityEquivalentSourcesForward:
+    """
+    Test the forward capabilities of the gravity equivalent sources.
+    """
+
+    def get_mesh_3d(self, mesh, top: float, bottom: float):
+        """
+        Build a 3D mesh analogous to the 2D mesh + the top and bottom bounds.
+        """
+        origin = (*mesh.origin, bottom)
+        h = (*mesh.h, np.array([top - bottom], dtype=np.float64))
+        mesh_3d = TensorMesh(h=h, origin=origin)
+        return mesh_3d
+
+    @pytest.mark.parametrize("engine", ("geoana", "choclo"))
+    @pytest.mark.parametrize("store_sensitivities", ("ram", "forward_only"))
+    def test_forward_vs_simulation(
+        self,
+        tensor_mesh,
+        mesh_bottom,
+        mesh_top,
+        gravity_survey,
+        engine,
+        store_sensitivities,
+    ):
+        """
+        Test forward of the eq sources vs. using the integral 3d simulation.
+        """
+        # Build 3D mesh that is analogous to the 2D mesh with bottom and top
+        mesh_3d = self.get_mesh_3d(tensor_mesh, top=mesh_top, bottom=mesh_bottom)
+        # Build simulations
+        mapping = get_mapping(tensor_mesh)
+        sim_3d = gravity.Simulation3DIntegral(
+            survey=gravity_survey, mesh=mesh_3d, rhoMap=mapping
+        )
+        eq_sources = gravity.SimulationEquivalentSourceLayer(
+            mesh=tensor_mesh,
+            cell_z_top=mesh_top,
+            cell_z_bottom=mesh_bottom,
+            survey=gravity_survey,
+            rhoMap=mapping,
+            engine=engine,
+            store_sensitivities=store_sensitivities,
+        )
+        # Compare predictions of both simulations
+        model = get_block_model(tensor_mesh, 2.67)
+        np.testing.assert_allclose(
+            sim_3d.dpred(model), eq_sources.dpred(model), atol=1e-7
+        )
+
+    @pytest.mark.parametrize("engine", ("geoana", "choclo"))
+    @pytest.mark.parametrize("store_sensitivities", ("ram", "forward_only"))
+    @pytest.mark.parametrize("components", COMPONENTS + [["gz", "gzz"]])
+    def test_forward_vs_simulation_with_components(
+        self,
+        coordinates,
+        tensor_mesh,
+        mesh_bottom,
+        mesh_top,
+        engine,
+        store_sensitivities,
+        components,
+    ):
+        """
+        Test forward vs simulation using different gravity components.
+        """
+        # Build survey
+        survey = build_gravity_survey(coordinates, components)
+        # Build 3D mesh that is analogous to the 2D mesh with bottom and top
+        mesh_3d = self.get_mesh_3d(tensor_mesh, top=mesh_top, bottom=mesh_bottom)
+        # Build simulations
+        mapping = get_mapping(tensor_mesh)
+        sim_3d = gravity.Simulation3DIntegral(
+            survey=survey, mesh=mesh_3d, rhoMap=mapping
+        )
+        eq_sources = gravity.SimulationEquivalentSourceLayer(
+            mesh=tensor_mesh,
+            cell_z_top=mesh_top,
+            cell_z_bottom=mesh_bottom,
+            survey=survey,
+            rhoMap=mapping,
+            engine=engine,
+            store_sensitivities=store_sensitivities,
+        )
+        # Compare predictions of both simulations
+        model = get_block_model(tensor_mesh, 2.67)
+        np.testing.assert_allclose(
+            sim_3d.dpred(model), eq_sources.dpred(model), atol=5e-6
+        )
+
+    @pytest.mark.parametrize("engine", ("geoana", "choclo"))
+    def test_forward_vs_simulation_on_disk(
+        self,
+        tensor_mesh,
+        mesh_bottom,
+        mesh_top,
+        gravity_survey,
+        engine,
+        tmp_path,
+    ):
+        """
+        Test forward vs simulation storing sensitivities on disk.
+        """
+        # Build 3D mesh that is analogous to the 2D mesh with bottom and top
+        mesh_3d = self.get_mesh_3d(tensor_mesh, top=mesh_top, bottom=mesh_bottom)
+        # Define sensitivity_dir
+        if engine == "geoana":
+            sensitivity_path = tmp_path / "sensitivities_geoana"
+            sensitivity_path.mkdir()
+        elif engine == "choclo":
+            sensitivity_path = tmp_path / "sensitivities_choclo"
+        # Build simulations
+        mapping = get_mapping(tensor_mesh)
+        sim_3d = gravity.Simulation3DIntegral(
+            survey=gravity_survey, mesh=mesh_3d, rhoMap=mapping
+        )
+        eq_sources = gravity.SimulationEquivalentSourceLayer(
+            mesh=tensor_mesh,
+            cell_z_top=mesh_top,
+            cell_z_bottom=mesh_bottom,
+            survey=gravity_survey,
+            rhoMap=mapping,
+            engine=engine,
+            store_sensitivities="disk",
+            sensitivity_path=str(sensitivity_path),
+        )
+        # Compare predictions of both simulations
+        model = get_block_model(tensor_mesh, 2.67)
+        np.testing.assert_allclose(sim_3d.dpred(model), eq_sources.dpred(model))
+
+    @pytest.mark.parametrize("engine", ("geoana", "choclo"))
+    @pytest.mark.parametrize("store_sensitivities", ("ram", "forward_only"))
+    def test_forward_vs_simulation_with_active_cells(
+        self,
+        tensor_mesh,
+        mesh_bottom,
+        mesh_top,
+        gravity_survey,
+        engine,
+        store_sensitivities,
+    ):
+        """
+        Test forward vs simulation using active cells.
+        """
+        model = get_block_model(tensor_mesh, 2.67)
+
+        # Define some inactive cells inside the block
+        block_cells_indices = np.indices(model.shape).ravel()[model != 0]
+        inactive_indices = block_cells_indices[
+            : block_cells_indices.size // 2
+        ]  # mark half of the cells in the block as inactive
+        active_cells = np.ones_like(model, dtype=bool)
+        active_cells[inactive_indices] = False
+        assert not np.all(active_cells)  # check we do have inactive cells
+
+        # Keep only values of the model in the active cells
+        model = model[active_cells]
+
+        # Build 3D mesh that is analogous to the 2D mesh with bottom and top
+        mesh_3d = self.get_mesh_3d(tensor_mesh, top=mesh_top, bottom=mesh_bottom)
+
+        # Build simulations
+        mapping = simpeg.maps.IdentityMap(nP=model.size)
+        sim_3d = gravity.Simulation3DIntegral(
+            survey=gravity_survey,
+            mesh=mesh_3d,
+            rhoMap=mapping,
+            active_cells=active_cells,
+        )
+        eq_sources = gravity.SimulationEquivalentSourceLayer(
+            mesh=tensor_mesh,
+            cell_z_top=mesh_top,
+            cell_z_bottom=mesh_bottom,
+            survey=gravity_survey,
+            rhoMap=mapping,
+            engine=engine,
+            store_sensitivities=store_sensitivities,
+            active_cells=active_cells,
+        )
+        # Compare predictions of both simulations
+        np.testing.assert_allclose(
+            sim_3d.dpred(model), eq_sources.dpred(model), atol=1e-7
+        )
+
+    @pytest.mark.parametrize("store_sensitivities", ("ram", "forward_only"))
+    def test_forward_geoana_choclo(
+        self, mesh, mesh_bottom, mesh_top, gravity_survey, store_sensitivities
+    ):
+        """Compare forwards using geoana and choclo."""
+        # Build simulations
+        mapping = get_mapping(mesh)
+        kwargs = dict(
+            mesh=mesh,
+            cell_z_top=mesh_top,
+            cell_z_bottom=mesh_bottom,
+            survey=gravity_survey,
+            rhoMap=mapping,
+            store_sensitivities=store_sensitivities,
+        )
+        sim_geoana = gravity.SimulationEquivalentSourceLayer(engine="geoana", **kwargs)
+        sim_choclo = gravity.SimulationEquivalentSourceLayer(engine="choclo", **kwargs)
+        model = get_block_model(mesh, 2.67)
+        np.testing.assert_allclose(sim_geoana.dpred(model), sim_choclo.dpred(model))
+
+    def test_forward_choclo_serial_parallel(
+        self, mesh, mesh_bottom, mesh_top, gravity_survey
+    ):
+        """Test forward using choclo in serial and in parallel."""
+        # Build simulations
+        mapping = get_mapping(mesh)
+        kwargs = dict(
+            mesh=mesh,
+            cell_z_top=mesh_top,
+            cell_z_bottom=mesh_bottom,
+            survey=gravity_survey,
+            rhoMap=mapping,
+            engine="choclo",
+        )
+        sim_parallel = gravity.SimulationEquivalentSourceLayer(
+            numba_parallel=True, **kwargs
+        )
+        sim_serial = gravity.SimulationEquivalentSourceLayer(
+            numba_parallel=False, **kwargs
+        )
+        model = get_block_model(mesh, 2.67)
+        np.testing.assert_allclose(sim_parallel.dpred(model), sim_serial.dpred(model))
+
+
+class TestGravityEquivalentSources:
+    """
+    Test fitting equivalent sources with synthetic data.
+    """
+
+    def get_mesh_top_bottom(self, mesh, array=False):
+        """Build the top and bottom boundaries of the mesh.
+
+        If array is True, the outputs are going to be arrays, otherwise they'll
+        be floats.
+        """
+        top, bottom = -20.0, -50.0
+        if array:
+            rng = np.random.default_rng(seed=42)
+            mesh_top = np.full(mesh.n_cells, fill_value=top) + rng.normal(
+                scale=0.5, size=mesh.n_cells
+            )
+            mesh_bottom = np.full(mesh.n_cells, fill_value=bottom) + rng.normal(
+                scale=0.5, size=mesh.n_cells
+            )
+        else:
+            mesh_top, mesh_bottom = top, bottom
+        return mesh_top, mesh_bottom
+
+    def build_synthetic_data(self, simulation, model):
+        data = simulation.make_synthetic_data(
+            model,
+            relative_error=0.0,
+            noise_floor=1e-3,
+            add_noise=True,
+            random_seed=1,
+        )
+        return data
+
+    def build_inversion(self, mesh, simulation, synthetic_data):
+        """Build inversion problem."""
+        # Build data misfit and regularization terms
+        data_misfit = simpeg.data_misfit.L2DataMisfit(
+            simulation=simulation, data=synthetic_data
+        )
+        regularization = simpeg.regularization.WeightedLeastSquares(mesh=mesh)
+        # Choose optimization
+        optimization = ProjectedGNCG(
+            maxIterLS=5,
+            maxIterCG=20,
+            tolCG=1e-4,
+        )
+        # Build inverse problem
+        inverse_problem = simpeg.inverse_problem.BaseInvProblem(
+            data_misfit, regularization, optimization
+        )
+        # Define directives
+        starting_beta = simpeg.directives.BetaEstimate_ByEig(
+            beta0_ratio=1e-1, random_seed=42
+        )
+        beta_schedule = simpeg.directives.BetaSchedule(coolingFactor=3, coolingRate=1)
+        update_jacobi = simpeg.directives.UpdatePreconditioner()
+        target_misfit = simpeg.directives.TargetMisfit(chifact=1)
+        sensitivity_weights = simpeg.directives.UpdateSensitivityWeights(
+            every_iteration=False
+        )
+        directives = [
+            sensitivity_weights,
+            starting_beta,
+            beta_schedule,
+            update_jacobi,
+            target_misfit,
+        ]
+        # Define inversion
+        inversion = simpeg.inversion.BaseInversion(inverse_problem, directives)
+        return inversion
+
+    @pytest.mark.parametrize("engine", ("geoana", "choclo"))
+    @pytest.mark.parametrize(
+        "top_bottom_as_array",
+        (False, True),
+        ids=("top-bottom-float", "top-bottom-array"),
+    )
+    def test_predictions_on_data_points(
+        self,
+        tree_mesh,
+        gravity_survey,
+        top_bottom_as_array,
+        engine,
+    ):
+        """
+        Test eq sources predictions on the same data points.
+
+        The equivalent sources should be able to reproduce the same data with
+        which they were trained.
+        """
+        # Get mesh top and bottom
+        mesh_top, mesh_bottom = self.get_mesh_top_bottom(
+            tree_mesh, array=top_bottom_as_array
+        )
+        # Build simulation
+        mapping = get_mapping(tree_mesh)
+        simulation = gravity.SimulationEquivalentSourceLayer(
+            tree_mesh,
+            mesh_top,
+            mesh_bottom,
+            survey=gravity_survey,
+            rhoMap=mapping,
+            engine=engine,
+        )
+        # Generate synthetic data
+        model = get_block_model(tree_mesh, 2.67)
+        synthetic_data = self.build_synthetic_data(simulation, model)
+        # Build inversion
+        inversion = self.build_inversion(tree_mesh, simulation, synthetic_data)
+        # Run inversion
+        starting_model = np.zeros(tree_mesh.n_cells)
+        recovered_model = inversion.run(starting_model)
+        # Predict data
+        prediction = simulation.dpred(recovered_model)
+        # Check if prediction is close to the synthetic data
+        atol, rtol = 0.005, 1e-5
+        np.testing.assert_allclose(
+            prediction, synthetic_data.dobs, atol=atol, rtol=rtol
+        )
+
+
+class TestMagneticEquivalentSources:
+    @pytest.fixture
+    def survey(self, coordinates):
+        """
+        Sample survey for the gravity equivalent sources.
+        """
+        return self._build_survey(coordinates, components="tmi")
+
+    def _build_survey(self, coordinates, components):
+        """
+        Build a magnetic survey.
+        """
+        receivers = magnetics.Point(coordinates, components=components)
+        source_field = magnetics.UniformBackgroundField(
+            [receivers], amplitude=50_000, inclination=35, declination=12
+        )
+        survey = magnetics.Survey(source_field)
+        return survey
+
+    def test_choclo_not_implemented(self, tensor_mesh, mesh_top, mesh_bottom, survey):
+        """
+        Test if error is raised when passing "choclo" to magnetic eq sources.
+        """
+        msg = (
+            "Magnetic equivalent sources with choclo as engine has not been"
+            " implemented yet. Use 'geoana' instead."
+        )
+        mapping = simpeg.maps.IdentityMap(nP=tensor_mesh.n_cells)
+        with pytest.raises(NotImplementedError, match=msg):
+            magnetics.SimulationEquivalentSourceLayer(
+                tensor_mesh,
+                mesh_top,
+                mesh_bottom,
+                survey=survey,
+                rhoMap=mapping,
+                engine="choclo",
+            )