Merge pull request #12 from slimgroup/dim

Support dimension input

.github/workflows/runtests.yml

@@ -22,7 +22,7 @@ jobs:
       fail-fast: false
 
       matrix:
-        version: ['1.6', '1.7', '1.8', '1']
+        version: ['1.6', '1.7', '1.8', '1.9', '1']
         os: [ubuntu-latest]
 
     steps:

Project.toml

@@ -1,14 +1,14 @@
 name = "ImageGather"
 uuid = "355d8124-6b2e-49b5-aab5-cdbc0a5fccbe"
 authors = ["mloubout <[email protected]>"]
-version = "0.2.10"
+version = "0.3.0"
 
 [deps]
 JUDI = "f3b833dc-6b2e-5b9c-b940-873ed6319979"
 LinearAlgebra = "37e2e46d-f89d-539d-b4ee-838fcccc9c8e"
 
 [compat]
-JUDI = "3.4.4"
+JUDI = "3.4.6"
 julia = "1"
 
 [extras]

examples/layers_sscig.jl

@@ -61,16 +61,21 @@ rtm = J0'*dD
 offsets = -40f0:model.d[1]:40f0
 # offsets = [0f0]
 J = judiExtendedJacobian(F(model0), q, offsets)
+Jz = judiExtendedJacobian(F(model0), q, offsets, dims=:z)
 
 ssodm = J'*dD
+ssodmz = Jz'*dD
 
 ssor = zeros(Float32, size(ssodm)...)
 for h=1:size(ssor, 1)
     ssor[h, :, :] .= dm.data
 end
 
 dDe = J*ssor
+dDez= Jz*ssor
 # @show norm(dDe - dD), norm(ssor[:] - dm[:])
 a, b = dot(dD, dDe), dot(ssodm[:], ssor[:])
 @printf(" <F x, y> : %2.5e, <x, F' y> : %2.5e, rel-error : %2.5e, ratio : %2.5e \n", a, b, (a - b)/(a + b), a/b)
 
+a, b = dot(dD, dDez), dot(ssodmz[:], ssor[:])
+@printf(" <F x, y> : %2.5e, <x, F' y> : %2.5e, rel-error : %2.5e, ratio : %2.5e \n", a, b, (a - b)/(a + b), a/b)

src/implementation.py

@@ -1,13 +1,11 @@
-from devito import Inc, Operator, Function, CustomDimension, norm
-from devito.finite_differences.differentiable import Add
+from devito import Inc, Operator, Function, CustomDimension
 from devito.builtins import initialize_function
-from devito.types.utils import DimensionTuple
 
 import numpy as np
 
 from propagators import forward, gradient
 from geom_utils import src_rec, geom_expr
-from sensitivity import grad_expr, lin_src
+from sensitivity import grad_expr
 from fields import wavefield
 from kernels import wave_kernel
 from utils import opt_op
@@ -16,29 +14,32 @@
 def double_rtm(model, wavelet, src_coords, res, res_o, rec_coords, ic="as"):
     """
     """
-    _, u, illum, _ = forward(model, src_coords, None, wavelet, illum=True, save=True, t_sub=4)
+    _, u, illum, _ = forward(model, src_coords, None, wavelet, illum=True,
+                             save=True, t_sub=4)
 
     # RTMs
     rtm = gradient(model, res, rec_coords, u, ic=ic)[0]
     rtmo = gradient(model, res_o, rec_coords, u, ic=ic)[0]
     return rtm.data, rtmo.data, illum.data
 
 
-def cig_grad(model, src_coords, wavelet, rec_coords, res, offsets, ic="as", space_order=8, omni=False, illum=False):
+def cig_grad(model, src_coords, wavelet, rec_coords, res, offsets, ic="as",
+             space_order=8, dims=None, illum=False):
     """
     """
