diff --git a/demos/with_automatic_differentiation/forward_circle_vp.py b/demos/with_automatic_differentiation/run_forward.py
similarity index 55%
rename from demos/with_automatic_differentiation/forward_circle_vp.py
rename to demos/with_automatic_differentiation/run_forward.py
index c4d6a10f..0867b8a6 100644
--- a/demos/with_automatic_differentiation/forward_circle_vp.py
+++ b/demos/with_automatic_differentiation/run_forward.py
@@ -7,7 +7,7 @@
 model = {}
 
 model["opts"] = {
-    "method": "CG",  # either CG or KMV
+    "method": "KMV",  # either CG or KMV
     "quadrature": "KMV",  # Equi or KMV
     "degree": 1,  # p order
     "dimension": 2,  # dimension
@@ -21,8 +21,8 @@
 
 # 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
+    "Lz": 1.0,  # depth in km - always positive
+    "Lx": 1.0,  # width in km - always positive
     "Ly": 0.0,  # thickness in km - always positive
     "meshfile": "not_used.msh",
     "initmodel": "not_used.hdf5",
@@ -34,17 +34,17 @@
     "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
+    "lz": 0.0,  # thickness of the ABL in the z-direction (km) - always positive
+    "lx": 0.0,  # 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": spyro.create_transect((0.2, 0.2), (0.2, 0.8), 1),
+    "source_pos": spyro.create_transect((0.2, 0.15), (0.8, 0.15), 3),
     "frequency": 10.0,
     "delay": 1.0,
-    "receiver_locations": spyro.create_transect((0.9, 0.2), (0.9, 0.8), 10),
+    "receiver_locations": spyro.create_transect((0.2, 0.2), (0.8, 0.2), 10),
 }
 model["aut_dif"] = {
     "status": True,
@@ -60,20 +60,48 @@
 }
 
 comm = spyro.utils.mpi_init(model)
-mesh = RectangleMesh(100, 100, 1.5, 1.5)  # to test FWI, mesh aligned with interface
+mesh = UnitSquareMesh(100, 100)  # 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)
+
+def make_vp_circle(vp_guess=False, plot_vp=False):
+    """Acoustic velocity model"""
+    x, z = SpatialCoordinate(mesh)
+    if vp_guess:
+        vp = Function(V).interpolate(1.5 + 0.0 * x)
+    else:
+        vp = Function(V).interpolate(
+            2.5
+            + 1 * tanh(100 * (0.125 - sqrt((x - 0.5) ** 2 + (z - 0.5) ** 2)))
+        )
+    if plot_vp:
+        outfile = File("acoustic_cp.pvd")
+        outfile.write(vp)
+    return vp
+
 wavelet = spyro.full_ricker_wavelet(
     dt=model["timeaxis"]["dt"],
     tf=model["timeaxis"]["tf"],
     freq=model["acquisition"]["frequency"],
 )
 
-spyro.solvers.forward_AD(model, mesh, comm, vp_exact, wavelet, debug=True)
+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)
+else:
+    raise NotImplementedError("Only implemented for 1 processor")
diff --git a/spyro/solvers/forward_AD.py b/spyro/solvers/forward_AD.py
index 60bba6c2..594f2428 100644
--- a/spyro/solvers/forward_AD.py
+++ b/spyro/solvers/forward_AD.py
@@ -8,7 +8,7 @@
 
 
 # @ensemble_forward
-def forward(model, mesh, comm, c, wavelet, source_num=0, fwi=False, **kwargs):
+def forward(model, mesh, comm, c, wavelet, source_number=0, fwi=False, **kwargs):
     """Secord-order in time fully-explicit scheme.
 
     Parameters
@@ -23,7 +23,7 @@ def forward(model, mesh, comm, c, wavelet, source_num=0, fwi=False, **kwargs):
         The velocity model interpolated onto the mesh.
     wavelet: array-like
         Time series data that's injected at the source location.
-    source_num: `int`, optional
+    source_number: `int`, optional
         The source number you wish to simulate
     fwi: `bool`, optional
         Whether this forward simulation is for FWI or not.
@@ -98,7 +98,7 @@ def forward(model, mesh, comm, c, wavelet, source_num=0, fwi=False, **kwargs):
         # cost function
         J = 0.0
     for step in range(nt):
-        f.assign(source.apply_source_based_in_vom(wavelet[step], source_num))
+        f.assign(source.apply_source_based_in_vom(1.0, source_number)*Constant(wavelet[step]))
         solver.solve()
         u_np1.assign(X)
         # receiver function
diff --git a/spyro/sources/Sources.py b/spyro/sources/Sources.py
index 9390be00..2a17b625 100644
--- a/spyro/sources/Sources.py
+++ b/spyro/sources/Sources.py
@@ -117,6 +117,8 @@ def apply_source_based_in_vom(self, wavelet, source_number):
         ----------
         wavelet : float
             Value of the wavelet at time t.
+        source_number : int
+            The source number.
 
         Returns
         -------