better fs

examples/scripts/modeling_basic_2D.jl

@@ -9,10 +9,12 @@
 #' This example is converted to a markdown file for the documentation.
 
 #' # Import JUDI, Linear algebra utilities and Plotting
-using JUDI, PyPlot, LinearAlgebra
+using JUDI, PyPlot, LinearAlgebra, SlimPlotting
 
 #+ echo = false; results = "hidden"
 close("all")
+imcmap = "cet_CET_L1"
+dcmap = "PuOr"
 
 #' # Create a JUDI model structure
 #' In JUDI, a `Model` structure contains the grid information (origin, spacing, number of gridpoints)
@@ -91,7 +93,7 @@ q = judiVector(srcGeometry, wavelet)
 #' condition for the propagation.
 
 # Setup options
-opt = Options(subsampling_factor=2, space_order=32)
+opt = Options(subsampling_factor=2, space_order=16, free_surface=false)
 
 #' Linear Operators
 #' The core idea behind JUDI is to abstract seismic inverse problems in term of linear algebra. In its simplest form, seismic inversion can be formulated as
@@ -119,10 +121,7 @@ dobs = Pr*F*adjoint(Ps)*q
 
 #' Plot the shot record
 fig = figure()
-imshow(dobs.data[1], vmin=-1, vmax=1, cmap="PuOr", extent=[xrec[1], xrec[end], timeD/1000, 0], aspect="auto")
-xlabel("Receiver position (m)")
-ylabel("Time (s)")
-title("Synthetic data")
+plot_sdata(dobs[1]; new_fig=false, name="Synthetic data", cmap=dcmap)
 display(fig)
 
 #' Because we have abstracted the linear algebra, we can solve the adjoint wave-equation as well 
@@ -152,19 +151,13 @@ rtm = adjoint(J)*dD
 
 #' We show the linearized data.
 fig = figure()
-imshow(dD.data[1], vmin=-1, vmax=1, cmap="PuOr", extent=[xrec[1], xrec[end], timeD/1000, 0], aspect="auto")
-xlabel("Receiver position (m)")
-ylabel("Time (s)")
-title("Linearized data")
+plot_sdata(dobs[1]; new_fig=false, name="Linearized data", cmap=dcmap)
 display(fig)
 
 
 #' And the RTM image
 fig = figure()
