Merge pull request #183 from slimgroup/elastic

Elastic forward

.github/workflows/ci-examples.yml

@@ -29,7 +29,7 @@ jobs:
         uses: actions/checkout@v3
 
       - id: set-matrix
-        run: echo "::set-output name=matrix::$(ls examples/scripts/*.jl examples/machine-learning/*.jl | xargs -n 1 | jq -R -s -c 'split("\n")[:-1]')"
+        run: echo "matrix=$(ls examples/scripts/*.jl examples/machine-learning/*.jl | xargs -n 1 | jq -R -s -c 'split("\n")[:-1]')" >> $GITHUB_OUTPUT
         shell: bash
 
   run-examples:

.github/workflows/ci-judi.yml

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

.github/workflows/ci-op.yml

@@ -31,7 +31,7 @@ jobs:
       matrix:
         os: [ubuntu-latest]
         op: ["ISO_OP", "ISO_OP_FS", "TTI_OP", "TTI_OP_FS", "BASICS"]
-        version: ['1.8']
+        version: ['1']
         omp: [2]
 
         include:
@@ -51,7 +51,7 @@ jobs:
             omp: 1
 
           - os: ubuntu-latest
-            version: '1.7'
+            version: '1.9'
             op: "VISCO_AC_OP"
             omp: 2
 

.gitignore

@@ -1,3 +1,5 @@
 *.jl.cov
 *.jl.*.cov
 *.jl.mem
+
+/.luarc.json

Project.toml

@@ -1,7 +1,7 @@
 name = "JUDI"
 uuid = "f3b833dc-6b2e-5b9c-b940-873ed6319979"
 authors = ["Philipp Witte, Mathias Louboutin"]
-version = "3.2.3"
+version = "3.3.0"
 
 [deps]
 ChainRulesCore = "d360d2e6-b24c-11e9-a2a3-2a2ae2dbcce4"
@@ -13,7 +13,6 @@ JOLI = "bb331ad6-a1cf-11e9-23da-9bcb53c69f6f"
 LinearAlgebra = "37e2e46d-f89d-539d-b4ee-838fcccc9c8e"
 PyCall = "438e738f-606a-5dbb-bf0a-cddfbfd45ab0"
 Random = "9a3f8284-a2c9-5f02-9a11-845980a1fd5c"
-Reexport = "189a3867-3050-52da-a836-e630ba90ab69"
 Requires = "ae029012-a4dd-5104-9daa-d747884805df"
 SegyIO = "157a0f19-4d44-4de5-a0d0-07e2f0ac4dfa"
 
@@ -26,11 +25,10 @@ FFTW = "1"
 Flux = "0.12, 0.13"
 JOLI = "0.7, 0.8"
 PyCall = "1.18, 1.90, 1.91, 1.62"
-Reexport = "0.2, 1"
 Requires = "1"
 SegyIO = "0.7.7, 0.8.3"
 TimerOutputs = "0.5"
-julia = "1"
+julia = "1.6"
 
 [extras]
 Aqua = "4c88cf16-eb10-579e-8560-4a9242c79595"

examples/machine-learning/example_extended_source_cnn.jl

@@ -29,7 +29,7 @@ zrec = range(50f0, stop=50f0, length=nxrec)
 
 # receiver sampling and recording time
 time = 1000f0   # receiver recording time [ms]
-dt = .75f0    # receiver sampling interval [ms]
+dt = 1f0    # receiver sampling interval [ms]
 
 # Set up receiver structure
 recGeometry = Geometry(xrec, yrec, zrec; dt=dt, t=time, nsrc=nsrc)

examples/scripts/fwi_example_ADloss.jl

@@ -61,15 +61,15 @@ function objective_function(x)
     grad = .125f0*grad/maximum(abs.(grad))  # scale for line search
 
     global count; count+= 1
-    return fval, vec(grad.data)
+    return fval, grad
 end
 
 # Bound projection
 proj(x) = reshape(median([vec(mmin) vec(x) vec(mmax)]; dims=2),model0.n)
 
 # FWI with SPG
 options = spg_options(verbose=3, maxIter=fevals, memory=3)
