From 821b88ccaf0695291fdbdfed2146d01efcb58b4d Mon Sep 17 00:00:00 2001 From: mloubout Date: Tue, 13 Aug 2024 18:30:31 -0400 Subject: [PATCH 1/2] support dimension input for sscig --- .github/workflows/runtests.yml | 2 +- Project.toml | 2 +- examples/layers_sscig.jl | 5 ++ src/implementation.py | 113 +++++++++++++++------------------ src/subsurface_gather.jl | 43 ++++++++++--- test/runtests.jl | 40 ++++++++---- 6 files changed, 118 insertions(+), 87 deletions(-) diff --git a/.github/workflows/runtests.yml b/.github/workflows/runtests.yml index 9308880..095fef1 100644 --- a/.github/workflows/runtests.yml +++ b/.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: diff --git a/Project.toml b/Project.toml index 359fad6..7bb26a9 100644 --- a/Project.toml +++ b/Project.toml @@ -8,7 +8,7 @@ JUDI = "f3b833dc-6b2e-5b9c-b940-873ed6319979" LinearAlgebra = "37e2e46d-f89d-539d-b4ee-838fcccc9c8e" [compat] -JUDI = "3.4.4" +JUDI = "3.4.6" julia = "1" [extras] diff --git a/examples/layers_sscig.jl b/examples/layers_sscig.jl index 0af266d..756dd17 100644 --- a/examples/layers_sscig.jl +++ b/examples/layers_sscig.jl @@ -61,8 +61,10 @@ 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) @@ -70,7 +72,10 @@ for h=1:size(ssor, 1) end dDe = J*ssor +dDez= Jz*ssor # @show norm(dDe - dD), norm(ssor[:] - dm[:]) a, b = dot(dD, dDe), dot(ssodm[:], ssor[:]) @printf(" : %2.5e, : %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(" : %2.5e, : %2.5e, rel-error : %2.5e, ratio : %2.5e \n", a, b, (a - b)/(a + b), a/b) diff --git a/src/implementation.py b/src/implementation.py index 608d25f..f4614ec 100644 --- a/src/implementation.py +++ 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,7 +14,8 @@ 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] @@ -24,21 +23,23 @@ def double_rtm(model, wavelet, src_coords, res, res_o, rec_coords, ic="as"): 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,38 +52,36 @@ 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) @@ -90,42 +89,48 @@ def cig_lin(model, src_coords, wavelet, rec_coords, dm_ext, offsets, ic="as", sp _, 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 diff --git a/src/subsurface_gather.jl b/src/subsurface_gather.jl index 45d6870..1ea9a31 100644 --- a/src/subsurface_gather.jl +++ b/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 diff --git a/test/runtests.jl b/test/runtests.jl index a320e3a..8bfd230 100644 --- a/test/runtests.jl +++ b/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 \ No newline at end of file From 79287bf2962d12b4592ec39d4e805ee0bf7b3949 Mon Sep 17 00:00:00 2001 From: mloubout Date: Tue, 13 Aug 2024 22:02:28 -0400 Subject: [PATCH 2/2] bump version --- Project.toml | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/Project.toml b/Project.toml index 7bb26a9..034769e 100644 --- a/Project.toml +++ b/Project.toml @@ -1,7 +1,7 @@ name = "ImageGather" uuid = "355d8124-6b2e-49b5-aab5-cdbc0a5fccbe" authors = ["mloubout "] -version = "0.2.10" +version = "0.3.0" [deps] JUDI = "f3b833dc-6b2e-5b9c-b940-873ed6319979"