fix dm input

slimgroup · Oct 24, 2024 · b5240a1 · b5240a1
1 parent 0eb6c4b
commit b5240a1
Show file tree

Hide file tree

Showing 21 changed files with 98 additions and 91 deletions.
diff --git a/.github/workflows/ci-examples.yml b/.github/workflows/ci-examples.yml
@@ -42,6 +42,7 @@ jobs:
       DEVITO_LANGUAGE: "openmp"
       DEVITO_LOGGING: "ERROR"
       OMP_NUM_THREADS: 1
+      JULIA_NUM_THREADS: 1
       NITER: 2
 
     strategy:
@@ -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

diff --git a/.github/workflows/ci-judi.yml b/.github/workflows/ci-judi.yml
@@ -22,6 +22,7 @@ jobs:
       DEVITO_ARCH: ${{ matrix.cc }}
       DEVITO_LANGUAGE: "openmp"
       OMP_NUM_THREADS: 4
+      JULIA_NUM_THREADS: 1
       GROUP: "JUDI"
 
     strategy:
@@ -35,7 +36,7 @@ jobs:
 
         include:
           - os: macos-15
-            version: '1'
+            version: 'lts'
             arch: ARM64
             cc: clang
 
@@ -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
@@ -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

diff --git a/.github/workflows/ci-op.yml b/.github/workflows/ci-op.yml
@@ -23,6 +23,7 @@ jobs:
       DEVITO_LANGUAGE: "openmp"
       DEVITO_LOGGING: "INFO"
       OMP_NUM_THREADS: ${{ matrix.omp }}
+      JULIA_NUM_THREADS: 1
       GROUP: ${{ matrix.op }}
 
     strategy:
@@ -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
@@ -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

diff --git a/examples/scripts/modeling_basic_2D.jl b/examples/scripts/modeling_basic_2D.jl
@@ -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,
@@ -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.
 #' 

diff --git a/src/TimeModeling/Modeling/losses.jl b/src/TimeModeling/Modeling/losses.jl
@@ -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)
 
@@ -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)
diff --git a/src/TimeModeling/Modeling/misfit_fg.jl b/src/TimeModeling/Modeling/misfit_fg.jl
@@ -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,

diff --git a/src/TimeModeling/Modeling/python_interface.jl b/src/TimeModeling/Modeling/python_interface.jl
@@ -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
 
@@ -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
 
@@ -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,
@@ -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
@@ -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,
@@ -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,

diff --git a/src/TimeModeling/Modeling/twri_objective.jl b/src/TimeModeling/Modeling/twri_objective.jl
@@ -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,

diff --git a/src/TimeModeling/Preconditioners/DataPreconditioners.jl b/src/TimeModeling/Preconditioners/DataPreconditioners.jl
@@ -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)
@@ -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)
@@ -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
@@ -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
@@ -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

diff --git a/src/TimeModeling/Preconditioners/base.jl b/src/TimeModeling/Preconditioners/base.jl
@@ -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))