diff --git a/.github/workflows/runtests.yml b/.github/workflows/runtests.yml
new file mode 100644
index 0000000..9308880
--- /dev/null
+++ b/.github/workflows/runtests.yml
@@ -0,0 +1,42 @@
+name: CI
+
+on:
+  push:
+    branches:
+      - main
+  pull_request:
+    branches:
+      - main
+
+jobs:
+  test:
+    name: Julia ${{ matrix.version }} - ${{ matrix.os }}
+    runs-on: ${{ matrix.os }}
+    env:
+      DEVITO_ARCH: gcc-9
+      DEVITO_LANGUAGE: "openmp"
+      DEVITO_LOGGING: "ERROR"
+      OMP_NUM_THREADS: 1
+
+    strategy:
+      fail-fast: false
+
+      matrix:
+        version: ['1.6', '1.7', '1.8', '1']
+        os: [ubuntu-latest]
+
+    steps:
+      - name: Checkout ImageGather.jl
+        uses: actions/checkout@v3
+
+      - name: Setup julia
+        uses: julia-actions/setup-julia@v1
+        with:
+          version: ${{ matrix.version }}
+
+      - name: Build ImageGather.jl
+        uses: julia-actions/julia-buildpkg@latest
+
+      - name: Run tests
+        id: test
+        uses: julia-actions/julia-runtest@latest
diff --git a/Project.toml b/Project.toml
index c0ab9d9..b17ec07 100644
--- a/Project.toml
+++ b/Project.toml
@@ -1,13 +1,14 @@
 name = "ImageGather"
 uuid = "355d8124-6b2e-49b5-aab5-cdbc0a5fccbe"
 authors = ["mloubout <mathias.louboutin@gmail.com>"]
-version = "0.2.4"
+version = "0.2.5"
 
 [deps]
 JUDI = "f3b833dc-6b2e-5b9c-b940-873ed6319979"
+LinearAlgebra = "37e2e46d-f89d-539d-b4ee-838fcccc9c8e"
 
 [compat]
-JUDI = "3.1.3"
+JUDI = "3.3.1"
 julia = "1"
 
 [extras]
diff --git a/src/implementation.py b/src/implementation.py
index eb68898..c57a7b7 100644
--- a/src/implementation.py
+++ b/src/implementation.py
@@ -16,24 +16,20 @@
 def double_rtm(model, wavelet, src_coords, res, res_o, rec_coords, space_order=8):
     """
     """
-    _, u, _ = forward(model, src_coords, None, wavelet, space_order=space_order,
-                      save=True)
-
-    # Illumination
-    illum = Function(name="I", grid=model.grid, space_order=0)
-    Operator(Inc(illum, u**2))()
+    _, u, illum, _ = forward(model, src_coords, None, wavelet, space_order=space_order,
+                             illum=True, save=True)
 
     # RTMs
-    rtm, _ = gradient(model, res, rec_coords, u, space_order=space_order)
-    rtmo, _ = gradient(model, res_o, rec_coords, u, space_order=space_order)
+    rtm = gradient(model, res, rec_coords, u, space_order=space_order)[0]
+    rtmo = gradient(model, res_o, rec_coords, u, space_order=space_order)[0]
     return rtm.data, rtmo.data, illum.data
 
 
-def cig_grad(model, src_coords, wavelet, rec_coords, res, offsets, isic=False, space_order=8, omni=False):
+def cig_grad(model, src_coords, wavelet, rec_coords, res, offsets, ic="as", space_order=8, omni=False, illum=False):
     """
     """
     u = forward(model, src_coords, None, wavelet, space_order=space_order, save=True)[1]
-    # Setting adjoint wavefieldgradient
+    # Setting adjoint wavefield
     v = wavefield(model, space_order, fw=False)
 
     # Set up PDE expression and rearrange
@@ -41,27 +37,29 @@ def cig_grad(model, src_coords, wavelet, rec_coords, res, offsets, isic=False, s
 
     # Setup source and receiver
     go_expr = geom_expr(model, v, src_coords=rec_coords,
-                       wavelet=res, fw=False)
+                        wavelet=res, fw=False)
     # Setup gradient wrt m with all offsets
     ohs = make_offsets(offsets, model, omni)
