Skip to content

Commit

Permalink
fix dm input
Browse files Browse the repository at this point in the history
  • Loading branch information
mloubout committed Oct 25, 2024
1 parent 0eb6c4b commit 572258a
Show file tree
Hide file tree
Showing 22 changed files with 136 additions and 107 deletions.
9 changes: 6 additions & 3 deletions .github/workflows/ci-examples.yml
Original file line number Diff line number Diff line change
Expand Up @@ -42,6 +42,7 @@ jobs:
DEVITO_LANGUAGE: "openmp"
DEVITO_LOGGING: "ERROR"
OMP_NUM_THREADS: 1
JULIA_NUM_THREADS: 1
NITER: 2

strategy:
Expand All @@ -60,21 +61,23 @@ jobs:
version: ${{ matrix.version }}
arch: x64

- uses: julia-actions/cache@v2

- name: Set julia python
run: |
python3 -m pip install -U pip
python3 -m pip install "matplotlib<3.9" seiscm colorcet
echo "JULIA_PYTHONCALL_EXE=$(which python3)" >> $GITHUB_ENV
echo "JULIA_CONDAPKG_BACKEND=\"Null\"" >> $GITHUB_ENV
python3 -m pip install devito[tests,extras]@git+https://github.com/devitocodes/devito.git
JULIA_PYTHONCALL_EXE=$(which python3) JULIA_CONDAPKG_BACKEND="Null" julia -e 'using Pkg;Pkg.add(["PythonCall", "PyPlot", "SlimPlotting"]);Pkg.build("PythonCall")'
echo "PYTHON=$(which python3)" >> $GITHUB_ENV
echo "PYCALL_JL_RUNTIME_PYTHON=$(which python3)" >> $GITHUB_ENV
- name: Build JUDI
uses: julia-actions/julia-buildpkg@latest

- name: Install packages
run: |
julia -e 'using Pkg;Pkg.add(["NLopt", "Flux", "JOLI", "Zygote", "IterativeSolvers", "SlimOptim", "HDF5", "SegyIO", "SetIntersectionProjection"])'
julia -e 'using Pkg;Pkg.add(["SlimPlotting", "PyPlot", "NLopt", "Flux", "JOLI", "Zygote", "IterativeSolvers", "SlimOptim", "HDF5", "SegyIO", "SetIntersectionProjection"])'
julia -e 'using Pkg; Pkg.develop(PackageSpec(path=pwd()))'
- name: Run examples
Expand Down
9 changes: 6 additions & 3 deletions .github/workflows/ci-judi.yml
Original file line number Diff line number Diff line change
Expand Up @@ -22,6 +22,7 @@ jobs:
DEVITO_ARCH: ${{ matrix.cc }}
DEVITO_LANGUAGE: "openmp"
OMP_NUM_THREADS: 4
JULIA_NUM_THREADS: 1
GROUP: "JUDI"

strategy:
Expand Down Expand Up @@ -49,11 +50,13 @@ jobs:
version: ${{ matrix.version }}
arch: ${{ matrix.arch }}

- uses: julia-actions/cache@v2

- name: Setup clang for osx
if: runner.os == 'macOS'
run: |
brew install llvm libomp
echo "/opt/homebrew/bin:/opt/homebrew/opt/llvm/bin" >> $GITHUB_PATH
echo "/opt/homebrew/opt/llvm/bin" >> $GITHUB_PATH
- name: Set up Python 3.9
uses: actions/setup-python@v5
Expand All @@ -62,10 +65,10 @@ jobs:

- name: Set julia python
run: |
echo "PYTHON=$(which python3)" >> $GITHUB_ENV
echo "PYCALL_JL_RUNTIME_PYTHON=$(which python3)" >> $GITHUB_ENV
echo "JULIA_PYTHONCALL_EXE=$(which python3)" >> $GITHUB_ENV
echo "JULIA_CONDAPKG_BACKEND=\"Null\"" >> $GITHUB_ENV
python3 -m pip install devito[tests,extras]@git+https://github.com/devitocodes/devito.git
JULIA_PYTHONCALL_EXE=$(which python3) JULIA_CONDAPKG_BACKEND="Null" julia -e 'using Pkg;Pkg.add("PythonCall");Pkg.build("PythonCall")'
- name: Build JUDI
uses: julia-actions/julia-buildpkg@latest
Expand Down
9 changes: 6 additions & 3 deletions .github/workflows/ci-op.yml
Original file line number Diff line number Diff line change
Expand Up @@ -23,6 +23,7 @@ jobs:
DEVITO_LANGUAGE: "openmp"
DEVITO_LOGGING: "INFO"
OMP_NUM_THREADS: ${{ matrix.omp }}
JULIA_NUM_THREADS: 1
GROUP: ${{ matrix.op }}

strategy:
Expand Down Expand Up @@ -82,11 +83,13 @@ jobs:
version: ${{ matrix.version }}
arch: ${{ matrix.arch }}

- uses: julia-actions/cache@v2

- name: Setup clang for osx
if: runner.os == 'macOS'
run: |
brew install llvm libomp
echo "/opt/homebrew/bin:/opt/homebrew/opt/llvm/bin" >> $GITHUB_PATH
echo "/opt/homebrew/opt/llvm/bin" >> $GITHUB_PATH
- name: Set up Python 3.9
uses: actions/setup-python@v5
Expand All @@ -95,10 +98,10 @@ jobs:

- name: Set julia python
run: |
echo "PYTHON=$(which python3)" >> $GITHUB_ENV
echo "PYCALL_JL_RUNTIME_PYTHON=$(which python3)" >> $GITHUB_ENV
echo "JULIA_PYTHONCALL_EXE=$(which python3)" >> $GITHUB_ENV
echo "JULIA_CONDAPKG_BACKEND=\"Null\"" >> $GITHUB_ENV
python3 -m pip install devito[tests,extras]@git+https://github.com/devitocodes/devito.git
JULIA_PYTHONCALL_EXE=$(which python3) JULIA_CONDAPKG_BACKEND="Null" julia -e 'using Pkg;Pkg.add("PythonCall");Pkg.build("PythonCall")'
- name: Build JUDI
uses: julia-actions/julia-buildpkg@latest
Expand Down
16 changes: 13 additions & 3 deletions examples/scripts/modeling_basic_2D.jl
Original file line number Diff line number Diff line change
Expand Up @@ -137,7 +137,7 @@ dot1 = dot(q, qad)
# <F x, y>
dot2 = dot(dobs, dobs)
# Compare
@show dot1, dot2, (dot2 - dot2)/(dot1 + dot2)
@show dot1, dot2, (dot1 - dot2)/(dot1 + dot2)

#' # Inversion
#' Our main goal is to provide an inversion framework for seismic inversion. To this end, as shown earlier,
Expand All @@ -151,15 +151,25 @@ rtm = adjoint(J)*dD

#' We show the linearized data.
fig = figure()
plot_sdata(dobs[2]; new_fig=false, name="Linearized data", cmap=dcmap)
plot_sdata(dD[2]; new_fig=false, name="Linearized data", cmap=dcmap)
display(fig)


