diff --git a/demos/with_automatic_differentiation/run_forward.py b/demos/with_automatic_differentiation/run_forward.py
index 0867b8a6..ad7ac5e1 100644
--- a/demos/with_automatic_differentiation/run_forward.py
+++ b/demos/with_automatic_differentiation/run_forward.py
@@ -1,5 +1,6 @@
 from firedrake import *
 import spyro
+from spyro.io import front_ensemble_run_forward
 import matplotlib.pyplot as plt
 import numpy as np
 
@@ -16,7 +17,10 @@
 }
 
 model["parallelism"] = {
-    "type": "spatial",  # options: automatic (same number of cores for evey processor) or spatial
+    # options: shots_parallelism (same number of cores for every processor)
+    # or automatic (same number of cores for every processor) or
+    # or spatial.
+    "type": "none",
 }
 
 # Define the domain size without the ABL.
@@ -59,13 +63,16 @@
     "fspool": 1,  # how frequently to save solution to RAM
 }
 
-comm = spyro.utils.mpi_init(model)
-mesh = UnitSquareMesh(100, 100)  # to test FWI, mesh aligned with interface
+comm, spatial_comm = spyro.utils.mpi_init(model)
+if model["parallelism"]["type"] == "shots_parallelism":
+    # Only shots parallelism.
+    mesh = UnitSquareMesh(50, 50, comm=spatial_comm)
+else:
+    mesh = UnitSquareMesh(50, 50)
 
 element = spyro.domains.space.FE_method(
     mesh, model["opts"]["method"], model["opts"]["degree"]
 )
-
 V = FunctionSpace(mesh, element)
 
 
@@ -89,19 +96,48 @@ def make_vp_circle(vp_guess=False, plot_vp=False):
     tf=model["timeaxis"]["tf"],
     freq=model["acquisition"]["frequency"],
 )
-
+# True acoustic velocity model
 vp_exact = make_vp_circle(plot_vp=True)
-if comm.comm.rank == 0:
-    for sn in range(len(model["acquisition"]["source_pos"])):
-        receiver_data = spyro.solvers.forward_AD(model, mesh, comm, vp_exact,
-                                                 wavelet, debug=True,
-                                                 source_number=sn)
 
-        # --- Plot the receiver data --- #
-        data = []
-        for _, rec in enumerate(receiver_data):
-            data.append(rec.dat.data_ro[:])
 
-        spyro.plots.plot_shots(model, comm, data, vmax=1e-08, vmin=-1e-08)
+def run_forward(source_number):
+    """Execute a acoustic wave equation.
+
+    Parameters
+    ----------
+    source_number: `int`, optional
+        The source number defined by the user.
+
+    Notes
+    -----
+    The forward solver (`forward_AD`) is implemented in spyro using firedrake's
+    functions that can be annotated by the algorithimic differentiation (AD).
+    This is because spyro is capable of executing Full Waveform Inversion (FWI),
+    which needs the computation of the gradient of the objective function with
+    respect to the velocity model through (AD).
+    """
+    receiver_data = spyro.solvers.forward_AD(model, mesh, comm, vp_exact,
+                                             wavelet, debug=True,
+                                             source_number=source_number)
+    # --- Plot the receiver data --- #
+    data = []
+    for _, rec in enumerate(receiver_data):
+        data.append(rec.dat.data_ro[:])
+
+    spyro.plots.plot_shots(model, comm, data, vmax=1e-08, vmin=-1e-08)
+
+
+# Processor number.
+rank = comm.ensemble_comm.rank
+# Number of processors used in the simulation.
+size = comm.ensemble_comm.size
+if size == 1:
+    for sn in range(len(model["acquisition"]["source_pos"])):
+        run_forward(sn)
+elif size == len(model["acquisition"]["source_pos"]):
+    # Only run the forward simulation for the source number that matches the
+    # processor number. 
+    run_forward(rank)
 else:
-    raise NotImplementedError("Only implemented for 1 processor")
+    raise NotImplementedError("`size` must be 1 or equal to `num_sources`."
+                              "Different values are not supported yet.")
diff --git a/demos/without_automatic_differentiation/run_forward.py b/demos/without_automatic_differentiation/run_forward.py
index 093c61c2..06713934 100644
--- a/demos/without_automatic_differentiation/run_forward.py
+++ b/demos/without_automatic_differentiation/run_forward.py
@@ -47,7 +47,7 @@
     "nspool": 100,  # how frequently to output solution to pvds
     "fspool": 99999,  # how frequently to save solution to RAM
 }
