Repository: simpeg/simpeg (python)

Issue:
BUG: Logic for getJ
### Describe the issue:

Right now, we cache `Jmatrix` in a lot of places (which is good for saving computation!) but we are not careful about checking if the `m` that is provided to `getJ` is indeed consistent with `self.model` and therefore with the `Jmatrix` that is returned. This can lead to some confusing behaviour 

For example, if you are performing a sensitivity analysis, you could get yourself into a confusing spot 

### Reproducible code example:

```python
model1 = halfspace 
model2 = model_with_structure 

sim = dc.Simulation2DNodal(mesh, survey, storeJ=True)  # because I later want to run an inversion 

J1 = sim.getJ(model1) 
J2 = sim.getJ(model2)  # this will return J1 -- wrong!
```


### Error message:

_No response_

### Runtime information:

latest SimPEG version 

### Context for the issue:

This is a conceptual bug. Although it could be resolved if we set `storeJ=False`, the behaviour is non-intuitive and is something that we should address. 

The issue lies here https://github.com/simpeg/simpeg/blob/716e221e83b986479037ae90ba9787f53d9b7334/SimPEG/electromagnetics/static/resistivity/simulation_2d.py#L267-L278

My suggestion, include a check that the model provided is basically the same as `self.model.` If this is computationally expensive enough that we don't want to be calling it every time, then we could also allow `m=None` and we return the cached J. Something like the following: 

```
def getJ(self, m=None, f=None):
    """
    Generate Full sensitivity matrix
    """
    if m is None:
        if self._Jmatrix is not None:
            return self._Jmatrix   # early return if no model provided and J has been created

    elif not np.all_close(m, self.model):
        self.model = m    # this will clear self._Jmatrix if it exists 

    if getattr(self, "_Jmatrix", None) is None:
        if self.verbose:
            print("Calculating J and storing")     
        if f is None:
            f = self.fields(m)
        self._Jmatrix = (self._Jtvec(self.model, v=None, f=f)).T
    return self._Jmatrix
```

Thanks @anna1963 for catching this! 

Note that I only highlighted this for DC, but it will be common to wherever we have `getJ`

Pull request:
Set the model when calling `getJ` in DC and SIP simulations
#### Summary
make sure we set the model when calling getJ. This will ensure that if the model is changed, the sensitivity matrix is cleared (for nonlinear problems) and the sensitivity recomputed. This will close #1358 

