Merge pull request #131 from NDF-Poli-USP/issue_0130_stacey

Implement Stacey boundary condition

spyro/examples/elastic_local_abc.py

@@ -0,0 +1,94 @@
+import numpy as np
+import spyro
+
+
+output_dir = "results"
+
+L = 3000  # Edge size [m]
+n = 60    # Number of elements in each direction
+h = L/n   # Element size [m]
+
+rho = 2700     # Density [kg/m3]
+Vp = 3000      # P wave velocity [m/s]
+Vs = 1732      # S wave velocity [m/s]
+
+smag = 1e6
+freq = 10  # Central frequency of Ricker wavelet [Hz]
+source_locations = [[-L/2, L/2]]
+
+time_step = 2e-4  # [s]
+final_time = 2.0  # [s]
+out_freq = int(0.01/time_step)
+
+
+def build_solver(local_abc, dt_scheme):
+    d = {}
+
+    d["options"] = {
+        "cell_type": "T",
+        "variant": "lumped",
+        "degree": 4,
+        "dimension": 2,
+    }
+
+    d["parallelism"] = {
+        "type": "automatic",
+    }
+
+    d["mesh"] = {
+        "Lz": L,
+        "Lx": L,
+        "Ly": 0,
+        "h": h,
+        "mesh_file": None,
+        "mesh_type": "firedrake_mesh",
+    }
+
+    d["acquisition"] = {
+        "source_type": "ricker",
+        "source_locations": source_locations,
+        "frequency": freq,
+        "delay": 1.5,
+        "delay_type": "multiples_of_minimun",
+        "amplitude": smag * np.array([0, 1]),
+        "receiver_locations": [],
+    }
+
+    d["synthetic_data"] = {
+        "type": "object",
+        "density": rho,
+        "p_wave_velocity": Vp,
+        "s_wave_velocity": Vs,
+        "real_velocity_file": None,
+    }
+
+    d["time_axis"] = {
+        "initial_time": 0,
+        "final_time": final_time,
+        "dt": time_step,
+        "output_frequency": out_freq,
+        "gradient_sampling_frequency": 1,
+    }
+
+    d["visualization"] = {
+        "forward_output": False,
+        "time": True,
+        "time_filename": f"{output_dir}/time.npy",
+        "mechanical_energy": True,
+        "mechanical_energy_filename": f"{output_dir}/mechanical_energy_{local_abc}_{dt_scheme}.npy",
+    }
+
+    if local_abc is not None:
+        d["absorving_boundary_conditions"] = {
+            "status": True,
+            "damping_type": "local",
+            "local": {
+                "type": local_abc,
+                "dt_scheme": dt_scheme,
+            },
+        }
+
+    wave = spyro.IsotropicWave(d)
+    wave.set_mesh(mesh_parameters={'dx': h})
+
+    return wave

spyro/io/field_logger.py

@@ -1,3 +1,4 @@
+import numpy as np
 import warnings
 
 from firedrake import VTKFile
@@ -15,6 +16,19 @@ def write(self, t):
         self.file.write(self.callback(), time=t, name=self.name)
 
 
+class Functional:
+    def __init__(self, filename, callback):
+        self.filename = filename
+        self.callback = callback
+        self.list = []
+
+    def sample(self):
+        self.list.append(self.callback())
+
+    def save(self):
+        np.save(self.filename, self.list)
+
+
 class FieldLogger:
     def __init__(self, comm, vis_dict):
         self.comm = comm
@@ -24,9 +38,24 @@ def __init__(self, comm, vis_dict):
         self.__enabled_fields = []
         self.__wave_data = []
 
+        self.__rank = comm.comm.Get_rank()
+        if self.__rank == 0:
+            self.__func_data = {}
+            self.__enabled_functionals = {}
+
+            self.__time_enabled = self.vis_dict.get("time", False)
+            if self.__time_enabled:
+                self.__time = []
+                self.__time_filename = self.vis_dict.get("time_filename", "time.npy")
+                print(f"Saving time in: {self.__time_filename}")
+
     def add_field(self, key, name, callback):
         self.__wave_data.append((key, name, callback))
 
+    def add_functional(self, key, callback):
+        if self.__rank == 0:
+            self.__func_data[key] = callback
+
     def start_logging(self, source_id):
         if self.__source_id is not None:
             warnings.warn("Started a new record without stopping the previous one")
@@ -45,9 +74,37 @@ def start_logging(self, source_id):
                 file = VTKFile(filename, comm=self.comm.comm)
                 self.__enabled_fields.append(Field(name, file, callback))
 
+        if self.__rank == 0:
+            if self.__time_enabled:
+                self.__time = []
+
+            for key, callback in self.__func_data.items():
+                enabled = self.vis_dict.get(key, False)
+                if enabled:
+                    filename = self.vis_dict.get(key + "_filename", key + ".npy")
+                    print(f"Saving {key} in: {filename}")
+                    self.__enabled_functionals[key] = Functional(filename, callback)
+
     def stop_logging(self):
         self.__source_id = None
 
