diff --git a/Project.toml b/Project.toml
index 77b2da7..359fad6 100644
--- a/Project.toml
+++ b/Project.toml
@@ -1,7 +1,7 @@
 name = "ImageGather"
 uuid = "355d8124-6b2e-49b5-aab5-cdbc0a5fccbe"
 authors = ["mloubout <mathias.louboutin@gmail.com>"]
-version = "0.2.9"
+version = "0.2.10"
 
 [deps]
 JUDI = "f3b833dc-6b2e-5b9c-b940-873ed6319979"
diff --git a/examples/offset_map.jl b/examples/offset_map.jl
index da2169b..83f4fb3 100644
--- a/examples/offset_map.jl
+++ b/examples/offset_map.jl
@@ -2,7 +2,7 @@
 # Date: June 2021
 #
 
-using JUDI, LinearAlgebra, Images, PyPlot, DSP, ImageGather
+using JUDI, LinearAlgebra, Images, PyPlot, DSP, ImageGather, SlimPlotting
 
 # Set up model structure
 n = (601, 333)   # (x,y,z) or (x,z)
@@ -54,7 +54,7 @@ wavelet = ricker_wavelet(timeS, dtS, f0)
 q = judiVector(srcGeometry, wavelet)
 
 ###################################################################################################
-opt = Options(space_order=16)
+opt = Options(space_order=16, IC="as")
 # Setup operators
 F = judiModeling(model, srcGeometry, recGeometry; options=opt)
 F0 = judiModeling(model0, srcGeometry, recGeometry; options=opt)
@@ -70,15 +70,17 @@ res = deepcopy(dD)
 mute!(res.data[1], offs)
 reso = deepcopy(res)
 
-rtm = J'*res
+I = inv(judiIllumination(J))
+
+rtm = I*J'*res
 
 omap = Array{Any}(undef, 2)
 i = 1
 
 # try a bunch of weighting functions
-for (wf, iwf) = zip([x-> x./4000f0 .+ 1, x-> log.(x .+ 10)], [x -> 4000f0.*(x .- 1), x-> exp.(x) .- 10])
+for (wf, iwf) = zip([x-> x .+ 5f3, x-> log.(x .+ 10)], [x -> x .- 5f3, x-> exp.(x) .- 10])
     reso.data[1] .= res.data[1] .* wf(offs)'
-    rtmo = J'*reso
+    rtmo = I*J'*reso
     omap[i] = iwf(offset_map(rtm.data, rtmo.data))
     global i+=1
 end
@@ -86,7 +88,8 @@ end
 figure()
 for (i, name)=enumerate(["shift", "log"])
     subplot(1,2,i)
-    imshow(omap[i]', vmin=0, vmax=4000, cmap="gist_ncar", aspect="auto")
+    plot_velocity(omap[i]', (1,1); cmap="jet", aspect="auto", perc=98, new_fig=false, vmax=5000)
     colorbar()
     title(name)
 end
+tight_layout()
diff --git a/src/subsurface_gather.jl b/src/subsurface_gather.jl
index d450657..45d6870 100644
--- a/src/subsurface_gather.jl
+++ b/src/subsurface_gather.jl
@@ -43,7 +43,7 @@ JUDI.process_input_data(::judiExtendedJacobian{D, :born, FT}, q::Vector{D}) wher
 
 ############################################################
 
-function propagate(J::judiExtendedJacobian{T, :born, O}, q::AbstractArray{T}) where {T, O}
+function propagate(J::judiExtendedJacobian{T, :born, O}, q::AbstractArray{T}, illum::Bool) where {T, O}
     srcGeometry, srcData, recGeometry, _, dm = make_input(J, q)
     # Load full geometry for out-of-core geometry containers
     recGeometry = Geometry(recGeometry)
@@ -74,7 +74,7 @@ function propagate(J::judiExtendedJacobian{T, :born, O}, q::AbstractArray{T}) wh
     return judiVector{Float32, Matrix{Float32}}(1, recGeometry, [dD])
 end
 
-function propagate(J::judiExtendedJacobian{T, :adjoint_born, O}, q::AbstractArray{T}) where {T, O}
+function propagate(J::judiExtendedJacobian{T, :adjoint_born, O}, q::AbstractArray{T}, illum::Bool) where {T, O}
     srcGeometry, srcData, recGeometry, recData, _ = make_input(J, q)
     # Load full geometry for out-of-core geometry containers
     recGeometry = Geometry(recGeometry)
diff --git a/test/runtests.jl b/test/runtests.jl
index f413fb0..a320e3a 100644
--- a/test/runtests.jl
+++ b/test/runtests.jl
@@ -62,11 +62,12 @@ nh = length(offsets)
 J = judiExtendedJacobian(F(model0), q, offsets)
 
 ssodm = J'*dD
-@test size(ssodm, 3) == nh
+@show size(ssodm)
+@test size(ssodm, 1) == nh
 
 ssor = zeros(Float32, size(ssodm)...)
-for h=1:size(ssor, 3)
-    ssor[:, :, h] .= dm.data
+for h=1:size(ssor, 1)
+    ssor[h, :, :] .= dm.data
 end
 
 dDe = J*ssor
@@ -76,4 +77,4 @@ 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
+@test norm(rtm.data - ssodm[div(nh, 2)+1, :, :])/norm(rtm) ≈ 0f0 atol=1f-5 rtol=0