#### PR Checklist
* [x] If this is a work in progress PR, set as a Draft PR
* [x] Linted my code according to the [style guides](https://docs.simpeg.xyz/content/getting_started/contributing/code-style.html).
* [ ] Added [tests](https://docs.simpeg.xyz/content/getting_started/practices.html#testing) to verify changes to the code.
* [ ] Added necessary documentation to any new functions/classes following the
      expect [style](https://docs.simpeg.xyz/content/getting_started/practices.html#documentation).
* [ ] Marked as ready for review (if this is was a draft PR), and converted 
      to a Pull Request
* [ ] Tagged ``@simpeg/simpeg-developers`` when ready for review.

#### Reference issue
Closes #1358 


Patch:
diff --git a/simpeg/electromagnetics/static/resistivity/simulation.py b/simpeg/electromagnetics/static/resistivity/simulation.py
index cd7c97016..f901f0390 100644
--- a/simpeg/electromagnetics/static/resistivity/simulation.py
+++ b/simpeg/electromagnetics/static/resistivity/simulation.py
@@ -112,6 +112,7 @@ class BaseDCSimulation(BaseElectricalPDESimulation):
         return f
 
     def getJ(self, m, f=None):
+        self.model = m
         if getattr(self, "_Jmatrix", None) is None:
             if f is None:
                 f = self.fields(m)
diff --git a/simpeg/electromagnetics/static/resistivity/simulation_2d.py b/simpeg/electromagnetics/static/resistivity/simulation_2d.py
index a405d0ccd..433d6f00f 100644
--- a/simpeg/electromagnetics/static/resistivity/simulation_2d.py
+++ b/simpeg/electromagnetics/static/resistivity/simulation_2d.py
@@ -278,10 +278,10 @@ class BaseDCSimulation2D(BaseElectricalPDESimulation):
         """
         Generate Full sensitivity matrix
         """
+        self.model = m
         if getattr(self, "_Jmatrix", None) is None:
             if self.verbose:
                 print("Calculating J and storing")
-            self.model = m
             if f is None:
                 f = self.fields(m)
             self._Jmatrix = (self._Jtvec(m, v=None, f=f)).T
diff --git a/simpeg/electromagnetics/static/spectral_induced_polarization/simulation.py b/simpeg/electromagnetics/static/spectral_induced_polarization/simulation.py
index 5abc4b4e8..b0b3344e9 100644
--- a/simpeg/electromagnetics/static/spectral_induced_polarization/simulation.py
+++ b/simpeg/electromagnetics/static/spectral_induced_polarization/simulation.py
@@ -359,6 +359,7 @@ class BaseSIPSimulation(BaseIPSimulation):
         """
         Generate Full sensitivity matrix
         """
+        self.model = m
 
         if self._Jmatrix is not None:
             return self._Jmatrix
diff --git a/simpeg/electromagnetics/static/spectral_induced_polarization/simulation_2d.py b/simpeg/electromagnetics/static/spectral_induced_polarization/simulation_2d.py
index 2a0faf490..7b1a3be00 100644
--- a/simpeg/electromagnetics/static/spectral_induced_polarization/simulation_2d.py
+++ b/simpeg/electromagnetics/static/spectral_induced_polarization/simulation_2d.py
@@ -17,6 +17,8 @@ class BaseSIPSimulation2D(BaseSIPSimulation):
         Generate Full sensitivity matrix
         """
 
+        self.model = m
+
         if self.verbose:
             print(">> Compute Sensitivity matrix")
 


Test patch:
diff --git a/tests/em/static/test_model_assignment.py b/tests/em/static/test_model_assignment.py
new file mode 100644
index 000000000..cc21f59bb
--- /dev/null
+++ b/tests/em/static/test_model_assignment.py
@@ -0,0 +1,180 @@
+"""
+Test model assignment to simulation classes
+
+Test if the `getJ` method of a few static EM simulations updates the `model`.
+These tests have been added as part of the bugfix in #1361.
+"""
+
+import pytest
+import numpy as np
+
+from discretize import TensorMesh
+from simpeg import utils
+from simpeg.maps import IdentityMap, Wires
+from simpeg.electromagnetics import resistivity as dc
+from simpeg.electromagnetics import spectral_induced_polarization as sip
+from simpeg.electromagnetics.static.utils import generate_dcip_sources_line
+
+
+class TestDCSimulations:
+    @pytest.fixture
+    def mesh_3d(self):
+        """Sample mesh."""
+        cell_size = 0.5
+        npad = 2
+        hx = [(cell_size, npad, -1.5), (cell_size, 10), (cell_size, npad, 1.5)]
+        hy = [(cell_size, npad, -1.5), (cell_size, 10), (cell_size, npad, 1.5)]
+        hz = [(cell_size, npad, -1.5), (cell_size, 10), (cell_size, npad, 1.5)]
+        mesh = TensorMesh([hx, hy, hz], x0="CCC")
+        return mesh
+
+    @pytest.fixture
+    def survey_3d(self, mesh_3d):
+        """Sample survey."""
+        xmin, xmax = mesh_3d.nodes_x.min(), mesh_3d.nodes_x.max()
+        ymin, ymax = mesh_3d.nodes_y.min(), mesh_3d.nodes_y.max()
+        x = mesh_3d.nodes_x[(mesh_3d.nodes_x > xmin) & (mesh_3d.nodes_x < xmax)]
+        y = mesh_3d.nodes_y[(mesh_3d.nodes_y > ymin) & (mesh_3d.nodes_y < ymax)]
+
+        Aloc = np.r_[1.25, 0.0, 0.0]
+        Bloc = np.r_[-1.25, 0.0, 0.0]
+        M = utils.ndgrid(x - 1.0, y, np.r_[0.0])
+        N = utils.ndgrid(x + 1.0, y, np.r_[0.0])
+        rx = dc.receivers.Dipole(M, N)
+        src = dc.sources.Dipole([rx], Aloc, Bloc)
+        survey = dc.survey.Survey([src])
+        return survey
+
+    @pytest.fixture
+    def mesh_2d(self):
+        """Sample mesh."""
+        cell_size = 0.5
+        width = 10.0
+        hx = [
+            (cell_size, 10, -1.3),
+            (cell_size, width / cell_size),
+            (cell_size, 10, 1.3),
+        ]
+        hy = [(cell_size, 3, -1.3), (cell_size, 3, 1.3)]
+        mesh = TensorMesh([hx, hy], "CN")
+        return mesh
+
+    @pytest.fixture
+    def survey_2d(self, mesh_2d):
+        """Sample survey."""
+        survey_end_points = np.array([-5.0, 5.0, 0, 0])
+
+        source_list = generate_dcip_sources_line(
+            "dipole-dipole", "volt", "2D", survey_end_points, 0.0, 5, 2.5
+        )
+        survey = dc.survey.Survey(source_list)
+        return survey
+
+    @pytest.mark.parametrize(
+        "simulation_class",
+        (dc.simulation.Simulation3DNodal, dc.simulation.Simulation3DCellCentered),
+    )
+    @pytest.mark.parametrize("storeJ", [True, False])
+    def test_simulation_3d(self, mesh_3d, survey_3d, simulation_class, storeJ):
+        """
+        Test model assignment on the ``getJ`` method of 3d simulations
+        """
+        mapping = IdentityMap(mesh_3d)
+        simulation = simulation_class(
+            mesh=mesh_3d, survey=survey_3d, sigmaMap=mapping, storeJ=storeJ
+        )
+        model_1 = np.ones(mesh_3d.nC) * 1e-2
+        model_2 = np.ones(mesh_3d.nC) * 1e-1
+        # Call `getJ` passing a model and check if it was correctly assigned
+        j_1 = simulation.getJ(model_1)
+        assert model_1 is simulation.model
+        # Call `getJ` passing a different model and check if it was correctly assigned
+        j_2 = simulation.getJ(model_2)
+        assert model_2 is simulation.model
+        # Check if the two Js are different
+        assert not np.allclose(j_1, j_2)
+
+    @pytest.mark.parametrize(
+        "simulation_class",
+        (dc.simulation_2d.Simulation2DNodal, dc.simulation_2d.Simulation2DCellCentered),
+    )
+    @pytest.mark.parametrize("storeJ", [True, False])
+    def test_simulation_2d(self, mesh_2d, survey_2d, simulation_class, storeJ):
+        """
+        Test model assignment on the ``getJ`` method of 2d simulations
+        """
+        mapping = IdentityMap(mesh_2d)
+        simulation = simulation_class(
+            mesh=mesh_2d, survey=survey_2d, sigmaMap=mapping, storeJ=storeJ
+        )
+        model_1 = np.ones(mesh_2d.nC) * 1e-2
+        model_2 = np.ones(mesh_2d.nC) * 1e-1
+        # Call `getJ` passing a model and check if it was correctly assigned
+        j_1 = simulation.getJ(model_1)
+        assert model_1 is simulation.model
+        # Call `getJ` passing a different model and check if it was correctly assigned
+        j_2 = simulation.getJ(model_2)
+        assert model_2 is simulation.model
+        # Check if the two Js are different
+        assert not np.allclose(j_1, j_2)
+
+
+class TestSIPSimulations:
+    @pytest.fixture
+    def mesh_3d(self):
+        """Sample mesh."""
+        cs = 25.0
+        hx = [(cs, 0, -1.3), (cs, 21), (cs, 0, 1.3)]
+        hy = [(cs, 0, -1.3), (cs, 21), (cs, 0, 1.3)]
+        hz = [(cs, 0, -1.3), (cs, 20)]
+        mesh = TensorMesh([hx, hy, hz], x0="CCN")
+        return mesh
+
+    @pytest.fixture
+    def survey_3d(self, mesh_3d):
+        """Sample survey."""
+        x = mesh_3d.cell_centers_x[
+            (mesh_3d.cell_centers_x > -155.0) & (mesh_3d.cell_centers_x < 155.0)
+        ]
+        y = mesh_3d.cell_centers_y[
+            (mesh_3d.cell_centers_y > -155.0) & (mesh_3d.cell_centers_y < 155.0)
+        ]
+        Aloc = np.r_[-200.0, 0.0, 0.0]
+        Bloc = np.r_[200.0, 0.0, 0.0]
+        M = utils.ndgrid(x - 25.0, y, np.r_[0.0])
+
+        times = np.arange(10) * 1e-3 + 1e-3
+        rx = sip.receivers.Pole(M, times)
+        src = sip.sources.Dipole([rx], Aloc, Bloc)
+        survey = sip.Survey([src])
+        return survey
+
+    @pytest.mark.xfail(
+        reason=(
+            "SIP simulation requires some care to pass this test. "
+            "See #1361 for more details."
+        )
+    )
+    def test_simulation_3d(self, mesh_3d, survey_3d):
+        """
+        Test model assignment on the ``getJ`` method of 3d simulations
+        """
+        wires = Wires(("eta", mesh_3d.nC), ("taui", mesh_3d.nC))
+        sigma = np.ones(mesh_3d.nC) * 1e-2
+        simulation = sip.Simulation3DNodal(
+            mesh_3d,
+            sigma=sigma,
+            survey=survey_3d,
+            etaMap=wires.eta,
+            tauiMap=wires.taui,
+        )
+        model_1 = np.r_[sigma, 1.0 / sigma]
+        model_2 = np.r_[sigma * 2, 1.0 / sigma]
+        # Call `getJ` passing a model and check if it was correctly assigned
+        j_1 = simulation.getJ(model_1)
+        assert model_1 is simulation.model
+        # Call `getJ` passing a different model and check if it was correctly assigned
+        j_2 = simulation.getJ(model_2)
+        assert model_2 is simulation.model
+        # Check if the two Js are different
+        assert not np.allclose(j_1, j_2)