Repository: simpeg/simpeg
Pull request #1622 (2025-03-08 00:09:35)

Implement `G` matrix as `LinearOperator` in gravity simulation
#### Summary

Make the `G` property in the gravity 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. The `getJ` method also returns either a 2D Numpy array or a `LinearOperator` formed by the product between `G` and the `rhoDeriv` 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. Deprecate the public `gtg_diagonal` property. 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
<!--Example: write "Closes #NNNN" to automatically close that issue on merge.-->

Related to #1618

(possibly fixes it?)

#### 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 the Jvec and Jtvec between ram and forward_only, so we are sure the two results are correct. Don't just test them independently with a simpler version of the same thing.
- [x] Extend docs for `getJ` and `G`.
- [x] Extend docs for `Jvec` and `Jtvec`
- [x] Extend docs for `getJtJdiag`
- [x] Improve the caching mechanism of `_gtgdiag`: store the hash of the weights.
- [x] Deprecate the public `gtgdiag` property
- [x] Extend docs: explain better what happens with `forward_only` in the constructor.



Patch:
diff --git a/simpeg/potential_fields/gravity/_numba_functions.py b/simpeg/potential_fields/gravity/_numba_functions.py
index fe2e69e20..b43a45dd7 100644
--- a/simpeg/potential_fields/gravity/_numba_functions.py
+++ b/simpeg/potential_fields/gravity/_numba_functions.py
@@ -163,6 +163,312 @@ def _sensitivity_gravity(
             )
 
 
