Repository: simpeg/simpeg (python)

Issue:
BUG: Issue with `rhoDeriv` in gravity simulation
### Describe the issue:

I noticed that the `rhoDeriv` attribute in the gravity simulation (integral form) assumes different values before and after computing `dpred`. This is specially relevant when trying to run methods like `getJ` and `Jvec`: running `dpred` before those methods would change their output.

The following example runs `getJ` after creating the simulation, obtaining a matrix full of `Zero`s. One expects this matrix to be the sensitivity matrix itself, since the mapping for the density is the `IdentityMap`. Then the script runs `dpred` and then `getJ`, obtaining not a matrix full of `Zero`s, but a matrix of sparse matrices.

Moreover, if we compute the predicted data as the dot product of the jacobian and the model vector, we get an array full of sparse arrays. One would expect this array to be the same as the one returned by the `dpred` method.

A part of this issue is originated by the fact that the `rhoDeriv` is `Zero` when initializing the simulation, and gets modified when `dpred` is called. Another part of the issue might be related to how the dense sensitivity matrix gets multiplied by `rhoDeriv` when computing `getJ`.

### Reproducible code example:

```python
import numpy as np

from discretize import TensorMesh
from simpeg.utils import model_builder
from simpeg import maps
from simpeg.potential_fields import gravity


def get_simulation_and_densities():
    # 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 = gravity.receivers.Point(coordinates, components="gz")
    source_field = gravity.sources.SourceField(receiver_list=[receivers])
    survey = gravity.survey.Survey(source_field)

    # Mesh
    dh = 5.0
    hx = [(dh, 4)]
    mesh = TensorMesh([hx, hx, hx], "CCN")

    # Define densities
    densities = 1e-10 * np.ones(mesh.n_cells)
    ind_sphere = model_builder.get_indices_sphere(
        np.r_[0.0, 0.0, -20.0], 10.0, mesh.cell_centers
    )
    densities[ind_sphere] = 0.2

    # Simulation
    model_map = maps.IdentityMap(nP=mesh.n_cells)
    simulation = gravity.simulation.Simulation3DIntegral(
        survey=survey,
        mesh=mesh,
        rhoMap=model_map,
        store_sensitivities="ram",
        engine="choclo",
    )
    return simulation, densities

# Define simulation and run getJ
simulation, densities = get_simulation_and_densities()
jac = simulation.getJ(densities)
print(f"{jac=}\n")

# Define simulation and run dpred before running getJ
simulation, densities = get_simulation_and_densities()
expected_dpred = simulation.dpred(densities)
jac_w_dpred = simulation.getJ(densities)
print(f"{jac_w_dpred=}")

# Compare the dpred with dpred computed as dotproduct of jacobian and model
dpred = jac_w_dpred @ densities
np.allclose(dpred, expected_dpred)
```

### Error message:

