Repository: simpeg/simpeg
Pull request #1321 (2024-01-10 19:11:58)

Magnetic simulation with Choclo as engine
#### Summary

Add new implementation of the magnetic simulation using the integral formulation based on [Choclo](https://www.fatiando.org/choclo) and Numba. Allow users to run the simulation with Choclo as engine through a new `engine` argument of the constructor of the simulation class. Add `numba_parallel` argument to control if the simulation should be run in parallel or in a single thread. Move common methods and properties of the gravity and magnetic simulation to the base class for potential field simulations. Add new `potential_fields/_numba_utils.py` file with functions that are needed both in the magnetic and gravity Numba implementations. Refactor and extend tests for the magnetic simulation, including tests that assert the accuracy of the new implementation against analytic solutions. Add new tests for the base class for potential field simulations. Update examples using Choclo as engine, include admonitions instructing how to use it.

#### 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
-->

**Todo:**

- [x] Move the `engine` and `numba_parallel` arguments to the base class.
- [x] Make the new mag simulation to properly use the `engine` and `numba_parallel` attributes of the parent.
- [x] Add check for `n_processes` when `engine=="choclo"` in mag simulation.
- [x] Make tests for `engine`, `numba_parallel` and `engine` + `sensitivity_path` for the base class (not only for the child classes).
- [x] Improve docstrings: explain the different options users have to choose a different engine and parallelization options.
- [x] Implement storing the sensitivity matrix in disk



codecov[bot]: ## [Codecov](https://app.codecov.io/gh/simpeg/simpeg/pull/1321?src=pr&el=h1&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=simpeg) Report
Attention: `218 lines` in your changes are missing coverage. Please review.
> Comparison is base [(`5231446`)](https://app.codecov.io/gh/simpeg/simpeg/commit/5231446ee858fe677cdcbdcfe65be081b61ebd20?el=desc&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=simpeg) 82.30% compared to head [(`d2abea3`)](https://app.codecov.io/gh/simpeg/simpeg/pull/1321?src=pr&el=desc&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=simpeg) 33.65%.

| [Files](https://app.codecov.io/gh/simpeg/simpeg/pull/1321?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=simpeg) | Patch % | Lines |
|---|---|---|
| [...PEG/potential\_fields/magnetics/\_numba\_functions.py](https://app.codecov.io/gh/simpeg/simpeg/pull/1321?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=simpeg#diff-U2ltUEVHL3BvdGVudGlhbF9maWVsZHMvbWFnbmV0aWNzL19udW1iYV9mdW5jdGlvbnMucHk=) | 11.53% | [138 Missing :warning: ](https://app.codecov.io/gh/simpeg/simpeg/pull/1321?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=simpeg) |
| [SimPEG/potential\_fields/magnetics/simulation.py](https://app.codecov.io/gh/simpeg/simpeg/pull/1321?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=simpeg#diff-U2ltUEVHL3BvdGVudGlhbF9maWVsZHMvbWFnbmV0aWNzL3NpbXVsYXRpb24ucHk=) | 17.52% | [80 Missing :warning: ](https://app.codecov.io/gh/simpeg/simpeg/pull/1321?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=simpeg) |

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


```diff
@@             Coverage Diff             @@
##             main    #1321       +/-   ##
===========================================
- Coverage   82.30%   33.65%   -48.65%     
===========================================
  Files         165      166        +1     
  Lines       25241    25492      +251     
===========================================
- Hits        20775     8580    -12195     
- Misses       4466    16912    +12446     
```



</details>

[:umbrella: View full report in Codecov by Sentry](https://app.codecov.io/gh/simpeg/simpeg/pull/1321?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: > Just a few notes while you're working on this.

Thanks @jcapriot for the early review!
santisoler: @thibaut-kobold @jedman I've been working towards finishing this PR. I think it's ready to be reviewed. Would you like to take a look at it, give it a try and review it?
thibaut-kobold: @santisoler 
question / feature request: Will it work with `SimulationEquivalentSourceLayer`? That'd be awesome.

I believe `SimulationEquivalentSourceLayer` work with a 2D mesh to which you set the top and bottom along topo through the simulation object (rather than the mesh). It seems this does not work with `choclo` at the moment (at least not on my end).
santisoler: > @santisoler question / feature request: Will it work with `SimulationEquivalentSourceLayer`? That'd be awesome.
> 
> I believe `SimulationEquivalentSourceLayer` work with a 2D mesh to which you set the top and bottom along topo through the simulation object (rather than the mesh). It seems this does not work with `choclo` at the moment (at least not on my end).

Thanks for pointing it out. It wasn't my goal on this PR to make sure that this implementation would work with the `SimulationEquivalentSourceLayer`. We should enable that! Though I'm hesitant to tackle that feature in this PR: it's already a little bit large, and since "geoana" is the default engine at the moment, users won't face any breaking change after this PR gets merged.

Would you like to open an issue requesting that feature? It would also be great if you can include a minimal example and the error you're facing when trying to use `SimulationEquivalentSourceLayer` with the Choclo implementation of the magnetic simulation.

Patch:
diff --git a/simpeg/potential_fields/_numba_utils.py b/simpeg/potential_fields/_numba_utils.py
new file mode 100644
index 000000000..2bdea6da2
--- /dev/null
+++ b/simpeg/potential_fields/_numba_utils.py
@@ -0,0 +1,43 @@
+"""
+Utility functions for Numba implementations
+
+These functions are meant to be used both in the Numba-based gravity and
+magnetic simulations.
+"""
+
+try:
+    from numba import jit
+except ImportError:
+    # Define dummy jit decorator
+    def jit(*args, **kwargs):
+        return lambda f: f
+
+
+@jit(nopython=True)
+def kernels_in_nodes_to_cell(kernels, nodes_indices):
+    """
+    Evaluate integral on a given cell from evaluation of kernels on nodes
+
+    Parameters
+    ----------
+    kernels : (n_active_nodes,) array
+        Array with kernel values on each one of the nodes in the mesh.
+    nodes_indices : (8,) array of int
+        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
diff --git a/simpeg/potential_fields/base.py b/simpeg/potential_fields/base.py
index d457483e3..db63463a1 100644
--- a/simpeg/potential_fields/base.py
+++ b/simpeg/potential_fields/base.py
@@ -10,6 +10,11 @@ from simpeg.utils import mkvc
 from ..simulation import LinearSimulation
 from ..utils import validate_active_indices, validate_integer, validate_string
 
+try:
+    import choclo
+except ImportError:
+    choclo = None
+
 ###############################################################################
 #                                                                             #
 #                             Base Potential Fields Simulation                #
@@ -47,7 +52,14 @@ class BasePFSimulation(LinearSimulation):
     n_processes : None or int, optional
         The number of processes to use in the internal multiprocessing pool for forward
         modeling. The default value of 1 will not use multiprocessing. Any other setting
-        will. `None` implies setting by the number of cpus.
+        will. `None` implies setting by the number of cpus. If engine is
+        ``"choclo"``, then this argument will be ignored.
+    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.
 
     Notes
     -----
@@ -73,6 +85,8 @@ class BasePFSimulation(LinearSimulation):
         store_sensitivities="ram",
         n_processes=1,
         sensitivity_dtype=np.float32,
+        engine="geoana",
+        numba_parallel=True,
         **kwargs,
     ):
         # If deprecated property set with kwargs
@@ -88,10 +102,18 @@ class BasePFSimulation(LinearSimulation):
 
         self.store_sensitivities = store_sensitivities
         self.sensitivity_dtype = sensitivity_dtype
+        self.engine = engine
+        self.numba_parallel = numba_parallel
         super().__init__(mesh, **kwargs)
         self.solver = None
         self.n_processes = n_processes
 
+        # Check sensitivity_path when engine is "choclo"
+        self._check_engine_and_sensitivity_path()
+
+        # Check dimensions of the mesh when engine is "choclo"
+        self._check_engine_and_mesh_dimensions()
+
         # Find non-zero cells indices
         if ind_active is None:
             ind_active = np.ones(mesh.n_cells, dtype=bool)
@@ -188,6 +210,54 @@ class BasePFSimulation(LinearSimulation):
             value = validate_integer("n_processes", value, min_val=1)
         self._n_processes = value
 
+    @property
+    def engine(self) -> str:
+        """
+        Engine that will be used to run the simulation.
+
+        It can be either ``"geoana"`` or "``choclo``".
+        """
+        return self._engine
+
+    @engine.setter
+    def engine(self, value: str):
+        validate_string(
+            "engine", value, string_list=("geoana", "choclo"), case_sensitive=True
+        )
+        if value == "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"
+            )
+        self._engine = value
+
+    @property
+    def numba_parallel(self) -> bool:
+        """
+        Run simulation in parallel or single-threaded when using Numba.
+
+        If True, the simulation will run in parallel. If False, it will
+        run in serial.
+
+        .. important::
+
+            If ``engine`` is not ``"choclo"`` this property will be ignored.
+        """
+        return self._numba_parallel
+
+    @numba_parallel.setter
+    def numba_parallel(self, value: bool):
+        if not isinstance(value, bool):
+            raise TypeError(
+                f"Invalid 'numba_parallel' value of type {type(value)}. Must be a bool."
+            )
+        self._numba_parallel = value
+
     @property
     def ind_active(self):
         """Active topography cells.
@@ -255,6 +325,71 @@ class BasePFSimulation(LinearSimulation):
             np.save(sens_name, kernel)
         return kernel
 
+    def _check_engine_and_sensitivity_path(self):
+        """
+        Check if sensitivity_path is a file if engine is set to "choclo"
+        """
+        if (
+            self.engine == "choclo"
+            and self.store_sensitivities == "disk"
+            and 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 _check_engine_and_mesh_dimensions(self):
+        """
+        Check dimensions of the mesh when using choclo as engine
+        """
+        if self.engine == "choclo" and self.mesh.dim != 3:
+            raise ValueError(
+                f"Invalid mesh with {self.mesh.dim} dimensions. "
+                "Only 3D meshes are supported when using 'choclo' as engine."
+            )
+
+    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.mesh.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 cell nodes 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 BaseEquivalentSourceLayerSimulation(BasePFSimulation):
     """Base equivalent source layer simulation class.
diff --git a/simpeg/potential_fields/gravity/_numba_functions.py b/simpeg/potential_fields/gravity/_numba_functions.py
index 1d6b363b2..3eb6ae9bf 100644
--- a/simpeg/potential_fields/gravity/_numba_functions.py
+++ b/simpeg/potential_fields/gravity/_numba_functions.py
@@ -15,6 +15,8 @@ except ImportError:
 else:
     from numba import jit, prange
 
+from .._numba_utils import kernels_in_nodes_to_cell
+
 
 def _forward_gravity(
     receivers,
@@ -30,7 +32,7 @@ def _forward_gravity(
 
     This function should be used with a `numba.jit` decorator, for example:
 
-    ..code::
+    .. code::
 
         from numba import jit
 
@@ -85,16 +87,9 @@ def _forward_gravity(
             fields[i] += (
                 constant_factor
                 * densities[k]
-                * _kernels_in_nodes_to_cell(
+                * 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],
+                    cell_nodes[k, :],
                 )
             )
 
@@ -112,7 +107,7 @@ def _sensitivity_gravity(
 
     This function should be used with a `numba.jit` decorator, for example:
 
-    ..code::
+    .. code::
 
         from numba import jit
 
@@ -162,16 +157,9 @@ def _sensitivity_gravity(
             )
         # 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(
+            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],
+                cell_nodes[k, :],
             )
 
 
@@ -204,46 +192,6 @@ def _evaluate_kernel(
     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)
diff --git a/simpeg/potential_fields/gravity/simulation.py b/simpeg/potential_fields/gravity/simulation.py
index 311b19dfc..1596b15ec 100644
--- a/simpeg/potential_fields/gravity/simulation.py
+++ b/simpeg/potential_fields/gravity/simulation.py
@@ -1,7 +1,5 @@
-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
@@ -85,13 +83,13 @@ class Simulation3DIntegral(BasePFSimulation):
         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)
+    rho : numpy.ndarray, optional
         Density array for the active cells in the mesh.
-    rhoMap : Mapping (optional)
+    rhoMap : Mapping, optional
         Model mapping.
     sensitivity_dtype : numpy.dtype, optional
         Data type that will be used to build the sensitivity matrix.
-    store_sensitivities : str
+    store_sensitivities : {"ram", "disk", "forward_only"}
         Options for storing sensitivity matrix. There are 3 options
 
         - 'ram': sensitivities are stored in the computer's RAM
@@ -102,9 +100,8 @@ class Simulation3DIntegral(BasePFSimulation):
     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``".
+    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
@@ -122,51 +119,13 @@ class Simulation3DIntegral(BasePFSimulation):
         numba_parallel=True,
         **kwargs,
     ):
-        super().__init__(mesh, **kwargs)
+        super().__init__(mesh, engine=engine, numba_parallel=numba_parallel, **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)
-        if self.engine == "choclo":
-            # Check dimensions of the mesh
-            if self.mesh.dim != 3:
-                raise ValueError(
-                    f"Invalid mesh with {self.mesh.dim} dimensions. "
-                    "Only 3D meshes are supported when using 'choclo' as engine."
-                )
-            # Define jit functions
-            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(
@@ -176,15 +135,15 @@ class Simulation3DIntegral(BasePFSimulation):
                 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."
-                )
+
+        # Define jit functions
+        if self.engine == "choclo":
+            if self.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 fields(self, m):
         """
@@ -455,45 +414,6 @@ class Simulation3DIntegral(BasePFSimulation):
             index_offset += n_rows
         return sensitivity_matrix
 
-    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.mesh.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 cell nodes 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/potential_fields/magnetics/_numba_functions.py b/simpeg/potential_fields/magnetics/_numba_functions.py
new file mode 100644
index 000000000..92a5d2eac
--- /dev/null
+++ b/simpeg/potential_fields/magnetics/_numba_functions.py
@@ -0,0 +1,659 @@
+"""
+Numba functions for magnetic simulation of rectangular prisms
+"""
+
+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
+
+from .._numba_utils import kernels_in_nodes_to_cell
+
+
+def _sensitivity_mag(
+    receivers,
+    nodes,
+    sensitivity_matrix,
+    cell_nodes,
+    regional_field,
+    kernel_x,
+    kernel_y,
+    kernel_z,
+    constant_factor,
+    scalar_model,
+):
+    r"""
+    Fill the sensitivity matrix for single mag component
+
+    This function should be used with a `numba.jit` decorator, for example:
+
+    .. code::
+
+        from numba import jit
+
+        jit_sensitivity_mag = jit(nopython=True, parallel=True)(_sensitivity_mag)
+
+    Parameters
+    ----------
+    receivers : (n_receivers, 3) array
+        Array with the locations of the receivers
+    nodes : (n_active_nodes, 3) array
+        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) array
+        Array of integers, where each row contains the indices of the nodes for
+        each active cell in the mesh.
+    regional_field : (3,) array
+        Array containing the x, y and z components of the regional magnetic
+        field (uniform background field).
+    kernel_x, kernel_y, kernel_z : callable
+        Kernels used to compute the desired magnetic component. For example,
+        for computing bx we need to use ``kernel_x=kernel_ee``,
+        ``kernel_y=kernel_en``, ``kernel_z=kernel_eu``.
+    constant_factor : float
+        Constant factor that will be used to multiply each element of the
+        sensitivity matrix.
+    scalar_model : bool
+        If True, the sensitivity matrix is built to work with scalar models
+        (susceptibilities).
+        If False, the sensitivity matrix is built to work with vector models
+        (effective susceptibilities).
+
+    Notes
+    -----
+
+    About the kernel functions
+    ^^^^^^^^^^^^^^^^^^^^^^^^^^
+
+    For computing the ``bx`` component of the magnetic field we need to use the
+    following kernels:
+
+    .. code::
+
+        kernel_x=kernel_ee, kernel_y=kernel_en, kernel_z=kernel_eu
+
+
+    For computing the ``by`` component of the magnetic field we need to use the
+    following kernels:
+
+    .. code::
+
+        kernel_x=kernel_en, kernel_y=kernel_nn, kernel_z=kernel_nu
+
+    For computing the ``bz`` component of the magnetic field we need to use the
+    following kernels:
+
+    .. code::
+
+        kernel_x=kernel_eu, kernel_y=kernel_nu, kernel_z=kernel_uu
+
+    About the model array
+    ^^^^^^^^^^^^^^^^^^^^^
+
+    The ``model`` must always be a 1d array:
+
+    * If ``scalar_model`` is ``True``, then ``model`` should be a 1d array with
+      the same number of elements as active cells in the mesh. It should store
+      the magnetic susceptibilities of each active cell in SI units.
+    * If ``scalar_model`` is ``False``, then ``model`` should be a 1d array
+      with a number of elements equal to three times the active cells in the
+      mesh. It should store the components of the magnetization vector of each
+      active cell in :math:`Am^{-1}`. The order in which the components should
+      be passed are:
+          * every _easting_ component of each active cell,
+          * then every _northing_ component of each active cell,
+          * and finally every _upward_ component of each active cell.
+
+    About the sensitivity matrix
+    ^^^^^^^^^^^^^^^^^^^^^^^^^^^^
+
+    Each row of the sensitivity matrix corresponds to a single receiver
+    location.
+
+    If ``scalar_model`` is True, then each element of the row will
+    correspond to the partial derivative of the selected magnetic component
+    with respect to the susceptibility of each cell in the mesh.
+
+    If ``scalar_model`` is False, then each row can be split in three sections
+    containing:
+
+    * the partial derivatives of the selected magnetic component with respect
+      to the _x_ component of the effective susceptibility of each cell; then
+    * the partial derivatives of the selected magnetic component with respect
+      to the _y_ component of the effective susceptibility of each cell; and then
+    * the partial derivatives of the selected magnetic component with respect
+      to the _z_ component of the effective susceptibility of each cell.
+
+    So, if we call :math:`B_j` the magnetic field component on the receiver
+    :math:`j`, and :math:`\bar{\chi}^{(i)} = (\chi_x^{(i)}, \chi_y^{(i)},
+    \chi_z^{(i)})` the effective susceptibility of the active cell :math:`i`,
+    then each row of the sensitivity matrix will be:
+
+    .. math::
+
+        \left[
+            \frac{\partial B_j}{\partial \chi_x^{(1)}},
+            \dots,
+            \frac{\partial B_j}{\partial \chi_x^{(N)}},
+            \frac{\partial B_j}{\partial \chi_y^{(1)}},
+            \dots,
+            \frac{\partial B_j}{\partial \chi_y^{(N)}},
+            \frac{\partial B_j}{\partial \chi_z^{(1)}},
+            \dots,
+            \frac{\partial B_j}{\partial \chi_z^{(N)}}
+        \right]
+
+    where :math:`N` is the total number of active cells.
+    """
+    n_receivers = receivers.shape[0]
+    n_nodes = nodes.shape[0]
+    n_cells = cell_nodes.shape[0]
+    fx, fy, fz = regional_field
+    regional_field_amplitude = np.sqrt(fx**2 + fy**2 + fz**2)
+    fx /= regional_field_amplitude
+    fy /= regional_field_amplitude
+    fz /= regional_field_amplitude
+    # Evaluate kernel function on each node, for each receiver location
+    for i in prange(n_receivers):
+        # Allocate vectors for kernels evaluated on mesh nodes
+        kx, ky, kz = np.empty(n_nodes), np.empty(n_nodes), np.empty(n_nodes)
+        # Allocate small vector for the nodes indices for a given cell
+        nodes_indices = np.empty(8, dtype=cell_nodes.dtype)
+        for j in range(n_nodes):
+            dx = nodes[j, 0] - receivers[i, 0]
+            dy = nodes[j, 1] - receivers[i, 1]
+            dz = nodes[j, 2] - receivers[i, 2]
+            distance = np.sqrt(dx**2 + dy**2 + dz**2)
+            kx[j] = kernel_x(dx, dy, dz, distance)
+            ky[j] = kernel_y(dx, dy, dz, distance)
+            kz[j] = kernel_z(dx, dy, dz, distance)
+        # Compute sensitivity matrix elements from the kernel values
+        for k in range(n_cells):
+            nodes_indices = cell_nodes[k, :]
+            ux = kernels_in_nodes_to_cell(kx, nodes_indices)
+            uy = kernels_in_nodes_to_cell(ky, nodes_indices)
+            uz = kernels_in_nodes_to_cell(kz, nodes_indices)
+            if scalar_model:
+                sensitivity_matrix[i, k] = (
+                    constant_factor
+                    * regional_field_amplitude
+                    * (ux * fx + uy * fy + uz * fz)
+                )
+            else:
+                sensitivity_matrix[i, k] = (
+                    constant_factor * regional_field_amplitude * ux
+                )
+                sensitivity_matrix[i, k + n_cells] = (
+                    constant_factor * regional_field_amplitude * uy
+                )
+                sensitivity_matrix[i, k + 2 * n_cells] = (
+                    constant_factor * regional_field_amplitude * uz
+                )
+
+
+def _sensitivity_tmi(
+    receivers,
+    nodes,
+    sensitivity_matrix,
+    cell_nodes,
+    regional_field,
+    constant_factor,
+    scalar_model,
+):
+    r"""
+    Fill the sensitivity matrix for TMI
+
+    This function should be used with a `numba.jit` decorator, for example:
+
+    .. code::
+
+        from numba import jit
+
+        jit_sensitivity_tmi = jit(nopython=True, parallel=True)(_sensitivity_tmi)
+
+    Parameters
+    ----------
+    receivers : (n_receivers, 3) array
+        Array with the locations of the receivers
+    nodes : (n_active_nodes, 3) array
+        Array with the location of the mesh nodes.
+    sensitivity_matrix : array
+        Empty 2d array where the sensitivity matrix elements will be filled.
+        This could be a preallocated empty array or a slice of it.
+        The array should have a shape of ``(n_receivers, n_active_nodes)``
+        if ``scalar_model`` is True.
+        The array should have a shape of ``(n_receivers, 3 * n_active_nodes)``
+        if ``scalar_model`` is False.
+    cell_nodes : (n_active_cells, 8) array
+        Array of integers, where each row contains the indices of the nodes for
+        each active cell in the mesh.
+    regional_field : (3,) array
+        Array containing the x, y and z components of the regional magnetic
+        field (uniform background field).
+    constant_factor : float
+        Constant factor that will be used to multiply each element of the
+        sensitivity matrix.
+    scalar_model : bool
+        If True, the sensitivity matrix is build to work with scalar models
+        (susceptibilities).
+        If False, the sensitivity matrix is build to work with vector models
+        (effective susceptibilities).
+
+    Notes
+    -----
+
+    About the model array
+    ^^^^^^^^^^^^^^^^^^^^^
+
+    The ``model`` must always be a 1d array:
+
+    * If ``scalar_model`` is ``True``, then ``model`` should be a 1d array with
+      the same number of elements as active cells in the mesh. It should store
+      the magnetic susceptibilities of each active cell in SI units.
+    * If ``scalar_model`` is ``False``, then ``model`` should be a 1d array
+      with a number of elements equal to three times the active cells in the
+      mesh. It should store the components of the magnetization vector of each
+      active cell in :math:`Am^{-1}`. The order in which the components should
+      be passed are:
+          * every _easting_ component of each active cell,
+          * then every _northing_ component of each active cell,
+          * and finally every _upward_ component of each active cell.
+
+    About the sensitivity matrix
+    ^^^^^^^^^^^^^^^^^^^^^^^^^^^^
+
+    Each row of the sensitivity matrix corresponds to a single receiver
+    location.
+
+    If ``scalar_model`` is True, then each element of the row will
+    correspond to the partial derivative of the tmi
+    with respect to the susceptibility of each cell in the mesh.
+
+    If ``scalar_model`` is False, then each row can be split in three sections
+    containing:
+
+    * the partial derivatives of the tmi with respect
+      to the _x_ component of the effective susceptibility of each cell; then
+    * the partial derivatives of the tmi with respect
+      to the _y_ component of the effective susceptibility of each cell; and then
+    * the partial derivatives of the tmi with respect
+      to the _z_ component of the effective susceptibility of each cell.
+
+    So, if we call :math:`T_j` the tmi on the receiver
+    :math:`j`, and :math:`\bar{\chi}^{(i)} = (\chi_x^{(i)}, \chi_y^{(i)},
+    \chi_z^{(i)})` the effective susceptibility of the active cell :math:`i`,
+    then each row of the sensitivity matrix will be:
+
+    .. math::
+
+        \left[
+            \frac{\partial T_j}{\partial \chi_x^{(1)}},
+            \dots,
+            \frac{\partial T_j}{\partial \chi_x^{(N)}},
+            \frac{\partial T_j}{\partial \chi_y^{(1)}},
+            \dots,
+            \frac{\partial T_j}{\partial \chi_y^{(N)}},
+            \frac{\partial T_j}{\partial \chi_z^{(1)}},
+            \dots,
+            \frac{\partial T_j}{\partial \chi_z^{(N)}}
+        \right]
+
+    where :math:`N` is the total number of active cells.
+    """
+    n_receivers = receivers.shape[0]
+    n_nodes = nodes.shape[0]
+    n_cells = cell_nodes.shape[0]
+    fx, fy, fz = regional_field
+    regional_field_amplitude = np.sqrt(fx**2 + fy**2 + fz**2)
+    fx /= regional_field_amplitude
+    fy /= regional_field_amplitude
+    fz /= regional_field_amplitude
+    # Evaluate kernel function on each node, for each receiver location
+    for i in prange(n_receivers):
+        # Allocate vectors for kernels evaluated on mesh nodes
+        kxx, kyy, kzz = np.empty(n_nodes), np.empty(n_nodes), np.empty(n_nodes)
+        kxy, kxz, kyz = np.empty(n_nodes), np.empty(n_nodes), np.empty(n_nodes)
+        # Allocate small vector for the nodes indices for a given cell
+        nodes_indices = np.empty(8, dtype=cell_nodes.dtype)
+        for j in range(n_nodes):
+            dx = nodes[j, 0] - receivers[i, 0]
+            dy = nodes[j, 1] - receivers[i, 1]
+            dz = nodes[j, 2] - receivers[i, 2]
+            distance = np.sqrt(dx**2 + dy**2 + dz**2)
+            kxx[j] = choclo.prism.kernel_ee(dx, dy, dz, distance)
+            kyy[j] = choclo.prism.kernel_nn(dx, dy, dz, distance)
+            kzz[j] = choclo.prism.kernel_uu(dx, dy, dz, distance)
+            kxy[j] = choclo.prism.kernel_en(dx, dy, dz, distance)
+            kxz[j] = choclo.prism.kernel_eu(dx, dy, dz, distance)
+            kyz[j] = choclo.prism.kernel_nu(dx, dy, dz, distance)
+        # Compute sensitivity matrix elements from the kernel values
+        for k in range(n_cells):
+            nodes_indices = cell_nodes[k, :]
+            uxx = kernels_in_nodes_to_cell(kxx, nodes_indices)
+            uyy = kernels_in_nodes_to_cell(kyy, nodes_indices)
+            uzz = kernels_in_nodes_to_cell(kzz, nodes_indices)
+            uxy = kernels_in_nodes_to_cell(kxy, nodes_indices)
+            uxz = kernels_in_nodes_to_cell(kxz, nodes_indices)
+            uyz = kernels_in_nodes_to_cell(kyz, nodes_indices)
+            bx = uxx * fx + uxy * fy + uxz * fz
+            by = uxy * fx + uyy * fy + uyz * fz
+            bz = uxz * fx + uyz * fy + uzz * fz
+            # Fill the sensitivity matrix element(s) that correspond to the
+            # current active cell
+            if scalar_model:
+                sensitivity_matrix[i, k] = (
+                    constant_factor
+                    * regional_field_amplitude
+                    * (bx * fx + by * fy + bz * fz)
+                )
+            else:
+                sensitivity_matrix[i, k] = (
+                    constant_factor * regional_field_amplitude * bx
+                )
+                sensitivity_matrix[i, k + n_cells] = (
+                    constant_factor * regional_field_amplitude * by
+                )
+                sensitivity_matrix[i, k + 2 * n_cells] = (
+                    constant_factor * regional_field_amplitude * bz
+                )
+
+
+def _forward_mag(
+    receivers,
+    nodes,
+    model,
+    fields,
+    cell_nodes,
+    regional_field,
+    kernel_x,
+    kernel_y,
+    kernel_z,
+    constant_factor,
+    scalar_model,
+):
+    """
+    Forward model single magnetic component
+
+    This function should be used with a `numba.jit` decorator, for example:
+
+    .. code::
+
+        from numba import jit
+
+        jit_forward = jit(nopython=True, parallel=True)(_forward_mag)
+
+    Parameters
+    ----------
+    receivers : (n_receivers, 3) array
+        Array with the locations of the receivers
+    nodes : (n_active_nodes, 3) array
+        Array with the location of the mesh nodes.
+    model : (n_active_cells) or (3 * n_active_cells) array
+        Array containing the susceptibilities (scalar) or effective
+        susceptibilities (vector) of the active cells in the mesh, in SI
+        units.
+        Susceptibilities are expected if ``scalar_model`` is True,
+        and the array should have ``n_active_cells`` elements.
+        Effective susceptibilities are expected if ``scalar_model`` is False,
+        and the array should have ``3 * n_active_cells`` elements.
+    fields : (n_receivers) array
+        Array full of zeros where the magnetic component on each receiver will
+        be stored. This could be a preallocated array or a slice of it.
+    cell_nodes : (n_active_cells, 8) array
+        Array of integers, where each row contains the indices of the nodes for
+        each active cell in the mesh.
+    regional_field : (3,) array
+        Array containing the x, y and z components of the regional magnetic
+        field (uniform background field).
+    kernel_x, kernel_y, kernel_z : callable
+        Kernels used to compute the desired magnetic component. For example,
+        for computing bx we need to use ``kernel_x=kernel_ee``,
+        ``kernel_y=kernel_en``, ``kernel_z=kernel_eu``.
+    constant_factor : float
+        Constant factor that will be used to multiply each element of the
+        sensitivity matrix.
+    scalar_model : bool
+        If True, the forward will be computing assuming that the ``model`` has
+        susceptibilities (scalar model) for each active cell.
+        If False, the forward will be computing assuming that the ``model`` has
+        effective susceptibilities (vector model) for each active cell.
+
+    Notes
+    -----
+
+    About the kernel functions
+    ^^^^^^^^^^^^^^^^^^^^^^^^^^
+
+    For computing the ``bx`` component of the magnetic field we need to use the
+    following kernels:
+
+    .. code::
+
+        kernel_x=kernel_ee, kernel_y=kernel_en, kernel_z=kernel_eu
+
+
+    For computing the ``by`` component of the magnetic field we need to use the
+    following kernels:
+
+    .. code::
+
+        kernel_x=kernel_en, kernel_y=kernel_nn, kernel_z=kernel_nu
+
+    For computing the ``bz`` component of the magnetic field we need to use the
+    following kernels:
+
+    .. code::
+
+        kernel_x=kernel_eu, kernel_y=kernel_nu, kernel_z=kernel_uu
+
+
+    About the model array
+    ^^^^^^^^^^^^^^^^^^^^^
+
+    The ``model`` must always be a 1d array:
+
+    * If ``scalar_model`` is ``True``, then ``model`` should be a 1d array with
+      the same number of elements as active cells in the mesh. It should store
+      the magnetic susceptibilities of each active cell in SI units.
+    * If ``scalar_model`` is ``False``, then ``model`` should be a 1d array
+      with a number of elements equal to three times the active cells in the
+      mesh. It should store the components of the magnetization vector of each
+      active cell in :math:`Am^{-1}`. The order in which the components should
+      be passed are:
+          * every _easting_ component of each active cell,
+          * then every _northing_ component of each active cell,
+          * and finally every _upward_ component of each active cell.
+
+    """
+    n_receivers = receivers.shape[0]
+    n_nodes = nodes.shape[0]
+    n_cells = cell_nodes.shape[0]
+    fx, fy, fz = regional_field
+    regional_field_amplitude = np.sqrt(fx**2 + fy**2 + fz**2)
+    fx /= regional_field_amplitude
+    fy /= regional_field_amplitude
+    fz /= regional_field_amplitude
+    # Evaluate kernel function on each node, for each receiver location
+    for i in prange(n_receivers):
+        # Allocate vectors for kernels evaluated on mesh nodes
+        kx, ky, kz = np.empty(n_nodes), np.empty(n_nodes), np.empty(n_nodes)
+        # Allocate small vector for the nodes indices for a given cell
+        nodes_indices = np.empty(8, dtype=cell_nodes.dtype)
+        for j in range(n_nodes):
+            dx = nodes[j, 0] - receivers[i, 0]
+            dy = nodes[j, 1] - receivers[i, 1]
+            dz = nodes[j, 2] - receivers[i, 2]
+            distance = np.sqrt(dx**2 + dy**2 + dz**2)
+            kx[j] = kernel_x(dx, dy, dz, distance)
+            ky[j] = kernel_y(dx, dy, dz, distance)
+            kz[j] = kernel_z(dx, dy, dz, distance)
+        # Compute sensitivity matrix elements from the kernel values
+        for k in range(n_cells):
+            nodes_indices = cell_nodes[k, :]
+            ux = kernels_in_nodes_to_cell(kx, nodes_indices)
+            uy = kernels_in_nodes_to_cell(ky, nodes_indices)
+            uz = kernels_in_nodes_to_cell(kz, nodes_indices)
+            if scalar_model:
+                fields[i] += (
+                    constant_factor
+                    * model[k]
+                    * regional_field_amplitude
+                    * (ux * fx + uy * fy + uz * fz)
+                )
+            else:
+                fields[i] += (
+                    constant_factor
+                    * regional_field_amplitude
+                    * (
+                        ux * model[k]
+                        + uy * model[k + n_cells]
+                        + uz * model[k + 2 * n_cells]
+                    )
+                )
+
+
+def _forward_tmi(
+    receivers,
+    nodes,
+    model,
+    fields,
+    cell_nodes,
+    regional_field,
+    constant_factor,
+    scalar_model,
+):
+    """
+    Forward model the TMI
+
+    This function should be used with a `numba.jit` decorator, for example:
+
+    .. code::
+
+        from numba import jit
+
+        jit_forward = jit(nopython=True, parallel=True)(_forward_tmi)
+
+    Parameters
+    ----------
+    receivers : (n_receivers, 3) array
+        Array with the locations of the receivers
+    nodes : (n_active_nodes, 3) array
+        Array with the location of the mesh nodes.
+    model : (n_active_cells) or (3 * n_active_cells)
+        Array with the susceptibility (scalar model) or the effective
+        susceptibility (vector model) of each active cell in the mesh.
+        If the model is scalar, the ``model`` array should have
+        ``n_active_cells`` elements and ``scalar_model`` should be True.
+        If the model is vector, the ``model`` array should have
+        ``3 * n_active_cells`` elements and ``scalar_model`` should be False.
+    fields : (n_receivers) array
+        Array full of zeros where the TMI on each receiver will be stored. This
+        could be a preallocated array or a slice of it.
+    cell_nodes : (n_active_cells, 8) array
+        Array of integers, where each row contains the indices of the nodes for
+        each active cell in the mesh.
+    regional_field : (3,) array
+        Array containing the x, y and z components of the regional magnetic
+        field (uniform background field).
+    constant_factor : float
+        Constant factor that will be used to multiply each element of the
+        sensitivity matrix.
+    scalar_model : bool
+        If True, the sensitivity matrix is build to work with scalar models
+        (susceptibilities).
+        If False, the sensitivity matrix is build to work with vector models
+        (effective susceptibilities).
+
+    Notes
+    -----
+
+    The ``model`` must always be a 1d array:
+
+    * If ``scalar_model`` is ``True``, then ``model`` should be a 1d array with
+      the same number of elements as active cells in the mesh. It should store
+      the magnetic susceptibilities of each active cell in SI units.
+    * If ``scalar_model`` is ``False``, then ``model`` should be a 1d array
+      with a number of elements equal to three times the active cells in the
+      mesh. It should store the components of the magnetization vector of each
+      active cell in :math:`Am^{-1}`. The order in which the components should
+      be passed are:
+          * every _easting_ component of each active cell,
+          * then every _northing_ component of each active cell,
+          * and finally every _upward_ component of each active cell.
+
+    """
+    n_receivers = receivers.shape[0]
+    n_nodes = nodes.shape[0]
+    n_cells = cell_nodes.shape[0]
+    fx, fy, fz = regional_field
+    regional_field_amplitude = np.sqrt(fx**2 + fy**2 + fz**2)
+    fx /= regional_field_amplitude
+    fy /= regional_field_amplitude
+    fz /= regional_field_amplitude
+    # Evaluate kernel function on each node, for each receiver location
+    for i in prange(n_receivers):
+        # Allocate vectors for kernels evaluated on mesh nodes
+        kxx, kyy, kzz = np.empty(n_nodes), np.empty(n_nodes), np.empty(n_nodes)
+        kxy, kxz, kyz = np.empty(n_nodes), np.empty(n_nodes), np.empty(n_nodes)
+        # Allocate small vector for the nodes indices for a given cell
+        nodes_indices = np.empty(8, dtype=cell_nodes.dtype)
+        for j in range(n_nodes):
+            dx = nodes[j, 0] - receivers[i, 0]
+            dy = nodes[j, 1] - receivers[i, 1]
+            dz = nodes[j, 2] - receivers[i, 2]
+            distance = np.sqrt(dx**2 + dy**2 + dz**2)
+            kxx[j] = choclo.prism.kernel_ee(dx, dy, dz, distance)
+            kyy[j] = choclo.prism.kernel_nn(dx, dy, dz, distance)
+            kzz[j] = choclo.prism.kernel_uu(dx, dy, dz, distance)
+            kxy[j] = choclo.prism.kernel_en(dx, dy, dz, distance)
+            kxz[j] = choclo.prism.kernel_eu(dx, dy, dz, distance)
+            kyz[j] = choclo.prism.kernel_nu(dx, dy, dz, distance)
+        # Compute sensitivity matrix elements from the kernel values
+        for k in range(n_cells):
+            nodes_indices = cell_nodes[k, :]
+            uxx = kernels_in_nodes_to_cell(kxx, nodes_indices)
+            uyy = kernels_in_nodes_to_cell(kyy, nodes_indices)
+            uzz = kernels_in_nodes_to_cell(kzz, nodes_indices)
+            uxy = kernels_in_nodes_to_cell(kxy, nodes_indices)
+            uxz = kernels_in_nodes_to_cell(kxz, nodes_indices)
+            uyz = kernels_in_nodes_to_cell(kyz, nodes_indices)
+            bx = uxx * fx + uxy * fy + uxz * fz
+            by = uxy * fx + uyy * fy + uyz * fz
+            bz = uxz * fx + uyz * fy + uzz * fz
+            if scalar_model:
+                fields[i] += (
+                    constant_factor
+                    * model[k]
+                    * regional_field_amplitude
+                    * (bx * fx + by * fy + bz * fz)
+                )
+            else:
+                fields[i] += (
+                    constant_factor
+                    * regional_field_amplitude
+                    * (
+                        bx * model[k]
+                        + by * model[k + n_cells]
+                        + bz * model[k + 2 * n_cells]
+                    )
+                )
+
+
+_sensitivity_tmi_serial = jit(nopython=True, parallel=False)(_sensitivity_tmi)
+_sensitivity_tmi_parallel = jit(nopython=True, parallel=True)(_sensitivity_tmi)
+_forward_tmi_serial = jit(nopython=True, parallel=False)(_forward_tmi)
+_forward_tmi_parallel = jit(nopython=True, parallel=True)(_forward_tmi)
+_forward_mag_serial = jit(nopython=True, parallel=False)(_forward_mag)
+_forward_mag_parallel = jit(nopython=True, parallel=True)(_forward_mag)
+_sensitivity_mag_serial = jit(nopython=True, parallel=False)(_sensitivity_mag)
+_sensitivity_mag_parallel = jit(nopython=True, parallel=True)(_sensitivity_mag)
diff --git a/simpeg/potential_fields/magnetics/simulation.py b/simpeg/potential_fields/magnetics/simulation.py
index 5d3d27171..85d6e5310 100644
--- a/simpeg/potential_fields/magnetics/simulation.py
+++ b/simpeg/potential_fields/magnetics/simulation.py
@@ -1,3 +1,4 @@
+import warnings
 import numpy as np
 import scipy.sparse as sp
 from geoana.kernels import (
@@ -20,11 +21,68 @@ from ..base import BaseEquivalentSourceLayerSimulation, BasePFSimulation
 from .analytics import CongruousMagBC
 from .survey import Survey
 
+from ._numba_functions import (
+    choclo,
+    _sensitivity_tmi_parallel,
+    _sensitivity_tmi_serial,
+    _sensitivity_mag_parallel,
+    _sensitivity_mag_serial,
+    _forward_tmi_parallel,
+    _forward_tmi_serial,
+    _forward_mag_parallel,
+    _forward_mag_serial,
+)
+
+if choclo is not None:
+    CHOCLO_SUPPORTED_COMPONENTS = {"tmi", "bx", "by", "bz"}
+    CHOCLO_KERNELS = {
+        "bx": (choclo.prism.kernel_ee, choclo.prism.kernel_en, choclo.prism.kernel_eu),
+        "by": (choclo.prism.kernel_en, choclo.prism.kernel_nn, choclo.prism.kernel_nu),
+        "bz": (choclo.prism.kernel_eu, choclo.prism.kernel_nu, choclo.prism.kernel_uu),
+    }
+
 
 class Simulation3DIntegral(BasePFSimulation):
     """
-    magnetic simulation in integral form.
+    Magnetic simulation in integral form.
 
+    Parameters
+    ----------
+    mesh : discretize.TreeMesh or discretize.TensorMesh
+        Mesh use to run the magnetic simulation.
+    survey : simpeg.potential_fields.magnetics.Survey
+        Magnetic survey with information of the receivers.
+    ind_active : (n_cells) numpy.ndarray, optional
+        Array that indicates which cells in ``mesh`` are active cells.
+    chi : numpy.ndarray, optional
+        Susceptibility array for the active cells in the mesh.
+    chiMap : Mapping, optional
+        Model mapping.
+    model_type : str, optional
+        Whether the model are susceptibilities of the cells (``"scalar"``),
+        or effective susceptibilities (``"vector"``).
+    is_amplitude_data : bool, optional
+        If True, the returned fields will be the amplitude of the magnetic
+        field. If False, the fields will be returned unmodified.
+    sensitivity_dtype : numpy.dtype, optional
+        Data type that will be used to build the sensitivity matrix.
+    store_sensitivities : {"ram", "disk", "forward_only"}
+        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 : {"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.
     """
 
     chi, chiMap, chiDeriv = props.Invertible("Magnetic Susceptibility (SI)")
@@ -36,10 +94,12 @@ class Simulation3DIntegral(BasePFSimulation):
         chiMap=None,
         model_type="scalar",
         is_amplitude_data=False,
-        **kwargs
+        engine="geoana",
+        numba_parallel=True,
+        **kwargs,
     ):
         self.model_type = model_type
-        super().__init__(mesh, **kwargs)
+        super().__init__(mesh, engine=engine, numba_parallel=numba_parallel, **kwargs)
         self.chi = chi
         self.chiMap = chiMap
 
@@ -49,6 +109,28 @@ class Simulation3DIntegral(BasePFSimulation):
         self.is_amplitude_data = is_amplitude_data
         self.modelMap = self.chiMap
 
+        # 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 magnetic simulation.",
+                UserWarning,
+                stacklevel=1,
+            )
+            self.n_processes = None
+
+        if self.engine == "choclo":
+            if self.numba_parallel:
+                self._sensitivity_tmi = _sensitivity_tmi_parallel
+                self._sensitivity_mag = _sensitivity_mag_parallel
+                self._forward_tmi = _forward_tmi_parallel
+                self._forward_mag = _forward_mag_parallel
+            else:
+                self._sensitivity_tmi = _sensitivity_tmi_serial
+                self._sensitivity_mag = _sensitivity_mag_serial
+                self._forward_tmi = _forward_tmi_serial
+                self._forward_mag = _forward_mag_serial
+
     @property
     def model_type(self):
         """Type of magnetization model
@@ -103,7 +185,10 @@ class Simulation3DIntegral(BasePFSimulation):
         self.model = model
         # model = self.chiMap * model
         if self.store_sensitivities == "forward_only":
-            fields = mkvc(self.linear_operator())
+            if self.engine == "choclo":
+                fields = self._forward(self.chi)
+            else:
+                fields = mkvc(self.linear_operator())
         else:
             fields = np.asarray(
                 self.G @ self.chi.astype(self.sensitivity_dtype, copy=False)
@@ -117,7 +202,10 @@ class Simulation3DIntegral(BasePFSimulation):
     @property
     def G(self):
         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
 
@@ -494,6 +582,151 @@ class Simulation3DIntegral(BasePFSimulation):
             deletes = deletes + ["_gtg_diagonal", "_ampDeriv"]
         return deletes
 
+    def _forward(self, model):
+        """
+        Forward model the fields of active cells in the mesh on receivers.
+
+        Parameters
+        ----------
+        model : (n_active_cells) or (3 * n_active_cells) array
+            Array containing the susceptibilities (scalar) or effective
+            susceptibilities (vector) of the active cells in the mesh, in SI
+            units.
+            Susceptibilities are expected if ``model_type`` is ``"scalar"``,
+            and the array should have ``n_active_cells`` elements.
+            Effective susceptibilities are expected if ``model_type`` is
+            ``"vector"``, and the array should have ``3 * n_active_cells``
+            elements.
+
+        Returns
+        -------
+        (nD, ) array
+            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()
+        # Get regional field
+        regional_field = self.survey.source_field.b0
+        # Allocate fields array
+        fields = np.zeros(self.survey.nD, dtype=self.sensitivity_dtype)
+        # Define the constant factor
+        constant_factor = 1 / 4 / np.pi
+        # Start computing the fields
+        index_offset = 0
+        scalar_model = self.model_type == "scalar"
+        for components, receivers in self._get_components_and_receivers():
+            if not CHOCLO_SUPPORTED_COMPONENTS.issuperset(components):
+                raise NotImplementedError(
+                    f"Other components besides {CHOCLO_SUPPORTED_COMPONENTS} "
+                    "aren't implemented yet."
+                )
+            n_components = len(components)
+            n_rows = n_components * receivers.shape[0]
+            for i, component in enumerate(components):
+                vector_slice = slice(
+                    index_offset + i, index_offset + n_rows, n_components
+                )
+                if component == "tmi":
+                    self._forward_tmi(
+                        receivers,
+                        active_nodes,
+                        model,
+                        fields[vector_slice],
+                        active_cell_nodes,
+                        regional_field,
+                        constant_factor,
+                        scalar_model,
+                    )
+                else:
+                    kernel_x, kernel_y, kernel_z = CHOCLO_KERNELS[component]
+                    self._forward_mag(
+                        receivers,
+                        active_nodes,
+                        model,
+                        fields[vector_slice],
+                        active_cell_nodes,
+                        regional_field,
+                        kernel_x,
+                        kernel_y,
+                        kernel_z,
+                        constant_factor,
+                        scalar_model,
+                    )
+            index_offset += n_rows
+        return fields
+
+    def _sensitivity_matrix(self):
+        """
+        Compute the sensitivity matrix G
+
+        Returns
+        -------
+        (nD, n_active_cells) array
+        """
+        # Gather active nodes and the indices of the nodes for each active cell
+        active_nodes, active_cell_nodes = self._get_active_nodes()
+        # Get regional field
+        regional_field = self.survey.source_field.b0
+        # Allocate sensitivity matrix
+        if self.model_type == "scalar":
+            n_columns = self.nC
+        else:
+            n_columns = 3 * self.nC
+        shape = (self.survey.nD, n_columns)
+        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)
+        # Define the constant factor
+        constant_factor = 1 / 4 / np.pi
+        # Start filling the sensitivity matrix
+        index_offset = 0
+        scalar_model = self.model_type == "scalar"
+        for components, receivers in self._get_components_and_receivers():
+            if not CHOCLO_SUPPORTED_COMPONENTS.issuperset(components):
+                raise NotImplementedError(
+                    f"Other components besides {CHOCLO_SUPPORTED_COMPONENTS} "
+                    "aren't implemented yet."
+                )
+            n_components = len(components)
+            n_rows = n_components * receivers.shape[0]
+            for i, component in enumerate(components):
+                matrix_slice = slice(
+                    index_offset + i, index_offset + n_rows, n_components
+                )
+                if component == "tmi":
+                    self._sensitivity_tmi(
+                        receivers,
+                        active_nodes,
+                        sensitivity_matrix[matrix_slice, :],
+                        active_cell_nodes,
+                        regional_field,
+                        constant_factor,
+                        scalar_model,
+                    )
+                else:
+                    kernel_x, kernel_y, kernel_z = CHOCLO_KERNELS[component]
+                    self._sensitivity_mag(
+                        receivers,
+                        active_nodes,
+                        sensitivity_matrix[matrix_slice, :],
+                        active_cell_nodes,
+                        regional_field,
+                        kernel_x,
+                        kernel_y,
+                        kernel_z,
+                        constant_factor,
+                        scalar_model,
+                    )
+            index_offset += n_rows
+        return sensitivity_matrix
+
 
 class SimulationEquivalentSourceLayer(
     BaseEquivalentSourceLayerSimulation, Simulation3DIntegral
diff --git a/tutorials/04-magnetics/plot_2a_magnetics_induced.py b/tutorials/04-magnetics/plot_2a_magnetics_induced.py
index d5b90274f..67889ba62 100644
--- a/tutorials/04-magnetics/plot_2a_magnetics_induced.py
+++ b/tutorials/04-magnetics/plot_2a_magnetics_induced.py
@@ -169,8 +169,10 @@ plt.show()
 # susceptibility model using the integral formulation.
 #
 
+###############################################################################
 # Define the forward simulation. By setting the 'store_sensitivities' keyword
 # argument to "forward_only", we simulate the data without storing the sensitivities
+
 simulation = magnetics.simulation.Simulation3DIntegral(
     survey=survey,
     mesh=mesh,
@@ -178,9 +180,21 @@ simulation = magnetics.simulation.Simulation3DIntegral(
     chiMap=model_map,
     ind_active=ind_active,
     store_sensitivities="forward_only",
+    engine="choclo",
 )
 
+###############################################################################
+# .. tip::
+#
+#    Since SimPEG v0.22.0 we can use `Choclo
+#    <https://www.fatiando.org/choclo>`_ as the engine for running the magnetic
+#    simulations, which results in faster and more memory efficient runs. Just
+#    pass ``engine="choclo"`` when constructing the simulation.
+#
+
+###############################################################################
 # Compute predicted data for a susceptibility model
+
 dpred = simulation.dpred(model)
 
 # Plot
diff --git a/tutorials/04-magnetics/plot_2b_magnetics_mvi.py b/tutorials/04-magnetics/plot_2b_magnetics_mvi.py
index 0007708f8..70c962122 100644
--- a/tutorials/04-magnetics/plot_2b_magnetics_mvi.py
+++ b/tutorials/04-magnetics/plot_2b_magnetics_mvi.py
@@ -225,8 +225,10 @@ cbar.set_label(
 # in the case of remanent magnetization.
 #
 
+###############################################################################
 # Define the forward simulation. By setting the 'store_sensitivities' keyword
 # argument to "forward_only", we simulate the data without storing the sensitivities
+
 simulation = magnetics.simulation.Simulation3DIntegral(
     survey=survey,
     mesh=mesh,
@@ -236,7 +238,10 @@ simulation = magnetics.simulation.Simulation3DIntegral(
     store_sensitivities="forward_only",
 )
 
+
+###############################################################################
 # Compute predicted data for some model
+
 dpred = simulation.dpred(model)
 n_data = len(dpred)
 
diff --git a/tutorials/04-magnetics/plot_inv_2a_magnetics_induced.py b/tutorials/04-magnetics/plot_inv_2a_magnetics_induced.py
index 24430e337..92ab0691e 100644
--- a/tutorials/04-magnetics/plot_inv_2a_magnetics_induced.py
+++ b/tutorials/04-magnetics/plot_inv_2a_magnetics_induced.py
@@ -229,15 +229,27 @@ starting_model = background_susceptibility * np.ones(nC)
 # class.
 #
 
+###############################################################################
 # Define the problem. Define the cells below topography and the mapping
+
 simulation = magnetics.simulation.Simulation3DIntegral(
     survey=survey,
     mesh=mesh,
     model_type="scalar",
     chiMap=model_map,
     ind_active=active_cells,
+    engine="choclo",
 )
 
+###############################################################################
+# .. tip::
+#
+#    Since SimPEG v0.22.0 we can use `Choclo
+#    <https://www.fatiando.org/choclo>`_ as the engine for running the magnetic
+#    simulations, which results in faster and more memory efficient runs. Just
+#    pass ``engine="choclo"`` when constructing the simulation.
+#
+
 
 #######################################################################
 # Define Inverse Problem


Test patch:
diff --git a/tests/pf/test_base_pf_simulation.py b/tests/pf/test_base_pf_simulation.py
new file mode 100644
index 000000000..9cf7964ee
--- /dev/null
+++ b/tests/pf/test_base_pf_simulation.py
@@ -0,0 +1,306 @@
+"""
+Test BasePFSimulation class
+"""
+
+import pytest
+import numpy as np
+from discretize import CylindricalMesh, TensorMesh, TreeMesh
+
+import simpeg
+from simpeg.potential_fields.base import BasePFSimulation
+from simpeg.survey import BaseSurvey
+from simpeg.potential_fields import gravity, magnetics
+
+
+@pytest.fixture
+def mock_simulation_class():
+    """
+    Mock simulation class as child of BasePFSimulation
+    """
+
+    class MockSimulation(BasePFSimulation):
+        @property
+        def G(self):
+            """Define a dummy G property to avoid warnings on tests."""
+            pass
+
+    return MockSimulation
+
+
+@pytest.fixture
+def tensor_mesh():
+    """
+    Return sample TensorMesh
+    """
+    h = (3, 3, 3)
+    return TensorMesh(h)
+
+
+@pytest.fixture
+def tree_mesh():
+    """
+    Return sample TensorMesh
+    """
+    h = (4, 4, 4)
+    mesh = TreeMesh(h)
+    mesh.refine_points(points=(0, 0, 0), level=2)
+    return mesh
+
+
+@pytest.fixture
+def mock_survey_class():
+    """
+    Mock survey class as child of BaseSurvey
+    """
+
+    class MockSurvey(BaseSurvey):
+        pass
+
+    return MockSurvey
+
+
+class TestEngine:
+    """
+    Test the engine property and some of its relations with other attributes
+    """
+
+    def test_invalid_engine(self, tensor_mesh, mock_simulation_class):
+        """
+        Test if error is raised after invalid engine
+        """
+        engine = "invalid engine"
+        msg = rf"'engine' must be in \('geoana', 'choclo'\). Got '{engine}'"
+        with pytest.raises(ValueError, match=msg):
+            mock_simulation_class(tensor_mesh, engine=engine)
+
+    def test_invalid_engine_without_choclo(
+        self, tensor_mesh, mock_simulation_class, monkeypatch
+    ):
+        """
+        Test error after choosing "choclo" as engine but not being installed
+        """
+        monkeypatch.setattr(simpeg.potential_fields.base, "choclo", None)
+        engine = "choclo"
+        msg = "The choclo package couldn't be found."
+        with pytest.raises(ImportError, match=msg):
+            mock_simulation_class(tensor_mesh, engine=engine)
+
+    def test_sensitivity_path_as_dir(self, tensor_mesh, mock_simulation_class, tmpdir):
+        """
+        Test error if the sensitivity_path is a dir
+
+        Error should be raised if using ``engine=="choclo"`` and setting
+        ``store_sensitivities="disk"``.
+        """
+        sensitivity_path = str(tmpdir.mkdir("sensitivities"))
+        msg = f"The passed sensitivity_path '{sensitivity_path}' is a directory."
+        with pytest.raises(ValueError, match=msg):
+            mock_simulation_class(
+                tensor_mesh,
+                engine="choclo",
+                store_sensitivities="disk",
+                sensitivity_path=sensitivity_path,
+            )
+
+
+class TestGetActiveNodes:
+    """
+    Tests _get_active_nodes private method
+    """
+
+    def test_invalid_mesh(self, tensor_mesh, mock_simulation_class):
+        """
+        Test error on invalid mesh class
+        """
+        # Initialize base simulation with valid mesh (so we don't trigger
+        # errors in the constructor)
+        simulation = mock_simulation_class(tensor_mesh)
+        # Assign an invalid mesh to the simulation
+        simulation.mesh = CylindricalMesh(tensor_mesh.h)
+        msg = "Invalid mesh of type CylindricalMesh."
+        with pytest.raises(TypeError, match=msg):
+            simulation._get_active_nodes()
+
+    def test_no_inactive_cells_tensor(self, tensor_mesh, mock_simulation_class):
+        """
+        Test _get_active_nodes when all cells are active on a tensor mesh
+        """
+        simulation = mock_simulation_class(tensor_mesh)
+        active_nodes, active_cell_nodes = simulation._get_active_nodes()
+        np.testing.assert_equal(active_nodes, tensor_mesh.nodes)
+        np.testing.assert_equal(active_cell_nodes, tensor_mesh.cell_nodes)
+
+    def test_no_inactive_cells_tree(self, tree_mesh, mock_simulation_class):
+        """
+        Test _get_active_nodes when all cells are active on a tree mesh
+        """
+        simulation = mock_simulation_class(tree_mesh)
+        active_nodes, active_cell_nodes = simulation._get_active_nodes()
+        np.testing.assert_equal(active_nodes, tree_mesh.total_nodes)
+        np.testing.assert_equal(active_cell_nodes, tree_mesh.cell_nodes)
+
+    def test_inactive_cells_tensor(self, tensor_mesh, mock_simulation_class):
+        """
+        Test _get_active_nodes with some inactive cells on a tensor mesh
+        """
+        # Define active cells: only the first cell is active
+        active_cells = np.zeros(tensor_mesh.n_cells, dtype=bool)
+        active_cells[0] = True
+        # Initialize simulation
+        simulation = mock_simulation_class(tensor_mesh, ind_active=active_cells)
+        # Build expected active_nodes and active_cell_nodes
+        expected_active_nodes = tensor_mesh.nodes[tensor_mesh[0].nodes]
+        expected_active_cell_nodes = np.atleast_2d(np.arange(8, dtype=int))
+        # Test method
+        active_nodes, active_cell_nodes = simulation._get_active_nodes()
+        np.testing.assert_equal(active_nodes, expected_active_nodes)
+        np.testing.assert_equal(active_cell_nodes, expected_active_cell_nodes)
+
+    def test_inactive_cells_tree(self, tree_mesh, mock_simulation_class):
+        """
+        Test _get_active_nodes with some inactive cells on a tensor mesh
+        """
+        # Define active cells: only the first cell is active
+        active_cells = np.zeros(tree_mesh.n_cells, dtype=bool)
+        active_cells[0] = True
+
+        # Initialize simulation
+        simulation = mock_simulation_class(tree_mesh, ind_active=active_cells)
+
+        # Build expected active_nodes (in the right order for a single cell)
+        expected_active_nodes = [
+            [0, 0, 0],
+            [0.25, 0, 0],
+            [0, 0.25, 0],
+            [0.25, 0.25, 0],
+            [0, 0, 0.25],
+            [0.25, 0, 0.25],
+            [0, 0.25, 0.25],
+            [0.25, 0.25, 0.25],
+        ]
+
+        # Run method
+        active_nodes, active_cell_nodes = simulation._get_active_nodes()
+
+        # Check shape of active nodes and check if all of them are there
+        assert active_nodes.shape == (8, 3)
+        for node in expected_active_nodes:
+            assert node in active_nodes
+
+        # Check shape of active_cell_nodes and check if they are in the right
+        # order
+        assert active_cell_nodes.shape == (1, 8)
+        for node, node_index in zip(expected_active_nodes, active_cell_nodes[0]):
+            np.testing.assert_equal(node, active_nodes[node_index])
+
+
+class TestGetComponentsAndReceivers:
+    """
+    Test _get_components_and_receivers private method
+    """
+
+    @pytest.fixture
+    def receiver_locations(self):
+        receiver_locations = np.array([[0, 0, 0], [1, 1, 1]], dtype=np.float64)
+        return receiver_locations
+
+    @pytest.fixture
+    def gravity_survey(self, receiver_locations):
+        receiver_locations = np.array([[0, 0, 0], [1, 1, 1]], dtype=np.float64)
+        components = ["gxy", "guv"]
+        receivers = gravity.receivers.Point(
+            receiver_locations,
+            components=components,
+        )
+        # Define the SourceField and the Survey
+        source_field = gravity.sources.SourceField(receiver_list=[receivers])
+        return gravity.Survey(source_field)
+
+    @pytest.fixture
+    def magnetic_survey(self, receiver_locations):
+        receiver_locations = np.array([[0, 0, 0], [1, 1, 1]], dtype=np.float64)
+        components = ["tmi", "bx"]
+        receivers = magnetics.receivers.Point(
+            receiver_locations,
+            components=components,
+        )
+        # Define the SourceField and the Survey
+        source_field = magnetics.sources.UniformBackgroundField(
+            receiver_list=[receivers],
+            amplitude=55_000,
+            inclination=45.0,
+            declination=12.0,
+        )
+        return magnetics.Survey(source_field)
+
+    def test_missing_source_field(
+        self, tensor_mesh, mock_survey_class, mock_simulation_class
+    ):
+        """
+        Test error after missing survey in simulation
+        """
+        survey = mock_survey_class(source_list=None)
+        simulation = mock_simulation_class(tensor_mesh, survey=survey)
+        msg = "The survey '(.*)' has no 'source_field' attribute."
+        with pytest.raises(AttributeError, match=msg):
+            # need to iterate over the generator to actually test its code
+            [item for item in simulation._get_components_and_receivers()]
+
+    def test_components_and_receivers_gravity(
+        self, tensor_mesh, gravity_survey, mock_simulation_class, receiver_locations
+    ):
+        """
+        Test method on a gravity survey
+        """
+        simulation = mock_simulation_class(tensor_mesh, survey=gravity_survey)
+        components_and_receivers = tuple(
+            items for items in simulation._get_components_and_receivers()
+        )
+        # Check we have a single element in the iterator
+        assert len(components_and_receivers) == 1
+        # Check if components and receiver locations are correct
+        components, receivers = components_and_receivers[0]
+        assert components == ["gxy", "guv"]
+        np.testing.assert_equal(receivers, receiver_locations)
+
+    def test_components_and_receivers_magnetics(
+        self, tensor_mesh, magnetic_survey, mock_simulation_class, receiver_locations
+    ):
+        """
+        Test method on a magnetic survey
+        """
+        simulation = mock_simulation_class(tensor_mesh, survey=magnetic_survey)
+        components_and_receivers = tuple(
+            items for items in simulation._get_components_and_receivers()
+        )
+        # Check we have a single element in the iterator
+        assert len(components_and_receivers) == 1
+        # Check if components and receiver locations are correct
+        components, receivers = components_and_receivers[0]
+        assert components == ["tmi", "bx"]
+        np.testing.assert_equal(receivers, receiver_locations)
+
+
+class TestInvalidMeshChoclo:
+    @pytest.fixture(params=("tensormesh", "treemesh"))
+    def mesh(self, request):
+        """Sample 2D mesh."""
+        hx, hy = [(0.1, 8)], [(0.1, 8)]
+        h = (hx, hy)
+        if request.param == "tensormesh":
+            mesh = TensorMesh(h, "CC")
+        else:
+            mesh = TreeMesh(h, origin="CC")
+            mesh.finalize()
+        return mesh
+
+    def test_invalid_mesh_with_choclo(self, mesh, mock_simulation_class):
+        """
+        Test if simulation raises error when passing an invalid mesh and using choclo
+        """
+        msg = (
+            "Invalid mesh with 2 dimensions. "
+            "Only 3D meshes are supported when using 'choclo' as engine."
+        )
+        with pytest.raises(ValueError, match=msg):
+            mock_simulation_class(mesh, engine="choclo")
diff --git a/tests/pf/test_forward_Grav_Linear.py b/tests/pf/test_forward_Grav_Linear.py
index 40aff12b5..32a964710 100644
--- a/tests/pf/test_forward_Grav_Linear.py
+++ b/tests/pf/test_forward_Grav_Linear.py
@@ -329,7 +329,8 @@ class TestsGravitySimulation:
     def test_invalid_engine(self, simple_mesh):
         """Test if error is raised after invalid engine."""
         engine = "invalid engine"
-        with pytest.raises(ValueError, match=f"Invalid engine '{engine}'"):
+        msg = rf"'engine' must be in \('geoana', 'choclo'\). Got '{engine}'"
+        with pytest.raises(ValueError, match=msg):
             gravity.Simulation3DIntegral(simple_mesh, engine=engine)
 
     def test_choclo_and_n_proceesses(self, simple_mesh):
@@ -405,7 +406,7 @@ class TestsGravitySimulation:
         Check if error is raised when choclo is missing and chosen as engine.
         """
         # Monkeypatch choclo in simpeg.potential_fields.base
-        monkeypatch.setattr(simpeg.potential_fields.gravity.simulation, "choclo", None)
+        monkeypatch.setattr(simpeg.potential_fields.base, "choclo", None)
         # Check if error is raised
         msg = "The choclo package couldn't be found."
         with pytest.raises(ImportError, match=msg):
diff --git a/tests/pf/test_forward_Mag_Linear.py b/tests/pf/test_forward_Mag_Linear.py
index 45ff17ba8..9cf349823 100644
--- a/tests/pf/test_forward_Mag_Linear.py
+++ b/tests/pf/test_forward_Mag_Linear.py
@@ -1,120 +1,667 @@
-import pytest
-import unittest
+from __future__ import annotations
 
 import discretize
 import numpy as np
+import pytest
 from geoana.em.static import MagneticPrism
 from scipy.constants import mu_0
 
+import simpeg
 from simpeg import maps, utils
 from simpeg.potential_fields import magnetics as mag
 
 
-def test_ana_mag_forward():
-    nx = 5
-    ny = 5
-
-    h0_amplitude, h0_inclination, h0_declination = (50000.0, 60.0, 250.0)
-    b0 = mag.analytics.IDTtoxyz(-h0_inclination, h0_declination, h0_amplitude)
-    chi1 = 0.01
-    chi2 = 0.02
-
-    # Define a mesh
-    cs = 0.2
-    hxind = [(cs, 41)]
-    hyind = [(cs, 41)]
-    hzind = [(cs, 41)]
-    mesh = discretize.TensorMesh([hxind, hyind, hzind], "CCC")
+def get_block_inds(grid: np.ndarray, block: np.ndarray) -> np.ndarray:
+    """
+    Get the indices for a block
+
+    Parameters
+    ----------
+    grid : np.ndarray
+        (n, 3) array of xyz locations
+    block : np.ndarray
+        (3, 2) array of (xmin, xmax), (ymin, ymax), (zmin, zmax) dimensions of
+        the block.
+
+    Returns
+    -------
+    np.ndarray
+        boolean array of indices corresponding to the 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])
+    )
 
-    # 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])
+def create_block_model(
+    mesh: discretize.TensorMesh,
+    blocks: tuple[np.ndarray, ...],
+    block_params: tuple[float, ...] | tuple[np.ndarray, ...],
+) -> tuple[np.ndarray, np.ndarray]:
+    """
+    Create a magnetic model from a sequence of blocks
+
+    Parameters
+    ----------
+    mesh : discretize.TensorMesh
+        TensorMesh object to put the model on
+    blocks : Tuple[np.ndarray, ...]
+        Tuple of block definitions (each element is (3, 2) array of
+        (xmin, xmax), (ymin, ymax), (zmin, zmax) dimensions of the block)
+    block_params : Tuple[float, ...]
+        Tuple of parameters to assign for each block. Must be the same length
+        as ``blocks``.
+
+    Returns
+    -------
+    Tuple[np.ndarray, np.ndarray]
+        Tuple of the magnetic model and active_cells (a boolean array)
+
+    Raises
+    ------
+    ValueError
+        if ``blocks`` and ``block_params`` have incompatible dimensions
+    """
+    if len(blocks) != len(block_params):
+        raise ValueError(
+            "'blocks' and 'block_params' must have the same number of elements"
+        )
+    model = np.zeros((mesh.n_cells, np.atleast_1d(block_params[0]).shape[0]))
+    for block, params in zip(blocks, block_params):
+        block_ind = get_block_inds(mesh.cell_centers, block)
+        model[block_ind] = params
+    active_cells = np.any(np.abs(model) > 0, axis=1)
+    return model.squeeze(), active_cells
+
+
+def create_mag_survey(
+    components: list[str],
+    receiver_locations: np.ndarray,
+    inducing_field_params: tuple[float, float, float],
+) -> mag.Survey:
+    """
+    create a magnetic Survey
+
+    Parameters
+    ----------
+    components : List[str]
+        List of components to model
+    receiver_locations : np.ndarray
+        (n, 3) array of xyz receiver locations
+    inducing_field_params : Tuple[float, float, float]
+        amplitude, inclination, and declination of the inducing field
+
+    Returns
+    -------
+    mag.Survey
+        a magnetic Survey instance
+    """
+
+    receivers = mag.Point(receiver_locations, components=components)
+    strenght, inclination, declination = inducing_field_params
+    source_field = mag.UniformBackgroundField(
+        receiver_list=[receivers],
+        amplitude=strenght,
+        inclination=inclination,
+        declination=declination,
+    )
+    return mag.Survey(source_field)
+
+
+class TestsMagSimulation:
+    """
+    Test mag simulation against the analytic solutions single prisms
+    """
+
+    @pytest.fixture
+    def mag_mesh(self) -> discretize.TensorMesh:
+        """
+        a small tensor mesh for testing magnetic simulations
+
+        Returns
+        -------
+        discretize.TensorMesh
+            the tensor mesh for testing
+        """
+        # Define a mesh
+        cs = 0.2
+        hxind = [(cs, 41)]
+        hyind = [(cs, 41)]
+        hzind = [(cs, 41)]
+        mesh = discretize.TensorMesh([hxind, hyind, hzind], "CCC")
+        return mesh
+
+    @pytest.fixture
+    def two_blocks(self) -> tuple[np.ndarray, np.ndarray]:
+        """
+        The parameters defining two blocks
+
+        Returns
+        -------
+        Tuple[np.ndarray, np.ndarray]
+            Tuple of (3, 2) arrays of (xmin, xmax), (ymin, ymax), (zmin, zmax)
+            dimensions of each block.
+        """
+        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]])
+        return block1, block2
+
+    @pytest.fixture
+    def receiver_locations(self) -> np.ndarray:
+        """
+        a grid of receivers for testing
+
+        Returns
+        -------
+        np.ndarray
+            (n, 3) array of receiver locations
+        """
+        # Create plane of observations
+        nx, ny = 5, 5
+        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
+        return np.c_[X.reshape(-1), Y.reshape(-1), Z.reshape(-1)]
+
+    @pytest.fixture
+    def inducing_field(
+        self,
+    ) -> tuple[tuple[float, float, float], tuple[float, float, float]]:
+        """
+        inducing field
+
+        Return inducing field as amplitude and angles and as vector components.
+
+        Returns
+        -------
+        tuple[tuple[float, float, float], tuple[float, float, float]]
+            (amplitude, inclination, declination), (b_x, b_y, b_z)
+        """
+        h0_amplitude, h0_inclination, h0_declination = (50000.0, 60.0, 250.0)
+        b0 = mag.analytics.IDTtoxyz(-h0_inclination, h0_declination, h0_amplitude)
+        return (h0_amplitude, h0_inclination, h0_declination), b0
+
+    @pytest.mark.parametrize(
+        "engine,parallel_kwargs",
+        [
+            ("geoana", {"n_processes": None}),
+            ("geoana", {"n_processes": 1}),
+            ("choclo", {"numba_parallel": False}),
+            ("choclo", {"numba_parallel": True}),
+        ],
+        ids=["geoana_serial", "geoana_parallel", "choclo_serial", "choclo_parallel"],
+    )
+    @pytest.mark.parametrize("store_sensitivities", ("ram", "disk", "forward_only"))
+    def test_magnetic_field_and_tmi_w_susceptibility(
+        self,
+        engine,
+        parallel_kwargs,
+        store_sensitivities,
+        tmp_path,
+        mag_mesh,
+        two_blocks,
+        receiver_locations,
+        inducing_field,
+    ):
+        """
+        Test forwarding the magnetic field and tmi (with susceptibility as model)
+        """
+        inducing_field_params, b0 = inducing_field
+
+        chi1 = 0.01
+        chi2 = 0.02
+        model, active_cells = create_block_model(mag_mesh, two_blocks, (chi1, chi2))
+        model_reduced = model[active_cells]
+        # Create reduced identity map for Linear Problem
+        identity_map = maps.IdentityMap(nP=int(sum(active_cells)))
+
+        survey = create_mag_survey(
+            components=["bx", "by", "bz", "tmi"],
+            receiver_locations=receiver_locations,
+            inducing_field_params=inducing_field_params,
         )
 
-    block1_inds = get_block_inds(mesh.cell_centers, block1)
-    block2_inds = get_block_inds(mesh.cell_centers, block2)
-
-    model = np.zeros(mesh.n_cells)
-    model[block1_inds] = chi1
-    model[block2_inds] = chi2
+        sim = mag.Simulation3DIntegral(
+            mag_mesh,
+            survey=survey,
+            chiMap=identity_map,
+            ind_active=active_cells,
+            sensitivity_path=str(tmp_path / f"{engine}"),
+            store_sensitivities=store_sensitivities,
+            engine=engine,
+            **parallel_kwargs,
+        )
 
-    active_cells = model != 0.0
-    model_reduced = model[active_cells]
+        data = sim.dpred(model_reduced)
+        d_x = data[0::4]
+        d_y = data[1::4]
+        d_z = data[2::4]
+        d_t = data[3::4]
+
+        # Compute analytical response from magnetic prism
+        block1, block2 = two_blocks
+        prism_1 = MagneticPrism(block1[:, 0], block1[:, 1], chi1 * b0 / mu_0)
+        prism_2 = MagneticPrism(block2[:, 0], block2[:, 1], -chi1 * b0 / mu_0)
+        prism_3 = MagneticPrism(block2[:, 0], block2[:, 1], chi2 * b0 / mu_0)
+
+        d = (
+            prism_1.magnetic_flux_density(receiver_locations)
+            + prism_2.magnetic_flux_density(receiver_locations)
+            + prism_3.magnetic_flux_density(receiver_locations)
+        )
 
-    # Create reduced identity map for Linear Problem
-    idenMap = maps.IdentityMap(nP=int(sum(active_cells)))
+        # TMI projection
+        tmi = sim.tmi_projection
+        d_t2 = d_x * tmi[0] + d_y * tmi[1] + d_z * tmi[2]
+
+        # Check results
+        rtol, atol = 1e-7, 1e-6
+        np.testing.assert_allclose(
+            d_t, d_t2, rtol=rtol, atol=atol
+        )  # double check internal projection
+        np.testing.assert_allclose(d_x, d[:, 0], rtol=rtol, atol=atol)
+        np.testing.assert_allclose(d_y, d[:, 1], rtol=rtol, atol=atol)
+        np.testing.assert_allclose(d_z, d[:, 2], rtol=rtol, atol=atol)
+        np.testing.assert_allclose(d_t, d @ tmi, rtol=rtol, atol=atol)
+
+    @pytest.mark.parametrize(
+        "engine, parallel_kwargs",
+        [
+            ("geoana", {"n_processes": None}),
+            ("geoana", {"n_processes": 1}),
+            ("choclo", {"numba_parallel": False}),
+            ("choclo", {"numba_parallel": True}),
+        ],
+        ids=["geoana_serial", "geoana_parallel", "choclo_serial", "choclo_parallel"],
+    )
+    @pytest.mark.parametrize("store_sensitivities", ("ram", "disk", "forward_only"))
+    def test_magnetic_gradiometry_w_susceptibility(
+        self,
+        engine,
+        parallel_kwargs,
+        store_sensitivities,
+        tmp_path,
+        mag_mesh,
+        two_blocks,
+        receiver_locations,
+        inducing_field,
+    ):
+        """
+        Test magnetic gradiometry components (with susceptibility as model)
+        """
+        inducing_field_params, b0 = inducing_field
+        chi1 = 0.01
+        chi2 = 0.02
+        model, active_cells = create_block_model(mag_mesh, two_blocks, (chi1, chi2))
+        model_reduced = model[active_cells]
+        # Create reduced identity map for Linear Problem
+        identity_map = maps.IdentityMap(nP=int(sum(active_cells)))
+
+        survey = create_mag_survey(
+            components=["bxx", "bxy", "bxz", "byy", "byz", "bzz"],
+            receiver_locations=receiver_locations,
+            inducing_field_params=inducing_field_params,
+        )
+        sim = mag.Simulation3DIntegral(
+            mag_mesh,
+            survey=survey,
+            chiMap=identity_map,
+            ind_active=active_cells,
+            sensitivity_path=str(tmp_path / f"{engine}"),
+            store_sensitivities=store_sensitivities,
+            engine=engine,
+            **parallel_kwargs,
+        )
+        if engine == "choclo":
+            # gradient simulation not implemented for choclo yet
+            with pytest.raises(NotImplementedError):
+                data = sim.dpred(model_reduced)
+        else:
+            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 magnetic prism
+            block1, block2 = two_blocks
+            prism_1 = MagneticPrism(block1[:, 0], block1[:, 1], chi1 * b0 / mu_0)
+            prism_2 = MagneticPrism(block2[:, 0], block2[:, 1], -chi1 * b0 / mu_0)
+            prism_3 = MagneticPrism(block2[:, 0], block2[:, 1], chi2 * b0 / mu_0)
+
+            d = (
+                prism_1.magnetic_field_gradient(receiver_locations)
+                + prism_2.magnetic_field_gradient(receiver_locations)
+                + prism_3.magnetic_field_gradient(receiver_locations)
+            ) * mu_0
+
+            # Check results
+            rtol, atol = 5e-7, 1e-6
+            np.testing.assert_allclose(d_xx, d[..., 0, 0], rtol=rtol, atol=atol)
+            np.testing.assert_allclose(d_xy, d[..., 0, 1], rtol=rtol, atol=atol)
+            np.testing.assert_allclose(d_xz, d[..., 0, 2], rtol=rtol, atol=atol)
+            np.testing.assert_allclose(d_yy, d[..., 1, 1], rtol=rtol, atol=atol)
+            np.testing.assert_allclose(d_yz, d[..., 1, 2], rtol=rtol, atol=atol)
+            np.testing.assert_allclose(d_zz, d[..., 2, 2], rtol=rtol, atol=atol)
+
+    @pytest.mark.parametrize(
+        "engine, parallel_kwargs",
+        [
+            ("geoana", {"n_processes": None}),
+            ("geoana", {"n_processes": 1}),
+            ("choclo", {"numba_parallel": False}),
+            ("choclo", {"numba_parallel": True}),
+        ],
+        ids=["geoana_serial", "geoana_parallel", "choclo_serial", "choclo_parallel"],
+    )
+    @pytest.mark.parametrize("store_sensitivities", ("ram", "disk", "forward_only"))
+    def test_magnetic_vector_and_tmi_w_magnetization(
+        self,
+        engine,
+        parallel_kwargs,
+        store_sensitivities,
+        tmp_path,
+        mag_mesh,
+        two_blocks,
+        receiver_locations,
+        inducing_field,
+    ):
+        """
+        Test magnetic vector and TMI (using magnetization vectors as model)
+        """
+        inducing_field_params, b0 = inducing_field
+        M1 = (utils.mat_utils.dip_azimuth2cartesian(45, -40) * 0.05).squeeze()
+        M2 = (utils.mat_utils.dip_azimuth2cartesian(120, 32) * 0.1).squeeze()
+
+        model, active_cells = create_block_model(mag_mesh, two_blocks, (M1, M2))
+        model_reduced = model[active_cells].reshape(-1, order="F")
+        # Create reduced identity map for Linear Problem
+        identity_map = maps.IdentityMap(nP=int(sum(active_cells)) * 3)
+
+        survey = create_mag_survey(
+            components=["bx", "by", "bz", "tmi"],
+            receiver_locations=receiver_locations,
+            inducing_field_params=inducing_field_params,
+        )
 
-    # 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)]
-    components = ["bx", "by", "bz", "tmi"]
+        sim = mag.Simulation3DIntegral(
+            mag_mesh,
+            survey=survey,
+            chiMap=identity_map,
+            ind_active=active_cells,
+            sensitivity_path=str(tmp_path / f"{engine}"),
+            store_sensitivities=store_sensitivities,
+            model_type="vector",
+            engine=engine,
+            **parallel_kwargs,
+        )
 
-    rxLoc = mag.Point(locXyz, components=components)
-    srcField = mag.UniformBackgroundField(
-        receiver_list=[rxLoc],
-        amplitude=h0_amplitude,
-        inclination=h0_inclination,
-        declination=h0_declination,
-    )
-    survey = mag.Survey(srcField)
+        data = sim.dpred(model_reduced).reshape(-1, 4)
 
-    # Creat reduced identity map for Linear Problem
-    idenMap = maps.IdentityMap(nP=int(sum(active_cells)))
+        # Compute analytical response from magnetic prism
+        block1, block2 = two_blocks
+        prism_1 = MagneticPrism(
+            block1[:, 0], block1[:, 1], M1 * np.linalg.norm(b0) / mu_0
+        )
+        prism_2 = MagneticPrism(
+            block2[:, 0], block2[:, 1], -M1 * np.linalg.norm(b0) / mu_0
+        )
+        prism_3 = MagneticPrism(
+            block2[:, 0], block2[:, 1], M2 * np.linalg.norm(b0) / mu_0
+        )
 
-    sim = mag.Simulation3DIntegral(
-        mesh,
-        survey=survey,
-        chiMap=idenMap,
-        ind_active=active_cells,
-        store_sensitivities="forward_only",
-        n_processes=None,
+        d = (
+            prism_1.magnetic_flux_density(receiver_locations)
+            + prism_2.magnetic_flux_density(receiver_locations)
+            + prism_3.magnetic_flux_density(receiver_locations)
+        )
+        tmi = sim.tmi_projection
+
+        # Check results
+        rtol, atol = 5e-7, 1e-6
+        np.testing.assert_allclose(data[:, 0], d[:, 0], rtol=rtol, atol=atol)
+        np.testing.assert_allclose(data[:, 1], d[:, 1], rtol=rtol, atol=atol)
+        np.testing.assert_allclose(data[:, 2], d[:, 2], rtol=rtol, atol=atol)
+        np.testing.assert_allclose(data[:, 3], d @ tmi, rtol=rtol, atol=atol)
+
+    @pytest.mark.parametrize(
+        "engine, parallel_kwargs",
+        [
+            ("geoana", {"n_processes": None}),
+            ("geoana", {"n_processes": 1}),
+            ("choclo", {"numba_parallel": False}),
+            ("choclo", {"numba_parallel": True}),
+        ],
+        ids=["geoana_serial", "geoana_parallel", "choclo_serial", "choclo_parallel"],
     )
+    @pytest.mark.parametrize("store_sensitivities", ("ram", "disk", "forward_only"))
+    def test_magnetic_field_amplitude_w_magnetization(
+        self,
+        engine,
+        parallel_kwargs,
+        store_sensitivities,
+        tmp_path,
+        mag_mesh,
+        two_blocks,
+        receiver_locations,
+        inducing_field,
+    ):
+        """
+        Test magnetic field amplitude (using magnetization vectors as model)
+        """
+        inducing_field_params, b0 = inducing_field
+        M1 = (utils.mat_utils.dip_azimuth2cartesian(45, -40) * 0.05).squeeze()
+        M2 = (utils.mat_utils.dip_azimuth2cartesian(120, 32) * 0.1).squeeze()
+
+        model, active_cells = create_block_model(mag_mesh, two_blocks, (M1, M2))
+        model_reduced = model[active_cells].reshape(-1, order="F")
+        # Create reduced identity map for Linear Problem
+        identity_map = maps.IdentityMap(nP=int(sum(active_cells)) * 3)
+
+        survey = create_mag_survey(
+            components=["bx", "by", "bz"],
+            receiver_locations=receiver_locations,
+            inducing_field_params=inducing_field_params,
+        )
 
-    data = sim.dpred(model_reduced)
-    d_x = data[0::4]
-    d_y = data[1::4]
-    d_z = data[2::4]
-    d_t = data[3::4]
-
-    tmi = sim.tmi_projection
-    d_t2 = d_x * tmi[0] + d_y * tmi[1] + d_z * tmi[2]
-    np.testing.assert_allclose(d_t, d_t2)  # double check internal projection
+        sim = mag.Simulation3DIntegral(
+            mag_mesh,
+            survey=survey,
+            chiMap=identity_map,
+            ind_active=active_cells,
+            sensitivity_path=str(tmp_path / f"{engine}"),
+            store_sensitivities=store_sensitivities,
+            model_type="vector",
+            is_amplitude_data=True,
+            engine=engine,
+            **parallel_kwargs,
+        )
 
-    # Compute analytical response from magnetic prism
-    prism_1 = MagneticPrism(block1[:, 0], block1[:, 1], chi1 * b0 / mu_0)
-    prism_2 = MagneticPrism(block2[:, 0], block2[:, 1], -chi1 * b0 / mu_0)
-    prism_3 = MagneticPrism(block2[:, 0], block2[:, 1], chi2 * b0 / mu_0)
+        data = sim.dpred(model_reduced)
 
-    d = (
-        prism_1.magnetic_flux_density(locXyz)
-        + prism_2.magnetic_flux_density(locXyz)
-        + prism_3.magnetic_flux_density(locXyz)
-    )
+        # Compute analytical response from magnetic prism
+        block1, block2 = two_blocks
+        prism_1 = MagneticPrism(
+            block1[:, 0], block1[:, 1], M1 * np.linalg.norm(b0) / mu_0
+        )
+        prism_2 = MagneticPrism(
+            block2[:, 0], block2[:, 1], -M1 * np.linalg.norm(b0) / mu_0
+        )
+        prism_3 = MagneticPrism(
+            block2[:, 0], block2[:, 1], M2 * np.linalg.norm(b0) / mu_0
+        )
 
-    np.testing.assert_allclose(d_x, d[:, 0])
-    np.testing.assert_allclose(d_y, d[:, 1])
-    np.testing.assert_allclose(d_z, d[:, 2])
-    np.testing.assert_allclose(d_t, d @ tmi)
+        d = (
+            prism_1.magnetic_flux_density(receiver_locations)
+            + prism_2.magnetic_flux_density(receiver_locations)
+            + prism_3.magnetic_flux_density(receiver_locations)
+        )
+        d_amp = np.linalg.norm(d, axis=1)
+
+        # Check results
+        rtol, atol = 5e-7, 1e-6
+        np.testing.assert_allclose(data, d_amp, 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,
+        mag_mesh,
+        receiver_locations,
+        tmp_path,
+    ):
+        """Test sensitivity_dtype."""
+        # Create survey
+        receivers = mag.Point(receiver_locations, components="tmi")
+        sources = mag.UniformBackgroundField(
+            [receivers], amplitude=50_000, inclination=45, declination=10
+        )
+        survey = mag.Survey(sources)
+        # Create reduced identity map for Linear Problem
+        active_cells = np.ones(mag_mesh.n_cells, dtype=bool)
+        idenMap = maps.IdentityMap(nP=mag_mesh.n_cells)
+        # Create simulation
+        sensitivity_path = tmp_path
+        if engine == "choclo":
+            sensitivity_path /= "dummy"
+        simulation = mag.Simulation3DIntegral(
+            mag_mesh,
+            survey=survey,
+            chiMap=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, mag_mesh, invalid_dtype):
+        """
+        Test invalid sensitivity_dtype assignment
+        """
+        simulation = mag.Simulation3DIntegral(mag_mesh)
+        # 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, mag_mesh):
+        """Test if error is raised after invalid engine."""
+        engine = "invalid engine"
+        msg = rf"'engine' must be in \('geoana', 'choclo'\). Got '{engine}'"
+        with pytest.raises(ValueError, match=msg):
+            mag.Simulation3DIntegral(mag_mesh, engine=engine)
+
+    def test_choclo_and_n_proceesses(self, mag_mesh):
+        """Check if warning is raised after passing n_processes with choclo engine."""
+        msg = "The 'n_processes' will be ignored when selecting 'choclo'"
+        with pytest.warns(UserWarning, match=msg):
+            simulation = mag.Simulation3DIntegral(
+                mag_mesh, 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, mag_mesh, tmp_path):
+        """
+        Check if error is raised when sensitivity_path is a dir with choclo engine.
+        """
+        # 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):
+            mag.Simulation3DIntegral(
+                mag_mesh,
+                store_sensitivities="disk",
+                sensitivity_path=str(sensitivity_path),
+                engine="choclo",
+            )
+
+    def test_sensitivities_on_disk(self, mag_mesh, receiver_locations, tmp_path):
+        """
+        Test if sensitivity matrix is correctly being stored in disk when asked
+        """
+        # Build survey
+        survey = create_mag_survey(
+            components=["tmi"],
+            receiver_locations=receiver_locations,
+            inducing_field_params=(50000.0, 20.0, 45.0),
+        )
+        # Build simulation
+        sensitivities_path = tmp_path / "sensitivities"
+        simulation = mag.Simulation3DIntegral(
+            mesh=mag_mesh,
+            survey=survey,
+            store_sensitivities="disk",
+            sensitivity_path=str(sensitivities_path),
+            engine="choclo",
+        )
+        simulation.G
+        # Check if sensitivity matrix was stored in disk and is a memmap
+        assert sensitivities_path.is_file()
+        assert type(simulation.G) is np.memmap
+
+    def test_sensitivities_on_ram(self, mag_mesh, receiver_locations, tmp_path):
+        """
+        Test if sensitivity matrix is correctly being allocated in memory when asked
+        """
+        # Build survey
+        survey = create_mag_survey(
+            components=["tmi"],
+            receiver_locations=receiver_locations,
+            inducing_field_params=(50000.0, 20.0, 45.0),
+        )
+        # Build simulation
+        simulation = mag.Simulation3DIntegral(
+            mesh=mag_mesh,
+            survey=survey,
+            store_sensitivities="ram",
+            engine="choclo",
+        )
+        simulation.G
+        # Check if sensitivity matrix is a Numpy array (stored in memory)
+        assert type(simulation.G) is np.ndarray
+
+    def test_choclo_missing(self, mag_mesh, monkeypatch):
+        """
+        Check if error is raised when choclo is missing and chosen as engine.
+        """
+        # Monkeypatch choclo in simpeg.potential_fields.base
+        monkeypatch.setattr(simpeg.potential_fields.base, "choclo", None)
+        # Check if error is raised
+        msg = "The choclo package couldn't be found."
+        with pytest.raises(ImportError, match=msg):
+            mag.Simulation3DIntegral(mag_mesh, engine="choclo")
 
 
 def test_ana_mag_tmi_grad_forward():
+    """
+    Test TMI gradiometry using susceptibilities as model
+    """
     nx = 61
     ny = 61
 
-    H0 = (50000.0, 60.0, 250.0)
-    b0 = mag.analytics.IDTtoxyz(-H0[1], H0[2], H0[0])
+    h0_amplitude, h0_inclination, h0_declination = (50000.0, 60.0, 250.0)
+    b0 = mag.analytics.IDTtoxyz(-h0_inclination, h0_declination, h0_amplitude)
     chi1 = 0.01
     chi2 = 0.02
 
@@ -164,11 +711,14 @@ def test_ana_mag_tmi_grad_forward():
 
     rxLoc = mag.Point(locXyz, components=components)
     srcField = mag.UniformBackgroundField(
-        [rxLoc], amplitude=H0[0], inclination=H0[1], declination=H0[2]
+        receiver_list=[rxLoc],
+        amplitude=h0_amplitude,
+        inclination=h0_inclination,
+        declination=h0_declination,
     )
     survey = mag.Survey(srcField)
 
-    # Creat reduced identity map for Linear Problem
+    # Create reduced identity map for Linear Problem
     idenMap = maps.IdentityMap(nP=int(sum(active_cells)))
 
     sim = mag.Simulation3DIntegral(
@@ -196,9 +746,15 @@ def test_ana_mag_tmi_grad_forward():
         + prism_2.magnetic_field_gradient(locXyz)
         + prism_3.magnetic_field_gradient(locXyz)
     ) * mu_0
-    tmi_x = (d[:, 0, 0] * b0[0] + d[:, 0, 1] * b0[1] + d[:, 0, 2] * b0[2]) / H0[0]
-    tmi_y = (d[:, 1, 0] * b0[0] + d[:, 1, 1] * b0[1] + d[:, 1, 2] * b0[2]) / H0[0]
-    tmi_z = (d[:, 2, 0] * b0[0] + d[:, 2, 1] * b0[1] + d[:, 2, 2] * b0[2]) / H0[0]
+    tmi_x = (
+        d[:, 0, 0] * b0[0] + d[:, 0, 1] * b0[1] + d[:, 0, 2] * b0[2]
+    ) / h0_amplitude
+    tmi_y = (
+        d[:, 1, 0] * b0[0] + d[:, 1, 1] * b0[1] + d[:, 1, 2] * b0[2]
+    ) / h0_amplitude
+    tmi_z = (
+        d[:, 2, 0] * b0[0] + d[:, 2, 1] * b0[1] + d[:, 2, 2] * b0[2]
+    ) / h0_amplitude
     np.testing.assert_allclose(d_x, tmi_x, rtol=1e-10, atol=1e-12)
     np.testing.assert_allclose(d_y, tmi_y, rtol=1e-10, atol=1e-12)
     np.testing.assert_allclose(d_z, tmi_z, rtol=1e-10, atol=1e-12)
@@ -219,283 +775,40 @@ def test_ana_mag_tmi_grad_forward():
     )
 
 
-def test_ana_mag_grad_forward():
-    nx = 5
-    ny = 5
-
-    h0_amplitude, h0_inclination, h0_declination = (50000.0, 60.0, 250.0)
-    b0 = mag.analytics.IDTtoxyz(-h0_inclination, h0_declination, h0_amplitude)
-    chi1 = 0.01
-    chi2 = 0.02
-
-    # 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])
-        )
-
-    block1_inds = get_block_inds(mesh.cell_centers, block1)
-    block2_inds = get_block_inds(mesh.cell_centers, block2)
-
-    model = np.zeros(mesh.n_cells)
-    model[block1_inds] = chi1
-    model[block2_inds] = chi2
-
-    active_cells = model != 0.0
-    model_reduced = model[active_cells]
-
-    # Create reduced identity map for Linear Problem
-    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)]
-    components = ["bxx", "bxy", "bxz", "byy", "byz", "bzz"]
-
-    rxLoc = mag.Point(locXyz, components=components)
-    srcField = mag.UniformBackgroundField(
-        [rxLoc],
-        amplitude=h0_amplitude,
-        inclination=h0_inclination,
-        declination=h0_declination,
-    )
-    survey = mag.Survey(srcField)
-
-    # Creat reduced identity map for Linear Problem
-    idenMap = maps.IdentityMap(nP=int(sum(active_cells)))
-
-    sim = mag.Simulation3DIntegral(
-        mesh,
-        survey=survey,
-        chiMap=idenMap,
-        ind_active=active_cells,
-        store_sensitivities="forward_only",
-        n_processes=None,
-    )
-
-    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 magnetic prism
-    prism_1 = MagneticPrism(block1[:, 0], block1[:, 1], chi1 * b0 / mu_0)
-    prism_2 = MagneticPrism(block2[:, 0], block2[:, 1], -chi1 * b0 / mu_0)
-    prism_3 = MagneticPrism(block2[:, 0], block2[:, 1], chi2 * b0 / mu_0)
-
-    d = (
-        prism_1.magnetic_field_gradient(locXyz)
-        + prism_2.magnetic_field_gradient(locXyz)
-        + prism_3.magnetic_field_gradient(locXyz)
-    ) * mu_0
-
-    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)
-
-
-def test_ana_mag_vec_forward():
-    nx = 5
-    ny = 5
-
-    h0_amplitude, h0_inclination, h0_declination = (50000.0, 60.0, 250.0)
-    b0 = mag.analytics.IDTtoxyz(-h0_inclination, h0_declination, h0_amplitude)
-
-    M1 = utils.mat_utils.dip_azimuth2cartesian(45, -40) * 0.05
-    M2 = utils.mat_utils.dip_azimuth2cartesian(120, 32) * 0.1
-
-    # 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])
+class TestInvalidMeshChoclo:
+    @pytest.fixture(params=("tensormesh", "treemesh"))
+    def mesh(self, request):
+        """Sample 2D mesh."""
+        hx, hy = [(0.1, 8)], [(0.1, 8)]
+        h = (hx, hy)
+        if request.param == "tensormesh":
+            mesh = discretize.TensorMesh(h, "CC")
+        else:
+            mesh = discretize.TreeMesh(h, origin="CC")
+            mesh.finalize()
+        return mesh
+
+    def test_invalid_mesh_with_choclo(self, mesh):
+        """
+        Test if simulation raises error when passing an invalid mesh and using choclo
+        """
+        # Build survey
+        receivers_locations = np.array([[0, 0, 0]])
+        receivers = mag.Point(receivers_locations)
+        sources = mag.UniformBackgroundField(
+            receiver_list=[receivers],
+            amplitude=50_000,
+            inclination=45.0,
+            declination=12.0,
         )
