Merge pull request #46 from Olender/issue_0043_adding_fwi

Issue 0043 adding fwi

.github/workflows/python-tests.yml

@@ -24,12 +24,22 @@ jobs:
           source /home/olender/Firedrakes/newest3/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/newest3/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/newest3/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/newest3/firedrake/bin/activate
+          mpiexec -n 6 pytest --cov-report=xml --cov-append --cov=spyro test_parallel/test_fwi.py
     # - name: Running serial tests for adjoint
     #   run: |
     #       source /home/olender/Firedrakes/main/firedrake/bin/activate

README.md

@@ -2,17 +2,17 @@
 [![Python tests](https://github.com/NDF-Poli-USP/spyro/actions/workflows/python-tests.yml/badge.svg)](https://github.com/NDF-Poli-USP/spyro/actions/workflows/python-tests.yml)
 [![codecov](https://codecov.io/gh/Olender/spyro-1/graph/badge.svg?token=69M30UMRFD)](https://codecov.io/gh/Olender/spyro-1)
 
-SPIRO: Seismic Parallel Inversion and Reconstruction Optimization framework
+spyro: seismic parallel inversion and reconstruction optimization framework
 ============================================
 
-Acoustic wave modeling in Firedrake
+Wave modeling in Firedrake
 
-SPIRO is a Python library for modeling acoustic waves. The main
+spyro is a Python library for modeling acoustic waves. The main
 functionality is a set of forward and adjoint wave propagators for solving the acoustic wave equation in the time domain.
-These wave propagators can be used to form complete full waveform inversion (FWI) applications. See the [demos](https://github.com/krober10nd/spyro/tree/main/demos).
-To implement these solvers, SPIRO uses the finite element package [Firedrake](https://www.firedrakeproject.org/index.html).
+These wave propagators can be used to form complete full waveform inversion (FWI) applications. See the [notebooks](https://github.com/Olender/spyro-1/tree/main/notebook_tutorials).
+To implement these solvers, spyro uses the finite element package [Firedrake](https://www.firedrakeproject.org/index.html).
 
-To use SPIRO, you'll need to have some knowledge of Python and some basic concepts in inverse modeling relevant to active-sourcce seismology.
+To use spyro, you'll need to have some knowledge of Python and some basic concepts in inverse modeling relevant to active-source seismology.
 
 Discussions about development take place on our Slack channel. Everyone is invited to join using the link: https://join.slack.com/t/spyroworkspace/shared_invite/zt-u87ih28m-2h9JobfkdArs4ku3a1wLLQ
 

spyro/examples/rectangle.py

@@ -73,12 +73,6 @@
     "cmax": 4.5,
     "R": 1e-6,
     "pad_length": 0.25,
-    "status": True,
-    "damping_type": "PML",
-    "exponent": 2,
-    "cmax": 4.5,
-    "R": 1e-6,
-    "pad_length": 0.25,
 }
 
 rectangle_dictionary["acquisition"] = {

spyro/io/basicio.py

@@ -339,7 +339,7 @@ def interpolate(Model, fname, V):
     """
     sd = V.mesh().geometric_dimension()
     m = V.ufl_domain()
-    if Model.abc_status:
+    if Model.abc_active:
         minz = -Model.length_z - Model.abc_pad_length
         maxz = 0.0
         minx = 0.0 - Model.abc_pad_length

spyro/io/model_parameters.py

@@ -192,7 +192,7 @@ class Model_parameters:
         User defined mesh.
     firedrake_mesh: firedrake.Mesh
         Firedrake mesh.
-    abc_status: bool
+    abc_active: bool
         Whether or not the absorbing boundary conditions are used.
     abc_exponent: int
         Exponent of the absorbing boundary conditions.
@@ -375,15 +375,15 @@ def _sanitize_absorving_boundary_condition(self):
                 "status": False
             }
         dictionary = self.input_dictionary["absorving_boundary_conditions"]
-        self.abc_status = dictionary["status"]
+        self.abc_active = dictionary["status"]
 
         BL_obj = io.boundary_layer_io.read_boundary_layer(dictionary)
         self.abc_exponent = BL_obj.abc_exponent
         self.abc_cmax = BL_obj.abc_cmax
         self.abc_R = BL_obj.abc_R
         self.abc_pad_length = BL_obj.abc_pad_length
         self.abc_boundary_layer_type = BL_obj.abc_boundary_layer_type
-        if self.abc_status:
+        if self.abc_active:
             self._correct_time_integrator_for_abc()
 
     def _correct_time_integrator_for_abc(self):

spyro/meshing/meshing_functions.py

@@ -280,8 +280,13 @@ def create_firedrake_2D_mesh(self):
         """
         Creates a 2D mesh based on Firedrake meshing utilities.
         """
-        nx = int(self.length_x / self.dx)
-        nz = int(self.length_z / self.dx)
+        if self.abc_pad:
+            nx = int( (self.length_x + 2*self.abc_pad) / self.dx)
+            nz = int( (self.length_z + self.abc_pad)/ self.dx)
+        else:
+            nx = int(self.length_x / self.dx)
+            nz = int(self.length_z / self.dx)
+
         comm = self.comm
         if self.cell_type == "quadrilateral":
             quadrilateral = True

spyro/plots/__init__.py

@@ -1,8 +1,11 @@
 from .plots import plot_shots, plot_mesh_sizes, plot_model, plot_function
+from .plots import debug_plot, debug_pvd
 
 __all__ = [
     "plot_shots",
     "plot_mesh_sizes",
     "plot_model",
     "plot_function",
+    "debug_plot",
+    "debug_pvd",
 ]

spyro/plots/plots.py

@@ -214,3 +214,13 @@ def plot_function(function):
     fig.set_figheight = 9.0
     firedrake.tricontourf(function, axes=axes)
     axes.axis('equal')
+
+
+def debug_plot(function, filename="debug.png"):
+    plot_function(function)
+    plt.savefig(filename)
+
+
+def debug_pvd(function, filename="debug.pvd"):
+    out = firedrake.File(filename)
+    out.write(function)

spyro/solvers/acoustic_solver_construction_with_pml.py

@@ -26,6 +26,7 @@ def construct_solver_or_matrix_with_pml_2d(Wave_object):
     V = Wave_object.function_space
     Z = Wave_object.vector_function_space
     W = V * Z
+    Wave_object.mixed_function_space = W
     dxlump = dx(scheme=Wave_object.quadrature_rule)
     dslump = ds(scheme=Wave_object.surface_quadrature_rule)
 
@@ -92,6 +93,7 @@ def construct_solver_or_matrix_with_pml_3d(Wave_object):
     V = Wave_object.function_space
     Z = Wave_object.vector_function_space
     W = V * V * Z
+    Wave_object.mixed_function_space = W
     dxlump = dx(scheme=Wave_object.quadrature_rule)
     dslump = ds(scheme=Wave_object.surface_quadrature_rule)
 

spyro/solvers/acoustic_wave.py

@@ -144,3 +144,10 @@ def reset_pressure(self):
             self.u_n.assign(0.0)
         except:
             warnings.warn("No pressure to reset")
+        if self.abc_active:
+            try:
+                self.X_n.assign(0.0)
+                self.X_nm1.assign(0.0)
+            except:
+                warnings.warn("No mixed space pressure to reset")
+

spyro/solvers/backward_time_integration.py

@@ -8,19 +8,39 @@ def backward_wave_propagator(Wave_obj, dt=None):
 
     Parameters:
     -----------
+    Wave_obj: Spyro wave object
+        Wave object that already propagated a forward wave.
     dt: Python 'float' (optional)
         Time step to be used explicitly. If not mentioned uses the default,
         that was estabilished in the wave object for the adjoint model.
-    final_time: Python 'float' (optional)
-        Time which simulation ends. If not mentioned uses the default,
-        that was estabilished in the wave object.
 
     Returns:
     --------
-    usol: Firedrake 'Function'
-        Pressure wavefield at the final time.
-    u_rec: numpy array
-        Pressure wavefield at the receivers across the timesteps.
+    dJ: Firedrake 'Function'
+        Calculated gradient
+    """
+    if Wave_obj.abc_active is False:
+        return backward_wave_propagator_no_pml(Wave_obj, dt=dt)
+    elif Wave_obj.abc_active:
+        return mixed_space_backward_wave_propagator(Wave_obj, dt=dt)
+
+
+def backward_wave_propagator_no_pml(Wave_obj, dt=None):
+    """Propagates the adjoint wave backwards in time.
+    Currently uses central differences. Does not have any PML.
+
+    Parameters:
+    -----------
+    Wave_obj: Spyro wave object
+        Wave object that already propagated a forward wave.
+    dt: Python 'float' (optional)
+        Time step to be used explicitly. If not mentioned uses the default,
+        that was estabilished in the wave object for the adjoint model.
+
+    Returns:
+    --------
+    dJ: Firedrake 'Function'
+        Calculated gradient
     """
     Wave_obj.reset_pressure()
     if dt is not None:
@@ -39,12 +59,14 @@ def backward_wave_propagator(Wave_obj, dt=None):
     comm.comm.barrier()
 
     X = fire.Function(Wave_obj.function_space)
-    dJ = fire.Function(Wave_obj.function_space)#, name="gradient")
+    dJ = fire.Function(Wave_obj.function_space)  # , name="gradient")
 
     final_time = Wave_obj.final_time
     dt = Wave_obj.dt
     t = Wave_obj.current_time
-    nt = int((final_time - 0) / dt) + 1  # number of timesteps
+    if t != final_time:
+        print(f"Current time of {t}, different than final_time of {final_time}. Setting final_time to current time in backwards propagation.", flush= True)
+    nt = int(t / dt) + 1  # number of timesteps
 
     u_nm1 = Wave_obj.u_nm1
     u_n = Wave_obj.u_n
@@ -129,3 +151,123 @@ def backward_wave_propagator(Wave_obj, dt=None):
 
     dJ.dat.data_with_halos[:] *= (dt/2)
     return dJ
+
+
+def mixed_space_backward_wave_propagator(Wave_obj, dt=None):
+    """Propagates the adjoint wave backwards in time.
+    Currently uses central differences. Based on the
+    mixed space implementation of PML.
+
+    Parameters:
+    -----------
+    Wave_obj: Spyro wave object
+        Wave object that already propagated a forward wave.
+    dt: Python 'float' (optional)
+        Time step to be used explicitly. If not mentioned uses the default,
+        that was estabilished in the wave object for the adjoint model.
+
+    Returns:
+    --------
+    dJ: Firedrake 'Function'
+        Calculated gradient
+    """
+    Wave_obj.reset_pressure()
+    if dt is not None:
+        Wave_obj.dt = dt
+
+    forward_solution = Wave_obj.forward_solution
+    receivers = Wave_obj.receivers
+    residual = Wave_obj.misfit
+    comm = Wave_obj.comm
+    temp_filename = Wave_obj.forward_output_file
+
+    filename, file_extension = temp_filename.split(".")
+    output_filename = "backward." + file_extension
+
+    output = fire.File(output_filename, comm=comm.comm)
+    comm.comm.barrier()
+
+    X = Wave_obj.X
+    dJ = fire.Function(Wave_obj.function_space)  # , name="gradient")
+
+    final_time = Wave_obj.final_time
+    dt = Wave_obj.dt
+    t = Wave_obj.current_time
+    if t != final_time:
+        print(f"Current time of {t}, different than final_time of {final_time}. Setting final_time to current time in backwards propagation.", flush= True)
+    nt = int(t / dt) + 1  # number of timesteps
+
+    X_nm1 = Wave_obj.X_nm1
+    X_n = Wave_obj.X_n
+    X_np1 = fire.Function(Wave_obj.mixed_function_space)
+
+    rhs_forcing = fire.Function(Wave_obj.function_space)
+
+    B = Wave_obj.B
+    rhs = Wave_obj.rhs
+
+    # Define a gradient problem
+    m_u = fire.TrialFunction(Wave_obj.function_space)
+    m_v = fire.TestFunction(Wave_obj.function_space)
+    mgrad = m_u * m_v * fire.dx(scheme=Wave_obj.quadrature_rule)
+
+    # dufordt2 = fire.Function(Wave_obj.function_space)
+    ufor = fire.Function(Wave_obj.function_space)
+    uadj = fire.Function(Wave_obj.function_space)  # auxiliarly function for the gradient compt.
+
+    # ffG = -2 * (Wave_obj.c)**(-3) * fire.dot(dufordt2, uadj) * m_v * fire.dx(scheme=Wave_obj.quadrature_rule)
+    ffG =  2.0 * Wave_obj.c * fire.dot(fire.grad(uadj), fire.grad(ufor)) * m_v * fire.dx(scheme=Wave_obj.quadrature_rule)
+
+    lhsG = mgrad
+    rhsG = ffG
+
+    gradi = fire.Function(Wave_obj.function_space)
+    grad_prob = fire.LinearVariationalProblem(lhsG, rhsG, gradi)
+    grad_solver = fire.LinearVariationalSolver(
+        grad_prob,
+        solver_parameters={
+            "ksp_type": "preonly",
+            "pc_type": "jacobi",
+            "mat_type": "matfree",
+        },
+    )
+
+    # assembly_callable = create_assembly_callable(rhs, tensor=B)
+
+    for step in range(nt-1, -1, -1):
+        rhs_forcing.assign(0.0)
+        B = fire.assemble(rhs, tensor=B)
+        f = receivers.apply_receivers_as_source(rhs_forcing, residual, step)
+        B0 = B.sub(0)
+        B0 += f
+        Wave_obj.solver.solve(X, B)
+
+        X_np1.assign(X)
+
+        if (step) % Wave_obj.output_frequency == 0:
+            if Wave_obj.forward_output:
+                output.write(X_n.sub(0), time=t, name="Pressure")
+
+            helpers.display_progress(Wave_obj.comm, t)
+
+        if step % Wave_obj.gradient_sampling_frequency == 0:
+            # duadjdt2.assign( ((u_np1 - 2.0 * u_n + u_nm1) / fire.Constant(dt**2)) )
+            uadj.assign(X_np1.sub(0))
+            ufor.assign(forward_solution.pop())
+
+            grad_solver.solve()
+            if step == nt-1 or step == 0:
+                dJ += gradi
+            else:
+                dJ += 2*gradi
+
+        X_nm1.assign(X_n)
+        X_n.assign(X_np1)
+
+        t = step * float(dt)
+
+    Wave_obj.current_time = t
+    helpers.display_progress(Wave_obj.comm, t)
+
+    dJ.dat.data_with_halos[:] *= (dt/2)
+    return dJ