Repository: simpeg/simpeg
Pull request #1285 (2023-09-19 20:04:52)

Gravity simulation with Choclo as engine
#### Summary

Add the option to use [Choclo](https://www.fatiando.org/choclo) as the engine for running the gravity simulations, considerably reducing the computation time of the forward simulations and the creation of the sensitivity matrix. Add Choclo as one of the optional dependencies. Update and extend tests for the gravity simulation using pytest. By default, the simulation will use `geoana` as the engine. Users can alternatively choose `choclo` instead if it's installed.


#### 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.-->

#### What does this implement/fix?
<!--Please explain your changes.-->

#### Additional information
<!--Any additional information you think is important.-->


<!--
Once all tests pass and the code has been reviewed and approved, it will be merged into main
-->



Patch:
diff --git a/SimPEG/potential_fields/gravity/_numba_functions.py b/SimPEG/potential_fields/gravity/_numba_functions.py
new file mode 100644
index 000000000..c84069f15
--- /dev/null
+++ b/SimPEG/potential_fields/gravity/_numba_functions.py
@@ -0,0 +1,250 @@
+"""
+Numba functions for gravity simulation using Choclo.
+"""
+import numpy as np
+
+try:
+    import choclo
+except ImportError:
+    # Define dummy jit decorator
+    def jit(*args, **kwargs):
+        return lambda f: f
+
+    choclo = None
+else:
+    from numba import jit, prange
+
+
+def _forward_gravity(
+    receivers,
+    nodes,
+    densities,
+    fields,
+    cell_nodes,
+    kernel_func,
+    constant_factor,
+):
+    """
+    Forward model the gravity field of active cells on receivers
+
+    This function should be used with a `numba.jit` decorator, for example:
+
+    ..code::
+
+        from numba import jit
+
+        jit_forward_gravity = jit(nopython=True, parallel=True)(_forward_gravity)
+
+    Parameters
+    ----------
+    receivers : (n_receivers, 3) numpy.ndarray
+        Array with the locations of the receivers
+    nodes : (n_active_nodes, 3) numpy.ndarray
+        Array with the location of the mesh nodes.
+    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.
+    cell_nodes : (n_active_cells, 8) numpy.ndarray
+        Array of integers, where each row contains the indices of the nodes for
+        each active cell in the mesh.
+    kernel_func : callable
+        Kernel function that will be evaluated on each node of the mesh. Choose
+        one of the kernel 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_nodes = nodes.shape[0]
+    n_cells = cell_nodes.shape[0]
+    # Evaluate kernel function on each node, for each receiver location
+    for i in prange(n_receivers):
+        # Allocate vector for kernels evaluated on mesh nodes
+        kernels = np.empty(n_nodes)
+        for j in range(n_nodes):
+            kernels[j] = _evaluate_kernel(
+                receivers[i, 0],
+                receivers[i, 1],
+                receivers[i, 2],
+                nodes[j, 0],
+                nodes[j, 1],
+                nodes[j, 2],
+                kernel_func,
+            )
+        # Compute fields from the kernel values
+        for k in range(n_cells):
+            fields[i] += (
+                constant_factor
+                * densities[k]
+                * _kernels_in_nodes_to_cell(
+                    kernels,
+                    cell_nodes[k, 0],
+                    cell_nodes[k, 1],
+                    cell_nodes[k, 2],
+                    cell_nodes[k, 3],
+                    cell_nodes[k, 4],
+                    cell_nodes[k, 5],
+                    cell_nodes[k, 6],
+                    cell_nodes[k, 7],
+                )
+            )
+
+
+def _sensitivity_gravity(
+    receivers,
+    nodes,
+    sensitivity_matrix,
+    cell_nodes,
+    kernel_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_sensitivity = jit(nopython=True, parallel=True)(_sensitivity_gravity)
+
+    Parameters
+    ----------
+    receivers : (n_receivers, 3) numpy.ndarray
+        Array with the locations of the receivers
+    nodes : (n_active_nodes, 3) numpy.ndarray
+        Array with the location of the mesh nodes.
+    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.
+    cell_nodes : (n_active_cells, 8) numpy.ndarray
+        Array of integers, where each row contains the indices of the nodes for
+        each active cell in the mesh.
+    kernel_func : callable
+        Kernel function that will be evaluated on each node of the mesh. Choose
+        one of the kernel 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_nodes = nodes.shape[0]
+    n_cells = cell_nodes.shape[0]
+    # Evaluate kernel function on each node, for each receiver location
+    for i in prange(n_receivers):
+        # Allocate vector for kernels evaluated on mesh nodes
+        kernels = np.empty(n_nodes)
+        for j in range(n_nodes):
+            kernels[j] = _evaluate_kernel(
+                receivers[i, 0],
+                receivers[i, 1],
+                receivers[i, 2],
+                nodes[j, 0],
+                nodes[j, 1],
+                nodes[j, 2],
+                kernel_func,
+            )
+        # Compute sensitivity matrix elements from the kernel values
+        for k in range(n_cells):
+            sensitivity_matrix[i, k] = constant_factor * _kernels_in_nodes_to_cell(
+                kernels,
+                cell_nodes[k, 0],
+                cell_nodes[k, 1],
+                cell_nodes[k, 2],
+                cell_nodes[k, 3],
+                cell_nodes[k, 4],
+                cell_nodes[k, 5],
+                cell_nodes[k, 6],
+                cell_nodes[k, 7],
+            )
+
+
+@jit(nopython=True)
+def _evaluate_kernel(
+    receiver_x, receiver_y, receiver_z, node_x, node_y, node_z, kernel_func
+):
+    """
+    Evaluate a kernel function for a single node and receiver
+
+    Parameters
+    ----------
+    receiver_x, receiver_y, receiver_z : floats
+        Coordinates of the receiver.
+    node_x, node_y, node_z : floats
+        Coordinates of the node.
+    kernel_func : callable
+        Kernel function that should be evaluated. For example, use one of the
+        kernel functions in ``choclo.prism``.
+
+    Returns
+    -------
+    float
+        Kernel evaluated on the given node and receiver.
+    """
+    dx = node_x - receiver_x
+    dy = node_y - receiver_y
+    dz = node_z - receiver_z
+    distance = np.sqrt(dx**2 + dy**2 + dz**2)
+    return kernel_func(dx, dy, dz, distance)
+
+
+@jit(nopython=True)
+def _kernels_in_nodes_to_cell(
+    kernels,
+    nodes_indices_0,
+    nodes_indices_1,
+    nodes_indices_2,
+    nodes_indices_3,
+    nodes_indices_4,
+    nodes_indices_5,
+    nodes_indices_6,
+    nodes_indices_7,
+):
+    """
+    Evaluate integral on a given cell from evaluation of kernels on nodes
+
+    Parameters
+    ----------
+    kernels : (n_active_nodes,) numpy.ndarray
+        Array with kernel values on each one of the nodes in the mesh.
+    nodes_indices : ints
+        Indices of the nodes for the current cell in "F" order (x changes
+        faster than y, and y faster than z).
+
+    Returns
+    -------
+    float
+    """
+    result = (
+        -kernels[nodes_indices_0]
+        + kernels[nodes_indices_1]
+        + kernels[nodes_indices_2]
+        - kernels[nodes_indices_3]
+        + kernels[nodes_indices_4]
+        - kernels[nodes_indices_5]
+        - kernels[nodes_indices_6]
+        + kernels[nodes_indices_7]
+    )
+    return result
+
+
+# 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)
diff --git a/SimPEG/potential_fields/gravity/simulation.py b/SimPEG/potential_fields/gravity/simulation.py
index c5f134e40..a41ea4811 100644
--- a/SimPEG/potential_fields/gravity/simulation.py
+++ b/SimPEG/potential_fields/gravity/simulation.py
@@ -1,4 +1,7 @@
+import os
+import warnings
 import numpy as np
+import discretize
 import scipy.constants as constants
 from geoana.kernels import prism_fz, prism_fzx, prism_fzy, prism_fzz
 from scipy.constants import G as NewtG
@@ -9,6 +12,52 @@ from SimPEG.utils import mkvc, sdiag
 from ...base import BasePDESimulation
 from ..base import BaseEquivalentSourceLayerSimulation, BasePFSimulation
 
+from ._numba_functions import (
+    choclo,
+    _sensitivity_gravity_serial,
+    _sensitivity_gravity_parallel,
+    _forward_gravity_serial,
+    _forward_gravity_parallel,
+)
+
+if choclo is not None:
+    from numba import jit
+
+    @jit(nopython=True)
+    def kernel_uv(easting, northing, upward, radius):
+        """Kernel for Guv gradiometry component."""
+        result = 0.5 * (
+            choclo.prism.kernel_nn(easting, northing, upward, radius)
+            - choclo.prism.kernel_ee(easting, northing, upward, radius)
+        )
+        return result
+
+    CHOCLO_KERNELS = {
+        "gx": choclo.prism.kernel_e,
+        "gy": choclo.prism.kernel_n,
+        "gz": choclo.prism.kernel_u,
+        "gxx": choclo.prism.kernel_ee,
+        "gyy": choclo.prism.kernel_nn,
+        "gzz": choclo.prism.kernel_uu,
+        "gxy": choclo.prism.kernel_en,
+        "gxz": choclo.prism.kernel_eu,
+        "gyz": choclo.prism.kernel_nu,
+        "guv": kernel_uv,
+    }
+
+
+def _get_conversion_factor(component):
+    """
+    Return conversion factor for the given component
+    """
+    if component in ("gx", "gy", "gz"):
+        conversion_factor = 1e8
+    elif component in ("gxx", "gyy", "gzz", "gxy", "gxz", "gyz", "guv"):
+        conversion_factor = 1e12
+    else:
+        raise ValueError(f"Invalid component '{component}'.")
+    return conversion_factor
+
 
 class Simulation3DIntegral(BasePFSimulation):
     """
@@ -27,27 +76,134 @@ class Simulation3DIntegral(BasePFSimulation):
 
         Gradient components ("gxx", "gyy", "gzz", "gxy", "gxz", "gyz") are
         returned in Eotvos (:math:`10^{-9} s^{-2}`).
+
+    Parameters
+    ----------
+    mesh : discretize.TreeMesh or discretize.TensorMesh
+        Mesh use to run the gravity simulation.
+    survey : SimPEG.potential_fields.gravity.Survey
+        Gravity survey with information of the receivers.
+    ind_active : (n_cells) numpy.ndarray, optional
+        Array that indicates which cells in ``mesh`` are active cells.
+    rho : numpy.ndarray (optional)
+        Density array for the active cells in the mesh.
+    rhoMap : Mapping (optional)
+        Model mapping.
+    sensitivity_dtype : numpy.dtype, optional
+        Data type that will be used to build the sensitivity matrix.
+    store_sensitivities : str
+        Options for storing sensitivity matrix. There are 3 options
+
+        - 'ram': sensitivities are stored in the computer's RAM
+        - 'disk': sensitivities are written to a directory
+        - 'forward_only': you intend only do perform a forward simulation and
+          sensitivities do not need to be stored
+
+    sensitivity_path : str, optional
+        Path to store the sensitivity matrix if ``store_sensitivities`` is set
+        to ``"disk"``. Default to "./sensitivities".
+    engine : str, optional
+       Choose which engine should be used to run the forward model:
+       ``"geoana"`` or "``choclo``".
+    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.
     """
 
     rho, rhoMap, rhoDeriv = props.Invertible("Density")
 
-    def __init__(self, mesh, rho=None, rhoMap=None, **kwargs):
+    def __init__(
+        self,
+        mesh,
+        rho=None,
+        rhoMap=None,
+        engine="geoana",
+        numba_parallel=True,
+        **kwargs,
+    ):
         super().__init__(mesh, **kwargs)
         self.rho = rho
         self.rhoMap = rhoMap
         self._G = None
         self._gtg_diagonal = None
         self.modelMap = self.rhoMap
+        self.numba_parallel = numba_parallel
+        self.engine = engine
+        self._sanity_checks_engine(kwargs)
+        # Define jit functions
+        if self.engine == "choclo":
+            if numba_parallel:
+                self._sensitivity_gravity = _sensitivity_gravity_parallel
+                self._forward_gravity = _forward_gravity_parallel
+            else:
+                self._sensitivity_gravity = _sensitivity_gravity_serial
+                self._forward_gravity = _forward_gravity_serial
+
+    def _sanity_checks_engine(self, kwargs):
+        """
+        Sanity checks for the engine parameter.
+
+        Needs the kwargs passed to the __init__ method to raise some warnings.
+        Will set n_processes to None if it's present in kwargs.
+        """
+        if self.engine not in ("choclo", "geoana"):
+            raise ValueError(
+                f"Invalid engine '{self.engine}'. Choose from 'geoana' or 'choclo'."
+            )
+        if self.engine == "choclo" and choclo is None:
+            raise ImportError(
+                "The choclo package couldn't be found."
+                "Running a gravity simulation with 'engine=\"choclo\"' needs "
+                "choclo to be installed."
+                "\nTry installing choclo with:"
+                "\n    pip install choclo"
+                "\nor:"
+                "\n    conda install choclo"
+            )
+        # Warn if n_processes has been passed
+        if self.engine == "choclo" and "n_processes" in kwargs:
+            warnings.warn(
+                "The 'n_processes' will be ignored when selecting 'choclo' as the "
+                "engine in the gravity simulation.",
+                UserWarning,
+                stacklevel=1,
+            )
+            self.n_processes = None
+        # Sanity checks for sensitivity_path when using choclo and storing in disk
+        if self.engine == "choclo" and self.store_sensitivities == "disk":
+            if os.path.isdir(self.sensitivity_path):
+                raise ValueError(
+                    f"The passed sensitivity_path '{self.sensitivity_path}' is "
+                    "a directory. "
+                    "When using 'choclo' as the engine, 'senstivity_path' "
+                    "should be the path to a new or existing file."
+                )
 
     def fields(self, m):
-        self.model = m
+        """
+        Forward model the gravity field of the mesh on the receivers in the survey
+
+        Parameters
+        ----------
+        m : (n_active_cells,) numpy.ndarray
+            Array with values for the model.
 
+        Returns
+        -------
+        (nD,) numpy.ndarray
+            Gravity fields generated by the given model on every receiver
+            location.
+        """
+        self.model = m
         if self.store_sensitivities == "forward_only":
             # Compute the linear operation without forming the full dense G
-            fields = mkvc(self.linear_operator())
+            if self.engine == "choclo":
+                fields = self._forward(self.rho)
+            else:
+                fields = mkvc(self.linear_operator())
         else:
             fields = self.G @ (self.rho).astype(self.sensitivity_dtype, copy=False)
-
         return np.asarray(fields)
 
     def getJtJdiag(self, m, W=None, f=None):
@@ -95,8 +251,10 @@ class Simulation3DIntegral(BasePFSimulation):
         Gravity forward operator
         """
         if getattr(self, "_G", None) is None:
-            self._G = self.linear_operator()
-
+            if self.engine == "choclo":
+                self._G = self._sensitivity_matrix()
+            else:
+                self._G = self.linear_operator()
         return self._G
 
     @property
@@ -205,6 +363,140 @@ class Simulation3DIntegral(BasePFSimulation):
             ]
         )
 
