Merge pull request #151 from slimgroup/illum-again

Illumination and precon

.github/workflows/deploy_doc.yaml

@@ -24,6 +24,11 @@ jobs:
 
       - uses: julia-actions/setup-julia@latest
 
+      - name: Set up Python 3.9
+        uses: actions/setup-python@v4
+        with:
+          python-version: 3.9
+
       - name: Install dependencies
         run: julia --project=docs/ -e 'using Pkg; Pkg.develop(PackageSpec(path=pwd())); Pkg.instantiate()'
 

.github/workflows/flake8.yml

@@ -18,6 +18,11 @@ jobs:
     steps:
     - uses: actions/checkout@v1
 
+    - name: Set up Python 3.9
+      uses: actions/setup-python@v4
+      with:
+        python-version: 3.9
+
     - name: Install dependencies
       run: |
         python -m pip install --upgrade pip

Project.toml

@@ -1,7 +1,7 @@
 name = "JUDI"
 uuid = "f3b833dc-6b2e-5b9c-b940-873ed6319979"
 authors = ["Philipp Witte, Mathias Louboutin"]
-version = "3.1.13"
+version = "3.2.0"
 
 [deps]
 ChainRulesCore = "d360d2e6-b24c-11e9-a2a3-2a2ae2dbcce4"

README.md

@@ -235,6 +235,9 @@ J = judiJacobian(F, q)
 # Right-hand preconditioners (model topmute)
 Mr = judiTopmute(model0.n, 52, 10)	# mute up to grid point 52, with 10 point taper
 
+# Left-hand preconditioners
+Ml = judiDataMute(q.geometry, dD.geometry; t0=.120)	# data topmute starting at 120ms (30 samples)
+
 # Stochastic gradient
 x = zeros(Float32, info.n)	# zero initial guess
 batchsize = 10	# use subset of 10 shots per iteration
@@ -246,11 +249,10 @@ for j=1:niter
 
 	# Select batch and set up left-hand preconditioner
 	i = randperm(dD.nsrc)[1:batchsize]
-	Ml = judiMarineTopmute2D(30, dD[i].geometry)	# data topmute starting at time sample 30
 
 	# Compute residual and gradient
-	r = Ml*J[i]*Mr*x - Ml*dD[i]
-	g = adjoint(Mr)*adjoint(J[i])*adjoint(Ml)*r
+	r = Ml[i]*J[i]*Mr*x - Ml[i]*dD[i]
+	g = adjoint(Mr)*adjoint(J[i])*adjoint(Ml[i])*r
 
 	# Step size and update variable
 	fval[j] = .5f0*norm(r)^2

docs/make.jl