```shell
jac=array([[Zero, Zero, Zero, ..., Zero, Zero, Zero],
       [Zero, Zero, Zero, ..., Zero, Zero, Zero],
       [Zero, Zero, Zero, ..., Zero, Zero, Zero],
       ...,
       [Zero, Zero, Zero, ..., Zero, Zero, Zero],
       [Zero, Zero, Zero, ..., Zero, Zero, Zero],
       [Zero, Zero, Zero, ..., Zero, Zero, Zero]], dtype=object)

jac_w_dpred=array([[<DIAgonal sparse matrix of dtype 'float64'
                with 64 stored elements (1 diagonals) and shape (64, 64)>,
        <DIAgonal sparse matrix of dtype 'float64'
                with 64 stored elements (1 diagonals) and shape (64, 64)>,
        <DIAgonal sparse matrix of dtype 'float64'
                with 64 stored elements (1 diagonals) and shape (64, 64)>, ...,
        <DIAgonal sparse matrix of dtype 'float64'
                with 64 stored elements (1 diagonals) and shape (64, 64)>,
        <DIAgonal sparse matrix of dtype 'float64'
                with 64 stored elements (1 diagonals) and shape (64, 64)>,
        <DIAgonal sparse matrix of dtype 'float64'
                with 64 stored elements (1 diagonals) and shape (64, 64)>],
       [<DIAgonal sparse matrix of dtype 'float64'
                with 64 stored elements (1 diagonals) and shape (64, 64)>,
        <DIAgonal sparse matrix of dtype 'float64'
                with 64 stored elements (1 diagonals) and shape (64, 64)>,
        <DIAgonal sparse matrix of dtype 'float64'
                with 64 stored elements (1 diagonals) and shape (64, 64)>, ...,
        <DIAgonal sparse matrix of dtype 'float64'
                with 64 stored elements (1 diagonals) and shape (64, 64)>,
        <DIAgonal sparse matrix of dtype 'float64'
                with 64 stored elements (1 diagonals) and shape (64, 64)>,
        <DIAgonal sparse matrix of dtype 'float64'
                with 64 stored elements (1 diagonals) and shape (64, 64)>],
       [<DIAgonal sparse matrix of dtype 'float64'
                with 64 stored elements (1 diagonals) and shape (64, 64)>,
        <DIAgonal sparse matrix of dtype 'float64'
                with 64 stored elements (1 diagonals) and shape (64, 64)>,
        <DIAgonal sparse matrix of dtype 'float64'
                with 64 stored elements (1 diagonals) and shape (64, 64)>, ...,
        <DIAgonal sparse matrix of dtype 'float64'
                with 64 stored elements (1 diagonals) and shape (64, 64)>,
        <DIAgonal sparse matrix of dtype 'float64'
                with 64 stored elements (1 diagonals) and shape (64, 64)>,
        <DIAgonal sparse matrix of dtype 'float64'
                with 64 stored elements (1 diagonals) and shape (64, 64)>],
       ...,
       [<DIAgonal sparse matrix of dtype 'float64'
                with 64 stored elements (1 diagonals) and shape (64, 64)>,
        <DIAgonal sparse matrix of dtype 'float64'
                with 64 stored elements (1 diagonals) and shape (64, 64)>,
        <DIAgonal sparse matrix of dtype 'float64'
                with 64 stored elements (1 diagonals) and shape (64, 64)>, ...,
        <DIAgonal sparse matrix of dtype 'float64'
                with 64 stored elements (1 diagonals) and shape (64, 64)>,
        <DIAgonal sparse matrix of dtype 'float64'
                with 64 stored elements (1 diagonals) and shape (64, 64)>,
        <DIAgonal sparse matrix of dtype 'float64'
                with 64 stored elements (1 diagonals) and shape (64, 64)>],
       [<DIAgonal sparse matrix of dtype 'float64'
                with 64 stored elements (1 diagonals) and shape (64, 64)>,
        <DIAgonal sparse matrix of dtype 'float64'
                with 64 stored elements (1 diagonals) and shape (64, 64)>,
        <DIAgonal sparse matrix of dtype 'float64'
                with 64 stored elements (1 diagonals) and shape (64, 64)>, ...,
        <DIAgonal sparse matrix of dtype 'float64'
                with 64 stored elements (1 diagonals) and shape (64, 64)>,
        <DIAgonal sparse matrix of dtype 'float64'
                with 64 stored elements (1 diagonals) and shape (64, 64)>,
        <DIAgonal sparse matrix of dtype 'float64'
                with 64 stored elements (1 diagonals) and shape (64, 64)>],
       [<DIAgonal sparse matrix of dtype 'float64'
                with 64 stored elements (1 diagonals) and shape (64, 64)>,
        <DIAgonal sparse matrix of dtype 'float64'
                with 64 stored elements (1 diagonals) and shape (64, 64)>,
        <DIAgonal sparse matrix of dtype 'float64'
                with 64 stored elements (1 diagonals) and shape (64, 64)>, ...,
        <DIAgonal sparse matrix of dtype 'float64'
                with 64 stored elements (1 diagonals) and shape (64, 64)>,
        <DIAgonal sparse matrix of dtype 'float64'
                with 64 stored elements (1 diagonals) and shape (64, 64)>,
        <DIAgonal sparse matrix of dtype 'float64'
                with 64 stored elements (1 diagonals) and shape (64, 64)>]],
      dtype=object)

Traceback (most recent call last):
  File "/home/santi/gravity_rhoderiv.py", line 55, in <module>
    np.allclose(dpred, expected_dpred)
  File "/home/santi/.miniforge3/envs/simpeg-dev/lib/python3.11/site-packages/numpy/_core/numeric.py", line 2307, in allclose
    res = all(isclose(a, b, rtol=rtol, atol=atol, equal_nan=equal_nan))
              ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/santi/.miniforge3/envs/simpeg-dev/lib/python3.11/site-packages/numpy/_core/numeric.py", line 2417, in isclose
    result = (less_equal(abs(x-y), atol + rtol * abs(y))
                             ~^~
  File "/home/santi/.miniforge3/envs/simpeg-dev/lib/python3.11/site-packages/scipy/sparse/_base.py", line 573, in __sub__
    raise NotImplementedError('subtracting a nonzero scalar from a '
NotImplementedError: subtracting a nonzero scalar from a sparse array is not supported
```