-sol = spg(objective_function, vec(model0.m), proj, options)
+sol = spg(objective_function, model0.m, proj, options)
 
 # Plot result
 imshow(reshape(sqrt.(1f0 ./ sol.x), model0.n)', extent=[0, 10, 3, 0])

examples/scripts/fwi_example_minConf.jl

@@ -48,15 +48,15 @@ function objective_function(x)
     grad = .125f0*grad/maximum(abs.(grad))  # scale for line search
 
     global count; count+= 1
-    return fval, vec(grad.data)
+    return fval, grad
 end
 
 # Bound projection
 proj(x) = reshape(median([vec(mmin) vec(x) vec(mmax)]; dims=2),model0.n)
 
 # FWI with SPG
 options = spg_options(verbose=3, maxIter=fevals, memory=3)
-sol = spg(objective_function, vec(model0.m), proj, options)
+sol = spg(objective_function, model0.m, proj, options)
 
 # Plot result
 imshow(reshape(sqrt.(1f0 ./ sol.x), model0.n)', extent=[0, 10, 3, 0])

examples/scripts/fwi_example_studentst.jl

@@ -61,18 +61,16 @@ function objective_function(x, misfit=mse)
 end
 
 # Bound projection
-proj(x) = reshape(median([vec(mmin) vec(x) vec(mmax)]; dims=2),model0.n)
-
+proj(x) = reshape(median([vec(mmin) vec(x) vec(mmax)]; dims=2), size(x))
 
 # Compare l2 with students t
 ϕmse = x->objective_function(x)
 ϕst = x->objective_function(x, studentst)
 
 # FWI with SPG
-m0 = vec(model0.m)
 options = spg_options(verbose=3, maxIter=fevals, memory=3)
-solmse = spg(ϕmse, m0, proj, options)
-solst = spg(ϕst, m0, proj, options)
+solmse = spg(ϕmse, vec(m0), proj, options)
+solst = spg(ϕst, vec(m0), proj, options)
 
 # Plot result
 figure(figsize=(10, 10))

examples/scripts/modeling_basic_elastic.jl

@@ -0,0 +1,64 @@
+# Example for basic 3D modeling:
+# The receiver positions and the source wavelets are the same for each of the four experiments.
+# Author: Philipp Witte, pwitte@eos.ubc.ca
+# Date: January 2017
+#
+
+using JUDI
+
+# Set up model structure
+n = (120, 100)    # (x,y,z) or (x,z)
+d = (10., 10.)
+o = (0., 0.)
+
+# Velocity [km/s]
+v = ones(Float32,n) .+ .5f0
+v[:, Int(round(end/2)):end] .= 4.0f0
+rho = ones(Float32, n)
+vs = zeros(Float32,n)
+vs[:, Int(round(end/2)):end] .= 2f0
+
+# Slowness squared [s^2/km^2]
+m = (1f0 ./ v).^2
+# Setup model structure
+nsrc = 1
+model = Model(n, d, o, m; rho=rho, vs=vs)
+
+# Set up 3D receiver geometry by defining one receiver vector in each x and y direction
+nxrec = 120
+nyrec = 100
+xrec = range(50f0, stop=1150f0, length=nxrec)
+yrec = 0f0
+zrec = range(0f0, stop=0f0, length=nxrec)
+
+# receiver sampling and recording time
+timeR = 1500f0   # receiver recording time [ms]
+dtR = 4f0    # receiver sampling interval
+
+# Set up receiver structure
+recGeometry = Geometry(xrec, yrec, zrec; dt=dtR, t=timeR, nsrc=nsrc)
+
+# Set up source geometry (cell array with source locations for each shot)
+xsrc = 600f0
+ysrc = 0f0
+zsrc = 0f0
+
+# source sampling and number of time steps
+timeS = 1500f0   # source length in [ms]
+dtS = 2f0    # source sampling interval
+
+# Set up source structure
+srcGeometry = Geometry(xsrc, ysrc, zsrc; dt=dtS, t=timeS)
+
+# setup wavelet
+f0 = 0.01f0
+wavelet = ricker_wavelet(timeS, dtS, f0)
+
+###################################################################################################
+
+# Setup operators
+F = judiModeling(model, srcGeometry, recGeometry)
+q = judiVector(srcGeometry, wavelet)
+
+# Nonlinear modeling
+dobs = F*q

examples/scripts/modeling_extended_source_3D.jl

@@ -37,7 +37,7 @@ zrec = 50f0
 (xrec, yrec, zrec) = setup_3D_grid(xrec, yrec, zrec)
 
 # receiver sampling and recording time
-time = 50f0   # receiver recording time [ms]
+time = 60f0   # receiver recording time [ms]
 dt = 4f0    # receiver sampling interval [ms]
 
 # Set up receiver structure

examples/scripts/modeling_medical_2D.jl

@@ -28,7 +28,7 @@ zrec = range(d[1], stop=d[1], length=nxrec)
 
 # receiver sampling and recording time
 timeR = 250f0   # receiver recording time [ms]
-dtR = 0.2f0    # receiver sampling interval [ms]
+dtR = 0.25f0    # receiver sampling interval [ms]
 
 # Set up receiver structure
 recGeometry = Geometry(xrec, yrec, zrec; dt=dtR, t=timeR, nsrc=nsrc)
@@ -40,7 +40,7 @@ zsrc = d[1]
 
 # source sampling and number of time steps
 timeS = 250f0     # ms
-dtS =  0.2f0    # ms
+dtS =  0.25f0    # ms
 
 # Set up source structure
 srcGeometry = Geometry(xsrc, ysrc, zsrc; dt=dtS, t=timeS)

examples/scripts/modeling_wavefields_2D.jl

@@ -88,7 +88,8 @@ qnew = Ps*adjoint(F)*u
 dtComp = get_dt(model)
 ntComp = get_computational_nt(q.geometry, model)
 u0 = zeros(Float32, ntComp[1], model.n[1] + 2*model.nb, model.n[2] + 2*model.nb)
-u0[:, 100, 45] = wavelet = -ricker_wavelet(timeS, dtComp, f0)
+wavelet = -ricker_wavelet(timeS, dtComp, f0)
+u0[1:length(wavelet), 100, 45] .= wavelet
 uf = judiWavefield(dtComp, u0)
 dobs2 = Pr*F*uf # same as dobs
 

src/JUDI.jl

@@ -54,6 +54,7 @@ import ChainRulesCore: rrule
 # Set python paths
 const pm = PyNULL()
 const ac = PyNULL()
+const pyut = PyNULL()
 
 # Create a lock for pycall FOR THREAD/TASK SAFETY
 # See discussion at
@@ -80,11 +81,12 @@ end
 
 _TFuture = Future
 _verbose = false
+_devices = []
 
 # Utility for data loading
 JUDI_DATA = joinpath(JUDIPATH, "../data")
-ftp_data(ftp::String, name::String) = run(`curl -L $(ftp) --create-dirs -o $(JUDI.JUDI_DATA)/$(name)`)
-ftp_data(ftp::String) = run(`curl -L $(ftp) --create-dirs -o $(JUDI.JUDI_DATA)/$(split(ftp, "/")[end])`)
+ftp_data(ftp::String, name::String) = Base.Downloads().download("$(ftp)/$(name)", "$(JUDI.JUDI_DATA)/$(name)")
+ftp_data(ftp::String) = Base.Downloads().download(ftp, "$(JUDI.JUDI_DATA)/$(split(ftp, "/")[end])")
 
 # Some usefull types
 const RangeOrVec = Union{AbstractRange, Vector}
@@ -116,12 +118,14 @@ function __init__()
     pushfirst!(PyVector(pyimport("sys")."path"), joinpath(JUDIPATH, "pysource"))
     copy!(pm, pyimport("models"))
     copy!(ac, pyimport("interface"))
+    copy!(pyut, pyimport("utils"))
     # Initialize lock at session start
     PYLOCK[] = ReentrantLock()
 
     if get(ENV, "DEVITO_PLATFORM", "") == "nvidiaX"
         @info "Initializing openacc/openmp offloading"
         devito_model(Model((21, 21, 21), (10., 10., 10.), (0., 0., 0.), randn(Float32, 21, 21, 21)), Options())
+        global _devices = parse.(Int, get(ENV, "CUDA_VISIBLE_DEVICES", "-1"))
     end
 
     @require Zygote="e88e6eb3-aa80-5325-afca-941959d7151f" begin
@@ -136,5 +140,4 @@ function __init__()
     end
 end
 
-
 end

src/TimeModeling/LinearOperators/basics.jl

@@ -41,7 +41,6 @@ Base.merge!(S1::AbstractSize, S2::AbstractSize) = merge!(S1.dims, S2.dims)
 
 Base.repr(S::AbstractSize) = "($(join(keys(S.dims), " * ")))"
 
-similar(x::judiMultiSourceVector, ::Type{ET}, dims::AbstractSize) where ET = similar(x, ET)
 similar(x::AbstractVector, ::Type{ET}, dims::AbstractSize) where ET = similar(x, ET, Int(dims))
 
 # Update size
@@ -68,5 +67,3 @@ time_space_src(nsrc::Integer, nt) = AbstractSize((:src, :time, :x, :y, :z), (nsr
 space_src(nsrc::Integer) = AbstractSize((:src, :x, :y, :z), (nsrc, 0, 0, 0))
 
 time_src(nsrc::Integer, nt) = AbstractSize((:src, :time), (nsrc, nt))
-
-rec_space(G::Geometry) = AbstractSize((:src, :time, :rec), (get_nsrc(G), G.nt, G.nrec))

src/TimeModeling/LinearOperators/callable.jl

@@ -0,0 +1,53 @@
+
+#############################################################################
+(F::judiModeling{D, O})(m::AbstractModel{D, N}) where {D, O, N} = judiModeling(m; options=F.options)
+(FS::judiDataSourceModeling{D, O})(F::judiModeling{D, O}) where {D, O} = judiDataSourceModeling{D, O}(FS.rInterpolation, F, FS.qInjection)
+(FS::judiDataModeling{D, O})(F::judiModeling{D, O}) where {D, O} = judiDataModeling{D, O}(FS.rInterpolation, F)
+(FS::judiPointSourceModeling{D, O})(F::judiModeling{D, O}) where {D, O} = judiPointSourceModeling{D, O}(F, FS.qInjection)
+
+(J::judiJacobian{D, O, FT})(F::judiModeling{D, O}) where {D, O, FT} = judiJacobian(J.F(F), J.q)
+
+for FT in [judiPointSourceModeling, judiDataModeling, judiDataSourceModeling, judiJacobian]
+    @eval begin
+        (F::$(FT))(m::AbstractModel) = F(F.F(m))
+    end
+end
+
+function (F::judiPropagator)(;kwargs...)
+    Fl = deepcopy(F)
+    for (k, v) in kwargs
+        k in _mparams(Fl.model) && getfield(Fl.model, k) .= v
+    end
+    Fl
+end
+
+function (F::judiPropagator)(m, q)
+    Fm = F(;m=m)
+    _track_illum(F.model, Fm.model)
+    return Fm*as_src(q)
+end
+
+function (F::judiPropagator)(m::AbstractArray)
+    @info "Assuming m to be squared slowness for $(typeof(F))"
+    return F(;m=m)
+end
+
+(F::judiPropagator)(m::AbstractModel, q) = F(m)*as_src(q)
+
+function (J::judiJacobian{D, O, FT})(q::judiMultiSourceVector) where {D, O, FT}
+    newJ = judiJacobian{D, O, FT}(J.m, J.n, J.F, q)
+    _track_illum(J.model, newJ.model)
+    return newJ
+end
+
+function (J::judiJacobian{D, O, FT})(x::Array{D, N}) where {D, O, FT, N}
+    if length(x) == prod(J.model.n)
+        return J(;m=m)
+    end
+    new_q = _as_src(J.qInjection.op, J.model, x)
+    newJ = judiJacobian{D, O, FT}(J.m, J.n, J.F, new_q)
+    _track_illum(J.model, newJ.model)
+    return newJ
+end
+
+

src/TimeModeling/LinearOperators/lazy.jl

@@ -194,32 +194,32 @@ get_nt(P::Projection) = P.n[:time]
 get_nt(P::jAdjoint{<:Projection}) = P.op.n[:time]
 
 """
-    reshape(x, P::judiProjection, m::Model; with_batch=false)
+    reshape(x, P::judiProjection, m::AbstractModel; with_batch=false)
 
 reshape the input `x` into an ML friendly format `HWCB` using the projection operator and model to infer dimensions sizes.
 """
-function reshape(x::Vector{T}, P::judiProjection{T}, ::Model; with_batch=false) where T
+function reshape(x::Vector{T}, P::judiProjection{T}, ::AbstractModel; with_batch=false) where T
     out = reshape(x, P.geometry)
     out = with_batch ? reshape(out, size(out, 1), size(out, 2), 1, size(out, 3)) : out
     return out
 end
 
-function reshape(x::Vector{T}, P::judiWavelet{T}, m::Model; with_batch=false) where T
+function reshape(x::Vector{T}, P::judiWavelet{T}, m::AbstractModel; with_batch=false) where T
     out = with_batch ? reshape(x, m.n..., 1,get_nsrc(P)) : reshape(x, m.n..., get_nsrc(P))
     return out
 end
 
-function _as_src(P::judiLRWF{T}, model::Model, q::Array{T}) where T
+function _as_src(P::judiLRWF{T}, model::AbstractModel, q::Array{T}) where T
     qcell = process_input_data(q, model, get_nsrc(P))
     return judiWeights{T}(get_nsrc(P), qcell)
 end
 
-function _as_src(P::judiProjection{T}, ::Model, q::Array{T}) where T
+function _as_src(P::judiProjection{T}, ::AbstractModel, q::Array{T}) where T
     qcell = process_input_data(q, P.geometry)
     return judiVector{T, Matrix{T}}(get_nsrc(P), P.geometry, qcell)
 end
 
-_as_src(::judiNoopOperator, ::Model, q::judiMultiSourceVector) = q
+_as_src(::judiNoopOperator, ::AbstractModel, q::judiMultiSourceVector) = q
 
 
 ############################################################################################################################