-    x = u.indices[1]
 
     # Subsurface offsets.
-    hs = (h.shape[0] for h in ohs.values())
-    hd = (h.indices[0] for h in ohs.values())
+    hs = tuple(h.shape[0] for h in ohs.values())
+    hd = tuple(h.indices[0] for h in ohs.values())
     gradm = Function(name="gradm", grid=model.grid, shape=(*hs, *u.shape[1:]),
-                     dimensions=(*hd, *model.grid.dimensions), space_order=0)
+                    dimensions=(*hd, *model.grid.dimensions))
+
     uh, vh = shifted_wf(u, v, ohs)
-    g_expr = grad_expr(gradm, uh, vh, model, ic=isic)
+    g_expr = grad_expr(gradm, uh, vh, model, ic=ic)
 
     # Create operator and run
-    subs = model.spacing_map
-    op = Operator(pde + go_expr + g_expr,subs=subs, name="cig_sso", opt=opt_op(model))
+    subs = model.grid.spacing_map
+    op = Operator(pde + go_expr + g_expr, subs=subs, name="cig_sso", opt=opt_op(model))
+
     try:
         op.cfunction
     except:
-        op = Operator(pde + go_expr + g_expr,subs=subs, name="cig_sso", opt='advanced')
+        op = Operator(pde + go_expr + g_expr, subs=subs, name="cig_sso", opt='advanced')
         op.cfunction
+
     # Get bounds from offsets
     dim_kw = make_kw(ohs, DimensionTuple(*u.shape[1:], getters=model.grid.dimensions))
     op(dt=model.critical_dt, **dim_kw)