+    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.
+        """
+        # Gather active nodes and the indices of the nodes for each active cell
+        active_nodes, active_cell_nodes = self._get_active_nodes()
+        # 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):
+                kernel_func = CHOCLO_KERNELS[component]
+                conversion_factor = _get_conversion_factor(component)
+                vector_slice = slice(
+                    index_offset + i, index_offset + n_elements, n_components
+                )
+                self._forward_gravity(
+                    receivers,
+                    active_nodes,
+                    densities,
+                    fields[vector_slice],
+                    active_cell_nodes,
+                    kernel_func,
+                    constants.G * conversion_factor,
+                )
+            index_offset += n_elements
+        return fields
+
+    def _sensitivity_matrix(self):
+        """
+        Compute the sensitivity matrix G
+
+        Returns
+        -------
+        (nD, n_active_cells) numpy.ndarray
+        """
+        # Gather active nodes and the indices of the nodes for each active cell
+        active_nodes, active_cell_nodes = self._get_active_nodes()
+        # 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):
+                kernel_func = CHOCLO_KERNELS[component]
+                conversion_factor = _get_conversion_factor(component)
+                matrix_slice = slice(
+                    index_offset + i, index_offset + n_rows, n_components
+                )
+                self._sensitivity_gravity(
+                    receivers,
+                    active_nodes,
+                    sensitivity_matrix[matrix_slice, :],
+                    active_cell_nodes,
+                    kernel_func,
+                    constants.G * conversion_factor,
+                )
+            index_offset += n_rows
+        return sensitivity_matrix
+
+    def _get_cell_nodes(self):
+        """
+        Return indices of nodes for each cell in the mesh.
+        """
+        if not isinstance(self.mesh, (discretize.TreeMesh, discretize.TensorMesh)):
+            raise TypeError(f"Invalid mesh of type {self.mesh.__class__.__name__}.")
+        cell_nodes = self.mesh.cell_nodes
+        return cell_nodes
+
+    def _get_active_nodes(self):
+        """
+        Return locations of nodes only for active cells
+
+        Also return an array containing the indices of the "active nodes" for
+        each active cell in the mesh
+        """
+        # Get all nodes in the mesh
+        if isinstance(self.mesh, discretize.TreeMesh):
+            nodes = self.mesh.total_nodes
+        elif isinstance(self.mesh, discretize.TensorMesh):
+            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
+        cell_nodes = self._get_cell_nodes()
+        # If all cells in the mesh are active, return nodes and cell_nodes
+        if self.nC == self.mesh.n_cells:
+            return nodes, cell_nodes
+        # Keep only the cell_nodes for active cells
+        cell_nodes = cell_nodes[self.ind_active]
+        # Get the unique indices of the nodes that belong to every active cell
+        # (these indices correspond to the original `nodes` array)
+        unique_nodes, active_cell_nodes = np.unique(cell_nodes, return_inverse=True)
+        # Select only the nodes that belong to the active cells (active nodes)
+        active_nodes = nodes[unique_nodes]
+        # Reshape indices of active cells for each active cell in the mesh
+        active_cell_nodes = active_cell_nodes.reshape(cell_nodes.shape)
+        return active_nodes, active_cell_nodes
+
+    def _get_components_and_receivers(self):
+        """Generator for receiver locations and their field components."""
+        if not hasattr(self.survey, "source_field"):
+            raise AttributeError(
+                f"The survey '{self.survey}' has no 'source_field' attribute."
+            )
+        for receiver_object in self.survey.source_field.receiver_list:
+            yield receiver_object.components, receiver_object.locations
+
 
 class SimulationEquivalentSourceLayer(
     BaseEquivalentSourceLayerSimulation, Simulation3DIntegral
diff --git a/SimPEG/simulation.py b/SimPEG/simulation.py
index 7e88483b9..16c9e3d7d 100644
--- a/SimPEG/simulation.py
+++ b/SimPEG/simulation.py
@@ -101,7 +101,7 @@ class BaseSimulation(props.HasModel):
 
     @property
     def sensitivity_path(self):
-        """Path to store the sensitivty.
+        """Path to store the sensitivity.
 
         Returns
         -------
diff --git a/SimPEG/utils/code_utils.py b/SimPEG/utils/code_utils.py
index b0c5623a4..87e4c1bed 100644
--- a/SimPEG/utils/code_utils.py
+++ b/SimPEG/utils/code_utils.py
@@ -491,6 +491,7 @@ class Report(ScoobyReport):
             "vtk",
             "utm",
             "memory_profiler",
+            "choclo",
         ]
 
         super().__init__(
diff --git a/setup.py b/setup.py
index 1a2459028..9e55ae825 100644
--- a/setup.py
+++ b/setup.py
@@ -40,7 +40,7 @@ setup(
         "scikit-learn>=1.2",
         "pymatsolver>=0.2",
         "matplotlib",
-        "discretize>=0.8",
+        "discretize>=0.10",
         "geoana>=0.5.0",
         "empymod>=2.0.0",
         "pandas",


Test patch:
diff --git a/environment_test.yml b/environment_test.yml
index b658b3042..888c2180a 100644
--- a/environment_test.yml
+++ b/environment_test.yml
@@ -7,7 +7,7 @@ dependencies:
   - scikit-learn>=1.2
   - pymatsolver>=0.2
   - matplotlib
-  - discretize>=0.8
+  - discretize>=0.10
   - geoana>=0.5.0
   - empymod>=2.0.0
   - setuptools_scm
@@ -36,6 +36,8 @@ dependencies:
   - pyvista
   - pip
   - python-kaleido
+  # Optional dependencies
+  - choclo
   # Linters and code style
   - pre-commit
   - black==23.1.0
diff --git a/tests/pf/test_forward_Grav_Linear.py b/tests/pf/test_forward_Grav_Linear.py
index 286e8ba6d..496e80b1f 100644
--- a/tests/pf/test_forward_Grav_Linear.py
+++ b/tests/pf/test_forward_Grav_Linear.py
@@ -1,29 +1,63 @@
-import unittest
+from unittest.mock import patch
 import pytest
 import discretize
 from SimPEG import maps
 from SimPEG.potential_fields import gravity
 from geoana.gravity import Prism
 import numpy as np
-import os
 
 
-def test_ana_grav_forward(tmp_path):
-    nx = 5
-    ny = 5
-
-    # Define a mesh
-    cs = 0.2
-    hxind = [(cs, 41)]
-    hyind = [(cs, 41)]
-    hzind = [(cs, 41)]
-    mesh = discretize.TensorMesh([hxind, hyind, hzind], "CCC")
-
-    # create a model of two blocks, 1 inside the other
-    block1 = np.array([[-1.5, 1.5], [-1.5, 1.5], [-1.5, 1.5]])
-    block2 = np.array([[-0.7, 0.7], [-0.7, 0.7], [-0.7, 0.7]])
-
-    def get_block_inds(grid, block):
+class TestsGravitySimulation:
+    """
+    Test gravity simulation.
+    """
+
+    @pytest.fixture
+    def blocks(self):
+        """Synthetic blocks to build the sample model."""
+        block1 = np.array([[-1.6, 1.6], [-1.6, 1.6], [-1.6, 1.6]])
+        block2 = np.array([[-0.8, 0.8], [-0.8, 0.8], [-0.8, 0.8]])
+        rho1 = 1.0
+        rho2 = 2.0
+        return (block1, block2), (rho1, rho2)
+
+    @pytest.fixture(params=("tensormesh", "treemesh"))
+    def mesh(self, blocks, request):
+        """Sample mesh."""
+        cs = 0.2
+        (block1, _), _ = blocks
+        if request.param == "tensormesh":
+            hxind, hyind, hzind = tuple([(cs, 42)] for _ in range(3))
+            mesh = discretize.TensorMesh([hxind, hyind, hzind], "CCC")
+        else:
+            h = cs * np.ones(64)
+            mesh = discretize.TreeMesh([h, h, h], origin="CCC")
+            x0, x1 = block1[:, 0], block1[:, 1]
+            mesh.refine_box(x0, x1, levels=9)
+        return mesh
+
+    @pytest.fixture
+    def simple_mesh(self):
+        """Simpler sample mesh, just to use it as a placeholder in some tests."""
+        return discretize.TensorMesh([5, 5, 5], "CCC")
+
+    @pytest.fixture
+    def density_and_active_cells(self, mesh, blocks):
+        """Sample density and active_cells arrays for the sample mesh."""
+        # create a model of two blocks, 1 inside the other
+        (block1, block2), (rho1, rho2) = blocks
+        block1_inds = self.get_block_inds(mesh.cell_centers, block1)
+        block2_inds = self.get_block_inds(mesh.cell_centers, block2)
+        # Define densities for each block
+        model = np.zeros(mesh.n_cells)
+        model[block1_inds] = rho1
+        model[block2_inds] = rho2
+        # Define active cells and reduce model
+        active_cells = model != 0.0
+        model_reduced = model[active_cells]
+        return model_reduced, active_cells
+
+    def get_block_inds(self, grid, block):
         return np.where(
             (grid[:, 0] > block[0, 0])
             & (grid[:, 0] < block[0, 1])
@@ -33,160 +67,395 @@ def test_ana_grav_forward(tmp_path):
             & (grid[:, 2] < block[2, 1])
         )
 
-    block1_inds = get_block_inds(mesh.cell_centers, block1)
-    block2_inds = get_block_inds(mesh.cell_centers, block2)
-
-    rho1 = 1.0
-    rho2 = 2.0
-
-    model = np.zeros(mesh.n_cells)
-    model[block1_inds] = rho1
-    model[block2_inds] = rho2
-
-    active_cells = model != 0.0
-    model_reduced = model[active_cells]
-
-    # Create reduced identity map for Linear Pproblem
-    idenMap = maps.IdentityMap(nP=int(sum(active_cells)))
-
-    # Create plane of observations
-    xr = np.linspace(-20, 20, nx)
-    yr = np.linspace(-20, 20, ny)
-    X, Y = np.meshgrid(xr, yr)
-    Z = np.ones_like(X) * 3.0
-    locXyz = np.c_[X.reshape(-1), Y.reshape(-1), Z.reshape(-1)]
-
-    receivers = gravity.Point(locXyz, components=["gx", "gy", "gz"])
-    sources = gravity.SourceField([receivers])
-    survey = gravity.Survey(sources)
-
-    sim = gravity.Simulation3DIntegral(
+    @pytest.fixture
+    def receivers_locations(self):
+        nx = 5
+        ny = 5
+        # Create plane of observations
+        xr = np.linspace(-20, 20, nx)
+        yr = np.linspace(-20, 20, ny)
+        x, y = np.meshgrid(xr, yr)
+        z = np.ones_like(x) * 3.0
+        receivers_locations = np.vstack([a.ravel() for a in (x, y, z)]).T
+        return receivers_locations
+
+    def get_analytic_solution(self, blocks, survey):
+        """Compute analytical response from dense prism."""
+        (block1, block2), (rho1, rho2) = blocks
+        # Build prisms (convert densities from g/cc to kg/m3)
+        prisms = [
+            Prism(block1[:, 0], block1[:, 1], rho1 * 1000),
+            Prism(block2[:, 0], block2[:, 1], -rho1 * 1000),
+            Prism(block2[:, 0], block2[:, 1], rho2 * 1000),
+        ]
+        # Forward model the prisms
+        components = survey.source_field.receiver_list[0].components
+        receivers_locations = survey.source_field.receiver_list[0].locations
+        if "gx" in components or "gy" in components or "gz" in components:
+            fields = sum(
+                prism.gravitational_field(receivers_locations) for prism in prisms
+            )
+            fields *= 1e5  # convert to mGal from m/s^2
+        else:
+            fields = sum(
+                prism.gravitational_gradient(receivers_locations) for prism in prisms
+            )
+            fields *= 1e9  # convert to Eotvos from 1/s^2
+        return fields
+
+    @pytest.mark.parametrize(
+        "engine, parallelism",
+        [("geoana", None), ("geoana", 1), ("choclo", False), ("choclo", True)],
+        ids=["geoana_serial", "geoana_parallel", "choclo_serial", "choclo_parallel"],
+    )
+    @pytest.mark.parametrize("store_sensitivities", ("ram", "disk", "forward_only"))
+    def test_accelerations_vs_analytic(
+        self,
+        engine,
+        parallelism,
+        store_sensitivities,
+        tmp_path,
+        blocks,
         mesh,
-        survey=survey,
-        rhoMap=idenMap,
-        ind_active=active_cells,
-        store_sensitivities="disk",
-        sensitivity_path=str(tmp_path) + os.sep,
+        density_and_active_cells,
+        receivers_locations,
+    ):
+        """
+        Test gravity acceleration components against analytic solutions of prisms.
+        """
+        components = ["gx", "gy", "gz"]
+        # Unpack fixtures
+        density, active_cells = density_and_active_cells
+        # Create survey
+        receivers = gravity.Point(receivers_locations, components=components)
+        sources = gravity.SourceField([receivers])
+        survey = gravity.Survey(sources)
+        # Create reduced identity map for Linear Problem
+        idenMap = maps.IdentityMap(nP=int(sum(active_cells)))
+        # Create simulation
+        if engine == "choclo":
+            sensitivity_path = tmp_path / "sensitivity_choclo"
+            kwargs = dict(numba_parallel=parallelism)
+        else:
+            sensitivity_path = tmp_path
+            kwargs = dict(n_processes=parallelism)
+        sim = gravity.Simulation3DIntegral(
+            mesh,
+            survey=survey,
+            rhoMap=idenMap,
+            ind_active=active_cells,
+            store_sensitivities=store_sensitivities,
+            engine=engine,
+            sensitivity_path=str(sensitivity_path),
+            sensitivity_dtype=np.float64,
+            **kwargs,
+        )
+        data = sim.dpred(density)
+        g_x, g_y, g_z = data[0::3], data[1::3], data[2::3]
+        solution = self.get_analytic_solution(blocks, survey)
+        # Check results
+        rtol, atol = 1e-9, 1e-6
+        np.testing.assert_allclose(g_x, solution[:, 0], rtol=rtol, atol=atol)
+        np.testing.assert_allclose(g_y, solution[:, 1], rtol=rtol, atol=atol)
+        np.testing.assert_allclose(g_z, solution[:, 2], rtol=rtol, atol=atol)
+
+    @pytest.mark.parametrize(
+        "engine, parallelism",
+        [("geoana", None), ("geoana", 1), ("choclo", False), ("choclo", True)],
+        ids=["geoana_serial", "geoana_parallel", "choclo_serial", "choclo_parallel"],
     )
-
-    with pytest.raises(TypeError):
-        sim.sensitivity_dtype = float
-
-    assert sim.sensitivity_dtype is np.float32
-
-    data = sim.dpred(model_reduced)
-    d_x = data[0::3]
-    d_y = data[1::3]
-    d_z = data[2::3]
-
-    # Compute analytical response from dense prism
-    prism_1 = Prism(block1[:, 0], block1[:, 1], rho1 * 1000)  # g/cc to kg/m**3
-    prism_2 = Prism(block2[:, 0], block2[:, 1], -rho1 * 1000)
-    prism_3 = Prism(block2[:, 0], block2[:, 1], rho2 * 1000)
-
-    d = (
-        prism_1.gravitational_field(locXyz)
-        + prism_2.gravitational_field(locXyz)
-        + prism_3.gravitational_field(locXyz)
-    ) * 1e5  # convert to mGal from m/s^2
-    d = d.astype(sim.sensitivity_dtype)
-    np.testing.assert_allclose(d_x, d[:, 0], rtol=1e-9, atol=1e-6)
-    np.testing.assert_allclose(d_y, d[:, 1], rtol=1e-9, atol=1e-6)
-    np.testing.assert_allclose(d_z, d[:, 2], rtol=1e-9, atol=1e-6)
-
-
-def test_ana_gg_forward():
-    nx = 5
-    ny = 5
-
-    # Define a mesh
-    cs = 0.2
-    hxind = [(cs, 41)]
-    hyind = [(cs, 41)]
-    hzind = [(cs, 41)]
-    mesh = discretize.TensorMesh([hxind, hyind, hzind], "CCC")
-
-    # create a model of two blocks, 1 inside the other
-    block1 = np.array([[-1.5, 1.5], [-1.5, 1.5], [-1.5, 1.5]])
-    block2 = np.array([[-0.7, 0.7], [-0.7, 0.7], [-0.7, 0.7]])
-
-    def get_block_inds(grid, block):
-        return np.where(
-            (grid[:, 0] > block[0, 0])
-            & (grid[:, 0] < block[0, 1])
-            & (grid[:, 1] > block[1, 0])
-            & (grid[:, 1] < block[1, 1])
-            & (grid[:, 2] > block[2, 0])
-            & (grid[:, 2] < block[2, 1])
+    @pytest.mark.parametrize("store_sensitivities", ("ram", "disk", "forward_only"))
+    def test_tensor_vs_analytic(
+        self,
+        engine,
+        parallelism,
+        store_sensitivities,
+        tmp_path,
+        blocks,
+        mesh,
+        density_and_active_cells,
+        receivers_locations,
+    ):
+        """
+        Test tensor components against analytic solutions of prisms.
+        """
+        components = ["gxx", "gxy", "gxz", "gyy", "gyz", "gzz"]
+        # Unpack fixtures
+        density, active_cells = density_and_active_cells
+        # Create survey
+        receivers = gravity.Point(receivers_locations, components=components)
+        sources = gravity.SourceField([receivers])
+        survey = gravity.Survey(sources)
+        # Create reduced identity map for Linear Problem
+        idenMap = maps.IdentityMap(nP=int(sum(active_cells)))
+        # Create simulation
+        if engine == "choclo":
+            sensitivity_path = tmp_path / "sensitivity_choclo"
+            kwargs = dict(numba_parallel=parallelism)
+        else:
+            sensitivity_path = tmp_path
+            kwargs = dict(n_processes=parallelism)
+        sim = gravity.Simulation3DIntegral(
+            mesh,
+            survey=survey,
+            rhoMap=idenMap,
+            ind_active=active_cells,
+            store_sensitivities=store_sensitivities,
+            engine=engine,
+            sensitivity_path=str(sensitivity_path),
+            sensitivity_dtype=np.float64,
+            **kwargs,
         )
-
-    block1_inds = get_block_inds(mesh.cell_centers, block1)
-    block2_inds = get_block_inds(mesh.cell_centers, block2)
-
-    rho1 = 1.0
-    rho2 = 2.0
-
-    model = np.zeros(mesh.n_cells)
-    model[block1_inds] = rho1
-    model[block2_inds] = rho2
-
-    active_cells = model != 0.0
-    model_reduced = model[active_cells]
-
-    # Create reduced identity map for Linear Pproblem
-    idenMap = maps.IdentityMap(nP=int(sum(active_cells)))
-
-    # Create plane of observations
-    xr = np.linspace(-20, 20, nx)
-    yr = np.linspace(-20, 20, ny)
-    X, Y = np.meshgrid(xr, yr)
-    Z = np.ones_like(X) * 3.0
-    locXyz = np.c_[X.reshape(-1), Y.reshape(-1), Z.reshape(-1)]
-
-    receivers = gravity.Point(
-        locXyz, components=["gxx", "gxy", "gxz", "gyy", "gyz", "gzz"]
+        data = sim.dpred(density)
+        g_xx, g_xy, g_xz = data[0::6], data[1::6], data[2::6]
+        g_yy, g_yz, g_zz = data[3::6], data[4::6], data[5::6]
+        solution = self.get_analytic_solution(blocks, survey)
+        # Check results
+        rtol, atol = 2e-6, 1e-6
+        np.testing.assert_allclose(g_xx, solution[..., 0, 0], rtol=rtol, atol=atol)
+        np.testing.assert_allclose(g_xy, solution[..., 0, 1], rtol=rtol, atol=atol)
+        np.testing.assert_allclose(g_xz, solution[..., 0, 2], rtol=rtol, atol=atol)
+        np.testing.assert_allclose(g_yy, solution[..., 1, 1], rtol=rtol, atol=atol)
+        np.testing.assert_allclose(g_yz, solution[..., 1, 2], rtol=rtol, atol=atol)
+        np.testing.assert_allclose(g_zz, solution[..., 2, 2], rtol=rtol, atol=atol)
+
+    @pytest.mark.parametrize(
+        "engine, parallelism",
+        [("geoana", 1), ("geoana", None), ("choclo", False), ("choclo", True)],
+        ids=["geoana_serial", "geoana_parallel", "choclo_serial", "choclo_parallel"],
     )
-    sources = gravity.SourceField([receivers])
-    survey = gravity.Survey(sources)
-
-    sim = gravity.Simulation3DIntegral(
+    @pytest.mark.parametrize("store_sensitivities", ("ram", "disk", "forward_only"))
+    def test_guv_vs_analytic(
+        self,
+        engine,
+        parallelism,
+        store_sensitivities,
+        tmp_path,
+        blocks,
         mesh,
-        survey=survey,
-        rhoMap=idenMap,
-        ind_active=active_cells,
-        store_sensitivities="forward_only",
-        n_processes=None,
+        density_and_active_cells,
+        receivers_locations,
+    ):
+        """
+        Test guv tensor component against analytic solutions of prisms.
+        """
+        components = ["guv"]
+        # Unpack fixtures
+        density, active_cells = density_and_active_cells
+        # Create survey
+        receivers = gravity.Point(receivers_locations, components=components)
+        sources = gravity.SourceField([receivers])
+        survey = gravity.Survey(sources)
+        # Create reduced identity map for Linear Problem
+        idenMap = maps.IdentityMap(nP=int(sum(active_cells)))
+        # Create simulation
+        if engine == "choclo":
+            sensitivity_path = tmp_path / "sensitivity_choclo"
+            kwargs = dict(numba_parallel=parallelism)
+        else:
+            sensitivity_path = tmp_path
+            kwargs = dict(n_processes=parallelism)
+        sim = gravity.Simulation3DIntegral(
+            mesh,
+            survey=survey,
+            rhoMap=idenMap,
+            ind_active=active_cells,
+            store_sensitivities=store_sensitivities,
+            engine=engine,
+            sensitivity_path=str(sensitivity_path),
+            sensitivity_dtype=np.float64,
+            **kwargs,
+        )
+        g_uv = sim.dpred(density)
+        solution = self.get_analytic_solution(blocks, survey)
+        g_xx_solution = solution[..., 0, 0]
+        g_yy_solution = solution[..., 1, 1]
+        g_uv_solution = 0.5 * (g_yy_solution - g_xx_solution)
+        # Check results
+        rtol, atol = 2e-6, 1e-6
+        np.testing.assert_allclose(g_uv, g_uv_solution, rtol=rtol, atol=atol)
+
+    @pytest.mark.parametrize("engine", ("choclo", "geoana"))
+    @pytest.mark.parametrize("store_sensitivities", ("ram", "disk", "forward_only"))
+    def test_sensitivity_dtype(
+        self,
+        engine,
+        store_sensitivities,
+        simple_mesh,
+        receivers_locations,
+        tmp_path,
+    ):
+        """Test sensitivity_dtype."""
+        # Create survey
+        receivers = gravity.Point(receivers_locations, components="gz")
+        sources = gravity.SourceField([receivers])
+        survey = gravity.Survey(sources)
+        # Create reduced identity map for Linear Problem
+        active_cells = np.ones(simple_mesh.n_cells, dtype=bool)
+        idenMap = maps.IdentityMap(nP=simple_mesh.n_cells)
+        # Create simulation
+        sensitivity_path = tmp_path
+        if engine == "choclo":
+            sensitivity_path /= "dummy"
+        simulation = gravity.Simulation3DIntegral(
+            simple_mesh,
+            survey=survey,
+            rhoMap=idenMap,
+            ind_active=active_cells,
+            engine=engine,
+            store_sensitivities=store_sensitivities,
+            sensitivity_path=str(sensitivity_path),
+        )
+        # sensitivity_dtype should be float64 when running forward only,
+        # but float32 in other cases
+        if store_sensitivities == "forward_only":
+            assert simulation.sensitivity_dtype is np.float64
+        else:
+            assert simulation.sensitivity_dtype is np.float32
+
+    @pytest.mark.parametrize("invalid_dtype", (float, np.float16))
+    def test_invalid_sensitivity_dtype_assignment(
+        self, simple_mesh, receivers_locations, invalid_dtype
+    ):
+        """
+        Test invalid sensitivity_dtype assignment
+        """
+        # Create survey
+        receivers = gravity.Point(receivers_locations, components="gz")
+        sources = gravity.SourceField([receivers])
+        survey = gravity.Survey(sources)
+        # Create reduced identity map for Linear Problem
+        active_cells = np.ones(simple_mesh.n_cells, dtype=bool)
+        idenMap = maps.IdentityMap(nP=simple_mesh.n_cells)
+        # Create simulation
+        simulation = gravity.Simulation3DIntegral(
+            simple_mesh,
+            survey=survey,
+            rhoMap=idenMap,
+            ind_active=active_cells,
+        )
+        # Check if error is raised
+        msg = "sensitivity_dtype must be either np.float32 or np.float64."
+        with pytest.raises(TypeError, match=msg):
+            simulation.sensitivity_dtype = invalid_dtype
+
+    def test_invalid_engine(self, simple_mesh, receivers_locations):
+        """Test if error is raised after invalid engine."""
+        # Create survey
+        receivers = gravity.Point(receivers_locations, components="gz")
+        sources = gravity.SourceField([receivers])
+        survey = gravity.Survey(sources)
+        # Create reduced identity map for Linear Problem
+        active_cells = np.ones(simple_mesh.n_cells, dtype=bool)
+        idenMap = maps.IdentityMap(nP=simple_mesh.n_cells)
+        # Check if error is raised after an invalid engine is passed
+        engine = "invalid engine"
+        with pytest.raises(ValueError, match=f"Invalid engine '{engine}'"):
+            gravity.Simulation3DIntegral(
+                simple_mesh,
+                survey=survey,
+                rhoMap=idenMap,
+                ind_active=active_cells,
+                engine=engine,
+            )
+
+    def test_choclo_and_n_proceesses(self, simple_mesh, receivers_locations):
+        """Check if warning is raised after passing n_processes with choclo engine."""
+        # Create survey
+        receivers = gravity.Point(receivers_locations, components="gz")
+        sources = gravity.SourceField([receivers])
+        survey = gravity.Survey(sources)
+        # Create reduced identity map for Linear Problem
+        active_cells = np.ones(simple_mesh.n_cells, dtype=bool)
+        idenMap = maps.IdentityMap(nP=simple_mesh.n_cells)
+        # Check if warning is raised
+        msg = "The 'n_processes' will be ignored when selecting 'choclo'"
+        with pytest.warns(UserWarning, match=msg):
+            simulation = gravity.Simulation3DIntegral(
+                simple_mesh,
+                survey=survey,
+                rhoMap=idenMap,
+                ind_active=active_cells,
+                engine="choclo",
+                n_processes=2,
+            )
+        # Check if n_processes was overwritten and set to None
+        assert simulation.n_processes is None
+
+    def test_choclo_and_sensitivity_path_as_dir(
+        self, simple_mesh, receivers_locations, tmp_path
+    ):
+        """
+        Check if error is raised when sensitivity_path is a dir with choclo engine.
+        """
+        # Create survey
+        receivers = gravity.Point(receivers_locations, components="gz")
+        sources = gravity.SourceField([receivers])
+        survey = gravity.Survey(sources)
+        # Create reduced identity map for Linear Problem
+        active_cells = np.ones(simple_mesh.n_cells, dtype=bool)
+        idenMap = maps.IdentityMap(nP=simple_mesh.n_cells)
+        # Create a sensitivity_path directory
+        sensitivity_path = tmp_path / "sensitivity_dummy"
+        sensitivity_path.mkdir()
+        # Check if error is raised
+        msg = f"The passed sensitivity_path '{str(sensitivity_path)}' is a directory"
+        with pytest.raises(ValueError, match=msg):
+            gravity.Simulation3DIntegral(
+                simple_mesh,
+                survey=survey,
+                rhoMap=idenMap,
+                ind_active=active_cells,
+                store_sensitivities="disk",
+                sensitivity_path=str(sensitivity_path),
+                engine="choclo",
+            )
+
+    @patch("SimPEG.potential_fields.gravity.simulation.choclo", None)
+    def test_choclo_missing(self, simple_mesh, receivers_locations):
+        """
+        Check if error is raised when choclo is missing and chosen as engine.
+        """
+        # Create survey
+        receivers = gravity.Point(receivers_locations, components="gz")
+        sources = gravity.SourceField([receivers])
+        survey = gravity.Survey(sources)
+        # Create reduced identity map for Linear Problem
+        active_cells = np.ones(simple_mesh.n_cells, dtype=bool)
+        idenMap = maps.IdentityMap(nP=simple_mesh.n_cells)
+        # Check if error is raised
+        msg = "The choclo package couldn't be found."
+        with pytest.raises(ImportError, match=msg):
+            gravity.Simulation3DIntegral(
+                simple_mesh,
+                survey=survey,
+                rhoMap=idenMap,
+                ind_active=active_cells,
+                engine="choclo",
+            )
+
+
+class TestConversionFactor:
+    """Test _get_conversion_factor function."""
+
+    @pytest.mark.parametrize(
+        "component",
+        ("gx", "gy", "gz", "gxx", "gyy", "gzz", "gxy", "gxz", "gyz", "guv"),
     )
-
-    # forward only should default to np.float64
-    assert sim.sensitivity_dtype is np.float64
-
-    data = sim.dpred(model_reduced)
-    d_xx = data[0::6]
-    d_xy = data[1::6]
-    d_xz = data[2::6]
-    d_yy = data[3::6]
-    d_yz = data[4::6]
-    d_zz = data[5::6]
-
-    # Compute analytical response from dense prism
-    prism_1 = Prism(block1[:, 0], block1[:, 1], rho1 * 1000)  # g/cc to kg/m**3
-    prism_2 = Prism(block2[:, 0], block2[:, 1], -rho1 * 1000)
-    prism_3 = Prism(block2[:, 0], block2[:, 1], rho2 * 1000)
-
-    d = (
-        prism_1.gravitational_gradient(locXyz)
-        + prism_2.gravitational_gradient(locXyz)
-        + prism_3.gravitational_gradient(locXyz)
-    ) * 1e9  # convert to Eotvos from 1/s^2
-
-    np.testing.assert_allclose(d_xx, d[..., 0, 0], rtol=1e-10, atol=1e-12)
-    np.testing.assert_allclose(d_xy, d[..., 0, 1], rtol=1e-10, atol=1e-12)
-    np.testing.assert_allclose(d_xz, d[..., 0, 2], rtol=1e-10, atol=1e-12)
-    np.testing.assert_allclose(d_yy, d[..., 1, 1], rtol=1e-10, atol=1e-12)
-    np.testing.assert_allclose(d_yz, d[..., 1, 2], rtol=1e-10, atol=1e-12)
-    np.testing.assert_allclose(d_zz, d[..., 2, 2], rtol=1e-10, atol=1e-12)
-
-
-if __name__ == "__main__":
-    unittest.main()
+    def test_conversion_factor(self, component):
+        """
+        Test _get_conversion_factor function with valid components
+        """
+        conversion_factor = gravity.simulation._get_conversion_factor(component)
+        if len(component) == 2:
+            assert conversion_factor == 1e5 * 1e3  # SI to mGal and g/cc to kg/m3
+        else:
+            assert conversion_factor == 1e9 * 1e3  # SI to Eotvos and g/cc to kg/m3
+
+    def test_invalid_conversion_factor(self):
+        """
+        Test invalid conversion factor _get_conversion_factor function
+        """
+        component = "invalid-component"
+        with pytest.raises(ValueError, match=f"Invalid component '{component}'"):
+            gravity.simulation._get_conversion_factor(component)
diff --git a/tests/pf/test_grav_inversion_linear.py b/tests/pf/test_grav_inversion_linear.py
index 4c98e6aec..c67a8ba47 100644
--- a/tests/pf/test_grav_inversion_linear.py
+++ b/tests/pf/test_grav_inversion_linear.py
@@ -1,5 +1,4 @@
-import shutil
-import unittest
+import pytest
 import numpy as np
 
 import discretize
@@ -17,125 +16,105 @@ from SimPEG import (
 from SimPEG.potential_fields import gravity
 
 
-class GravInvLinProblemTest(unittest.TestCase):
-    def setUp(self):
-        # Create a self.mesh
-        dx = 5.0
-        hxind = [(dx, 5, -1.3), (dx, 5), (dx, 5, 1.3)]
-        hyind = [(dx, 5, -1.3), (dx, 5), (dx, 5, 1.3)]
-        hzind = [(dx, 5, -1.3), (dx, 6)]
-        self.mesh = discretize.TensorMesh([hxind, hyind, hzind], "CCC")
-
-        # Get index of the center
-        midx = int(self.mesh.shape_cells[0] / 2)
-        midy = int(self.mesh.shape_cells[1] / 2)
-
-        # Lets create a simple Gaussian topo and set the active cells
-        [xx, yy] = np.meshgrid(self.mesh.nodes_x, self.mesh.nodes_y)
-        zz = -np.exp((xx**2 + yy**2) / 75**2) + self.mesh.nodes_z[-1]
-
-        # Go from topo to actv cells
-        topo = np.c_[utils.mkvc(xx), utils.mkvc(yy), utils.mkvc(zz)]
-        actv = active_from_xyz(self.mesh, topo, "N")
-
-        # Create active map to go from reduce space to full
-        self.actvMap = maps.InjectActiveCells(self.mesh, actv, -100)
-        nC = int(actv.sum())
-
-        # Create and array of observation points
-        xr = np.linspace(-20.0, 20.0, 20)
-        yr = np.linspace(-20.0, 20.0, 20)
-        X, Y = np.meshgrid(xr, yr)
-
-        # Move the observation points 5m above the topo
-        Z = -np.exp((X**2 + Y**2) / 75**2) + self.mesh.nodes_z[-1] + 5.0
-
-        # Create a MAGsurvey
-        locXYZ = np.c_[utils.mkvc(X.T), utils.mkvc(Y.T), utils.mkvc(Z.T)]
-        rxLoc = gravity.Point(locXYZ)
-        srcField = gravity.SourceField([rxLoc])
-        survey = gravity.Survey(srcField)
-
-        # We can now create a density model and generate data
-        # Here a simple block in half-space
-        model = np.zeros(
-            (
-                self.mesh.shape_cells[0],
-                self.mesh.shape_cells[1],
-                self.mesh.shape_cells[2],
-            )
+@pytest.mark.parametrize("engine", ("geoana", "choclo"))
+def test_gravity_inversion_linear(engine):
+    """Test gravity inversion."""
+    # Create a mesh
+    dx = 5.0
+    hxind = [(dx, 5, -1.3), (dx, 5), (dx, 5, 1.3)]
+    hyind = [(dx, 5, -1.3), (dx, 5), (dx, 5, 1.3)]
+    hzind = [(dx, 5, -1.3), (dx, 6)]
+    mesh = discretize.TensorMesh([hxind, hyind, hzind], "CCC")
+
+    # Get index of the center
+    midx = int(mesh.shape_cells[0] / 2)
+    midy = int(mesh.shape_cells[1] / 2)
+
+    # Lets create a simple Gaussian topo and set the active cells
+    [xx, yy] = np.meshgrid(mesh.nodes_x, mesh.nodes_y)
+    zz = -np.exp((xx**2 + yy**2) / 75**2) + mesh.nodes_z[-1]
+
+    # Go from topo to actv cells
+    topo = np.c_[utils.mkvc(xx), utils.mkvc(yy), utils.mkvc(zz)]
+    actv = active_from_xyz(mesh, topo, "N")
+    nC = int(actv.sum())
+
+    # Create and array of observation points
+    xr = np.linspace(-20.0, 20.0, 20)
+    yr = np.linspace(-20.0, 20.0, 20)
+    X, Y = np.meshgrid(xr, yr)
+
+    # Move the observation points 5m above the topo
+    Z = -np.exp((X**2 + Y**2) / 75**2) + mesh.nodes_z[-1] + 5.0
+
+    # Create a MAGsurvey
+    locXYZ = np.c_[utils.mkvc(X.T), utils.mkvc(Y.T), utils.mkvc(Z.T)]
+    rxLoc = gravity.Point(locXYZ)
+    srcField = gravity.SourceField([rxLoc])
+    survey = gravity.Survey(srcField)
+
+    # We can now create a density model and generate data
+    # Here a simple block in half-space
+    model = np.zeros(
+        (
+            mesh.shape_cells[0],
+            mesh.shape_cells[1],
+            mesh.shape_cells[2],
         )
-        model[(midx - 2) : (midx + 2), (midy - 2) : (midy + 2), -6:-2] = 0.5
-        model = utils.mkvc(model)
-        self.model = model[actv]
-
-        # Create reduced identity map
-        idenMap = maps.IdentityMap(nP=nC)
-
-        # Create the forward model operator
-        sim = gravity.Simulation3DIntegral(
-            self.mesh,
-            survey=survey,
-            rhoMap=idenMap,
-            ind_active=actv,
-            store_sensitivities="ram",
-            n_processes=None,
-        )
-
-        # Compute linear forward operator and compute some data
-        data = sim.make_synthetic_data(
-            self.model,
-            relative_error=0.0,
-            noise_floor=0.0005,
-            add_noise=True,
-            random_seed=2,
-        )
-
-        # Create a regularization
-        reg = regularization.Sparse(self.mesh, active_cells=actv, mapping=idenMap)
-        reg.norms = [0, 0, 0, 0]
-        reg.gradientType = "components"
-
-        # Data misfit function
-        dmis = data_misfit.L2DataMisfit(simulation=sim, data=data)
-
-        # Add directives to the inversion
-        opt = optimization.ProjectedGNCG(
-            maxIter=100, lower=-1.0, upper=1.0, maxIterLS=20, maxIterCG=10, tolCG=1e-3
-        )
-        invProb = inverse_problem.BaseInvProblem(dmis, reg, opt)
-
-        # Here is where the norms are applied
-        starting_beta = directives.BetaEstimateMaxDerivative(10.0)
-        IRLS = directives.Update_IRLS()
-        update_Jacobi = directives.UpdatePreconditioner()
-        sensitivity_weights = directives.UpdateSensitivityWeights(everyIter=False)
-        self.inv = inversion.BaseInversion(
-            invProb,
-            directiveList=[IRLS, sensitivity_weights, starting_beta, update_Jacobi],
-        )
-        self.sim = sim
-
-    def test_grav_inverse(self):
-        # Run the inversion
-        mrec = self.inv.run(self.model)
-        residual = np.linalg.norm(mrec - self.model) / np.linalg.norm(self.model)
-
-        # import matplotlib.pyplot as plt
-        # plt.figure()
-        # ax = plt.subplot(1, 2, 1)
-        # self.mesh.plot_slice(self.actvMap*mrec, ax=ax, clim=(0, 0.5), normal="Y")
-        # ax = plt.subplot(1, 2, 2)
-        # self.mesh.plot_slice(self.actvMap*self.model, ax=ax, clim=(0, 0.5), normal="Y")
-        # plt.show()
-
-        self.assertTrue(residual < 0.05)
-
-    def tearDown(self):
-        # Clean up the working directory
-        if self.sim.store_sensitivities == "disk":
-            shutil.rmtree(self.sim.sensitivity_path)
-
-
-if __name__ == "__main__":
-    unittest.main()
+    )
+    model[(midx - 2) : (midx + 2), (midy - 2) : (midy + 2), -6:-2] = 0.5
+    model = utils.mkvc(model)
+    model = model[actv]
+
+    # Create reduced identity map
+    idenMap = maps.IdentityMap(nP=nC)
+
+    # Create the forward model operator
+    kwargs = dict()
+    if engine == "geoana":
+        kwargs["n_processes"] = None
+    sim = gravity.Simulation3DIntegral(
+        mesh,
+        survey=survey,
+        rhoMap=idenMap,
+        ind_active=actv,
+        store_sensitivities="ram",
+        engine=engine,
+        **kwargs,
+    )
+
+    # Compute linear forward operator and compute some data
+    data = sim.make_synthetic_data(
+        model, relative_error=0.0, noise_floor=0.0005, add_noise=True, random_seed=2
+    )
+
+    # Create a regularization
+    reg = regularization.Sparse(mesh, active_cells=actv, mapping=idenMap)
+    reg.norms = [0, 0, 0, 0]
+    reg.gradient_type = "components"
+
+    # Data misfit function
+    dmis = data_misfit.L2DataMisfit(simulation=sim, data=data)
+
+    # Add directives to the inversion
+    opt = optimization.ProjectedGNCG(
+        maxIter=100, lower=-1.0, upper=1.0, maxIterLS=20, maxIterCG=10, tolCG=1e-3
+    )
+    invProb = inverse_problem.BaseInvProblem(dmis, reg, opt)
+
+    # Here is where the norms are applied
+    starting_beta = directives.BetaEstimateMaxDerivative(10.0)
+    IRLS = directives.Update_IRLS()
+    update_Jacobi = directives.UpdatePreconditioner()
+    sensitivity_weights = directives.UpdateSensitivityWeights(every_iteration=False)
+    inv = inversion.BaseInversion(
+        invProb,
+        directiveList=[IRLS, sensitivity_weights, starting_beta, update_Jacobi],
+    )
+
+    # Run the inversion
+    mrec = inv.run(model)
+    residual = np.linalg.norm(mrec - model) / np.linalg.norm(model)
+
+    # Assert result
+    assert np.all(residual < 0.05)
diff --git a/tests/utils/test_report.py b/tests/utils/test_report.py
index e9ebe09bb..9d5091545 100644
--- a/tests/utils/test_report.py
+++ b/tests/utils/test_report.py
@@ -34,6 +34,7 @@ class TestReport(unittest.TestCase):
                 "vtk",
                 "utm",
                 "memory_profiler",
+                "choclo",
             ],
             ncol=3,
             text_width=80,