+@jit(nopython=True, parallel=False)
+def _diagonal_G_T_dot_G_serial(
+    receivers,
+    nodes,
+    cell_nodes,
+    kernel_func,
+    constant_factor,
+    weights,
+    diagonal,
+):
+    """
+    Diagonal of ``G.T @ W.T @ W @ G`` without storing ``G``, in serial.
+
+    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.
+    kernel_func : callable
+        Kernel function that will be evaluated on each node of the mesh. Choose
+        one of the kernel functions in ``choclo.prism``.
+    constant_factor : float
+        Constant factor that will be used to multiply each element of the
+        sensitivity matrix.
+    weights : (n_receivers,) numpy.ndarray
+        Array with data weights. It should be the diagonal of the ``W`` matrix,
+        squared.
+    diagonal : (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_parallel`` one for parallelized computations.
+    """
+    n_receivers = receivers.shape[0]
+    n_nodes = nodes.shape[0]
+    n_cells = cell_nodes.shape[0]
+    # Evaluate kernel function on each node, for each receiver location
+    for i in prange(n_receivers):
+        # Allocate vector for kernels evaluated on mesh nodes
+        kernels = np.empty(n_nodes)
+        for j in range(n_nodes):
+            kernels[j] = _evaluate_kernel(
+                receivers[i, 0],
+                receivers[i, 1],
+                receivers[i, 2],
+                nodes[j, 0],
+                nodes[j, 1],
+                nodes[j, 2],
+                kernel_func,
+            )
+        # Add diagonal components in the running result.
+        for k in range(n_cells):
+            diagonal[k] += (
+                weights[i]
+                * (
+                    constant_factor
+                    * kernels_in_nodes_to_cell(
+                        kernels,
+                        cell_nodes[k, :],
+                    )
+                )
+                ** 2
+            )
+
+
+@jit(nopython=True, parallel=True)
+def _diagonal_G_T_dot_G_parallel(
+    receivers,
+    nodes,
+    cell_nodes,
+    kernel_func,
+    constant_factor,
+    weights,
+    diagonal,
+):
+    """
+    Diagonal of ``G.T @ W.T @ W @ G`` without storing ``G``, in parallel.
+
+    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.
+    kernel_func : callable
+        Kernel function that will be evaluated on each node of the mesh. Choose
+        one of the kernel functions in ``choclo.prism``.
+    constant_factor : float
+        Constant factor that will be used to multiply each element of the
+        sensitivity matrix.
+    weights : (n_receivers,) numpy.ndarray
+        Array with data weights. It should be the diagonal of the ``W`` matrix,
+        squared.
+    diagonal : (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_serial`` one for serialized computations.
+
+    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.
+    """
+    n_receivers = receivers.shape[0]
+    n_nodes = nodes.shape[0]
+    n_cells = cell_nodes.shape[0]
+    # Evaluate kernel function on each node, for each receiver location
+    for i in prange(n_receivers):
+        # Allocate vector for kernels evaluated on mesh nodes
+        kernels = np.empty(n_nodes)
+        for j in range(n_nodes):
+            kernels[j] = _evaluate_kernel(
+                receivers[i, 0],
+                receivers[i, 1],
+                receivers[i, 2],
+                nodes[j, 0],
+                nodes[j, 1],
+                nodes[j, 2],
+                kernel_func,
+            )
+        # Allocate array for the diagonal elements for the current receiver.
+        local_diagonal = np.empty(n_cells)
+        for k in range(n_cells):
+            local_diagonal[k] = (
+                weights[i]
+                * (
+                    constant_factor
+                    * kernels_in_nodes_to_cell(
+                        kernels,
+                        cell_nodes[k, :],
+                    )
+                )
+                ** 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 _sensitivity_gravity_t_dot_v_serial(
+    receivers,
+    nodes,
+    cell_nodes,
+    kernel_func,
+    constant_factor,
+    vector,
+    result,
+):
+    """
+    Compute ``G.T @ v`` in serial, without building G.
+
+    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.
+    kernel_func : callable
+        Kernel function that will be evaluated on each node of the mesh. Choose
+        one of the kernel functions in ``choclo.prism``.
+    constant_factor : float
+        Constant factor that will be used to multiply each element of the
+        sensitivity matrix.
+    vector : (n_receivers) numpy.ndarray
+        Array that represents the vector used in the dot product.
+    result : (n_active_cells) numpy.ndarray
+        Running result array where the output of the dot product will be added to.
+
+    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
+    ``_sensitivity_gravity_t_dot_v_parallel``.
+    """
+    n_receivers = receivers.shape[0]
+    n_nodes = nodes.shape[0]
+    n_cells = cell_nodes.shape[0]
+    # Evaluate kernel function on each node, for each receiver location
+    for i in range(n_receivers):
+        # Allocate vector for kernels evaluated on mesh nodes
+        kernels = np.empty(n_nodes)
+        for j in range(n_nodes):
+            kernels[j] = _evaluate_kernel(
+                receivers[i, 0],
+                receivers[i, 1],
+                receivers[i, 2],
+                nodes[j, 0],
+                nodes[j, 1],
+                nodes[j, 2],
+                kernel_func,
+            )
+        # Compute the i-th row of the sensitivity matrix and multiply it by the
+        # i-th element of the vector.
+        for k in range(n_cells):
+            result[k] += (
+                constant_factor
+                * vector[i]
+                * kernels_in_nodes_to_cell(
+                    kernels,
+                    cell_nodes[k, :],
+                )
+            )
+
+
+@jit(nopython=True, parallel=True)
+def _sensitivity_gravity_t_dot_v_parallel(
+    receivers,
+    nodes,
+    cell_nodes,
+    kernel_func,
+    constant_factor,
+    vector,
+    result,
+):
+    """
+    Compute ``G.T @ v`` in parallel without building G.
+
+    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.
+    kernel_func : callable
+        Kernel function that will be evaluated on each node of the mesh. Choose
+        one of the kernel functions in ``choclo.prism``.
+    constant_factor : float
+        Constant factor that will be used to multiply each element of the
+        sensitivity matrix.
+    vector : (n_receivers) numpy.ndarray
+        Array that represents the vector used in the dot product.
+    result : (n_active_cells) numpy.ndarray
+        Running result array where the output of the dot product will be added to.
+
+    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
+    ``_sensitivity_gravity_t_dot_v_serial``.
+    """
+    n_receivers = receivers.shape[0]
+    n_nodes = nodes.shape[0]
+    n_cells = cell_nodes.shape[0]
+    # Evaluate kernel function on each node, for each receiver location
+    for i in prange(n_receivers):
+        # Allocate vector for kernels evaluated on mesh nodes
+        kernels = np.empty(n_nodes)
+        # Allocate array for the current row of the sensitivity matrix
+        local_row = np.empty(n_cells)
+        for j in range(n_nodes):
+            kernels[j] = _evaluate_kernel(
+                receivers[i, 0],
+                receivers[i, 1],
+                receivers[i, 2],
+                nodes[j, 0],
+                nodes[j, 1],
+                nodes[j, 2],
+                kernel_func,
+            )
+        # Compute fields from the kernel values
+        for k in range(n_cells):
+            local_row[k] = (
+                constant_factor
+                * vector[i]
+                * kernels_in_nodes_to_cell(
+                    kernels,
+                    cell_nodes[k, :],
+                )
+            )
+        # 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)
 def _evaluate_kernel(
     receiver_x, receiver_y, receiver_z, node_x, node_y, node_z, kernel_func
diff --git a/simpeg/potential_fields/gravity/simulation.py b/simpeg/potential_fields/gravity/simulation.py
index 93df595c2..9df88fbe7 100644
--- a/simpeg/potential_fields/gravity/simulation.py
+++ b/simpeg/potential_fields/gravity/simulation.py
@@ -1,9 +1,12 @@
 from __future__ import annotations
+import hashlib
 import warnings
 import numpy as np
+from numpy.typing import NDArray
 import scipy.constants as constants
 from geoana.kernels import prism_fz, prism_fzx, prism_fzy, prism_fzz
 from scipy.constants import G as NewtG
+from scipy.sparse.linalg import LinearOperator, aslinearoperator
 
 from simpeg import props
 from simpeg.utils import mkvc, sdiag
@@ -17,12 +20,24 @@ from ._numba_functions import (
     _sensitivity_gravity_parallel,
     _forward_gravity_serial,
     _forward_gravity_parallel,
+    _sensitivity_gravity_t_dot_v_serial,
+    _sensitivity_gravity_t_dot_v_parallel,
     _forward_gravity_2d_mesh_serial,
     _forward_gravity_2d_mesh_parallel,
     _sensitivity_gravity_2d_mesh_serial,
     _sensitivity_gravity_2d_mesh_parallel,
+    _diagonal_G_T_dot_G_serial,
+    _diagonal_G_T_dot_G_parallel,
 )
 
+try:
+    from warnings import deprecated
+except ImportError:
+    # Use the deprecated decorator provided by typing_extensions (which
+    # supports older versions of Python) if it cannot be imported from
+    # warnings.
+    from typing_extensions import deprecated
+
 if choclo is not None:
     from numba import jit
 
@@ -155,7 +170,9 @@ class Simulation3DIntegral(BasePFSimulation):
         - 'ram': sensitivities are stored in the computer's RAM
         - 'disk': sensitivities are written to a directory
         - 'forward_only': you intend only do perform a forward simulation and
-          sensitivities do not need to be stored
+          sensitivities do not need to be stored. The sensitivity matrix ``G``
+          is never created, but it'll be defined as
+          a :class:`~scipy.sparse.linalg.LinearOperator`.
 
     sensitivity_path : str, optional
         Path to store the sensitivity matrix if ``store_sensitivities`` is set
@@ -188,8 +205,6 @@ class Simulation3DIntegral(BasePFSimulation):
         super().__init__(mesh, engine=engine, numba_parallel=numba_parallel, **kwargs)
         self.rho = rho
         self.rhoMap = rhoMap
-        self._G = None
-        self._gtg_diagonal = None
         self.modelMap = self.rhoMap
 
         # Warn if n_processes has been passed
@@ -207,9 +222,13 @@ class Simulation3DIntegral(BasePFSimulation):
             if self.numba_parallel:
                 self._sensitivity_gravity = _sensitivity_gravity_parallel
                 self._forward_gravity = _forward_gravity_parallel
+                self._sensitivity_t_dot_v = _sensitivity_gravity_t_dot_v_parallel
+                self._diagonal_G_T_dot_G = _diagonal_G_T_dot_G_parallel
             else:
                 self._sensitivity_gravity = _sensitivity_gravity_serial
                 self._forward_gravity = _forward_gravity_serial
+                self._sensitivity_t_dot_v = _sensitivity_gravity_t_dot_v_serial
+                self._diagonal_G_T_dot_G = _diagonal_G_T_dot_G_serial
 
     def fields(self, m):
         """
@@ -217,8 +236,8 @@ class Simulation3DIntegral(BasePFSimulation):
 
         Parameters
         ----------
-        m : (n_active_cells,) numpy.ndarray
-            Array with values for the model.
+        m : (n_param,) numpy.ndarray
+            The model parameters.
 
         Returns
         -------
@@ -239,38 +258,154 @@ class Simulation3DIntegral(BasePFSimulation):
         return np.asarray(fields)
 
     def getJtJdiag(self, m, W=None, f=None):
-        """
-        Return the diagonal of JtJ
+        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
+        -------
+        (n_active_cells) 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 rhoDeriv can be computed (if the
         # model is None, the rhoDeriv 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])
-            for i in range(len(W)):
-                diag += W[i] * (self.G[i] * self.G[i])
-            self._gtg_diagonal = diag
-        else:
-            diag = self._gtg_diagonal
-        return mkvc((sdiag(np.sqrt(diag)) @ self.rhoDeriv).power(2).sum(axis=0))
+        # 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
 
-    def getJ(self, m, f=None):
+        # Multiply the gtg_diagonal by the derivative of the mapping
+        diagonal = mkvc(
+            (sdiag(np.sqrt(self._gtg_diagonal)) @ self.rhoDeriv).power(2).sum(axis=0)
+        )
+        return diagonal
+
+    def _get_gtg_diagonal(self, weights: NDArray) -> NDArray:
         """
-        Sensitivity matrix
+        Compute the diagonal of ``G.T @ W.T @ W @ G``.
+
+        Parameters
+        ----------
+        weights : np.ndarray
+            Weights array: diagonal of ``W.T @ W``.
+
+        Returns
+        -------
+        np.ndarray
+        """
+        gtg_diagonal = (
+            self._gtg_diagonal_without_building_g(weights)
+            if self.store_sensitivities == "forward_only"
+            else
+            # In Einstein notation, the j-th element of the diagonal is:
+            #   d_j = w_i * G_{ij} * G_{ij}
+            np.asarray(np.einsum("i,ij,ij->j", weights, self.G, self.G))
+        )
+        return gtg_diagonal
+
+    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_active_cells) 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.
         """
         # Need to assign the model, so the rhoDeriv can be computed (if the
         # model is None, the rhoDeriv is going to be Zero).
         self.model = m
-        return self.G @ self.rhoDeriv
+        rhoDeriv = (
+            self.rhoDeriv
+            if not isinstance(self.G, LinearOperator)
+            else aslinearoperator(self.rhoDeriv)
+        )
+        return self.G @ rhoDeriv
 
     def Jvec(self, m, v, f=None):
         """
-        Sensitivity times a vector
+        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 rhoDeriv can be computed (if the
         # model is None, the rhoDeriv is going to be Zero).
@@ -280,7 +415,29 @@ class Simulation3DIntegral(BasePFSimulation):
 
     def Jtvec(self, m, v, f=None):
         """
-        Sensitivity transposed times a vector
+        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 rhoDeriv can be computed (if the
         # model is None, the rhoDeriv is going to be Zero).
@@ -289,26 +446,40 @@ class Simulation3DIntegral(BasePFSimulation):
         return np.asarray(self.rhoDeriv.T @ Jtvec)
 
     @property
-    def G(self):
+    def G(self) -> NDArray | np.memmap | LinearOperator:
         """
-        Gravity forward operator
+        Gravity forward operator.
         """
-        if getattr(self, "_G", None) is None:
-            if self.engine == "choclo":
-                self._G = self._sensitivity_matrix()
-            else:
-                self._G = self.linear_operator()
+        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
 
     @property
+    @deprecated(
+        "The `gtg_diagonal` property has been deprecated. "
+        "It will be removed in SimPEG v0.25.0.",
+        category=FutureWarning,
+    )
     def gtg_diagonal(self):
         """
         Diagonal of GtG
         """
-        if getattr(self, "_gtg_diagonal", None) is None:
-            return None
-
-        return self._gtg_diagonal
+        return getattr(self, "_gtg_diagonal", None)
 
     def evaluate_integral(self, receiver_location, components):
         """
@@ -450,7 +621,7 @@ class Simulation3DIntegral(BasePFSimulation):
 
     def _sensitivity_matrix(self):
         """
-        Compute the sensitivity matrix G
+        Compute the sensitivity matrix ``G``.
 
         Returns
         -------
@@ -492,6 +663,97 @@ class Simulation3DIntegral(BasePFSimulation):
             index_offset += n_rows
         return sensitivity_matrix
 
+    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) numpy.ndarray
+        """
+        # Gather active nodes and the indices of the nodes for each active cell
+        active_nodes, active_cell_nodes = self._get_active_nodes()
+        # Allocate resulting array
+        result = np.zeros(self.nC)
+        # Start filling the result array
+        index_offset = 0
+        for components, receivers in self._get_components_and_receivers():
+            n_components = len(components)
+            n_rows = n_components * receivers.shape[0]
+            for i, component in enumerate(components):
+                kernel_func = CHOCLO_KERNELS[component]
+                conversion_factor = _get_conversion_factor(component)
+                vector_slice = slice(
+                    index_offset + i, index_offset + n_rows, n_components
+                )
+                self._sensitivity_t_dot_v(
+                    receivers,
+                    active_nodes,
+                    active_cell_nodes,
+                    kernel_func,
+                    constants.G * conversion_factor,
+                    vector[vector_slice],
+                    result,
+                )
+            index_offset += n_rows
+        return result
+
+    def _sensitivity_matrix_as_operator(self):
+        """
+        Create a LinearOperator for the sensitivity matrix G.
+
+        Returns
+        -------
+        scipy.sparse.linalg.LinearOperator
+        """
+        shape = (self.survey.nD, self.nC)
+        linear_op = LinearOperator(
+            shape=shape,
+            matvec=self._forward,
+            rmatvec=self._sensitivity_matrix_transpose_dot_vec,
+            dtype=np.float64,
+        )
+        return linear_op
+
+    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()
+        # Allocate array for the diagonal of G.T @ G
+        diagonal = np.zeros(self.nC, dtype=np.float64)
+        # Start filling the diagonal array
+        for components, receivers in self._get_components_and_receivers():
+            for component in components:
+                kernel_func = CHOCLO_KERNELS[component]
+                conversion_factor = _get_conversion_factor(component)
+                self._diagonal_G_T_dot_G(
+                    receivers,
+                    active_nodes,
+                    active_cell_nodes,
+                    kernel_func,
+                    constants.G * conversion_factor,
+                    weights,
+                    diagonal,
+                )
+        return diagonal
+
 
 class SimulationEquivalentSourceLayer(
     BaseEquivalentSourceLayerSimulation, Simulation3DIntegral


Test patch:
diff --git a/tests/pf/test_forward_Grav_Linear.py b/tests/pf/test_forward_Grav_Linear.py
index 49f757f70..cc7513852 100644
--- a/tests/pf/test_forward_Grav_Linear.py
+++ b/tests/pf/test_forward_Grav_Linear.py
@@ -1,6 +1,8 @@
 import re
 
 import pytest
+from scipy.sparse import diags
+from scipy.sparse.linalg import LinearOperator, aslinearoperator
 import discretize
 import simpeg
 from simpeg import maps
@@ -418,13 +420,11 @@ class TestsGravitySimulation:
             gravity.Simulation3DIntegral(simple_mesh, engine="choclo")
 
 
-class TestJacobianGravity:
+class BaseFixtures:
     """
-    Test methods related to Jacobian matrix in gravity simulation.
+    Base test class with some fixtures.
     """
 
-    atol_ratio = 1e-7
-
     @pytest.fixture
     def survey(self):
         # Observation points
@@ -455,6 +455,14 @@ class TestJacobianGravity:
         densities[ind_sphere] = 0.2
         return densities
 
+
+class TestJacobianGravity(BaseFixtures):
+    """
+    Test methods related to Jacobian matrix in gravity simulation.
+    """
+
+    atol_ratio = 1e-7
+
     @pytest.fixture(params=["identity_map", "exp_map"])
     def mapping(self, mesh, request):
         mapping = (
@@ -465,9 +473,9 @@ class TestJacobianGravity:
         return mapping
 
     @pytest.mark.parametrize("engine", ["choclo", "geoana"])
-    def test_getJ(self, survey, mesh, densities, mapping, engine):
+    def test_getJ_as_array(self, survey, mesh, densities, mapping, engine):
         """
-        Test the getJ method.
+        Test the getJ method when J is an array in memory.
         """
         simulation = gravity.simulation.Simulation3DIntegral(
             survey=survey,
@@ -478,6 +486,7 @@ class TestJacobianGravity:
         )
         model = mapping * densities
         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.
         identity_map = type(mapping) is maps.IdentityMap
@@ -486,8 +495,60 @@ class TestJacobianGravity:
         )
         np.testing.assert_allclose(jac, expected_jac)
 
-    @pytest.mark.parametrize("engine", ["choclo", "geoana"])
-    def test_Jvec(self, survey, mesh, densities, mapping, engine):
+    def test_getJ_as_linear_operator(self, survey, mesh, densities, mapping):
+        """
+        Test the getJ method when J is a linear operator.
+        """
+        simulation = gravity.simulation.Simulation3DIntegral(
+            survey=survey,
+            mesh=mesh,
+            rhoMap=mapping,
+            store_sensitivities="forward_only",
+            engine="choclo",
+        )
+        model = mapping * densities
+        jac = simulation.getJ(model)
+        assert isinstance(jac, LinearOperator)
+        result = jac @ model
+        expected_result = simulation.G @ (mapping.deriv(model).diagonal() * model)
+        np.testing.assert_allclose(result, expected_result)
+
+    def test_getJ_as_linear_operator_not_implemented(
+        self, survey, mesh, densities, mapping
+    ):
+        """
+        Test getJ raises NotImplementedError when forward only with geoana.
+        """
+        simulation = gravity.simulation.Simulation3DIntegral(
+            survey=survey,
+            mesh=mesh,
+            rhoMap=mapping,
+            store_sensitivities="forward_only",
+            engine="geoana",
+        )
+        model = mapping * densities
+        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.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, densities, mapping, engine, store_sensitivities):
         """
         Test the Jvec method.
         """
@@ -495,7 +556,7 @@ class TestJacobianGravity:
             survey=survey,
             mesh=mesh,
             rhoMap=mapping,
-            store_sensitivities="ram",
+            store_sensitivities=store_sensitivities,
             engine=engine,
         )
         model = mapping * densities
@@ -505,15 +566,29 @@ class TestJacobianGravity:
 
         identity_map = type(mapping) is maps.IdentityMap
         expected_jac = (
-            simulation.G if identity_map else simulation.G @ mapping.deriv(model)
+            simulation.G
+            if identity_map
+            else simulation.G @ aslinearoperator(mapping.deriv(model))
         )
         expected_dpred = expected_jac @ vector
 
         atol = np.max(np.abs(expected_dpred)) * self.atol_ratio
         np.testing.assert_allclose(dpred, expected_dpred, atol=atol)
 
-    @pytest.mark.parametrize("engine", ["choclo", "geoana"])
-    def test_Jtvec(self, survey, mesh, densities, mapping, engine):
+    @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, densities, mapping, engine, store_sensitivities):
         """
         Test the Jtvec method.
         """
@@ -521,7 +596,7 @@ class TestJacobianGravity:
             survey=survey,
             mesh=mesh,
             rhoMap=mapping,
-            store_sensitivities="ram",
+            store_sensitivities=store_sensitivities,
             engine=engine,
         )
         model = mapping * densities
