add lsqr example for lsrtm

slimgroup · Jun 24, 2022 · a2360a4 · a2360a4
1 parent addf8d5
commit a2360a4
Show file tree

Hide file tree

Showing 2 changed files with 57 additions and 16 deletions.
diff --git a/examples/scripts/lsrtm_2D.jl b/examples/scripts/lsrtm_2D.jl
@@ -0,0 +1,55 @@
+# LS-RTM of the 2D Marmousi model using LSQR
+# Author: [email protected]
+# Date: June 2022
+
+using Statistics, Random, LinearAlgebra, JOLI
+using JUDI, SegyIO, HDF5, PyPlot, IterativeSolvers
+
+
+# Load migration velocity model
+if ~isfile("$(JUDI.JUDI_DATA)/marmousi_model.h5")
+    ftp_data("ftp://slim.gatech.edu/data/SoftwareRelease/Imaging.jl/2DLSRTM/marmousi_model.h5")
+end
+n, d, o, m0 = read(h5open("$(JUDI.JUDI_DATA)/marmousi_model.h5", "r"), "n", "d", "o", "m0")
+
+# Set up model structure
+model0 = Model((n[1], n[2]), (d[1], d[2]), (o[1], o[2]), m0)
+
+# Load data
+if ~isfile("$(JUDI.JUDI_DATA)/marmousi_2D.segy")
+    ftp_data("ftp://slim.gatech.edu/data/SoftwareRelease/Imaging.jl/2DLSRTM/marmousi_2D.segy")
+end
+block = segy_read("$(JUDI.JUDI_DATA)/marmousi_2D.segy")
+d_lin = judiVector(block)   # linearized observed data
+
+# Set up wavelet
+src_geometry = Geometry(block; key = "source", segy_depth_key = "SourceDepth")
+wavelet = ricker_wavelet(src_geometry.t[1], src_geometry.dt[1], 0.03)    # 30 Hz wavelet
+q = judiVector(src_geometry, wavelet)
+
+###################################################################################################
+# Infer subsampling based on free memory
+mem = Sys.free_memory()/(1024^3)
+t_sub = max(1, ceil(Int, 40/mem))
+# Setup operators
+opt = Options(subsampling_factor=t_sub, isic=true)  # ~40 GB of memory per source without subsampling
+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)
+
+#' LSQR
+niter = parse(Int, get(ENV, "NITER", "10"))
+lsqr_sol = zeros(Float32, prod(n))
+Ml = judiMarineTopmute2D(30, d_lin[1].geometry)
+lsqr!(lsqr_sol, Ml*J[1]*Mr, Ml*d_lin[1]; maxiter=niter)     # only invert the first shot record
+
+# Save final velocity model, function value and history
+h5open("lsrtm_marmousi_lsqr_result.h5", "w") do file
+    write(file, "x", reshape(lsqr_sol, model0.n))
+end
+
+# Plot final image
+figure(); imshow(copy(adjoint(reshape(lsqr_sol, model0.n))), extent = (0, 7.99, 3.19, 0), cmap = "gray", vmin = -3e-2, vmax = 3e-2)
+title("LS-RTM with LSQR"); xlabel("Lateral position [km]"); ylabel("Depth [km]")
diff --git a/examples/scripts/splsrtm_2D.jl b/examples/scripts/splsrtm_2D.jl
@@ -3,7 +3,7 @@
 # Date: April 2022
 
 using Statistics, Random, LinearAlgebra, JOLI
-using JUDI, SegyIO, HDF5, PyPlot, SlimOptim, IterativeSolvers
+using JUDI, SegyIO, HDF5, PyPlot, SlimOptim
 
 
 # Load migration velocity model
@@ -73,18 +73,4 @@ end
 
 # Plot final image
 figure(); imshow(copy(adjoint(reshape(solb.x, model0.n))), extent = (0, 7.99, 3.19, 0), cmap = "gray", vmin = -3e-2, vmax = 3e-2)
-title("SPLS-RTM with SGD"); xlabel("Lateral position [km]"); ylabel("Depth [km]")
-
-#' LSQR
-lsqr_sol = 0f0 .* solb.x
-Ml = judiMarineTopmute2D(30, d_lin[1].geometry)
-lsqr!(lsqr_sol, Ml*J[1]*Mr, Ml*d_lin[1]; maxiter=1)
-
-# Save final velocity model, function value and history
-h5open("lsrtm_marmousi_lsqr_result.h5", "w") do file
-    write(file, "x", reshape(lsqr_sol, model0.n))
-end
-
-# Plot final image
-figure(); imshow(copy(adjoint(reshape(lsqr_sol, model0.n))), extent = (0, 7.99, 3.19, 0), cmap = "gray", vmin = -3e-2, vmax = 3e-2)
-title("SPLS-RTM with LSQR"); xlabel("Lateral position [km]"); ylabel("Depth [km]")
+title("SPLS-RTM with Linearized Bregman"); xlabel("Lateral position [km]"); ylabel("Depth [km]")