From 4510cfc816e42ce4ec5efaa1c2815f4a7706b719 Mon Sep 17 00:00:00 2001 From: Iglesia Dolci Date: Tue, 3 Sep 2024 17:07:53 +0100 Subject: [PATCH 01/13] Forward solver adapted for automatic differentiation. --- .../run_forward_ad.py | 118 ++++++ .../with_automatic_differentiation/run_fwi.py | 183 ++++++++++ spyro/solvers/forward_AD.py | 336 ++++++++---------- 3 files changed, 457 insertions(+), 180 deletions(-) create mode 100644 demos/with_automatic_differentiation/run_forward_ad.py create mode 100644 demos/with_automatic_differentiation/run_fwi.py diff --git a/demos/with_automatic_differentiation/run_forward_ad.py b/demos/with_automatic_differentiation/run_forward_ad.py new file mode 100644 index 00000000..ffed5404 --- /dev/null +++ b/demos/with_automatic_differentiation/run_forward_ad.py @@ -0,0 +1,118 @@ +from firedrake import * +import spyro +import matplotlib.pyplot as plt +import numpy as np + +import spyro.solvers + +# --- Basid setup to run a forward simulation with AD --- # +model = {} + +model["opts"] = { + "method": "KMV", # either CG or KMV + "quadrature": "KMV", # Equi or KMV + "degree": 1, # p order + "dimension": 2, # dimension + "regularization": False, # regularization is on? + "gamma": 1e-5, # regularization parameter +} + +model["parallelism"] = { + # options: + # `shots_parallelism`. Shots parallelism. + # None - no shots parallelism. + "type": "shots_parallelism", + "num_spacial_cores": 1, # Number of cores to use in the spatial parallelism. +} + +# Define the domain size without the ABL. +model["mesh"] = { + "Lz": 1.0, # depth in km - always positive + "Lx": 1.0, # width in km - always positive + "Ly": 0.0, # thickness in km - always positive + "meshfile": "not_used.msh", + "initmodel": "not_used.hdf5", + "truemodel": "not_used.hdf5", +} + +# Specify a 250-m Absorbing Boundary Layer (ABL) on the three sides of the domain to damp outgoing waves. +model["BCs"] = { + "status": False, # True or False, used to turn on any type of BC + "outer_bc": "non-reflective", # none or non-reflective (outer boundary condition) + "abl_bc": "none", # none, gaussian-taper, or alid + "lz": 0.0, # thickness of the ABL in the z-direction (km) - always positive + "lx": 0.0, # thickness of the ABL in the x-direction (km) - always positive + "ly": 0.0, # thickness of the ABL in the y-direction (km) - always positive +} + +model["acquisition"] = { + "source_type": "Ricker", + "source_pos": spyro.create_transect((0.2, 0.15), (0.8, 0.15), 3), + "frequency": 7.0, + "delay": 1.0, + "receiver_locations": spyro.create_transect((0.2, 0.2), (0.8, 0.2), 10), +} +model["aut_dif"] = { + "status": True, +} + +model["timeaxis"] = { + "t0": 0.0, # Initial time for event + "tf": 0.8, # Final time for event (for test 7) + "dt": 0.001, # timestep size (divided by 2 in the test 4. dt for test 3 is 0.00050) + "amplitude": 1, # the Ricker has an amplitude of 1. + "nspool": 20, # (20 for dt=0.00050) how frequently to output solution to pvds + "fspool": 1, # how frequently to save solution to RAM +} + + +# Use emsemble parallelism. +M = model["parallelism"]["num_spacial_cores"] +my_ensemble = Ensemble(COMM_WORLD, M) +mesh = UnitSquareMesh(50, 50, comm=my_ensemble.comm) +element = spyro.domains.space.FE_method( + mesh, model["opts"]["method"], model["opts"]["degree"] +) +V = FunctionSpace(mesh, element) + + +def make_vp_circle(vp_guess=False, plot_vp=False): + """Acoustic velocity model""" + x, z = SpatialCoordinate(mesh) + if vp_guess: + vp = Function(V).interpolate(1.5 + 0.0 * x) + else: + vp = Function(V).interpolate( + 2.5 + + 1 * tanh(100 * (0.125 - sqrt((x - 0.5) ** 2 + (z - 0.5) ** 2))) + ) + if plot_vp: + outfile = File("acoustic_cp.pvd") + outfile.write(vp) + return vp + + +forward_solver = spyro.solvers.forward_ad.ForwardSolver( + model, mesh +) + +c_true = make_vp_circle() +# Ricker wavelet +wavelet = spyro.full_ricker_wavelet( + dt=model["timeaxis"]["dt"], + tf=model["timeaxis"]["tf"], + freq=model["acquisition"]["frequency"], +) + +if model["parallelism"]["type"] is None: + for sn in range(len(model["acquisition"]["source_pos"])): + rec_data, _ = forward_solver.execute(c_true, sn, wavelet) + spyro.plots.plot_shots( + model, my_ensemble.comm, rec_data, vmax=1e-08, vmin=-1e-08) +else: + # source_number based on the ensemble.ensemble_comm.rank + source_number = my_ensemble.ensemble_comm.rank + rec_data, _ = forward_solver.execute_acoustic( + c_true, source_number, wavelet) + spyro.plots.plot_shots( + model, my_ensemble.comm, rec_data, vmax=1e-08, vmin=-1e-08) diff --git a/demos/with_automatic_differentiation/run_fwi.py b/demos/with_automatic_differentiation/run_fwi.py new file mode 100644 index 00000000..e525f89c --- /dev/null +++ b/demos/with_automatic_differentiation/run_fwi.py @@ -0,0 +1,183 @@ +from firedrake import * +from firedrake.adjoint import * +from checkpoint_schedules import Revolve, SingleMemoryStorageSchedule +import spyro +import matplotlib.pyplot as plt +import numpy as np +from scipy.optimize import minimize as scipy_minimize +from mpi4py import MPI +# clear cache constantly to measure memory usage +import gc +import warnings +warnings.filterwarnings("ignore") + +# --- Basid setup to run a forward simulation with AD --- # +model = {} + +model["opts"] = { + "method": "KMV", # either CG or KMV + "quadrature": "KMV", # Equi or KMV + "degree": 1, # p order + "dimension": 2, # dimension + "regularization": False, # regularization is on? + "gamma": 1e-5, # regularization parameter +} + +model["parallelism"] = { + # options: + # `"shots_parallelism"` (same number of cores for every processor. Apply only + # shots parallelism, i.e., the spatial domain is not parallelised.) + # `"automatic"` (same number of cores for every processor. Apply shots and + # spatial parallelism.) + # `"spatial"` (Only spatial parallelisation). + # `None` (No parallelisation). + "type": "None", +} + +# Define the domain size without the ABL. +model["mesh"] = { + "Lz": 1.0, # depth in km - always positive + "Lx": 1.0, # width in km - always positive + "Ly": 0.0, # thickness in km - always positive + "meshfile": "not_used.msh", + "initmodel": "not_used.hdf5", + "truemodel": "not_used.hdf5", +} + +# Specify a 250-m Absorbing Boundary Layer (ABL) on the three sides of the domain to damp outgoing waves. +model["BCs"] = { + "status": False, # True or False, used to turn on any type of BC + "outer_bc": "non-reflective", # none or non-reflective (outer boundary condition) + "abl_bc": "none", # none, gaussian-taper, or alid + "lz": 0.0, # thickness of the ABL in the z-direction (km) - always positive + "lx": 0.0, # thickness of the ABL in the x-direction (km) - always positive + "ly": 0.0, # thickness of the ABL in the y-direction (km) - always positive +} + +model["acquisition"] = { + "source_type": "Ricker", + "source_pos": spyro.create_transect((0.3, 0.15), (0.7, 0.15), 1), + "frequency": 7.0, + "delay": 1.0, + "receiver_locations": spyro.create_transect((0.2, 0.8), (0.8, 0.8), 10), +} +model["aut_dif"] = { + "status": True, +} + +model["timeaxis"] = { + "t0": 0.0, # Initial time for event + "tf": 0.8, # Final time for event (for test 7) + "dt": 0.002, # timestep size (divided by 2 in the test 4. dt for test 3 is 0.00050) + "amplitude": 1, # the Ricker has an amplitude of 1. + "nspool": 20, # (20 for dt=0.00050) how frequently to output solution to pvds + "fspool": 1, # how frequently to save solution to RAM +} + + +def make_vp_circle(vp_guess=False, plot_vp=False): + """Acoustic velocity model""" + x, z = SpatialCoordinate(mesh) + if vp_guess: + vp = Function(V).interpolate(1.5) + else: + vp = Function(V).interpolate( + 2.5 + + 1 * tanh(100 * (0.125 - sqrt((x - 0.5) ** 2 + (z - 0.5) ** 2))) + ) + if plot_vp: + outfile = File("acoustic_cp.pvd") + outfile.write(vp) + return vp + + +true_receiver_data = [] +iterations = 0 +nt = int(model["timeaxis"]["tf"] / model["timeaxis"]["dt"]) # number of timesteps +def J(mesh, comm, vp_exact, wavelet, vom, source_number, vp_guess, + plot_receiver_data=False): + global true_receiver_data + guess_receiver_data, functional = spyro.solvers.forward_AD( + model, mesh, comm, vp_guess, wavelet, vom, source_number=source_number, + true_receiver_data=true_receiver_data[source_number], fwi=True + ) + if plot_receiver_data: + data = [rec.dat.data_ro[:] for rec in guess_receiver_data] + spyro.plots.plot_shots(model, comm, data, vmax=1e-08, vmin=-1e-08) + return functional + +save_J = [] +def run_fwi(vp_guess_data): + global iterations + if iterations == 0: + true_data = [] + for sn in range(len(model["acquisition"]["source_pos"])): + true_receiver_data.append(spyro.solvers.forward_AD(model, mesh, comm, vp_exact, + wavelet, vom, debug=True, + source_number=sn)) + J_total = 0.0 + dJ_total = Function(V) + vp_guess = Function(V) + vp_guess.dat.data[:] = vp_guess_data + File("vp_end" + str(iterations) + ".pvd").write(vp_guess) + if size == 1: + for sn in range(len(model["acquisition"]["source_pos"])): + continue_annotation() + tape = get_working_tape() + tape.progress_bar = ProgressBar + get_working_tape().enable_checkpointing(SingleMemoryStorageSchedule()) + Js = J(mesh, comm, vp_exact, wavelet, vom, sn, vp_guess) + print(Js) + dJ_total += compute_gradient(Js, Control(vp_guess)) + J_total += Js + get_working_tape().clear_tape() + elif size == len(model["acquisition"]["source_pos"]): + J_local = J(mesh, comm, vp_exact, wavelet, vom, rank, vp_guess) + dJ_local = compute_gradient(J_local, Control(vp_guess)) + J_total = COMM_WORLD.allreduce(J_local, op=MPI.SUM) + dJ_total = comm.allreduce(dJ_local, dJ_total, op=MPI.SUM) + else: + raise NotImplementedError("`size` must be 1 or equal to `num_sources`." + "Different values are not supported yet.") + iterations += 1 + return J_total, dJ_total.dat.data[:] + + +comm, spatial_comm = spyro.utils.mpi_init(model) +if model["parallelism"]["type"] == "shots_parallelism": + # Only shots parallelism. + mesh = UnitSquareMesh(60, 60, comm=spatial_comm) +else: + mesh = UnitSquareMesh(60, 60) + +element = spyro.domains.space.FE_method( + mesh, model["opts"]["method"], model["opts"]["degree"] +) + +V = FunctionSpace(mesh, element) +# Receiver mesh. +vom = VertexOnlyMesh(mesh, model["acquisition"]["receiver_locations"]) + +wavelet = spyro.full_ricker_wavelet( + dt=model["timeaxis"]["dt"], + tf=model["timeaxis"]["tf"], + freq=model["acquisition"]["frequency"], +) +# True acoustic velocity model +vp_exact = make_vp_circle(plot_vp=True) +vp_guess = make_vp_circle(plot_vp=True, vp_guess=True) +# Processor number. +rank = comm.ensemble_comm.rank +# Number of processors used in the simulation. +size = comm.ensemble_comm.size +vmax = 3.5 +vmin = 1.5 +vp_0 = vp_guess.vector().gather() +bounds = [(vmin, vmax) for _ in range(len(vp_0))] +result_data = scipy_minimize(run_fwi, vp_0, method='L-BFGS-B', + jac=True, tol=1e-15, bounds=bounds, + options={"disp": True, "eps": 1e-15, + "gtol": 1e-15, "maxiter": 20}) +vp_end = Function(V) +vp_end.dat.data[:] = result_data.x +File("vp_end.pvd").write(vp_end) \ No newline at end of file diff --git a/spyro/solvers/forward_AD.py b/spyro/solvers/forward_AD.py index a8e257f3..6bfb0c81 100644 --- a/spyro/solvers/forward_AD.py +++ b/spyro/solvers/forward_AD.py @@ -1,180 +1,156 @@ -# from firedrake import * - -# # from .. import utils -# from ..domains import quadrature, space - -# # from ..pml import damping -# # from ..io import ensemble_forward -# from . import helpers - -# # Note this turns off non-fatal warnings -# set_log_level(ERROR) - - -# # @ensemble_forward -# def forward( -# model, -# mesh, -# comm, -# c, -# excitations, -# wavelet, -# receivers, -# source_num=0, -# output=False, -# **kwargs -# ): -# """Secord-order in time fully-explicit scheme -# with implementation of a Perfectly Matched Layer (PML) using -# CG FEM with or without higher order mass lumping (KMV type elements). - -# Parameters -# ---------- -# model: Python `dictionary` -# Contains model options and parameters -# mesh: Firedrake.mesh object -# The 2D/3D triangular mesh -# comm: Firedrake.ensemble_communicator -# The MPI communicator for parallelism -# c: Firedrake.Function -# The velocity model interpolated onto the mesh. -# excitations: A list Firedrake.Functions -# wavelet: array-like -# Time series data that's injected at the source location. -# receivers: A :class:`spyro.Receivers` object. -# Contains the receiver locations and sparse interpolation methods. -# source_num: `int`, optional -# The source number you wish to simulate -# output: `boolean`, optional -# Whether or not to write results to pvd files. - -# Returns -# ------- -# usol: list of Firedrake.Functions -# The full field solution at `fspool` timesteps -# usol_recv: array-like -# The solution interpolated to the receivers at all timesteps - -# """ - -# method = model["opts"]["method"] -# degree = model["opts"]["degree"] -# dim = model["opts"]["dimension"] -# dt = model["timeaxis"]["dt"] -# tf = model["timeaxis"]["tf"] -# nspool = model["timeaxis"]["nspool"] -# nt = int(tf / dt) # number of timesteps -# excitations.current_source = source_num -# params = set_params(method) -# element = space.FE_method(mesh, method, degree) - -# V = FunctionSpace(mesh, element) - -# qr_x, qr_s, _ = quadrature.quadrature_rules(V) - -# if dim == 2: -# z, x = SpatialCoordinate(mesh) -# elif dim == 3: -# z, x, y = SpatialCoordinate(mesh) - -# u = TrialFunction(V) -# v = TestFunction(V) - -# u_nm1 = Function(V) -# u_n = Function(V) -# u_np1 = Function(V) - -# if output: -# outfile = helpers.create_output_file("forward.pvd", comm, source_num) - -# t = 0.0 -# m = 1 / (c * c) -# m1 = ( -# m * ((u - 2.0 * u_n + u_nm1) / Constant(dt**2)) * v * dx(scheme=qr_x) -# ) -# a = dot(grad(u_n), grad(v)) * dx(scheme=qr_x) # explicit -# f = Function(V) -# nf = 0 - -# if model["BCs"]["outer_bc"] == "non-reflective": -# nf = c * ((u_n - u_nm1) / dt) * v * ds(scheme=qr_s) - -# h = CellSize(mesh) -# FF = ( -# m1 + a + nf - (1 / (h / degree * h / degree)) * f * v * dx(scheme=qr_x) -# ) -# X = Function(V) - -# lhs_ = lhs(FF) -# rhs_ = rhs(FF) - -# problem = LinearVariationalProblem(lhs_, rhs_, X) -# solver = LinearVariationalSolver(problem, solver_parameters=params) - -# usol_recv = [] - -# P = FunctionSpace(receivers, "DG", 0) -# interpolator = Interpolator(u_np1, P) -# J0 = 0.0 - -# for step in range(nt): -# excitations.apply_source(f, wavelet[step]) - -# solver.solve() -# u_np1.assign(X) - -# rec = Function(P) -# interpolator.interpolate(output=rec) - -# fwi = kwargs.get("fwi") -# p_true_rec = kwargs.get("true_rec") - -# usol_recv.append(rec.dat.data) - -# if fwi: -# J0 += calc_objective_func(rec, p_true_rec[step], step, dt, P) - -# if step % nspool == 0: -# assert ( -# norm(u_n) < 1 -# ), "Numerical instability. Try reducing dt or building the mesh differently" -# if output: -# outfile.write(u_n, time=t, name="Pressure") -# if t > 0: -# helpers.display_progress(comm, t) - -# u_nm1.assign(u_n) -# u_n.assign(u_np1) - -# t = step * float(dt) - -# if fwi: -# return usol_recv, J0 -# else: -# return usol_recv - - -# def calc_objective_func(p_rec, p_true_rec, IT, dt, P): -# true_rec = Function(P) -# true_rec.dat.data[:] = p_true_rec -# J = 0.5 * assemble(inner(true_rec - p_rec, true_rec - p_rec) * dx) -# return J - - -# def set_params(method): -# if method == "KMV": -# params = {"ksp_type": "preonly", "pc_type": "jacobi"} -# elif ( -# method == "CG" -# and mesh.ufl_cell() != quadrilateral -# and mesh.ufl_cell() != hexahedron -# ): -# params = {"ksp_type": "cg", "pc_type": "jacobi"} -# elif method == "CG" and ( -# mesh.ufl_cell() == quadrilateral or mesh.ufl_cell() == hexahedron -# ): -# params = {"ksp_type": "preonly", "pc_type": "jacobi"} -# else: -# raise ValueError("method is not yet supported") - -# return params +from firedrake import * +from firedrake.adjoint import * +from ..domains import quadrature, space +from firedrake.__future__ import interpolate +import finat +# Note this turns off non-fatal warnings +set_log_level(ERROR) + + +class ForwardSolver: + """Forward solver for the acoustic wave equation. + + This forward solver is prepared to work with the automatic + differentiation. + + Parameters + ---------- + model : dict + Dictionary containing the model parameters. + mesh : Mesh + Firedrake mesh object. + """ + + def __init__(self, model, mesh): + self.model = model + self.mesh = mesh + self.element = space.FE_method( + self.mesh, self.model["opts"]["method"], + self.model["opts"]["degree"] + ) + self.V = FunctionSpace(self.mesh, self.element) + self.receiver_mesh = VertexOnlyMesh( + self.mesh, self.model["acquisition"]["receiver_locations"]) + + def execute_acoustic( + self, c, source_number, wavelet, compute_functional=False, + true_data_receivers=None, annotate=False + ): + """Time-stepping acoustic forward solver. + + Parameters + ---------- + c : firedrake.Function + Velocity field. + source_number : int + Number of the source. This is used to select the source + location. + wavelet : list + Time-dependent wavelet. + compute_functional : bool, optional + Whether to compute the functional. If True, the true receiver + data must be provided. + true_data_receivers : list, optional + True receiver data. This is used to compute the functional. + annotate : bool, optional + Whether to annotate the forward solver. Annotated solvers are + used for automated adjoint computations from automatic + differentiation. + + Returns + ------- + (receiver_data : list, J_val : float) + Receiver data and functional value. + + Raises + ------ + ValueError + If true_data_receivers is not provided when computing the + functional. + """ + if annotate: + continue_annotation() + # RHS + source_function = Cofunction(self.V.dual()) + solver, u_np1, u_n, u_nm1 = self._acoustic_lvs( + c, source_function) + # Sources. + source_mesh = VertexOnlyMesh( + self.mesh, + [self.model["acquisition"]["source_locations"][source_number]] + ) + # Source function space. + V_s = FunctionSpace(source_mesh, "DG", 0) + d_s = Function(V_s) + d_s.assign(1.0) + source_cofunction = assemble(d_s * TestFunction(V_s) * dx) + # Interpolate from the source function space to the velocity function space. + q_s = Cofunction(self.V.dual()).interpolate(source_cofunction) + + # Receivers + V_r = FunctionSpace(self.receiver_mesh, "DG", 0) + # Interpolate object. + interpolate_receivers = interpolate(u_np1, V_r) + + # Time execution. + J_val = 0.0 + receiver_data = [] + for step in range(self.model["timeaxis"]["nt"]): + source_cofunction.assign(wavelet(step) * q_s) + solver.solve() + u_nm1.assign(u_n) + u_n.assign(u_np1) + receiver_data.append(assemble(interpolate_receivers)) + if compute_functional: + if not true_data_receivers: + raise ValueError("True receiver data is required for" + "computing the functional.") + misfit = receiver_data - true_data_receivers[step] + J_val += assemble(0.5 * inner(misfit, misfit) * dx) + + return receiver_data, J_val + + def _acoustic_lvs(self, c, source_function): + # Acoustic linear variational solver. + V = self.V + dt = self.model["timeaxis"]["dt"] + u = TrialFunction(V) + v = TestFunction(V) + u_np1 = Function(V) # timestep n+1 + u_n = Function(V) # timestep n + u_nm1 = Function(V) # timestep n-1 + + qr_x, qr_s, _ = quadrature.quadrature_rules(V) + time_term = (1 / (c * c)) * (u - 2.0 * u_n + u_nm1) / \ + Constant(dt**2) * v * dx(scheme=quad_rule) + + nf = 0 + if self.model["BCs"]["outer_bc"] == "non-reflective": + nf = (1/c) * ((u_n - u_nm1) / dt) * v * ds(scheme=qr_s) + + a = dot(grad(u_n), grad(v)) * dx(scheme=qr_x) + F = time_term + a + nf + lin_var = LinearVariationalProblem( + lhs(F), rhs(F) + source_function, u_np1) + solver = LinearVariationalSolver( + lin_var,solver_parameters=self._solver_parameters()) + return solver, u_np1, u_n, u_nm1 + + def _solver_parameters(self): + if self.model["opts"]["method"] == "KMV": + params = {"ksp_type": "preonly", "pc_type": "jacobi"} + elif ( + self.model["opts"]["method"] == "CG" + and self.mesh.ufl_cell() != quadrilateral + and self.mesh.ufl_cell() != hexahedron + ): + params = {"ksp_type": "cg", "pc_type": "jacobi"} + elif self.model["opts"]["method"] == "CG" and ( + self.mesh.ufl_cell() == quadrilateral + or self.mesh.ufl_cell() == hexahedron + ): + params = {"ksp_type": "preonly", "pc_type": "jacobi"} + else: + raise ValueError("method is not yet supported") + + return params From c8025eeb5d5125c41c1256eeada40df88e47f1f0 Mon Sep 17 00:00:00 2001 From: Iglesia Dolci Date: Tue, 3 Sep 2024 18:46:04 +0100 Subject: [PATCH 02/13] forward working --- .../run_forward_ad.py | 51 +++++------ spyro/solvers/__init__.py | 2 + spyro/solvers/forward_AD.py | 88 +++++++------------ spyro/solvers/time_integration_ad.py | 42 +++++++++ 4 files changed, 104 insertions(+), 79 deletions(-) create mode 100644 spyro/solvers/time_integration_ad.py diff --git a/demos/with_automatic_differentiation/run_forward_ad.py b/demos/with_automatic_differentiation/run_forward_ad.py index ffed5404..b3022eb2 100644 --- a/demos/with_automatic_differentiation/run_forward_ad.py +++ b/demos/with_automatic_differentiation/run_forward_ad.py @@ -1,4 +1,4 @@ -from firedrake import * +import firedrake as fire import spyro import matplotlib.pyplot as plt import numpy as np @@ -9,8 +9,8 @@ model = {} model["opts"] = { - "method": "KMV", # either CG or KMV - "quadrature": "KMV", # Equi or KMV + "method": "KMV", # either CG or mass_lumped_triangle + "quadrature": "KMV", # Equi or mass_lumped_triangle "degree": 1, # p order "dimension": 2, # dimension "regularization": False, # regularization is on? @@ -58,7 +58,7 @@ model["timeaxis"] = { "t0": 0.0, # Initial time for event - "tf": 0.8, # Final time for event (for test 7) + "tf": 0.2, # Final time for event (for test 7) "dt": 0.001, # timestep size (divided by 2 in the test 4. dt for test 3 is 0.00050) "amplitude": 1, # the Ricker has an amplitude of 1. "nspool": 20, # (20 for dt=0.00050) how frequently to output solution to pvds @@ -68,51 +68,52 @@ # Use emsemble parallelism. M = model["parallelism"]["num_spacial_cores"] -my_ensemble = Ensemble(COMM_WORLD, M) -mesh = UnitSquareMesh(50, 50, comm=my_ensemble.comm) -element = spyro.domains.space.FE_method( - mesh, model["opts"]["method"], model["opts"]["degree"] -) -V = FunctionSpace(mesh, element) +my_ensemble = fire.Ensemble(fire.COMM_WORLD, M) +mesh = fire.UnitSquareMesh(50, 50, comm=my_ensemble.comm) +element = fire.FiniteElement( + model["opts"]["method"], mesh.ufl_cell(), degree=model["opts"]["degree"], + variant=model["opts"]["quadrature"] + ) +V = fire.FunctionSpace(mesh, element) def make_vp_circle(vp_guess=False, plot_vp=False): """Acoustic velocity model""" - x, z = SpatialCoordinate(mesh) + x, z = fire.SpatialCoordinate(mesh) if vp_guess: - vp = Function(V).interpolate(1.5 + 0.0 * x) + vp = fire.Function(V).interpolate(1.5 + 0.0 * x) else: - vp = Function(V).interpolate( + vp = fire.Function(V).interpolate( 2.5 - + 1 * tanh(100 * (0.125 - sqrt((x - 0.5) ** 2 + (z - 0.5) ** 2))) + + 1 * fire.tanh(100 * (0.125 - fire.sqrt((x - 0.5) ** 2 + (z - 0.5) ** 2))) ) if plot_vp: - outfile = File("acoustic_cp.pvd") + outfile = fire.File("acoustic_cp.pvd") outfile.write(vp) return vp -forward_solver = spyro.solvers.forward_ad.ForwardSolver( - model, mesh -) +forward_solver = spyro.solvers.forward_ad.ForwardSolver(model, mesh, V) c_true = make_vp_circle() # Ricker wavelet wavelet = spyro.full_ricker_wavelet( - dt=model["timeaxis"]["dt"], - tf=model["timeaxis"]["tf"], - freq=model["acquisition"]["frequency"], + model["timeaxis"]["dt"], model["timeaxis"]["tf"], + model["acquisition"]["frequency"], ) if model["parallelism"]["type"] is None: + outfile = fire.VTKFile("solution.pvd") for sn in range(len(model["acquisition"]["source_pos"])): rec_data, _ = forward_solver.execute(c_true, sn, wavelet) - spyro.plots.plot_shots( - model, my_ensemble.comm, rec_data, vmax=1e-08, vmin=-1e-08) + sol = forward_solver.solution + outfile.write(sol) else: # source_number based on the ensemble.ensemble_comm.rank source_number = my_ensemble.ensemble_comm.rank rec_data, _ = forward_solver.execute_acoustic( c_true, source_number, wavelet) - spyro.plots.plot_shots( - model, my_ensemble.comm, rec_data, vmax=1e-08, vmin=-1e-08) + sol = forward_solver.solution + fire.VTKFile( + "solution_" + str(source_number) + ".pvd", comm=my_ensemble.comm + ).write(sol) diff --git a/spyro/solvers/__init__.py b/spyro/solvers/__init__.py index 5e7a19f6..6b8120ec 100644 --- a/spyro/solvers/__init__.py +++ b/spyro/solvers/__init__.py @@ -2,10 +2,12 @@ from .acoustic_wave import AcousticWave from .mms_acoustic import AcousticWaveMMS from .inversion import FullWaveformInversion +from .forward_ad import ForwardSolver __all__ = [ "Wave", "AcousticWave", "AcousticWaveMMS", "FullWaveformInversion", + "ForwardSolver", ] diff --git a/spyro/solvers/forward_AD.py b/spyro/solvers/forward_AD.py index 6bfb0c81..591436ae 100644 --- a/spyro/solvers/forward_AD.py +++ b/spyro/solvers/forward_AD.py @@ -1,17 +1,18 @@ -from firedrake import * -from firedrake.adjoint import * -from ..domains import quadrature, space +import firedrake as fire +import firedrake.adjoint as fire_ad +from ..domains import quadrature +from .time_integration_ad import central_difference_acoustic from firedrake.__future__ import interpolate import finat # Note this turns off non-fatal warnings -set_log_level(ERROR) +fire.set_log_level(fire.ERROR) class ForwardSolver: - """Forward solver for the acoustic wave equation. + """Wave equation forward solver. This forward solver is prepared to work with the automatic - differentiation. + differentiation. Only the acoustic wave equation is implemented. Parameters ---------- @@ -21,16 +22,13 @@ class ForwardSolver: Firedrake mesh object. """ - def __init__(self, model, mesh): + def __init__(self, model, mesh, function_space): self.model = model self.mesh = mesh - self.element = space.FE_method( - self.mesh, self.model["opts"]["method"], - self.model["opts"]["degree"] - ) - self.V = FunctionSpace(self.mesh, self.element) - self.receiver_mesh = VertexOnlyMesh( + self.V = function_space + self.receiver_mesh = fire.VertexOnlyMesh( self.mesh, self.model["acquisition"]["receiver_locations"]) + self.solution = None def execute_acoustic( self, c, source_number, wavelet, compute_functional=False, @@ -38,6 +36,8 @@ def execute_acoustic( ): """Time-stepping acoustic forward solver. + The time integration is done using a central difference scheme. + Parameters ---------- c : firedrake.Function @@ -68,73 +68,53 @@ def execute_acoustic( If true_data_receivers is not provided when computing the functional. """ + self.solution = None if annotate: - continue_annotation() + fire_ad.continue_annotation() # RHS - source_function = Cofunction(self.V.dual()) - solver, u_np1, u_n, u_nm1 = self._acoustic_lvs( - c, source_function) + source_function = fire.Cofunction(self.V.dual()) + solver, u_np1, u_n, u_nm1 = central_difference_acoustic( + self, c, source_function) # Sources. - source_mesh = VertexOnlyMesh( + source_mesh = fire.VertexOnlyMesh( self.mesh, - [self.model["acquisition"]["source_locations"][source_number]] + [self.model["acquisition"]["source_pos"][source_number]] ) # Source function space. - V_s = FunctionSpace(source_mesh, "DG", 0) - d_s = Function(V_s) + V_s = fire.FunctionSpace(source_mesh, "DG", 0) + d_s = fire.Function(V_s) d_s.assign(1.0) - source_cofunction = assemble(d_s * TestFunction(V_s) * dx) + source_d_s = fire.assemble(d_s * fire.TestFunction(V_s) * fire.dx) # Interpolate from the source function space to the velocity function space. - q_s = Cofunction(self.V.dual()).interpolate(source_cofunction) + q_s = fire.Cofunction(self.V.dual()).interpolate(source_d_s) # Receivers - V_r = FunctionSpace(self.receiver_mesh, "DG", 0) + V_r = fire.FunctionSpace(self.receiver_mesh, "DG", 0) # Interpolate object. interpolate_receivers = interpolate(u_np1, V_r) # Time execution. J_val = 0.0 receiver_data = [] - for step in range(self.model["timeaxis"]["nt"]): - source_cofunction.assign(wavelet(step) * q_s) + total_steps = int(self.model["timeaxis"]["tf"] / self.model["timeaxis"]["dt"]) + for step in range(total_steps): + source_function.assign(wavelet[step] * q_s) solver.solve() u_nm1.assign(u_n) u_n.assign(u_np1) - receiver_data.append(assemble(interpolate_receivers)) + receiver_data.append(fire.assemble(interpolate_receivers)) if compute_functional: if not true_data_receivers: raise ValueError("True receiver data is required for" "computing the functional.") misfit = receiver_data - true_data_receivers[step] - J_val += assemble(0.5 * inner(misfit, misfit) * dx) - + J_val += fire.assemble(0.5 * fire.inner(misfit, misfit) * dx) + self.solution = u_np1 return receiver_data, J_val - def _acoustic_lvs(self, c, source_function): - # Acoustic linear variational solver. - V = self.V - dt = self.model["timeaxis"]["dt"] - u = TrialFunction(V) - v = TestFunction(V) - u_np1 = Function(V) # timestep n+1 - u_n = Function(V) # timestep n - u_nm1 = Function(V) # timestep n-1 - - qr_x, qr_s, _ = quadrature.quadrature_rules(V) - time_term = (1 / (c * c)) * (u - 2.0 * u_n + u_nm1) / \ - Constant(dt**2) * v * dx(scheme=quad_rule) - - nf = 0 - if self.model["BCs"]["outer_bc"] == "non-reflective": - nf = (1/c) * ((u_n - u_nm1) / dt) * v * ds(scheme=qr_s) - - a = dot(grad(u_n), grad(v)) * dx(scheme=qr_x) - F = time_term + a + nf - lin_var = LinearVariationalProblem( - lhs(F), rhs(F) + source_function, u_np1) - solver = LinearVariationalSolver( - lin_var,solver_parameters=self._solver_parameters()) - return solver, u_np1, u_n, u_nm1 + def execute_elastic(self): + raise NotImplementedError("Elastic wave equation is not yet implemented" + "for the automatic differentiation based FWI.") def _solver_parameters(self): if self.model["opts"]["method"] == "KMV": diff --git a/spyro/solvers/time_integration_ad.py b/spyro/solvers/time_integration_ad.py new file mode 100644 index 00000000..72d0c93b --- /dev/null +++ b/spyro/solvers/time_integration_ad.py @@ -0,0 +1,42 @@ +import firedrake as fire +from ..domains import quadrature + + +def central_difference_acoustic(forwardsolver, c, source_function): + """ + Perform central difference time integration for wave propagation. + + Parameters: + ----------- + forwardsolver: Spyro object + The Wave object containing the necessary data and parameters. + + Returns: + -------- + (receiver_data : list, J_val : float) + Receiver data and functional value. + """ + # Acoustic linear variational solver. + V = forwardsolver.V + dt = forwardsolver.model["timeaxis"]["dt"] + u = fire.TrialFunction(V) + v = fire.TestFunction(V) + u_np1 = fire.Function(V) # timestep n+1 + u_n = fire.Function(V) # timestep n + u_nm1 = fire.Function(V) # timestep n-1 + + qr_x, qr_s, _ = quadrature.quadrature_rules(V) + time_term = (1 / (c * c)) * (u - 2.0 * u_n + u_nm1) / \ + fire.Constant(dt**2) * v * fire.dx(scheme=qr_x) + + nf = 0 + if forwardsolver.model["BCs"]["outer_bc"] == "non-reflective": + nf = (1/c) * ((u_n - u_nm1) / dt) * v * fire.ds(scheme=qr_s) + + a = fire.dot(fire.grad(u_n), fire.grad(v)) * fire.dx(scheme=qr_x) + F = time_term + a + nf + lin_var = fire.LinearVariationalProblem( + fire.lhs(F), fire.rhs(F) + source_function, u_np1) + solver = fire.LinearVariationalSolver( + lin_var, solver_parameters=forwardsolver._solver_parameters()) + return solver, u_np1, u_n, u_nm1 From 803a60da53af2ab681e8dca20977bdf6d7f9f528 Mon Sep 17 00:00:00 2001 From: Iglesia Dolci Date: Wed, 4 Sep 2024 07:08:50 +0100 Subject: [PATCH 03/13] testing 2D; fwi demo working --- .../run_forward_ad.py | 24 +-- .../with_automatic_differentiation/run_fwi.py | 183 ------------------ .../run_fwi_ad.py | 153 +++++++++++++++ spyro/solvers/forward_AD.py | 15 +- spyro/tools/gradient_test_ad.py | 149 -------------- test/test_gradient_ad.py | 131 +++++++++++++ test_ad/__init__.py | 0 test_ad/test_gradient_3d_AD.py | 86 -------- test_ad/test_gradient_AD.py | 77 -------- 9 files changed, 301 insertions(+), 517 deletions(-) delete mode 100644 demos/with_automatic_differentiation/run_fwi.py create mode 100644 demos/with_automatic_differentiation/run_fwi_ad.py delete mode 100644 spyro/tools/gradient_test_ad.py create mode 100644 test/test_gradient_ad.py delete mode 100644 test_ad/__init__.py delete mode 100644 test_ad/test_gradient_3d_AD.py delete mode 100644 test_ad/test_gradient_AD.py diff --git a/demos/with_automatic_differentiation/run_forward_ad.py b/demos/with_automatic_differentiation/run_forward_ad.py index b3022eb2..1c595070 100644 --- a/demos/with_automatic_differentiation/run_forward_ad.py +++ b/demos/with_automatic_differentiation/run_forward_ad.py @@ -66,17 +66,6 @@ } -# Use emsemble parallelism. -M = model["parallelism"]["num_spacial_cores"] -my_ensemble = fire.Ensemble(fire.COMM_WORLD, M) -mesh = fire.UnitSquareMesh(50, 50, comm=my_ensemble.comm) -element = fire.FiniteElement( - model["opts"]["method"], mesh.ufl_cell(), degree=model["opts"]["degree"], - variant=model["opts"]["quadrature"] - ) -V = fire.FunctionSpace(mesh, element) - - def make_vp_circle(vp_guess=False, plot_vp=False): """Acoustic velocity model""" x, z = fire.SpatialCoordinate(mesh) @@ -88,11 +77,22 @@ def make_vp_circle(vp_guess=False, plot_vp=False): + 1 * fire.tanh(100 * (0.125 - fire.sqrt((x - 0.5) ** 2 + (z - 0.5) ** 2))) ) if plot_vp: - outfile = fire.File("acoustic_cp.pvd") + outfile = fire.VTKFile("acoustic_cp.pvd") outfile.write(vp) return vp +# Use emsemble parallelism. +M = model["parallelism"]["num_spacial_cores"] +my_ensemble = fire.Ensemble(fire.COMM_WORLD, M) +mesh = fire.UnitSquareMesh(50, 50, comm=my_ensemble.comm) +element = fire.FiniteElement( + model["opts"]["method"], mesh.ufl_cell(), degree=model["opts"]["degree"], + variant=model["opts"]["quadrature"] + ) +V = fire.FunctionSpace(mesh, element) + + forward_solver = spyro.solvers.forward_ad.ForwardSolver(model, mesh, V) c_true = make_vp_circle() diff --git a/demos/with_automatic_differentiation/run_fwi.py b/demos/with_automatic_differentiation/run_fwi.py deleted file mode 100644 index e525f89c..00000000 --- a/demos/with_automatic_differentiation/run_fwi.py +++ /dev/null @@ -1,183 +0,0 @@ -from firedrake import * -from firedrake.adjoint import * -from checkpoint_schedules import Revolve, SingleMemoryStorageSchedule -import spyro -import matplotlib.pyplot as plt -import numpy as np -from scipy.optimize import minimize as scipy_minimize -from mpi4py import MPI -# clear cache constantly to measure memory usage -import gc -import warnings -warnings.filterwarnings("ignore") - -# --- Basid setup to run a forward simulation with AD --- # -model = {} - -model["opts"] = { - "method": "KMV", # either CG or KMV - "quadrature": "KMV", # Equi or KMV - "degree": 1, # p order - "dimension": 2, # dimension - "regularization": False, # regularization is on? - "gamma": 1e-5, # regularization parameter -} - -model["parallelism"] = { - # options: - # `"shots_parallelism"` (same number of cores for every processor. Apply only - # shots parallelism, i.e., the spatial domain is not parallelised.) - # `"automatic"` (same number of cores for every processor. Apply shots and - # spatial parallelism.) - # `"spatial"` (Only spatial parallelisation). - # `None` (No parallelisation). - "type": "None", -} - -# Define the domain size without the ABL. -model["mesh"] = { - "Lz": 1.0, # depth in km - always positive - "Lx": 1.0, # width in km - always positive - "Ly": 0.0, # thickness in km - always positive - "meshfile": "not_used.msh", - "initmodel": "not_used.hdf5", - "truemodel": "not_used.hdf5", -} - -# Specify a 250-m Absorbing Boundary Layer (ABL) on the three sides of the domain to damp outgoing waves. -model["BCs"] = { - "status": False, # True or False, used to turn on any type of BC - "outer_bc": "non-reflective", # none or non-reflective (outer boundary condition) - "abl_bc": "none", # none, gaussian-taper, or alid - "lz": 0.0, # thickness of the ABL in the z-direction (km) - always positive - "lx": 0.0, # thickness of the ABL in the x-direction (km) - always positive - "ly": 0.0, # thickness of the ABL in the y-direction (km) - always positive -} - -model["acquisition"] = { - "source_type": "Ricker", - "source_pos": spyro.create_transect((0.3, 0.15), (0.7, 0.15), 1), - "frequency": 7.0, - "delay": 1.0, - "receiver_locations": spyro.create_transect((0.2, 0.8), (0.8, 0.8), 10), -} -model["aut_dif"] = { - "status": True, -} - -model["timeaxis"] = { - "t0": 0.0, # Initial time for event - "tf": 0.8, # Final time for event (for test 7) - "dt": 0.002, # timestep size (divided by 2 in the test 4. dt for test 3 is 0.00050) - "amplitude": 1, # the Ricker has an amplitude of 1. - "nspool": 20, # (20 for dt=0.00050) how frequently to output solution to pvds - "fspool": 1, # how frequently to save solution to RAM -} - - -def make_vp_circle(vp_guess=False, plot_vp=False): - """Acoustic velocity model""" - x, z = SpatialCoordinate(mesh) - if vp_guess: - vp = Function(V).interpolate(1.5) - else: - vp = Function(V).interpolate( - 2.5 - + 1 * tanh(100 * (0.125 - sqrt((x - 0.5) ** 2 + (z - 0.5) ** 2))) - ) - if plot_vp: - outfile = File("acoustic_cp.pvd") - outfile.write(vp) - return vp - - -true_receiver_data = [] -iterations = 0 -nt = int(model["timeaxis"]["tf"] / model["timeaxis"]["dt"]) # number of timesteps -def J(mesh, comm, vp_exact, wavelet, vom, source_number, vp_guess, - plot_receiver_data=False): - global true_receiver_data - guess_receiver_data, functional = spyro.solvers.forward_AD( - model, mesh, comm, vp_guess, wavelet, vom, source_number=source_number, - true_receiver_data=true_receiver_data[source_number], fwi=True - ) - if plot_receiver_data: - data = [rec.dat.data_ro[:] for rec in guess_receiver_data] - spyro.plots.plot_shots(model, comm, data, vmax=1e-08, vmin=-1e-08) - return functional - -save_J = [] -def run_fwi(vp_guess_data): - global iterations - if iterations == 0: - true_data = [] - for sn in range(len(model["acquisition"]["source_pos"])): - true_receiver_data.append(spyro.solvers.forward_AD(model, mesh, comm, vp_exact, - wavelet, vom, debug=True, - source_number=sn)) - J_total = 0.0 - dJ_total = Function(V) - vp_guess = Function(V) - vp_guess.dat.data[:] = vp_guess_data - File("vp_end" + str(iterations) + ".pvd").write(vp_guess) - if size == 1: - for sn in range(len(model["acquisition"]["source_pos"])): - continue_annotation() - tape = get_working_tape() - tape.progress_bar = ProgressBar - get_working_tape().enable_checkpointing(SingleMemoryStorageSchedule()) - Js = J(mesh, comm, vp_exact, wavelet, vom, sn, vp_guess) - print(Js) - dJ_total += compute_gradient(Js, Control(vp_guess)) - J_total += Js - get_working_tape().clear_tape() - elif size == len(model["acquisition"]["source_pos"]): - J_local = J(mesh, comm, vp_exact, wavelet, vom, rank, vp_guess) - dJ_local = compute_gradient(J_local, Control(vp_guess)) - J_total = COMM_WORLD.allreduce(J_local, op=MPI.SUM) - dJ_total = comm.allreduce(dJ_local, dJ_total, op=MPI.SUM) - else: - raise NotImplementedError("`size` must be 1 or equal to `num_sources`." - "Different values are not supported yet.") - iterations += 1 - return J_total, dJ_total.dat.data[:] - - -comm, spatial_comm = spyro.utils.mpi_init(model) -if model["parallelism"]["type"] == "shots_parallelism": - # Only shots parallelism. - mesh = UnitSquareMesh(60, 60, comm=spatial_comm) -else: - mesh = UnitSquareMesh(60, 60) - -element = spyro.domains.space.FE_method( - mesh, model["opts"]["method"], model["opts"]["degree"] -) - -V = FunctionSpace(mesh, element) -# Receiver mesh. -vom = VertexOnlyMesh(mesh, model["acquisition"]["receiver_locations"]) - -wavelet = spyro.full_ricker_wavelet( - dt=model["timeaxis"]["dt"], - tf=model["timeaxis"]["tf"], - freq=model["acquisition"]["frequency"], -) -# True acoustic velocity model -vp_exact = make_vp_circle(plot_vp=True) -vp_guess = make_vp_circle(plot_vp=True, vp_guess=True) -# Processor number. -rank = comm.ensemble_comm.rank -# Number of processors used in the simulation. -size = comm.ensemble_comm.size -vmax = 3.5 -vmin = 1.5 -vp_0 = vp_guess.vector().gather() -bounds = [(vmin, vmax) for _ in range(len(vp_0))] -result_data = scipy_minimize(run_fwi, vp_0, method='L-BFGS-B', - jac=True, tol=1e-15, bounds=bounds, - options={"disp": True, "eps": 1e-15, - "gtol": 1e-15, "maxiter": 20}) -vp_end = Function(V) -vp_end.dat.data[:] = result_data.x -File("vp_end.pvd").write(vp_end) \ No newline at end of file diff --git a/demos/with_automatic_differentiation/run_fwi_ad.py b/demos/with_automatic_differentiation/run_fwi_ad.py new file mode 100644 index 00000000..232dcd19 --- /dev/null +++ b/demos/with_automatic_differentiation/run_fwi_ad.py @@ -0,0 +1,153 @@ +import firedrake as fire +import firedrake.adjoint as fire_ad +import spyro + + +# --- Basid setup to run a forward simulation with AD --- # +model = {} + +model["opts"] = { + "method": "KMV", # either CG or mass_lumped_triangle + "quadrature": "KMV", # Equi or mass_lumped_triangle + "degree": 1, # p order + "dimension": 2, # dimension + "regularization": False, # regularization is on? + "gamma": 1e-5, # regularization parameter +} + +model["parallelism"] = { + # options: + # `shots_parallelism`. Shots parallelism. + # None - no shots parallelism. + "type": "shots_parallelism", + "num_spacial_cores": 1, # Number of cores to use in the spatial + # parallelism. +} + +# Define the domain size without the ABL. +model["mesh"] = { + "Lz": 1.0, # depth in km - always positive + "Lx": 1.0, # width in km - always positive + "Ly": 0.0, # thickness in km - always positive + "meshfile": "not_used.msh", + "initmodel": "not_used.hdf5", + "truemodel": "not_used.hdf5", +} + +# Specify a 250-m Absorbing Boundary Layer (ABL) on the three sides of the domain to damp outgoing waves. +model["BCs"] = { + "status": False, # True or False, used to turn on any type of BC + "outer_bc": "non-reflective", # none or non-reflective (outer boundary condition) + "abl_bc": "none", # none, gaussian-taper, or alid + "lz": 0.0, # thickness of the ABL in the z-direction (km) - always positive + "lx": 0.0, # thickness of the ABL in the x-direction (km) - always positive + "ly": 0.0, # thickness of the ABL in the y-direction (km) - always positive +} + +model["acquisition"] = { + "source_type": "Ricker", + "source_pos": spyro.create_transect((0.2, 0.15), (0.8, 0.15), 3), + "frequency": 7.0, + "delay": 1.0, + "receiver_locations": spyro.create_transect((0.2, 0.2), (0.8, 0.2), 10), +} +model["aut_dif"] = { + "status": True, +} + +model["timeaxis"] = { + "t0": 0.0, # Initial time for event + "tf": 0.6, # Final time for event (for test 7) + "dt": 0.001, # timestep size (divided by 2 in the test 4. dt for test 3 is 0.00050) + "amplitude": 1, # the Ricker has an amplitude of 1. + "nspool": 20, # (20 for dt=0.00050) how frequently to output solution to pvds + "fspool": 1, # how frequently to save solution to RAM +} + + +def make_c_camembert(c_guess=False, plot_c=False): + """Acoustic velocity model""" + x, z = fire.SpatialCoordinate(mesh) + if c_guess: + c = fire.Function(V).interpolate(1.5 + 0.0 * x) + else: + c = fire.Function(V).interpolate( + 2.5 + + 1 * fire.tanh(100 * (0.125 - fire.sqrt((x - 0.5) ** 2 + (z - 0.5) ** 2))) + ) + if plot_c: + outfile = fire.VTKFile("acoustic_cp.pvd") + outfile.write(c) + return c + + +def forward( + c, compute_functional=False, true_data_receivers=None, annotate=False +): + if annotate: + fire_ad.continue_annotation() + if model["parallelism"]["type"] is None: + outfile = fire.VTKFile("solution.pvd") + receiver_data = [] + J = 0.0 + for sn in range(len(model["acquisition"]["source_pos"])): + rec_data, J_val = forward_solver.execute(c, sn, wavelet) + receiver_data.append(rec_data) + J += J_val + sol = forward_solver.solution + outfile.write(sol) + + else: + # source_number based on the ensemble.ensemble_comm.rank + source_number = my_ensemble.ensemble_comm.rank + receiver_data, J = forward_solver.execute_acoustic( + c, source_number, wavelet, + compute_functional=compute_functional, + true_data_receivers=true_data_receivers + ) + sol = forward_solver.solution + fire.VTKFile( + "solution_" + str(source_number) + ".pvd", comm=my_ensemble.comm + ).write(sol) + + return receiver_data, J + + +# Use emsemble parallelism. +M = model["parallelism"]["num_spacial_cores"] +my_ensemble = fire.Ensemble(fire.COMM_WORLD, M) +mesh = fire.UnitSquareMesh(50, 50, comm=my_ensemble.comm) +element = fire.FiniteElement( + model["opts"]["method"], mesh.ufl_cell(), degree=model["opts"]["degree"], + variant=model["opts"]["quadrature"] + ) +V = fire.FunctionSpace(mesh, element) + + +forward_solver = spyro.solvers.forward_ad.ForwardSolver(model, mesh, V) +c_true = make_c_camembert() +# Ricker wavelet +wavelet = spyro.full_ricker_wavelet( + model["timeaxis"]["dt"], model["timeaxis"]["tf"], + model["acquisition"]["frequency"], +) + +true_rec, _ = forward(c_true) + +# --- FWI with AD --- # +c_guess = make_c_camembert(c_guess=True) +guess_rec, J = forward( + c_guess, compute_functional=True, true_data_receivers=true_rec, + annotate=True + ) + +# :class:`~.EnsembleReducedFunctional` is employed to recompute in +# parallel the functional and its gradient associated with the multiple sources +# (3 in this case). +J_hat = fire_ad.EnsembleReducedFunctional( + J, fire_ad.Control(c_guess), my_ensemble) +fire_ad.taylor_test(J_hat, c_guess, fire.Function(V).assign(1.0)) +c_optimised = fire_ad.minimize(J_hat, method="L-BFGS-B", + options={"disp": True, "maxiter": 10}, + bounds=(1.5, 3.5), + derivative_options={"riesz_representation": 'l2'}) \ No newline at end of file diff --git a/spyro/solvers/forward_AD.py b/spyro/solvers/forward_AD.py index 591436ae..9d3f5f03 100644 --- a/spyro/solvers/forward_AD.py +++ b/spyro/solvers/forward_AD.py @@ -32,7 +32,7 @@ def __init__(self, model, mesh, function_space): def execute_acoustic( self, c, source_number, wavelet, compute_functional=False, - true_data_receivers=None, annotate=False + true_data_receivers=None ): """Time-stepping acoustic forward solver. @@ -52,10 +52,6 @@ def execute_acoustic( data must be provided. true_data_receivers : list, optional True receiver data. This is used to compute the functional. - annotate : bool, optional - Whether to annotate the forward solver. Annotated solvers are - used for automated adjoint computations from automatic - differentiation. Returns ------- @@ -69,8 +65,6 @@ def execute_acoustic( functional. """ self.solution = None - if annotate: - fire_ad.continue_annotation() # RHS source_function = fire.Cofunction(self.V.dual()) solver, u_np1, u_n, u_nm1 = central_difference_acoustic( @@ -102,13 +96,14 @@ def execute_acoustic( solver.solve() u_nm1.assign(u_n) u_n.assign(u_np1) - receiver_data.append(fire.assemble(interpolate_receivers)) + rec_data = fire.assemble(interpolate_receivers) + receiver_data.append(rec_data) if compute_functional: if not true_data_receivers: raise ValueError("True receiver data is required for" "computing the functional.") - misfit = receiver_data - true_data_receivers[step] - J_val += fire.assemble(0.5 * fire.inner(misfit, misfit) * dx) + misfit = rec_data - true_data_receivers[step] + J_val += fire.assemble(0.5 * fire.inner(misfit, misfit) * fire.dx) self.solution = u_np1 return receiver_data, J_val diff --git a/spyro/tools/gradient_test_ad.py b/spyro/tools/gradient_test_ad.py deleted file mode 100644 index ab159962..00000000 --- a/spyro/tools/gradient_test_ad.py +++ /dev/null @@ -1,149 +0,0 @@ -# import numpy as np -# from firedrake import * -# from pyadjoint import enlisting -# import spyro - -# forward = spyro.solvers.forward_AD - - -# def gradient_test_acoustic( -# model, mesh, V, comm, vp_exact, vp_guess, mask=None -# ): -# """Gradient test for the acoustic FWI problem - -# Parameters -# ---------- -# model : `dictionary` -# Contains simulation parameters and options. -# mesh : a Firedrake.mesh -# 2D/3D simplicial mesh read in by Firedrake.Mesh -# V : Firedrake.FunctionSpace object -# The space of the finite elements -# comm : Firedrake.ensemble_communicator -# An ensemble communicator -# vp_exact : Firedrake.Function -# The exact velocity model -# vp_guess : Firedrake.Function -# The guess velocity model -# mask : Firedrake.Function, optional -# A mask for the gradient test - -# Returns -# ------- -# None -# """ -# import firedrake_adjoint as fire_adj - -# with fire_adj.stop_annotating(): -# if comm.comm.rank == 0: -# print("######## Starting gradient test ########", flush=True) - -# sources = spyro.Sources(model, mesh, V, comm) -# receivers = spyro.Receivers(model, mesh, V, comm) - -# wavelet = spyro.full_ricker_wavelet( -# model["timeaxis"]["dt"], -# model["timeaxis"]["tf"], -# model["acquisition"]["frequency"], -# ) -# point_cloud = receivers.set_point_cloud(comm) -# # simulate the exact model -# if comm.comm.rank == 0: -# print("######## Running the exact model ########", flush=True) -# p_exact_recv = forward( -# model, mesh, comm, vp_exact, sources, wavelet, point_cloud -# ) - -# # simulate the guess model -# if comm.comm.rank == 0: -# print("######## Running the guess model ########", flush=True) -# p_guess_recv, Jm = forward( -# model, -# mesh, -# comm, -# vp_guess, -# sources, -# wavelet, -# point_cloud, -# fwi=True, -# true_rec=p_exact_recv, -# ) -# if comm.comm.rank == 0: -# print( -# "\n Cost functional at fixed point : " + str(Jm) + " \n ", -# flush=True, -# ) - -# # compute the gradient of the control (to be verified) -# if comm.comm.rank == 0: -# print( -# "######## Computing the gradient by automatic differentiation ########", -# flush=True, -# ) -# control = fire_adj.Control(vp_guess) -# dJ = fire_adj.compute_gradient(Jm, control) -# if mask: -# dJ *= mask - -# # File("gradient.pvd").write(dJ) - -# # steps = [1e-3, 1e-4, 1e-5, 1e-6, 1e-7] # step length -# # steps = [1e-4, 1e-5, 1e-6, 1e-7] # step length -# steps = [1e-5, 1e-6, 1e-7, 1e-8] # step length -# with fire_adj.stop_annotating(): -# delta_m = Function(V) # model direction (random) -# delta_m.assign(dJ) -# Jhat = fire_adj.ReducedFunctional(Jm, control) -# derivative = enlisting.Enlist(Jhat.derivative()) -# hs = enlisting.Enlist(delta_m) - -# projnorm = sum(hi._ad_dot(di) for hi, di in zip(hs, derivative)) - -# # this deepcopy is important otherwise pertubations accumulate -# vp_original = vp_guess.copy(deepcopy=True) - -# if comm.comm.rank == 0: -# print( -# "######## Computing the gradient by finite diferences ########", -# flush=True, -# ) -# errors = [] -# for step in steps: # range(3): -# # steps.append(step) -# # perturb the model and calculate the functional (again) -# # J(m + delta_m*h) -# vp_guess = vp_original + step * delta_m -# p_guess_recv, Jp = forward( -# model, -# mesh, -# comm, -# vp_guess, -# sources, -# wavelet, -# point_cloud, -# fwi=True, -# true_rec=p_exact_recv, -# ) - -# fd_grad = (Jp - Jm) / step -# if comm.comm.rank == 0: -# print( -# "\n Cost functional for step " -# + str(step) -# + " : " -# + str(Jp) -# + ", fd approx.: " -# + str(fd_grad) -# + ", grad'*dir : " -# + str(projnorm) -# + " \n ", -# flush=True, -# ) - -# errors.append(100 * ((fd_grad - projnorm) / projnorm)) - -# fire_adj.get_working_tape().clear_tape() - -# # all errors less than 1 % -# errors = np.array(errors) -# assert (np.abs(errors) < 5.0).all() diff --git a/test/test_gradient_ad.py b/test/test_gradient_ad.py new file mode 100644 index 00000000..3c008bf1 --- /dev/null +++ b/test/test_gradient_ad.py @@ -0,0 +1,131 @@ +import firedrake as fire +import firedrake.adjoint as fire_ad +import spyro + + +# --- Basid setup to run a forward simulation with AD --- # +model = {} + +model["opts"] = { + "method": "KMV", # either CG or mass_lumped_triangle + "quadrature": "KMV", # Equi or mass_lumped_triangle + "degree": 1, # p order + "dimension": 2, # dimension + "regularization": False, # regularization is on? + "gamma": 1e-5, # regularization parameter +} + +model["parallelism"] = { + # options: + # `shots_parallelism`. Shots parallelism. + # None - no shots parallelism. + "type": None, + "num_spacial_cores": 1, # Number of cores to use in the spatial parallelism. +} + +# Define the domain size without the ABL. +model["mesh"] = { + "Lz": 1.0, # depth in km - always positive + "Lx": 1.0, # width in km - always positive + "Ly": 0.0, # thickness in km - always positive + "meshfile": "not_used.msh", + "initmodel": "not_used.hdf5", + "truemodel": "not_used.hdf5", +} + +# Specify a 250-m Absorbing Boundary Layer (ABL) on the three sides of the domain to damp outgoing waves. +model["BCs"] = { + "status": False, # True or False, used to turn on any type of BC + "outer_bc": "non-reflective", # none or non-reflective (outer boundary condition) + "lz": 0.0, # thickness of the ABL in the z-direction (km) - always positive + "lx": 0.0, # thickness of the ABL in the x-direction (km) - always positive + "ly": 0.0, # thickness of the ABL in the y-direction (km) - always positive +} + +model["acquisition"] = { + "source_type": "Ricker", + "source_pos": spyro.create_transect((0.2, 0.15), (0.8, 0.15), 1), + "frequency": 7.0, + "delay": 1.0, + "receiver_locations": spyro.create_transect((0.2, 0.2), (0.8, 0.2), 10), +} +model["aut_dif"] = { + "status": True, +} + +model["timeaxis"] = { + "t0": 0.0, # Initial time for event + "tf": 0.4, # Final time for event (for test 7) + "dt": 0.004, # timestep size (divided by 2 in the test 4. dt for test 3 is 0.00050) + "amplitude": 1, # the Ricker has an amplitude of 1. + "nspool": 20, # (20 for dt=0.00050) how frequently to output solution to pvds + "fspool": 1, # how frequently to save solution to RAM +} + + +def make_vp_circle(V, mesh, vp_guess=False, plot_vp=False): + """Acoustic velocity model""" + x, z = fire.SpatialCoordinate(mesh) + if vp_guess: + vp = fire.Function(V).interpolate(1.5 + 0.0 * x) + else: + vp = fire.Function(V).interpolate( + 2.5 + + 1 * fire.tanh(100 * (0.125 - fire.sqrt((x - 0.5) ** 2 + (z - 0.5) ** 2))) + ) + if plot_vp: + outfile = fire.VTKFile("acoustic_cp.pvd") + outfile.write(vp) + return vp + + +def forward( + c, fwd_solver, wavelet, ensemble, + compute_functional=False, true_data_receivers=None, annotate=False +): + if annotate: + fire_ad.continue_annotation() + # source_number based on the ensemble.ensemble_comm.rank + source_number = ensemble.ensemble_comm.rank + receiver_data, J = fwd_solver.execute_acoustic( + c, source_number, wavelet, + compute_functional=compute_functional, + true_data_receivers=true_data_receivers + ) + return receiver_data, J + + +def test_taylor(): + # Test only serial for now. + M = model["parallelism"]["num_spacial_cores"] + my_ensemble = fire.Ensemble(fire.COMM_WORLD, M) + mesh = fire.UnitSquareMesh(20, 20, comm=my_ensemble.comm) + element = fire.FiniteElement( + model["opts"]["method"], mesh.ufl_cell(), degree=model["opts"]["degree"], + variant=model["opts"]["quadrature"] + ) + V = fire.FunctionSpace(mesh, element) + + fwd_solver = spyro.solvers.forward_ad.ForwardSolver(model, mesh, V) + # Ricker wavelet + wavelet = spyro.full_ricker_wavelet( + model["timeaxis"]["dt"], model["timeaxis"]["tf"], + model["acquisition"]["frequency"], + ) + c_true = make_vp_circle(V, mesh) + true_rec, _ = forward(c_true, fwd_solver, wavelet, my_ensemble) + + # --- Gradient with AD --- # + c_guess = make_vp_circle(V, mesh, vp_guess=True) + _, J = forward( + c_guess, fwd_solver, wavelet, my_ensemble, + compute_functional=True, + true_data_receivers=true_rec, annotate=True + ) + + # :class:`~.EnsembleReducedFunctional` is employed to recompute in + # parallel the functional and its gradient associated. + J_hat = fire_ad.EnsembleReducedFunctional( + J, fire_ad.Control(c_guess), my_ensemble) + + assert fire_ad.taylor_test(J_hat, c_guess, fire.Function(V).assign(1.0)) > 1.9 diff --git a/test_ad/__init__.py b/test_ad/__init__.py deleted file mode 100644 index e69de29b..00000000 diff --git a/test_ad/test_gradient_3d_AD.py b/test_ad/test_gradient_3d_AD.py deleted file mode 100644 index 05a6f58d..00000000 --- a/test_ad/test_gradient_3d_AD.py +++ /dev/null @@ -1,86 +0,0 @@ -from firedrake import * -from scipy.optimize import * -import spyro -import pytest - - -# from ..domains import quadrature, space -@pytest.mark.skip(reason="no way of currently testing this with cicleCI resources") -def test_gradient_3d_AD(): - model = {} - - model["opts"] = { - "method": "KMV", # either CG or KMV - "quadrature": "KMV", # Equi or KMV - "degree": 2, # p order - "dimension": 3, # dimension - "regularization": False, # regularization is on? - "gamma": 1e-5, # regularization parameter - } - - model["parallelism"] = { - "type": "automatic", # options: automatic (same number of cores for evey processor) or spatial - } - - # Define the domain size without the ABL. - model["mesh"] = { - "Lz": 1.0, # depth in km - always positive - "Lx": 1.0, # width in km - always positive - "Ly": 1.0, # thickness in km - always positive - "meshfile": "test/meshes/Uniform3D.msh", - "initmodel": "not_used.hdf5", - "truemodel": "not_used.hdf5", - } - - # Specify a 250-m Absorbing Boundary Layer (ABL) on the three sides of the domain to damp outgoing waves. - model["BCs"] = { - "status": False, # True or False, used to turn on any type of BC - "outer_bc": "non-reflective", # none or non-reflective (outer boundary condition) - "abl_bc": "none", # none, gaussian-taper, or alid - "lz": 0.25, # thickness of the ABL in the z-direction (km) - always positive - "lx": 0.25, # thickness of the ABL in the x-direction (km) - always positive - "ly": 0.25, # thickness of the ABL in the y-direction (km) - always positive - } - # receivers = spyro.insert_fixed_value(spyro.create_2d_grid(0.1,0.9,0.1,0.9,2),-0.9, 0) - - model["acquisition"] = { - "source_type": "Ricker", - "source_pos": [(-0.5, 0.5, 0.5)], - "frequency": 10.0, - "delay": 1.0, - "num_rec_x_columns": 5, - "num_rec_y_columns": 1, - "num_rec_z_columns": 1, - # first and final points of the receivers columns (z, x, y) - "receiver_locations": [(-0.1, 0.1, 0.5), (-0.1, 0.9, 0.5)], - } - model["aut_dif"] = { - "status": True, - } - - model["timeaxis"] = { - "t0": 0.0, # Initial time for event - "tf": 0.8, # Final time for event (for test 7) - "dt": 0.001, # timestep size (divided by 2 in the test 4. dt for test 3 is 0.00050) - "amplitude": 1, # the Ricker has an amplitude of 1. - "nspool": 20, # (20 for dt=0.00050) how frequently to output solution to pvds - "fspool": 1, # how frequently to save solution to RAM - } - - comm = spyro.utils.mpi_init(model) - mesh, V = spyro.io.read_mesh(model, comm) - - element = spyro.domains.space.FE_method( - mesh, model["opts"]["method"], model["opts"]["degree"] - ) - - V = FunctionSpace(mesh, element) - z, x, y = SpatialCoordinate(mesh) - - vp_exact = Function(V).interpolate(1.0 + 0.0 * x) - vp_guess = Function(V).interpolate(0.8 + 0.0 * x) - - spyro.tools.gradient_test_acoustic_ad(model, mesh, V, comm, vp_exact, vp_guess) - - -# test_gradient_3d_AD() diff --git a/test_ad/test_gradient_AD.py b/test_ad/test_gradient_AD.py deleted file mode 100644 index 34f5839a..00000000 --- a/test_ad/test_gradient_AD.py +++ /dev/null @@ -1,77 +0,0 @@ -from firedrake import * -from scipy.optimize import * -import spyro - - -# from ..domains import quadrature, space -# @pytest.mark.skip(reason="no way of currently testing this") -def test_gradient_AD(): - model = {} - - model["opts"] = { - "method": "KMV", # either CG or KMV - "quadrature": "KMV", # Equi or KMV - "degree": 1, # p order - "dimension": 2, # dimension - "regularization": False, # regularization is on? - "gamma": 1e-5, # regularization parameter - } - - model["parallelism"] = { - "type": "automatic", # options: automatic (same number of cores for evey processor) or spatial - } - - # Define the domain size without the ABL. - model["mesh"] = { - "Lz": 1.5, # depth in km - always positive - "Lx": 1.5, # width in km - always positive - "Ly": 0.0, # thickness in km - always positive - "meshfile": "not_used.msh", - "initmodel": "not_used.hdf5", - "truemodel": "not_used.hdf5", - } - - # Specify a 250-m Absorbing Boundary Layer (ABL) on the three sides of the domain to damp outgoing waves. - model["BCs"] = { - "status": False, # True or False, used to turn on any type of BC - "outer_bc": "non-reflective", # none or non-reflective (outer boundary condition) - "abl_bc": "none", # none, gaussian-taper, or alid - "lz": 0.25, # thickness of the ABL in the z-direction (km) - always positive - "lx": 0.25, # thickness of the ABL in the x-direction (km) - always positive - "ly": 0.0, # thickness of the ABL in the y-direction (km) - always positive - } - - model["acquisition"] = { - "source_type": "Ricker", - "source_pos": [(0.75, 0.75)], - "frequency": 10.0, - "delay": 1.0, - "receiver_locations": spyro.create_transect((0.9, 0.2), (0.9, 0.8), 10), - } - model["aut_dif"] = { - "status": True, - } - - model["timeaxis"] = { - "t0": 0.0, # Initial time for event - "tf": 0.8, # Final time for event (for test 7) - "dt": 0.001, # timestep size (divided by 2 in the test 4. dt for test 3 is 0.00050) - "amplitude": 1, # the Ricker has an amplitude of 1. - "nspool": 20, # (20 for dt=0.00050) how frequently to output solution to pvds - "fspool": 1, # how frequently to save solution to RAM - } - - comm = spyro.utils.mpi_init(model) - mesh = RectangleMesh(100, 100, 1.5, 1.5) # to test FWI, mesh aligned with interface - - element = spyro.domains.space.FE_method( - mesh, model["opts"]["method"], model["opts"]["degree"] - ) - - V = FunctionSpace(mesh, element) - z, x = SpatialCoordinate(mesh) - - vp_exact = Function(V).interpolate(1.0 + 0.0 * x) - vp_guess = Function(V).interpolate(0.8 + 0.0 * x) - - spyro.tools.gradient_test_acoustic_ad(model, mesh, V, comm, vp_exact, vp_guess) From 6d43c59b15d1fb7849c833a6a361fa6e68f080a0 Mon Sep 17 00:00:00 2001 From: Iglesia Dolci Date: Wed, 4 Sep 2024 07:17:47 +0100 Subject: [PATCH 04/13] minor changes --- .../run_forward_ad.py | 16 ++++++++-------- test/test_gradient_ad.py | 18 +++++++++--------- 2 files changed, 17 insertions(+), 17 deletions(-) diff --git a/demos/with_automatic_differentiation/run_forward_ad.py b/demos/with_automatic_differentiation/run_forward_ad.py index 1c595070..ab90bc7c 100644 --- a/demos/with_automatic_differentiation/run_forward_ad.py +++ b/demos/with_automatic_differentiation/run_forward_ad.py @@ -66,20 +66,20 @@ } -def make_vp_circle(vp_guess=False, plot_vp=False): +def make_c_camembert(c_guess=False, plot_c=False): """Acoustic velocity model""" x, z = fire.SpatialCoordinate(mesh) - if vp_guess: - vp = fire.Function(V).interpolate(1.5 + 0.0 * x) + if c_guess: + c = fire.Function(V).interpolate(1.5 + 0.0 * x) else: - vp = fire.Function(V).interpolate( + c = fire.Function(V).interpolate( 2.5 + 1 * fire.tanh(100 * (0.125 - fire.sqrt((x - 0.5) ** 2 + (z - 0.5) ** 2))) ) - if plot_vp: + if plot_c: outfile = fire.VTKFile("acoustic_cp.pvd") - outfile.write(vp) - return vp + outfile.write(c) + return c # Use emsemble parallelism. @@ -95,7 +95,7 @@ def make_vp_circle(vp_guess=False, plot_vp=False): forward_solver = spyro.solvers.forward_ad.ForwardSolver(model, mesh, V) -c_true = make_vp_circle() +c_true = make_c_camembert() # Ricker wavelet wavelet = spyro.full_ricker_wavelet( model["timeaxis"]["dt"], model["timeaxis"]["tf"], diff --git a/test/test_gradient_ad.py b/test/test_gradient_ad.py index 3c008bf1..9c9e3d09 100644 --- a/test/test_gradient_ad.py +++ b/test/test_gradient_ad.py @@ -63,20 +63,20 @@ } -def make_vp_circle(V, mesh, vp_guess=False, plot_vp=False): +def make_c_camembert(V, mesh, c_guess=False, plot_c=False): """Acoustic velocity model""" x, z = fire.SpatialCoordinate(mesh) - if vp_guess: - vp = fire.Function(V).interpolate(1.5 + 0.0 * x) + if c_guess: + c = fire.Function(V).interpolate(1.5 + 0.0 * x) else: - vp = fire.Function(V).interpolate( + c = fire.Function(V).interpolate( 2.5 + 1 * fire.tanh(100 * (0.125 - fire.sqrt((x - 0.5) ** 2 + (z - 0.5) ** 2))) ) - if plot_vp: + if plot_c: outfile = fire.VTKFile("acoustic_cp.pvd") - outfile.write(vp) - return vp + outfile.write(c) + return c def forward( @@ -112,11 +112,11 @@ def test_taylor(): model["timeaxis"]["dt"], model["timeaxis"]["tf"], model["acquisition"]["frequency"], ) - c_true = make_vp_circle(V, mesh) + c_true = make_c_camembert(V, mesh) true_rec, _ = forward(c_true, fwd_solver, wavelet, my_ensemble) # --- Gradient with AD --- # - c_guess = make_vp_circle(V, mesh, vp_guess=True) + c_guess = make_c_camembert(V, mesh, c_guess=True) _, J = forward( c_guess, fwd_solver, wavelet, my_ensemble, compute_functional=True, From 4bc2af92ba9342dfee223f523ab213343e6d410f Mon Sep 17 00:00:00 2001 From: Daiane Iglesia Dolci <63597005+Ig-dolci@users.noreply.github.com> Date: Wed, 4 Sep 2024 07:24:49 +0100 Subject: [PATCH 05/13] Update demos/with_automatic_differentiation/run_fwi_ad.py --- demos/with_automatic_differentiation/run_fwi_ad.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/demos/with_automatic_differentiation/run_fwi_ad.py b/demos/with_automatic_differentiation/run_fwi_ad.py index 232dcd19..7fc2d778 100644 --- a/demos/with_automatic_differentiation/run_fwi_ad.py +++ b/demos/with_automatic_differentiation/run_fwi_ad.py @@ -91,7 +91,7 @@ def forward( receiver_data = [] J = 0.0 for sn in range(len(model["acquisition"]["source_pos"])): - rec_data, J_val = forward_solver.execute(c, sn, wavelet) + rec_data, J_val = forward_solver.execute_acoustic(c, sn, wavelet) receiver_data.append(rec_data) J += J_val sol = forward_solver.solution From f0d32d412d18adeeda9f239e8d505c81380bab4f Mon Sep 17 00:00:00 2001 From: Iglesia Dolci Date: Wed, 4 Sep 2024 07:33:48 +0100 Subject: [PATCH 06/13] Random perturbation --- test/test_gradient_ad.py | 6 ++++-- 1 file changed, 4 insertions(+), 2 deletions(-) diff --git a/test/test_gradient_ad.py b/test/test_gradient_ad.py index 9c9e3d09..59a9882c 100644 --- a/test/test_gradient_ad.py +++ b/test/test_gradient_ad.py @@ -1,6 +1,7 @@ import firedrake as fire import firedrake.adjoint as fire_ad import spyro +from numpy.random import rand # --- Basid setup to run a forward simulation with AD --- # @@ -127,5 +128,6 @@ def test_taylor(): # parallel the functional and its gradient associated. J_hat = fire_ad.EnsembleReducedFunctional( J, fire_ad.Control(c_guess), my_ensemble) - - assert fire_ad.taylor_test(J_hat, c_guess, fire.Function(V).assign(1.0)) > 1.9 + h = fire.Function(V) + h.dat.data[:] = rand(V.dim()) + assert fire_ad.taylor_test(J_hat, c_guess, h) > 1.9 From 58036a0c2ceaa605a083bf1c156b8f06a5ec997d Mon Sep 17 00:00:00 2001 From: Iglesia Dolci Date: Wed, 4 Sep 2024 07:35:11 +0100 Subject: [PATCH 07/13] checkpoint option --- demos/with_automatic_differentiation/run_fwi_ad.py | 1 + 1 file changed, 1 insertion(+) diff --git a/demos/with_automatic_differentiation/run_fwi_ad.py b/demos/with_automatic_differentiation/run_fwi_ad.py index 7fc2d778..6d7b08e3 100644 --- a/demos/with_automatic_differentiation/run_fwi_ad.py +++ b/demos/with_automatic_differentiation/run_fwi_ad.py @@ -53,6 +53,7 @@ } model["aut_dif"] = { "status": True, + "checkpointing": False, } model["timeaxis"] = { From cd35e51b4cf1a92400e90ce454ece74b95e14205 Mon Sep 17 00:00:00 2001 From: Iglesia Dolci Date: Wed, 4 Sep 2024 08:03:42 +0100 Subject: [PATCH 08/13] Add checkpointing --- .../with_automatic_differentiation/run_fwi_ad.py | 16 ++++++++++++---- spyro/solvers/forward_AD.py | 11 ++++++++++- 2 files changed, 22 insertions(+), 5 deletions(-) diff --git a/demos/with_automatic_differentiation/run_fwi_ad.py b/demos/with_automatic_differentiation/run_fwi_ad.py index 6d7b08e3..b247b246 100644 --- a/demos/with_automatic_differentiation/run_fwi_ad.py +++ b/demos/with_automatic_differentiation/run_fwi_ad.py @@ -1,5 +1,6 @@ import firedrake as fire import firedrake.adjoint as fire_ad +from checkpoint_schedules import Revolve import spyro @@ -53,12 +54,12 @@ } model["aut_dif"] = { "status": True, - "checkpointing": False, + "checkpointing": True, } model["timeaxis"] = { "t0": 0.0, # Initial time for event - "tf": 0.6, # Final time for event (for test 7) + "tf": 0.8, # Final time for event (for test 7) "dt": 0.001, # timestep size (divided by 2 in the test 4. dt for test 3 is 0.00050) "amplitude": 1, # the Ricker has an amplitude of 1. "nspool": 20, # (20 for dt=0.00050) how frequently to output solution to pvds @@ -87,6 +88,12 @@ def forward( ): if annotate: fire_ad.continue_annotation() + if model["aut_dif"]["checkpointing"]: + total_steps = int(model["timeaxis"]["tf"] / model["timeaxis"]["dt"]) + steps_store = int(total_steps / 10) # Store 10% of the steps. + tape = fire_ad.get_working_tape() + tape.progress_bar = fire.ProgressBar + tape.enable_checkpointing(Revolve(total_steps, steps_store)) if model["parallelism"]["type"] is None: outfile = fire.VTKFile("solution.pvd") receiver_data = [] @@ -147,8 +154,9 @@ def forward( # (3 in this case). J_hat = fire_ad.EnsembleReducedFunctional( J, fire_ad.Control(c_guess), my_ensemble) -fire_ad.taylor_test(J_hat, c_guess, fire.Function(V).assign(1.0)) c_optimised = fire_ad.minimize(J_hat, method="L-BFGS-B", options={"disp": True, "maxiter": 10}, bounds=(1.5, 3.5), - derivative_options={"riesz_representation": 'l2'}) \ No newline at end of file + derivative_options={"riesz_representation": 'l2'}) + +fire.VTKFile("c_optimised.pvd").write(c_optimised) \ No newline at end of file diff --git a/spyro/solvers/forward_AD.py b/spyro/solvers/forward_AD.py index 9d3f5f03..db21bedb 100644 --- a/spyro/solvers/forward_AD.py +++ b/spyro/solvers/forward_AD.py @@ -91,7 +91,16 @@ def execute_acoustic( J_val = 0.0 receiver_data = [] total_steps = int(self.model["timeaxis"]["tf"] / self.model["timeaxis"]["dt"]) - for step in range(total_steps): + if ( + fire_ad.get_working_tape()._checkpoint_manager + and self.model["aut_dif"]["checkpointing"] + ): + time_range = fire_ad.get_working_tape().timestepper( + iter(range(total_steps))) + else: + time_range = range(total_steps) + + for step in time_range: source_function.assign(wavelet[step] * q_s) solver.solve() u_nm1.assign(u_n) From be3bdad5437ab180fef2a679e367de5f554746d1 Mon Sep 17 00:00:00 2001 From: Iglesia Dolci Date: Wed, 4 Sep 2024 08:14:17 +0100 Subject: [PATCH 09/13] test checkpointing --- .../with_automatic_differentiation/run_fwi_ad.py | 2 +- test/test_gradient_ad.py | 15 +++++++++++++++ 2 files changed, 16 insertions(+), 1 deletion(-) diff --git a/demos/with_automatic_differentiation/run_fwi_ad.py b/demos/with_automatic_differentiation/run_fwi_ad.py index b247b246..e1540f9a 100644 --- a/demos/with_automatic_differentiation/run_fwi_ad.py +++ b/demos/with_automatic_differentiation/run_fwi_ad.py @@ -159,4 +159,4 @@ def forward( bounds=(1.5, 3.5), derivative_options={"riesz_representation": 'l2'}) -fire.VTKFile("c_optimised.pvd").write(c_optimised) \ No newline at end of file +fire.VTKFile("c_optimised.pvd").write(c_optimised) diff --git a/test/test_gradient_ad.py b/test/test_gradient_ad.py index 59a9882c..55b8b286 100644 --- a/test/test_gradient_ad.py +++ b/test/test_gradient_ad.py @@ -2,6 +2,7 @@ import firedrake.adjoint as fire_ad import spyro from numpy.random import rand +from checkpoint_schedules import Revolve # --- Basid setup to run a forward simulation with AD --- # @@ -52,6 +53,7 @@ } model["aut_dif"] = { "status": True, + "checkpointing": False, } model["timeaxis"] = { @@ -86,6 +88,12 @@ def forward( ): if annotate: fire_ad.continue_annotation() + if model["aut_dif"]["checkpointing"]: + total_steps = int(model["timeaxis"]["tf"] / model["timeaxis"]["dt"]) + steps_store = int(total_steps / 10) # Store 10% of the steps. + tape = fire_ad.get_working_tape() + tape.progress_bar = fire.ProgressBar + tape.enable_checkpointing(Revolve(total_steps, steps_store)) # source_number based on the ensemble.ensemble_comm.rank source_number = ensemble.ensemble_comm.rank receiver_data, J = fwd_solver.execute_acoustic( @@ -131,3 +139,10 @@ def test_taylor(): h = fire.Function(V) h.dat.data[:] = rand(V.dim()) assert fire_ad.taylor_test(J_hat, c_guess, h) > 1.9 + fire_ad.get_working_tape().clear_tape() + fire_ad.pause_annotation() + + +def test_taylor_checkpointing(): + model["aut_dif"]["checkpointing"] = True + test_taylor() From c4ecc41503ddbbe5baeff0dae46903e70976675c Mon Sep 17 00:00:00 2001 From: Iglesia Dolci Date: Wed, 4 Sep 2024 08:44:34 +0100 Subject: [PATCH 10/13] organize --- .../__init__.py | 0 .../run_forward_ad.py | 85 +------------- .../run_fwi_ad.py | 106 +++++------------- demos/with_automatic_differentiation/utils.py | 104 +++++++++++++++++ 4 files changed, 136 insertions(+), 159 deletions(-) create mode 100644 demos/with_automatic_differentiation/__init__.py create mode 100644 demos/with_automatic_differentiation/utils.py diff --git a/demos/with_automatic_differentiation/__init__.py b/demos/with_automatic_differentiation/__init__.py new file mode 100644 index 00000000..e69de29b diff --git a/demos/with_automatic_differentiation/run_forward_ad.py b/demos/with_automatic_differentiation/run_forward_ad.py index ab90bc7c..47a789eb 100644 --- a/demos/with_automatic_differentiation/run_forward_ad.py +++ b/demos/with_automatic_differentiation/run_forward_ad.py @@ -1,86 +1,13 @@ import firedrake as fire import spyro -import matplotlib.pyplot as plt -import numpy as np - -import spyro.solvers +from demos.with_automatic_differentiation.utils import \ + model_settings, make_c_camembert +import os +os.environ["OMP_NUM_THREADS"] = "1" # --- Basid setup to run a forward simulation with AD --- # -model = {} - -model["opts"] = { - "method": "KMV", # either CG or mass_lumped_triangle - "quadrature": "KMV", # Equi or mass_lumped_triangle - "degree": 1, # p order - "dimension": 2, # dimension - "regularization": False, # regularization is on? - "gamma": 1e-5, # regularization parameter -} - -model["parallelism"] = { - # options: - # `shots_parallelism`. Shots parallelism. - # None - no shots parallelism. - "type": "shots_parallelism", - "num_spacial_cores": 1, # Number of cores to use in the spatial parallelism. -} - -# Define the domain size without the ABL. -model["mesh"] = { - "Lz": 1.0, # depth in km - always positive - "Lx": 1.0, # width in km - always positive - "Ly": 0.0, # thickness in km - always positive - "meshfile": "not_used.msh", - "initmodel": "not_used.hdf5", - "truemodel": "not_used.hdf5", -} - -# Specify a 250-m Absorbing Boundary Layer (ABL) on the three sides of the domain to damp outgoing waves. -model["BCs"] = { - "status": False, # True or False, used to turn on any type of BC - "outer_bc": "non-reflective", # none or non-reflective (outer boundary condition) - "abl_bc": "none", # none, gaussian-taper, or alid - "lz": 0.0, # thickness of the ABL in the z-direction (km) - always positive - "lx": 0.0, # thickness of the ABL in the x-direction (km) - always positive - "ly": 0.0, # thickness of the ABL in the y-direction (km) - always positive -} - -model["acquisition"] = { - "source_type": "Ricker", - "source_pos": spyro.create_transect((0.2, 0.15), (0.8, 0.15), 3), - "frequency": 7.0, - "delay": 1.0, - "receiver_locations": spyro.create_transect((0.2, 0.2), (0.8, 0.2), 10), -} -model["aut_dif"] = { - "status": True, -} - -model["timeaxis"] = { - "t0": 0.0, # Initial time for event - "tf": 0.2, # Final time for event (for test 7) - "dt": 0.001, # timestep size (divided by 2 in the test 4. dt for test 3 is 0.00050) - "amplitude": 1, # the Ricker has an amplitude of 1. - "nspool": 20, # (20 for dt=0.00050) how frequently to output solution to pvds - "fspool": 1, # how frequently to save solution to RAM -} - - -def make_c_camembert(c_guess=False, plot_c=False): - """Acoustic velocity model""" - x, z = fire.SpatialCoordinate(mesh) - if c_guess: - c = fire.Function(V).interpolate(1.5 + 0.0 * x) - else: - c = fire.Function(V).interpolate( - 2.5 - + 1 * fire.tanh(100 * (0.125 - fire.sqrt((x - 0.5) ** 2 + (z - 0.5) ** 2))) - ) - if plot_c: - outfile = fire.VTKFile("acoustic_cp.pvd") - outfile.write(c) - return c +model = model_settings() # Use emsemble parallelism. M = model["parallelism"]["num_spacial_cores"] @@ -95,7 +22,7 @@ def make_c_camembert(c_guess=False, plot_c=False): forward_solver = spyro.solvers.forward_ad.ForwardSolver(model, mesh, V) -c_true = make_c_camembert() +c_true = make_c_camembert(mesh, V) # Ricker wavelet wavelet = spyro.full_ricker_wavelet( model["timeaxis"]["dt"], model["timeaxis"]["tf"], diff --git a/demos/with_automatic_differentiation/run_fwi_ad.py b/demos/with_automatic_differentiation/run_fwi_ad.py index e1540f9a..824ed3d8 100644 --- a/demos/with_automatic_differentiation/run_fwi_ad.py +++ b/demos/with_automatic_differentiation/run_fwi_ad.py @@ -2,90 +2,35 @@ import firedrake.adjoint as fire_ad from checkpoint_schedules import Revolve import spyro +from demos.with_automatic_differentiation import utils - -# --- Basid setup to run a forward simulation with AD --- # -model = {} - -model["opts"] = { - "method": "KMV", # either CG or mass_lumped_triangle - "quadrature": "KMV", # Equi or mass_lumped_triangle - "degree": 1, # p order - "dimension": 2, # dimension - "regularization": False, # regularization is on? - "gamma": 1e-5, # regularization parameter -} - -model["parallelism"] = { - # options: - # `shots_parallelism`. Shots parallelism. - # None - no shots parallelism. - "type": "shots_parallelism", - "num_spacial_cores": 1, # Number of cores to use in the spatial - # parallelism. -} - -# Define the domain size without the ABL. -model["mesh"] = { - "Lz": 1.0, # depth in km - always positive - "Lx": 1.0, # width in km - always positive - "Ly": 0.0, # thickness in km - always positive - "meshfile": "not_used.msh", - "initmodel": "not_used.hdf5", - "truemodel": "not_used.hdf5", -} - -# Specify a 250-m Absorbing Boundary Layer (ABL) on the three sides of the domain to damp outgoing waves. -model["BCs"] = { - "status": False, # True or False, used to turn on any type of BC - "outer_bc": "non-reflective", # none or non-reflective (outer boundary condition) - "abl_bc": "none", # none, gaussian-taper, or alid - "lz": 0.0, # thickness of the ABL in the z-direction (km) - always positive - "lx": 0.0, # thickness of the ABL in the x-direction (km) - always positive - "ly": 0.0, # thickness of the ABL in the y-direction (km) - always positive -} - -model["acquisition"] = { - "source_type": "Ricker", - "source_pos": spyro.create_transect((0.2, 0.15), (0.8, 0.15), 3), - "frequency": 7.0, - "delay": 1.0, - "receiver_locations": spyro.create_transect((0.2, 0.2), (0.8, 0.2), 10), -} -model["aut_dif"] = { - "status": True, - "checkpointing": True, -} - -model["timeaxis"] = { - "t0": 0.0, # Initial time for event - "tf": 0.8, # Final time for event (for test 7) - "dt": 0.001, # timestep size (divided by 2 in the test 4. dt for test 3 is 0.00050) - "amplitude": 1, # the Ricker has an amplitude of 1. - "nspool": 20, # (20 for dt=0.00050) how frequently to output solution to pvds - "fspool": 1, # how frequently to save solution to RAM -} - - -def make_c_camembert(c_guess=False, plot_c=False): - """Acoustic velocity model""" - x, z = fire.SpatialCoordinate(mesh) - if c_guess: - c = fire.Function(V).interpolate(1.5 + 0.0 * x) - else: - c = fire.Function(V).interpolate( - 2.5 - + 1 * fire.tanh(100 * (0.125 - fire.sqrt((x - 0.5) ** 2 + (z - 0.5) ** 2))) - ) - if plot_c: - outfile = fire.VTKFile("acoustic_cp.pvd") - outfile.write(c) - return c +model = utils.model_settings() def forward( c, compute_functional=False, true_data_receivers=None, annotate=False ): + """Time-stepping acoustic forward solver. + + The time integration is done using a central difference scheme. + + Parameters + ---------- + c : firedrake.Function + Velocity field. + compute_functional : bool, optional + Whether to compute the functional. If True, the true receiver + data must be provided. + true_data_receivers : list, optional + True receiver data. This is used to compute the functional. + annotate : bool, optional + If True, the forward model is annotated for automatic differentiation. + + Returns + ------- + (receiver_data : list, J_val : float) + Receiver data and functional value. + """ if annotate: fire_ad.continue_annotation() if model["aut_dif"]["checkpointing"]: @@ -94,6 +39,7 @@ def forward( tape = fire_ad.get_working_tape() tape.progress_bar = fire.ProgressBar tape.enable_checkpointing(Revolve(total_steps, steps_store)) + if model["parallelism"]["type"] is None: outfile = fire.VTKFile("solution.pvd") receiver_data = [] @@ -133,7 +79,7 @@ def forward( forward_solver = spyro.solvers.forward_ad.ForwardSolver(model, mesh, V) -c_true = make_c_camembert() +c_true = utils.make_c_camembert(mesh, V) # Ricker wavelet wavelet = spyro.full_ricker_wavelet( model["timeaxis"]["dt"], model["timeaxis"]["tf"], @@ -143,7 +89,7 @@ def forward( true_rec, _ = forward(c_true) # --- FWI with AD --- # -c_guess = make_c_camembert(c_guess=True) +c_guess = utils.make_c_camembert(mesh, V, c_guess=True) guess_rec, J = forward( c_guess, compute_functional=True, true_data_receivers=true_rec, annotate=True diff --git a/demos/with_automatic_differentiation/utils.py b/demos/with_automatic_differentiation/utils.py new file mode 100644 index 00000000..6d89fcf7 --- /dev/null +++ b/demos/with_automatic_differentiation/utils.py @@ -0,0 +1,104 @@ +# --- Basid setup to run a forward simulation with AD --- # +import firedrake as fire +import spyro + +def model_settings(): + """Model settings for forward and Full Waveform Inversion (FWI) + simulations. + + Returns + ------- + model : dict + Dictionary containing the model settings. + """ + + model = {} + + model["opts"] = { + "method": "KMV", # either CG or mass_lumped_triangle + "quadrature": "KMV", # Equi or mass_lumped_triangle + "degree": 1, # p order + "dimension": 2, # dimension + "regularization": False, # regularization is on? + "gamma": 1e-5, # regularization parameter + } + + model["parallelism"] = { + # options: + # `shots_parallelism`. Shots parallelism. + # None - no shots parallelism. + "type": "shots_parallelism", + "num_spacial_cores": 1, # Number of cores to use in the spatial + # parallelism. + } + + # Define the domain size without the ABL. + model["mesh"] = { + "Lz": 1.0, # depth in km - always positive + "Lx": 1.0, # width in km - always positive + "Ly": 0.0, # thickness in km - always positive + "meshfile": "not_used.msh", + "initmodel": "not_used.hdf5", + "truemodel": "not_used.hdf5", + } + + # Specify a 250-m Absorbing Boundary Layer (ABL) on the three sides of the domain to damp outgoing waves. + model["BCs"] = { + "status": False, # True or False, used to turn on any type of BC + "outer_bc": "non-reflective", # none or non-reflective (outer boundary condition) + "abl_bc": "none", # none, gaussian-taper, or alid + "lz": 0.0, # thickness of the ABL in the z-direction (km) - always positive + "lx": 0.0, # thickness of the ABL in the x-direction (km) - always positive + "ly": 0.0, # thickness of the ABL in the y-direction (km) - always positive + } + + model["acquisition"] = { + "source_type": "Ricker", + "source_pos": spyro.create_transect((0.2, 0.15), (0.8, 0.15), 3), + "frequency": 7.0, + "delay": 1.0, + "receiver_locations": spyro.create_transect((0.2, 0.2), (0.8, 0.2), 10), + } + model["aut_dif"] = { + "status": True, + "checkpointing": True, + } + + model["timeaxis"] = { + "t0": 0.0, # Initial time for event + "tf": 0.8, # Final time for event (for test 7) + "dt": 0.001, # timestep size (divided by 2 in the test 4. dt for test 3 is 0.00050) + "amplitude": 1, # the Ricker has an amplitude of 1. + "nspool": 20, # (20 for dt=0.00050) how frequently to output solution to pvds + "fspool": 1, # how frequently to save solution to RAM + } + + return model + + +def make_c_camembert(mesh, function_space, c_guess=False, plot_c=False): + """Acoustic velocity model. + + Parameters + ---------- + mesh : firedrake.Mesh + Mesh. + function_space : firedrake.FunctionSpace + Function space. + c_guess : bool, optional + If True, the initial guess for the velocity field is returned. + plot_c : bool, optional + If True, the velocity field is saved to a VTK file. + """ + x, z = fire.SpatialCoordinate(mesh) + if c_guess: + c = fire.Function(function_space).interpolate(1.5 + 0.0 * x) + else: + c = fire.Function(function_space).interpolate( + 2.5 + + 1 * fire.tanh(100 * (0.125 - fire.sqrt((x - 0.5) ** 2 + (z - 0.5) ** 2))) + ) + if plot_c: + outfile = fire.VTKFile("acoustic_cp.pvd") + outfile.write(c) + return c From cc5b4d17a5a0a50462a93b619fcf503f025ac624 Mon Sep 17 00:00:00 2001 From: Ig-dolci Date: Wed, 4 Sep 2024 14:45:31 +0100 Subject: [PATCH 11/13] rename forward file --- spyro/solvers/{forward_AD.py => forward_ad.py} | 0 1 file changed, 0 insertions(+), 0 deletions(-) rename spyro/solvers/{forward_AD.py => forward_ad.py} (100%) diff --git a/spyro/solvers/forward_AD.py b/spyro/solvers/forward_ad.py similarity index 100% rename from spyro/solvers/forward_AD.py rename to spyro/solvers/forward_ad.py From b557319e9551db1b588d763d54f3f6b9aa8a965a Mon Sep 17 00:00:00 2001 From: Iglesia Dolci Date: Wed, 4 Sep 2024 15:28:17 +0100 Subject: [PATCH 12/13] OMP_NUM_THREADS 1 --- demos/with_automatic_differentiation/run_fwi_ad.py | 4 ++++ 1 file changed, 4 insertions(+) diff --git a/demos/with_automatic_differentiation/run_fwi_ad.py b/demos/with_automatic_differentiation/run_fwi_ad.py index 824ed3d8..84fb2393 100644 --- a/demos/with_automatic_differentiation/run_fwi_ad.py +++ b/demos/with_automatic_differentiation/run_fwi_ad.py @@ -3,7 +3,10 @@ from checkpoint_schedules import Revolve import spyro from demos.with_automatic_differentiation import utils +import os +os.environ["OMP_NUM_THREADS"] = "1" +# --- Basid setup to run a FWI --- # model = utils.model_settings() @@ -79,6 +82,7 @@ def forward( forward_solver = spyro.solvers.forward_ad.ForwardSolver(model, mesh, V) +# Camembert model. c_true = utils.make_c_camembert(mesh, V) # Ricker wavelet wavelet = spyro.full_ricker_wavelet( From 3ed4881fbed31652e9abd7d3a364649b26c1a723 Mon Sep 17 00:00:00 2001 From: olender Date: Wed, 11 Sep 2024 13:25:37 -0300 Subject: [PATCH 13/13] updating firedrake location --- .github/workflows/python-tests.yml | 10 +++++----- 1 file changed, 5 insertions(+), 5 deletions(-) diff --git a/.github/workflows/python-tests.yml b/.github/workflows/python-tests.yml index cee05c72..1d47b323 100644 --- a/.github/workflows/python-tests.yml +++ b/.github/workflows/python-tests.yml @@ -17,28 +17,28 @@ jobs: - uses: actions/checkout@v3 - name: Running serial tests run: | - source /home/olender/firedrakes/2024_07_19/firedrake/bin/activate + source /home/olender/firedrakes/2024_09_11/firedrake/bin/activate pytest --cov-report=xml --cov=spyro test/ - name: Running parallel 3D forward test run: | - source /home/olender/firedrakes/2024_07_19/firedrake/bin/activate + source /home/olender/firedrakes/2024_09_11/firedrake/bin/activate mpiexec -n 6 pytest test_3d/test_hexahedral_convergence.py mpiexec -n 6 pytest test_parallel/test_forward.py mpiexec -n 6 pytest test_parallel/test_fwi.py - name: Covering parallel 3D forward test continue-on-error: true run: | - source /home/olender/firedrakes/2024_07_19/firedrake/bin/activate + source /home/olender/firedrakes/2024_09_11/firedrake/bin/activate mpiexec -n 6 pytest --cov-report=xml --cov-append --cov=spyro test_3d/test_hexahedral_convergence.py - name: Covering parallel forward test continue-on-error: true run: | - source /home/olender/firedrakes/2024_07_19/firedrake/bin/activate + source /home/olender/firedrakes/2024_09_11/firedrake/bin/activate mpiexec -n 6 pytest --cov-report=xml --cov-append --cov=spyro test_parallel/test_forward.py - name: Covering parallel fwi test continue-on-error: true run: | - source /home/olender/firedrakes/2024_07_19/firedrake/bin/activate + source /home/olender/firedrakes/2024_09_11/firedrake/bin/activate mpiexec -n 6 pytest --cov-report=xml --cov-append --cov=spyro test_parallel/test_fwi.py - name: Uploading coverage to Codecov run: export CODECOV_TOKEN="6cd21147-54f7-4b77-94ad-4b138053401d" && bash <(curl -s https://codecov.io/bash)