Repository: simpeg/simpeg
Pull request #1634 (2025-03-24 22:28:34)

Implement magnetic G as linear operator
#### Summary

Make the `G` property in the magnetic survey to be either a 2D Numpy array, or a `scipy.sparse.LinearOperator` if `store_sensitivities` is set to `"forward_only"`. The `LinearOperator` is capable of performing operations like `G @ m` and `G.T @ v` without allocating the full `G` matrix on memory or disk. This allows `Jvec` and `Jtvec` to be called and not allocate the full `G` matrix. Overwrite the `getJ` method to return either a 2D Numpy array or a `LinearOperator` formed by the product between `G` and the `chiDeriv` as a `LinearOperator` as well. Reimplement the `getJtJdiag` method: allow it to get the diagonal without building `G` if `"forward_only"`, optimize the computation of the diagonal when `G` is fully formed through `einsum`s, and fix the cache of `gtg_diagonal` by checking the hash of the `W` argument. Raise `NotImplementedError`s in some methods when `is_amplitude_data` is `True`, specially the `"forward_only"` cases where we don't want to allocate the `G` matrix. Extend docstrings of relevant methods. Add tests for all the new features.


#### 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/latest/content/getting_started/contributing/code-style.html).
* [ ] Added [tests](https://docs.simpeg.xyz/latest/content/getting_started/contributing/testing.html) to verify changes to the code.
* [ ] Added necessary documentation to any new functions/classes following the
      expect [style](https://docs.simpeg.xyz/latest/content/getting_started/contributing/documentation.html).
* [ ] 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

Part of #1618

#### 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] test when `is_amplitude_data`. Maybe we should deprecate this and allow a new `b_norm` component in the receivers, but for another PR.
- [x] override the `getJ` method provided by `LinearSimulation`. The mag simulation has a `getJ` method, although it might not work...
- [x] implement diagonal of `G.T @ G` when `forward_only`


codecov[bot]: ## [Codecov](https://app.codecov.io/gh/simpeg/simpeg/pull/1634?dropdown=coverage&src=pr&el=h1&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=simpeg) Report
Attention: Patch coverage is `60.62500%` with `63 lines` in your changes missing coverage. Please review.
> Project coverage is 86.47%. Comparing base [(`1148ef1`)](https://app.codecov.io/gh/simpeg/simpeg/commit/1148ef1e9e8bf6388ce6ff766c5c663e69e50635?dropdown=coverage&el=desc&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=simpeg) to head [(`bf6f370`)](https://app.codecov.io/gh/simpeg/simpeg/commit/bf6f370ced726591aac289a96e617ac496b5c920?dropdown=coverage&el=desc&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=simpeg).

| [Files with missing lines](https://app.codecov.io/gh/simpeg/simpeg/pull/1634?dropdown=coverage&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/1634?src=pr&el=tree&filepath=simpeg%2Fpotential_fields%2Fmagnetics%2F_numba_functions.py&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=simpeg#diff-c2ltcGVnL3BvdGVudGlhbF9maWVsZHMvbWFnbmV0aWNzL19udW1iYV9mdW5jdGlvbnMucHk=) | 6.25% | [60 Missing :warning: ](https://app.codecov.io/gh/simpeg/simpeg/pull/1634?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/1634?src=pr&el=tree&filepath=simpeg%2Fpotential_fields%2Fmagnetics%2Fsimulation.py&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=simpeg#diff-c2ltcGVnL3BvdGVudGlhbF9maWVsZHMvbWFnbmV0aWNzL3NpbXVsYXRpb24ucHk=) | 93.18% | [1 Missing and 2 partials :warning: ](https://app.codecov.io/gh/simpeg/simpeg/pull/1634?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    #1634      +/-   ##
==========================================
- Coverage   86.55%   86.47%   -0.09%     
==========================================
  Files         409      409              
  Lines       53118    53264     +146     
  Branches     5024     5037      +13     
==========================================
+ Hits        45975    46058      +83     
- Misses       5724     5785      +61     
- Partials     1419     1421       +2     
```

</details>

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

<details><summary> :rocket: New features to boost your workflow: </summary>

- :snowflake: [Test Analytics](https://docs.codecov.com/docs/test-analytics): Detect flaky tests, report on failures, and find test suite problems.
</details>

Patch:
diff --git a/pyproject.toml b/pyproject.toml
index 75dd690fb..a62d394c0 100644
--- a/pyproject.toml
+++ b/pyproject.toml
@@ -261,4 +261,5 @@ filterwarnings = [
     "ignore::simpeg.utils.solver_utils.DefaultSolverWarning",
     "error:You are running a pytest without setting a random seed.*:UserWarning",
     "error:The `index_dictionary` property has been deprecated:FutureWarning",
-]
\ No newline at end of file
+]
+xfail_strict = true
diff --git a/simpeg/potential_fields/magnetics/_numba_functions.py b/simpeg/potential_fields/magnetics/_numba_functions.py
index dfda67145..031b057a2 100644
--- a/simpeg/potential_fields/magnetics/_numba_functions.py
+++ b/simpeg/potential_fields/magnetics/_numba_functions.py
@@ -47,9 +47,13 @@ def _sensitivity_mag(
         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
+    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_cells)``
+        if ``scalar_model`` is True.
+        The array should have a shape of ``(n_receivers, 3 * n_active_cells)``
+        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.
@@ -180,135 +184,1560 @@ def _sensitivity_mag(
         # 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)
+            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_cells)``
+        if ``scalar_model`` is True.
+        The array should have a shape of ``(n_receivers, 3 * n_active_cells)``
+        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 _sensitivity_tmi_derivative(
+    receivers,
+    nodes,
+    sensitivity_matrix,
+    cell_nodes,
+    regional_field,
+    kernel_xx,
+    kernel_yy,
+    kernel_zz,
+    kernel_xy,
+    kernel_xz,
+    kernel_yz,
+    constant_factor,
+    scalar_model,
+):
+    r"""
+    Fill the sensitivity matrix for a TMI derivative.
+
+    This function should be used with a `numba.jit` decorator, for example:
+
+    .. code::
+
+        from numba import jit
+
+        jit_sens = jit(nopython=True, parallel=True)(_sensitivity_tmi_derivative)
+
+    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_cells)``
+        if ``scalar_model`` is True.
+        The array should have a shape of ``(n_receivers, 3 * n_active_cells)``
+        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).
+    kernel_xx, kernel_yy, kernel_zz, kernel_xy, kernel_xz, kernel_yz : callables
+        Kernel functions used for computing the desired TMI derivative.
+    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 kernel functions
+    ^^^^^^^^^^^^^^^^^^^^^^^^^^
+
+    To compute the :math:`\alpha` derivative of the TMI :math:`\Delta T` (with
+    :math:`\alpha \in \{x, y, z\}` we need to evaluate third order kernels
+    functions for the prism. The kernels we need to evaluate can be obtained by
+    fixing one of the subindices to the direction of the derivative
+    (:math:`\alpha`) and cycle through combinations of the other two.
+
+    For ``tmi_x`` we need to pass:
+
+    .. code::
+
+        kernel_xx=kernel_eee, kernel_yy=kernel_enn, kernel_zz=kernel_euu,
+        kernel_xy=kernel_een, kernel_xz=kernel_eeu, kernel_yz=kernel_enu
+
+    For ``tmi_y`` we need to pass:
+
+    .. code::
+
+        kernel_xx=kernel_een, kernel_yy=kernel_nnn, kernel_zz=kernel_nuu,
+        kernel_xy=kernel_enn, kernel_xz=kernel_enu, kernel_yz=kernel_nnu
+
+    For ``tmi_z`` we need to pass:
+
+    .. code::
+
+        kernel_xx=kernel_eeu, kernel_yy=kernel_nnu, kernel_zz=kernel_uuu,
+        kernel_xy=kernel_enu, kernel_xz=kernel_euu, kernel_yz=kernel_nuu
+
+
+    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 derivative (spatial)
+    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 derivative with respect
+      to the _x_ component of the effective susceptibility of each cell; then
+    * the partial derivatives of the tmi derivative with respect
+      to the _y_ component of the effective susceptibility of each cell; and then
+    * the partial derivatives of the tmi derivative with respect
+      to the _z_ component of the effective susceptibility of each 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] = kernel_xx(dx, dy, dz, distance)
+            kyy[j] = kernel_yy(dx, dy, dz, distance)
+            kzz[j] = kernel_zz(dx, dy, dz, distance)
+            kxy[j] = kernel_xy(dx, dy, dz, distance)
+            kxz[j] = kernel_xz(dx, dy, dz, distance)
+            kyz[j] = kernel_yz(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
+                )
+
+
+@jit(nopython=True, parallel=False)
+def _mag_sensitivity_t_dot_v_serial(
+    receivers,
+    nodes,
+    cell_nodes,
+    regional_field,
+    kernel_x,
+    kernel_y,
+    kernel_z,
+    constant_factor,
+    scalar_model,
+    vector,
+    result,
+):
+    r"""
+    Compute ``G.T @ v`` in serial, without building G, for a single magnetic component.
+
+    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.
+    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).
+    vector : (n_receivers) numpy.ndarray
+        Array that represents the vector used in the dot product.
+    result : (n_active_cells) or (3 * n_active_cells) numpy.ndarray
+        Running result array where the output of the dot product will be added to.
+        The array should have ``n_active_cells`` elements if ``scalar_model``
+        is True, or ``3 * n_active_cells`` otherwise.
+
+    Notes
+    -----
+    This function is meant to be run in serial. Writing to the ``result`` array
+    inside a parallel loop over the receivers generates a race condition that
+    leads to corrupted outputs.
+
+    A parallel implementation of this function is available in
+    ``_mag_sensitivity_t_dot_v_parallel``.
+
+    See also
+    --------
+    _sensitivity_mag
+        Compute the sensitivity matrix for a single magnetic component by
+        allocating it in memory.
+    """
+    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:
+                result[k] += (
+                    constant_factor
+                    * vector[i]
+                    * regional_field_amplitude
+                    * (ux * fx + uy * fy + uz * fz)
+                )
+            else:
+                result[k] += constant_factor * vector[i] * regional_field_amplitude * ux
+                result[k + n_cells] += (
+                    constant_factor * vector[i] * regional_field_amplitude * uy
+                )
+                result[k + 2 * n_cells] += (
+                    constant_factor * vector[i] * regional_field_amplitude * uz
+                )
+
+
+@jit(nopython=True, parallel=True)
+def _mag_sensitivity_t_dot_v_parallel(
+    receivers,
+    nodes,
+    cell_nodes,
+    regional_field,
+    kernel_x,
+    kernel_y,
+    kernel_z,
+    constant_factor,
+    scalar_model,
+    vector,
+    result,
+):
+    r"""
+    Compute ``G.T @ v`` in parallel, without building G, for a single magnetic component
+
+    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.
+    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).
+    vector : (n_receivers) numpy.ndarray
+        Array that represents the vector used in the dot product.
+    result : (n_active_cells) or (3 * n_active_cells) numpy.ndarray
+        Running result array where the output of the dot product will be added to.
+        The array should have ``n_active_cells`` elements if ``scalar_model``
+        is True, or ``3 * n_active_cells`` otherwise.
+
+    Notes
+    -----
+    This function is meant to be run in parallel.
+    This implementation instructs each thread to allocate their own array for
+    the current row of the sensitivity matrix. After computing the elements of
+    that row, it gets added to the running ``result`` array through a reduction
+    operation handled by Numba.
+
+    A serialized implementation of this function is available in
+    ``_mag_sensitivity_t_dot_v_serial``.
+
+    See also
+    --------
+    _sensitivity_mag
+        Compute the sensitivity matrix for a single magnetic component by
+        allocating it in memory.
+    """
+    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
+    result_size = result.size
+    # 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 array for the current row of the sensitivity matrix
+        local_row = np.empty(result_size)
+        # 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:
+                local_row[k] = (
+                    constant_factor
+                    * vector[i]
+                    * regional_field_amplitude
+                    * (ux * fx + uy * fy + uz * fz)
+                )
+            else:
+                local_row[k] = (
+                    constant_factor * vector[i] * regional_field_amplitude * ux
+                )
+                local_row[k + n_cells] = (
+                    constant_factor * vector[i] * regional_field_amplitude * uy
+                )
+                local_row[k + 2 * n_cells] = (
+                    constant_factor * vector[i] * regional_field_amplitude * uz
+                )
+        # Apply reduction operation to add the values of the row to the running
+        # result. Avoid slicing the `result` array when updating it to avoid
+        # racing conditions, just add the `local_row` to the `results`
+        # variable.
+        result += local_row
+
+
+@jit(nopython=True, parallel=False)
+def _tmi_sensitivity_t_dot_v_serial(
+    receivers,
+    nodes,
+    cell_nodes,
+    regional_field,
+    constant_factor,
+    scalar_model,
+    vector,
+    result,
+):
+    r"""
+    Compute ``G.T @ v`` in serial, without building G, for 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.
+    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).
+    vector : (n_receivers) numpy.ndarray
+        Array that represents the vector used in the dot product.
+    result : (n_active_cells) or (3 * n_active_cells) numpy.ndarray
+        Running result array where the output of the dot product will be added to.
+        The array should have ``n_active_cells`` elements if ``scalar_model``
+        is True, or ``3 * n_active_cells`` otherwise.
+
+    Notes
+    -----
+    This function is meant to be run in serial. Writing to the ``result`` array
+    inside a parallel loop over the receivers generates a race condition that
+    leads to corrupted outputs.
+
+    A parallel implementation of this function is available in
+    ``_tmi_sensitivity_t_dot_v_parallel``.
+
+    See also
+    --------
+    _sensitivity_tmi
+        Compute the sensitivity matrix for TMI by allocating it in memory.
+    """
+    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:
+                result[k] += (
+                    constant_factor
+                    * vector[i]
+                    * regional_field_amplitude
+                    * (bx * fx + by * fy + bz * fz)
+                )
+            else:
+                result[k] += constant_factor * vector[i] * regional_field_amplitude * bx
+                result[k + n_cells] += (
+                    constant_factor * vector[i] * regional_field_amplitude * by
+                )
+                result[k + 2 * n_cells] += (
+                    constant_factor * vector[i] * regional_field_amplitude * bz
+                )
+
+
+@jit(nopython=True, parallel=True)
+def _tmi_sensitivity_t_dot_v_parallel(
+    receivers,
+    nodes,
+    cell_nodes,
+    regional_field,
+    constant_factor,
+    scalar_model,
+    vector,
+    result,
+):
+    r"""
+    Compute ``G.T @ v`` in parallel, without building G, for 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.
+    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).
+    vector : (n_receivers) numpy.ndarray
+        Array that represents the vector used in the dot product.
+    result : (n_active_cells) or (3 * n_active_cells) numpy.ndarray
+        Running result array where the output of the dot product will be added to.
+        The array should have ``n_active_cells`` elements if ``scalar_model``
+        is True, or ``3 * n_active_cells`` otherwise.
+
+    Notes
+    -----
+    This function is meant to be run in parallel.
+    This implementation instructs each thread to allocate their own array for
+    the current row of the sensitivity matrix. After computing the elements of
+    that row, it gets added to the running ``result`` array through a reduction
+    operation handled by Numba.
+
+    A serialized implementation of this function is available in
+    ``_tmi_sensitivity_t_dot_v_serial``.
+
+    See also
+    --------
+    _sensitivity_tmi
+        Compute the sensitivity matrix for TMI by allocating it in memory.
+    """
+    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
+    result_size = result.size
+    # 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 array for the current row of the sensitivity matrix
+        local_row = np.empty(result_size)
+        # 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:
+                local_row[k] = (
+                    constant_factor
+                    * vector[i]
+                    * regional_field_amplitude
+                    * (bx * fx + by * fy + bz * fz)
+                )
+            else:
+                local_row[k] = (
+                    constant_factor * vector[i] * regional_field_amplitude * bx
+                )
+                local_row[k + n_cells] = (
+                    constant_factor * vector[i] * regional_field_amplitude * by
+                )
+                local_row[k + 2 * n_cells] = (
+                    constant_factor * vector[i] * regional_field_amplitude * bz
+                )
+        # Apply reduction operation to add the values of the row to the running
+        # result. Avoid slicing the `result` array when updating it to avoid
+        # racing conditions, just add the `local_row` to the `results`
+        # variable.
+        result += local_row
+
+
+@jit(nopython=True, parallel=False)
+def _tmi_derivative_sensitivity_t_dot_v_serial(
+    receivers,
+    nodes,
+    cell_nodes,
+    regional_field,
+    kernel_xx,
+    kernel_yy,
+    kernel_zz,
+    kernel_xy,
+    kernel_xz,
+    kernel_yz,
+    constant_factor,
+    scalar_model,
+    vector,
+    result,
+):
+    r"""
+    Compute ``G.T @ v`` in serial, without building G, for a spatial TMI derivative.
+
+    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.
+    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_xx, kernel_yy, kernel_zz, kernel_xy, kernel_xz, kernel_yz : callables
+        Kernel functions used for computing the desired TMI derivative.
+    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).
+    vector : (n_receivers) numpy.ndarray
+        Array that represents the vector used in the dot product.
+    result : (n_active_cells) or (3 * n_active_cells) numpy.ndarray
+        Running result array where the output of the dot product will be added to.
+        The array should have ``n_active_cells`` elements if ``scalar_model``
+        is True, or ``3 * n_active_cells`` otherwise.
+
+    Notes
+    -----
+    This function is meant to be run in serial. Writing to the ``result`` array
+    inside a parallel loop over the receivers generates a race condition that
+    leads to corrupted outputs.
+
+    A parallel implementation of this function is available in
+    ``_tmi_derivative_sensitivity_t_dot_v_parallel``.
+
+    See also
+    --------
+    _sensitivity_tmi_derivative
+        Compute the sensitivity matrix for a TMI derivative by allocating it in memory.
+    """
+    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] = kernel_xx(dx, dy, dz, distance)
+            kyy[j] = kernel_yy(dx, dy, dz, distance)
+            kzz[j] = kernel_zz(dx, dy, dz, distance)
+            kxy[j] = kernel_xy(dx, dy, dz, distance)
+            kxz[j] = kernel_xz(dx, dy, dz, distance)
+            kyz[j] = kernel_yz(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:
+                result[k] += (
+                    constant_factor
+                    * vector[i]
+                    * regional_field_amplitude
+                    * (bx * fx + by * fy + bz * fz)
+                )
+            else:
+                result[k] += constant_factor * vector[i] * regional_field_amplitude * bx
+                result[k + n_cells] += (
+                    constant_factor * vector[i] * regional_field_amplitude * by
+                )
+                result[k + 2 * n_cells] += (
+                    constant_factor * vector[i] * regional_field_amplitude * bz
+                )
+
+
+@jit(nopython=True, parallel=True)
+def _tmi_derivative_sensitivity_t_dot_v_parallel(
+    receivers,
+    nodes,
+    cell_nodes,
+    regional_field,
+    kernel_xx,
+    kernel_yy,
+    kernel_zz,
+    kernel_xy,
+    kernel_xz,
+    kernel_yz,
+    constant_factor,
+    scalar_model,
+    vector,
+    result,
+):
+    r"""
+    Compute ``G.T @ v`` in parallel, without building G, for a spatial TMI derivative.
+
+    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.
+    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_xx, kernel_yy, kernel_zz, kernel_xy, kernel_xz, kernel_yz : callables
+        Kernel functions used for computing the desired TMI derivative.
+    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).
+    vector : (n_receivers) numpy.ndarray
+        Array that represents the vector used in the dot product.
+    result : (n_active_cells) or (3 * n_active_cells) numpy.ndarray
+        Running result array where the output of the dot product will be added to.
+        The array should have ``n_active_cells`` elements if ``scalar_model``
+        is True, or ``3 * n_active_cells`` otherwise.
+
+    Notes
+    -----
+    This function is meant to be run in parallel.
+    This implementation instructs each thread to allocate their own array for
+    the current row of the sensitivity matrix. After computing the elements of
+    that row, it gets added to the running ``result`` array through a reduction
+    operation handled by Numba.
+
+    A serialized implementation of this function is available in
+    ``_tmi_derivative_sensitivity_t_dot_v_serial``.
+
+    See also
+    --------
+    _sensitivity_tmi_derivative
+        Compute the sensitivity matrix for a TMI derivative by allocating it in memory.
+    """
+    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
+    result_size = result.size
+    # 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 array for the current row of the sensitivity matrix
+        local_row = np.empty(result_size)
+        # 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] = kernel_xx(dx, dy, dz, distance)
+            kyy[j] = kernel_yy(dx, dy, dz, distance)
+            kzz[j] = kernel_zz(dx, dy, dz, distance)
+            kxy[j] = kernel_xy(dx, dy, dz, distance)
+            kxz[j] = kernel_xz(dx, dy, dz, distance)
+            kyz[j] = kernel_yz(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:
+                local_row[k] = (
+                    constant_factor
+                    * vector[i]
+                    * regional_field_amplitude
+                    * (bx * fx + by * fy + bz * fz)
+                )
+            else:
+                local_row[k] = (
+                    constant_factor * vector[i] * regional_field_amplitude * bx
+                )
+                local_row[k + n_cells] = (
+                    constant_factor * vector[i] * regional_field_amplitude * by
+                )
+                local_row[k + 2 * n_cells] = (
+                    constant_factor * vector[i] * regional_field_amplitude * bz
+                )
+        # Apply reduction operation to add the values of the row to the running
+        # result. Avoid slicing the `result` array when updating it to avoid
+        # racing conditions, just add the `local_row` to the `results`
+        # variable.
+        result += local_row
+
+
+@jit(nopython=True, parallel=False)
+def _diagonal_G_T_dot_G_mag_serial(
+    receivers,
+    nodes,
+    cell_nodes,
+    regional_field,
+    kernel_x,
+    kernel_y,
+    kernel_z,
+    constant_factor,
+    scalar_model,
+    weights,
+    diagonal,
+):
+    """
+    Diagonal of ``G.T @ W.T @ W @ G`` for single magnetic component, in serial.
+
+    This function doesn't store the full ``G`` matrix in memory.
+
+    Parameters
+    ----------
+    receivers : (n_receivers, 3) numpy.ndarray
+        Array with the locations of the receivers
+    nodes : (n_active_nodes, 3) numpy.ndarray
+        Array with the location of the mesh nodes.
+    cell_nodes : (n_active_cells, 8) numpy.ndarray
+        Array of integers, where each row contains the indices of the nodes for
+        each active cell in the mesh.
+    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).
+    weights : (n_receivers,) numpy.ndarray
+        Array with data weights. It should be the diagonal of the ``W`` matrix,
+        squared.
+    diagonal : (n_active_cells) or (3 * n_active_cells) numpy.ndarray
+        Array where the diagonal of ``G.T @ G`` will be added to.
+
+    Notes
+    -----
+    This function is meant to be run in serial. Use the
+    ``_diagonal_G_T_dot_G_mag_parallel`` one for parallelized computations.
+    """
+    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:
+                g_element = (
+                    constant_factor
+                    * regional_field_amplitude
+                    * (ux * fx + uy * fy + uz * fz)
+                )
+                diagonal[k] += weights[i] * g_element**2
+            else:
+                const = constant_factor * regional_field_amplitude
+                diagonal[k] += weights[i] * (const * ux) ** 2
+                diagonal[k + n_cells] += weights[i] * (const * uy) ** 2
+                diagonal[k + 2 * n_cells] += weights[i] * (const * uz) ** 2
+
+
+@jit(nopython=True, parallel=True)
+def _diagonal_G_T_dot_G_mag_parallel(
+    receivers,
+    nodes,
+    cell_nodes,
+    regional_field,
+    kernel_x,
+    kernel_y,
+    kernel_z,
+    constant_factor,
+    scalar_model,
+    weights,
+    diagonal,
+):
+    """
+    Diagonal of ``G.T @ W.T @ W @ G`` for single magnetic component, in parallel.
+
+    This function doesn't store the full ``G`` matrix in memory.
+
+    Parameters
+    ----------
+    receivers : (n_receivers, 3) numpy.ndarray
+        Array with the locations of the receivers
+    nodes : (n_active_nodes, 3) numpy.ndarray
+        Array with the location of the mesh nodes.
+    cell_nodes : (n_active_cells, 8) numpy.ndarray
+        Array of integers, where each row contains the indices of the nodes for
+        each active cell in the mesh.
+    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).
+    weights : (n_receivers,) numpy.ndarray
+        Array with data weights. It should be the diagonal of the ``W`` matrix,
+        squared.
+    diagonal : (n_active_cells) or (3 * n_active_cells) numpy.ndarray
+        Array where the diagonal of ``G.T @ G`` will be added to.
+
+    Notes
+    -----
+    This function is meant to be run in parallel. Use the
+    ``_diagonal_G_T_dot_G_mag_serial`` one for serialized computations.
+    """
+    n_receivers = receivers.shape[0]
+    n_nodes = nodes.shape[0]
+    n_cells = cell_nodes.shape[0]
+    diagonal_size = diagonal.size
+    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 array for the diagonal elements for the current receiver.
+        local_diagonal = np.empty(diagonal_size)
+        # 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:
+                g_element = (
+                    constant_factor
+                    * regional_field_amplitude
+                    * (ux * fx + uy * fy + uz * fz)
+                )
+                local_diagonal[k] = weights[i] * g_element**2
+            else:
+                const = constant_factor * regional_field_amplitude
+                local_diagonal[k] = weights[i] * (const * ux) ** 2
+                local_diagonal[k + n_cells] = weights[i] * (const * uy) ** 2
+                local_diagonal[k + 2 * n_cells] = weights[i] * (const * uz) ** 2
+        # Add the result to the diagonal.
+        # Apply reduction operation to add the values of the local diagonal to
+        # the running diagonal array. Avoid slicing the `diagonal` array when
+        # updating it to avoid racing conditions, just add the `local_diagonal`
+        # to the `diagonal` variable.
+        diagonal += local_diagonal
+
+
+@jit(nopython=True, parallel=False)
+def _diagonal_G_T_dot_G_tmi_serial(
+    receivers,
+    nodes,
+    cell_nodes,
+    regional_field,
+    constant_factor,
+    scalar_model,
+    weights,
+    diagonal,
+):
+    """
+    Diagonal of ``G.T @ W.T @ W @ G`` for TMI, in serial.
+
+    This function doesn't store the full ``G`` matrix in memory.
+
+    Parameters
+    ----------
+    receivers : (n_receivers, 3) numpy.ndarray
+        Array with the locations of the receivers
+    nodes : (n_active_nodes, 3) numpy.ndarray
+        Array with the location of the mesh nodes.
+    cell_nodes : (n_active_cells, 8) numpy.ndarray
+        Array of integers, where each row contains the indices of the nodes for
+        each active cell in the mesh.
+    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 built to work with scalar models
+        (susceptibilities).
+        If False, the sensitivity matrix is built to work with vector models
+        (effective susceptibilities).
+    weights : (n_receivers,) numpy.ndarray
+        Array with data weights. It should be the diagonal of the ``W`` matrix,
+        squared.
+    diagonal : (n_active_cells) or (3 * n_active_cells) numpy.ndarray
+        Array where the diagonal of ``G.T @ G`` will be added to.
+
+    Notes
+    -----
+    This function is meant to be run in serial. Use the
+    ``_diagonal_G_T_dot_G_tmi_parallel`` one for parallelized computations.
+    """
+    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:
+                g_element = (
+                    constant_factor
+                    * regional_field_amplitude
+                    * (bx * fx + by * fy + bz * fz)
+                )
+                diagonal[k] += weights[i] * g_element**2
+            else:
+                const = constant_factor * regional_field_amplitude
+                diagonal[k] += weights[i] * (const * bx) ** 2
+                diagonal[k + n_cells] += weights[i] * (const * by) ** 2
+                diagonal[k + 2 * n_cells] += weights[i] * (const * bz) ** 2
+
+
+@jit(nopython=True, parallel=True)
+def _diagonal_G_T_dot_G_tmi_parallel(
+    receivers,
+    nodes,
+    cell_nodes,
+    regional_field,
+    constant_factor,
+    scalar_model,
+    weights,
+    diagonal,
+):
+    """
+    Diagonal of ``G.T @ W.T @ W @ G`` for TMI, in parallel.
+
+    This function doesn't store the full ``G`` matrix in memory.
+
+    Parameters
+    ----------
+    receivers : (n_receivers, 3) numpy.ndarray
+        Array with the locations of the receivers
+    nodes : (n_active_nodes, 3) numpy.ndarray
+        Array with the location of the mesh nodes.
+    cell_nodes : (n_active_cells, 8) numpy.ndarray
+        Array of integers, where each row contains the indices of the nodes for
+        each active cell in the mesh.
+    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 built to work with scalar models
+        (susceptibilities).
+        If False, the sensitivity matrix is built to work with vector models
+        (effective susceptibilities).
+    weights : (n_receivers,) numpy.ndarray
+        Array with data weights. It should be the diagonal of the ``W`` matrix,
+        squared.
+    diagonal : (n_active_cells) or (3 * n_active_cells) numpy.ndarray
+        Array where the diagonal of ``G.T @ G`` will be added to.
+
+    Notes
+    -----
+    This function is meant to be run in parallel. Use the
+    ``_diagonal_G_T_dot_G_tmi_serial`` one for serialized computations.
+    """
+    n_receivers = receivers.shape[0]
+    n_nodes = nodes.shape[0]
+    n_cells = cell_nodes.shape[0]
+    diagonal_size = diagonal.size
+    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 array for the diagonal elements for the current receiver.
+        local_diagonal = np.empty(diagonal_size)
+        # 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:
-                sensitivity_matrix[i, k] = (
+                g_element = (
                     constant_factor
                     * regional_field_amplitude
-                    * (ux * fx + uy * fy + uz * fz)
+                    * (bx * fx + by * fy + bz * fz)
                 )
+                local_diagonal[k] = weights[i] * g_element**2
             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
-                )
+                const = constant_factor * regional_field_amplitude
+                local_diagonal[k] = weights[i] * (const * bx) ** 2
+                local_diagonal[k + n_cells] = weights[i] * (const * by) ** 2
+                local_diagonal[k + 2 * n_cells] = weights[i] * (const * bz) ** 2
 
+        # Add the result to the diagonal.
+        # Apply reduction operation to add the values of the local diagonal to
+        # the running diagonal array. Avoid slicing the `diagonal` array when
+        # updating it to avoid racing conditions, just add the `local_diagonal`
+        # to the `diagonal` variable.
+        diagonal += local_diagonal
 
-def _sensitivity_tmi(
+
+@jit(nopython=True, parallel=False)
+def _diagonal_G_T_dot_G_tmi_deriv_serial(
     receivers,
     nodes,
-    sensitivity_matrix,
     cell_nodes,
     regional_field,
+    kernel_xx,
+    kernel_yy,
+    kernel_zz,
+    kernel_xy,
+    kernel_xz,
+    kernel_yz,
     constant_factor,
     scalar_model,
+    weights,
+    diagonal,
 ):
-    r"""
-    Fill the sensitivity matrix for TMI
-
-    This function should be used with a `numba.jit` decorator, for example:
-
-    .. code::
-
-        from numba import jit
+    """
+    Diagonal of ``G.T @ W.T @ W @ G`` for TMI derivatives, in serial.
 
-        jit_sensitivity_tmi = jit(nopython=True, parallel=True)(_sensitivity_tmi)
+    This function doesn't store the full ``G`` matrix in memory.
 
     Parameters
     ----------
-    receivers : (n_receivers, 3) array
+    receivers : (n_receivers, 3) numpy.ndarray
         Array with the locations of the receivers
-    nodes : (n_active_nodes, 3) array
+    nodes : (n_active_nodes, 3) numpy.ndarray
         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
+    cell_nodes : (n_active_cells, 8) numpy.ndarray
         Array of integers, where each row contains the indices of the nodes for
         each active cell in the mesh.
     regional_field : (3,) array
         Array containing the x, y and z components of the regional magnetic
         field (uniform background field).
+    kernel_xx, kernel_yy, kernel_zz, kernel_xy, kernel_xz, kernel_yz : callables
+        Kernel functions used for computing the desired TMI derivative.
     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
+        If True, the sensitivity matrix is built to work with scalar models
         (susceptibilities).
-        If False, the sensitivity matrix is build to work with vector models
+        If False, the sensitivity matrix is built to work with vector models
         (effective susceptibilities).
+    weights : (n_receivers,) numpy.ndarray
+        Array with data weights. It should be the diagonal of the ``W`` matrix,
+        squared.
+    diagonal : (n_active_cells) or (3 * n_active_cells) numpy.ndarray
+        Array where the diagonal of ``G.T @ G`` will be added to.
 
     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.
+    This function is meant to be run in serial. Use the
+    ``_diagonal_G_T_dot_G_tmi_deriv_parallel`` one for parallelized computations.
     """
     n_receivers = receivers.shape[0]
     n_nodes = nodes.shape[0]
@@ -330,12 +1759,12 @@ def _sensitivity_tmi(
             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)
+            kxx[j] = kernel_xx(dx, dy, dz, distance)
+            kyy[j] = kernel_yy(dx, dy, dz, distance)
+            kzz[j] = kernel_zz(dx, dy, dz, distance)
+            kxy[j] = kernel_xy(dx, dy, dz, distance)
+            kxz[j] = kernel_xz(dx, dy, dz, distance)
+            kyz[j] = kernel_yz(dx, dy, dz, distance)
         # Compute sensitivity matrix elements from the kernel values
         for k in range(n_cells):
             nodes_indices = cell_nodes[k, :]
@@ -348,30 +1777,25 @@ def _sensitivity_tmi(
             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] = (
+                g_element = (
                     constant_factor
                     * regional_field_amplitude
                     * (bx * fx + by * fy + bz * fz)
                 )
+                diagonal[k] += weights[i] * g_element**2
             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
-                )
+                const = constant_factor * regional_field_amplitude
+                diagonal[k] += weights[i] * (const * bx) ** 2
+                diagonal[k + n_cells] += weights[i] * (const * by) ** 2
+                diagonal[k + 2 * n_cells] += weights[i] * (const * bz) ** 2
 
 
-def _sensitivity_tmi_derivative(
+@jit(nopython=True, parallel=True)
+def _diagonal_G_T_dot_G_tmi_deriv_parallel(
     receivers,
     nodes,
-    sensitivity_matrix,
     cell_nodes,
     regional_field,
     kernel_xx,
@@ -382,32 +1806,21 @@ def _sensitivity_tmi_derivative(
     kernel_yz,
     constant_factor,
     scalar_model,
+    weights,
+    diagonal,
 ):
-    r"""
-    Fill the sensitivity matrix for a TMI derivative.
-
-    This function should be used with a `numba.jit` decorator, for example:
-
-    .. code::
-
-        from numba import jit
+    """
+    Diagonal of ``G.T @ W.T @ W @ G`` for TMI derivatives, in parallel.
 
-        jit_sens = jit(nopython=True, parallel=True)(_sensitivity_tmi_derivative)
+    This function doesn't store the full ``G`` matrix in memory.
 
     Parameters
     ----------
-    receivers : (n_receivers, 3) array
+    receivers : (n_receivers, 3) numpy.ndarray
         Array with the locations of the receivers
-    nodes : (n_active_nodes, 3) array
+    nodes : (n_active_nodes, 3) numpy.ndarray
         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
+    cell_nodes : (n_active_cells, 8) numpy.ndarray
         Array of integers, where each row contains the indices of the nodes for
         each active cell in the mesh.
     regional_field : (3,) array
@@ -419,85 +1832,25 @@ def _sensitivity_tmi_derivative(
         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
+        If True, the sensitivity matrix is built to work with scalar models
         (susceptibilities).
-        If False, the sensitivity matrix is build to work with vector models
+        If False, the sensitivity matrix is built to work with vector models
         (effective susceptibilities).
+    weights : (n_receivers,) numpy.ndarray
+        Array with data weights. It should be the diagonal of the ``W`` matrix,
+        squared.
+    diagonal : (n_active_cells) or (3 * n_active_cells) numpy.ndarray
+        Array where the diagonal of ``G.T @ G`` will be added to.
 
     Notes
     -----
-
-    About the kernel functions
-    ^^^^^^^^^^^^^^^^^^^^^^^^^^
-
-    To compute the :math:`\alpha` derivative of the TMI :math:`\Delta T` (with
-    :math:`\alpha \in \{x, y, z\}` we need to evaluate third order kernels
-    functions for the prism. The kernels we need to evaluate can be obtained by
-    fixing one of the subindices to the direction of the derivative
-    (:math:`\alpha`) and cycle through combinations of the other two.
-
-    For ``tmi_x`` we need to pass:
-
-    .. code::
-
-        kernel_xx=kernel_eee, kernel_yy=kernel_enn, kernel_zz=kernel_euu,
-        kernel_xy=kernel_een, kernel_xz=kernel_eeu, kernel_yz=kernel_enu
-
-    For ``tmi_y`` we need to pass:
-
-    .. code::
-
-        kernel_xx=kernel_een, kernel_yy=kernel_nnn, kernel_zz=kernel_nuu,
-        kernel_xy=kernel_enn, kernel_xz=kernel_enu, kernel_yz=kernel_nnu
-
-    For ``tmi_z`` we need to pass:
-
-    .. code::
-
-        kernel_xx=kernel_eeu, kernel_yy=kernel_nnu, kernel_zz=kernel_uuu,
-        kernel_xy=kernel_enu, kernel_xz=kernel_euu, kernel_yz=kernel_nuu
-
-
-    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 derivative (spatial)
-    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 derivative with respect
-      to the _x_ component of the effective susceptibility of each cell; then
-    * the partial derivatives of the tmi derivative with respect
-      to the _y_ component of the effective susceptibility of each cell; and then
-    * the partial derivatives of the tmi derivative with respect
-      to the _z_ component of the effective susceptibility of each cell.
+    This function is meant to be run in parallel. Use the
+    ``_diagonal_G_T_dot_G_tmi_serial`` one for serialized computations.
     """
     n_receivers = receivers.shape[0]
     n_nodes = nodes.shape[0]
     n_cells = cell_nodes.shape[0]
+    diagonal_size = diagonal.size
     fx, fy, fz = regional_field
     regional_field_amplitude = np.sqrt(fx**2 + fy**2 + fz**2)
     fx /= regional_field_amplitude
@@ -508,6 +1861,8 @@ def _sensitivity_tmi_derivative(
         # 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 array for the diagonal elements for the current receiver.
+        local_diagonal = np.empty(diagonal_size)
         # 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):
@@ -533,24 +1888,26 @@ def _sensitivity_tmi_derivative(
             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] = (
+                g_element = (
                     constant_factor
                     * regional_field_amplitude
                     * (bx * fx + by * fy + bz * fz)
                 )
+                local_diagonal[k] = weights[i] * g_element**2
             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
-                )
+                const = constant_factor * regional_field_amplitude
+                local_diagonal[k] = weights[i] * (const * bx) ** 2
+                local_diagonal[k + n_cells] = weights[i] * (const * by) ** 2
+                local_diagonal[k + 2 * n_cells] = weights[i] * (const * bz) ** 2
+
+        # Add the result to the diagonal.
+        # Apply reduction operation to add the values of the local diagonal to
+        # the running diagonal array. Avoid slicing the `diagonal` array when
+        # updating it to avoid racing conditions, just add the `local_diagonal`
+        # to the `diagonal` variable.
+        diagonal += local_diagonal
 
 
 def _forward_mag(
@@ -1463,9 +2820,13 @@ def _sensitivity_mag_2d_mesh(
         Array with the top boundaries of each active cell in the 2D mesh.
     bottom : (n_active_cells) np.ndarray
         Array with the bottom boundaries of each active cell in the 2D mesh.
-    sensitivity_matrix : (n_receivers, n_active_nodes) array
+    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_cells)``
+        if ``scalar_model`` is True.
+        The array should have a shape of ``(n_receivers, 3 * n_active_cells)``
+        if ``scalar_model`` is False.
     regional_field : (3,) array
         Array containing the x, y and z components of the regional magnetic
         field (uniform background field).
@@ -1651,9 +3012,9 @@ def _sensitivity_tmi_2d_mesh(
     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)``
+        The array should have a shape of ``(n_receivers, n_active_cells)``
         if ``scalar_model`` is True.
-        The array should have a shape of ``(n_receivers, 3 * n_active_nodes)``
+        The array should have a shape of ``(n_receivers, 3 * n_active_cells)``
         if ``scalar_model`` is False.
     regional_field : (3,) array
         Array containing the x, y and z components of the regional magnetic
@@ -1833,9 +3194,9 @@ def _sensitivity_tmi_derivative_2d_mesh(
     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)``
+        The array should have a shape of ``(n_receivers, n_active_cells)``
         if ``scalar_model`` is True.
-        The array should have a shape of ``(n_receivers, 3 * n_active_nodes)``
+        The array should have a shape of ``(n_receivers, 3 * n_active_cells)``
         if ``scalar_model`` is False.
     regional_field : (3,) array
         Array containing the x, y and z components of the regional magnetic
diff --git a/simpeg/potential_fields/magnetics/simulation.py b/simpeg/potential_fields/magnetics/simulation.py
index c103fc9cf..2f3314dc3 100644
--- a/simpeg/potential_fields/magnetics/simulation.py
+++ b/simpeg/potential_fields/magnetics/simulation.py
@@ -1,5 +1,7 @@
+import hashlib
 import warnings
 import numpy as np
+from numpy.typing import NDArray
 import scipy.sparse as sp
 from geoana.kernels import (
     prism_fxxy,
@@ -11,6 +13,7 @@ from geoana.kernels import (
     prism_fzzz,
 )
 from scipy.constants import mu_0
+from scipy.sparse.linalg import LinearOperator, aslinearoperator
 
 from simpeg import props, utils
 from simpeg.utils import mat_utils, mkvc, sdiag
@@ -48,6 +51,18 @@ from ._numba_functions import (
     _sensitivity_tmi_derivative_serial,
     _sensitivity_tmi_derivative_2d_mesh_serial,
     _sensitivity_tmi_derivative_2d_mesh_parallel,
+    _mag_sensitivity_t_dot_v_serial,
+    _mag_sensitivity_t_dot_v_parallel,
+    _tmi_sensitivity_t_dot_v_serial,
+    _tmi_sensitivity_t_dot_v_parallel,
+    _tmi_derivative_sensitivity_t_dot_v_serial,
+    _tmi_derivative_sensitivity_t_dot_v_parallel,
+    _diagonal_G_T_dot_G_mag_serial,
+    _diagonal_G_T_dot_G_mag_parallel,
+    _diagonal_G_T_dot_G_tmi_serial,
+    _diagonal_G_T_dot_G_tmi_parallel,
+    _diagonal_G_T_dot_G_tmi_deriv_serial,
+    _diagonal_G_T_dot_G_tmi_deriv_parallel,
 )
 
 if choclo is not None:
@@ -205,9 +220,7 @@ class Simulation3DIntegral(BasePFSimulation):
         self.chi = chi
         self.chiMap = chiMap
 
-        self._G = None
         self._M = None
-        self._gtg_diagonal = None
         self.is_amplitude_data = is_amplitude_data
         self.modelMap = self.chiMap
 
@@ -229,6 +242,16 @@ class Simulation3DIntegral(BasePFSimulation):
                 self._forward_mag = _forward_mag_parallel
                 self._forward_tmi_derivative = _forward_tmi_derivative_parallel
                 self._sensitivity_tmi_derivative = _sensitivity_tmi_derivative_parallel
+                self._mag_sensitivity_t_dot_v = _mag_sensitivity_t_dot_v_parallel
+                self._tmi_sensitivity_t_dot_v = _tmi_sensitivity_t_dot_v_parallel
+                self._tmi_derivative_sensitivity_t_dot_v = (
+                    _tmi_derivative_sensitivity_t_dot_v_parallel
+                )
+                self._diagonal_G_T_dot_G_mag = _diagonal_G_T_dot_G_mag_parallel
+                self._diagonal_G_T_dot_G_tmi = _diagonal_G_T_dot_G_tmi_parallel
+                self._diagonal_G_T_dot_G_tmi_deriv = (
+                    _diagonal_G_T_dot_G_tmi_deriv_parallel
+                )
             else:
                 self._sensitivity_tmi = _sensitivity_tmi_serial
                 self._sensitivity_mag = _sensitivity_mag_serial
@@ -236,6 +259,16 @@ class Simulation3DIntegral(BasePFSimulation):
                 self._forward_mag = _forward_mag_serial
                 self._forward_tmi_derivative = _forward_tmi_derivative_serial
                 self._sensitivity_tmi_derivative = _sensitivity_tmi_derivative_serial
+                self._mag_sensitivity_t_dot_v = _mag_sensitivity_t_dot_v_serial
+                self._tmi_sensitivity_t_dot_v = _tmi_sensitivity_t_dot_v_serial
+                self._tmi_derivative_sensitivity_t_dot_v = (
+                    _tmi_derivative_sensitivity_t_dot_v_serial
+                )
+                self._diagonal_G_T_dot_G_mag = _diagonal_G_T_dot_G_mag_serial
+                self._diagonal_G_T_dot_G_tmi = _diagonal_G_T_dot_G_tmi_serial
+                self._diagonal_G_T_dot_G_tmi_deriv = (
+                    _diagonal_G_T_dot_G_tmi_deriv_serial
+                )
 
     @property
     def model_type(self):
@@ -306,13 +339,24 @@ class Simulation3DIntegral(BasePFSimulation):
         return fields
 
     @property
-    def G(self):
-        if getattr(self, "_G", None) is None:
-            if self.engine == "choclo":
-                self._G = self._sensitivity_matrix()
-            else:
-                self._G = self.linear_operator()
-
+    def G(self) -> NDArray | np.memmap | LinearOperator:
+        if not hasattr(self, "_G"):
+            match self.engine, self.store_sensitivities:
+                case ("choclo", "forward_only"):
+                    self._G = self._sensitivity_matrix_as_operator()
+                case ("choclo", _):
+                    self._G = self._sensitivity_matrix()
+                case ("geoana", "forward_only"):
+                    msg = (
+                        "Accessing matrix G with "
+                        'store_sensitivities="forward_only" and engine="geoana" '
+                        "hasn't been implemented yet."
+                        'Choose store_sensitivities="ram" or "disk", '
+                        'or another engine, like "choclo".'
+                    )
+                    raise NotImplementedError(msg)
+                case ("geoana", _):
+                    self._G = self.linear_operator()
         return self._G
 
     modelType = deprecate_property(
@@ -338,42 +382,195 @@ class Simulation3DIntegral(BasePFSimulation):
 
         return self._tmi_projection
 
-    def getJtJdiag(self, m, W=None, f=None):
+    def getJ(self, m, f=None) -> NDArray[np.float64 | np.float32] | LinearOperator:
+        r"""
+        Sensitivity matrix :math:`\mathbf{J}`.
+
+        Parameters
+        ----------
+        m : (n_param,) numpy.ndarray
+            The model parameters.
+        f : Ignored
+            Not used, present here for API consistency by convention.
+
+        Returns
+        -------
+        (nD, n_params) np.ndarray or scipy.sparse.linalg.LinearOperator.
+            Array or :class:`~scipy.sparse.linalg.LinearOperator` for the
+            :math:`\mathbf{J}` matrix.
+            A :class:`~scipy.sparse.linalg.LinearOperator` will be returned if
+            ``store_sensitivities`` is ``"forward_only"``, otherwise a dense
+            array will be returned.
+
+        Notes
+        -----
+        If ``store_sensitivities`` is ``"ram"`` or ``"disk"``, a dense array
+        for the ``J`` matrix is returned.
+        A :class:`~scipy.sparse.linalg.LinearOperator` is returned if
+        ``store_sensitivities`` is ``"forward_only"``. This object can perform
+        operations like ``J @ m`` or ``J.T @ v`` without allocating the full
+        ``J`` matrix in memory.
         """
-        Return the diagonal of JtJ
+        if self.is_amplitude_data:
+            msg = (
+                "The `getJ` method is not yet implemented to work with "
+                "`is_amplitude_data`."
+            )
+            raise NotImplementedError(msg)
+
+        # Need to assign the model, so the chiDeriv can be computed (if the
+        # model is None, the chiDeriv is going to be Zero).
+        self.model = m
+        chiDeriv = (
+            self.chiDeriv
+            if not isinstance(self.G, LinearOperator)
+            else aslinearoperator(self.chiDeriv)
+        )
+        return self.G @ chiDeriv
+
+    def getJtJdiag(self, m, W=None, f=None):
+        r"""
+        Compute diagonal of :math:`\mathbf{J}^T \mathbf{J}``.
+
+        Parameters
+        ----------
+        m : (n_param,) numpy.ndarray
+            The model parameters.
+        W : (nD, nD) np.ndarray or scipy.sparse.sparray, optional
+            Diagonal matrix with the square root of the weights. If not None,
+            the function returns the diagonal of
+            :math:`\mathbf{J}^T \mathbf{W}^T \mathbf{W} \mathbf{J}``.
+        f : Ignored
+            Not used, present here for API consistency by convention.
+
+        Returns
+        -------
+        (nparam) np.ndarray
+            Array with the diagonal of ``J.T @ J``.
+
+        Notes
+        -----
+        If ``store_sensitivities`` is ``"forward_only"``, the ``G`` matrix is
+        never allocated in memory, and the diagonal is obtained by
+        accumulation, computing each element of the ``G`` matrix on the fly.
+
+        This method caches the diagonal ``G.T @ W.T @ W @ G`` and the sha256
+        hash of the diagonal of the ``W`` matrix. This way, if same weights are
+        passed to it, it reuses the cached diagonal so it doesn't need to be
+        recomputed.
+        If new weights are passed, the cache is updated with the latest
+        diagonal of ``G.T @ W.T @ W @ G``.
         """
+        # Need to assign the model, so the chiDeriv can be computed (if the
+        # model is None, the chiDeriv is going to be Zero).
         self.model = m
 
-        if W is None:
-            W = np.ones(self.survey.nD)
-        else:
-            W = W.diagonal() ** 2
-        if getattr(self, "_gtg_diagonal", None) is None:
-            diag = np.zeros(self.G.shape[1])
-            if not self.is_amplitude_data:
-                for i in range(len(W)):
-                    diag += W[i] * (self.G[i] * self.G[i])
-            else:
+        # We should probably check that W is diagonal. Let's assume it for now.
+        weights = (
+            W.diagonal() ** 2
+            if W is not None
+            else np.ones(self.survey.nD, dtype=np.float64)
+        )
+
+        # Compute gtg (G.T @ W.T @ W @ G) if it's not cached, or if the
+        # weights are not the same.
+        weights_sha256 = hashlib.sha256(weights)
+        use_cached_gtg = (
+            hasattr(self, "_gtg_diagonal")
+            and hasattr(self, "_weights_sha256")
+            and self._weights_sha256.digest() == weights_sha256.digest()
+        )
+        if not use_cached_gtg:
+            self._gtg_diagonal = self._get_gtg_diagonal(weights)
+            self._weights_sha256 = weights_sha256
+
+        # Multiply the gtg_diagonal by the derivative of the mapping
+        diagonal = mkvc(
+            (sdiag(np.sqrt(self._gtg_diagonal)) @ self.chiDeriv).power(2).sum(axis=0)
+        )
+        return diagonal
+
+    def _get_gtg_diagonal(self, weights: NDArray) -> NDArray:
+        """
+        Compute the diagonal of ``G.T @ W.T @ W @ G``.
+
+        Parameters
+        ----------
+        weights : np.ndarray
+            Weights array: diagonal of ``W.T @ W``.
+
+        Returns
+        -------
+        np.ndarray
+        """
+        match (self.engine, self.store_sensitivities, self.is_amplitude_data):
+            case ("geoana", "forward_only", _):
+                msg = (
+                    "Computing the diagonal of `G.T @ G` using "
+                    "`'forward_only'` and `'geoana'` as engine hasn't been "
+                    "implemented yet."
+                )
+                raise NotImplementedError(msg)
+            case ("choclo", "forward_only", True):
+                msg = (
+                    "Computing the diagonal of `G.T @ G` using "
+                    "`'forward_only'` and `is_amplitude_data` hasn't been "
+                    "implemented yet."
+                )
+                raise NotImplementedError(msg)
+            case ("choclo", "forward_only", False):
+                gtg_diagonal = self._gtg_diagonal_without_building_g(weights)
+            case (_, _, False):
+                # In Einstein notation, the j-th element of the diagonal is:
+                #   d_j = w_i * G_{ij} * G_{ij}
+                gtg_diagonal = np.asarray(
+                    np.einsum("i,ij,ij->j", weights, self.G, self.G)
+                )
+            case (_, _, True):
                 ampDeriv = self.ampDeriv
                 Gx = self.G[::3]
                 Gy = self.G[1::3]
                 Gz = self.G[2::3]
-                for i in range(len(W)):
+                gtg_diagonal = np.zeros(self.G.shape[1])
+                for i in range(weights.size):
                     row = (
                         ampDeriv[0, i] * Gx[i]
                         + ampDeriv[1, i] * Gy[i]
                         + ampDeriv[2, i] * Gz[i]
                     )
-                    diag += W[i] * (row * row)
-            self._gtg_diagonal = diag
-        else:
-            diag = self._gtg_diagonal
-        return mkvc((sdiag(np.sqrt(diag)) @ self.chiDeriv).power(2).sum(axis=0))
+                    gtg_diagonal += weights[i] * (row * row)
+        return gtg_diagonal
 
     def Jvec(self, m, v, f=None):
+        """
+        Dot product between sensitivity matrix and a vector.
+
+        Parameters
+        ----------
+        m : (n_param,) numpy.ndarray
+            The model parameters. This array is used to compute the ``J``
+            matrix.
+        v : (n_param,) numpy.ndarray
+            Vector used in the matrix-vector multiplication.
+        f : Ignored
+            Not used, present here for API consistency by convention.
+
+        Returns
+        -------
+        (nD,) numpy.ndarray
+
+        Notes
+        -----
+        If ``store_sensitivities`` is set to ``"forward_only"``, then the
+        matrix `G` is never fully constructed, and the dot product is computed
+        by accumulation, computing the matrix elements on the fly. Otherwise,
+        the full matrix ``G`` is constructed and stored either in memory or
+        disk.
+        """
+        # Need to assign the model, so the chiDeriv can be computed (if the
+        # model is None, the chiDeriv is going to be Zero).
         self.model = m
         dmu_dm_v = self.chiDeriv @ v
-
         Jvec = self.G @ dmu_dm_v.astype(self.sensitivity_dtype, copy=False)
 
         if self.is_amplitude_data:
@@ -381,10 +578,37 @@ class Simulation3DIntegral(BasePFSimulation):
             Jvec = Jvec.reshape((-1, 3)).T  # reshape((3, -1), order="F")
             ampDeriv_Jvec = self.ampDeriv * Jvec
             return ampDeriv_Jvec[0] + ampDeriv_Jvec[1] + ampDeriv_Jvec[2]
-        else:
-            return Jvec
+
+        return Jvec
 
     def Jtvec(self, m, v, f=None):
+        """
+        Dot product between transposed sensitivity matrix and a vector.
+
+        Parameters
+        ----------
+        m : (n_param,) numpy.ndarray
+            The model parameters. This array is used to compute the ``J``
+            matrix.
+        v : (nD,) numpy.ndarray
+            Vector used in the matrix-vector multiplication.
+        f : Ignored
+            Not used, present here for API consistency by convention.
+
+        Returns
+        -------
+        (nD,) numpy.ndarray
+
+        Notes
+        -----
+        If ``store_sensitivities`` is set to ``"forward_only"``, then the
+        matrix `G` is never fully constructed, and the dot product is computed
+        by accumulation, computing the matrix elements on the fly. Otherwise,
+        the full matrix ``G`` is constructed and stored either in memory or
+        disk.
+        """
+        # Need to assign the model, so the chiDeriv can be computed (if the
+        # model is None, the chiDeriv is going to be Zero).
         self.model = m
 
         if self.is_amplitude_data:
@@ -793,10 +1017,8 @@ class Simulation3DIntegral(BasePFSimulation):
         # 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
+        scalar_model = self.model_type == "scalar"
+        n_columns = self.nC if scalar_model else 3 * self.nC
         shape = (self.survey.nD, n_columns)
         if self.store_sensitivities == "disk":
             sensitivity_matrix = np.memmap(
@@ -812,7 +1034,6 @@ class Simulation3DIntegral(BasePFSimulation):
         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(
@@ -871,6 +1092,194 @@ class Simulation3DIntegral(BasePFSimulation):
             index_offset += n_rows
         return sensitivity_matrix
 
+    def _sensitivity_matrix_as_operator(self):
+        """
+        Create a LinearOperator for the sensitivity matrix G.
+
+        Returns
+        -------
+        scipy.sparse.linalg.LinearOperator
+        """
+        n_columns = self.nC if self.model_type == "scalar" else self.nC * 3
+        shape = (self.survey.nD, n_columns)
+        linear_op = LinearOperator(
+            shape=shape,
+            matvec=self._forward,
+            rmatvec=self._sensitivity_matrix_transpose_dot_vec,
+            dtype=np.float64,
+        )
+        return linear_op
+
+    def _sensitivity_matrix_transpose_dot_vec(self, vector):
+        """
+        Compute ``G.T @ v`` without building ``G``.
+
+        Parameters
+        ----------
+        vector : (nD) numpy.ndarray
+            Vector used in the dot product.
+
+        Returns
+        -------
+        (n_active_cells) or (3 * n_active_cells) numpy.ndarray
+        """
+        # Gather active nodes and the indices of the nodes for each active cell
+        active_nodes, active_cell_nodes = self._get_active_nodes()
+        # Get regional field
+        regional_field = self.survey.source_field.b0
+
+        # Allocate resulting array.
+        scalar_model = self.model_type == "scalar"
+        result = np.zeros(self.nC if scalar_model else 3 * self.nC)
+
+        # Define the constant factor
+        constant_factor = 1 / 4 / np.pi
+
+        # Fill the result array
+        index_offset = 0
+        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._tmi_sensitivity_t_dot_v(
+                        receivers,
+                        active_nodes,
+                        active_cell_nodes,
+                        regional_field,
+                        constant_factor,
+                        scalar_model,
+                        vector[vector_slice],
+                        result,
+                    )
+                elif component in ("tmi_x", "tmi_y", "tmi_z"):
+                    kernel_xx, kernel_yy, kernel_zz, kernel_xy, kernel_xz, kernel_yz = (
+                        CHOCLO_KERNELS[component]
+                    )
+                    self._tmi_derivative_sensitivity_t_dot_v(
+                        receivers,
+                        active_nodes,
+                        active_cell_nodes,
+                        regional_field,
+                        kernel_xx,
+                        kernel_yy,
+                        kernel_zz,
+                        kernel_xy,
+                        kernel_xz,
+                        kernel_yz,
+                        constant_factor,
+                        scalar_model,
+                        vector[vector_slice],
+                        result,
+                    )
+                else:
+                    kernel_x, kernel_y, kernel_z = CHOCLO_KERNELS[component]
+                    self._mag_sensitivity_t_dot_v(
+                        receivers,
+                        active_nodes,
+                        active_cell_nodes,
+                        regional_field,
+                        kernel_x,
+                        kernel_y,
+                        kernel_z,
+                        constant_factor,
+                        scalar_model,
+                        vector[vector_slice],
+                        result,
+                    )
+            index_offset += n_rows
+        return result
+
+    def _gtg_diagonal_without_building_g(self, weights):
+        """
+        Compute the diagonal of ``G.T @ G`` without building the ``G`` matrix.
+
+        Parameters
+        -----------
+        weights : (nD,) array
+            Array with data weights. It should be the diagonal of the ``W``
+            matrix, squared.
+
+        Returns
+        -------
+        (n_active_cells) numpy.ndarray
+        """
+        # Gather active nodes and the indices of the nodes for each active cell
+        active_nodes, active_cell_nodes = self._get_active_nodes()
+        # Get regional field
+        regional_field = self.survey.source_field.b0
+        # Define the constant factor
+        constant_factor = 1 / 4 / np.pi
+
+        # Allocate array for the diagonal
+        scalar_model = self.model_type == "scalar"
+        n_columns = self.nC if scalar_model else 3 * self.nC
+        diagonal = np.zeros(n_columns, dtype=np.float64)
+
+        # Start filling the diagonal array
+        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."
+                )
+            for component in components:
+                if component == "tmi":
+                    self._diagonal_G_T_dot_G_tmi(
+                        receivers,
+                        active_nodes,
+                        active_cell_nodes,
+                        regional_field,
+                        constant_factor,
+                        scalar_model,
+                        weights,
+                        diagonal,
+                    )
+                elif component in ("tmi_x", "tmi_y", "tmi_z"):
+                    kernel_xx, kernel_yy, kernel_zz, kernel_xy, kernel_xz, kernel_yz = (
+                        CHOCLO_KERNELS[component]
+                    )
+                    self._diagonal_G_T_dot_G_tmi_deriv(
+                        receivers,
+                        active_nodes,
+                        active_cell_nodes,
+                        regional_field,
+                        kernel_xx,
+                        kernel_yy,
+                        kernel_zz,
+                        kernel_xy,
+                        kernel_xz,
+                        kernel_yz,
+                        constant_factor,
+                        scalar_model,
+                        weights,
+                        diagonal,
+                    )
+                else:
+                    kernel_x, kernel_y, kernel_z = CHOCLO_KERNELS[component]
+                    self._diagonal_G_T_dot_G_mag(
+                        receivers,
+                        active_nodes,
+                        active_cell_nodes,
+                        regional_field,
+                        kernel_x,
+                        kernel_y,
+                        kernel_z,
+                        constant_factor,
+                        scalar_model,
+                        weights,
+                        diagonal,
+                    )
+        return diagonal
+
 
 class SimulationEquivalentSourceLayer(
     BaseEquivalentSourceLayerSimulation, Simulation3DIntegral
@@ -1037,10 +1446,8 @@ class SimulationEquivalentSourceLayer(
         # 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
+        scalar_model = self.model_type == "scalar"
+        n_columns = self.nC if scalar_model else 3 * self.nC
         shape = (self.survey.nD, n_columns)
         if self.store_sensitivities == "disk":
             sensitivity_matrix = np.memmap(
@@ -1054,7 +1461,6 @@ class SimulationEquivalentSourceLayer(
             sensitivity_matrix = np.empty(shape, dtype=self.sensitivity_dtype)
         # 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(


Test patch:
diff --git a/tests/pf/test_forward_Mag_Linear.py b/tests/pf/test_forward_Mag_Linear.py
index e593c7187..9ddc45ac8 100644
--- a/tests/pf/test_forward_Mag_Linear.py
+++ b/tests/pf/test_forward_Mag_Linear.py
@@ -5,8 +5,10 @@ import re
 import discretize
 import numpy as np
 import pytest
+from scipy.sparse import diags
 from geoana.em.static import MagneticPrism
 from scipy.constants import mu_0
+from scipy.sparse.linalg import LinearOperator, aslinearoperator
 
 import simpeg
 from simpeg import maps, utils
@@ -1034,3 +1036,706 @@ class TestnD(BaseFixtures):
         assert simulation.nD == n_data
         dpred = simulation.dpred(susceptibilities)
         assert dpred.size == simulation.nD
+
+
+@pytest.mark.parametrize(
+    "scalar_model", [True, False], ids=["scalar_model", "vector_model"]
+)
+class TestGLinearOperator(BaseFixtures):
+    """
+    Test G as a linear operator.
+    """
+
+    @pytest.fixture
+    def mapping(self, mesh, scalar_model):
+        nparams = mesh.n_cells if scalar_model else 3 * mesh.n_cells
+        return maps.IdentityMap(nP=nparams)
+
+    @pytest.mark.parametrize("parallel", [True, False], ids=["parallel", "serial"])
+    def test_G_dot_m(
+        self, survey, mesh, mapping, susceptibilities, scalar_model, parallel
+    ):
+        """Test G @ m."""
+        model_type = "scalar" if scalar_model else "vector"
+        simulation, simulation_ram = (
+            mag.simulation.Simulation3DIntegral(
+                survey=survey,
+                mesh=mesh,
+                chiMap=mapping,
+                store_sensitivities=store,
+                engine="choclo",
+                numba_parallel=parallel,
+                model_type=model_type,
+            )
+            for store in ("forward_only", "ram")
+        )
+        assert isinstance(simulation.G, LinearOperator)
+        assert isinstance(simulation_ram.G, np.ndarray)
+
+        expected = simulation_ram.G @ susceptibilities
+
+        atol = np.max(np.abs(expected)) * 1e-8
+        np.testing.assert_allclose(simulation.G @ susceptibilities, expected, atol=atol)
+
+    @pytest.mark.parametrize("parallel", [True, False], ids=["parallel", "serial"])
+    def test_G_t_dot_v(self, survey, mesh, mapping, scalar_model, parallel):
+        """Test G.T @ v."""
+        model_type = "scalar" if scalar_model else "vector"
+        simulation, simulation_ram = (
+            mag.simulation.Simulation3DIntegral(
+                survey=survey,
+                mesh=mesh,
+                chiMap=mapping,
+                store_sensitivities=store,
+                engine="choclo",
+                numba_parallel=parallel,
+                model_type=model_type,
+            )
+            for store in ("forward_only", "ram")
+        )
+        assert isinstance(simulation.G, LinearOperator)
+        assert isinstance(simulation_ram.G, np.ndarray)
+
+        vector = np.random.default_rng(seed=42).uniform(size=survey.nD)
+        expected = simulation_ram.G.T @ vector
+
+        atol = np.max(np.abs(expected)) * 1e-7
+        np.testing.assert_allclose(simulation.G.T @ vector, expected, atol=atol)
+
+    def test_not_implemented(self, survey, mesh, mapping, scalar_model):
+        """
+        Test NotImplementedError when forward_only and geoana as engine.
+        """
+        engine, store = "geoana", "forward_only"
+        model_type = "scalar" if scalar_model else "vector"
+        simulation = mag.simulation.Simulation3DIntegral(
+            survey=survey,
+            mesh=mesh,
+            chiMap=mapping,
+            store_sensitivities=store,
+            engine=engine,
+            model_type=model_type,
+        )
+        msg = re.escape(
+            "Accessing matrix G with "
+            'store_sensitivities="forward_only" and engine="geoana" '
+            "hasn't been implemented yet."
+        )
+        with pytest.raises(NotImplementedError, match=msg):
+            simulation.G
+
+
+@pytest.mark.parametrize(
+    "scalar_model", [True, False], ids=["scalar_model", "vector_model"]
+)
+class TestJacobian(BaseFixtures):
+    """
+    Test methods related to Jacobian matrix in magnetic simulation.
+    """
+
+    atol_ratio = 1e-7
+
+    @pytest.fixture(params=["identity_map", "exp_map"])
+    def mapping(self, mesh, scalar_model: bool, request):
+        nparams = mesh.n_cells if scalar_model else 3 * mesh.n_cells
+        mapping = (
+            maps.IdentityMap(nP=nparams)
+            if request.param == "identity_map"
+            else maps.ExpMap(nP=nparams)
+        )
+        return mapping
+
+    @pytest.mark.parametrize("engine", ["choclo", "geoana"])
+    def test_getJ_as_array(
+        self, survey, mesh, mapping, susceptibilities, scalar_model, engine
+    ):
+        """
+        Test the getJ method when J is an array in memory.
+        """
+        model_type = "scalar" if scalar_model else "vector"
+        simulation = mag.simulation.Simulation3DIntegral(
+            survey=survey,
+            mesh=mesh,
+            chiMap=mapping,
+            store_sensitivities="ram",
+            engine=engine,
+            model_type=model_type,
+        )
+        is_identity_map = type(mapping) is maps.IdentityMap
+        model = susceptibilities if is_identity_map else np.log(susceptibilities)
+
+        jac = simulation.getJ(model)
+        assert isinstance(jac, np.ndarray)
+        # With an identity mapping, the jacobian should be the same as G.
+        # With an exp mapping, the jacobian should be G @ the mapping derivative.
+        expected_jac = (
+            simulation.G if is_identity_map else simulation.G @ mapping.deriv(model)
+        )
+        np.testing.assert_allclose(jac, expected_jac)
+
+    @pytest.mark.parametrize(
+        "engine",
+        [
+            "choclo",
+            pytest.param(
+                "geoana",
+                marks=pytest.mark.xfail(
+                    reason="not implemented", raises=NotImplementedError
+                ),
+            ),
+        ],
+    )
+    @pytest.mark.parametrize("transpose", [False, True], ids=["J @ m", "J.T @ v"])
+    def test_getJ_as_linear_operator(
+        self, survey, mesh, mapping, susceptibilities, scalar_model, engine, transpose
+    ):
+        """
+        Test the getJ method when J is a linear operator.
+        """
+        model_type = "scalar" if scalar_model else "vector"
+        simulation = mag.simulation.Simulation3DIntegral(
+            survey=survey,
+            mesh=mesh,
+            chiMap=mapping,
+            store_sensitivities="forward_only",
+            engine=engine,
+            model_type=model_type,
+        )
+        is_identity_map = type(mapping) is maps.IdentityMap
+        model = susceptibilities if is_identity_map else np.log(susceptibilities)
+
+        jac = simulation.getJ(model)
+        assert isinstance(jac, LinearOperator)
+        if transpose:
+            vector = np.random.default_rng(seed=42).uniform(size=survey.nD)
+            result = jac.T @ vector
+            expected_result = mapping.deriv(model).T @ (simulation.G.T @ vector)
+        else:
+            result = jac @ model
+            expected_result = simulation.G @ (mapping.deriv(model).diagonal() * model)
+        np.testing.assert_allclose(result, expected_result)
+
+    @pytest.mark.parametrize("engine", ["choclo", "geoana"])
+    def test_getJ_not_implemented(
+        self, survey, mesh, mapping, susceptibilities, scalar_model, engine
+    ):
+        model_type = "scalar" if scalar_model else "vector"
+        simulation = mag.simulation.Simulation3DIntegral(
+            survey=survey,
+            mesh=mesh,
+            chiMap=mapping,
+            store_sensitivities="ram",
+            engine=engine,
+            model_type=model_type,
+            is_amplitude_data=True,
+        )
+        is_identity_map = type(mapping) is maps.IdentityMap
+        model = susceptibilities if is_identity_map else np.log(susceptibilities)
+        with pytest.raises(NotImplementedError):
+            simulation.getJ(model)
+
+    @pytest.mark.parametrize(
+        ("engine", "store_sensitivities"),
+        [
+            ("choclo", "ram"),
+            ("choclo", "forward_only"),
+            ("geoana", "ram"),
+            pytest.param(
+                "geoana",
+                "forward_only",
+                marks=pytest.mark.xfail(reason="not implemented"),
+            ),
+        ],
+    )
+    def test_Jvec(
+        self,
+        survey,
+        mesh,
+        mapping,
+        susceptibilities,
+        scalar_model,
+        engine,
+        store_sensitivities,
+    ):
+        """
+        Test the Jvec method.
+        """
+        model_type = "scalar" if scalar_model else "vector"
+        simulation = mag.simulation.Simulation3DIntegral(
+            survey=survey,
+            mesh=mesh,
+            chiMap=mapping,
+            store_sensitivities=store_sensitivities,
+            engine=engine,
+            model_type=model_type,
+            sensitivity_dtype=np.float64,
+        )
+        is_identity_map = type(mapping) is maps.IdentityMap
+        model = susceptibilities if is_identity_map else np.log(susceptibilities)
+
+        vector = np.random.default_rng(seed=42).uniform(size=susceptibilities.size)
+        result = simulation.Jvec(model, vector)
+
+        expected_jac = (
+            simulation.G
+            if is_identity_map
+            else simulation.G @ aslinearoperator(mapping.deriv(model))
+        )
+        expected = expected_jac @ vector
+
+        atol = np.max(np.abs(expected)) * self.atol_ratio
+        np.testing.assert_allclose(result, expected, atol=atol)
+
+    @pytest.mark.parametrize(
+        ("engine", "store_sensitivities"),
+        [
+            ("choclo", "ram"),
+            ("choclo", "forward_only"),
+            ("geoana", "ram"),
+            pytest.param(
+                "geoana",
+                "forward_only",
+                marks=pytest.mark.xfail(reason="not implemented"),
+            ),
+        ],
+    )
+    def test_Jtvec(
+        self,
+        survey,
+        mesh,
+        mapping,
+        susceptibilities,
+        scalar_model,
+        engine,
+        store_sensitivities,
+    ):
+        """
+        Test the Jtvec method.
+        """
+        model_type = "scalar" if scalar_model else "vector"
+        simulation = mag.simulation.Simulation3DIntegral(
+            survey=survey,
+            mesh=mesh,
+            chiMap=mapping,
+            store_sensitivities=store_sensitivities,
+            engine=engine,
+            model_type=model_type,
+            sensitivity_dtype=np.float64,
+        )
+        is_identity_map = type(mapping) is maps.IdentityMap
+        model = susceptibilities if is_identity_map else np.log(susceptibilities)
+
+        vector = np.random.default_rng(seed=42).uniform(size=survey.nD)
+        result = simulation.Jtvec(model, vector)
+
+        expected_jac = (
+            simulation.G
+            if is_identity_map
+            else simulation.G @ aslinearoperator(mapping.deriv(model))
+        )
+        expected = expected_jac.T @ vector
+
+        atol = np.max(np.abs(result)) * self.atol_ratio
+        np.testing.assert_allclose(result, expected, atol=atol)
+
+    @pytest.mark.parametrize(
+        "engine",
+        [
+            "choclo",
+            pytest.param("geoana", marks=pytest.mark.xfail(reason="not implemented")),
+        ],
+    )
+    @pytest.mark.parametrize("method", ["Jvec", "Jtvec"])
+    def test_array_vs_linear_operator(
+        self, survey, mesh, mapping, susceptibilities, scalar_model, engine, method
+    ):
+        """
+        Test methods when using "ram" and "forward_only".
+
+        They should give the same results.
+        """
+        model_type = "scalar" if scalar_model else "vector"
+        simulation_lo, simulation_ram = (
+            mag.simulation.Simulation3DIntegral(
+                survey=survey,
+                mesh=mesh,
+                chiMap=mapping,
+                store_sensitivities=store,
+                engine=engine,
+                model_type=model_type,
+                sensitivity_dtype=np.float64,
+            )
+            for store in ("forward_only", "ram")
+        )
+        match method:
+            case "Jvec":
+                vector_size = susceptibilities.size
+            case "Jtvec":
+                vector_size = survey.nD
+            case _:  # pragma: no cover
+                raise ValueError(f"Invalid method '{method}'")  # pragma: no cover
+        is_identity_map = type(mapping) is maps.IdentityMap
+        model = susceptibilities if is_identity_map else np.log(susceptibilities)
+
+        vector = np.random.default_rng(seed=42).uniform(size=vector_size)
+        result_lo = getattr(simulation_lo, method)(model, vector)
+        result_ram = getattr(simulation_ram, method)(model, vector)
+        atol = np.max(np.abs(result_ram)) * self.atol_ratio
+        np.testing.assert_allclose(result_lo, result_ram, atol=atol)
+
+    @pytest.mark.parametrize("engine", ["choclo", "geoana"])
+    @pytest.mark.parametrize("weights", [True, False])
+    def test_getJtJdiag(
+        self, survey, mesh, mapping, susceptibilities, scalar_model, engine, weights
+    ):
+        """
+        Test the ``getJtJdiag`` method with G as an array in memory.
+        """
+        model_type = "scalar" if scalar_model else "vector"
+        simulation = mag.simulation.Simulation3DIntegral(
+            survey=survey,
+            mesh=mesh,
+            chiMap=mapping,
+            store_sensitivities="ram",
+            engine=engine,
+            model_type=model_type,
+            sensitivity_dtype=np.float64,
+        )
+        is_identity_map = type(mapping) is maps.IdentityMap
+        model = susceptibilities if is_identity_map else np.log(susceptibilities)
+
+        kwargs = {}
+        if weights:
+            w_matrix = diags(np.random.default_rng(seed=42).uniform(size=survey.nD))
+            kwargs = {"W": w_matrix}
+        jtj_diag = simulation.getJtJdiag(model, **kwargs)
+
+        expected_jac = (
+            simulation.G if is_identity_map else simulation.G @ mapping.deriv(model)
+        )
+        if weights:
+            expected = np.diag(expected_jac.T @ w_matrix.T @ w_matrix @ expected_jac)
+        else:
+            expected = np.diag(expected_jac.T @ expected_jac)
+
+        atol = np.max(np.abs(jtj_diag)) * self.atol_ratio
+        np.testing.assert_allclose(jtj_diag, expected, atol=atol)
+
+    @pytest.mark.parametrize(
+        ("engine", "is_amplitude_data"),
+        [("geoana", True), ("geoana", False), ("choclo", True)],
+        ids=("geoana-amplitude_data", "geoana-regular_data", "choclo-amplitude_data"),
+    )
+    def test_getJtJdiag_not_implemented(
+        self,
+        survey,
+        mesh,
+        mapping,
+        susceptibilities,
+        scalar_model,
+        engine,
+        is_amplitude_data,
+    ):
+        """
+        Test NotImplementedErrors on ``getJtJdiag``.
+        """
+        model_type = "scalar" if scalar_model else "vector"
+        simulation = mag.simulation.Simulation3DIntegral(
+            survey=survey,
+            mesh=mesh,
+            chiMap=mapping,
+            store_sensitivities="forward_only",
+            engine=engine,
+            is_amplitude_data=is_amplitude_data,
+            model_type=model_type,
+        )
+        is_identity_map = type(mapping) is maps.IdentityMap
+        model = susceptibilities if is_identity_map else np.log(susceptibilities)
+        with pytest.raises(NotImplementedError):
+            simulation.getJtJdiag(model)
+
+    @pytest.mark.parametrize("parallel", [True, False], ids=("parallel", "serial"))
+    def test_getJtJdiag_forward_only(
+        self, survey, mesh, mapping, susceptibilities, scalar_model, parallel
+    ):
+        """
+        Test the ``getJtJdiag`` method with ``"forward_only"`` and choclo.
+        """
+        model_type = "scalar" if scalar_model else "vector"
+        simulation, simulation_ram = (
+            mag.simulation.Simulation3DIntegral(
+                survey=survey,
+                mesh=mesh,
+                chiMap=mapping,
+                store_sensitivities=store,
+                engine="choclo",
+                numba_parallel=parallel,
+                model_type=model_type,
+            )
+            for store in ("forward_only", "ram")
+        )
+        is_identity_map = type(mapping) is maps.IdentityMap
+        model = susceptibilities if is_identity_map else np.log(susceptibilities)
+
+        expected = simulation_ram.getJtJdiag(model)
+        result = simulation.getJtJdiag(model)
+
+        atol = np.max(np.abs(expected)) * 1e-8
+        np.testing.assert_allclose(result, expected, atol=atol)
+
+    @pytest.mark.parametrize("engine", ["choclo", "geoana"])
+    def test_getJtJdiag_caching(
+        self, survey, mesh, mapping, susceptibilities, scalar_model, engine
+    ):
+        """
+        Test the caching behaviour of the ``getJtJdiag`` method.
+        """
+        model_type = "scalar" if scalar_model else "vector"
+        simulation = mag.simulation.Simulation3DIntegral(
+            survey=survey,
+            mesh=mesh,
+            chiMap=mapping,
+            store_sensitivities="ram",
+            engine=engine,
+            model_type=model_type,
+        )
+
+        # Get diagonal of J.T @ J without any weight
+        is_identity_map = type(mapping) is maps.IdentityMap
+        model = susceptibilities if is_identity_map else np.log(susceptibilities)
+        jtj_diagonal_1 = simulation.getJtJdiag(model)
+        assert hasattr(simulation, "_gtg_diagonal")
+        assert hasattr(simulation, "_weights_sha256")
+        gtg_diagonal_1 = simulation._gtg_diagonal
+        weights_sha256_1 = simulation._weights_sha256
+
+        # Compute it again and make sure we get the same result
+        np.testing.assert_allclose(jtj_diagonal_1, simulation.getJtJdiag(model))
+
+        # Get a new diagonal with weights
+        weights_matrix = diags(
+            np.random.default_rng(seed=42).uniform(size=simulation.survey.nD)
+        )
+        jtj_diagonal_2 = simulation.getJtJdiag(model, W=weights_matrix)
+        assert hasattr(simulation, "_gtg_diagonal")
+        assert hasattr(simulation, "_weights_sha256")
+        gtg_diagonal_2 = simulation._gtg_diagonal
+        weights_sha256_2 = simulation._weights_sha256
+
+        # The two results should be different
+        assert not np.array_equal(jtj_diagonal_1, jtj_diagonal_2)
+        assert not np.array_equal(gtg_diagonal_1, gtg_diagonal_2)
+        assert weights_sha256_1.digest() != weights_sha256_2.digest()
+
+
+@pytest.mark.parametrize(
+    "scalar_model", [True, False], ids=["scalar_model", "vector_model"]
+)
+class TestJacobianAmplitudeData(BaseFixtures):
+    """
+    Test Jacobian related methods with ``is_amplitude_data``.
+    """
+
+    atol_ratio = 1e-7
+
+    @pytest.fixture
+    def survey(self):
+        """
+        Sample survey with fixed components bx, by, bz.
+
+        These components are assumed when working with ``is_amplitude_data=True``.
+        """
+        # Observation points
+        x = np.linspace(-20.0, 20.0, 4)
+        x, y = np.meshgrid(x, x)
+        z = 5.0 * np.ones_like(x)
+        coordinates = np.vstack((x.ravel(), y.ravel(), z.ravel())).T
+        receivers = mag.receivers.Point(coordinates, components=["bx", "by", "bz"])
+        source_field = mag.UniformBackgroundField(
+            receiver_list=[receivers],
+            amplitude=55_000,
+            inclination=12,
+            declination=-35,
+        )
+        survey = mag.survey.Survey(source_field)
+        return survey
+
+    @pytest.fixture(params=["identity_map", "exp_map"])
+    def mapping(self, mesh, scalar_model: bool, request):
+        nparams = mesh.n_cells if scalar_model else 3 * mesh.n_cells
+        mapping = (
+            maps.IdentityMap(nP=nparams)
+            if request.param == "identity_map"
+            else maps.ExpMap(nP=nparams)
+        )
+        return mapping
+
+    @pytest.mark.parametrize("engine", ["choclo", "geoana"])
+    def test_getJ_not_implemented(
+        self, survey, mesh, mapping, susceptibilities, scalar_model, engine
+    ):
+        """
+        Test the getJ method when J is an array in memory.
+        """
+        model_type = "scalar" if scalar_model else "vector"
+        simulation = mag.simulation.Simulation3DIntegral(
+            survey=survey,
+            mesh=mesh,
+            chiMap=mapping,
+            store_sensitivities="ram",
+            engine=engine,
+            model_type=model_type,
+            is_amplitude_data=True,
+            sensitivity_dtype=np.float64,
+        )
+        is_identity_map = type(mapping) is maps.IdentityMap
+        model = susceptibilities if is_identity_map else np.log(susceptibilities)
+        with pytest.raises(NotImplementedError):
+            simulation.getJ(model)
+
+    @pytest.mark.parametrize(
+        ("engine", "store_sensitivities"),
+        [
+            ("choclo", "ram"),
+            ("choclo", "forward_only"),
+            ("geoana", "ram"),
+            pytest.param(
+                "geoana",
+                "forward_only",
+                marks=pytest.mark.xfail(reason="not implemented"),
+            ),
+        ],
+    )
+    def test_Jvec(
+        self,
+        survey,
+        mesh,
+        mapping,
+        susceptibilities,
+        scalar_model,
+        engine,
+        store_sensitivities,
+    ):
+        r"""
+        Test the Jvec method.
+
+        Test the Jvec method through an alternative implementation.
+        Define a :math:`f(\chi)` forward model function that returns the norm of the
+        magnetic field given the susceptibility values of :math:`\chi`:
+
+        .. math::
+
+            f(\chi)
+                = \lvert \mathbf{B} \rvert
+                = \sqrt{B_x^2(\chi) + B_y^2(\chi) + B_z^2(\chi)}
+
+        The gradient of :math:`f(\chi)` (jacobian matrix :math:`\mathbf{J}`) can be
+        written as:
+
+        .. math::
+
+            \mathbf{J} = \mathbf{J}_x + \mathbf{J}_y + \mathbf{J}_z
+
+        where:
+
+        .. math::
+
+            \mathbf{J}_x =
+                \frac{1}{\lvert \mathbf{B} \rvert}
+                B_x(\chi)
+                \frac{\partial B_x}{\partial \chi}
+                \frac{\partial \chi}{\partial \mathbf{m}}
+
+        and :math:`\mathbf{m}` is the model.
+        """
+        model_type = "scalar" if scalar_model else "vector"
+        simulation = mag.simulation.Simulation3DIntegral(
+            survey=survey,
+            mesh=mesh,
+            chiMap=mapping,
+            store_sensitivities=store_sensitivities,
+            engine=engine,
+            model_type=model_type,
+            is_amplitude_data=True,
+            sensitivity_dtype=np.float64,
+        )
+        is_identity_map = type(mapping) is maps.IdentityMap
+        model = susceptibilities if is_identity_map else np.log(susceptibilities)
+        vector = np.random.default_rng(seed=42).uniform(size=susceptibilities.size)
+        result = simulation.Jvec(model, vector)
+
+        magnetic_field = simulation.G @ susceptibilities
+        bx, by, bz = (magnetic_field[i::3] for i in (0, 1, 2))
+        inv_amplitude = 1 / np.sqrt(bx**2 + by**2 + bz**2)
+
+        g_dot_chideriv_v = (
+            simulation.G @ aslinearoperator(mapping.deriv(model)) @ vector
+        )
+        jac_x_dot_v = diags(inv_amplitude) @ diags(bx) @ g_dot_chideriv_v[0::3]
+        jac_y_dot_v = diags(inv_amplitude) @ diags(by) @ g_dot_chideriv_v[1::3]
+        jac_z_dot_v = diags(inv_amplitude) @ diags(bz) @ g_dot_chideriv_v[2::3]
+        expected = jac_x_dot_v + jac_y_dot_v + jac_z_dot_v
+
+        atol = np.max(np.abs(expected)) * self.atol_ratio
+        np.testing.assert_allclose(result, expected, atol=atol)
+
+    @pytest.mark.parametrize(
+        ("engine", "store_sensitivities"),
+        [
+            ("choclo", "ram"),
+            ("choclo", "forward_only"),
+            ("geoana", "ram"),
+            pytest.param(
+                "geoana",
+                "forward_only",
+                marks=pytest.mark.xfail(reason="not implemented"),
+            ),
+        ],
+    )
+    def test_Jtvec(
+        self,
+        survey,
+        mesh,
+        mapping,
+        susceptibilities,
+        scalar_model,
+        engine,
+        store_sensitivities,
+    ):
+        """
+        Test the Jtvec method.
+
+        Test it similarly to Jvec, but computing the transpose of the matrices.
+        """
+        model_type = "scalar" if scalar_model else "vector"
+        simulation = mag.simulation.Simulation3DIntegral(
+            survey=survey,
+            mesh=mesh,
+            chiMap=mapping,
+            store_sensitivities=store_sensitivities,
+            engine=engine,
+            model_type=model_type,
+            is_amplitude_data=True,
+            sensitivity_dtype=np.float64,
+        )
+        is_identity_map = type(mapping) is maps.IdentityMap
+        model = susceptibilities if is_identity_map else np.log(susceptibilities)
+
+        # Need to set size as survey.nD / 3 because there's a bug in simulation.nD.
+        vector = np.random.default_rng(seed=42).uniform(size=survey.nD // 3)
+        result = simulation.Jtvec(model, vector)
+
+        magnetic_field = simulation.G @ susceptibilities
+        bx, by, bz = (magnetic_field[i::3] for i in (0, 1, 2))
+        inv_amplitude = 1 / np.sqrt(bx**2 + by**2 + bz**2)
+        v = np.array(
+            (
+                bx * inv_amplitude * vector,
+                by * inv_amplitude * vector,
+                bz * inv_amplitude * vector,
+            )
+        ).T.ravel()  # interleave the values for bx, by, bz
+        expected = mapping.deriv(model).T @ (simulation.G.T @ v)
+
+        atol = np.max(np.abs(result)) * self.atol_ratio
+        np.testing.assert_allclose(result, expected, atol=atol)