### Runtime information:

```
--------------------------------------------------------------------------------
  Date: Wed Mar 05 09:49:54 2025 PST

                OS : Linux (Debian GNU/Linux 12)
            CPU(s) : 8
           Machine : x86_64
      Architecture : 64bit
               RAM : 15.5 GiB
       Environment : Python
       File system : ext4

  Python 3.11.10 | packaged by conda-forge | (main, Oct 16 2024, 01:27:36)
  [GCC 13.3.0]

            simpeg : 0.23.1.dev15+g2e3689625.d20250218
        discretize : 0.11.2
       pymatsolver : 0.3.1
             numpy : 2.0.2
             scipy : 1.15.2
        matplotlib : 3.10.1
            geoana : 0.7.2
            libdlf : 0.3.0
            pydiso : 0.1.2
             numba : 0.60.0
              dask : 2024.10.0
           sklearn : 1.6.1
            pandas : 2.2.3
             sympy : 1.13.3
           IPython : 8.28.0
        ipywidgets : 8.1.5
            plotly : 6.0.0
               vtk : 9.3.1
   memory_profiler : 0.61.0
            choclo : 0.3.2
--------------------------------------------------------------------------------
```

### Context for the issue:

_No response_

Pull request description:
Fix bug on `getJ` of gravity simulation
#### Summary

Assign the model before returning the Jacobian matrix. Don't use Numpy's `dot` function: it doesn't work propertly when mixing Numpy arrays and Scipy's sparse arrays. Use the matmul (`@`) operator instead. Add tests that catch the bug.

#### 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

Fixes #1617 

#### What does this implement/fix?
<!--Please explain your changes.-->

#### Additional information
<!--Any additional information you think is important.-->


<!--
Once all tests pass and the code has been reviewed and approved, it will be merged into main
-->


