fix parallel illum

slimgroup · Feb 1, 2024 · b3f3772 · b3f3772
1 parent 34067f0
commit b3f3772
Show file tree

Hide file tree

Showing 6 changed files with 42 additions and 40 deletions.
diff --git a/src/TimeModeling/Modeling/misfit_fg.jl b/src/TimeModeling/Modeling/misfit_fg.jl
@@ -1,8 +1,8 @@
 
 export fwi_objective, lsrtm_objective, fwi_objective!, lsrtm_objective!
 
-function multi_src_fg(model_full::AbstractModel, source::judiVector, dObs::judiVector, dm, options::JUDIOptions, nlind::Bool, lin::Bool,
-                      misfit::Function)
+function multi_src_fg(model_full::AbstractModel, source::judiVector, dObs::judiVector, dm, options::JUDIOptions,
+                      nlind::Bool, lin::Bool, misfit::Function, illum::Bool)
     # Setup time-domain linear or nonlinear foward and adjoint modeling and interface to devito
     GC.gc(true)
     devito.clear_cache()
@@ -14,9 +14,6 @@ function multi_src_fg(model_full::AbstractModel, source::judiVector, dObs::judiV
     dObs.geometry = Geometry(dObs.geometry)
     source.geometry = Geometry(source.geometry)
 
-    # Compute illumination ?
-    illum = compute_illum(model_full, :adjoint_born)
-
     # Limit model to area with sources/receivers
     if options.limit_m == true
         @juditime "Limit model to geometry" begin

diff --git a/src/TimeModeling/Modeling/propagation.jl b/src/TimeModeling/Modeling/propagation.jl
@@ -4,9 +4,9 @@
 Base propagation interfaces that calls the devito `mode` propagator (forward/adjoint/..)
 with `q` as a source. The return type is infered from `F`.
 """
-function propagate(F::judiPropagator{T, O}, q::AbstractArray{T}) where {T, O}
+function propagate(F::judiPropagator{T, O}, q::AbstractArray{T}, illum::Bool) where {T, O}
     srcGeometry, srcData, recGeometry, recData, dm = make_input(F, q)
-    return time_modeling(F.model, srcGeometry, srcData, recGeometry, recData, dm, O, F.options, _prop_fw(F))
+    return time_modeling(F.model, srcGeometry, srcData, recGeometry, recData, dm, O, F.options, _prop_fw(F), illum)
 end
 
 propagate(t::Tuple{judiPropagator, AbstractArray}) = propagate(t[1], t[2])
@@ -80,7 +80,8 @@ function multi_src_propagate(F::judiPropagator{T, O}, q::AbstractArray{T}) where
     # Number of sources and init result
     nsrc = get_nsrc(F, q)
     pool = _worker_pool()
-    arg_func = i -> (F[i], src_i(F, q, i))
+    illum = compute_illum(F.model, O)
+    arg_func = i -> (F[i], src_i(F, q, i), illum)
     # Distribute source
     res = run_and_reduce(propagate, pool, nsrc, arg_func)
     res = update_illum(res, F)
@@ -101,8 +102,9 @@ function multi_src_fg!(G, model, q, dobs, dm; options=Options(), kw...)
     # Number of sources and init result
     nsrc = try q.nsrc catch; dobs.nsrc end
     pool = _worker_pool()
+    illum = compute_illum(model, :adjoint_born)
     # Distribute source
-    arg_func = i -> (model, q[i], dobs[i], dm, options[i], values(kw)...)
+    arg_func = i -> (model, q[i], dobs[i], dm, options[i], values(kw)..., illum)
     # Distribute source
     res = run_and_reduce(multi_src_fg, pool, nsrc, arg_func)
     f, g = update_illum(res, model, :adjoint_born)

diff --git a/src/TimeModeling/Modeling/time_modeling_serial.jl b/src/TimeModeling/Modeling/time_modeling_serial.jl
@@ -8,7 +8,7 @@ PhysOrNot = Union{PhysicalParameter, Array, Nothing}
 # Setup time-domain linear or nonlinear foward and adjoint modeling and interface to devito
 function time_modeling(model_full::AbstractModel, srcGeometry::GeomOrNot, srcData::ArrayOrNot,
                        recGeometry::GeomOrNot, recData::ArrayOrNot, dm::PhysOrNot,
-                       op::Symbol, options::JUDIOptions, fw::Bool)
+                       op::Symbol, options::JUDIOptions, fw::Bool, illum::Bool)
     GC.gc(true)
     devito.clear_cache()
 
@@ -21,9 +21,6 @@ function time_modeling(model_full::AbstractModel, srcGeometry::GeomOrNot, srcDat
         return judiVector(recGeometry, zeros(Float32, recGeometry.nt[1], length(recGeometry.xloc[1])))
     end
 
-    # Compute illumination ?
-    illum = compute_illum(model_full, op)
-
     # limit model to area with sources/receivers
     if options.limit_m == true
         @juditime "Limit model to geometry" begin

diff --git a/src/TimeModeling/Preconditioners/ModelPreconditioners.jl b/src/TimeModeling/Preconditioners/ModelPreconditioners.jl
@@ -185,6 +185,8 @@ end
 function set_val(I::judiIllumination{T, M, K, R}, mode, v) where {T, M, K, R}
     key = mode ∈ [:forward, :born] ? "u" : "v"
     if key in keys(I.illums)
+        # For safety when propagation doesn't reach all the model
+        v[v .== 0] .= 1f0
         combine!(I.illums[key], v)
     end
 end

diff --git a/src/TimeModeling/Types/GeometryStructure.jl b/src/TimeModeling/Types/GeometryStructure.jl
@@ -242,6 +242,7 @@ function timings_from_segy(data, dt=nothing, t=nothing, t0=nothing)
     return [range(t0[j], step=dtCell[j], stop=tCell[j]) for j=1:nsrc]
 end
 
+segy_t0(b::Vector{SeisCon}, i) = segy_t0(b[i], 1)
 segy_t0(b::SeisBlock, i) = segy_t0(b)
 segy_t0(b::SeisCon, i) = segy_t0(b.blocks[i])
 segy_t0(b::SeisBlock) = segy_t0(b.fileheader.bfh)
@@ -556,7 +557,7 @@ Merge all the sub-geometries `1:get_nsrc(Geometry)` into a single supershot geom
 """
 function super_shot_geometry(G::Geometry{T}) where T
     as_set = coords_from_set(as_coord_set(G))
-    return GeometryIC{T}(as_set..., [G.t[1]])
+    return GeometryIC{T}(as_set..., [G.taxis[1]])
 end
 
 

diff --git a/test/test_geometry.jl b/test/test_geometry.jl
@@ -23,12 +23,13 @@ datapath = joinpath(dirname(pathof(JUDI)))*"/../data/"
         geometry_t =  Geometry(xsrc, ysrc, zsrc; dt=2, t=1000)
         @test geometry_t == geometry
 
-        @test isequal(typeof(geometry.xloc), Array{Array{Float32,1}, 1})
-        @test isequal(typeof(geometry.yloc), Array{Array{Float32,1}, 1})
-        @test isequal(typeof(geometry.zloc), Array{Array{Float32,1}, 1})
-        @test isequal(typeof(geometry.nt), Array{Integer, 1})
-        @test isequal(typeof(geometry.dt), Array{Float32, 1})
-        @test isequal(typeof(geometry.t), Array{Float32, 1})
+        @test isa(geometry.xloc, Vector{Array{Float32,1}})
+        @test isa(geometry.yloc, Vector{Array{Float32,1}})
+        @test isa(geometry.zloc, Vector{Array{Float32,1}})
+        @test isa(geometry.nt, Vector{Int})
+        @test isa(geometry.dt, Vector{Float32})
+        @test isa(geometry.t, Vector{Float32})
+        @test isa(geometry.taxis, Vector{<:StepRangeLen{Float32}})
 
         # Constructor if coordinates are not passed as a cell arrays
         xsrc = range(100f0, stop=1100f0, length=2)[1:nsrc]
@@ -37,20 +38,22 @@ datapath = joinpath(dirname(pathof(JUDI)))*"/../data/"
 
         geometry = Geometry(xsrc, ysrc, zsrc; dt=4f0, t=1000f0, nsrc=nsrc)
 
-        @test isequal(typeof(geometry.xloc), Array{Array{Float32,1}, 1})
-        @test isequal(typeof(geometry.yloc), Array{Array{Float32,1}, 1})
-        @test isequal(typeof(geometry.zloc), Array{Array{Float32,1}, 1})
-        @test isequal(typeof(geometry.nt), Array{Integer, 1})
-        @test isequal(typeof(geometry.dt), Array{Float32, 1})
-        @test isequal(typeof(geometry.t), Array{Float32, 1})
-
+        @test isa(geometry.xloc, Vector{Array{Float32,1}})
+        @test isa(geometry.yloc, Vector{Array{Float32,1}})
+        @test isa(geometry.zloc, Vector{Array{Float32,1}})
+        @test isa(nt(geometry), Vector{Int})
+        @test isa(dt(geometry), Vector{Float32})
+        @test isa(t(geometry), Vector{Float32})
+        @test isa(t0(geometry), Vector{Float32})
+        @test isa(geometry.taxis, Vector{<:StepRangeLen{Float32}})
+
         # Set up source geometry object from in-core data container
         block = segy_read(datapath*"unit_test_shot_records_$(nsrc).segy"; warn_user=false)
         src_geometry = Geometry(block; key="source", segy_depth_key="SourceSurfaceElevation")
         rec_geometry = Geometry(block; key="receiver", segy_depth_key="RecGroupElevation")
 
-        @test isequal(typeof(src_geometry), GeometryIC{Float32})
-        @test isequal(typeof(rec_geometry), GeometryIC{Float32})
+        @test isa(src_geometry, GeometryIC{Float32})
+        @test isa(rec_geometry, GeometryIC{Float32})
         @test isequal(get_header(block, "SourceSurfaceElevation")[1], src_geometry.zloc[1][1])
         @test isequal(get_header(block, "RecGroupElevation")[1], rec_geometry.zloc[1][1])
         @test isequal(get_header(block, "SourceX")[1], src_geometry.xloc[1][1])
@@ -61,8 +64,8 @@ datapath = joinpath(dirname(pathof(JUDI)))*"/../data/"
         src_geometry = Geometry(container; key="source", segy_depth_key="SourceSurfaceElevation")
         rec_geometry = Geometry(container; key="receiver", segy_depth_key="RecGroupElevation")
 
-        @test isequal(typeof(src_geometry), GeometryOOC{Float32})
-        @test isequal(typeof(rec_geometry), GeometryOOC{Float32})
+        @test isa(src_geometry, GeometryOOC{Float32})
+        @test isa(rec_geometry, GeometryOOC{Float32})
         @test isequal(src_geometry.key, "source")
         @test isequal(rec_geometry.key, "receiver")
         @test isequal(src_geometry.segy_depth_key, "SourceSurfaceElevation")
@@ -78,8 +81,8 @@ datapath = joinpath(dirname(pathof(JUDI)))*"/../data/"
         src_geometry = Geometry(container_cell; key="source", segy_depth_key="SourceSurfaceElevation")
         rec_geometry = Geometry(container_cell; key="receiver", segy_depth_key="RecGroupElevation")
 
-        @test isequal(typeof(src_geometry), GeometryOOC{Float32})
-        @test isequal(typeof(rec_geometry), GeometryOOC{Float32})
+        @test isa(src_geometry, GeometryOOC{Float32})
+        @test isa(rec_geometry, GeometryOOC{Float32})
         @test isequal(src_geometry.key, "source")
         @test isequal(rec_geometry.key, "receiver")
         @test isequal(src_geometry.segy_depth_key, "SourceSurfaceElevation")
@@ -90,34 +93,34 @@ datapath = joinpath(dirname(pathof(JUDI)))*"/../data/"
         src_geometry_ic = Geometry(src_geometry)
         rec_geometry_ic = Geometry(rec_geometry)
 
-        @test isequal(typeof(src_geometry_ic), GeometryIC{Float32})
-        @test isequal(typeof(rec_geometry_ic), GeometryIC{Float32})
+        @test isa(src_geometry_ic, GeometryIC{Float32})
+        @test isa(rec_geometry_ic, GeometryIC{Float32})
         @test isequal(get_header(block, "SourceSurfaceElevation")[1], src_geometry_ic.zloc[1][1])
         @test isequal(get_header(block, "RecGroupElevation")[1], rec_geometry_ic.zloc[1][1])
         @test isequal(get_header(block, "SourceX")[1], src_geometry_ic.xloc[1][1])
         @test isequal(get_header(block, "GroupX")[1], rec_geometry_ic.xloc[1][1])
 
         # Subsample in-core geometry structure
         src_geometry_sub = subsample(src_geometry_ic, 1)
-        @test isequal(typeof(src_geometry_sub), GeometryIC{Float32})
+        @test isa(src_geometry_sub, GeometryIC{Float32})
         @test isequal(length(src_geometry_sub.xloc), 1)
         src_geometry_sub = subsample(src_geometry_ic, 1:1)
-        @test isequal(typeof(src_geometry_sub), GeometryIC{Float32})
+        @test isa(src_geometry_sub, GeometryIC{Float32})
         @test isequal(length(src_geometry_sub.xloc), 1)
 
         inds = nsrc > 1 ? (1:nsrc) : 1
         src_geometry_sub = subsample(src_geometry_ic, inds)
-        @test isequal(typeof(src_geometry_sub), GeometryIC{Float32})
+        @test isa(src_geometry_sub, GeometryIC{Float32})
         @test isequal(length(src_geometry_sub.xloc), nsrc)
 
         # Subsample out-of-core geometry structure
         src_geometry_sub = subsample(src_geometry, 1)
-        @test isequal(typeof(src_geometry_sub), GeometryOOC{Float32})
+        @test isa(src_geometry_sub, GeometryOOC{Float32})
         @test isequal(length(src_geometry_sub.dt), 1)
         @test isequal(src_geometry_sub.segy_depth_key, "SourceSurfaceElevation")
 
         src_geometry_sub = subsample(src_geometry, inds)
-        @test isequal(typeof(src_geometry_sub), GeometryOOC{Float32})
+        @test isa(src_geometry_sub, GeometryOOC{Float32})
         @test isequal(length(src_geometry_sub.dt), nsrc)
         @test isequal(src_geometry_sub.segy_depth_key, "SourceSurfaceElevation")