-    u = forward(model, src_coords, None, wavelet, space_order=space_order, save=True)[1]
+    so = max(space_order, np.max(np.abs(offsets)) // model.grid.spacing[0])
+    u = forward(model, src_coords, None, wavelet,
+                space_order=(space_order, so, so), save=True)[1]
     # Setting adjoint wavefield
-    v = wavefield(model, space_order, fw=False)
+    v = wavefield(model, space_order, fw=False, tfull=True)
 
     # Set up PDE expression and rearrange
-    pde = wave_kernel(model, v, fw=False)
+    pde, extra = wave_kernel(model, v, fw=False)
 
     # Setup source and receiver
-    go_expr = geom_expr(model, v, src_coords=rec_coords,
-                        wavelet=res, fw=False)
+    go_expr = geom_expr(model, v, src_coords=rec_coords, wavelet=res, fw=False)
     # Setup gradient wrt m with all offsets
-    ohs = make_offsets(offsets, model, omni)
+    ohs = make_offsets(offsets, model, dims)
 
     # Subsurface offsets.
     hs = tuple(h.shape[0] for h in ohs.values())
@@ -51,81 +52,85 @@ def cig_grad(model, src_coords, wavelet, rec_coords, res, offsets, ic="as", spac
 
     # Create operator and run
     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.cfunction
+    op = Operator(pde + go_expr + g_expr, subs=subs, name="cig_sso",
+                  opt=opt_op(model))
+    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)
+    op(dt=model.critical_dt)
 
     # Output
     return gradm.data
 
 
-def cig_lin(model, src_coords, wavelet, rec_coords, dm_ext, offsets, ic="as", space_order=8, omni=False):
+def cig_lin(model, src_coords, wavelet, rec_coords, dm_ext, offsets,
+            ic="as", space_order=8, dims=None):
     """
     """
+    so = max(space_order, np.max(np.abs(offsets)) // model.grid.spacing[0])
     nt = wavelet.shape[0]
     dt = model.grid.time_dim.spacing
-    oh = make_offsets(offsets, model, omni)
+    oh = make_offsets(offsets, model, dims)
 
     # Setting wavefield
-    u = wavefield(model, space_order, nt=nt)
-    ul = wavefield(model, space_order, name="l")
+    u = wavefield(model, (space_order, 2*so, so), nt=nt, tfull=True)
+    ul = wavefield(model, (space_order, so, so), name="l")
 
     # Set up PDE expression and rearrange
-    pde = wave_kernel(model, u)
-    qlin = ext_src(model, u, dm_ext, oh, ic=ic)
+    pde, extra = wave_kernel(model, u)
+    qlin = ext_src(model, u, dm_ext, oh, ul, ic=ic)
     fact = 1 / (model.damp/dt + (model.m * model.irho)/dt**2)
-    pdel = wave_kernel(model, ul) + [Inc(ul.forward, fact * qlin)]
+    pdel, extral = wave_kernel(model, ul)
+    pdel += [Inc(ul.forward, fact * qlin)]
 
     # Setup source and receiver
     go_expr = geom_expr(model, u, src_coords=src_coords, wavelet=wavelet)
     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.grid.spacing_map
-    op = Operator(pde + go_expr + pdel + go_exprl,
-                  subs=subs, name="extborn", opt=opt_op(model))
+    mode, opt = opt_op(model)
+    opt['min-storage'] = True
+    op = Operator(pde + go_expr + extra + pdel + extral + go_exprl,
+                  subs=subs, name="extborn", opt=(mode, opt))
     op.cfunction
 
     # Remove edge for offsets
-    dim_kw = make_kw(oh, DimensionTuple(*u.shape[1:], getters=model.grid.dimensions),
-                     born=True)
-    op(dt=model.critical_dt, rcvul=rcvl, **dim_kw)
+    op(dt=model.critical_dt, rcvul=rcvl)
 
     # Output
     return rcvl.data
 
 
-def ext_src(model, u, dm_ext, oh, ic="as"):
+def ext_src(model, u, dm_ext, oh, ul, ic="as"):
     # Extended perturbation
     hs = (h.shape[0] for h in oh.values())
     hd = (h.indices[0] for h in oh.values())
     dm = Function(name="gradm", grid=model.grid, shape=(*hs, *u.shape[1:]),
                   dimensions=(*hd, *model.grid.dimensions),
                   space_order=u.space_order)
-    initialize_function(dm, dm_ext, (*((0, 0) for _ in oh.values()), *model.padsizes))
+    initialize_function(dm, dm_ext, (*((0, 0) for _ in oh.values()),
+                                     *model.padsizes))
 
     # extended source
-    uh = _shift(u, oh)
-    ql = -model.irho * _shift(uh.dt2 * dm, oh)
+    uh = u
+    for k, v in oh.items():
+        uh = uh._subs(k, k - 2 * v)
+        dm = dm._subs(k, k - v)
+
+    ql = -model.irho * uh.dt2 * dm
     return ql
 
 
-def make_offsets(offsets, model, omni):
-    dims = model.grid.dimensions
-    dims = dims if omni else (dims[0],)
+def make_offsets(offsets, model, dims):
+    dims = [d for d in model.grid.dimensions if d.name in dims]
+    x, z = model.grid.dimensions[0], model.grid.dimensions[-1]
     ohs = dict()
+    nh = offsets.shape[0]
     for d in dims:
-        nh = offsets.shape[0]
-        offs = CustomDimension("off%s" % d.name, 0, nh-1, nh)
-        oh = Function(name="offv%s" % d.name, grid=model.grid, space_order=0,
+        parent = x if (len(dims) == 1 and dims[0] == z) else None
+        offs = CustomDimension("off%s" % d.name, 0, nh-1, nh, parent=parent)
+        oh = Function(name="offv%s" % d.name, space_order=0,
                       dimensions=(offs,), shape=(nh,), dtype=np.int32)
         oh.data[:] = offsets // model.grid.spacing[0]
         ohs[d] = oh
@@ -135,25 +140,9 @@ def make_offsets(offsets, model, omni):
 def shifted_wf(u, v, ohs):
     uh = u
     vh = v
-    for k, v in ohs.items():
+    for k, oh in ohs.items():
         if v.shape[0] == 1:
             continue
-        uh = uh._subs(k, k-v)
-        vh = vh._subs(k, k+v)
+        uh = uh._subs(k, k - oh)
+        vh = vh._subs(k, k + oh)
     return uh, vh
-
-
-def _shift(u, oh):
-    uh = u 
-    for k, v in oh.items():
-        uh = uh._subs(k, k-v)
-    return uh
-
-
-def make_kw(ohs, shape, born=False):
-    kw = dict()
-    scale = 2 if born else 1
-    for d, v in ohs.items():
-        kw['%s_m' % d.name] = -scale*np.min(v.data.view(np.ndarray))
-        kw['%s_M' % d.name] = shape[d] - scale*np.max(v.data.view(np.ndarray))
-    return kw

src/subsurface_gather.jl

@@ -6,24 +6,48 @@ struct judiExtendedJacobian{D, O, FT} <: judiAbstractJacobian{D, O, FT}
     F::FT
     q::judiMultiSourceVector
     offsets::Vector{D}
-    omni::Bool
+    dims::Vector{Symbol}
 end
 
 """
-    J = judiExtendedJacobian(F, q, offsets; options::JUDIOptions)
+    J = judiExtendedJacobian(F, q, offsets; options::JUDIOptions, omni=false, dims=nothing)
 
 Extended jacobian (extended Born modeling operator) for subsurface horsizontal offsets `offsets`. Its adjoint
 comput the subsurface common offset volume. In succint way, the extened born modeling Operator can summarized in a linear algebra frmaework as:
 
+Options structure for seismic modeling.
+
+`omni`: If `true`, the extended jacobian will be computed for all dimensions.
+`dims`: If `omni` is `false`, the extended jacobian will be computed for the dimension(s) specified in `dims`.
+
 """
-function judiExtendedJacobian(F::judiComposedPropagator{D, O}, q::judiMultiSourceVector, offsets; options=nothing, omni=false) where {D, O}
+function judiExtendedJacobian(F::judiComposedPropagator{D, O}, q::judiMultiSourceVector, offsets;
+                              options=nothing, omni=false, dims=nothing) where {D, O}
     JUDI.update!(F.options, options)
     offsets = Vector{D}(offsets)
-    return judiExtendedJacobian{D, :born, typeof(F)}(F.m, space(F.model.n), F, q, offsets, omni)
+    ndim = length(F.model.n)
+    if omni
+        dims = [:x, :y, :z][1:ndim]
+    else
+        if isnothing(dims)
+            dims = [:x]
+        else
+            dims = symvec(dims)
+            if ndim == 2
+                dims[dims .== :z] .= :y
+            end
+        end
+    end
+
+    return judiExtendedJacobian{D, :born, typeof(F)}(F.m, space(F.model.n), F, q, offsets, dims)
 end
 
-adjoint(J::judiExtendedJacobian{D, O, FT}) where {D, O, FT} = judiExtendedJacobian{D, adjoint(O), FT}(J.n, J.m, J.F, J.q, J.offsets, J.omni)
-getindex(J::judiExtendedJacobian{D, O, FT}, i) where {D, O, FT} = judiExtendedJacobian{D, O, FT}(J.m[i], J.n[i], J.F[i], J.q[i], J.offsets, J.omni)
+symvec(s::Symbol) = [s]
+symvec(s::Tuple) = [symvec(ss)[1] for ss in s]::Vector{Symbol}
+symvec(s::Vector) = [symvec(ss)[1] for ss in s]::Vector{Symbol}
+
+adjoint(J::judiExtendedJacobian{D, O, FT}) where {D, O, FT} = judiExtendedJacobian{D, adjoint(O), FT}(J.n, J.m, J.F, J.q, J.offsets, J.dims)
+getindex(J::judiExtendedJacobian{D, O, FT}, i) where {D, O, FT} = judiExtendedJacobian{D, O, FT}(J.m[i], J.n[i], J.F[i], J.q[i], J.offsets, J.dims)
 
 function make_input(J::judiExtendedJacobian{D, :adjoint_born, FT}, q) where {D, FT}
     srcGeom, srcData = JUDI.make_src(J.q, J.F.qInjection)
@@ -56,7 +80,8 @@ function propagate(J::judiExtendedJacobian{T, :born, O}, q::AbstractArray{T}, il
 
     # Set up Python model structure
     modelPy = devito_model(J.model, J.options)
-    dmd = reshape(dm, length(J.offsets), J.model.n...)
+    nh = [length(J.offsets) for _=1:length(J.dims)]
+    dmd = reshape(dm, nh..., J.model.n...)
     dtComp = convert(Float32, modelPy."critical_dt")
 
     # Extrapolate input data to computational grid
@@ -68,7 +93,7 @@ function propagate(J::judiExtendedJacobian{T, :born, O}, q::AbstractArray{T}, il
 
     # Devito interface
     dD = JUDI.wrapcall_data(impl."cig_lin", modelPy, src_coords, qIn, rec_coords,
-                            dmd, J.offsets, ic=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, dims=J.dims)
     dD = time_resample(dD, dtComp, recGeometry)
     # Output shot record as judiVector
     return judiVector{Float32, Matrix{Float32}}(1, recGeometry, [dD])
@@ -95,7 +120,7 @@ function propagate(J::judiExtendedJacobian{T, :adjoint_born, O}, q::AbstractArra
     # Devito
     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)
+               illum=false, ic=J.options.IC, space_order=J.options.space_order, dims=J.dims)
     end
     g = remove_padding_cig(g, modelPy.padsizes; true_adjoint=J.options.sum_padding)
     return g

test/runtests.jl

@@ -59,22 +59,34 @@ rtm = J0'*dD
 offsets = -40f0:model.d[1]:40f0
 nh = length(offsets)
 
-J = judiExtendedJacobian(F(model0), q, offsets)
+for dims in ((:x, :z), :z, :x)
 
-ssodm = J'*dD
-@show size(ssodm)
-@test size(ssodm, 1) == nh
+    J = judiExtendedJacobian(F(model0), q, offsets, dims=dims)
 
-ssor = zeros(Float32, size(ssodm)...)
-for h=1:size(ssor, 1)
-    ssor[h, :, :] .= dm.data
-end
+    ssodm = J'*dD
+    @show size(ssodm)
+    @test size(ssodm, 1) == nh
 
-dDe = J*ssor
-# @show norm(dDe - dD), norm(ssor[:] - dm[:])
-a, b = dot(dD, dDe), dot(ssodm[:], ssor[:])
+    ssor = zeros(Float32, size(ssodm)...)
+    for h=1:size(ssor, 1)
+        if dims == (:x, :z)
+            for h2=1:size(ssor, 2)
+                ssor[h, h2, :, :] .= dm.data
+            end
+        else
+            ssor[h, :, :] .= dm.data
+        end
+    end
 
-@test (a-b)/(a+b) ≈ 0 atol=sqrt(eps(1f0)) rtol=0
+    dDe = J*ssor
+    # @show norm(dDe - dD), norm(ssor[:] - dm[:])
+    a, b = dot(dD, dDe), dot(ssodm[:], ssor[:])
 
-# 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 (a-b)/(a+b) ≈ 0 atol=1f-3 rtol=0
+
+    # Make sure zero offset is the rtm, remove the sumpadding
+    ih = div(nh, 2)+1
+    rtmc = dims == (:x, :z) ? ssodm[ih, ih, :, :] : ssodm[ih, :, :]
+
+    @test norm(rtm.data - rtmc, Inf) ≈ 0f0 atol=1f-4 rtol=0
+end