-comm = spyro.utils.mpi_init(model)
+comm, _ = spyro.utils.mpi_init(model)
 mesh, V = spyro.io.read_mesh(model, comm)
 vp = spyro.io.interpolate(model, mesh, V, guess=False)
 sources = spyro.Sources(model, mesh, V, comm)
diff --git a/demos/without_automatic_differentiation/run_fwi.py b/demos/without_automatic_differentiation/run_fwi.py
index a52bf09c..e0c800aa 100644
--- a/demos/without_automatic_differentiation/run_fwi.py
+++ b/demos/without_automatic_differentiation/run_fwi.py
@@ -61,7 +61,7 @@
     "nspool": 1000,  # how frequently to output solution to pvds
     "fspool": 10,  # how frequently to save solution to RAM
 }
-comm = spyro.utils.mpi_init(model)
+comm, _ = spyro.utils.mpi_init(model)
 # if comm.comm.rank == 0 and comm.ensemble_comm.rank == 0:
 #    fil = open("FUNCTIONAL_FWI_P5.txt", "w")
 mesh, V = spyro.io.read_mesh(model, comm)
diff --git a/spyro/io/__init__.py b/spyro/io/__init__.py
index add90fb3..0b1b59d4 100644
--- a/spyro/io/__init__.py
+++ b/spyro/io/__init__.py
@@ -9,7 +9,6 @@
 from .io import (
     interpolate,
     ensemble_forward,
-    ensemble_forward_ad,
     ensemble_forward_elastic_waves,
     ensemble_gradient,
     ensemble_gradient_elastic_waves,
@@ -26,7 +25,6 @@
     "read_mesh",
     "interpolate",
     "ensemble_forward",
-    "ensemble_forward_ad",
     "ensemble_forward_elastic_waves",
     "ensemble_gradient",
     "ensemble_gradient_elastic_waves",
diff --git a/spyro/io/io.py b/spyro/io/io.py
index 166e6f17..c8faffc4 100644
--- a/spyro/io/io.py
+++ b/spyro/io/io.py
@@ -115,25 +115,6 @@ def wrapper(*args, **kwargs):
     return wrapper
 
 
-def ensemble_forward_ad(func):
-    """Decorator for forward_ad to distribute shots for ensemble parallelism"""
-
-    def wrapper(*args, **kwargs):
-        acq = args[0].get("acquisition")
-        num = len(acq["source_pos"])
-        fwi = kwargs.get("fwi")
-        _comm = args[2]
-        for snum in range(num):
-            if is_owner(_comm, snum):
-                if fwi:
-                    u_r, J = func(*args, **dict(kwargs, source_num=snum))
-                    return u_r, J
-                else:
-                    u_r = func(*args, **dict(kwargs, source_num=snum))
-
-    return wrapper
-
-
 def ensemble_forward_elastic_waves(func):
     """Decorator for forward elastic waves to distribute shots for ensemble parallelism"""
 
diff --git a/spyro/solvers/forward_AD.py b/spyro/solvers/forward_AD.py
index 594f2428..cbe250b5 100644
--- a/spyro/solvers/forward_AD.py
+++ b/spyro/solvers/forward_AD.py
@@ -7,7 +7,6 @@
 set_log_level(ERROR)
 
 
-# @ensemble_forward
 def forward(model, mesh, comm, c, wavelet, source_number=0, fwi=False, **kwargs):
     """Secord-order in time fully-explicit scheme.
 
diff --git a/spyro/utils/utils.py b/spyro/utils/utils.py
index 2fcc9e5b..7620d21c 100644
--- a/spyro/utils/utils.py
+++ b/spyro/utils/utils.py
@@ -89,24 +89,37 @@ def mysize(COMM=COMM_SELF):
     return COMM.Get_size()
 
 
-def mpi_init(model):
+def mpi_init(model, spatial_core_parallelism=None):
     """Initialize computing environment"""
     # rank = myrank()
     # size = mysize()
-    available_cores = COMM_WORLD.size
-    if model["parallelism"]["type"] == "automatic":
-        num_cores_per_shot = available_cores / len(model["acquisition"]["source_pos"])
-        if available_cores % len(model["acquisition"]["source_pos"]) != 0:
+    if (
+        model["parallelism"]["type"] == "shots_parallelism"
+        or model["parallelism"]["type"] == "automatic"
+    ):
+        
+        num_cores_per_shot = COMM_WORLD.size / len(model["acquisition"]["source_pos"])
+        if COMM_WORLD.size % len(model["acquisition"]["source_pos"]) != 0:
             raise ValueError(
                 "Available cores cannot be divided between sources equally."
             )
     elif model["parallelism"]["type"] == "spatial":
-        num_cores_per_shot = available_cores
+        # Parallellism is only over spatial domain. No shots parallelism.
+        num_cores_per_shot = COMM_WORLD.size
     elif model["parallelism"]["type"] == "custom":
         raise ValueError("Custom parallelism not yet implemented")
+    elif model["parallelism"]["type"] == "none":
+        num_cores_per_shot = 1
 
+    if model["parallelism"]["type"] == "shots_parallelism":
+        # Parrallellism is over shots. No spatial parallelism.
+        spatial_comm = COMM_WORLD.Split(
+            COMM_WORLD.rank % len(model["acquisition"]["source_pos"])
+            )
+    else:
+        spatial_comm = None
     comm_ens = Ensemble(COMM_WORLD, num_cores_per_shot)
-    return comm_ens
+    return comm_ens, spatial_comm
 
 
 def communicate(array, my_ensemble):
@@ -137,6 +150,26 @@ def communicate(array, my_ensemble):
     return array_reduced
 
 
+def communicate_receiver_vom(array, comm):
+    """Gather all coordinates from all ranks.
+
+    Parameters
+    ----------
+    array: numpy array
+        Array of data to gather across all ranks.
+    comm: Firedrake.comm
+        Firedrake ensemble communicator.
+
+    Returns
+    -------
+    array: numpy array
+        Array of data gathered across all ranks.
+    """
+    array = array.copy()
+    array = comm.allgather(array)
+    array = np.concatenate(array)
+    return array
+
 def analytical_solution_for_pressure_based_on_MMS(model, mesh, time):
     degree = model["opts"]["degree"]
     V = FunctionSpace(mesh, "CG", degree)
diff --git a/test/not_a_test.py b/test/not_a_test.py
index f88a0ddd..f1394f1b 100644
--- a/test/not_a_test.py
+++ b/test/not_a_test.py
@@ -35,7 +35,7 @@ def test_gradient_talyor_remainder_v2():
     from ROL.firedrake_vector import FiredrakeVector as FeVector
     import ROL
 
-    comm = spyro.utils.mpi_init(model)
+    comm, _ = spyro.utils.mpi_init(model)
 
     mesh, V = spyro.io.read_mesh(model, comm)
 
diff --git a/test/test_MMS.py b/test/test_MMS.py
index facbecec..0bfedf20 100644
--- a/test/test_MMS.py
+++ b/test/test_MMS.py
@@ -86,7 +86,7 @@ def run_solve(timestep_method, method, model, mesh, expr):
     elif method == "KMV":
         variant = "KMV"
 
-    comm = spyro.utils.mpi_init(testmodel)
+    comm, _ = spyro.utils.mpi_init(testmodel)
 
     element = FiniteElement(method, mesh.ufl_cell(), degree=1, variant=variant)
     V = FunctionSpace(mesh, element)
diff --git a/test/test_gradient.py b/test/test_gradient.py
index f80277fc..392342ca 100644
--- a/test/test_gradient.py
+++ b/test/test_gradient.py
@@ -53,7 +53,7 @@ def test_gradient_pml():
 
 def _test_gradient(options, pml=False):
 
-    comm = spyro.utils.mpi_init(options)
+    comm, _ = spyro.utils.mpi_init(options)
 
     mesh, V = spyro.io.read_mesh(options, comm)
 
diff --git a/test/test_gradient_3d.py b/test/test_gradient_3d.py
index d7722f2e..55626cbb 100644
--- a/test/test_gradient_3d.py
+++ b/test/test_gradient_3d.py
@@ -39,7 +39,7 @@ def test_gradient_3d():
 
 def _test_gradient(options, pml=False):
 
-    comm = spyro.utils.mpi_init(options)
+    comm, _ = spyro.utils.mpi_init(options)
 
     mesh, V = spyro.io.read_mesh(options, comm)
 
diff --git a/test/test_newat.py b/test/test_newat.py
index c915afb6..80077243 100644
--- a/test/test_newat.py
+++ b/test/test_newat.py
@@ -19,7 +19,7 @@ def triangle_area(p1, p2, p3):
 
 def test_correct_receiver_location_generation2D():
     """Tests if receiver locations where generated correctly"""
-    comm = spyro.utils.mpi_init(model)
+    comm, _ = spyro.utils.mpi_init(model)
     mesh, V = spyro.io.read_mesh(model, comm)
 
     receivers = spyro.create_transect((-0.1, 0.3), (-0.1, 0.9), 3)
@@ -30,7 +30,7 @@ def test_correct_receiver_location_generation2D():
 
 def test_correct_receiver_to_cell_location2D():
     """Tests if the receivers where located in the correct cell"""
-    comm = spyro.utils.mpi_init(model)
+    comm, _ = spyro.utils.mpi_init(model)
     model["opts"]["degree"] = 3
     mesh, V = spyro.io.read_mesh(model, comm)
 
@@ -89,7 +89,7 @@ def test_correct_receiver_to_cell_location2D():
 
 
 def test_correct_at_value2D():
-    comm = spyro.utils.mpi_init(model)
+    comm, _ = spyro.utils.mpi_init(model)
     model["opts"]["degree"] = 3
     mesh, V = spyro.io.read_mesh(model, comm)
     pz = -0.1
@@ -122,7 +122,7 @@ def test_correct_at_value2D():
 
 def test_correct_at_value2D_quad():
     model_quad = deepcopy(model)
-    comm = spyro.utils.mpi_init(model_quad)
+    comm, _ = spyro.utils.mpi_init(model_quad)
     model_quad["opts"]["degree"] = 3
     model_quad["opts"]["degree"] = 3
     mesh, V = spyro.io.read_mesh(model_quad, comm)
@@ -174,7 +174,7 @@ def test_correct_receiver_location_generation3D():
     """Tests if receiver locations where generated correctly"""
 
     test_model = deepcopy(model3D)
-    comm = spyro.utils.mpi_init(test_model)
+    comm, _ = spyro.utils.mpi_init(test_model)
     mesh, V = spyro.io.read_mesh(test_model, comm)
     test_model["acquisition"]["num_receivers"] = 3
     receivers = spyro.create_transect((-0.05, 0.3, 0.5), (-0.05, 0.9, 0.5), 3)
@@ -189,7 +189,7 @@ def test_correct_receiver_to_cell_location3D():
     """Tests if the receivers where located in the correct cell"""
 
     test_model1 = deepcopy(model3D)
-    comm = spyro.utils.mpi_init(test_model1)
+    comm, _ = spyro.utils.mpi_init(test_model1)
     mesh, V = spyro.io.read_mesh(test_model1, comm)
     rec = spyro.create_transect((-0.05, 0.1, 0.5), (-0.05, 0.9, 0.5), 3)
     test_model1["acquisition"]["receiver_locations"] = rec
@@ -263,7 +263,7 @@ def test_correct_at_value3D():
     test_model2 = deepcopy(model3D)
     test_model2["acquisition"]["num_receivers"] = 3
     test_model2["opts"]["degree"] = 3
-    comm = spyro.utils.mpi_init(test_model2)
+    comm, _ = spyro.utils.mpi_init(test_model2)
     mesh, V = spyro.io.read_mesh(test_model2, comm)
     x_start = 0.09153949331982138
     x_end = 0.09153949331982138
diff --git a/test/test_parallel_source.py b/test/test_parallel_source.py
index e46289d0..48fbcb6f 100644
--- a/test/test_parallel_source.py
+++ b/test/test_parallel_source.py
@@ -18,7 +18,7 @@
 @pytest.mark.skip(reason="no way of currently testing this")
 def test_parallel_source():
 
-    comm = spyro.utils.mpi_init(options)
+    comm, _ = spyro.utils.mpi_init(options)
 
     mesh, V = spyro.io.read_mesh(options, comm)
 
diff --git a/test/test_sources.py b/test/test_sources.py
index dd707b09..fc1c81e1 100644
--- a/test/test_sources.py
+++ b/test/test_sources.py
@@ -52,7 +52,7 @@ def test_ricker_varies_in_time():
     # modelRicker["opts"]["method"] = "CG"
     # modelRicker["opts"]["degree"] = 2
     # modelRicker["opts"]["dimension"] = 2
-    # comm = spyro.utils.mpi_init(modelRicker)
+    # comm, _ = spyro.utils.mpi_init(modelRicker)
     # mesh = fire.UnitSquareMesh(10, 10)
     # element = fire.FiniteElement("CG", mesh.ufl_cell(), 2, variant="equispaced")
     # V = fire.FunctionSpace(mesh, element)
diff --git a/test/test_tools.py b/test/test_tools.py
index 3c79a9f8..c628264e 100644
--- a/test/test_tools.py
+++ b/test/test_tools.py
@@ -55,7 +55,7 @@ def test_mesh_generation_for_grid_calc():
     model = spyro.tools.create_model_for_grid_point_calculation(
         grid_point_calculator_parameters, degree_reference
     )
-    comm = spyro.utils.mpi_init(model)
+    comm, _ = spyro.utils.mpi_init(model)
     for G in Gs:
         model["mesh"]["meshfile"] = "meshes/2Dhomogeneous" + str(G) + ".msh"
         model = spyro.tools.generate_mesh(model, G, comm)
diff --git a/test_ad/__init__.py b/test_ad/__init__.py
deleted file mode 100644
index e69de29b..00000000
diff --git a/test_ad/test_gradient_3d_AD.py b/test_ad/test_gradient_3d_AD.py
deleted file mode 100644
index 05a6f58d..00000000
--- a/test_ad/test_gradient_3d_AD.py
+++ /dev/null
@@ -1,86 +0,0 @@
-from firedrake import *
-from scipy.optimize import *
-import spyro
-import pytest
-
-
-# from ..domains import quadrature, space
-@pytest.mark.skip(reason="no way of currently testing this with cicleCI resources")
-def test_gradient_3d_AD():
-    model = {}
-
-    model["opts"] = {
-        "method": "KMV",  # either CG or KMV
-        "quadrature": "KMV",  # Equi or KMV
-        "degree": 2,  # p order
-        "dimension": 3,  # dimension
-        "regularization": False,  # regularization is on?
-        "gamma": 1e-5,  # regularization parameter
-    }
-
-    model["parallelism"] = {
-        "type": "automatic",  # options: automatic (same number of cores for evey processor) or spatial
-    }
-
-    # Define the domain size without the ABL.
-    model["mesh"] = {
-        "Lz": 1.0,  # depth in km - always positive
-        "Lx": 1.0,  # width in km - always positive
-        "Ly": 1.0,  # thickness in km - always positive
-        "meshfile": "test/meshes/Uniform3D.msh",
-        "initmodel": "not_used.hdf5",
-        "truemodel": "not_used.hdf5",
-    }
-
-    # Specify a 250-m Absorbing Boundary Layer (ABL) on the three sides of the domain to damp outgoing waves.
-    model["BCs"] = {
-        "status": False,  # True or False, used to turn on any type of BC
-        "outer_bc": "non-reflective",  # none or non-reflective (outer boundary condition)
-        "abl_bc": "none",  # none, gaussian-taper, or alid
-        "lz": 0.25,  # thickness of the ABL in the z-direction (km) - always positive
-        "lx": 0.25,  # thickness of the ABL in the x-direction (km) - always positive
-        "ly": 0.25,  # thickness of the ABL in the y-direction (km) - always positive
-    }
-    # receivers = spyro.insert_fixed_value(spyro.create_2d_grid(0.1,0.9,0.1,0.9,2),-0.9, 0)
-
-    model["acquisition"] = {
-        "source_type": "Ricker",
-        "source_pos": [(-0.5, 0.5, 0.5)],
-        "frequency": 10.0,
-        "delay": 1.0,
-        "num_rec_x_columns": 5,
-        "num_rec_y_columns": 1,
-        "num_rec_z_columns": 1,
-        # first and final points of the receivers columns (z, x, y)
-        "receiver_locations": [(-0.1, 0.1, 0.5), (-0.1, 0.9, 0.5)],
-    }
-    model["aut_dif"] = {
-        "status": True,
-    }
-
-    model["timeaxis"] = {
-        "t0": 0.0,  # Initial time for event
-        "tf": 0.8,  # Final time for event (for test 7)
-        "dt": 0.001,  # timestep size (divided by 2 in the test 4. dt for test 3 is 0.00050)
-        "amplitude": 1,  # the Ricker has an amplitude of 1.
-        "nspool": 20,  # (20 for dt=0.00050) how frequently to output solution to pvds
-        "fspool": 1,  # how frequently to save solution to RAM
-    }
-
-    comm = spyro.utils.mpi_init(model)
-    mesh, V = spyro.io.read_mesh(model, comm)
-
-    element = spyro.domains.space.FE_method(
-        mesh, model["opts"]["method"], model["opts"]["degree"]
-    )
-
-    V = FunctionSpace(mesh, element)
-    z, x, y = SpatialCoordinate(mesh)
-
-    vp_exact = Function(V).interpolate(1.0 + 0.0 * x)
-    vp_guess = Function(V).interpolate(0.8 + 0.0 * x)
-
-    spyro.tools.gradient_test_acoustic_ad(model, mesh, V, comm, vp_exact, vp_guess)
-
-
-# test_gradient_3d_AD()
diff --git a/test_ad/test_gradient_AD.py b/test_ad/test_gradient_AD.py
deleted file mode 100644
index 34f5839a..00000000
--- a/test_ad/test_gradient_AD.py
+++ /dev/null
@@ -1,77 +0,0 @@
-from firedrake import *
-from scipy.optimize import *
-import spyro
-
-
-# from ..domains import quadrature, space
-# @pytest.mark.skip(reason="no way of currently testing this")
-def test_gradient_AD():
-    model = {}
-
-    model["opts"] = {
-        "method": "KMV",  # either CG or KMV
-        "quadrature": "KMV",  # Equi or KMV
-        "degree": 1,  # p order
-        "dimension": 2,  # dimension
-        "regularization": False,  # regularization is on?
-        "gamma": 1e-5,  # regularization parameter
-    }
-
-    model["parallelism"] = {
-        "type": "automatic",  # options: automatic (same number of cores for evey processor) or spatial
-    }
-
-    # Define the domain size without the ABL.
-    model["mesh"] = {
-        "Lz": 1.5,  # depth in km - always positive
-        "Lx": 1.5,  # width in km - always positive
-        "Ly": 0.0,  # thickness in km - always positive
-        "meshfile": "not_used.msh",
-        "initmodel": "not_used.hdf5",
-        "truemodel": "not_used.hdf5",
-    }
-
-    # Specify a 250-m Absorbing Boundary Layer (ABL) on the three sides of the domain to damp outgoing waves.
-    model["BCs"] = {
-        "status": False,  # True or False, used to turn on any type of BC
-        "outer_bc": "non-reflective",  # none or non-reflective (outer boundary condition)
-        "abl_bc": "none",  # none, gaussian-taper, or alid
-        "lz": 0.25,  # thickness of the ABL in the z-direction (km) - always positive
-        "lx": 0.25,  # thickness of the ABL in the x-direction (km) - always positive
-        "ly": 0.0,  # thickness of the ABL in the y-direction (km) - always positive
-    }
-
-    model["acquisition"] = {
-        "source_type": "Ricker",
-        "source_pos": [(0.75, 0.75)],
-        "frequency": 10.0,
-        "delay": 1.0,
-        "receiver_locations": spyro.create_transect((0.9, 0.2), (0.9, 0.8), 10),
-    }
-    model["aut_dif"] = {
-        "status": True,
-    }
-
-    model["timeaxis"] = {
-        "t0": 0.0,  # Initial time for event
-        "tf": 0.8,  # Final time for event (for test 7)
-        "dt": 0.001,  # timestep size (divided by 2 in the test 4. dt for test 3 is 0.00050)
-        "amplitude": 1,  # the Ricker has an amplitude of 1.
-        "nspool": 20,  # (20 for dt=0.00050) how frequently to output solution to pvds
-        "fspool": 1,  # how frequently to save solution to RAM
-    }
-
-    comm = spyro.utils.mpi_init(model)
-    mesh = RectangleMesh(100, 100, 1.5, 1.5)  # to test FWI, mesh aligned with interface
-
-    element = spyro.domains.space.FE_method(
-        mesh, model["opts"]["method"], model["opts"]["degree"]
-    )
-
-    V = FunctionSpace(mesh, element)
-    z, x = SpatialCoordinate(mesh)
-
-    vp_exact = Function(V).interpolate(1.0 + 0.0 * x)
-    vp_guess = Function(V).interpolate(0.8 + 0.0 * x)
-
-    spyro.tools.gradient_test_acoustic_ad(model, mesh, V, comm, vp_exact, vp_guess)