@@ -531,17 +606,60 @@ class TestJacobianGravity:
 
         identity_map = type(mapping) is maps.IdentityMap
         expected_jac = (
-            simulation.G if identity_map else simulation.G @ mapping.deriv(model)
+            simulation.G
+            if 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, densities, mapping, engine, method
+    ):
+        """
+        Test methods when using "ram" and "forward_only".
+
+        They should give the same results.
+        """
+        simulation_lo, simulation_ram = (
+            gravity.simulation.Simulation3DIntegral(
+                survey=survey,
+                mesh=mesh,
+                rhoMap=mapping,
+                store_sensitivities=store,
+                engine=engine,
+            )
+            for store in ("forward_only", "ram")
+        )
+        match method:
+            case "Jvec":
+                vector_size = densities.size
+            case "Jtvec":
+                vector_size = survey.nD
+            case _:
+                raise ValueError(f"Invalid method '{method}'")
+        vector = np.random.default_rng(seed=42).uniform(size=vector_size)
+        model = mapping * densities
+        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"])
-    def test_getJtJdiag(self, survey, mesh, densities, mapping, engine):
+    @pytest.mark.parametrize("weights", [True, False])
+    def test_getJtJdiag(self, survey, mesh, densities, mapping, engine, weights):
         """
-        Test the getJtJdiag method.
+        Test the ``getJtJdiag`` method with G as an array in memory.
         """
         simulation = gravity.simulation.Simulation3DIntegral(
             survey=survey,
@@ -551,17 +669,188 @@ class TestJacobianGravity:
             engine=engine,
         )
         model = mapping * densities