-
-    block1_inds = get_block_inds(mesh.cell_centers, block1)
-    block2_inds = get_block_inds(mesh.cell_centers, block2)
-
-    model = np.zeros((mesh.n_cells, 3))
-    model[block1_inds] = M1
-    model[block2_inds] = M2
-
-    active_cells = np.any(model != 0.0, axis=1)
-    model_reduced = model[active_cells].reshape(-1, order="F")
-
-    # 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)]
-    components = ["bx", "by", "bz", "tmi"]
-
-    rxLoc = mag.Point(locXyz, components=components)
-    srcField = mag.UniformBackgroundField(
-        receiver_list=[rxLoc],
-        amplitude=h0_amplitude,
-        inclination=h0_inclination,
-        declination=h0_declination,
-    )
-    survey = mag.Survey(srcField)
-
-    # Create reduced identity map for Linear Problem
-    idenMap = maps.IdentityMap(nP=int(sum(active_cells)) * 3)
-
-    sim = mag.Simulation3DIntegral(
-        mesh,
-        survey=survey,
-        chiMap=idenMap,
-        ind_active=active_cells,
-        store_sensitivities="forward_only",
-        model_type="vector",
-        n_processes=None,
-    )
-
-    data = sim.dpred(model_reduced).reshape(-1, 4)
-
-    # Compute analytical response from magnetic prism
-    prism_1 = MagneticPrism(block1[:, 0], block1[:, 1], M1 * np.linalg.norm(b0) / mu_0)
-    prism_2 = MagneticPrism(block2[:, 0], block2[:, 1], -M1 * np.linalg.norm(b0) / mu_0)
-    prism_3 = MagneticPrism(block2[:, 0], block2[:, 1], M2 * np.linalg.norm(b0) / mu_0)
-
-    d = (
-        prism_1.magnetic_flux_density(locXyz)
-        + prism_2.magnetic_flux_density(locXyz)
-        + prism_3.magnetic_flux_density(locXyz)
-    )
-    tmi = sim.tmi_projection
-
-    np.testing.assert_allclose(data[:, 0], d[:, 0])
-    np.testing.assert_allclose(data[:, 1], d[:, 1])
-    np.testing.assert_allclose(data[:, 2], d[:, 2])
-    np.testing.assert_allclose(data[:, 3], d @ tmi)
-
-
-def test_ana_mag_amp_forward():
-    nx = 5
-    ny = 5
-
-    h0_amplitude, h0_inclination, h0_declination = (50000.0, 60.0, 250.0)
-    b0 = mag.analytics.IDTtoxyz(-h0_inclination, h0_declination, h0_amplitude)
-
-    M1 = utils.mat_utils.dip_azimuth2cartesian(45, -40) * 0.05
-    M2 = utils.mat_utils.dip_azimuth2cartesian(120, 32) * 0.1
-
-    # 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])
+        survey = mag.Survey(sources)
+        # Check if error is raised
+        msg = (
+            "Invalid mesh with 2 dimensions. "
+            "Only 3D meshes are supported when using 'choclo' as engine."
         )