#' And the RTM image
fig = figure()
plot_simage(rtm'; new_fig=false, name="RTM image", cmap=imcmap)
display(fig)

#' We can easily now again test the adjointness of our operator with the standard dot test. Because we
#' intend to conserve our linear algebra abstraction, `judiVector` implements all the necessary linear
#' algebra functions such as dot product or norm to be used directly.
# <x, J'y>
dot3 = dot(dm, rtm)
# <J x, y>
dot4 = dot(dD, dD)
# Compare
@show dot3, dot4, (dot3 - dot4)/(dot3 + dot4)


#' ## Inversion utility functions
#' We currently introduced the lineaar operators that allow to write seismic modeling and inversion in a high-level, linear algebra way. These linear operators allow the script to closely follow the mathematics and to be readable and understandable.
#'
Expand Down
22 changes: 4 additions & 18 deletions src/TimeModeling/Modeling/losses.jl
Original file line number Diff line number Diff line change
Expand Up @@ -12,19 +12,12 @@ and its derivative w.r.t `x`
`x-y`
"""
function _mse(x::AT, y::AbstractArray{T}) where {T<:Number, AT<:AbstractArray{T}}
f = .5f0 * norm(x - y, 2)^2
function mse(x::AbstractArray{T}, y::AbstractArray{T}) where {T<:Number}
f = T(.5) * norm(x - y, 2)^2
r = x - y
return f, r
end

function mse(x::PyArray{T}, y::AbstractArray{T}) where {T<:Number}
f, r = _mse(x, y)
return f, Py(r).to_numpy()
end

mse(x::Matrix{T}, y::Matrix{T}) where {T<:Number} = _mse(x, y)

"""
studentst(x, y)
Expand All @@ -37,18 +30,11 @@ and its derivative w.r.t x
`(k + 1) * (x - y) / (k + (x - y)^2)`
"""
function _studentst(x::AT, y::AbstractArray{T}; k=T(2)) where {T<:Number, AT<:AbstractArray{T}}
function studentst(x::AbstractArray{T}, y::AbstractArray{T}; k=T(2)) where {T<:Number}
k = convert(T, k)
f = sum(_studentst_loss.(x, y, k))
r = (k + 1) .* (x - y) ./ (k .+ (x - y).^2)
return f::T, r::AT{T}
end

function studentst(x::PyArray{T}, y::AbstractArray{T}; k=T(2)) where {T<:Number}
f, r = _studentst(x, y, k)
return f, Py(r).to_numpy()
return f, r
end

studentst(x::Matrix{T}, y::Matrix{T}; k=T(2)) where {T<:Number} = _studentst(x, y, k)

_studentst_loss(x::T, y::T, k::T) where {T<:Number} = T(1/2) * (k + 1) * log(1 + (x-y)^2 / k)
3 changes: 0 additions & 3 deletions src/TimeModeling/Modeling/misfit_fg.jl
Original file line number Diff line number Diff line change
Expand Up @@ -65,9 +65,6 @@ function _multi_src_fg(model_full::AbstractModel, source::Dtypes, dObs::Dtypes,

length(options.frequencies) == 0 ? freqs = nothing : freqs = options.frequencies

# Make dObserved/Y numpy to avoid indexing issues
dIn = Py(dObserved).to_numpy()

@juditime "Python call to J_adjoint" begin
argout = wrapcall_data(ac.J_adjoint, modelPy, src_coords, qIn, rec_coords, dObserved, t_sub=options.subsampling_factor,
checkpointing=options.optimal_checkpointing,
Expand Down
24 changes: 9 additions & 15 deletions src/TimeModeling/Modeling/python_interface.jl
Original file line number Diff line number Diff line change
Expand Up @@ -66,7 +66,6 @@ function devito_interface(modelPy::Py, srcGeometry::Nothing, srcData::Array, rec
rec_coords = setup_grid(recGeometry, modelPy.shape)

# Devito call
srcData = Py(srcData).to_numpy()
return wrapcall_data(ac.forward_wf_src, modelPy, srcData, rec_coords, fw=fw, f0=options.f0, illum=illum)
end

Expand All @@ -77,7 +76,6 @@ function devito_interface(modelPy::Py, srcGeometry::Nothing, srcData::Array, rec
dtComp = pyconvert(Float32, modelPy.critical_dt)

# Devito call
srcData = Py(srcData).to_numpy()
return wrapcall_data(ac.forward_wf_src_norec, modelPy, srcData, fw=fw, illum=illum)
end

Expand Down Expand Up @@ -115,9 +113,6 @@ function devito_interface(modelPy::Py, srcGeometry::Geometry, srcData::Array, re
rec_coords = setup_grid(recGeometry, modelPy.shape)
length(options.frequencies) == 0 ? freqs = nothing : freqs = options.frequencies

# Make dIn numpy to avoid indexing issues
dIn = Py(dIn).to_numpy()

return wrapcall_data(ac.J_adjoint, modelPy,
src_coords, qIn, rec_coords, dIn, fw=fw, t_sub=options.subsampling_factor,
checkpointing=options.optimal_checkpointing,
Expand All @@ -131,16 +126,16 @@ end
function devito_interface(modelPy::Py, weights::Array, srcData::Array, recGeometry::Geometry, recData::Nothing,
dm::Nothing, options::JUDIOptions, illum::Bool, fw::Bool)
judilog("Pr*$(_op_str(fw))*Pw'*w")
weights = pad_array(reshape(weights, modelPy.shape), modelPy.padsizes; mode=:zeros)
shape = pyconvert(Tuple, modelPy.shape)
weights = pad_array(reshape(weights, shape), modelPy.padsizes; mode=:zeros)
# Interpolate input data to computational grid
dtComp = pyconvert(Float32, modelPy.critical_dt)
qIn = time_resample(srcData, recGeometry, dtComp)

# Set up coordinates with devito dimensions
rec_coords = setup_grid(recGeometry, modelPy.shape)
rec_coords = setup_grid(recGeometry, shape)

# Devito call
weights = Py(weights).to_numpy()
return wrapcall_data(ac.forward_rec_w, modelPy, weights, qIn, rec_coords,
fw=fw, f0=options.f0, illum=illum)
end
Expand All @@ -166,13 +161,14 @@ end
function devito_interface(modelPy::Py, weights::Array, srcData::Array, recGeometry::Geometry, recData::Nothing,
dm::Union{PhysicalParameter, Array}, options::JUDIOptions, illum::Bool, fw::Bool)
judilog("Jw($(_op_str(fw)), q)*dm")
weights = pad_array(reshape(weights, modelPy.shape), modelPy.padsizes; mode=:zeros)
shape = pyconvert(Tuple, modelPy.shape)
weights = pad_array(reshape(weights, shape), modelPy.padsizes; mode=:zeros)
# Interpolate input data to computational grid
dtComp = pyconvert(Float32, modelPy.critical_dt)
qIn = time_resample(srcData, recGeometry, dtComp)

# Set up coordinates with devito dimensions
rec_coords = setup_grid(recGeometry, modelPy.shape)
rec_coords = setup_grid(recGeometry, shape)

# Devito call
return wrapcall_data(ac.born_rec_w, modelPy, weights, qIn, rec_coords,
Expand All @@ -182,20 +178,18 @@ end
# Adjoint Jacobian of extended source modeling: dm = J'*d_lin
function devito_interface(modelPy::Py, weights::Array, srcData::Array, recGeometry::Geometry, recData::Array, dm::Nothing, options::JUDIOptions, illum::Bool, fw::Bool)
judilog("Jw($(_op_str(fw)), q)'*d_lin")
weights = pad_array(reshape(weights, modelPy.shape), modelPy.padsizes; mode=:zeros)
shape = pyconvert(Tuple, modelPy.shape)
weights = pad_array(reshape(weights, shape), modelPy.padsizes; mode=:zeros)
# Interpolate input data to computational grid
dtComp = pyconvert(Float32, modelPy.critical_dt)
qIn = time_resample(srcData, recGeometry, dtComp)
dIn = time_resample(recData, recGeometry, dtComp)
qIn, dIn = _maybe_pad_t0(qIn, recGeometry, dIn, recGeometry, dtComp)

# Set up coordinates with devito dimensions
rec_coords = setup_grid(recGeometry, modelPy.shape)
rec_coords = setup_grid(recGeometry, shape)
length(options.frequencies) == 0 ? freqs = nothing : freqs = options.frequencies

# Make dIn numpy to avoid indexing issues
dIn = Py(dIn).to_numpy()

return wrapcall_data(ac.J_adjoint, modelPy,
nothing, qIn, rec_coords, dIn, fw=fw, t_sub=options.subsampling_factor,
checkpointing=options.optimal_checkpointing,
Expand Down
7 changes: 0 additions & 7 deletions src/TimeModeling/Modeling/twri_objective.jl
Original file line number Diff line number Diff line change
Expand Up @@ -94,13 +94,6 @@ function _twri_objective(model_full::AbstractModel, source::judiVector, dObs::ju
~isempty(options.frequencies) ? freqs = options.frequencies : freqs = nothing
~isempty(options.frequencies) ? (wfilt, freqs) = filter_w(qIn, dtComp, freqs) : wfilt = nothing

# Make dObserved/Y numpy to avoid indexing issues
dObserved = Py(dObserved).to_numpy()
qIn = Py(qIn).to_numpy()
if !isnothing(Y)
Y = Py(Y).to_numpy()
end

argout = wrapcall_data(ac.wri_func, modelPy,
src_coords, qIn, rec_coords, dObserved, Y, t_sub=options.subsampling_factor,
grad=optionswri.params, grad_corr=optionswri.grad_corr, eps=optionswri.eps,
Expand Down
10 changes: 5 additions & 5 deletions src/TimeModeling/Preconditioners/DataPreconditioners.jl
Original file line number Diff line number Diff line change
Expand Up @@ -118,7 +118,7 @@ function muteshot!(shot::AbstractMatrix{T}, rGeom::Geometry, srcGeom::Geometry;
end
end

function muteshot(shot::VecOrMat{T}, srcGeom::Geometry, recGeom::Geometry;
function muteshot(shot::AbstractVecOrMat{T}, srcGeom::Geometry, recGeom::Geometry;
vp=1500, t0=.1, mode=:reflection, taperwidth=floor(Int, 2/t0)) where {T<:Number}
sr = reshape(shot, recGeom)
for s=1:get_nsrc(recGeom)
Expand Down Expand Up @@ -158,7 +158,7 @@ judiFilter(geometry::Geometry, fmin::T, fmax::T) where T = judiFilter(geometry,
judiFilter(geometry::Geometry, fmin::Float32, fmax::Float32) = FrequencyFilter{Float32, fmin, fmax}(n_samples(geometry), geometry)
judiFilter(v::judiVector, fmin, fmax) = judiFilter(v.geometry, fmin, fmax)

function matvec(D::FrequencyFilter{T, fm, FM}, x::VecOrMat{T}) where {T, fm, FM}
function matvec(D::FrequencyFilter{T, fm, FM}, x::AbstractVecOrMat{T}) where {T, fm, FM}
dr = reshape(x, D.recGeom)
for j=1:get_nsrc(D.recGeom)
dri = view(dr, :, :, j)
Expand Down Expand Up @@ -227,7 +227,7 @@ filter_data(Din::judiVector, ::Any; fmin=0, fmax=Inf) = filter_data(Din; fmin=fm
Performs a causal filtering [fmin, fmax] on the input data bases on its sampling rate `dt`.
Automatically perfroms a lowpass if fmin=0 (default)
"""
function filter_data(Din::Matrix{T}, dt_in; fmin=0, fmax=Inf) where T
function filter_data(Din::AbstractMatrix{T}, dt_in; fmin=0, fmax=Inf) where T
out = similar(Din)
filter!(out, Din, dt_in; fmin=T(fmin), fmax=T(fmax))
return out
Expand Down Expand Up @@ -304,7 +304,7 @@ function matvec(D::TimeDifferential{T, K}, x::judiVector{T, AT}) where {T, AT, K
return out
end

function matvec(D::TimeDifferential{T, K}, x::Array{T}) where {T, K}
function matvec(D::TimeDifferential{T, K}, x::AbstractArray{T}) where {T, K}
xr = reshape(x, D.recGeom)
out = similar(xr)
# make omega^K
Expand Down Expand Up @@ -371,7 +371,7 @@ function matvec(D::TimeGain{T, K}, x::judiVector{T, AT}) where {T, AT, K}
return out
end

function matvec(D::TimeGain{T, K}, x::Array{T}) where {T, K}
function matvec(D::TimeGain{T, K}, x::AbstractArray{T}) where {T, K}
xr = reshape(x, D.recGeom)
# make time^K
timek = (D.recGeom.taxis[1]).^K
Expand Down
1 change: 1 addition & 0 deletions src/TimeModeling/Preconditioners/base.jl
Original file line number Diff line number Diff line change
Expand Up @@ -28,6 +28,7 @@ getproperty(J::Preconditioner, s::Symbol) = _get_property(J, Val{s}())
*(J::Preconditioner, ms::judiMultiSourceVector) = matvec(J, ms)
*(J::Preconditioner, ms::PhysicalParameter) = matvec(J, ms)
*(J::Preconditioner, v::VecOrMat{T}) where T = matvec(J, v)
*(J::Preconditioner, v::PyArray{T}) where T = matvec(J, v)

mul!(out::judiMultiSourceVector, J::Preconditioner, ms::judiMultiSourceVector) = copyto!(out, matvec(J, ms))
mul!(out::PhysicalParameter, J::Preconditioner, ms::PhysicalParameter) = copyto!(out, matvec(J, ms))
Expand Down
12 changes: 7 additions & 5 deletions src/TimeModeling/Utils/auxiliaryFunctions.jl
Original file line number Diff line number Diff line change
Expand Up @@ -25,7 +25,8 @@ Parameters
"""
function devito_model(model::MT, options::JUDIOptions, dm) where {MT<:AbstractModel}
# Set up Python model structure
physpar = Dict((n, isa(v, PhysicalParameter) ? Py(v.data).to_numpy() : v) for (n, v) in _params(model))
physpar = Dict((n, isa(v, PhysicalParameter) ? v.data : v) for (n, v) in _params(model))
dm = isnothing(dm) ? nothing : (isa(dm, PhysicalParameter) ? dm.data : dm)

modelPy = pm.Model(origin(model), spacing(model), size(model), fs=options.free_surface,
nbl=nbl(model), space_order=options.space_order, dt=options.dt_comp, dm=dm;
Expand All @@ -52,16 +53,17 @@ function pad_array(m::Array{DT}, nb::NTuple{N, Tuple{Int64, Int64}}; mode::Symbo
n = size(m)
new_size = Tuple([n[i] + sum(nb[i]) for i=1:length(nb)])
Ei = []
for i=1:length(nb)
for i=length(nb):-1:1
left, right = nb[i]
push!(Ei, joExtend(n[i], mode;pad_upper=right, pad_lower=left, RDT=DT, DDT=DT))
end
padded = joKron(Ei...) * PermutedDimsArray(m, length(n):-1:1)[:]
return PyReverseDims(reshape(padded, reverse(new_size)))
padded = joKron(Ei...) * m[:]
return reshape(padded, new_size)
end

pad_array(::Nothing, ::NTuple{N, Tuple{Int64, Int64}}; s::Symbol=:border) where N = nothing
pad_array(m::Number, ::NTuple{N, Tuple{Int64, Int64}}; s::Symbol=:border) where N = m
pad_array(m::Array{DT}, nb::Py; mode::Symbol=:border) where DT = pad_array(m, pyconvert(Tuple, nb), mode=mode)

"""
remove_padding(m, nb; true_adjoint=False)
Expand All @@ -86,7 +88,7 @@ function remove_padding(gradient::AbstractArray{DT}, nb::NTuple{ND, Tuple{Int64,
return out
end

remove_padding(g::PyArray, nb::Py; true_adjoint::Bool=false) = remove_padding(g, pyconvert(Tuple, nb), true_adjoint=true)
remove_padding(g::PyArray, nb::Py; true_adjoint::Bool=false) = remove_padding(g, pyconvert(Tuple, nb), true_adjoint=true_adjoint)

"""
limit_model_to_receiver_area(srcGeometry, recGeometry, model, buffer; pert=nothing)
Expand Down
Loading

0 comments on commit 572258a

Please sign in to comment.