slimgroup · mloubout · Mar 28, 2024 · Mar 27, 2024
diff --git a/Project.toml b/Project.toml
@@ -1,7 +1,7 @@
 name = "TimeProbeSeismic"
 uuid = "cd417b76-fc6b-4228-be91-94f5ea77264b"
 authors = ["Mathias Louboutin", "Felix J. Herrmann"]
-version = "1.1.2"
+version = "1.1.3"
 
 [deps]
 JUDI = "f3b833dc-6b2e-5b9c-b940-873ed6319979"
@@ -11,6 +11,6 @@ Reexport = "189a3867-3050-52da-a836-e630ba90ab69"
 Statistics = "10745b16-79ce-11e8-11f9-7d13ad32a3b2"
 
 [compat]
-JUDI = "^3.3.3"
+JUDI = "3.4.1"
 Reexport = "1"
 julia = "1.6"
diff --git a/scripts/compare_vectors.jl b/scripts/compare_vectors.jl
@@ -2,11 +2,13 @@
 # Author: mlouboutin3@gatech.edu
 # Date: February 2021
 #
-using TimeProbeSeismic, PyPlot, LinearAlgebra, HDF5, Images
+using TimeProbeSeismic, PyPlot, LinearAlgebra, HDF5, Images, SlimPlotting
+
+# set_devito_config("log-level", "PERF")
 
 # Load starting model
-~isfile(datadir("models/", "overthrust_model.h5")) && run(`curl -L ftp://slim.gatech.edu/data/SoftwareRelease/WaveformInversion.jl/2DFWI/overthrust_model_2D.h5 --create-dirs -o $(datadir("models", "overthrust_model.h5"))`)
-n, d, o, m0, m = read(h5open(datadir("models", "overthrust_model.h5"), "r"), "n", "d", "o", "m0", "m")
+~isfile(datadir("models/", "overthrust_model.h5")) && run(`curl -L ftp://slim.gatech.edu/data/SoftwareRelease/WaveformInversion.jl/2DFWI/overthrust_model_2D.h5 --create-dirs -o $(datadir("models/", "overthrust_model.h5"))`)
+n, d, o, m0, m = read(h5open(datadir("models/", "overthrust_model.h5"), "r"), "n", "d", "o", "m0", "m")
 
 n = Tuple(n)
 o = Tuple(o)
@@ -33,7 +35,7 @@ dtD = 4f0  # receiver sampling interval [ms]
 recGeometry = Geometry(xrec, yrec, zrec; dt=dtD, t=timeD, nsrc=nsrc)
 
 # Set up source geometry (cell array with source locations for each shot)
-xsrc = div(n[1], 2) * d[1]
+xsrc = div(n[1], 4) * d[1]
 ysrc = 0f0
 zsrc = 12.5f0
 
@@ -45,25 +47,19 @@ f0 = 0.010f0     # kHz
 wavelet = ricker_wavelet(timeD, dtD, f0)
 q = judiVector(srcGeometry, wavelet)
 
-# Set up info structure for linear operators
-ntComp = get_computational_nt(srcGeometry, recGeometry, model)
-info = Info(prod(n), nsrc, ntComp)
-
 ###################################################################################################
-# Write shots as segy files to disk
 opt = Options(space_order=16)
 
 # Setup operators
-F = judiModeling(info, model, srcGeometry, recGeometry; options=opt)
-F0 = judiModeling(info, model0, srcGeometry, recGeometry; options=opt)
+F = judiModeling(model, srcGeometry, recGeometry; options=opt)
+F0 = judiModeling(model0, srcGeometry, recGeometry; options=opt)
 J = judiJacobian(F0, q)
 
 # data
 # Nonlinear modeling
 dobs = F*q
 d0 = F0*q
 residual = d0 - dobs
-δd = J*dm
 
 # gradient
 g = J'*residual
@@ -80,25 +76,25 @@ function fourierJ(F0, q, ps)
     return J
 end
 