Patch:
diff --git a/simpeg/potential_fields/gravity/simulation.py b/simpeg/potential_fields/gravity/simulation.py
index 5d52ab06a..93df595c2 100644
--- a/simpeg/potential_fields/gravity/simulation.py
+++ b/simpeg/potential_fields/gravity/simulation.py
@@ -226,6 +226,7 @@ class Simulation3DIntegral(BasePFSimulation):
             Gravity fields generated by the given model on every receiver
             location.
         """
+        # Need to assign the model, so the rho property can be accessed.
         self.model = m
         if self.store_sensitivities == "forward_only":
             # Compute the linear operation without forming the full dense G
@@ -241,6 +242,8 @@ class Simulation3DIntegral(BasePFSimulation):
         """
         Return the diagonal of JtJ
         """
+        # 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:
@@ -260,12 +263,18 @@ class Simulation3DIntegral(BasePFSimulation):
         """
         Sensitivity matrix
         """
-        return self.G.dot(self.rhoDeriv)
+        # 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
 
     def Jvec(self, m, v, f=None):
         """
         Sensitivity times a vector
         """
+        # 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
         dmu_dm_v = self.rhoDeriv @ v
         return self.G @ dmu_dm_v.astype(self.sensitivity_dtype, copy=False)
 
@@ -273,6 +282,9 @@ class Simulation3DIntegral(BasePFSimulation):
         """
         Sensitivity transposed times a vector
         """
+        # 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
         Jtvec = self.G.T @ v.astype(self.sensitivity_dtype, copy=False)
         return np.asarray(self.rhoDeriv.T @ Jtvec)
 


Test patch:
diff --git a/tests/pf/test_forward_Grav_Linear.py b/tests/pf/test_forward_Grav_Linear.py
index 063fff990..49f757f70 100644
--- a/tests/pf/test_forward_Grav_Linear.py
+++ b/tests/pf/test_forward_Grav_Linear.py
@@ -5,6 +5,7 @@ import discretize
 import simpeg
 from simpeg import maps
 from simpeg.potential_fields import gravity
+from simpeg.utils import model_builder
 from geoana.gravity import Prism
 import numpy as np
 
@@ -417,6 +418,151 @@ class TestsGravitySimulation:
             gravity.Simulation3DIntegral(simple_mesh, engine="choclo")
 
 
+class TestJacobianGravity:
+    """
+    Test methods related to Jacobian matrix in gravity simulation.
+    """
+
+    atol_ratio = 1e-7
+
+    @pytest.fixture
+    def survey(self):
+        # 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 = gravity.receivers.Point(coordinates, components="gz")
+        source_field = gravity.sources.SourceField(receiver_list=[receivers])
+        survey = gravity.survey.Survey(source_field)
+        return survey
+
+    @pytest.fixture
+    def mesh(self):
+        # Mesh
+        dh = 5.0
+        hx = [(dh, 4)]
+        mesh = discretize.TensorMesh([hx, hx, hx], "CCN")
+        return mesh
+
+    @pytest.fixture
+    def densities(self, mesh):
+        # Define densities
+        densities = 1e-10 * np.ones(mesh.n_cells)
+        ind_sphere = model_builder.get_indices_sphere(
+            np.r_[0.0, 0.0, -20.0], 10.0, mesh.cell_centers
+        )
+        densities[ind_sphere] = 0.2
+        return densities
+
+    @pytest.fixture(params=["identity_map", "exp_map"])
+    def mapping(self, mesh, request):
+        mapping = (
+            maps.IdentityMap(nP=mesh.n_cells)
+            if request.param == "identity_map"
+            else maps.ExpMap(nP=mesh.n_cells)
+        )
+        return mapping
+
+    @pytest.mark.parametrize("engine", ["choclo", "geoana"])
+    def test_getJ(self, survey, mesh, densities, mapping, engine):
+        """
+        Test the getJ method.
+        """
+        simulation = gravity.simulation.Simulation3DIntegral(
+            survey=survey,
+            mesh=mesh,
+            rhoMap=mapping,
+            store_sensitivities="ram",
+            engine=engine,
+        )
+        model = mapping * densities
+        jac = simulation.getJ(model)
+        # 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
+        expected_jac = (
+            simulation.G if identity_map else simulation.G @ mapping.deriv(model)
+        )
+        np.testing.assert_allclose(jac, expected_jac)
+
+    @pytest.mark.parametrize("engine", ["choclo", "geoana"])
+    def test_Jvec(self, survey, mesh, densities, mapping, engine):
+        """
+        Test the Jvec method.
+        """
+        simulation = gravity.simulation.Simulation3DIntegral(
+            survey=survey,
+            mesh=mesh,
+            rhoMap=mapping,
+            store_sensitivities="ram",
+            engine=engine,
+        )
+        model = mapping * densities
+
+        vector = np.random.default_rng(seed=42).uniform(size=densities.size)
+        dpred = simulation.Jvec(model, vector)
+
+        identity_map = type(mapping) is maps.IdentityMap
+        expected_jac = (
+            simulation.G if identity_map else simulation.G @ 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):
+        """
+        Test the Jtvec method.
+        """
+        simulation = gravity.simulation.Simulation3DIntegral(
+            survey=survey,
+            mesh=mesh,
+            rhoMap=mapping,
+            store_sensitivities="ram",
+            engine=engine,
+        )
+        model = mapping * densities
+
+        vector = np.random.default_rng(seed=42).uniform(size=survey.nD)
+        result = simulation.Jtvec(model, vector)
+
+        identity_map = type(mapping) is maps.IdentityMap
+        expected_jac = (
+            simulation.G if identity_map else simulation.G @ 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", "geoana"])
+    def test_getJtJdiag(self, survey, mesh, densities, mapping, engine):
+        """
+        Test the getJtJdiag method.
+        """
+        simulation = gravity.simulation.Simulation3DIntegral(
+            survey=survey,
+            mesh=mesh,
+            rhoMap=mapping,
+            store_sensitivities="ram",
+            engine=engine,
+        )
+        model = mapping * densities
+        jtj_diag = simulation.getJtJdiag(model)
+
+        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)
+
+        atol = np.max(np.abs(jtj_diag)) * self.atol_ratio
+        np.testing.assert_allclose(jtj_diag, expected, atol=atol)
+
+
 class TestConversionFactor:
     """Test _get_conversion_factor function."""