+        if self.__rank == 0:
+            if self.__time_enabled:
+                np.save(self.__time_filename, self.__time)
+
+            for functional in self.__enabled_functionals.values():
+                functional.save()
+
     def log(self, t):
         for field in self.__enabled_fields:
             field.write(t)
+
+        if self.__rank == 0:
+            if self.__time_enabled:
+                self.__time.append(t)
+
+            for functional in self.__enabled_functionals.values():
+                functional.sample()
+
+    def get(self, key):
+        return self.__enabled_functionals[key].list[-1]

spyro/solvers/elastic_wave/forms.py

@@ -1,7 +1,7 @@
 from firedrake import (assemble, Cofunction, Constant, div, dot, dx, grad,
                        inner, lhs, LinearSolver, rhs, TestFunction, TrialFunction)
 
-from .local_abc import clayton_engquist_A1
+from .local_abc import local_abc_form
 
 
 def isotropic_elastic_without_pml(wave):
@@ -30,15 +30,7 @@ def isotropic_elastic_without_pml(wave):
     if b is not None:
         F_s += dot(b, v)*dx(scheme=quad_rule)
 
-    abc_dict = wave.input_dictionary.get("absorving_boundary_conditions", None)
-    if abc_dict is None:
-        F_t = 0
-    else:
-        abc_active = abc_dict.get("status", False)
-        if abc_active:
-            F_t = clayton_engquist_A1(wave)
-        else:
-            F_t = 0
+    F_t = local_abc_form(wave)
 
     F = F_m + F_k - F_s - F_t
 

spyro/solvers/elastic_wave/functionals.py

@@ -0,0 +1,21 @@
+from firedrake import (Constant, div, dx, grad, inner)
+
+
+def mechanical_energy_form(wave):
+    u_nm1 = wave.u_nm1
+    u_n = wave.u_n
+
+    dt = Constant(wave.dt)
+    rho = wave.rho
+    lmbda = wave.lmbda
+    mu = wave.mu
+
+    # Kinetic energy
+    v = (u_n - u_nm1)/dt
+    K = (rho/2)*inner(v, v)*dx
+
+    # Strain energy
+    eps = lambda v: 0.5*(grad(v) + grad(v).T)
+    U = (lmbda*div(u_n)*div(u_n) + 2*mu*inner(eps(u_n), eps(u_n)))*dx
+
+    return K + U

spyro/solvers/elastic_wave/isotropic_wave.py

@@ -1,11 +1,12 @@
 import numpy as np
 
-from firedrake import (Constant, curl, DirichletBC, div, Function,
+from firedrake import (assemble, Constant, curl, DirichletBC, div, Function,
                        FunctionSpace, project, VectorFunctionSpace)
 
 from .elastic_wave import ElasticWave
 from .forms import (isotropic_elastic_without_pml,
                     isotropic_elastic_with_pml)
+from .functionals import mechanical_energy_form
 from ...domains.space import FE_method
 from ...utils.typing import override
 
@@ -23,6 +24,7 @@ def __init__(self, dictionary, comm=None):
 
         self.u_n = None   # Current displacement field
         self.u_nm1 = None  # Displacement field in previous iteration
+        self.u_nm2 = None  # Displacement field at iteration n-2
         self.u_np1 = None  # Displacement field in next iteration
 
         # Volumetric sourcers (defined through UFL)
@@ -43,6 +45,10 @@ def __init__(self, dictionary, comm=None):
         self.field_logger.add_field("s-wave", "S-wave",
                                     lambda: self.update_s_wave())
 
+        self.mechanical_energy = None
+        self.field_logger.add_functional("mechanical_energy",
+                                         lambda: assemble(self.mechanical_energy))
+
     @override
     def initialize_model_parameters_from_object(self, synthetic_data_dict: dict):
         def constant_wrapper(value):
@@ -108,6 +114,8 @@ def _get_vstate(self):
 
     @override
     def _set_prev_vstate(self, vstate):
+        if self.u_nm2 is not None:
+            self.u_nm2.assign(self.u_nm1)
         self.u_nm1.assign(vstate)
 
     @override
@@ -149,6 +157,17 @@ def matrix_building(self):
         self.u_np1 = Function(self.function_space,
                               name=self.get_function_name())
 
+        abc_dict = self.input_dictionary.get("absorving_boundary_conditions", None)
+        if abc_dict is not None:
+            abc_active = abc_dict.get("status", False)
+            if abc_active:
+                dt_scheme = abc_dict.get("local", {}).get("dt_scheme", None)
+                if dt_scheme == "backward_2nd":
+                    self.u_nm2 = Function(self.function_space,
+                                          name=self.get_function_name())
+
+        self.mechanical_energy = mechanical_energy_form(self)
+
         self.parse_initial_conditions()
         self.parse_boundary_conditions()
         self.parse_volumetric_forces()