-
-    block1_inds = get_block_inds(mesh.cell_centers, block1)
-    block2_inds = get_block_inds(mesh.cell_centers, block2)
-
-    model = np.zeros((mesh.n_cells, 3))
-    model[block1_inds] = M1
-    model[block2_inds] = M2
-
-    active_cells = np.any(model != 0.0, axis=1)
-    model_reduced = model[active_cells].reshape(-1, order="F")
-
-    # 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)]
-    components = ["bx", "by", "bz"]
-
-    rxLoc = mag.Point(locXyz, components=components)
-    srcField = mag.UniformBackgroundField(
-        receiver_list=[rxLoc],
-        amplitude=h0_amplitude,
-        inclination=h0_inclination,
-        declination=h0_declination,
-    )
-    survey = mag.Survey(srcField)
-
-    # Create reduced identity map for Linear Problem
-    idenMap = maps.IdentityMap(nP=int(sum(active_cells)) * 3)
-
-    sim = mag.Simulation3DIntegral(
-        mesh,
-        survey=survey,
-        chiMap=idenMap,
-        ind_active=active_cells,
-        store_sensitivities="forward_only",
-        model_type="vector",
-        is_amplitude_data=True,
-        n_processes=None,
-    )
-
-    data = sim.dpred(model_reduced)
-
-    # Compute analytical response from magnetic prism
-    prism_1 = MagneticPrism(block1[:, 0], block1[:, 1], M1 * np.linalg.norm(b0) / mu_0)
-    prism_2 = MagneticPrism(block2[:, 0], block2[:, 1], -M1 * np.linalg.norm(b0) / mu_0)
-    prism_3 = MagneticPrism(block2[:, 0], block2[:, 1], M2 * np.linalg.norm(b0) / mu_0)
-
-    d = (
-        prism_1.magnetic_flux_density(locXyz)
-        + prism_2.magnetic_flux_density(locXyz)
-        + prism_3.magnetic_flux_density(locXyz)
-    )
-    d_amp = np.linalg.norm(d, axis=1)
-
-    np.testing.assert_allclose(data, d_amp)
+        with pytest.raises(ValueError, match=msg):
+            mag.Simulation3DIntegral(mesh, survey, engine="choclo")
 
 
 def test_removed_modeltype():
@@ -513,7 +826,3 @@ def test_removed_modeltype():
     message = "modelType has been removed, please use model_type."
     with pytest.raises(NotImplementedError, match=message):
         sim.modelType
-
-
-if __name__ == "__main__":
-    unittest.main()