From e45081a59e669cbc96b643197f64986646573c17 Mon Sep 17 00:00:00 2001 From: Keith Roberts Date: Wed, 14 Apr 2021 15:43:13 -0300 Subject: [PATCH] sigsbee2b sharp interface problem --- run_forward_sharp_interface_sigsbee.py | 10 +- run_sharp_interface_optimization_sigsbee.py | 406 ++++++++++++++ spyro/io/io.py | 571 ++++++++++---------- spyro/solvers/leapfrog_adjoint_level_set.py | 11 +- 4 files changed, 706 insertions(+), 292 deletions(-) create mode 100644 run_sharp_interface_optimization_sigsbee.py diff --git a/run_forward_sharp_interface_sigsbee.py b/run_forward_sharp_interface_sigsbee.py index 712fdcb6..29f602f5 100644 --- a/run_forward_sharp_interface_sigsbee.py +++ b/run_forward_sharp_interface_sigsbee.py @@ -15,7 +15,6 @@ "quadrature": "KMV", "dimension": 2, # dimension } -# bbox = (-9.15162, 0.0, 0.0, 27.43962) model["mesh"] = { "Lz": 9.15162, # depth in km - always positive "Lx": 27.43962, # width in km - always positive @@ -37,7 +36,7 @@ } """The three surveys shared the same acquisition geometry. Each receiver recorded data every .008 seconds for 1 500 timesteps resulting in 12 seconds of data. A 7 950 m (26 100 ft) long streamer cable was deployed with 348 hydrophones spaced 22.86 m (75 ft) apart. Shots were fired every 45.72 m (150 ft) starting at 3 330 m (10 925 ft). Table 5 shows the values that Sigsbee shot headers should contain.""" # We do 1/10 of the total number of true shots -sources = spyro.create_transect((-0.01, 3.33), (-0.01, 19.4896), 50) +sources = spyro.create_transect((-0.01, 3.33), (-0.01, 19.4896), 40) model["acquisition"] = { "source_type": "Ricker", "num_sources": len(sources), @@ -56,12 +55,11 @@ "fspool": 9999, # how frequently to save solution to ram } - comm = spyro.utils.mpi_init(model) mesh, V = spyro.io.read_mesh(model, comm) -vp_exact = spyro.io.interpolate(model, mesh, V, guess=True) +vp_exact = spyro.io.interpolate(model, mesh, V, guess=False) File("exact_vp.pvd").write(vp_exact, name="true_velocity") @@ -89,10 +87,10 @@ sources, receivers, source_num=sn, - output=True, + output=False, ) print(time.time() - t1) - spyro.io.save_shots("shots/forward_exact_level_set" + str(sn) + ".dat", p_recv) + spyro.io.save_shots("shots/sigsbee2b_true_" + str(sn) + ".dat", p_recv) spyro.plots.plot_shotrecords( model, p_recv, diff --git a/run_sharp_interface_optimization_sigsbee.py b/run_sharp_interface_optimization_sigsbee.py new file mode 100644 index 00000000..7a4cbc52 --- /dev/null +++ b/run_sharp_interface_optimization_sigsbee.py @@ -0,0 +1,406 @@ +from mpi4py import MPI +import numpy as np +from firedrake import * + +import spyro + +model = {} +model["parallelism"] = { + # "type": "automatic", # options: automatic (same number of cores for evey processor), custom, off + "type": "off", + "custom_cores_per_shot": [], # only if the user wants a different number of cores for every shot. +} +model["opts"] = { + "method": "KMV", + "degree": 1, # p order + "quadrature": "KMV", + "dimension": 2, # dimension +} +model["mesh"] = { + "Lz": 9.15162, # depth in km - always positive + "Lx": 27.43962, # width in km - always positive + "Ly": 0.0, # thickness in km - always positive + "meshfile": "meshes/sigsbee2b_guess.msh", + "initmodel": "velocity_models/sigsbee2b_guess.hdf5", + "truemodel": "velocity_models/sigsbee2b_true.hdf5", + "background": "velocity_models/sigsbee2b_true_background.hdf5", +} +model["PML"] = { + "status": True, # true, # true or false + "outer_bc": "non-reflective", # dirichlet, neumann, non-reflective (outer boundary condition) + "damping_type": "polynomial", # polynomial, hyperbolic, shifted_hyperbolic + "exponent": 2, + "cmax": 4.5, # maximum acoustic wave velocity in pml - km/s + "R": 0.001, # theoretical reflection coefficient + "lz": 0.50, # thickness of the pml in the z-direction (km) - always positive + "lx": 0.50, # thickness of the pml in the x-direction (km) - always positive + "ly": 0.0, # thickness of the pml in the y-direction (km) - always positive +} +sources = spyro.create_transect((-0.01, 3.33), (-0.01, 19.4896), 3) +model["acquisition"] = { + "source_type": "Ricker", + "num_sources": len(sources), + "source_pos": sources, + "frequency": 5.0, + "delay": 1.0, + "amplitude": 1.0, + "num_receivers": 348, + "receiver_locations": None, +} +model["timeaxis"] = { + "t0": 0.0, # initial time for event + "tf": 6.0, # final time for event + "dt": 0.001, # timestep size + "nspool": 100, # how frequently to output solution to pvds + "fspool": 9999, # how frequently to save solution to ram +} + +VP_1 = 4.4 # inside subdomain to be optimized +#### end of options #### + + +def calculate_indicator_from_vp(vp): + """Create an indicator function""" + dgV = FunctionSpace(mesh, "DG", 0) + # the salt body is 1 everything else is 2 + cond = conditional(vp > (VP_1 - 0.1), 1, 2) + indicator = Function(dgV, name="indicator").interpolate(cond) + return indicator + + +def update_velocity(V, q, vp_background): + """Update the velocity (material properties) + based on the indicator function + """ + sd1 = SubDomainData(q < 1.5) + + # place background field everywhere + vp_new = Function(V, name="velocity") + vp_new.assign(vp_background) + vp_new.interpolate(Constant(VP_1), subset=sd1) + + return vp_new + + +def create_weighting_function(V, const=100.0, M=5, width=0.1, show=False): + """Create a weighting function g, which is large near the + boundary of the domain and a constant smaller value in the interior + + Inputs + ------ + V: Firedrake.FunctionSpace + const: the weight function is equal to this constant value, except close to the boundary + M: maximum value on the boundary will be M**2 + width: the decimal fraction of the domain where the weight is > constant + show: Visualize the weighting function + + Outputs + ------- + wei: a Firedrake.Function containing the weights + + """ + + # get coordinates of DoFs + m = V.ufl_domain() + W2 = VectorFunctionSpace(m, V.ufl_element()) + coords = interpolate(m.coordinates, W2) + Z, X = coords.dat.data[:, 0], coords.dat.data[:, 1] + + a0 = np.amin(X) + a1 = np.amax(X) + b0 = np.amin(Z) + b1 = np.amax(Z) + + cx = a1 - a0 # x-coordinate of center of rectangle + cz = b1 - b0 # z-coordinate of center of rectangle + + def h(t, d): + L = width * d # fraction of the domain where the weight is > constant + return (np.maximum(0.0, M / L * t + M)) ** 2 + + w = const * ( + 1.0 + + np.maximum( + h(X - a1, a1 - a0) + h(a0 - X, a1 - a0), + h(b0 - Z, b1 - b0) + h(Z - b1, b1 - b0), + ) + ) + if show: + import matplotlib.pyplot as plt + + plt.scatter(Z, X, 5, c=w) + plt.colorbar() + plt.show() + + wei = Function(V, w, name="weighting_function") + + # File("weighting_function.pvd").write(wei) + return wei + + +def calculate_functional(model, mesh, comm, vp, sources, receivers, iter_num): + """Calculate the l2-norm functional""" + if comm.ensemble_comm.rank == 0: + print("Computing the cost functional", flush=True) + J_local = np.zeros((1)) + J_total = np.zeros((1)) + XMIN = 0.01 + XMAX = 27.43 + for sn in range(model["acquisition"]["num_sources"]): + if spyro.io.is_owner(comm, sn): + # calculate receivers for each source + offset = 7.950 # offset of 7.950 km + _, xsrc = model["acquisition"]["source_pos"][sn] + xmin = xsrc - offset if xsrc - offset > XMIN else 0.01 + xmax = xsrc + offset if xsrc + offset < XMAX else 27.43 + model["acquisition"]["receiver_locations"] = spyro.create_transect( + (-0.01, xmin), (-0.01, xmax), 348 + ) + receivers = spyro.Receivers(model, mesh, V, comm).create() + + guess, guess_dt, guess_recv = spyro.solvers.Leapfrog_level_set( + model, mesh, comm, vp, sources, receivers, source_num=sn + ) + f = "shots/sigsbee2b_true_" + str(sn) + ".dat" + p_exact_recv = spyro.io.load_shots(f) + # DEBUG + # viz the signal at receiver # 100 + if comm.comm.rank == 0: + import matplotlib.pyplot as plt + + plt.plot(p_exact_recv[:-1:1, 100], "k-") + plt.plot(guess_recv[:, 100], "r-") + plt.ylim(-5e-5, 5e-5) + plt.title("Receiver #100") + plt.savefig( + "comparison_" + + str(comm.ensemble_comm.rank) + + "_iter_" + + str(iter_num) + + ".png" + ) + plt.close() + # END DEBUG + + residual = spyro.utils.evaluate_misfit( + model, + comm, + guess_recv, + p_exact_recv, + ) + J_local[0] += spyro.utils.compute_functional(model, comm, residual) + if comm.ensemble_comm.size > 1: + COMM_WORLD.Allreduce(J_local, J_total, op=MPI.SUM) + J_total[0] /= comm.ensemble_comm.size + if comm.ensemble_comm.rank == 0: + print(f"The cost functional is: {J_total[0]}") + return ( + J_total[0], + guess, + guess_dt, + residual, + ) + + +def calculate_gradient( + model, mesh, comm, vp, vp_background, guess, guess_dt, weighting, residual +): + """Calculate the shape gradient""" + if comm.ensemble_comm.rank == 0: + print("Computing the shape derivative...", flush=True) + VF = VectorFunctionSpace(mesh, model["opts"]["method"], model["opts"]["degree"]) + theta = Function(VF, name="gradient") + for sn in range(model["acquisition"]["num_sources"]): + if spyro.io.is_owner(comm, sn): + theta_local = spyro.solvers.Leapfrog_adjoint_level_set( + model, + mesh, + comm, + vp, + vp_background, + guess, + guess_dt, + weighting, + residual, + source_num=sn, + output=False, + piecewise_smooth=True, + ) + # sum shape gradient if ensemble parallelism here + if comm.ensemble_comm.size > 1: + comm.ensemble_comm.Allreduce( + theta_local.dat.data[:], theta.dat.data[:], op=MPI.SUM + ) + else: + theta = theta_local + # scale factor + # theta.dat.data[:] *= -1 + return theta + + +def model_update(mesh, indicator, theta, step, timesteps=10): + """Solve a transport equation to move the subdomains around based + on the shape gradient which hopefully minimizes the functional. + """ + if comm.ensemble_comm.rank == 0: + print("Updating the shape...", flush=True) + indicator_new = spyro.solvers.advect( + mesh, + indicator, + step * theta, + number_of_timesteps=timesteps, + output=False, + ) + return indicator_new + + +def optimization( + model, mesh, V, comm, vp, vp_background, sources, receivers, max_iter=10 +): + """Optimization with steepest descent using a line search algorithm""" + beta0 = beta0_init = 1.5 + max_ls = 10 + gamma = gamma2 = 0.8 + timesteps = 10 + + # the file that contains the shape gradient each iteration + if comm.ensemble_comm.rank == 0: + grad_file = File("theta.pvd", comm=comm.comm) + + weighting = create_weighting_function(V, width=0.1, M=10, const=1e-9) + + ls_iter = 0 + iter_num = 0 + # calculate the new functional for the new model + J_old, guess, guess_dt, residual = calculate_functional( + model, mesh, comm, vp, sources, receivers, iter_num + ) + while iter_num < max_iter: + if comm.ensemble_comm.rank == 0 and iter_num == 0 and ls_iter == 0: + print("Commencing the inversion...") + + if comm.ensemble_comm.rank == 0: + print(f"The step size is: {beta0}", flush=True) + + # compute the shape gradient for the new domain + if ls_iter == 0: + theta = calculate_gradient( + model, + mesh, + comm, + vp, + vp_background, + guess, + guess_dt, + weighting, + residual, + piecewise_smooth=True, + ) + # write the gradient to a vtk file + if comm.ensemble_comm.rank == 0: + grad_file.write(theta, name="gradient") + # calculate the so-called indicator function by thresholding vp + indicator = calculate_indicator_from_vp(vp) + # update the new shape by solving the transport equation with the indicator field + indicator_new = model_update(mesh, indicator, theta, beta0, timesteps=timesteps) + # update the velocity according to the new indicator + vp_new = update_velocity(V, indicator_new, vp_background) + # write ALL velocity updates to a vtk file + if comm.ensemble_comm.rank == 0: + evolution_of_velocity.write(vp_new, name="velocity") + # compute the new functional + J_new, guess_new, guess_dt_new, residual_new = calculate_functional( + model, mesh, comm, vp_new, sources, receivers, iter_num + ) + # write the new velocity to a vtk file + # using a line search to attempt to reduce the functional + if J_new < J_old: + if comm.ensemble_comm.rank == 0: + print( + f"Iteration {iter_num}: Functional was {J_old}. Accepting shape update...new functional is: {J_new}", + flush=True, + ) + + iter_num += 1 + # accept new domain + J_old = J_new + guess = guess_new + guess_dt = guess_dt_new + residual = residual_new + indicator = indicator_new + vp = vp_new + # update step + if ls_iter == max_ls: + beta0 = max(beta0 * gamma2, 0.1 * beta0_init) + elif ls_iter == 0: + beta0 = min(beta0 / gamma2, 1.0) + else: + # no change to step + beta0 = beta0 + ls_iter = 0 + elif ls_iter < max_ls: + if comm.ensemble_comm.rank == 0: + print( + f"Previous cost functional {J_old}...new cost functional {J_new}", + flush=True, + ) + print( + f"Line search number {ls_iter}...reducing step size...", flush=True + ) + # advance the line search counter + ls_iter += 1 + if abs(J_new - J_old) < 1e-16: + print( + f"Line search number {ls_iter}...increasing step size...", + flush=True, + ) + beta0 = min(beta0 / gamma2, 10.0) + else: + print( + f"Line search number {ls_iter}...reducing step size...", flush=True + ) + # reduce step length by gamma + beta0 *= gamma + # now solve the transport equation over again + # but with the reduced step + # Need to recompute guess_dt since was discarded above + # compute the new functional (using old velocity field) + J, guess, guess_dt, residual = calculate_functional( + model, mesh, comm, vp, sources, receivers, iter_num + ) + else: + print( + f"Failed to reduce the functional after {ls_iter} line searches...", + flush=True, + ) + break + + return vp + + +# run the script + +comm = spyro.utils.mpi_init(model) + +mesh, V = spyro.io.read_mesh(model, comm) + +vp = spyro.io.interpolate(model, mesh, V, guess=True) + +vp_background = spyro.io.interpolate(model, mesh, V, background=True) + +# visualize the updates with this file +if comm.ensemble_comm.rank == 0: + evolution_of_velocity = File("evolution_of_velocity.pvd", comm=comm.comm) + evolution_of_velocity.write(vp, name="velocity") + + File("vp_background.pvd").write(vp_background) + +# Configure the sources and receivers +sources = spyro.Sources(model, mesh, V, comm).create() + +receivers = spyro.Receivers(model, mesh, V, comm).create() + +# run the optimization based on a line search for max_iter iterations +vp = optimization( + model, mesh, V, comm, vp, vp_background, sources, receivers, max_iter=50 +) diff --git a/spyro/io/io.py b/spyro/io/io.py index e298a4fb..3a575fdf 100644 --- a/spyro/io/io.py +++ b/spyro/io/io.py @@ -1,282 +1,289 @@ -from __future__ import with_statement - -import os -import pickle - -import firedrake as fire -import h5py -import numpy as np -from scipy.interpolate import RegularGridInterpolator -from scipy.interpolate import griddata -import segyio - -from .. import domains - -__all__ = ["write_function_to_grid", "create_segy", "is_owner", "save_shots", "load_shots", "read_mesh", "interpolate"] - - - -def write_function_to_grid(function, V, grid_spacing): - """Interpolate a Firedrake function to a structured grid""" - # get DoF coordinates - m = V.ufl_domain() - W = VectorFunctionSpace(m, V.ufl_element()) - coords = interpolate(m.coordinates, W) - x, y = coords.dat.data[:, 0], coords.dat.data[:, 1] - - # add buffer to avoid NaN when calling griddata - min_x = np.amin(x) + 0.01 - max_x = np.amax(x) - 0.01 - min_y = np.amin(y) + 0.01 - max_y = np.amax(y) - 0.01 - - z = function.dat.data[:] * 1000.0 # convert from km/s to m/s - - # target grid to interpolate to - xi = np.arange(min_x, max_x, grid_spacing) - yi = np.arange(min_y, max_y, grid_spacing) - xi, yi = np.meshgrid(xi, yi) - - # interpolate - zi = griddata((x, y), z, (xi, yi), method="linear") - - return xi, yi, zi - - -def create_segy(velocity, filename): - """Write the velocity data into a segy file named filename""" - spec = segyio.spec() - - velocity = np.flipud(velocity.T) - - spec.sorting = 2 # not sure what this means - spec.format = 1 # not sure what this means - spec.samples = range(velocity.shape[0]) - spec.ilines = range(velocity.shape[1]) - spec.xlines = range(velocity.shape[0]) - - assert np.sum(np.isnan(velocity[:])) == 0 - - with segyio.create(filename, spec) as f: - for tr, il in enumerate(spec.ilines): - f.trace[tr] = velocity[:, tr] - - -def save_shots(filename, array): - """Save a `numpy.ndarray` to a `pickle`. - - Parameters - ---------- - filename: str - The filename to save the data as a `pickle` - array: `numpy.ndarray` - The data to save a pickle (e.g., a shot) - - Returns - ------- - None - - """ - with open(filename, "wb") as f: - pickle.dump(array, f) - return None - - -def load_shots(filename): - """Load a `pickle` to a `numpy.ndarray`. - - Parameters - ---------- - filename: str - The filename to save the data as a `pickle` - - Returns - ------- - array: `numpy.ndarray` - The data - - """ - - with open(filename, "rb") as f: - array = np.asarray(pickle.load(f), dtype=float) - return array - - -def is_owner(ens_comm, rank): - """Distribute shots between processors in using a modulus operator - - Parameters - ---------- - ens_comm: Firedrake.ensemble_communicator - An ensemble communicator - rank: int - The rank of the core - - Returns - ------- - boolean - `True` if `rank` owns this shot - - """ - return ens_comm.ensemble_comm.rank == (rank % ens_comm.ensemble_comm.size) - - -def _check_units(c): - if min(c.dat.data[:]) > 100.0: - # data is in m/s but must be in km/s - if fire.COMM_WORLD.rank == 0: - print("INFO: converting from m/s to km/s", flush=True) - c.assign(c / 1000.0) # meters to kilometers - return c - - -def interpolate(model, mesh, V, guess=False): - """Read and interpolate a seismic velocity model stored - in a HDF5 file onto the nodes of a finite element space. - - Parameters - ---------- - model: `dictionary` - Model options and parameters. - mesh: Firedrake.mesh object - A mesh object read in by Firedrake. - V: Firedrake.FunctionSpace object - The space of the finite elements. - guess: boolean, optinal - Is it a guess model or a `exact` model? - - Returns - ------- - c: Firedrake.Function - P-wave seismic velocity interpolated onto the nodes of the finite elements. - - """ - sd = V.mesh().geometric_dimension() - m = V.ufl_domain() - if model["PML"]["status"]: - minz = -model["mesh"]["Lz"] - model["PML"]["lz"] - maxz = 0.0 - minx = 0.0 - model["PML"]["lx"] - maxx = model["mesh"]["Lx"] + model["PML"]["lx"] - miny = 0.0 - model["PML"]["ly"] - maxy = model["mesh"]["Ly"] + model["PML"]["ly"] - else: - minz = -model["mesh"]["Lz"] - maxz = 0.0 - minx = 0.0 - maxx = model["mesh"]["Lx"] - miny = 0.0 - maxy = model["mesh"]["Ly"] - - W = fire.VectorFunctionSpace(m, V.ufl_element()) - coords = fire.interpolate(m.coordinates, W) - # (z,x) or (z,x,y) - if sd == 2: - qp_z, qp_x = coords.dat.data[:, 0], coords.dat.data[:, 1] - elif sd == 3: - qp_z, qp_x, qp_y = ( - coords.dat.data[:, 0], - coords.dat.data[:, 1], - coords.dat.data[:, 2], - ) - else: - raise NotImplementedError - - if guess: - fname = model["mesh"]["initmodel"] - else: - fname = model["mesh"]["truemodel"] - - with h5py.File(fname, "r") as f: - Z = np.asarray(f.get("velocity_model")[()]) - - if sd == 2: - nrow, ncol = Z.shape - z = np.linspace(minz, maxz, nrow) - x = np.linspace(minx, maxx, ncol) - - # make sure no out-of-bounds - qp_z2 = [minz if z < minz else maxz if z > maxz else z for z in qp_z] - qp_x2 = [minx if x < minx else maxx if x > maxx else x for x in qp_x] - - interpolant = RegularGridInterpolator((z, x), Z) - tmp = interpolant((qp_z2, qp_x2)) - elif sd == 3: - nrow, ncol, ncol2 = Z.shape - z = np.linspace(minz, maxz, nrow) - x = np.linspace(minx, maxx, ncol) - y = np.linspace(miny, maxy, ncol2) - - # make sure no out-of-bounds - qp_z2 = [minz if z < minz else maxz if z > maxz else z for z in qp_z] - qp_x2 = [minx if x < minx else maxx if x > maxx else x for x in qp_x] - qp_y2 = [miny if y < miny else maxy if y > maxy else y for y in qp_y] - - interpolant = RegularGridInterpolator((z, x, y), Z) - tmp = interpolant((qp_z2, qp_x2, qp_y2)) - - c = fire.Function(V, name="velocity") - c.dat.data[:] = tmp - c = _check_units(c) - return c - - -def read_mesh(model, ens_comm): - """Reads in an external mesh and scatters it between cores. - - Parameters - ---------- - model: `dictionary` - Model options and parameters. - ens_comm: Firedrake.ensemble_communicator - An ensemble communicator - - Returns - ------- - mesh: Firedrake.Mesh object - The distributed mesh across `ens_comm`. - V: Firedrake.FunctionSpace object - The space of the finite elements - - """ - - method = model["opts"]["method"] - degree = model["opts"]["degree"] - - num_sources = model["acquisition"]["num_sources"] - mshname = model["mesh"]["meshfile"] - - if method == "CG" or method == "KMV": - mesh = fire.Mesh( - mshname, - comm=ens_comm.comm, - distribution_parameters={ - "overlap_type": (fire.DistributedMeshOverlapType.NONE, 0) - }, - ) - else: - mesh = fire.Mesh(mshname, comm=ens_comm.comm) - if ens_comm.comm.rank == 0 and ens_comm.ensemble_comm.rank == 0: - print( - "INFO: Distributing %d shot(s) across %d core(s). Each shot is using %d cores" - % ( - num_sources, - fire.COMM_WORLD.size, - fire.COMM_WORLD.size / ens_comm.ensemble_comm.size, - ), - flush=True, - ) - print( - " rank %d on ensemble %d owns %d elements and can access %d vertices" - % ( - mesh.comm.rank, - ens_comm.ensemble_comm.rank, - mesh.num_cells(), - mesh.num_vertices(), - ), - flush=True, - ) - # Element type - element = domains.space.FE_method(mesh, method, degree) - # Space of problem - return mesh, fire.FunctionSpace(mesh, element) +from __future__ import with_statement + +import os +import pickle + +import firedrake as fire +import h5py +import numpy as np +from scipy.interpolate import RegularGridInterpolator +from scipy.interpolate import griddata +import segyio + +from .. import domains + +__all__ = ["write_function_to_grid", "create_segy", "is_owner", "save_shots", "load_shots", "read_mesh", "interpolate"] + + + +def write_function_to_grid(function, V, grid_spacing): + """Interpolate a Firedrake function to a structured grid""" + # get DoF coordinates + m = V.ufl_domain() + W = VectorFunctionSpace(m, V.ufl_element()) + coords = interpolate(m.coordinates, W) + x, y = coords.dat.data[:, 0], coords.dat.data[:, 1] + + # add buffer to avoid NaN when calling griddata + min_x = np.amin(x) + 0.01 + max_x = np.amax(x) - 0.01 + min_y = np.amin(y) + 0.01 + max_y = np.amax(y) - 0.01 + + z = function.dat.data[:] * 1000.0 # convert from km/s to m/s + + # target grid to interpolate to + xi = np.arange(min_x, max_x, grid_spacing) + yi = np.arange(min_y, max_y, grid_spacing) + xi, yi = np.meshgrid(xi, yi) + + # interpolate + zi = griddata((x, y), z, (xi, yi), method="linear") + + return xi, yi, zi + + +def create_segy(velocity, filename): + """Write the velocity data into a segy file named filename""" + spec = segyio.spec() + + velocity = np.flipud(velocity.T) + + spec.sorting = 2 # not sure what this means + spec.format = 1 # not sure what this means + spec.samples = range(velocity.shape[0]) + spec.ilines = range(velocity.shape[1]) + spec.xlines = range(velocity.shape[0]) + + assert np.sum(np.isnan(velocity[:])) == 0 + + with segyio.create(filename, spec) as f: + for tr, il in enumerate(spec.ilines): + f.trace[tr] = velocity[:, tr] + + +def save_shots(filename, array): + """Save a `numpy.ndarray` to a `pickle`. + + Parameters + ---------- + filename: str + The filename to save the data as a `pickle` + array: `numpy.ndarray` + The data to save a pickle (e.g., a shot) + + Returns + ------- + None + + """ + with open(filename, "wb") as f: + pickle.dump(array, f) + return None + + +def load_shots(filename): + """Load a `pickle` to a `numpy.ndarray`. + + Parameters + ---------- + filename: str + The filename to save the data as a `pickle` + + Returns + ------- + array: `numpy.ndarray` + The data + + """ + + with open(filename, "rb") as f: + array = np.asarray(pickle.load(f), dtype=float) + return array + + +def is_owner(ens_comm, rank): + """Distribute shots between processors in using a modulus operator + + Parameters + ---------- + ens_comm: Firedrake.ensemble_communicator + An ensemble communicator + rank: int + The rank of the core + + Returns + ------- + boolean + `True` if `rank` owns this shot + + """ + return ens_comm.ensemble_comm.rank == (rank % ens_comm.ensemble_comm.size) + + +def _check_units(c): + if min(c.dat.data[:]) > 100.0: + # data is in m/s but must be in km/s + if fire.COMM_WORLD.rank == 0: + print("INFO: converting from m/s to km/s", flush=True) + c.assign(c / 1000.0) # meters to kilometers + return c + + +def interpolate(model, mesh, V, guess=False, background=False): + """Read and interpolate a seismic velocity model stored + in a HDF5 file onto the nodes of a finite element space. + + Parameters + ---------- + model: `dictionary` + Model options and parameters. + mesh: Firedrake.mesh object + A mesh object read in by Firedrake. + V: Firedrake.FunctionSpace object + The space of the finite elements. + guess: boolean, optional + Is it a guess model or a `exact` model? + background: boolean, optional + background velocity model for sharp interface modeling + + Returns + ------- + c: Firedrake.Function + P-wave seismic velocity interpolated onto the nodes of the finite elements. + + """ + sd = V.mesh().geometric_dimension() + m = V.ufl_domain() + if model["PML"]["status"]: + minz = -model["mesh"]["Lz"] - model["PML"]["lz"] + maxz = 0.0 + minx = 0.0 - model["PML"]["lx"] + maxx = model["mesh"]["Lx"] + model["PML"]["lx"] + miny = 0.0 - model["PML"]["ly"] + maxy = model["mesh"]["Ly"] + model["PML"]["ly"] + else: + minz = -model["mesh"]["Lz"] + maxz = 0.0 + minx = 0.0 + maxx = model["mesh"]["Lx"] + miny = 0.0 + maxy = model["mesh"]["Ly"] + + W = fire.VectorFunctionSpace(m, V.ufl_element()) + coords = fire.interpolate(m.coordinates, W) + # (z,x) or (z,x,y) + if sd == 2: + qp_z, qp_x = coords.dat.data[:, 0], coords.dat.data[:, 1] + elif sd == 3: + qp_z, qp_x, qp_y = ( + coords.dat.data[:, 0], + coords.dat.data[:, 1], + coords.dat.data[:, 2], + ) + else: + raise NotImplementedError + + if guess: + fname = model["mesh"]["initmodel"] + else: + fname = model["mesh"]["truemodel"] + + if background: + fname = model["mesh"]["background"] + + print(fname) + + with h5py.File(fname, "r") as f: + Z = np.asarray(f.get("velocity_model")[()]) + + if sd == 2: + nrow, ncol = Z.shape + z = np.linspace(minz, maxz, nrow) + x = np.linspace(minx, maxx, ncol) + + # make sure no out-of-bounds + qp_z2 = [minz if z < minz else maxz if z > maxz else z for z in qp_z] + qp_x2 = [minx if x < minx else maxx if x > maxx else x for x in qp_x] + + interpolant = RegularGridInterpolator((z, x), Z) + tmp = interpolant((qp_z2, qp_x2)) + elif sd == 3: + nrow, ncol, ncol2 = Z.shape + z = np.linspace(minz, maxz, nrow) + x = np.linspace(minx, maxx, ncol) + y = np.linspace(miny, maxy, ncol2) + + # make sure no out-of-bounds + qp_z2 = [minz if z < minz else maxz if z > maxz else z for z in qp_z] + qp_x2 = [minx if x < minx else maxx if x > maxx else x for x in qp_x] + qp_y2 = [miny if y < miny else maxy if y > maxy else y for y in qp_y] + + interpolant = RegularGridInterpolator((z, x, y), Z) + tmp = interpolant((qp_z2, qp_x2, qp_y2)) + + c = fire.Function(V, name="velocity") + c.dat.data[:] = tmp + c = _check_units(c) + return c + + +def read_mesh(model, ens_comm): + """Reads in an external mesh and scatters it between cores. + + Parameters + ---------- + model: `dictionary` + Model options and parameters. + ens_comm: Firedrake.ensemble_communicator + An ensemble communicator + + Returns + ------- + mesh: Firedrake.Mesh object + The distributed mesh across `ens_comm`. + V: Firedrake.FunctionSpace object + The space of the finite elements + + """ + + method = model["opts"]["method"] + degree = model["opts"]["degree"] + + num_sources = model["acquisition"]["num_sources"] + mshname = model["mesh"]["meshfile"] + + if method == "CG" or method == "KMV": + mesh = fire.Mesh( + mshname, + comm=ens_comm.comm, + distribution_parameters={ + "overlap_type": (fire.DistributedMeshOverlapType.NONE, 0) + }, + ) + else: + mesh = fire.Mesh(mshname, comm=ens_comm.comm) + if ens_comm.comm.rank == 0 and ens_comm.ensemble_comm.rank == 0: + print( + "INFO: Distributing %d shot(s) across %d core(s). Each shot is using %d cores" + % ( + num_sources, + fire.COMM_WORLD.size, + fire.COMM_WORLD.size / ens_comm.ensemble_comm.size, + ), + flush=True, + ) + print( + " rank %d on ensemble %d owns %d elements and can access %d vertices" + % ( + mesh.comm.rank, + ens_comm.ensemble_comm.rank, + mesh.num_cells(), + mesh.num_vertices(), + ), + flush=True, + ) + # Element type + element = domains.space.FE_method(mesh, method, degree) + # Space of problem + return mesh, fire.FunctionSpace(mesh, element) diff --git a/spyro/solvers/leapfrog_adjoint_level_set.py b/spyro/solvers/leapfrog_adjoint_level_set.py index 35ad2f42..bab3353c 100644 --- a/spyro/solvers/leapfrog_adjoint_level_set.py +++ b/spyro/solvers/leapfrog_adjoint_level_set.py @@ -23,6 +23,7 @@ def Leapfrog_adjoint_level_set( mesh, comm, c, + c_background, guess, guess_dt, weighting, @@ -181,14 +182,16 @@ def Leapfrog_adjoint_level_set( if piecewise_smooth: # \tilde\nabla c = H*nabla c_salt + (1-H)*nabla c_background - c_salt = Function(V).assign(4.5) # velocity in salt + c_salt = Function(V).assign(4.4) # velocity in salt # analytical background gradient - Z, _ = SpatialCoordinate(mesh) - c_background = Function(V).interpolate(Min(1.5 + 4.0 * abs(Z), 4.1)) - indicator = Function(V).interpolate(c > 4.4) # 1 inside shape, 0 outside + # Z, _ = SpatialCoordinate(mesh) + # c_background = Function(V).interpolate(Min(1.5 + 4.0 * abs(Z), 4.1)) + indicator = Function(V).interpolate(c > 3.9) # 1 inside shape, 0 outside + gradc = Function(VF, name="grad_c").interpolate( indicator * grad(c_salt) + (1 - indicator) * grad(c_background) ) + # File("gradc.pvd").write(gradc) # ------------------------------------------------------- m1 = ((u - 2.0 * u_n + u_nm1) / Constant(dt ** 2)) * v * dx(rule=qr_x)