@@ -70,7 +68,7 @@ def cig_grad(model, src_coords, wavelet, rec_coords, res, offsets, isic=False, s
     return gradm.data
 
 
-def cig_lin(model, src_coords, wavelet, rec_coords, dm_ext, offsets, isic=False, space_order=8, omni=False):
+def cig_lin(model, src_coords, wavelet, rec_coords, dm_ext, offsets, ic="as", space_order=8, omni=False):
     """
     """
     nt = wavelet.shape[0]
@@ -83,7 +81,7 @@ def cig_lin(model, src_coords, wavelet, rec_coords, dm_ext, offsets, isic=False,
 
     # Set up PDE expression and rearrange
     pde = wave_kernel(model, u)
-    qlin = ext_src(model, u, dm_ext, oh, isic=isic)
+    qlin = ext_src(model, u, dm_ext, oh, ic=ic)
     fact = 1 / (model.damp/dt + (model.m * model.irho)/dt**2)
     pdel = wave_kernel(model, ul) + [Inc(ul.forward, fact * qlin)]
 
@@ -92,7 +90,7 @@ def cig_lin(model, src_coords, wavelet, rec_coords, dm_ext, offsets, isic=False,
     go_exprl = geom_expr(model, ul, rec_coords=rec_coords, nt=nt)
     _, rcvl = src_rec(model, ul, rec_coords=rec_coords, nt=nt)
     # Create operator and run
-    subs = model.spacing_map
+    subs = model.grid.spacing_map
     op = Operator(pde + go_expr + pdel + go_exprl,
                   subs=subs, name="extborn", opt=opt_op(model))
     op.cfunction
@@ -104,7 +102,7 @@ def cig_lin(model, src_coords, wavelet, rec_coords, dm_ext, offsets, isic=False,
     return rcvl.data
 
 
-def ext_src(model, u, dm_ext, oh, isic=False):
+def ext_src(model, u, dm_ext, oh, ic="as"):
     # Extended perturbation
     hs = (h.shape[0] for h in oh.values())
     hd = (h.indices[0] for h in oh.values())
@@ -137,6 +135,8 @@ def shifted_wf(u, v, ohs):
     uh = u
     vh = v
     for k, v in ohs.items():
+        if v.shape[0] == 1:
+            continue
         uh = uh._subs(k, k-v)
         vh = vh._subs(k, k+v)
     return uh, vh
@@ -152,6 +152,6 @@ def _shift(u, oh):
 def make_kw(ohs, shape):
     kw = dict()
     for d, v in ohs.items():
-        kw['%s_m' % d.name] = -np.min(v.data)
-        kw['%s_M' % d.name] = shape[d] - np.max(v.data)
+        kw['%s_m' % d.name] = -np.min(v.data.view(np.ndarray))
+        kw['%s_M' % d.name] = shape[d] - np.max(v.data.view(np.ndarray))
     return kw
diff --git a/src/subsurface_gather.jl b/src/subsurface_gather.jl
index 21e2d58..d450657 100644
--- a/src/subsurface_gather.jl
+++ b/src/subsurface_gather.jl
@@ -34,7 +34,8 @@ end
 function make_input(J::judiExtendedJacobian{D, :born, FT}, dm) where {D<:Number, FT}
     srcGeom, srcData = JUDI.make_src(J.q, J.F.qInjection)
     return srcGeom, srcData, J.F.rInterpolation.data[1], nothing, dm
-end 
+end
+
 *(J::judiExtendedJacobian{T, :born, O}, dm::Array{T, 3}) where {T, O} = J*vec(dm)
 *(J::judiExtendedJacobian{T, :born, O}, dm::Array{T, 4}) where {T, O} = J*vec(dm)
 
@@ -59,7 +60,7 @@ function propagate(J::judiExtendedJacobian{T, :born, O}, q::AbstractArray{T}) wh
     dtComp = convert(Float32, modelPy."critical_dt")
 
     # Extrapolate input data to computational grid
-    qIn = time_resample(srcData, srcGeometry, dtComp)[1]
+    qIn = time_resample(srcData, srcGeometry, dtComp)
 
     # Set up coordinates
     src_coords = setup_grid(srcGeometry, J.model.n)  # shifts source coordinates by origin
@@ -67,7 +68,7 @@ function propagate(J::judiExtendedJacobian{T, :born, O}, q::AbstractArray{T}) wh
 
     # Devito interface
     dD = JUDI.wrapcall_data(impl."cig_lin", modelPy, src_coords, qIn, rec_coords,
-                            dmd, J.offsets, isic=J.options.IC, space_order=J.options.space_order, omni=J.omni)
+                            dmd, J.offsets, ic=J.options.IC, space_order=J.options.space_order, omni=J.omni)
     dD = time_resample(dD, dtComp, recGeometry)
     # Output shot record as judiVector
     return judiVector{Float32, Matrix{Float32}}(1, recGeometry, [dD])
@@ -84,16 +85,18 @@ function propagate(J::judiExtendedJacobian{T, :adjoint_born, O}, q::AbstractArra
     dtComp = convert(Float32, modelPy."critical_dt")
 
     # Extrapolate input data to computational grid
-    qIn = time_resample(srcData, srcGeometry, dtComp)[1]
-    dObserved = time_resample(recData, recGeometry, dtComp)[1]
+    qIn = time_resample(srcData, srcGeometry, dtComp)
+    dObserved = time_resample(recData, recGeometry, dtComp)
 
     # Set up coordinates
     src_coords = setup_grid(srcGeometry, J.model.n)  # shifts source coordinates by origin
     rec_coords = setup_grid(recGeometry, J.model.n)  # shifts rec coordinates by origin
 
     # Devito
-    g = pycall(impl."cig_grad", PyArray, modelPy, src_coords, qIn, rec_coords, dObserved,
-               J.offsets, isic=J.options.IC, space_order=J.options.space_order, omni=J.omni)
+    g = JUDI.pylock() do 
+        pycall(impl."cig_grad", PyArray, modelPy, src_coords, qIn, rec_coords, dObserved, J.offsets,
+               illum=false, ic=J.options.IC, space_order=J.options.space_order, omni=J.omni)
+    end
     g = remove_padding_cig(g, modelPy.padsizes; true_adjoint=J.options.sum_padding)
     return g
 end
diff --git a/src/surface_gather.jl b/src/surface_gather.jl
index 80d7b4b..6cbf2b1 100644
--- a/src/surface_gather.jl
+++ b/src/surface_gather.jl
@@ -74,8 +74,11 @@ function double_rtm_cig(model_full, q::judiVector, data::judiVector, offs, optio
     mute!(res, off_r .- scale; dt=dtComp/1f3, t0=.25)
     res_o = res .* off_r'
     # Double rtm
-    rtm, rtmo, illum = pycall(impl."double_rtm", Tuple{PyArray, PyArray, PyArray},
-                              modelPy, qIn, src_coords, res, res_o, rec_coords, space_order=options.space_order)
+
+    rtm, rtmo, illum = JUDI.pylock() do
+        pycall(impl."double_rtm", Tuple{PyArray, PyArray, PyArray},
+               modelPy, qIn, src_coords, res, res_o, rec_coords, space_order=options.space_order)
+    end
     rtm = remove_padding(rtm, modelPy.padsizes)
     rtmo = remove_padding(rtmo, modelPy.padsizes)
     illum = remove_padding(illum, modelPy.padsizes)
diff --git a/test/runtests.jl b/test/runtests.jl
index 5af770e..08fad29 100644
--- a/test/runtests.jl
+++ b/test/runtests.jl
@@ -1,3 +1,79 @@
 using ImageGather, Test
 
-@test true
\ No newline at end of file
+using JUDI, LinearAlgebra
+
+# Set up model structure
+n = (301, 151)   # (x,y,z) or (x,z)
+d = (10., 10.)
+o = (0., 0.)
+
+# Velocity [km/s]
+v =  1.5f0 .* ones(Float32,n)
+v[:, 76:end] .= 2.5f0
+v0 = 1.5f0 .* ones(Float32,n)
+# Slowness squared [s^2/km^2]
+m = (1f0 ./ v).^2
+m0 = (1f0 ./ v0).^2
+
+# Setup info and model structure
+nsrc = 1	# number of sources
+model = Model(n, d, o, m; nb=40)
+model0 = Model(n, d, o, m0; nb=40)
+
+dm = model.m - model0.m
+
+# Set up receiver geometry
+nxrec = 151
+xrec = range(0f0, stop=(n[1] -1)*d[1], length=nxrec)
+yrec = 0f0
+zrec = range(20f0, stop=20f0, length=nxrec)
+
+# receiver sampling and recording time
+timeD = 2000f0   # receiver recording time [ms]
+dtD = 4f0    # receiver sampling interval [ms]
+
+# Set up receiver structure
+recGeometry = Geometry(xrec, yrec, zrec; dt=dtD, t=timeD, nsrc=nsrc)
+# Set up source geometry (cell array with source locations for each shot)
+xsrc = 1500f0
+ysrc = 0f0
+zsrc = 20f0
+
+# Set up source structure
+srcGeometry = Geometry(xsrc, ysrc, zsrc; dt=dtD, t=timeD)
+# setup wavelet
+f0 = 0.015f0     # kHz
+wavelet = ricker_wavelet(timeD, dtD, f0)
+q = diff(judiVector(srcGeometry, wavelet))
+
+###################################################################################################
+opt = Options()
+# Setup operators
+F = judiModeling(model, srcGeometry, recGeometry; options=opt)
+J0 = judiJacobian(F(model0), q)
+# Nonlinear modeling
+dD = J0*dm
+rtm = J0'*dD
+
+# Common surface offset image gather
+offsets = -40f0:model.d[1]:40f0
+nh = length(offsets)
+
+J = judiExtendedJacobian(F(model0), q, offsets)
+
+ssodm = J'*dD
+@test size(ssodm, 1) == nh
+
+ssor = zeros(Float32, size(ssodm)...)
+for h=1:size(ssor, 1)
+    ssor[h, :, :] .= dm.data
+end
+
+dDe = J*ssor
+# @show norm(dDe - dD), norm(ssor[:] - dm[:])
+a, b = dot(dD, dDe), dot(ssodm[:], ssor[:])
+
+@test (a-b)/(a+b) ≈ 0 atol=sqrt(eps(1f0)) rtol=0
+
+# Make sure zero offset is the rtm, remove the sumpadding
+@test norm(rtm.data - ssodm[div(nh, 2)+1, :, :])/norm(rtm) ≈ 0f0 atol=1f-5 rtol=0
\ No newline at end of file