-imshow(rtm', vmin=-1e2, vmax=1e2, cmap="Greys", extent=[0, (n[1]-1)*d[1], (n[2]-1)*d[2], 0 ], aspect="auto")
-xlabel("Lateral position(m)")
-ylabel("Depth (m)")
-title("RTM image")
+plot_simage(rtm'; new_fig=false, name="RTM image", cmap=imcmap)
 display(fig)
 
 #' ## Inversion utility functions
@@ -185,10 +178,7 @@ f, g = fwi_objective(model0, q, dobs; options=opt)
 
 #' Plot gradient
 fig = figure()
-imshow(g', vmin=-1e2, vmax=1e2, cmap="Greys", extent=[0, (n[1]-1)*d[1], (n[2]-1)*d[2], 0 ], aspect="auto")
-xlabel("Lateral position(m)")
-ylabel("Depth (m)")
-title("FWI gradient")
+plot_simage(g'; new_fig=false, name="FWI gradient", cmap=imcmap)
 display(fig)
 
 
@@ -199,17 +189,11 @@ fjn, gjn = lsrtm_objective(model0, q, dobs, dm; nlind=true, options=opt)
 
 #' Plot gradients
 fig = figure()
-imshow(gj', vmin=-1, vmax=1, cmap="Greys", extent=[0, (n[1]-1)*d[1], (n[2]-1)*d[2], 0 ], aspect="auto")
-xlabel("Lateral position(m)")
-ylabel("Depth (m)")
-title("LSRTM gradient")
+plot_simage(gj'; new_fig=false, name="LSRTM gradient", cmap=imcmap, cbar=true)
 display(fig)
 
 fig = figure()
-imshow(gjn', vmin=-1, vmax=1, cmap="Greys", extent=[0, (n[1]-1)*d[1], (n[2]-1)*d[2], 0 ], aspect="auto")
-xlabel("Lateral position(m)")
-ylabel("Depth (m)")
-title("LSRTM gradient with background data substracted")
+plot_simage(gjn'; new_fig=false, name="LSRTM gradient with background data substracted", cmap=imcmap, cbar=true)
 display(fig)
 
 #' By extension, lsrtm_objective is the same as fwi_objecive when `dm` is zero
@@ -218,13 +202,10 @@ display(fig)
 #' OMP_NUM_THREADS=1 (no parllelism) produces the exact (difference == 0) same result
 #' gjn2 == g
 fjn2, gjn2 = lsrtm_objective(model0, q, dobs, 0f0.*dm; nlind=true, options=opt)
-fig = figure()
 
 #' Plot gradient
-imshow(gjn2', vmin=-1e2, vmax=1e2, cmap="Greys", extent=[0, (n[1]-1)*d[1], (n[2]-1)*d[2], 0 ], aspect="auto")
-xlabel("Lateral position(m)")
-ylabel("Depth (m)")
-title("LSRTM gradient with zero perturbation")
+fig = figure()
+plot_simage(gjn2'; new_fig=false, name="LSRTM gradient with zero perturbation", cmap=imcmap)
 display(fig)
 
 
@@ -236,15 +217,9 @@ f, gmf = twri_objective(model0, q, dobs, nothing; options=Options(frequencies=[[
 
 #' Plot gradients
 fig = figure()
-imshow(gm', vmin=-1, vmax=1, cmap="Greys", extent=[0, (n[1]-1)*d[1], (n[2]-1)*d[2], 0 ], aspect="auto")
-xlabel("Lateral position(m)")
-ylabel("Depth (m)")
-title("TWRI gradient w.r.t m")
+plot_simage(gm'; new_fig=false, name="TWRI gradient w.r.t m", cmap=imcmap)
 display(fig)
 
 fig = figure()
-imshow(gy.data[1], vmin=-1e2, vmax=1e2, cmap="PuOr", extent=[xrec[1], xrec[end], timeD/1000, 0], aspect="auto")
-xlabel("Receiver position (m)")
-ylabel("Time (s)")
-title("TWRI gradient w.r.t y")
+plot_sdata(gy[1]; new_fig=false, name="TWRI gradient w.r.t y", cmap=dcmap)
 display(fig)

src/TimeModeling/Modeling/misfit_fg.jl

@@ -11,6 +11,7 @@ function _multi_src_fg(model_full::AbstractModel, source::Dtypes, dObs::Dtypes,
                       data_precon=nothing, model_precon=LinearAlgebra.I)
     GC.gc(true)
     devito.clear_cache()
+
     # assert this is for single source LSRTM
     @assert source.nsrc == 1 "Multiple sources are used in a single-source fwi_objective"
     @assert dObs.nsrc == 1 "Multiple-source data is used in a single-source fwi_objective"    
@@ -63,6 +64,7 @@ function _multi_src_fg(model_full::AbstractModel, source::Dtypes, dObs::Dtypes,
 
     length(options.frequencies) == 0 ? freqs = nothing : freqs = options.frequencies
     IT = illum ? (PyArray, PyArray) : (PyObject, PyObject)
+
     @juditime "Python call to J_adjoint" begin
         argout = rlock_pycall(ac."J_adjoint", Tuple{Float32, PyArray, IT...}, modelPy,
                 src_coords, qIn, rec_coords, dObserved, t_sub=options.subsampling_factor,

src/TimeModeling/Modeling/propagation.jl

@@ -44,9 +44,9 @@ function run_and_reduce(func, ::Nothing, nsrc, arg_func::Function; kw=nothing)
             kw_loc = isnothing(kw) ? Dict() : kw(i)
             next = func(arg_func(i)...; kw_loc...)
         end
-	@juditime "Reducting $(func) for src $(i)" begin
+        @juditime "Reducting $(func) for src $(i)" begin
             single_reduce!(out, next)
-	end
+        end
     end
     out
 end
@@ -114,7 +114,7 @@ function multi_src_fg!(G, model, q, dobs, dm; options=Options(), ms_func=multi_s
     kw_func = i -> Dict(:illum=> illum, Dict(k => kw_i(v, i) for (k, v) in kw)...)
     # Distribute source
     res = run_and_reduce(ms_func, pool, nsrc, arg_func; kw=kw_func)
-    f, g = update_illum(res, model, :adjoint_born)
+    res = update_illum(res, model, :adjoint_born)
     f, g = as_vec(res, Val(options.return_array))
     G .+= g
     return f

src/TimeModeling/Modeling/twri_objective.jl

@@ -71,7 +71,7 @@ function _twri_objective(model_full::AbstractModel, source::judiVector, dObs::ju
     dtComp = convert(Float32, modelPy."critical_dt")
 
     # Extrapolate input data to computational grid
-    qIn = time_resample(source.data[1], source.geometry, dtComp)
+    qIn = time_resample(make_input(source), source.geometry, dtComp)
     dObserved = time_resample(make_input(dObs), dObs.geometry, dtComp)
 
     if isnothing(y)

src/TimeModeling/Utils/auxiliaryFunctions.jl

@@ -26,8 +26,7 @@ Parameters
 function devito_model(model::MT, options::JUDIOptions, dm) where {MT<:AbstractModel}
     pad = pad_sizes(model, options)
     # Set up Python model structure
-    dm = pad_array(dm, pad)
-    physpar = Dict((n, isa(v, PhysicalParameter) ? pad_array(v.data, pad) : v) for (n, v) in _params(model))
+    physpar = Dict((n, isa(v, PhysicalParameter) ? v.data : v) for (n, v) in _params(model))
 
     modelPy = rlock_pycall(pm."Model", PyObject, origin(model), spacing(model), size(model), fs=options.free_surface,
                    nbl=nbl(model), space_order=options.space_order, dt=options.dt_comp, dm=dm;

src/pysource/fields.py

@@ -3,14 +3,16 @@
 from devito import (TimeFunction, ConditionalDimension, Function,
                     DefaultDimension, Dimension, VectorTimeFunction,
                     TensorTimeFunction)
-from devito.data.allocators import ExternalAllocator
+from devito.builtins import initialize_function
 from devito.tools import as_tuple
 
 try:
     import devitopro as dvp  # noqa
 except ImportError:
     import devito as dvp  # noqa
 
+from utils import compression_mode
+
 
 def wavefield(model, space_order, save=False, nt=None, fw=True, name='', t_sub=1):
     """
@@ -141,7 +143,7 @@ def wavefield_subsampled(model, u, nt, t_sub, space_order=8):
     for wf in as_tuple(u):
         usave = dvp.TimeFunction(name='us_%s' % wf.name, grid=model.grid, time_order=2,
                                  space_order=space_order, time_dim=time_subsampled,
-                                 save=nsave)
+                                 save=nsave, compression=compression_mode())
         wf_s.append(usave)
     return wf_s
 
@@ -174,9 +176,8 @@ def lr_src_fields(model, weight, wavelet, empty_w=False, rec=False):
     if empty_w:
         source_weight = Function(name='%s_weight' % wn, grid=model.grid, space_order=0)
     else:
-        source_weight = Function(name='%s_weight' % wn, grid=model.grid, space_order=0,
-                                 allocator=ExternalAllocator(weight),
-                                 initializer=lambda x: None)
+        source_weight = Function(name='%s_weight' % wn, grid=model.grid, space_order=0)
+        initialize_function(source_weight, weight, 0)
     return source_weight, wavelett
 
 

src/pysource/fields_exprs.py

@@ -1,8 +1,7 @@
 import numpy as np
 
-from devito import Inc, Eq, ConditionalDimension, cos, sin, sign
+from devito import Inc, Eq, ConditionalDimension, cos, sin
 from devito.tools import as_tuple
-from devito.symbolics import retrieve_functions, INT, retrieve_derivatives
 
 from fields import (wavefield_subsampled, lr_src_fields, fourier_modes, norm_holder,
                     illumination)
@@ -31,7 +30,7 @@ def save_subsampled(model, u, nt, t_sub, space_order=8):
         return []
     eq_save = []
     for (wfs, wf) in zip(wf_s, as_tuple(u)):
-        eq_save.append(Eq(wfs, wf))
+        eq_save.append(Eq(wfs, wf, subdomain=model.physical))
     return eq_save
 
 
@@ -102,7 +101,7 @@ def extended_rec(model, wavelet, v):
     return [Inc(ws, model.grid.time_dim.spacing * wf * wt)]
 
 
-def freesurface(model, eq, mfuncs=None, fd_only=False):
+def freesurface(model, fields):
     """
     Generate the stencil that mirrors the field as a free surface modeling for
     the acoustic wave equation
@@ -113,42 +112,20 @@ def freesurface(model, eq, mfuncs=None, fd_only=False):
         Physical model
     eq: Eq or List of Eq
         Equation to apply mirror to
-    mfuncs: List of Functions
-        List of functions to mirror (default=None). Mirrors all functions if not provided
     """
-    fs_eq = []
-    fsdomain = model.fsdomain
-    zfs = model.grid.subdomains['fsdomain'].dimensions[-1]
-    z = zfs.parent
-
-    for eq_i in eq:
-        for p in eq_i._flatten:
-            # Add modulo replacements to to rhs
-            lhs, rhs = p.lhs, p.rhs
-
-            Dz = {d for d in retrieve_derivatives(rhs) if z in d.dims}
-            # Remove inner duplicate
-            Dz = Dz - {d for D in Dz for d in retrieve_derivatives(D.expr) if z in d.dims}
-            Dz = {d: d._eval_at(lhs).evaluate for d in Dz}
-
-            funcs = {f for f in retrieve_functions(Dz.values())}
-
-            if mfuncs:
-                funcs = {f for f in funcs if f.function in mfuncs}
-
-            mapper = {}
-            for f in funcs:
-                zind = f.indices[-1]
-                if (zind - z).as_coeff_Mul()[0] < 0:
-                    s = sign(zind._subs(z.spacing, 1))
-                    mapper.update({f: s * f._subs(zind, INT(abs(zind)))})
-
-            dzmapper = {d: v.subs(mapper) for d, v in Dz.items()}
-            fs_eq.append(p.func(lhs, rhs.subs(dzmapper), subdomain=fsdomain))
-            if not fd_only:
-                fs_eq.append(p.func(lhs._subs(z, 0), 0, subdomain=fsdomain))
-
-    return fs_eq
+    eqs = []
+
+    z = model.grid.dimensions[-1]
+    zfs = model.fs_dim
+
+    for u in as_tuple(fields):
+        if u == 0:
+            continue
+
+        eqs.extend([Eq(u._subs(z, - zfs), - u._subs(z, zfs)),
+                    Eq(u._subs(z, 0), 0)])
+
+    return eqs
 
 
 def otf_dft(u, freq, dt, factor=None):

src/pysource/interface.py

@@ -466,6 +466,7 @@ def J_adjoint_standard(model, src_coords, wavelet, rec_coords, recin,
 
     if return_obj:
         return f, g.data, getattr(Iu, "data", None), getattr(Iv, "data", None)
+
     return g.data, getattr(Iu, "data", None), getattr(Iv, "data", None)
 
 
@@ -618,6 +619,7 @@ def wri_func(model, src_coords, wavelet, rec_coords, recin, yin,
             grady = c2 * recin - rcv.data[:]
             if norm_y != 0:
                 grady -= np.abs(eps) * ydat / norm_y
+            grady = grady.astype(model.dtype)
 
         # Correcting for reduced gradient
         if not grad_corr:

src/pysource/kernels.py

@@ -59,13 +59,13 @@ def acoustic_kernel(model, u, fw=True, q=None):
     damp = model.damp
     stencil = solve(wmr * u.dt2 + damp * udt - ulaplace - q, u_n)
 
-    if 'nofsdomain' in model.grid.subdomains:
-        pde = [Eq(u_n, stencil, subdomain=model.physical)]
-        pde += freesurface(model, pde, (u,))
+    pde = [Eq(u_n, stencil, subdomain=model.physical)]
+    if model.fs:
+        fseq = freesurface(model, u_n)
     else:
-        pde = [Eq(u_n, stencil, subdomain=model.physical)]
+        fseq = []
 
-    return pde
+    return pde, fseq
 
 
 def SLS_2nd_order(model, p, fw=True, q=None, f0=0.015):
@@ -130,7 +130,7 @@ def SLS_2nd_order(model, p, fw=True, q=None, f0=0.015):
 
         u_p = Eq(p.backward, solve(pde_p, p.backward))
 
-        return [u_r, u_p]
+        return [u_r, u_p], []
 
 
 def tti_kernel(model, u1, u2, fw=True, q=None):
@@ -163,20 +163,15 @@ def tti_kernel(model, u1, u2, fw=True, q=None):
     stencilp = solve(wmr * u1.dt2 + damp * udt1 - H0 - q[0], u1_n)
     stencilr = solve(wmr * u2.dt2 + damp * udt2 - H1 - q[1], u2_n)
 
-    if 'nofsdomain' in model.grid.subdomains:
-        # Water at free surface, no anisotrpy
-        acout_ttip = Eq(u1_n, stencilp.subs(model.zero_thomsen))
-        acout_ttir = Eq(u2_n, stencilr.subs(model.zero_thomsen))
-        pdea = freesurface(model, (acout_ttip, acout_ttir), (u1, u2))
-        # Standard PDE in subsurface
-        first_stencil = Eq(u1_n, stencilp, subdomain=model.physical)
-        second_stencil = Eq(u2_n, stencilr, subdomain=model.physical)
+    first_stencil = Eq(u1_n, stencilp, subdomain=model.physical)
+    second_stencil = Eq(u2_n, stencilr, subdomain=model.physical)
+
+    if model.fs:
+        fseq = freesurface(model, (u1_n, u2_n))
     else:
-        pdea = []
-        first_stencil = Eq(u1_n, stencilp, subdomain=model.physical)
-        second_stencil = Eq(u2_n, stencilr, subdomain=model.physical)
+        fseq = []
 
-    return [first_stencil, second_stencil] + pdea
+    return [first_stencil, second_stencil], fseq
 
 
 def elastic_kernel(model, v, tau, fw=True, q=None):
@@ -224,4 +219,9 @@ def elastic_kernel(model, v, tau, fw=True, q=None):
     u_t = Eq(tau.forward, model.damp * solve(eq_tau, tau.forward),
              subdomain=model.physical)
 
-    return [u_v, u_t]
+    if model.fs:
+        fseq = freesurface(model, tau.forward.diagonal())
+    else:
+        fseq = []
+
+    return [u_v, u_t], fseq