@@ -14,13 +14,14 @@ makedocs(sitename="JUDI documentation",
              "Home" => "index.md",
              "About" => "about.md",
              "Installation" => "installation.md",
+             "Getting Started" => "basics.md",
              "JUDI API" => Any[
                 "Abstract vectors" => "abstract_vectors.md",
                 "Data Structures" => "data_structures.md",
                 "Linear Operators" => "linear_operators.md",
+                "Preconditioners" => "preconditioners.md",
                 "Input/Output" => "io.md",
                 "Helper Functions" => "helper.md"],
-             "Getting Started" => "basics.md",
              "Inversion" => "inversion.md",
              "Tutorials" => map(
                 s -> "tutorials/$(s)",

docs/src/inversion.md

@@ -147,7 +147,9 @@ F = judiModeling(model0, q.geometry, dD.geometry; options=opt)
 J = judiJacobian(F, q)
 
 # Right-hand preconditioners (model topmute)
-Mr = judiTopmute(model0.n, 52, 10)	# mute up to grid point 52, with 10 point taper
+Mr = judiTopmute(model0; taperwidth=10)	# mute up to grid point 52, with 10 point taper
+# Left-hand side preconditioners
+Ml = judiDatMute(q.geometry, dD.geometry; t0=.120)	# data topmute starting at time 120ms
 
 # Stochastic gradient
 x = zeros(Float32, info.n)	# zero initial guess
@@ -160,11 +162,10 @@ for j=1:niter
 
 	# Select batch and set up left-hand preconditioner
 	i = randperm(dD.nsrc)[1:batchsize]
-	Ml = judiMarineTopmute2D(30, dD[i].geometry)	# data topmute starting at time sample 30
 
 	# Compute residual and gradient
-	r = Ml*J[i]*Mr*x - Ml*dD[i]
-	g = adjoint(Mr)*adjoint(J[i])*adjoint(Ml)*r
+	r = Ml[i]*J[i]*Mr*x - Ml[i]*dD[i]
+	g = adjoint(Mr)*adjoint(J[i])*adjoint(Ml[i])*r
 
 	# Step size and update variable
 	fval[j] = .5f0*norm(r)^2

docs/src/preconditioners.md

@@ -1,20 +1,22 @@
 # Seismic Preconditioners
 
-## Model topmute
+JUDI provides a selected number of preconditioners known to be beneficial to FWI and RTM. We welcome additional preconditionners from the community. Additionnaly, any JOLI operator can be used as a preconditiner in conbination with JUDI operator thanks to the fundamental interface between JUDI and JOLI.
 
-Create a linear operator for a 2D model topmute, i.e. for muting the water column:
-
-```julia
-Mr = judiTopmute(n, mute_start, length)
+```@contents
+Pages = ["precinditioners.md"]
 ```
 
-**Parameters:**
+## Model domain preconditioners
+
+Model space preconditioners acts on model size arrays such as the velocity model or the FWI/RTM gradient. These preconditioners are indepenedent of the number of sources and therefore should not be indexed.
 
- * `n`: Tuple of model dimensions (e.g. from `model.n`)
+### Water column muting (top mute)
 
- * `mute_start`: First grid point in z-direction from where on to mute the image. Can be a single integer or a vector of length `nx`, where `nx` is the number of grid points in x direction.
+Create a linear operator for a 2D model topmute, i.e. for muting the water column:
 
- * `length`: The mask is created with a linear taper from 0 to 1. The width of the taper is `length`.
+```@docs
+TopMute
+```
 
 **Usage:**
 
@@ -28,34 +30,68 @@ m_mute = Mr'*vec(m)
 
 As `Mr` is self adjoint, `Mr` is equal to `Mr'`.
 
-## Model depth scaling
+**legacy:**
+
+The legacy constructor `judiTopmute(n, wb, taperwidth)` is still available to construct a muting operator with user specified muting depth.
+
 
-Create a 2D model depth scaling:
+### Model depth scaling
+
+Create a 2D model depth scaling. This preconditionenr is the most simple form of inverse Hessain approximation compensating for illumination in the subsurface. We also describe below a more accurate diagonal approximation of the Hessian with the illlumination operator. Additionnaly, as a simple diagonal approximation, this operator is invertible and can be inverted with the standard julia `inv` function.
+
+```@docs
+DepthScaling
+```
+
+**Usage:**
 
 ```julia
-Mr = judiDepthScaling(model)
+# Forward
+m_mute = Mr*vec(m)
+
+# Adjoint
+m_mute = Mr'*vec(m)
 ```
 
-**Parameters:**
+### Illumination
 
- * `model`: JUDI `Model` structure.
 
+The illumination computed the energy of the wavefield along time for each grid point. This provides a first order diagonal approximation of the Hessian of FWI/LSRTM helping the ocnvergence and quality of an update.
 
-## Data topmute (experimental)
+```@docs
+judiIllumination
+```
 
-Create a data topmute for a 2D marine shot record (i.e. for a shot record with an end-on-spread acquisition geometry).
+**Usage:**
 
 ```julia
-Ml = judiMarineTopmute2D(muteStart, geometry; flipmask=false)
+# Forward
+m_mute = I*vec(m)
+
+# Adjoint
+m_mute = I'*vec(m)
 ```
 
-**Parameters:**
+## Data preconditioners
+
+These preconditioners are design to act on the shot records (data). These preconditioners are indexable by source number so that working with a subset of shot is trivial to implement.
 
- * `muteStart`: Vertical index of the apex of the shot record (i.e. the earliest point from where to mute).
 
- * `geometry`: A JUDI `Geometry` object with the receiver geometry.
+### Data topmute
 
- * `flipmask`: If the source is on the left side, set to `false` (default). If the source is on the right side, set to `true`.
+Create a data topmute for the data based on the source and receiver geometry (i.e based on the offsets between each souurce-receiver pair). THis operator allows two modes, `:reflection` for the standard "top-mute" direct wave muting and `:turning` for its opposite muting the reflection to compute gradients purely based on the turning waves. The muting operators uses a cosine taper at the mask limit to allow for a smooth transition.
+
+```@docs
+DataMute
+```
 
+### Band pass filter
+
+While not purely a preconditioner, because this operator acts on the data and is traditionally used fro frequency continuation in FWI, we implemented this operator as a source indexable linear operator as well. Additionally, the filtering function is available as a standalone julia function for general usage
+
+```@docs
+FrequencyFilter
+low_filter
+```
 
 

examples/compressive_splsrtm/BP_synthetic_2004/rtm_bp_2004_freq.jl

@@ -64,7 +64,7 @@ for j=1:d_lin.nsrc
 end
 
 # Topmute
-Ml = judiMarineTopmute2D(35, d_lin.geometry)    # data topmute
+Ml = judiDataMute(q.geometry, d_lin.geometry)
 d_lin = Ml*d_lin
 
 # RTM

examples/compressive_splsrtm/BP_synthetic_2004/splsrtm_bp_2004_freq.jl

@@ -58,6 +58,9 @@ D = judiDepthScaling(model0)
 T = judiTopmute(model0.n, (1 .- water_bottom), [])
 Mr = D*T
 
+# Left-hand preconditioners
+Ml = judiDataMute(q.geometry, d_obs.geometry)
+
 # Linearized Bregman parameters
 x = zeros(Float32, prod(model0.n))
 z = zeros(Float32, prod(model0.n))
@@ -87,18 +90,17 @@ for j=1:niter
     # Select batch and set up left-hand preconditioner
     i = randperm(d_obs.nsrc)[1:batchsize]
     d_sub = get_data(d_obs[i])
-    Ml = judiMarineTopmute2D(35, d_sub.geometry)
 
     # Compute residual and estimate source
     if j > 1
         d_pred = J[i]*Mr*x
-        r = Ml*d_pred - Ml*d_sub
+        r = Ml[i]*d_pred - Ml[i]*d_sub
     else
-        r = Ml*d_sub*(-1f0)    # skip forward modeling in first iteration
+        r = Ml[i]*d_sub*(-1f0)    # skip forward modeling in first iteration
     end
 
     # Residual and gradient
-    g = adjoint(Mr)*adjoint(J[i])*adjoint(Ml)*r
+    g = adjoint(Mr)*adjoint(J[i])*adjoint(Ml[i])*r
 
     # Step size and update variable
     fval[j] = .5*norm(r)^2

examples/scripts/illum.jl

@@ -0,0 +1,100 @@
+# Example for basic 2D modeling:
+# The receiver positions and the source wavelets are the same for each of the four experiments.
+# Author: Mathias Louboutin, mlouboutin3@gatech.edu
+# February 2022
+
+using JUDI, HDF5, SlimPlotting, PyPlot, LinearAlgebra, Downloads, SlimPlotting, Distributed
+
+# Load model
+data_path = dirname(pathof(JUDI))*"/../data/"
+# Load migration velocity model
+if ~isfile(data_path*"marmousi_model.h5")
+    Downloads.download("ftp://slim.gatech.edu/data/SoftwareRelease/Imaging.jl/2DLSRTM/marmousi_model.h5", data_path*"marmousi_model.h5")
+end
+n, d, o, m0, m = read(h5open(data_path*"marmousi_model.h5","r"), "n", "d", "o", "m0", "m")
+
+# Subsample tso that runs ok with ci
+fact = 3
+m = m[1:fact:end, 1:fact:end]
+m0 = m0[1:fact:end, 1:fact:end]
+n = size(m)
+d = (fact*d[1], fact*d[2]) # Base is 5m spacing
+# Setup info and model structure
+nsrc = 21	# number of sources
+model = Model(n, d, o, m)
+model0 = Model(n, d, o, m0)
+
+# Set up receiver geometry
+nxrec = n[1]
+xrec = range(0f0, stop=(n[1] - 1)*d[1], length=nxrec)
+yrec = 0f0
+zrec = range(d[end], stop=d[end], length=nxrec)
+
+# receiver sampling and recording time
+td = 3000f0   # receiver recording time [ms]
+dtd = 2f0    # receiver sampling interval [ms]
+
+# Set up receiver structure
+recGeometry = Geometry(xrec, yrec, zrec; dt=dtd, t=td, nsrc=nsrc)
+
+# Set up source geometry (cell array with source locations for each shot)
+xsrc = convertToCell(range(0f0, stop=(n[1] - 1)*d[1], length=nsrc))
+ysrc = convertToCell(range(0f0, stop=0f0, length=nsrc))
+zsrc = convertToCell(range(d[2], stop=d[2], length=nsrc))
+
+# Set up source structure
+srcGeometry = Geometry(xsrc, ysrc, zsrc; dt=dtd, t=td)
+
+# setup wavelet
+f0 = 0.015f0     # kHz
+wavelet = ricker_wavelet(td, dtd, f0)
+q = judiVector(srcGeometry, wavelet)
+
+###################################################################################################
+# Infer subsampling based on free memory
+mem = Sys.free_memory()/(1024^3)
+grad_mem = 40
+t_sub = max(1, ceil(Int, nworkers()*grad_mem/mem))
+
+# Write shots as segy files to disk
+opt = Options(IC="isic", subsampling_factor=t_sub)
+
+# Setup operators
+F = judiModeling(model, srcGeometry, recGeometry; options=opt)
+F0 = judiModeling(model0, srcGeometry, recGeometry; options=opt)
+J = judiJacobian(F0, q)
+
+# Nonlinear modeling
+dobs = F*q
+d0 = F0*q
+
+I = inv(judiIllumination(J; mode="uv"))
+
+rtm = J'*(d0 - dobs)
+
+# Plot illum
+c1, c2 = maximum(I.illums["u"])/10, maximum(I.illums["v"])/2
+
+figure()
+subplot(311)
+plot_velocity(I.illums["u"]'; new_fig=false, cmap="gist_ncar", vmax=c1)
+subplot(312)
+plot_velocity(I.illums["v"]'; new_fig=false, cmap="gist_ncar", vmax=c2)
+subplot(313)
+plot_velocity(I.illums["u"]'.*I.illums["v"]'; new_fig=false, cmap="gist_ncar", vmax=c1*c2)
+# savefig("Illums.png", bbox_inches="tight")
+tight_layout()
+
+# Plot rtm with corrections
+figure(figsize=(10, 5))
+subplot(221)
+plot_simage(rtm'; new_fig=false, name="rtm", aspect="auto")
+subplot(222)
+plot_simage((I*rtm)'; new_fig=false, name="rtmu", aspect="auto")
+subplot(223)
+plot_simage((I("v")*rtm)'; new_fig=false, name="rtmv", aspect="auto")
+subplot(224)
+plot_simage((I("uv")*rtm)'; new_fig=false, name="rtmuv", aspect="auto")
+tight_layout()
+# savefig("rtms.png", bbox_inches="tight")
+tight_layout()

examples/scripts/lsrtm_2D.jl

@@ -44,8 +44,11 @@ opt = Options(subsampling_factor=t_sub, isic=true)  # ~40 GB of memory per sourc
 M = judiModeling(model0, q.geometry, d_lin.geometry; options=opt)
 J = judiJacobian(M, q)
 
+
 # Right-hand preconditioners (model topmute)
-Mr = judiTopmute(model0.n, 52, 10)
+Mr = judiTopmute(model0; taperwidth=10)
+# Left-hand Preconditionners (data top mute)
+Ml = judiDataMute(q.geometry, d_lin.geometry)
 
 #' set up number of iterations
 niter = parse(Int, get(ENV, "NITER", "10"))
@@ -58,8 +61,7 @@ lsqr_sol = zeros(Float32, prod(model0.n))
 dinv = d_lin[indsrc]
 Jinv = J[indsrc]
 
-Ml = judiMarineTopmute2D(30, dinv.geometry)
-lsqr!(lsqr_sol, Ml*Jinv*Mr, Ml*dinv; maxiter=niter)
+lsqr!(lsqr_sol, Ml[indsrc]*Jinv*Mr, Ml[indsrc]*dinv; maxiter=niter)
 
 # Save final velocity model, function value and history
 h5open("lsrtm_marmousi_lsqr_result.h5", "w") do file