-ge = Array{Any}(undef, 3, 8)
+ge = Array{Any}(undef, 4, 8)
 
-for (p, mode) in enumerate([:Fourier, :QR, :Rademacher])
+for (p, mode) in enumerate([:Fourier, :QR, :Rademacher, :hutchpp])
     for ps=1:8
         Jp = mode == :Fourier ? fourierJ(F0, q, 2^ps) : judiJacobian(F0, q, 2^ps, dobs; mode=mode)
         ge[p, ps] = Jp'*residual
     end
 end
 
-for (p, mode) in enumerate([:Fourier, :QR, :Rademacher])
+for (p, mode) in enumerate([:Fourier, :QR, :Rademacher, :hutchpp])
     clip = maximum(g)/10
 
     figure(figsize=(15, 10))
     subplot(331)
-    imshow(g', cmap="seismic", vmin=-clip, vmax=clip, aspect="auto")
+    plot_simage(g'; cmap="seismic", new_fig=false)
     title("Reference")
     for ps=1:8
         subplot(3,3,ps+1)
-        imshow(ge[p, ps]', cmap="seismic", vmin=-clip, vmax=clip, aspect="auto")
+        plot_simage(ge[p, ps]'; cmap="seismic", new_fig=false)
         title("$(mode) r=$(2^ps)")
     end
     tight_layout()
@@ -108,13 +104,14 @@ for (p, mode) in enumerate([:Fourier, :QR, :Rademacher])
 
     figure(figsize=(15, 10))
     subplot(331)
-    imshow(g', cmap="seismic", vmin=-clip, vmax=clip, aspect="auto")
+    plot_simage(g'; cmap="seismic", new_fig=false)
     title("Reference")
     for ps=1:8
         subplot(3,3,ps+1)
-        imshow(g' - ge[p, ps]', cmap="seismic", vmin=-clip, vmax=clip, aspect="auto")
+        plot_simage(g' - ge[p, ps]'; cmap="seismic", new_fig=false)
         title("$(mode) Difference r=$(2^ps)")
     end
     tight_layout()
     savefig(plotsdir("geo-pros-paper/", "Error-$(mode).pdf"), bbox_inches=:tight, dpi=150)
 end
+
diff --git a/scripts/overthrust_probing_vectors.jl b/scripts/overthrust_probing_vectors.jl
@@ -61,23 +61,27 @@ title(L"$d_{obs} d_{obs}^\top$")
 
 timea = 0:4f0:3000f0
 
-fig, axs = subplots(5, 3, sharex=:col, sharey=:row, gridspec_kw=Dict(:hspace=> 0, :wspace=> 0))
+fig, axs = subplots(5, 4, sharex=:col, sharey=:row, gridspec_kw=Dict(:hspace=> 0, :wspace=> 0))
 for (i, ps)=enumerate([4, 16, 32, 64, 256])
     @show i
     # QR
-    Q = qr_data(dobs.data[1], ps)
+    Q = qr_data(dobs.data[1], ps; mode=:QR)
     axs[i, 1].imshow(Q*Q',vmin=-.05, vmax=.05, cmap="seismic", aspect="auto")
     i == 1 && axs[i, 1].set_title(L"$\mathbf{Q}\mathbf{Q}^\top$")
     axs[i, 1].set_ylabel("r=$(ps)", rotation=0, labelpad=20)
+    # hutchpp
+    Q = qr_data(dobs.data[1], ps; mode=:hutchpp)
+    axs[i, 2].imshow(Q*Q',vmin=-.05, vmax=.05, cmap="seismic", aspect="auto")
+    i == 1 && axs[i, 2].set_title(L"$\mathbf{Q}\mathbf{Q}^\top$")
     # gaussian
-    E = rand([-1f0, 1f0], length(timea), ps)/sqrt(ps)
-    axs[i, 2].imshow(E*E',vmin=-.25, vmax=.25, cmap="seismic", aspect="auto")
-    i == 1 && axs[i, 2].set_title(L"$\mathbf{Z}\mathbf{Z}^\top$")
+    Q = qr_data(dobs.data[1], ps; mode=:Gaussian)
+    axs[i, 3].imshow(Q*Q',vmin=-.25, vmax=.25, cmap="seismic", aspect="auto")
+    i == 1 && axs[i, 3].set_title(L"$\mathbf{Z}\mathbf{Z}^\top$")
     # DFT
     freq = select_frequencies(q_dist; fmin=0.002, fmax=0.04, nf=ps)
     F = exp.(-2*im*pi*timea*freq')/sqrt(ps)
-    axs[i, 3].imshow(real.(F*F'),vmin=-.5, vmax=.5, cmap="seismic", aspect="auto")
-    i == 1 && axs[i, 3].set_title(L"$\mathbf{F}\mathbf{F}^\top$")
+    axs[i, 4].imshow(real.(F*F'),vmin=-.5, vmax=.5, cmap="seismic", aspect="auto")
+    i == 1 && axs[i, 4].set_title(L"$\mathbf{F}\mathbf{F}^\top$")
 end
 plt.setp(plt.gcf().get_axes(), xticks=[], yticks=[])
 

diff --git a/src/TimeProbeSeismic.jl b/src/TimeProbeSeismic.jl
@@ -34,6 +34,7 @@ function __init__()
     copy!(si, pyimport("sensitivity"))
     copy!(ut, pyimport("utils"))
     copy!(np, pyimport("numpy"))
+    set!(dv."configuration", "autopadding", false)
 end
 
 # Propagators

diff --git a/src/interface.jl b/src/interface.jl
@@ -1,8 +1,8 @@
 
 function forward(model::AbstractModel, q::judiVector, dobs::judiVector; options=Options(), ps=16, modelPy=nothing, mode=:QR)
-    dt_Comp = get_dt(model)
     # Python model
     isnothing(modelPy) && (modelPy = devito_model(model, options))
+    dt_Comp = convert(Float32, modelPy."critical_dt")
     # Interpolate input data to computational grid
     q_data = time_resample(make_input(q), q.geometry, dt_Comp)
     d_data = time_resample(make_input(dobs), dobs.geometry, dt_Comp)
@@ -21,9 +21,9 @@ function forward(model::AbstractModel, q::judiVector, dobs::judiVector; options=
 end
 
 function born(model::AbstractModel, q::judiVector, dobs::judiVector, dm; options=Options(), ps=16, modelPy=nothing)
-    dt_Comp = get_dt(model)
     # Python model
     isnothing(modelPy) && (modelPy = devito_model(model, options; dm=dm))
+    dt_Comp = convert(Float32, modelPy."critical_dt")
 
     # Interpolate input data to computational grid
     q_data = time_resample(make_input(q), q.geometry, dt_Comp)
@@ -46,9 +46,10 @@ end
 
 function backprop(model::AbstractModel, residual::judiVector, Q::Array{Float32}, eu::PyObject;
                  options=Options(), ps=16, modelPy=nothin, offsets=0f0, pe=nothing)
-    dt_Comp = get_dt(model)
     # Python mode
     isnothing(modelPy) && (modelPy = devito_model(model, options))
+    dt_Comp = convert(Float32, modelPy."critical_dt")
+
     # Interpolate input data to computational grid
     d_data = time_resample(make_input(residual), residual.geometry, dt_Comp)
 

diff --git a/src/judi.jl b/src/judi.jl
@@ -1,30 +1,30 @@
-time_resample(q::judiVector, dt) = begin @assert q.nsrc == 1 ;time_resample(make_input(q), q.geometry.dt[1], dt) end
-
 #########################
 struct judiJacobianP{D, O, FT} <: judiAbstractJacobian{D, O, FT}
     m::AbstractSize
     n::AbstractSize
     F::FT
     q::judiMultiSourceVector
     r::Integer
-    mode::Symbol
+    probing_mode::Symbol
     dobs::judiVector
     offsets::Union{D, Vector{D}}
 end
 
+_accepted_modes = (:QR, :Rademacher, :Gaussian, :hutchpp)
+
 # Jacobian with probing
 function judiJacobian(F::judiComposedPropagator{D, O}, q::judiMultiSourceVector, r::Integer, dobs::judiMultiSourceVector;
                       options=nothing, mode=:QR, offsets=0f0) where {D, O}
-    mode ∈ [:QR, :Rademacher, :Gaussian] || throw(ArgumentError("Probing vector mode unrecognized, must be `:QR` or `:Rademacher`"))
+    mode ∈ _accepted_modes || throw(ArgumentError("Probing vector mode unrecognized, must be in $(_accepted_modes)"))
     update!(F.F.options, options)
     return judiJacobianP{D, :born, typeof(F)}(F.m, space(F.model.n), F, q, r, mode, dobs, asvec(offsets))
 end
 
 asvec(x::T) where T<:Number = x
 asvec(x) = collect(x)
 
-adjoint(J::judiJacobianP{D, O, FT}) where {D, O, FT} = judiJacobianP{D, adjoint(O), FT}(J.n, J.m, J.F, J.q, J.r, J.mode, J.dobs, J.offsets)
-getindex(J::judiJacobianP{D, O, FT}, i) where {D, O, FT} = judiJacobianP{D, O, FT}(J.m[i], J.n[i], J.F[i], J.q[i], J.r, J.mode, J.dobs[i], J.offsets)
+adjoint(J::judiJacobianP{D, O, FT}) where {D, O, FT} = judiJacobianP{D, adjoint(O), FT}(J.n, J.m, J.F, J.q, J.r, J.probing_mode, J.dobs, J.offsets)
+getindex(J::judiJacobianP{D, O, FT}, i) where {D, O, FT} = judiJacobianP{D, O, FT}(J.m[i], J.n[i], J.F[i], J.q[i], J.r, J.probing_mode, J.dobs[i], J.offsets)
 
 process_input_data(::judiJacobianP{D, :born, FT}, q::dmType{D}) where {D<:Number, FT} = q
 
@@ -56,7 +56,7 @@ function propagate(J::judiJacobianP{D, :adjoint_born, FT}, residual::AbstractArr
 
     model_loc = get_model(J.q.geometry, residual.geometry, J.model, J.options)
     modelPy = devito_model(model_loc, J.options)
-    _, Q, eu = forward(model_loc, J.q, J.dobs; ps=J.r, options=J.options, modelPy=modelPy, mode=J.mode)
+    _, Q, eu = forward(model_loc, J.q, J.dobs; ps=J.r, options=J.options, modelPy=modelPy, mode=J.probing_mode)
     ge = backprop(model_loc, residual, Q, eu; options=J.options, modelPy=modelPy, offsets=J.offsets)
     J.options.limit_m && (ge = extend_gradient(J.model, model_loc, ge))
 
@@ -72,12 +72,15 @@ function propagate(J::judiJacobianP{D, :adjoint_born, FT}, residual::AbstractArr
 end
 
 fwi_objective(model::MTypes, q::Dtypes, dobs::Dtypes, r::Integer; options=Options(), offsets=0f0) =
-    fwi_objective(model, q, dobs; options=options, r=r, offsets=offsets)
+    fwi_objective(model, q, dobs; options=options, r=r, offsets=offsets, func=multi_src_fg_r)
 
 lsrtm_objective(model::MTypes, q::Dtypes, dobs::Dtypes, dm::dmType, r::Integer; options=Options(), nlind=false, offsets=0f0) =
-    lsrtm_objective(model, q, dobs, dm; options=options, nlind=nlind, r=r, offsets=offsets)
+    lsrtm_objective(model, q, dobs, dm; options=options, nlind=nlind, r=r, offsets=offsets, func=multi_src_fg_r)
 
-function multi_src_fg(model::AbstractModel, q::judiVector, dobs::judiVector, dm, options::JUDIOptions, nlind::Bool, lin::Bool, r::Integer, offsets)
+
+function multi_src_fg_r(model::AbstractModel, q::judiVector, dobs::judiVector, dm, options::JUDIOptions;
+                      nlind::Bool=false, lin::Bool=false, misfit::Function=mse, illum::Bool=false, r::Integer=32,
+                      data_precon=nothing, model_precon=LinearAlgebra.I)
     q.geometry = Geometry(q.geometry)
     dobs.geometry = Geometry(dobs.geometry)
 
@@ -92,4 +95,4 @@ function multi_src_fg(model::AbstractModel, q::judiVector, dobs::judiVector, dm,
     ge = backprop(model_loc, residual, Q, eu; options=options, modelPy=modelPy, offsets=offsets, pe=eu.indices[1])
     options.limit_m && (ge = extend_gradient(model, model_loc, ge))
     return Ref{Float32}(.5f0*norm(residual)^2), PhysicalParameter(ge, model.d, model.o)
-end
+end
diff --git a/src/propagators.jl b/src/propagators.jl
@@ -1,3 +1,16 @@
+p_dims = Dict()
+
+# Basic caching of probing dim
+function pdim(ne::Integer)
+    if ne in keys(p_dims)
+        p_e = p_dims[ne]
+    else
+        p_e = dv.CustomDimension(name="p_e", 0, ne-1, ne)
+        p_dims[ne] = p_e
+    end
+    return p_e
+end
+
 # Forward propagation
 function forward(model::PyObject, src_coords::Array{Float32}, rcv_coords::Array{Float32},
                  wavelet::Array{Float32}, e::Array{Float32}, space_order::Integer=8, isic::Bool=false)
@@ -95,18 +108,18 @@ end
 
 function time_probe(e::Array{Float32, 2}, wf::PyObject; fw=true, isic=false, pe=nothing)
     ne = size(e, 2)
-    p_e = dv.DefaultDimension(name="p_e", default_value=ne)
+    p_e = pdim(ne)
     # sub_time
     t_sub = wf.grid.time_dim
     s = fw ? t_sub.spacing : 1
     # Probing vector
     nt = size(e, 1)
     q = dv.TimeFunction(name="Q", grid=wf.grid, dimensions=(t_sub, p_e),
-                        shape=(nt, ne), time_order=0, initializer=e)
+                        shape=(nt, ne), time_order=0, initializer=e, space_order=0)
     # Probed output
     if isnothing(pe)
-        pe = dv.Function(name="$(wf.name)e", grid=wf.grid, dimensions=(wf.grid.dimensions..., p_e),
-                        shape=(wf.grid.shape..., ne), space_order=wf.space_order)
+        pe = dv.Function(name="$(wf.name)e", grid=wf.grid, dimensions=(p_e, wf.grid.dimensions...),
+                        shape=(ne, wf.grid.shape...), space_order=wf.space_order)
     end
     probing = [dv.Eq(pe, pe + s*q*ic(wf, fw, isic))]
     return pe, probing
@@ -115,18 +128,18 @@ end
 
 function time_probe(e::Array{Float32, 2}, wf::Tuple{PyCall.PyObject, PyCall.PyObject}; fw=true, isic=false, pe=nothing)
     ne = size(e, 2)
-    p_e = dv.DefaultDimension(name="p_e", default_value=ne)
+    p_e = pdim(ne)
     # sub_time
     t_sub = wf[1].grid.time_dim
     s = fw ? t_sub.spacing : 1
     # Probing vector
     nt = size(e, 1)
     q = dv.TimeFunction(name="Q", grid=wf[1].grid, dimensions=(t_sub, p_e),
-                        shape=(nt, ne), time_order=0, initializer=e)
+                        shape=(nt, ne), time_order=0, initializer=e, space_order=0)
     # Probed output
     if isnothing(pe)
-        pe = dv.Function(name="$(wf[1].name)e", grid=wf[1].grid, dimensions=(wf[1].grid.dimensions..., p_e),
-                            shape=(wf[1].grid.shape..., ne), space_order=wf[1].space_order)
+        pe = dv.Function(name="$(wf[1].name)e", grid=wf[1].grid, dimensions=(p_e, wf[1].grid.dimensions...),
+                            shape=(ne, wf[1].grid.shape...), space_order=wf[1].space_order)
     end
     probing = [dv.Eq(pe, pe + s*q*ic(wf, fw, isic))]
     return pe, probing
@@ -146,7 +159,6 @@ end
 function combine(eu::PyObject, ev::PyObject, offsets::Number, isic::Bool=false)
     @assert offsets == 0
     g = dv.Function(name="ge", grid=eu.grid, space_order=0)
-    ev = ev._subs(ev.indices[end], eu.indices[end])
     eq = isic ? (eu * ev.laplace + si.inner_grad(eu, ev)) : eu * ev
     op = dv.Operator(dv.Inc(g, -eq), subs=eu.grid.spacing_map)
     op()
@@ -165,7 +177,6 @@ function combine(eu::PyObject, ev::PyObject, offsets::Vector{Int64}, isic::Bool=
     # IC with space shifts
     g = dv.Function(name="ge", grid=eu.grid, shape=(2*nh+1, eu.grid.shape...),
                     dimensions=(offs, eu.grid.dimensions...), space_order=0)
-    ev = ev._subs(ev.indices[end], eu.indices[end])
     x = eu.grid.dimensions[1]
 
     ev = ev._subs(x, x+oh)

diff --git a/src/utils.jl b/src/utils.jl
@@ -3,13 +3,31 @@ typedict(x::T) where {T} = Dict(fn=>getfield(x, fn) for fn ∈ fieldnames(T))
 
 function qr_data(d::Array{Float32,2}, ps::Integer; seed=nothing, ts=0, mode=:QR)
     !isnothing(seed) && Random.seed!(seed)
-    mode == :Gaussian && (return randn(Float32, size(d, 1), ps) ./ Float32(sqrt(ps)))
     S = rand([-1f0, 1f0], size(d, 1), ps)
-    mode == :Rademacher && (return S ./ Float32(sqrt(ps)))
-    ts > 0 && (d = circshift(d, (-ts, 0)))
-    AS = d * (d' * S) 
-    Q, _ = qr(AS)
-    return Matrix(Q)
+    if mode == :Gaussian
+        Q = randn(Float32, size(d, 1), ps) ./ Float32(sqrt(ps))
+    elseif mode == :Rademacher
+        return S ./ Float32(sqrt(ps))
+    elseif mode == :QR
+        ts > 0 && (d = circshift(d, (-ts, 0)))
+        AS = d * (d' * S) 
+        Q, _ = qr(AS)
+        return Matrix(Q)
+    elseif mode == :hutchpp
+        ts > 0 && (d = circshift(d, (-ts, 0)))
+        # Split S
+        m = 3 * ps / 2
+        mid = div(ps, 2)
+        S1, S2 = S[:, 1:mid], S[:, mid+1:end]
+        AS = d * (d' * S2)
+        Q, _ = qr(AS)
+        Q = Matrix(Q)
+        G = sqrt(3 / m) * (S1 - Q * (Q'* S1))   
+        Qhutch = hcat(Q, G)
+        return convert(Matrix{Float32}, Qhutch)
+    else
+        throw(ArgumentError("Probing vector mode unrecognized: $(mode), must be `:QR`, `:Rademacher`, `:Gaussian` or `:hutchpp`"))
+    end
 end
 
 simil(x, y) = dot(x, y)/(norm(x)*norm(y))