Repository: simpeg/simpeg
Pull request #1638: Fix `getJ` method in TDEM and FDEM 1D simulations
Author: login: ghwilliams
Created: 2025-03-31T12:14:45Z | Merged: 2025-04-24T20:52:49Z

#### Summary
@santisoler  Here we're gonna work on a solution for issue #1611.
Initially, I've included a test in tests/em/tdem/test_TDEM_Simulation1DLayered_getJ.py.
In the following steps, we're gonna work on making the necessary changes to the relevant classes/methods as already discussed during some previous dev meetings. 

#### Squash and Merge summary

Standardize the `getJ` methods of TDEM and FDEM 1D simulations so they return a dense array instead of a dictionary with blocks of the `J` matrix. Make the old `getJ` method private by renaming it to `_getJ`. The new `getJ` method reuses the `_getJ` and builds the while `J` as a dense array before returning it. Update `Jvec` and `Jtvec` so they use the new private method. Add tests to check the behaviour of the new public method.

#### 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/latest/content/getting_started/contributing/code-style.html).
* [x] 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.


Diff:
diff --git a/simpeg/electromagnetics/base_1d.py b/simpeg/electromagnetics/base_1d.py
index afba170629..22609872e1 100644
--- a/simpeg/electromagnetics/base_1d.py
+++ b/simpeg/electromagnetics/base_1d.py
@@ -334,7 +334,7 @@ def compute_complex_mu(self, frequencies):
             return mu_complex
 
     def Jvec(self, m, v, f=None):
-        Js = self.getJ(m, f=f)
+        Js = self._getJ(m, f=f)
         out = 0.0
         if self.hMap is not None:
             out = out + Js["dh"] @ (self.hDeriv @ v)
@@ -347,7 +347,7 @@ def Jvec(self, m, v, f=None):
         return out
 
     def Jtvec(self, m, v, f=None):
-        Js = self.getJ(m, f=f)
+        Js = self._getJ(m, f=f)
         out = 0.0
         if self.hMap is not None:
             out = out + self.hDeriv.T @ (Js["dh"].T @ v)
@@ -601,7 +601,7 @@ def get_threshold(self, uncert):
 
     def getJtJdiag(self, m, W=None, f=None):
         if getattr(self, "_gtgdiag", None) is None:
-            Js = self.getJ(m, f=f)
+            Js = self._getJ(m, f=f)
             if W is None:
                 W = np.ones(self.survey.nD)
             else:
diff --git a/simpeg/electromagnetics/frequency_domain/simulation_1d.py b/simpeg/electromagnetics/frequency_domain/simulation_1d.py
index f7fe4d867e..57806b6037 100644
--- a/simpeg/electromagnetics/frequency_domain/simulation_1d.py
+++ b/simpeg/electromagnetics/frequency_domain/simulation_1d.py
@@ -127,7 +127,27 @@ def fields(self, m):
 
         return self._project_to_data(v)
 
-    def getJ(self, m, f=None):
+    def _getJ(self, m, f=None):
+        """Build Jacobian matrix by blocks.
+
+        This method builds the Jacobian matrix by blocks, each block for a particular
+        invertible property (receiver height, conductivity, permeability, layer
+        thickness). Each block of the Jacobian matrix is stored within a dictionary.
+
+        Parameters
+        ----------
+        m : (n_param,) numpy.ndarray
+            The model parameters.
+        f : Ignored
+            Not used, present here for API consistency by convention.
+
+        Returns
+        -------
+        dict
+            Dictionary containing the blocks of the Jacobian matrix for the invertible
+            properties. The keys of the dictionary can be `"dh"`, `"ds"`, `"dmu"`, and
+            `"dthick"`.
+        """
         self.model = m
         if getattr(self, "_J", None) is None:
             self._J = {}
@@ -244,6 +264,44 @@ def getJ(self, m, f=None):
                     self._J["dthick"] = self._project_to_data(v_dthick)
         return self._J
 
+    def getJ(self, m, f=None):
+        r"""Get the Jacobian matrix.
+
+        This method generates and stores the full Jacobian matrix for the
+        model provided. I.e.:
+
+        .. math::
+            \mathbf{J} = \dfrac{\partial f(\mu(\mathbf{m}))}{\partial \mathbf{m}}
+
+        where :math:`f()` is the forward modelling function, :math:`\mu()` is the
+        mapping, and :math:`\mathbf{m}` is the model vector.
+
+        Parameters
+        ----------
+        m : (n_param,) numpy.ndarray
+            The model parameters.
+        f : Ignored
+            Not used, present here for API consistency by convention.
+
+        Returns
+        -------
+        (n_data, n_param) numpy.ndarray
+            The full Jacobian matrix.
+        """
+        Js = self._getJ(m, f=f)
+        # Map parameters with their corresponding derivatives
+        param_and_derivs = {
+            "dh": self.hDeriv,
+            "ds": self.sigmaDeriv,
+            "dmu": self.muDeriv,
+            "dthick": self.thicknessesDeriv,
+        }
+
+        # Compute J matrix
+        J = sum(Js[param] @ param_and_derivs[param] for param in Js)
+
+        return J
+
     def _project_to_data(self, v):
         i_dat = 0
         i_v = 0
diff --git a/simpeg/electromagnetics/time_domain/simulation_1d.py b/simpeg/electromagnetics/time_domain/simulation_1d.py
index e2f5b8a6bb..1e72af8527 100644
--- a/simpeg/electromagnetics/time_domain/simulation_1d.py
+++ b/simpeg/electromagnetics/time_domain/simulation_1d.py
@@ -249,7 +249,27 @@ def fields(self, m):
 
         return self._project_to_data(v.T)
 
-    def getJ(self, m, f=None):
+    def _getJ(self, m, f=None):
+        """Build Jacobian matrix by blocks.
+
+        This method builds the Jacobian matrix by blocks, each block for a particular
+        invertible property (receiver height, conductivity, permeability, layer
+        thickness). Each block of the Jacobian matrix is stored within a dictionary.
+
+        Parameters
+        ----------
+        m : (n_param,) numpy.ndarray
+            The model parameters.
+        f : Ignored
+            Not used, present here for API consistency by convention.
+
+        Returns
+        -------
+        dict
+            Dictionary containing the blocks of the Jacobian matrix for the invertible
+            properties. The keys of the dictionary can be `"dh"`, `"ds"`, `"dmu"`, and
+            `"dthick"`.
+        """
         self.model = m
         if getattr(self, "_J", None) is None:
             self._J = {}
@@ -346,6 +366,45 @@ def getJ(self, m, f=None):
                     self._J["dthick"] = self._project_to_data(v_dthick)
         return self._J
 
+    def getJ(self, m, f=None):
+        r"""Get the Jacobian matrix.
+
+        This method generates and stores the full Jacobian matrix for the
+        model provided. I.e.:
+
+        .. math::
+            \mathbf{J} = \dfrac{\partial f(\mu(\mathbf{m}))}{\partial \mathbf{m}}
+
+        where :math:`f()` is the forward modelling function, :math:`\mu()` is the
+        mapping, and :math:`\mathbf{m}` is the model vector.
+
+        Parameters
+        ----------
+        m : (n_param,) numpy.ndarray
+            The model parameters.
+        f : Ignored
+            Not used, present here for API consistency by convention.
+
+        Returns
+        -------
+        (n_data, n_param) numpy.ndarray
+            The full Jacobian matrix.
+        """
+        Js = self._getJ(m, f=f)
+
+        # Map parameters with their corresponding derivatives
+        param_and_derivs = {
+            "dh": self.hDeriv,
+            "ds": self.sigmaDeriv,
+            "dmu": self.muDeriv,
+            "dthick": self.thicknessesDeriv,
+        }
+
+        # Compute J matrix
+        J = sum(Js[param] @ param_and_derivs[param] for param in Js)
+
+        return J
+
     def _project_to_data(self, v):
         As = self._As
         if v.ndim == 3:
diff --git a/tests/em/em1d/test_EM1D_FD_fwd.py b/tests/em/em1d/test_EM1D_FD_fwd.py
index 4986024335..88a7498d5a 100644
--- a/tests/em/em1d/test_EM1D_FD_fwd.py
+++ b/tests/em/em1d/test_EM1D_FD_fwd.py
@@ -592,7 +592,7 @@ def test_rx_loc_shapes(rx_class, n_locs1, n_locs2, orientation, component):
 
     sim.sigmaMap = maps.IdentityMap(nP=1)
     # make sure forming J works
-    J = sim.getJ(np.ones(1))["ds"]
+    J = sim.getJ(np.ones(1))
     assert J.shape == (n_d, 1)
 
     # and all of its values are the same too:
diff --git a/tests/em/em1d/test_EM1D_FD_getJ.py b/tests/em/em1d/test_EM1D_FD_getJ.py
new file mode 100644
index 0000000000..cf18744e2f
--- /dev/null
+++ b/tests/em/em1d/test_EM1D_FD_getJ.py
@@ -0,0 +1,107 @@
+"""
+Test the getJ method of FDEM 1D simulation.
+"""
+
+import numpy as np
+import simpeg.electromagnetics.frequency_domain as fdem
+from simpeg import maps
+
+
+def create_simulation_and_conductivities(identity_mapping: bool):
+    # Create Survey
+    # -------------
+    # Source properties
+    frequencies = np.r_[382, 1822, 7970, 35920, 130100]  # frequencies in Hz
+    source_location = np.array([0.0, 0.0, 30.0])  # (3, ) numpy.array_like
+    source_orientation = "z"  # "x", "y" or "z"
+    moment = 1.0  # dipole moment in Am^2
+
+    # Receiver properties
+    receiver_locations = np.array([10.0, 0.0, 30.0])  # or (N, 3) numpy.ndarray
+    receiver_orientation = "z"  # "x", "y" or "z"
+    data_type = "ppm"  # "secondary", "total" or "ppm"
+
+    source_list = []  # create empty list for source objects
+
+    # loop over all sources
+    for freq in frequencies:
+        # Define receivers that measure real and imaginary component
+        # magnetic field data in ppm.
+        receiver_list = []
+        receiver_list.append(
+            fdem.receivers.PointMagneticFieldSecondary(
+                receiver_locations,
+                orientation=receiver_orientation,
+                data_type=data_type,
+                component="real",
+            )
+        )
+        receiver_list.append(
+            fdem.receivers.PointMagneticFieldSecondary(
+                receiver_locations,
+                orientation=receiver_orientation,
+                data_type=data_type,
+                component="imag",
+            )
+        )
+
+        # Define a magnetic dipole source at each frequency
+        source_list.append(
+            fdem.sources.MagDipole(
+                receiver_list=receiver_list,
+                frequency=freq,
+                location=source_location,
+                orientation=source_orientation,
+                moment=moment,
+            )
+        )
+
+    # Define the survey
+    survey = fdem.survey.Survey(source_list)
+
+    # Defining a 1D Layered Earth Model
+    # ---------------------------------
+    # Define layer thicknesses (m)
+    thicknesses = np.array([20.0, 40.0])
+
+    # Define layer conductivities (S/m)
+    conductivities = np.r_[0.1, 1.0, 0.1]
+
+    # Define a mapping
+    n_layers = len(conductivities)
+    model_mapping = (
+        maps.IdentityMap(nP=n_layers) if identity_mapping else maps.ExpMap(nP=n_layers)
+    )
+
+    # Define the Forward Simulation, Predict Data and Plot
+    # ----------------------------------------------------
+    simulation = fdem.Simulation1DLayered(
+        survey=survey,
+        thicknesses=thicknesses,
+        sigmaMap=model_mapping,
+    )
+
+    return simulation, conductivities
+
+
+def test_getJ():
+    """
+    Test if getJ returns different J matrices after passing different maps.
+    """
+    dpreds, jacobians = [], []
+
+    # Compute dpred and J using an identity map and an exp map
+    for identity_mapping in (True, False):
+        simulation, conductivities = create_simulation_and_conductivities(
+            identity_mapping
+        )
+        model = conductivities if identity_mapping else np.log(conductivities)
+        dpreds.append(simulation.dpred(model))
+        jac = simulation.getJ(model)
+        jacobians.append(jac)
+
+    # The two dpreds should be equal
+    assert np.allclose(*dpreds)
+
+    # The two J matrices should not be equal
+    assert not np.allclose(*jacobians, atol=0.0)
diff --git a/tests/em/em1d/test_EM1D_FD_jac_layers2.py b/tests/em/em1d/test_EM1D_FD_jac_layers2.py
new file mode 100644
index 0000000000..dbbe2154b8
--- /dev/null
+++ b/tests/em/em1d/test_EM1D_FD_jac_layers2.py
@@ -0,0 +1,490 @@
+from simpeg import maps
+from discretize import tests, TensorMesh
+import simpeg.electromagnetics.frequency_domain as fdem
+import numpy as np
+from scipy.constants import mu_0
+from scipy.sparse import diags
+
+
+class TestEM1D_FD_Jacobian_MagDipole:
+
+    # Tests 2nd order convergence of Jvec and Jtvec for magnetic dipole sources.
+    # - All src and rx orientations
+    # - All rx components
+    # - Span many frequencies
+    # - Tests derivatives wrt sigma, mu, thicknesses and h
+    def setup_class(self):
+        # Layers and topography
+        nearthick = np.logspace(-1, 1, 5)
+        deepthick = np.logspace(1, 2, 10)
+        thicknesses = np.r_[nearthick, deepthick]
+        topo = np.r_[0.0, 0.0, 100.0]
+
+        # Survey Geometry
+        height = 1e-5
+        src_location = np.array([0.0, 0.0, 100.0 + height])
+        rx_location = np.array([5.0, 5.0, 100.0 + height])
+        frequencies = np.logspace(1, 8, 9)
+        orientations = ["x", "y", "z"]
+        components = ["real", "imag", "both"]
+
+        # Define sources and receivers
+        source_list = []
+        for f in frequencies:
+            for tx_orientation in orientations:
+                receiver_list = []
+
+                for rx_orientation in orientations:
+                    for comp in components:
+                        receiver_list.append(
+                            fdem.receivers.PointMagneticFieldSecondary(
+                                rx_location, orientation=rx_orientation, component=comp
+                            )
+                        )
+
+                source_list.append(
+                    fdem.sources.MagDipole(
+                        receiver_list,
+                        frequency=f,
+                        location=src_location,
+                        orientation=tx_orientation,
+                    )
+                )
+
+        # Survey
+        survey = fdem.Survey(source_list)
+
+        self.topo = topo
+        self.survey = survey
+        self.showIt = False
+        self.height = height
+        self.frequencies = frequencies
+        self.thicknesses = thicknesses
+        self.nlayers = len(thicknesses) + 1
+
+        wire_map = maps.Wires(
+            ("mu", self.nlayers),
+            ("sigma", self.nlayers),
+            ("h", 1),
+            ("thicknesses", self.nlayers - 1),
+        )
+        self.sigma_map = maps.ExpMap(nP=self.nlayers) * wire_map.sigma
+        self.mu_map = maps.ExpMap(nP=self.nlayers) * wire_map.mu
+        self.thicknesses_map = maps.ExpMap(nP=self.nlayers - 1) * wire_map.thicknesses
+        nP = len(source_list)
+        surject_mesh = TensorMesh([np.ones(nP)])
+        self.h_map = maps.SurjectFull(surject_mesh) * maps.ExpMap(nP=1) * wire_map.h
+
+        sim = fdem.Simulation1DLayered(
+            survey=self.survey,
+            sigmaMap=self.sigma_map,
+            muMap=self.mu_map,
+            thicknessesMap=self.thicknesses_map,
+            hMap=self.h_map,
+            topo=self.topo,
+        )
+
+        self.sim = sim
+
+    def test_EM1DFDJvec_Layers(self):
+        # Conductivity
+        sigma_half = 0.01
+        sigma_blk = 0.1
+        sig = np.ones(self.nlayers) * sigma_half
+        sig[3] = sigma_blk
+
+        # Permeability
+        mu_half = mu_0
+        mu_blk = 2 * mu_0
+        mu = np.ones(self.nlayers) * mu_half
+        mu[3] = mu_blk
+
+        # General model
+        m_1D = np.r_[
+            np.log(mu), np.log(sig), np.log(self.height), np.log(self.thicknesses)
+        ]
+
+        def fwdfun(m):
+            resp = self.sim.dpred(m)
+            return resp
+            # return Hz
+
+        def jacfun(m, dm):
+            Jvec = self.sim.Jvec(m, dm)
+            return Jvec
+
+        def derChk(m):
+            return [fwdfun(m), lambda mx: jacfun(m, mx)]
+
+        dm = m_1D * 0.5
+
+        passed = tests.check_derivative(
+            derChk, m_1D, num=4, dx=dm, plotIt=False, eps=1e-15, random_seed=9186724
+        )
+        assert passed
+
+    def test_EM1DFDJtvec_Layers(self):
+        # Conductivity
+        sigma_half = 0.01
+        sigma_blk = 0.1
+        sig = np.ones(self.nlayers) * sigma_half
+        sig[3] = sigma_blk
+
+        # Permeability
+        mu_half = mu_0
+        mu_blk = 2 * mu_0
+        mu = np.ones(self.nlayers) * mu_half
+        mu[3] = mu_blk
+
+        # General model
+        m_true = np.r_[
+            np.log(mu), np.log(sig), np.log(self.height), np.log(self.thicknesses)
+        ]
+
+        dobs = self.sim.dpred(m_true)
+
+        m_ini = np.r_[
+            np.log(np.ones(self.nlayers) * 1.5 * mu_half),
+            np.log(np.ones(self.nlayers) * sigma_half),
+            np.log(0.5 * self.height),
+            np.log(self.thicknesses) * 0.9,
+        ]
+        resp_ini = self.sim.dpred(m_ini)
+        dr = resp_ini - dobs
+
+        def misfit(m, dobs):
+            dpred = self.sim.dpred(m)
+            misfit = np.linalg.norm(dpred - dobs) ** 2
+            dmisfit = 2.0 * self.sim.Jtvec(
+                m, dr
+            )  # derivative of ||dpred - dobs||^2 gives factor of 2
+            return misfit, dmisfit
+
+        def derChk(m):
+            return misfit(m, dobs)
+
+        passed = tests.check_derivative(
+            derChk, m_ini, num=4, plotIt=False, eps=1e-27, random_seed=2345
+        )
+        assert passed
+
+    def test_jtjdiag(self):
+        # Conductivity
+        sigma_half = 0.01
+        sigma_blk = 0.1
+        sig = np.ones(self.nlayers) * sigma_half
+        sig[3] = sigma_blk
+
+        # Permeability
+        mu_half = mu_0
+        mu_blk = 2 * mu_0
+        mu = np.ones(self.nlayers) * mu_half
+        mu[3] = mu_blk
+
+        # General model
+        model = np.r_[
+            np.log(mu), np.log(sig), np.log(self.height), np.log(self.thicknesses)
+        ]
+
+        rng = np.random.default_rng(seed=42)
+        weights_matrix = diags(rng.random(size=self.sim.survey.nD))
+        jtj_diag = self.sim.getJtJdiag(model, W=weights_matrix)
+
+        J = self.sim.getJ(model)
+        expected = np.diag(J.T @ weights_matrix.T @ weights_matrix @ J)
+        np.testing.assert_allclose(expected, jtj_diag)
+
+
+class TestEM1D_FD_Jacobian_CircularLoop:
+    # Tests 2nd order convergence of Jvec and Jtvec for horizontal loop sources.
+    # - All rx orientations
+    # - All rx components
+    # - Span many frequencies
+    # - Tests derivatives wrt sigma, mu, thicknesses and h
+    def setup_class(self):
+        nearthick = np.logspace(-1, 1, 5)
+        deepthick = np.logspace(1, 2, 10)
+        thicknesses = np.r_[nearthick, deepthick]
+        topo = np.r_[0.0, 0.0, 100.0]
+        height = 1e-5
+
+        src_location = np.array([0.0, 0.0, 100.0 + height])
+        rx_location = np.array([0.0, 0.0, 100.0 + height])
+        frequencies = np.logspace(1, 8, 9)
+        orientations = ["x", "y", "z"]
+        components = ["real", "imag", "both"]
+        I = 1.0
+        a = 10.0
+
+        # Define sources and receivers
+        source_list = []
+        for f in frequencies:
+            receiver_list = []
+
+            for rx_orientation in orientations:
+                for comp in components:
+                    receiver_list.append(
+                        fdem.receivers.PointMagneticFieldSecondary(
+                            rx_location, orientation=rx_orientation, component=comp
+                        )
+                    )
+
+            source_list.append(
+                fdem.sources.CircularLoop(
+                    receiver_list, f, src_location, radius=a, current=I
+                )
+            )
+
+        # Survey
+        survey = fdem.Survey(source_list)
+
+        self.topo = topo
+        self.survey = survey
+        self.showIt = False
+        self.height = height
+        self.frequencies = frequencies
+        self.thicknesses = thicknesses
+        self.nlayers = len(thicknesses) + 1
+
+        nP = len(source_list)
+
+        wire_map = maps.Wires(
+            ("sigma", self.nlayers),
+            ("mu", self.nlayers),
+            ("thicknesses", self.nlayers - 1),
+            ("h", 1),
+        )
+        self.sigma_map = maps.ExpMap(nP=self.nlayers) * wire_map.sigma
+        self.mu_map = maps.ExpMap(nP=self.nlayers) * wire_map.mu
+        self.thicknesses_map = maps.ExpMap(nP=self.nlayers - 1) * wire_map.thicknesses
+        surject_mesh = TensorMesh([np.ones(nP)])
+        self.h_map = maps.SurjectFull(surject_mesh) * maps.ExpMap(nP=1) * wire_map.h
+
+        sim = fdem.Simulation1DLayered(
+            survey=self.survey,
+            sigmaMap=self.sigma_map,
+            muMap=self.mu_map,
+            thicknessesMap=self.thicknesses_map,
+            hMap=self.h_map,
+            topo=self.topo,
+        )
+
+        self.sim = sim
+
+    def test_EM1DFDJvec_Layers(self):
+        # Conductivity
+        sigma_half = 0.01
+        sigma_blk = 0.1
+        sig = np.ones(self.nlayers) * sigma_half
+        sig[3] = sigma_blk
+
+        # Permeability
+        mu_half = mu_0
+        mu_blk = 2 * mu_0
+        mu = np.ones(self.nlayers) * mu_half
+        mu[3] = mu_blk
+
+        # General model
+        m_1D = np.r_[
+            np.log(sig), np.log(mu), np.log(self.thicknesses), np.log(self.height)
+        ]
+
+        def fwdfun(m):
+            resp = self.sim.dpred(m)
+            return resp
+            # return Hz
+
+        def jacfun(m, dm):
+            Jvec = self.sim.Jvec(m, dm)
+            return Jvec
+
+        def derChk(m):
+            return [fwdfun(m), lambda mx: jacfun(m, mx)]
+
+        dm = m_1D * 0.5
+
+        passed = tests.check_derivative(
+            derChk, m_1D, num=4, dx=dm, plotIt=False, eps=1e-15, random_seed=664
+        )
+        assert passed
+
+    def test_EM1DFDJtvec_Layers(self):
+        # Conductivity
+        sigma_half = 0.01
+        sigma_blk = 0.1
+        sig = np.ones(self.nlayers) * sigma_half
+        sig[3] = sigma_blk
+
+        # Permeability
+        mu_half = mu_0
+        mu_blk = 2 * mu_0
+        mu = np.ones(self.nlayers) * mu_half
+        mu[3] = mu_blk
+
+        # General model
+        m_true = np.r_[
+            np.log(sig), np.log(mu), np.log(self.thicknesses), np.log(self.height)
+        ]
+
+        dobs = self.sim.dpred(m_true)
+
+        m_ini = np.r_[
+            np.log(np.ones(self.nlayers) * sigma_half),
+            np.log(np.ones(self.nlayers) * 1.5 * mu_half),
+            np.log(self.thicknesses) * 0.9,
+            np.log(0.5 * self.height),
+        ]
+        resp_ini = self.sim.dpred(m_ini)
+        dr = resp_ini - dobs
+
+        def misfit(m, dobs):
+            dpred = self.sim.dpred(m)
+            misfit = np.linalg.norm(dpred - dobs) ** 2
+            dmisfit = 2 * self.sim.Jtvec(
+                m, dr
+            )  # derivative of ||dpred - dobs||^2 gives factor of 2
+            return misfit, dmisfit
+
+        def derChk(m):
+            return misfit(m, dobs)
+
+        passed = tests.check_derivative(
+            derChk, m_ini, num=4, plotIt=False, eps=1e-27, random_seed=42
+        )
+        assert passed
+
+
+class TestEM1D_FD_Jacobian_LineCurrent:
+    # Tests 2nd order convergence of Jvec and Jtvec for piecewise linear loop.
+    # - All rx orientations
+    # - All rx components
+    # - Span many frequencies
+    # - Tests derivatives wrt sigma, mu, thicknesses and h
+    def setup_class(self):
+        x_path = np.array([-2, -2, 2, 2, -2])
+        y_path = np.array([-1, 1, 1, -1, -1])
+        frequencies = np.logspace(0, 4)
+
+        wire_paths = np.c_[x_path, y_path, np.ones(5) * 0.5]
+        source_list = []
+        receiver_list = []
+        receiver_location = np.array([9.28, 0.0, 0.45])
+        orientations = ["x", "y", "z"]
+        components = ["real", "imag", "both"]
+
+        # Define sources and receivers
+        source_list = []
+        for f in frequencies:
+            receiver_list = []
+
+            for rx_orientation in orientations:
+                for comp in components:
+                    receiver_list.append(
+                        fdem.receivers.PointMagneticFieldSecondary(
+                            receiver_location,
+                            orientation=rx_orientation,
+                            component=comp,
+                        )
+                    )
+
+            source_list.append(fdem.sources.LineCurrent(receiver_list, f, wire_paths))
+
+        # Survey
+        survey = fdem.Survey(source_list)
+        self.thicknesses = np.array([20.0, 40.0])
+
+        self.nlayers = len(self.thicknesses) + 1
+        wire_map = maps.Wires(
+            ("sigma", self.nlayers),
+            ("mu", self.nlayers),
+            ("thicknesses", self.nlayers - 1),
+        )
+        self.sigma_map = maps.ExpMap(nP=self.nlayers) * wire_map.sigma
+        self.mu_map = maps.ExpMap(nP=self.nlayers) * wire_map.mu
+        self.thicknesses_map = maps.ExpMap(nP=self.nlayers - 1) * wire_map.thicknesses
+
+        sim = fdem.Simulation1DLayered(
+            survey=survey,
+            sigmaMap=self.sigma_map,
+            muMap=self.mu_map,
+            thicknessesMap=self.thicknesses_map,
+        )
+
+        self.sim = sim
+
+    def test_EM1DFDJvec_Layers(self):
+        # Conductivity
+        sigma_half = 0.01
+        sigma_blk = 0.1
+        sig = np.ones(self.nlayers) * sigma_half
+        sig[1] = sigma_blk
+
+        # Permeability
+        mu_half = mu_0
+        mu_blk = 1.1 * mu_0
+        mu = np.ones(self.nlayers) * mu_half
+        mu[1] = mu_blk
+
+        # General model
+        m_1D = np.r_[np.log(sig), np.log(mu), np.log(self.thicknesses)]
+
+        def fwdfun(m):
+            resp = self.sim.dpred(m)
+            return resp
+            # return Hz
+
+        def jacfun(m, dm):
+            Jvec = self.sim.Jvec(m, dm)
+            return Jvec
+
+        dm = m_1D * 0.5
+
+        def derChk(m):
+            return [fwdfun(m), lambda mx: jacfun(m, mx)]
+
+        passed = tests.check_derivative(
+            derChk, m_1D, num=4, dx=dm, plotIt=False, eps=1e-15, random_seed=1123
+        )
+        assert passed
+
+    def test_EM1DFDJtvec_Layers(self):
+        # Conductivity
+        sigma_half = 0.01
+        sigma_blk = 0.1
+        sig = np.ones(self.nlayers) * sigma_half
+        sig[1] = sigma_blk
+
+        # Permeability
+        mu_half = mu_0
+        mu_blk = 1.1 * mu_0
+        mu = np.ones(self.nlayers) * mu_half
+        mu[1] = mu_blk
+
+        # General model
+        m_true = np.r_[np.log(sig), np.log(mu), np.log(self.thicknesses)]
+
+        dobs = self.sim.dpred(m_true)
+
+        m_ini = np.r_[
+            np.log(np.ones(self.nlayers) * sigma_half),
+            np.log(np.ones(self.nlayers) * mu_half),
+            np.log(self.thicknesses) * 0.9,
+        ]
+        resp_ini = self.sim.dpred(m_ini)
+        dr = resp_ini - dobs
+
+        def misfit(m, dobs):
+            dpred = self.sim.dpred(m)
+            misfit = np.linalg.norm(dpred - dobs) ** 2
+            dmisfit = 2 * self.sim.Jtvec(
+                m, dr
+            )  # derivative of ||dpred - dobs||^2 gives factor of 2
+            return misfit, dmisfit
+
+        def derChk(m):
+            return misfit(m, dobs)
+
+        passed = tests.check_derivative(
+            derChk, m_ini, num=4, plotIt=False, eps=1e-27, random_seed=124
+        )
+        assert passed
diff --git a/tests/em/em1d/test_EM1D_FD_jac_layers3.py b/tests/em/em1d/test_EM1D_FD_jac_layers3.py
new file mode 100644
index 0000000000..bc49ac3432
--- /dev/null
+++ b/tests/em/em1d/test_EM1D_FD_jac_layers3.py
@@ -0,0 +1,461 @@
+from simpeg import maps
+from discretize import tests, TensorMesh
+import simpeg.electromagnetics.frequency_domain as fdem
+import numpy as np
+from scipy.constants import mu_0
+from scipy.sparse import diags
+
+
+class TestEM1D_FD_Jacobian_MagDipole:
+    # Tests 2nd order convergence of Jvec and Jtvec for magnetic dipole sources.
+    # - All src and rx orientations
+    # - All rx components
+    # - Span many frequencies
+    # - Tests derivatives wrt sigma, mu, thicknesses and h
+    def setup_class(self):
+        # Layers and topography
+        nearthick = np.logspace(-1, 1, 5)
+        deepthick = np.logspace(1, 2, 10)
+        thicknesses = np.r_[nearthick, deepthick]
+        topo = np.r_[0.0, 0.0, 100.0]
+
+        # Survey Geometry
+        height = 1e-5
+        src_location = np.array([0.0, 0.0, 100.0 + height])
+        rx_location = np.array([5.0, 5.0, 100.0 + height])
+        frequencies = np.logspace(1, 8, 9)
+        orientations = ["x", "y", "z"]
+        components = ["real", "imag", "both"]
+
+        # Define sources and receivers
+        source_list = []
+        for f in frequencies:
+            for tx_orientation in orientations:
+                receiver_list = []
+
+                for rx_orientation in orientations:
+                    for comp in components:
+                        receiver_list.append(
+                            fdem.receivers.PointMagneticFieldSecondary(
+                                rx_location, orientation=rx_orientation, component=comp
+                            )
+                        )
+
+                source_list.append(
+                    fdem.sources.MagDipole(
+                        receiver_list,
+                        frequency=f,
+                        location=src_location,
+                        orientation=tx_orientation,
+                    )
+                )
+
+        # Survey
+        survey = fdem.Survey(source_list)
+
+        self.topo = topo
+        self.survey = survey
+        self.showIt = False
+        self.height = height
+        self.frequencies = frequencies
+        self.thicknesses = thicknesses
+        self.nlayers = len(thicknesses) + 1
+
+        wire_map = maps.Wires(
+            ("sigma", self.nlayers),
+            ("h", 1),
+            ("thicknesses", self.nlayers - 1),
+        )
+        self.sigma_map = maps.ExpMap(nP=self.nlayers) * wire_map.sigma
+        self.thicknesses_map = maps.ExpMap(nP=self.nlayers - 1) * wire_map.thicknesses
+        nP = len(source_list)
+        surject_mesh = TensorMesh([np.ones(nP)])
+        self.h_map = maps.SurjectFull(surject_mesh) * maps.ExpMap(nP=1) * wire_map.h
+
+        sim = fdem.Simulation1DLayered(
+            survey=self.survey,
+            sigmaMap=self.sigma_map,
+            thicknessesMap=self.thicknesses_map,
+            hMap=self.h_map,
+            topo=self.topo,
+        )
+
+        self.sim = sim
+
+    def test_EM1DFDJvec_Layers(self):
+        # Conductivity
+        sigma_half = 0.01
+        sigma_blk = 0.1
+        sig = np.ones(self.nlayers) * sigma_half
+        sig[3] = sigma_blk
+
+        # General model
+        m_1D = np.r_[np.log(sig), np.log(self.height), np.log(self.thicknesses)]
+
+        def fwdfun(m):
+            resp = self.sim.dpred(m)
+            return resp
+            # return Hz
+
+        def jacfun(m, dm):
+            Jvec = self.sim.Jvec(m, dm)
+            return Jvec
+
+        def derChk(m):
+            return [fwdfun(m), lambda mx: jacfun(m, mx)]
+
+        dm = m_1D * 0.5
+
+        passed = tests.check_derivative(
+            derChk, m_1D, num=4, dx=dm, plotIt=False, eps=1e-15, random_seed=9186724
+        )
+        assert passed
+
+    def test_EM1DFDJtvec_Layers(self):
+        # Conductivity
+        sigma_half = 0.01
+        sigma_blk = 0.1
+        sig = np.ones(self.nlayers) * sigma_half
+        sig[3] = sigma_blk
+
+        # General model
+        m_true = np.r_[np.log(sig), np.log(self.height), np.log(self.thicknesses)]
+
+        dobs = self.sim.dpred(m_true)
+
+        m_ini = np.r_[
+            np.log(np.ones(self.nlayers) * sigma_half),
+            np.log(0.5 * self.height),
+            np.log(self.thicknesses) * 0.9,
+        ]
+        resp_ini = self.sim.dpred(m_ini)
+        dr = resp_ini - dobs
+
+        def misfit(m, dobs):
+            dpred = self.sim.dpred(m)
+            misfit = np.linalg.norm(dpred - dobs) ** 2
+            dmisfit = 2.0 * self.sim.Jtvec(
+                m, dr
+            )  # derivative of ||dpred - dobs||^2 gives factor of 2
+            return misfit, dmisfit
+
+        def derChk(m):
+            return misfit(m, dobs)
+
+        passed = tests.check_derivative(
+            derChk, m_ini, num=4, plotIt=False, eps=1e-27, random_seed=2345
+        )
+        assert passed
+
+    def test_jtjdiag(self):
+        # Conductivity
+        sigma_half = 0.01
+        sigma_blk = 0.1
+        sig = np.ones(self.nlayers) * sigma_half
+        sig[3] = sigma_blk
+
+        # General model
+        model = np.r_[np.log(sig), np.log(self.height), np.log(self.thicknesses)]
+
+        rng = np.random.default_rng(seed=42)
+        weights_matrix = diags(rng.random(size=self.sim.survey.nD))
+        jtj_diag = self.sim.getJtJdiag(model, W=weights_matrix)
+
+        J = self.sim.getJ(model)
+        expected = np.diag(J.T @ weights_matrix.T @ weights_matrix @ J)
+        np.testing.assert_allclose(expected, jtj_diag)
+
+
+class TestEM1D_FD_Jacobian_CircularLoop:
+    # Tests 2nd order convergence of Jvec and Jtvec for horizontal loop sources.
+    # - All rx orientations
+    # - All rx components
+    # - Span many frequencies
+    # - Tests derivatives wrt sigma, mu, thicknesses and h
+    def setup_class(self):
+        nearthick = np.logspace(-1, 1, 5)
+        deepthick = np.logspace(1, 2, 10)
+        thicknesses = np.r_[nearthick, deepthick]
+        topo = np.r_[0.0, 0.0, 100.0]
+        height = 1e-5
+
+        src_location = np.array([0.0, 0.0, 100.0 + height])
+        rx_location = np.array([0.0, 0.0, 100.0 + height])
+        frequencies = np.logspace(1, 8, 9)
+        orientations = ["x", "y", "z"]
+        components = ["real", "imag", "both"]
+        I = 1.0
+        a = 10.0
+
+        # Define sources and receivers
+        source_list = []
+        for f in frequencies:
+            receiver_list = []
+
+            for rx_orientation in orientations:
+                for comp in components:
+                    receiver_list.append(
+                        fdem.receivers.PointMagneticFieldSecondary(
+                            rx_location, orientation=rx_orientation, component=comp
+                        )
+                    )
+
+            source_list.append(
+                fdem.sources.CircularLoop(
+                    receiver_list, f, src_location, radius=a, current=I
+                )
+            )
+
+        # Survey
+        survey = fdem.Survey(source_list)
+
+        self.topo = topo
+        self.survey = survey
+        self.showIt = False
+        self.height = height
+        self.frequencies = frequencies
+        self.thicknesses = thicknesses
+        self.nlayers = len(thicknesses) + 1
+
+        nP = len(source_list)
+
+        wire_map = maps.Wires(
+            ("sigma", self.nlayers),
+            ("mu", self.nlayers),
+            ("thicknesses", self.nlayers - 1),
+            ("h", 1),
+        )
+        self.sigma_map = maps.ExpMap(nP=self.nlayers) * wire_map.sigma
+        self.mu_map = maps.ExpMap(nP=self.nlayers) * wire_map.mu
+        self.thicknesses_map = maps.ExpMap(nP=self.nlayers - 1) * wire_map.thicknesses
+        surject_mesh = TensorMesh([np.ones(nP)])
+        self.h_map = maps.SurjectFull(surject_mesh) * maps.ExpMap(nP=1) * wire_map.h
+
+        sim = fdem.Simulation1DLayered(
+            survey=self.survey,
+            sigmaMap=self.sigma_map,
+            muMap=self.mu_map,
+            thicknessesMap=self.thicknesses_map,
+            hMap=self.h_map,
+            topo=self.topo,
+        )
+
+        self.sim = sim
+
+    def test_EM1DFDJvec_Layers(self):
+        # Conductivity
+        sigma_half = 0.01
+        sigma_blk = 0.1
+        sig = np.ones(self.nlayers) * sigma_half
+        sig[3] = sigma_blk
+
+        # Permeability
+        mu_half = mu_0
+        mu_blk = 2 * mu_0
+        mu = np.ones(self.nlayers) * mu_half
+        mu[3] = mu_blk
+
+        # General model
+        m_1D = np.r_[
+            np.log(sig), np.log(mu), np.log(self.thicknesses), np.log(self.height)
+        ]
+
+        def fwdfun(m):
+            resp = self.sim.dpred(m)
+            return resp
+            # return Hz
+
+        def jacfun(m, dm):
+            Jvec = self.sim.Jvec(m, dm)
+            return Jvec
+
+        def derChk(m):
+            return [fwdfun(m), lambda mx: jacfun(m, mx)]
+
+        dm = m_1D * 0.5
+
+        passed = tests.check_derivative(
+            derChk, m_1D, num=4, dx=dm, plotIt=False, eps=1e-15, random_seed=664
+        )
+        assert passed
+
+    def test_EM1DFDJtvec_Layers(self):
+        # Conductivity
+        sigma_half = 0.01
+        sigma_blk = 0.1
+        sig = np.ones(self.nlayers) * sigma_half
+        sig[3] = sigma_blk
+
+        # Permeability
+        mu_half = mu_0
+        mu_blk = 2 * mu_0
+        mu = np.ones(self.nlayers) * mu_half
+        mu[3] = mu_blk
+
+        # General model
+        m_true = np.r_[
+            np.log(sig), np.log(mu), np.log(self.thicknesses), np.log(self.height)
+        ]
+
+        dobs = self.sim.dpred(m_true)
+
+        m_ini = np.r_[
+            np.log(np.ones(self.nlayers) * sigma_half),
+            np.log(np.ones(self.nlayers) * 1.5 * mu_half),
+            np.log(self.thicknesses) * 0.9,
+            np.log(0.5 * self.height),
+        ]
+        resp_ini = self.sim.dpred(m_ini)
+        dr = resp_ini - dobs
+
+        def misfit(m, dobs):
+            dpred = self.sim.dpred(m)
+            misfit = np.linalg.norm(dpred - dobs) ** 2
+            dmisfit = 2 * self.sim.Jtvec(
+                m, dr
+            )  # derivative of ||dpred - dobs||^2 gives factor of 2
+            return misfit, dmisfit
+
+        def derChk(m):
+            return misfit(m, dobs)
+
+        passed = tests.check_derivative(
+            derChk, m_ini, num=4, plotIt=False, eps=1e-27, random_seed=42
+        )
+        assert passed
+
+
+class TestEM1D_FD_Jacobian_LineCurrent:
+    # Tests 2nd order convergence of Jvec and Jtvec for piecewise linear loop.
+    # - All rx orientations
+    # - All rx components
+    # - Span many frequencies
+    # - Tests derivatives wrt sigma, mu, thicknesses and h
+    def setup_class(self):
+        x_path = np.array([-2, -2, 2, 2, -2])
+        y_path = np.array([-1, 1, 1, -1, -1])
+        frequencies = np.logspace(0, 4)
+
+        wire_paths = np.c_[x_path, y_path, np.ones(5) * 0.5]
+        source_list = []
+        receiver_list = []
+        receiver_location = np.array([9.28, 0.0, 0.45])
+        orientations = ["x", "y", "z"]
+        components = ["real", "imag", "both"]
+
+        # Define sources and receivers
+        source_list = []
+        for f in frequencies:
+            receiver_list = []
+
+            for rx_orientation in orientations:
+                for comp in components:
+                    receiver_list.append(
+                        fdem.receivers.PointMagneticFieldSecondary(
+                            receiver_location,
+                            orientation=rx_orientation,
+                            component=comp,
+                        )
+                    )
+
+            source_list.append(fdem.sources.LineCurrent(receiver_list, f, wire_paths))
+
+        # Survey
+        survey = fdem.Survey(source_list)
+        self.thicknesses = np.array([20.0, 40.0])
+
+        self.nlayers = len(self.thicknesses) + 1
+        wire_map = maps.Wires(
+            ("sigma", self.nlayers),
+            ("mu", self.nlayers),
+            ("thicknesses", self.nlayers - 1),
+        )
+        self.sigma_map = maps.ExpMap(nP=self.nlayers) * wire_map.sigma
+        self.mu_map = maps.ExpMap(nP=self.nlayers) * wire_map.mu
+        self.thicknesses_map = maps.ExpMap(nP=self.nlayers - 1) * wire_map.thicknesses
+
+        sim = fdem.Simulation1DLayered(
+            survey=survey,
+            sigmaMap=self.sigma_map,
+            muMap=self.mu_map,
+            thicknessesMap=self.thicknesses_map,
+        )
+
+        self.sim = sim
+
+    def test_EM1DFDJvec_Layers(self):
+        # Conductivity
+        sigma_half = 0.01
+        sigma_blk = 0.1
+        sig = np.ones(self.nlayers) * sigma_half
+        sig[1] = sigma_blk
+
+        # Permeability
+        mu_half = mu_0
+        mu_blk = 1.1 * mu_0
+        mu = np.ones(self.nlayers) * mu_half
+        mu[1] = mu_blk
+
+        # General model
+        m_1D = np.r_[np.log(sig), np.log(mu), np.log(self.thicknesses)]
+
+        def fwdfun(m):
+            resp = self.sim.dpred(m)
+            return resp
+            # return Hz
+
+        def jacfun(m, dm):
+            Jvec = self.sim.Jvec(m, dm)
+            return Jvec
+
+        dm = m_1D * 0.5
+
+        def derChk(m):
+            return [fwdfun(m), lambda mx: jacfun(m, mx)]
+
+        passed = tests.check_derivative(
+            derChk, m_1D, num=4, dx=dm, plotIt=False, eps=1e-15, random_seed=1123
+        )
+        assert passed
+
+    def test_EM1DFDJtvec_Layers(self):
+        # Conductivity
+        sigma_half = 0.01
+        sigma_blk = 0.1
+        sig = np.ones(self.nlayers) * sigma_half
+        sig[1] = sigma_blk
+
+        # Permeability
+        mu_half = mu_0
+        mu_blk = 1.1 * mu_0
+        mu = np.ones(self.nlayers) * mu_half
+        mu[1] = mu_blk
+
+        # General model
+        m_true = np.r_[np.log(sig), np.log(mu), np.log(self.thicknesses)]
+
+        dobs = self.sim.dpred(m_true)
+
+        m_ini = np.r_[
+            np.log(np.ones(self.nlayers) * sigma_half),
+            np.log(np.ones(self.nlayers) * mu_half),
+            np.log(self.thicknesses) * 0.9,
+        ]
+        resp_ini = self.sim.dpred(m_ini)
+        dr = resp_ini - dobs
+
+        def misfit(m, dobs):
+            dpred = self.sim.dpred(m)
+            misfit = np.linalg.norm(dpred - dobs) ** 2
+            dmisfit = 2 * self.sim.Jtvec(
+                m, dr
+            )  # derivative of ||dpred - dobs||^2 gives factor of 2
+            return misfit, dmisfit
+
+        def derChk(m):
+            return misfit(m, dobs)
+
+        passed = tests.check_derivative(
+            derChk, m_ini, num=4, plotIt=False, eps=1e-27, random_seed=124
+        )
+        assert passed
diff --git a/tests/em/em1d/test_EM1D_TD_general_jac_layers.py b/tests/em/em1d/test_EM1D_TD_general_jac_layers.py
index 0580d5407e..cc9cfddda9 100644
--- a/tests/em/em1d/test_EM1D_TD_general_jac_layers.py
+++ b/tests/em/em1d/test_EM1D_TD_general_jac_layers.py
@@ -342,7 +342,8 @@ def test_rx_loc_shapes(rx_class, n_locs1, n_locs2, orientation, waveform, compar
         d = sim.dpred(None)
     else:
         sim.sigmaMap = maps.IdentityMap(nP=1)
-        d = sim.getJ(np.ones(1))["ds"][:, 0]
+        J = sim.getJ(np.ones(1))
+        d = J[:, 0]
 
     # assert the shape is correct
     assert d.shape == (n_d,)
diff --git a/tests/em/em1d/test_EM1D_TD_getJ.py b/tests/em/em1d/test_EM1D_TD_getJ.py
new file mode 100644
index 0000000000..42273fbedf
--- /dev/null
+++ b/tests/em/em1d/test_EM1D_TD_getJ.py
@@ -0,0 +1,161 @@
+"""
+Test the getJ method of FDEM 1D simulation.
+"""
+
+import pytest
+import numpy as np
+import simpeg.electromagnetics.time_domain as tdem
+from simpeg import maps
+from scipy.sparse import diags
+
+
+def create_simulation_and_conductivities(identity_mapping: bool):
+    # Create Survey
+    # -------------
+    # Source properties
+    source_location = np.array([0.0, 0.0, 20.0])
+    source_orientation = "z"  # "x", "y" or "z"
+    source_current = 1.0  # maximum on-time current
+    source_radius = 6.0  # source loop radius
+
+    # Receiver properties
+    receiver_location = np.array([0.0, 0.0, 20.0])
+    receiver_orientation = "z"  # "x", "y" or "z"
+    times = np.logspace(-5, -2, 31)  # time channels (s)
+
+    # Define receiver list. In our case, we have only a single receiver for each source.
+    # When simulating the response for multiple component and/or field orientations,
+    # the list consists of multiple receiver objects.
+    receiver_list = []
+    receiver_list.append(
+        tdem.receivers.PointMagneticFluxDensity(
+            receiver_location, times, orientation=receiver_orientation
+        )
+    )
+
+    # Define the source waveform. Here we define a unit step-off. The definition
+    # of other waveform types is covered in a separate tutorial.
+    waveform = tdem.sources.StepOffWaveform()
+
+    # Define source list. In our case, we have only a single source.
+    source_list = [
+        tdem.sources.CircularLoop(
+            receiver_list=receiver_list,
+            location=source_location,
+            orientation=source_orientation,
+            waveform=waveform,
+            current=source_current,
+            radius=source_radius,
+        )
+    ]
+
+    # Define the survey
+    survey = tdem.Survey(source_list)
+
+    # Defining a 1D Layered Earth Model
+    # ---------------------------------
+    # Physical properties
+    background_conductivity = 1e-1
+    layer_conductivity = 1e0
+
+    # Layer thicknesses
+    thicknesses = np.array([40.0, 40.0])
+    n_layer = len(thicknesses) + 1
+
+    # Conductivities
+    conductivities = background_conductivity * np.ones(n_layer)
+    conductivities[1] = layer_conductivity
+
+    # Define a mapping
+    model_mapping = (
+        maps.IdentityMap(nP=n_layer) if identity_mapping else maps.ExpMap(nP=n_layer)
+    )
+
+    # Define the Forward Simulation, Predict Data and Plot
+    # ----------------------------------------------------
+    simulation = tdem.Simulation1DLayered(
+        survey=survey,
+        thicknesses=thicknesses,
+        sigmaMap=model_mapping,
+    )
+
+    return simulation, conductivities
+
+
+def test_getJ():
+    """
+    Test if getJ returns different J matrices after passing different maps.
+    """
+    dpreds, jacobians = [], []
+
+    # Compute dpred and J using an identity map and an exp map
+    for identity_mapping in (True, False):
+        simulation, conductivities = create_simulation_and_conductivities(
+            identity_mapping
+        )
+        model = conductivities if identity_mapping else np.log(conductivities)
+        dpreds.append(simulation.dpred(model))
+        jac = simulation.getJ(model)
+        jacobians.append(jac)
+
+    # The two dpreds should be equal
+    assert np.allclose(*dpreds)
+
+    # The two J matrices should not be equal
+    assert not np.allclose(*jacobians, atol=0.0)
+
+
+@pytest.mark.parametrize("mapping", ["identity", "expmap"])
+def test_JtJdiag(mapping):
+    """
+    Test the getJtJdiag method of the simulation.
+    """
+    identity_mapping = mapping == "identity"
+    simulation, conductivities = create_simulation_and_conductivities(identity_mapping)
+
+    model = conductivities if identity_mapping else np.log(conductivities)
+    rng = np.random.default_rng(seed=42)
+    weights_matrix = diags(rng.random(size=simulation.survey.nD))
+    jtj_diag = simulation.getJtJdiag(model, W=weights_matrix)
+
+    J = simulation.getJ(model)
+    expected = np.diag(J.T @ weights_matrix.T @ weights_matrix @ J)
+    np.testing.assert_allclose(expected, jtj_diag)
+
+
+@pytest.mark.parametrize("mapping", ["identity", "expmap"])
+def test_Jvec(mapping):
+    """
+    Test the Jvec method of the simulation.
+    """
+    identity_mapping = mapping == "identity"
+    simulation, conductivities = create_simulation_and_conductivities(identity_mapping)
+
+    model = conductivities if identity_mapping else np.log(conductivities)
+    rng = np.random.default_rng(seed=42)
+    vector = rng.random(size=model.size)
+    jvec = simulation.Jvec(model, vector)
+
+    J = simulation.getJ(model)
+    expected = J @ vector
+
+    np.testing.assert_allclose(expected, jvec)
+
+
+@pytest.mark.parametrize("mapping", ["identity", "expmap"])
+def test_Jtvec(mapping):
+    """
+    Test the Jtvec method of the simulation.
+    """
+    identity_mapping = mapping == "identity"
+    simulation, conductivities = create_simulation_and_conductivities(identity_mapping)
+
+    model = conductivities if identity_mapping else np.log(conductivities)
+    rng = np.random.default_rng(seed=42)
+    vector = rng.random(size=simulation.survey.nD)
+    jtvec = simulation.Jtvec(model, vector)
+
+    J = simulation.getJ(model)
+    expected = J.T @ vector
+
+    np.testing.assert_allclose(expected, jtvec)


Reviews:
user: login: santisoler
body: Awesome, thanks @ghwilliams for starting this!

I just left a minor comment regarding the test class you wrote.
state: COMMENTED
author_association: MEMBER
submitted_at: 2025-03-31T19:01:46Z
user: login: jcapriot
state: CHANGES_REQUESTED
author_association: MEMBER
submitted_at: 2025-04-03T13:46:24Z
user: login: ghwilliams
state: COMMENTED
author_association: CONTRIBUTOR
submitted_at: 2025-04-03T13:48:50Z
user: login: santisoler
body: Thanks @ghwilliams! I'm glad to see that the ordering of the J rows is easily solved. I left a comment below regarding that implementation. Let me know what do you think.

Thanks also for adding the tests in `tests/em/em1d/test_EM1D_FD_jac_layers2.py`. If you have the time, it would be nice to switch those tests from using `unittest` to using `pytest`. It's not a huge blocker, but such change would align with the current way of writing tests in the project. I can totally do it afterwards if you don't have the time or interest, no worries.
state: COMMENTED
author_association: MEMBER
submitted_at: 2025-04-08T22:08:39Z
user: login: ghwilliams
state: COMMENTED
author_association: CONTRIBUTOR
submitted_at: 2025-04-09T14:34:34Z
user: login: santisoler
state: COMMENTED
author_association: MEMBER
submitted_at: 2025-04-09T16:43:34Z

Review comments:
diff_hunk: @@ -0,0 +1,106 @@
+import numpy as np
+import simpeg.electromagnetics.time_domain as tdem
+from simpeg import maps
+import unittest
+
+
+def create_simulation_and_conductivities(identity_mapping: bool):
+    # Create Survey
+    # -------------
+    # Source properties
+    source_location = np.array([0.0, 0.0, 20.0])
+    source_orientation = "z"  # "x", "y" or "z"
+    source_current = 1.0  # maximum on-time current
+    source_radius = 6.0  # source loop radius
+
+    # Receiver properties
+    receiver_location = np.array([0.0, 0.0, 20.0])
+    receiver_orientation = "z"  # "x", "y" or "z"
+    times = np.logspace(-5, -2, 31)  # time channels (s)
+
+    # Define receiver list. In our case, we have only a single receiver for each source.
+    # When simulating the response for multiple component and/or field orientations,
+    # the list consists of multiple receiver objects.
+    receiver_list = []
+    receiver_list.append(
+        tdem.receivers.PointMagneticFluxDensity(
+            receiver_location, times, orientation=receiver_orientation
+        )
+    )
+
+    # Define the source waveform. Here we define a unit step-off. The definition of
+    # other waveform types is covered in a separate tutorial.
+    waveform = tdem.sources.StepOffWaveform()
+
+    # Define source list. In our case, we have only a single source.
+    source_list = [
+        tdem.sources.CircularLoop(
+            receiver_list=receiver_list,
+            location=source_location,
+            orientation=source_orientation,
+            waveform=waveform,
+            current=source_current,
+            radius=source_radius,
+        )
+    ]
+
+    # Define the survey
+    survey = tdem.Survey(source_list)
+
+    # Defining a 1D Layered Earth Model
+    # ---------------------------------
+    # Physical properties
+    background_conductivity = 1e-1
+    layer_conductivity = 1e0
+
+    # Layer thicknesses
+    thicknesses = np.array([40.0, 40.0])
+    n_layer = len(thicknesses) + 1
+
+    # Conductivities
+    conductivities = background_conductivity * np.ones(n_layer)
+    conductivities[1] = layer_conductivity
+
+    # Define a mapping
+    model_mapping = (
+        maps.IdentityMap(nP=n_layer) if identity_mapping else maps.ExpMap(nP=n_layer)
+    )
+
+    # Define the Forward Simulation, Predict Data and Plot
+    # ----------------------------------------------------
+    simulation = tdem.Simulation1DLayered(
+        survey=survey,
+        thicknesses=thicknesses,
+        sigmaMap=model_mapping,
+    )
+
+    return simulation, conductivities
+
+
+class TestSimulation1DLayeredGetJ(unittest.TestCase):
path: tests/em/tdem/test_TDEM_Simulation1DLayered_getJ.py
user: login: santisoler
body: We are moving away from using `unittest` and using `pytest` instead. `pytest` offers some features that make writing tests simpler and easier to extend. For example, there's no need to inherit from any class when testing with `pytest`:
```suggestion
class TestSimulation1DLayeredGetJ:
```
created_at: 2025-03-31T18:55:52Z
author_association: MEMBER
diff_hunk: @@ -335,29 +335,11 @@ def compute_complex_mu(self, frequencies):
 
     def Jvec(self, m, v, f=None):
         Js = self.getJ(m, f=f)
path: simpeg/electromagnetics/base_1d.py
user: login: jcapriot
body: The only change to this function should be to change `getJ` to `_getJ`
created_at: 2025-04-03T13:41:27Z
author_association: MEMBER
diff_hunk: @@ -335,29 +335,11 @@ def compute_complex_mu(self, frequencies):
 
     def Jvec(self, m, v, f=None):
         Js = self.getJ(m, f=f)
-        out = 0.0
-        if self.hMap is not None:
-            out = out + Js["dh"] @ (self.hDeriv @ v)
-        if self.sigmaMap is not None:
-            out = out + Js["ds"] @ (self.sigmaDeriv @ v)
-        if self.muMap is not None:
-            out = out + Js["dmu"] @ (self.muDeriv @ v)
-        if self.thicknessesMap is not None:
-            out = out + Js["dthick"] @ (self.thicknessesDeriv @ v)
-        return out
+        return Js @ v
 
     def Jtvec(self, m, v, f=None):
         Js = self.getJ(m, f=f)
path: simpeg/electromagnetics/base_1d.py
user: login: jcapriot
body: Only change to this function should be `getJ` to `_getJ`
created_at: 2025-04-03T13:42:19Z
author_association: MEMBER
diff_hunk: @@ -244,6 +245,28 @@ def getJ(self, m, f=None):
                     self._J["dthick"] = self._project_to_data(v_dthick)
         return self._J
 
+    def getJ(self, m, f=None):
path: simpeg/electromagnetics/frequency_domain/simulation_1d.py
user: login: jcapriot
body: All you have to do is sum them…

`Js[‘ds’] @ sigmaDeriv + Js[`dmu`] @ muDeriv +…`

Of course you’ll need to conditionally sum them if they’re there.
created_at: 2025-04-03T13:46:12Z
author_association: MEMBER
diff_hunk: @@ -244,6 +245,28 @@ def getJ(self, m, f=None):
                     self._J["dthick"] = self._project_to_data(v_dthick)
         return self._J
 
+    def getJ(self, m, f=None):
path: simpeg/electromagnetics/frequency_domain/simulation_1d.py
user: login: ghwilliams
body: Ok. I'll do it right now and see if it works. Thanks @jcapriot .
created_at: 2025-04-03T13:48:50Z
author_association: CONTRIBUTOR
diff_hunk: @@ -244,6 +244,32 @@ def getJ(self, m, f=None):
                     self._J["dthick"] = self._project_to_data(v_dthick)
         return self._J
 
+    def getJ(self, m, f=None):
+        Js = self._getJ(m, f=f)
+
+        J = None
+        if self.hMap is not None:
+            J = Js["dh"] @ self.hDeriv
+
+        if self.sigmaMap is not None:
+            J = (
+                Js["ds"] @ self.sigmaDeriv
+                if J is None
+                else J + Js["ds"] @ self.sigmaDeriv
+            )
+
+        if self.muMap is not None:
+            J = Js["dmu"] @ self.muDeriv if J is None else J + Js["dmu"] @ self.muDeriv
+
+        if self.thicknessesMap is not None:
+            J = (
+                Js["dthick"] @ self.thicknessesDeriv
+                if J is None
+                else J + Js["dthick"] @ self.thicknessesDeriv
+            )
+
+        return J
+
path: simpeg/electromagnetics/frequency_domain/simulation_1d.py
user: login: santisoler
body: Nice! I'm glad to see that this works as easily as this! I'm just leaving a minor refactor suggestion to simplify the code. I think this is easier to read, but let me know if you disagree.

This code assumes that the `_getJ` method only adds pieces of the jacobian for the parameters that have a mapping assigned. That's why I'm not checking if mappings are None here.

```python
    def getJ(self, m, f=None):
        # Get Js matrices (we need to set the model before accessing derivatives)
        Js = self._getJ(m, f=f)
        # Map parameters with their corresponding derivatives
        param_and_derivs = {
            "dh": self.hDeriv,
            "ds": self.sigmaDeriv,
            "dmu": self.muDeriv,
            "dthick": self.thicknessesDeriv,
        }
        # Compute J matrix
        J = sum(Js[param] @ param_and_derivs[param] for param in Js)
        return J
```
created_at: 2025-04-08T22:02:23Z
author_association: MEMBER
diff_hunk: @@ -244,6 +244,32 @@ def getJ(self, m, f=None):
                     self._J["dthick"] = self._project_to_data(v_dthick)
         return self._J
 
+    def getJ(self, m, f=None):
+        Js = self._getJ(m, f=f)
+
+        J = None
+        if self.hMap is not None:
+            J = Js["dh"] @ self.hDeriv
+
+        if self.sigmaMap is not None:
+            J = (
+                Js["ds"] @ self.sigmaDeriv
+                if J is None
+                else J + Js["ds"] @ self.sigmaDeriv
+            )
+
+        if self.muMap is not None:
+            J = Js["dmu"] @ self.muDeriv if J is None else J + Js["dmu"] @ self.muDeriv
+
+        if self.thicknessesMap is not None:
+            J = (
+                Js["dthick"] @ self.thicknessesDeriv
+                if J is None
+                else J + Js["dthick"] @ self.thicknessesDeriv
+            )
+
+        return J
+
path: simpeg/electromagnetics/frequency_domain/simulation_1d.py
user: login: ghwilliams
body: @santisoler Ok. I'll try those suggestions. Regarding rewriting the tests using pytest, I could open another PR for that because all the tests under tests\em\em1d use unittest. Let's discuss this further during today's dev meeting.
created_at: 2025-04-09T14:34:33Z
author_association: CONTRIBUTOR
diff_hunk: @@ -244,6 +244,32 @@ def getJ(self, m, f=None):
                     self._J["dthick"] = self._project_to_data(v_dthick)
         return self._J
 
+    def getJ(self, m, f=None):
+        Js = self._getJ(m, f=f)
+
+        J = None
+        if self.hMap is not None:
+            J = Js["dh"] @ self.hDeriv
+
+        if self.sigmaMap is not None:
+            J = (
+                Js["ds"] @ self.sigmaDeriv
+                if J is None
+                else J + Js["ds"] @ self.sigmaDeriv
+            )
+
+        if self.muMap is not None:
+            J = Js["dmu"] @ self.muDeriv if J is None else J + Js["dmu"] @ self.muDeriv
+
+        if self.thicknessesMap is not None:
+            J = (
+                Js["dthick"] @ self.thicknessesDeriv
+                if J is None
+                else J + Js["dthick"] @ self.thicknessesDeriv
+            )
+
+        return J
+
path: simpeg/electromagnetics/frequency_domain/simulation_1d.py
user: login: santisoler
body: Sure! I didn't mean to change _all_ tests in `em1d` to use pytest in this PR. I meant to write _new_ tests using pytest and avoid using unittest in those.
created_at: 2025-04-09T16:43:33Z
author_association: MEMBER

PR comments:
user: login: codecov[bot]
created_at: 2025-03-31T12:48:30Z
body: ## [Codecov](https://app.codecov.io/gh/simpeg/simpeg/pull/1638?dropdown=coverage&src=pr&el=h1&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=simpeg) Report
All modified and coverable lines are covered by tests :white_check_mark:
> Project coverage is 86.20%. Comparing base [(`24bf81e`)](https://app.codecov.io/gh/simpeg/simpeg/commit/24bf81e9343b63f7cb4ca48a34b5dda9c7709219?dropdown=coverage&el=desc&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=simpeg) to head [(`3b6b98f`)](https://app.codecov.io/gh/simpeg/simpeg/commit/3b6b98feddf990d2b9458012a362a27f03b237e6?dropdown=coverage&el=desc&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=simpeg).
> Report is 1 commits behind head on main.

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


```diff
@@            Coverage Diff             @@
##             main    #1638      +/-   ##
==========================================
+ Coverage   86.03%   86.20%   +0.16%     
==========================================
  Files         412      416       +4     
  Lines       54356    54970     +614     
  Branches     5106     5129      +23     
==========================================
+ Hits        46767    47388     +621     
+ Misses       6169     6164       -5     
+ Partials     1420     1418       -2     
```

</details>

[:umbrella: View full report in Codecov by Sentry](https://app.codecov.io/gh/simpeg/simpeg/pull/1638?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>
author_association: NONE
user: login: jcapriot
created_at: 2025-04-02T18:44:55Z
body: Quick note on how this should be done:

1: change current `getJ` to `_getJ`
2: change internal calls inside `Jvec`, `JTvec` and  `getJtJdiag` from `getJ` to `_getJ`
3: add new `getJ` method that multiplies each object in the `_getJ` dictionary by the respective mapping derivative, then sums them all together.
author_association: MEMBER
user: login: ghwilliams
created_at: 2025-04-03T10:21:30Z
body: > Quick note on how this should be done:
> 
> 1: change current `getJ` to `_getJ` 2: change internal calls inside `Jvec`, `JTvec` and `getJtJdiag` from `getJ` to `_getJ` 3: add new `getJ` method that multiplies each object in the `_getJ` dictionary by the respective mapping derivative, then sums them all together.

@jcapriot I've already made many of these modifications. Yesterday, during the dev meeting, we discussed (@santisoler and I) the implementation of the new getJ. We found out that there is a problem related to the order in which the full sensitivity matrix must be assembled from the individual blocks corresponding to the sensitivity related to all the different parameters composing the full model parameter vector, whose order is defined by the user when creating the simulation object. It is unclear how we can obtain this information from inside getJ. We know the user may have defined a composed mapping, for example, by multiplying an ExpMap by a Wire map. I'm looking for ways to get the specific order in which the user composed the full model parameter vector from this map (inside getJ). Please look at the code and let me know if you think getJ could be implemented differently, possibly avoiding this order issue altogether.
author_association: CONTRIBUTOR
user: login: jcapriot
created_at: 2025-04-03T13:47:52Z
body: > > Quick note on how this should be done:
> > 1: change current `getJ` to `_getJ` 2: change internal calls inside `Jvec`, `JTvec` and `getJtJdiag` from `getJ` to `_getJ` 3: add new `getJ` method that multiplies each object in the `_getJ` dictionary by the respective mapping derivative, then sums them all together.
> 
> @jcapriot I've already made many of these modifications. Yesterday, during the dev meeting, we discussed (@santisoler and I) the implementation of the new getJ. We found out that there is a problem related to the order in which the full sensitivity matrix must be assembled from the individual blocks corresponding to the sensitivity related to all the different parameters composing the full model parameter vector, whose order is defined by the user when creating the simulation object. It is unclear how we can obtain this information from inside getJ. We know the user may have defined a composed mapping, for example, by multiplying an ExpMap by a Wire map. I'm looking for ways to get the specific order in which the user composed the full model parameter vector from this map (inside getJ). Please look at the code and let me know if you think getJ could be implemented differently, possibly avoiding this order issue altogether.

You haven’t made those changes, and I’ve marked them directly on the review for you now. You’re overthinking it, if you implement it as I have requested, it is minimal changes overall.
author_association: MEMBER
user: login: ghwilliams
created_at: 2025-04-03T14:30:56Z
body: @jcapriot @santisoler I updated getJ following Joseph's suggestions. I ran the tests in tests/em/em1d, and they all passed. Looks like it will work.
author_association: CONTRIBUTOR
user: login: ghwilliams
created_at: 2025-04-08T14:13:27Z
body: @jcapriot @santisoler There was a concern about the correct assembly of J when multiple model parameters were defined. I added another test where I changed the order in which the model vector is specified. The test passed.
author_association: CONTRIBUTOR
user: login: jcapriot
created_at: 2025-04-08T20:04:59Z
body: Yeah, the getJ works as I expected. Just FYI the way mappings work is that they map from the entire model vector to a specific physical property. Their derivatives naturally include the necessary information about which model parameters the physical properties are derivatives of.

Also please implement my second requested point:

> 2: change internal calls inside `Jvec`, `Jtvec` and `getJtJdiag` from `getJ` to `_getJ`

and revert all your other changes to these functions.

I don’t want them to unnecessarily create matrices full of zeros.
author_association: MEMBER
user: login: santisoler
created_at: 2025-04-23T22:47:51Z
body: **TODO Before merging**:

- [x] Include a docstring to the `getJ` method
author_association: MEMBER
user: login: santisoler
created_at: 2025-04-24T01:28:07Z
body: @ghwilliams, I pushed some docstrings to the `getJ` and `_getJ` methods of the TDEM and FDEM simulations. I also changed the title of the PR and wrote a small summary I'll use when merging it.

The CIs pass, so I think I'll merge this tomorrow morning (local time). If you have the chance, could you take another look to my changes?

Thanks again for the contributions!
author_association: MEMBER
user: login: ghwilliams
created_at: 2025-04-24T12:27:28Z
body: > @ghwilliams, I pushed some docstrings to the `getJ` and `_getJ` methods of the TDEM and FDEM simulations. I also changed the title of the PR and wrote a small summary I'll use when merging it.
> 
> The CIs pass, so I think I'll merge this tomorrow morning (local time). If you have the chance, could you take another look to my changes?
> 
> Thanks again for the contributions!

Everything is fine. The documentation for both _getJ and getJ is ok.
author_association: CONTRIBUTOR
user: login: jcapriot
created_at: 2025-04-24T17:21:36Z
body: Should add a test for the JtJdiag method
author_association: MEMBER