-        jtj_diag = simulation.getJtJdiag(model)
+        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)
 
         identity_map = type(mapping) is maps.IdentityMap
         expected_jac = (
             simulation.G if identity_map else simulation.G @ mapping.deriv(model)
         )
-        expected = np.diag(expected_jac.T @ expected_jac)
+        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",
+        [
+            "choclo",
+            pytest.param("geoana", marks=pytest.mark.xfail(reason="not implemented")),
+        ],
+    )
+    @pytest.mark.parametrize("weights", [True, False])
+    def test_getJtJdiag_forward_only(
+        self, survey, mesh, densities, mapping, engine, weights
+    ):
+        """
+        Test the ``getJtJdiag`` method without building G.
+        """
+        simulation, simulation_ram = (
+            gravity.simulation.Simulation3DIntegral(
+                survey=survey,
+                mesh=mesh,
+                rhoMap=mapping,
+                store_sensitivities=store,
+                engine=engine,
+            )
+            for store in ("forward_only", "ram")
+        )
+        model = mapping * densities
+        kwargs = {}
+        if weights:
+            weights = np.random.default_rng(seed=42).uniform(size=survey.nD)
+            kwargs = {"W": diags(np.sqrt(weights))}
+        jtj_diag = simulation.getJtJdiag(model, **kwargs)
+        jtj_diag_ram = simulation_ram.getJtJdiag(model, **kwargs)
+
+        atol = np.max(np.abs(jtj_diag)) * self.atol_ratio
+        np.testing.assert_allclose(jtj_diag, jtj_diag_ram, atol=atol)
+
+    @pytest.mark.parametrize("engine", ("choclo", "geoana"))
+    def test_getJtJdiag_caching(self, survey, mesh, densities, mapping, engine):
+        """
+        Test the caching behaviour of the ``getJtJdiag`` method.
+        """
+        simulation = gravity.simulation.Simulation3DIntegral(
+            survey=survey,
+            mesh=mesh,
+            rhoMap=mapping,
+            store_sensitivities="ram",
+            engine=engine,
+        )
+        # Get diagonal of J.T @ J without any weight
+        model = mapping * densities
+        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()
+
+
+class TestGLinearOperator(BaseFixtures):
+    """
+    Test G as a linear operator.
+    """
+
+    @pytest.fixture
+    def mapping(self, mesh):
+        return maps.IdentityMap(nP=mesh.n_cells)
+
+    def test_not_implemented(self, survey, mesh, mapping):
+        """
+        Test NotImplementedError when using geoana as engine.
+        """
+        simulation = gravity.simulation.Simulation3DIntegral(
+            survey=survey,
+            mesh=mesh,
+            rhoMap=mapping,
+            store_sensitivities="forward_only",
+            engine="geoana",
+        )
+        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("parallel", [True, False])
+    def test_G_dot_m(self, survey, mesh, mapping, densities, parallel):
+        """Test G @ m."""
+        simulation, simulation_ram = (
+            gravity.simulation.Simulation3DIntegral(
+                survey=survey,
+                mesh=mesh,
+                rhoMap=mapping,
+                store_sensitivities=store,
+                engine="choclo",
+                numba_parallel=parallel,
+            )
+            for store in ("forward_only", "ram")
+        )
+        assert isinstance(simulation.G, LinearOperator)
+        assert isinstance(simulation_ram.G, np.ndarray)
+        np.testing.assert_allclose(
+            simulation.G @ densities, simulation_ram.G @ densities
+        )
+
+    @pytest.mark.parametrize("parallel", [True, False])
+    def test_G_t_dot_v(self, survey, mesh, mapping, parallel):
+        """Test G.T @ v."""
+        simulation, simulation_ram = (
+            gravity.simulation.Simulation3DIntegral(
+                survey=survey,
+                mesh=mesh,
+                rhoMap=mapping,
+                store_sensitivities=store,
+                engine="choclo",
+                numba_parallel=parallel,
+            )
+            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)
+        np.testing.assert_allclose(simulation.G.T @ vector, simulation_ram.G.T @ vector)
+
+
+class TestDeprecationWarning(BaseFixtures):
+    """
+    Test warnings after deprecated properties or methods of the simulation class.
+    """
+
+    def test_gtg_diagonal(self, survey, mesh):
+        """Test deprecation warning on gtg_diagonal property."""
+        mapping = maps.IdentityMap(nP=mesh.n_cells)
+        simulation = gravity.simulation.Simulation3DIntegral(
+            survey=survey,
+            mesh=mesh,
+            rhoMap=mapping,
+            store_sensitivities="ram",
+            engine="choclo",
+        )
+        msg = re.escape(
+            "The `gtg_diagonal` property has been deprecated. "
+            "It will be removed in SimPEG v0.25.0.",
+        )
+        with pytest.warns(FutureWarning, match=msg):
+            simulation.gtg_diagonal
+
 
 class TestConversionFactor:
     """Test _get_conversion_factor function."""