From b9be29efb8a45761b90302af9fdb26a0cb309d5f Mon Sep 17 00:00:00 2001
From: Eduardo Moscatelli de Souza <5752216+SouzaEM@users.noreply.github.com>
Date: Mon, 28 Oct 2024 15:05:10 -0300
Subject: [PATCH 1/7] Add Stacey boundary condition

---
 spyro/solvers/elastic_wave/forms.py     |  12 +--
 spyro/solvers/elastic_wave/local_abc.py | 126 +++++++++++++++++++++---
 2 files changed, 116 insertions(+), 22 deletions(-)

diff --git a/spyro/solvers/elastic_wave/forms.py b/spyro/solvers/elastic_wave/forms.py
index 598078ae..3df48f7a 100644
--- a/spyro/solvers/elastic_wave/forms.py
+++ b/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
 
diff --git a/spyro/solvers/elastic_wave/local_abc.py b/spyro/solvers/elastic_wave/local_abc.py
index a81744b4..8178d9b9 100644
--- a/spyro/solvers/elastic_wave/local_abc.py
+++ b/spyro/solvers/elastic_wave/local_abc.py
@@ -1,12 +1,20 @@
 from firedrake import (Constant, ds, TestFunction)
 
 
-def clayton_engquist_A1(wave):
+def local_abc_form(wave):
     '''
     Returns the linear form associated with the traction loads
-    when combined with the Clayton-Engquist A1 relations.
+    when combined with local absorbing boundary conditions.
     '''
-    F_t = 0  # linear form
+    abc_dict = wave.input_dictionary.get("absorving_boundary_conditions", None)
+    if abc_dict is None:
+        return 0
+    else:
+        abc_active = abc_dict.get("status", False)
+        if abc_active:
+            abc_type = abc_dict.get("local", {}).get("type", "Stacey")
+        else:
+            return 0
 
     V = wave.function_space
     v = TestFunction(V)
@@ -39,48 +47,78 @@ def clayton_engquist_A1(wave):
         uy_dz = u_n[iy].dx(iz)
         uy_dx = u_n[iy].dx(ix)
         uy_dy = u_n[iy].dx(iy)
+    else:
+        uy_dt = None
+        uz_dy = None
+        ux_dy = None
+        uy_dz = None
+        uy_dx = None
+        uy_dy = None
+
+    if abc_type == "Stacey":
+        callback = stacey_terms
+    elif abc_type == "CE_A1":
+        callback = clayton_engquist_A1_terms
+    else:
+        raise NotImplementedError(f"Unsupported local ABC: {abc_type}")
+
+    return callback(wave.dimension, rho, c_p, c_s,
+                    v, iz, ix, iy, qr_s,
+                    uz_dt, ux_dt, uy_dt,
+                    uz_dz, ux_dz, uy_dz,
+                    uz_dx, ux_dx, uy_dx,
+                    uz_dy, ux_dy, uy_dy)
+
+def clayton_engquist_A1_terms(ndim, rho, c_p, c_s,
+                              v, iz, ix, iy, qr_s,
+                              uz_dt, ux_dt, uy_dt,
+                              uz_dz, ux_dz, uy_dz,
+                              uz_dx, ux_dx, uy_dx,
+                              uz_dy, ux_dy, uy_dy):
+
+    F_t = 0
 
     # Plane z = -(Lz + pad)
     sig_zz = rho*c_p*uz_dt + rho*(c_p**2 - 2*c_s**2)*ux_dx
-    if wave.dimension == 3:
+    if ndim == 3:
         sig_zz += rho*(c_p**2 - 2*c_s**2)*uy_dy
     sig_xz = rho*c_s*ux_dt + rho*(c_s**2)*uz_dx
     F_t += -(sig_zz*v[iz] + sig_xz*v[ix])*ds(1, scheme=qr_s)
-    if wave.dimension == 3:
+    if ndim == 3:
         sig_yz = rho*c_s*uy_dt + rho*(c_s**2)*uz_dy
         F_t += -sig_yz*v[iy]*ds(1, scheme=qr_s)
 
     # Plane z = 0
     sig_zz = -rho*c_p*uz_dt + rho*(c_p**2 - 2*c_s**2)*ux_dx
-    if wave.dimension == 3:
+    if ndim == 3:
         sig_zz += rho*(c_p**2 - 2*c_s**2)*uy_dy
     sig_xz = -rho*c_s*ux_dt + rho*(c_s**2)*uz_dx
     F_t += (sig_zz*v[iz] + sig_xz*v[ix])*ds(2, scheme=qr_s)
-    if wave.dimension == 3:
+    if ndim == 3:
         sig_yz = -rho*c_s*uy_dt + rho*(c_s**2)*uz_dy
         F_t += sig_yz*v[iy]*ds(2, scheme=qr_s)
 
     # Plane x = -pad
     sig_zx = rho*c_s*uz_dt + rho*(c_s**2)*ux_dz
     sig_xx = rho*c_p*ux_dt + rho*(c_p**2 - 2*c_s**2)*uz_dz
-    if wave.dimension == 3:
+    if ndim == 3:
         sig_xx += rho*(c_p**2 - 2*c_s**2)*uy_dy
     F_t += -(sig_zx*v[iz] + sig_xx*v[ix])*ds(3, scheme=qr_s)
-    if wave.dimension == 3:
+    if ndim == 3:
         sig_yx = rho*c_s*uy_dt + rho*(c_s**2)*ux_dy
         F_t += -sig_yx*v[iy]*ds(3, scheme=qr_s)
 
     # Plane x = Lx + pad
     sig_zx = -rho*c_s*uz_dt + rho*(c_s**2)*ux_dz
     sig_xx = -rho*c_p*ux_dt + rho*(c_p**2 - 2*c_s**2)*uz_dz
-    if wave.dimension == 3:
+    if ndim == 3:
         sig_xx += rho*(c_p**2 - 2*c_s**2)*uy_dy
     F_t += (sig_zx*v[iz] + sig_xx*v[ix])*ds(4, scheme=qr_s)
-    if wave.dimension == 3:
+    if ndim == 3:
         sig_yx = -rho*c_s*uy_dt + rho*(c_s**2)*ux_dy
         F_t += sig_yx*v[iy]*ds(4, scheme=qr_s)
 
-    if wave.dimension == 3:
+    if ndim == 3:
         # Plane y = 0
         sig_zy = rho*c_s*uz_dt + rho*(c_s**2)*uy_dz
         sig_xy = rho*c_s*ux_dt + rho*(c_s**2)*uy_dx
@@ -94,3 +132,67 @@ def clayton_engquist_A1(wave):
         F_t += (sig_zy*v[iz] + sig_xy*v[ix] + sig_yy*v[iy])*ds(6, scheme=qr_s)
 
     return F_t
+
+def stacey_terms(ndim, rho, c_p, c_s,
+                 v, iz, ix, iy, qr_s,
+                 uz_dt, ux_dt, uy_dt,
+                 uz_dz, ux_dz, uy_dz,
+                 uz_dx, ux_dx, uy_dx,
+                 uz_dy, ux_dy, uy_dy):
+
+    F_t = 0
+
+    # Plane z = -(Lz + pad)
+    sig_zz = rho*c_p*uz_dt + rho*c_s*(c_p - 2*c_s)*ux_dx
+    if ndim == 3:
+        sig_zz += rho*c_s*(c_p - 2*c_s)*uy_dy
+    sig_xz = rho*c_s*ux_dt + rho*c_s*(2*c_s - c_p)*uz_dx
+    F_t += -(sig_zz*v[iz] + sig_xz*v[ix])*ds(1, scheme=qr_s)
+    if ndim == 3:
+        sig_yz = rho*c_s*uy_dt + rho*c_s*(2*c_s - c_p)*uz_dy
+        F_t += -sig_yz*v[iy]*ds(1, scheme=qr_s)
+
+    # Plane z = 0
+    sig_zz = -rho*c_p*uz_dt + rho*c_s*(c_p - 2*c_s)*ux_dx
+    if ndim == 3:
+        sig_zz += rho*c_s*(c_p - 2*c_s)*uy_dy
+    sig_xz = -rho*c_s*ux_dt - rho*c_s*(c_p - 2*c_s)*uz_dx
+    F_t += (sig_zz*v[iz] + sig_xz*v[ix])*ds(2, scheme=qr_s)
+    if ndim == 3:
+        sig_yz = -rho*c_s*uy_dt - rho*c_s*(c_p - 2*c_s)*uz_dy
+        F_t += sig_yz*v[iy]*ds(2, scheme=qr_s)
+
+    # Plane x = -pad
+    sig_zx = rho*c_s*uz_dt - rho*c_s*(c_p - 2*c_s)*ux_dz
+    sig_xx = rho*c_p*ux_dt + rho*c_s*(c_p - 2*c_s)*uz_dz
+    if ndim == 3:
+        sig_xx += rho*c_s*(c_p - 2*c_s)*uy_dy
+    F_t += -(sig_zx*v[iz] + sig_xx*v[ix])*ds(3, scheme=qr_s)
+    if ndim == 3:
+        sig_yx = rho*c_s*uy_dt - rho*c_s*(c_p - 2*c_s)*ux_dy
+        F_t += -sig_yx*v[iy]*ds(3, scheme=qr_s)
+
+    # Plane x = Lx + pad
+    sig_zx = -rho*c_s*uz_dt - rho*c_s*(c_p - 2*c_s)*ux_dz
+    sig_xx = -rho*c_p*ux_dt + rho*c_s*(c_p - 2*c_s)*uz_dz
+    if ndim == 3:
+        sig_xx += rho*c_s*(c_p - 2*c_s)*uy_dy
+    F_t += -(sig_zx*v[iz] + sig_xx*v[ix])*ds(3, scheme=qr_s)
+    if ndim == 3:
+        sig_yx = -rho*c_s*uy_dt - rho*c_s*(c_p - 2*c_s)*ux_dy
+        F_t += -sig_yx*v[iy]*ds(3, scheme=qr_s)
+
+    if ndim == 3:
+        # Plane y = 0
+        sig_zy = rho*c_s*uz_dt - rho*c_s*(c_p - 2*c_s)*uy_dz
+        sig_xy = rho*c_s*ux_dt - rho*c_s*(c_p - 2*c_s)*uy_dx
+        sig_yy = rho*c_p*uy_dt + rho*c_s*(c_p - 2*c_s)*(uz_dz + ux_dx)
+        F_t += -(sig_zy*v[iz] + sig_xy*v[ix] + sig_yy*v[iy])*ds(5, scheme=qr_s)
+
+        # Plane y = L_y + 2*pad
+        sig_zy = -rho*c_s*uz_dt - rho*c_s*(c_p - 2*c_s)*uy_dz
+        sig_xy = -rho*c_s*ux_dt - rho*c_s*(c_p - 2*c_s)*uy_dx
+        sig_yy = -rho*c_p*uy_dt + rho*c_s*(c_p - 2*c_s)*(uz_dz + ux_dx)
+        F_t += (sig_zy*v[iz] + sig_xy*v[ix] + sig_yy*v[iy])*ds(6, scheme=qr_s)
+
+    return F_t
\ No newline at end of file

From 00e9a524caa42d0942e88e58bac9e0ed9c923865 Mon Sep 17 00:00:00 2001
From: Eduardo Moscatelli de Souza <5752216+SouzaEM@users.noreply.github.com>
Date: Wed, 30 Oct 2024 09:43:53 -0300
Subject: [PATCH 2/7] Fix sign of Stacey BC and add central difference option
 to local ABCs

---
 spyro/solvers/elastic_wave/local_abc.py | 27 +++++++++++++++++--------
 1 file changed, 19 insertions(+), 8 deletions(-)

diff --git a/spyro/solvers/elastic_wave/local_abc.py b/spyro/solvers/elastic_wave/local_abc.py
index 8178d9b9..0991715a 100644
--- a/spyro/solvers/elastic_wave/local_abc.py
+++ b/spyro/solvers/elastic_wave/local_abc.py
@@ -1,4 +1,4 @@
-from firedrake import (Constant, ds, TestFunction)
+from firedrake import (Constant, ds, TestFunction, TrialFunction)
 
 
 def local_abc_form(wave):
@@ -13,6 +13,7 @@ def local_abc_form(wave):
         abc_active = abc_dict.get("status", False)
         if abc_active:
             abc_type = abc_dict.get("local", {}).get("type", "Stacey")
+            dt_scheme = abc_dict.get("local", {}).get("dt_scheme", "backward")
         else:
             return 0
 
@@ -34,14 +35,24 @@ def local_abc_form(wave):
     iy = 2
 
     # Partial derivatives
-    uz_dt = (u_n[iz] - u_nm1[iz])/dt
-    ux_dt = (u_n[ix] - u_nm1[ix])/dt
+    if dt_scheme == "backward":
+        uz_dt = (u_n[iz] - u_nm1[iz])/dt
+        ux_dt = (u_n[ix] - u_nm1[ix])/dt
+    elif dt_scheme == "central":
+        u = TrialFunction(V)
+        uz_dt = (u[iz] - u_nm1[iz])/(2*dt)
+        ux_dt = (u[ix] - u_nm1[ix])/(2*dt)
+    else:
+        raise NotImplementedError(f"Unsupported time discretization: {dt_scheme}")
     uz_dz = u_n[iz].dx(iz)
     uz_dx = u_n[iz].dx(ix)
     ux_dz = u_n[ix].dx(iz)
     ux_dx = u_n[ix].dx(ix)
     if wave.dimension == 3:
-        uy_dt = (u_n[iy] - u_nm1[iy])/dt
+        if dt_scheme == "backward":
+            uy_dt = (u_n[iy] - u_nm1[iy])/dt
+        elif dt_scheme == "central":
+            uy_dt = (u[iy] - u_nm1[iy])/(2*dt)
         uz_dy = u_n[iz].dx(iy)
         ux_dy = u_n[ix].dx(iy)
         uy_dz = u_n[iy].dx(iz)
@@ -146,10 +157,10 @@ def stacey_terms(ndim, rho, c_p, c_s,
     sig_zz = rho*c_p*uz_dt + rho*c_s*(c_p - 2*c_s)*ux_dx
     if ndim == 3:
         sig_zz += rho*c_s*(c_p - 2*c_s)*uy_dy
-    sig_xz = rho*c_s*ux_dt + rho*c_s*(2*c_s - c_p)*uz_dx
+    sig_xz = rho*c_s*ux_dt - rho*c_s*(c_p - 2*c_s)*uz_dx
     F_t += -(sig_zz*v[iz] + sig_xz*v[ix])*ds(1, scheme=qr_s)
     if ndim == 3:
-        sig_yz = rho*c_s*uy_dt + rho*c_s*(2*c_s - c_p)*uz_dy
+        sig_yz = rho*c_s*uy_dt - rho*c_s*(c_p - 2*c_s)*uz_dy
         F_t += -sig_yz*v[iy]*ds(1, scheme=qr_s)
 
     # Plane z = 0
@@ -177,10 +188,10 @@ def stacey_terms(ndim, rho, c_p, c_s,
     sig_xx = -rho*c_p*ux_dt + rho*c_s*(c_p - 2*c_s)*uz_dz
     if ndim == 3:
         sig_xx += rho*c_s*(c_p - 2*c_s)*uy_dy
-    F_t += -(sig_zx*v[iz] + sig_xx*v[ix])*ds(3, scheme=qr_s)
+    F_t += (sig_zx*v[iz] + sig_xx*v[ix])*ds(3, scheme=qr_s)
     if ndim == 3:
         sig_yx = -rho*c_s*uy_dt - rho*c_s*(c_p - 2*c_s)*ux_dy
-        F_t += -sig_yx*v[iy]*ds(3, scheme=qr_s)
+        F_t += sig_yx*v[iy]*ds(3, scheme=qr_s)
 
     if ndim == 3:
         # Plane y = 0

From 472e8b36374d2d6b0b267a140f40fc1d04c45e36 Mon Sep 17 00:00:00 2001
From: Eduardo Moscatelli de Souza <5752216+SouzaEM@users.noreply.github.com>
Date: Thu, 31 Oct 2024 12:38:09 -0300
Subject: [PATCH 3/7] Add mechanical energy log

---
 spyro/io/field_logger.py                     | 55 ++++++++++++++++++++
 spyro/solvers/elastic_wave/functionals.py    | 21 ++++++++
 spyro/solvers/elastic_wave/isotropic_wave.py |  8 ++-
 3 files changed, 83 insertions(+), 1 deletion(-)
 create mode 100644 spyro/solvers/elastic_wave/functionals.py

diff --git a/spyro/io/field_logger.py b/spyro/io/field_logger.py
index 194984b7..b93dfc68 100644
--- a/spyro/io/field_logger.py
+++ b/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.append((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,35 @@ 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 = []
+
+            self.__enabled_functionals = []
+            for key, callback in self.__func_data:
+                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.append(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:
+                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:
+                functional.sample()
\ No newline at end of file
diff --git a/spyro/solvers/elastic_wave/functionals.py b/spyro/solvers/elastic_wave/functionals.py
new file mode 100644
index 00000000..5f90c6ce
--- /dev/null
+++ b/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
\ No newline at end of file
diff --git a/spyro/solvers/elastic_wave/isotropic_wave.py b/spyro/solvers/elastic_wave/isotropic_wave.py
index 565db458..dd25d7bb 100644
--- a/spyro/solvers/elastic_wave/isotropic_wave.py
+++ b/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
 
@@ -43,6 +44,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):
@@ -148,6 +153,7 @@ def matrix_building(self):
                               name=self.get_function_name())
         self.u_np1 = Function(self.function_space,
                               name=self.get_function_name())
+        self.mechanical_energy = mechanical_energy_form(self)
 
         self.parse_initial_conditions()
         self.parse_boundary_conditions()

From 30ecd3885c4315f035d10aa7274d983b692b96dd Mon Sep 17 00:00:00 2001
From: Eduardo Moscatelli de Souza <5752216+SouzaEM@users.noreply.github.com>
Date: Thu, 31 Oct 2024 15:00:14 -0300
Subject: [PATCH 4/7] Fix xmax edge index in Stacey ABC

---
 spyro/solvers/elastic_wave/local_abc.py | 6 +++---
 1 file changed, 3 insertions(+), 3 deletions(-)

diff --git a/spyro/solvers/elastic_wave/local_abc.py b/spyro/solvers/elastic_wave/local_abc.py
index 0991715a..fd0a7467 100644
--- a/spyro/solvers/elastic_wave/local_abc.py
+++ b/spyro/solvers/elastic_wave/local_abc.py
@@ -188,10 +188,10 @@ def stacey_terms(ndim, rho, c_p, c_s,
     sig_xx = -rho*c_p*ux_dt + rho*c_s*(c_p - 2*c_s)*uz_dz
     if ndim == 3:
         sig_xx += rho*c_s*(c_p - 2*c_s)*uy_dy
-    F_t += (sig_zx*v[iz] + sig_xx*v[ix])*ds(3, scheme=qr_s)
+    F_t += (sig_zx*v[iz] + sig_xx*v[ix])*ds(4, scheme=qr_s)
     if ndim == 3:
         sig_yx = -rho*c_s*uy_dt - rho*c_s*(c_p - 2*c_s)*ux_dy
-        F_t += sig_yx*v[iy]*ds(3, scheme=qr_s)
+        F_t += sig_yx*v[iy]*ds(4, scheme=qr_s)
 
     if ndim == 3:
         # Plane y = 0
@@ -206,4 +206,4 @@ def stacey_terms(ndim, rho, c_p, c_s,
         sig_yy = -rho*c_p*uy_dt + rho*c_s*(c_p - 2*c_s)*(uz_dz + ux_dx)
         F_t += (sig_zy*v[iz] + sig_xy*v[ix] + sig_yy*v[iy])*ds(6, scheme=qr_s)
 
-    return F_t
\ No newline at end of file
+    return F_t

From 550bdb69a84d526dbb303b5945fb49d80cc49765 Mon Sep 17 00:00:00 2001
From: Eduardo Moscatelli de Souza <5752216+SouzaEM@users.noreply.github.com>
Date: Thu, 31 Oct 2024 15:53:31 -0300
Subject: [PATCH 5/7] Add option to run local ABCs with 2nd order backward
 velocity discretization

---
 spyro/solvers/elastic_wave/isotropic_wave.py | 13 +++++++++++++
 spyro/solvers/elastic_wave/local_abc.py      |  6 ++++++
 2 files changed, 19 insertions(+)

diff --git a/spyro/solvers/elastic_wave/isotropic_wave.py b/spyro/solvers/elastic_wave/isotropic_wave.py
index dd25d7bb..fd9b11f6 100644
--- a/spyro/solvers/elastic_wave/isotropic_wave.py
+++ b/spyro/solvers/elastic_wave/isotropic_wave.py
@@ -24,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)
@@ -113,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
@@ -153,6 +156,16 @@ def matrix_building(self):
                               name=self.get_function_name())
         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()
diff --git a/spyro/solvers/elastic_wave/local_abc.py b/spyro/solvers/elastic_wave/local_abc.py
index fd0a7467..86bfb98e 100644
--- a/spyro/solvers/elastic_wave/local_abc.py
+++ b/spyro/solvers/elastic_wave/local_abc.py
@@ -38,6 +38,10 @@ def local_abc_form(wave):
     if dt_scheme == "backward":
         uz_dt = (u_n[iz] - u_nm1[iz])/dt
         ux_dt = (u_n[ix] - u_nm1[ix])/dt
+    elif dt_scheme == "backward_2nd":
+        u_nm2 = wave.u_nm2
+        uz_dt = (3*u_n[iz] - 4*u_nm1[iz] + u_nm2[iz])/(2*dt)
+        ux_dt = (3*u_n[ix] - 4*u_nm1[ix] + u_nm2[ix])/(2*dt)
     elif dt_scheme == "central":
         u = TrialFunction(V)
         uz_dt = (u[iz] - u_nm1[iz])/(2*dt)
@@ -51,6 +55,8 @@ def local_abc_form(wave):
     if wave.dimension == 3:
         if dt_scheme == "backward":
             uy_dt = (u_n[iy] - u_nm1[iy])/dt
+        elif dt_scheme == "backward_2nd":
+            uy_dt = (3*u_n[iy] - 4*u_nm1[iy] + u_nm2[iy])/(2*dt)
         elif dt_scheme == "central":
             uy_dt = (u[iy] - u_nm1[iy])/(2*dt)
         uz_dy = u_n[iz].dx(iy)

From f542957b8a3a651d5a26917907aa40062b0c9753 Mon Sep 17 00:00:00 2001
From: Eduardo Moscatelli de Souza <5752216+SouzaEM@users.noreply.github.com>
Date: Fri, 1 Nov 2024 14:20:22 -0300
Subject: [PATCH 6/7] Add tests for local ABCs

---
 spyro/examples/elastic_local_abc.py | 94 +++++++++++++++++++++++++++++
 spyro/io/field_logger.py            | 20 +++---
 test/test_elastic_local_abc.py      | 33 ++++++++++
 3 files changed, 138 insertions(+), 9 deletions(-)
 create mode 100644 spyro/examples/elastic_local_abc.py
 create mode 100644 test/test_elastic_local_abc.py

diff --git a/spyro/examples/elastic_local_abc.py b/spyro/examples/elastic_local_abc.py
new file mode 100644
index 00000000..a403b4c9
--- /dev/null
+++ b/spyro/examples/elastic_local_abc.py
@@ -0,0 +1,94 @@
+import firedrake as fire
+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
\ No newline at end of file
diff --git a/spyro/io/field_logger.py b/spyro/io/field_logger.py
index b93dfc68..4194bdb7 100644
--- a/spyro/io/field_logger.py
+++ b/spyro/io/field_logger.py
@@ -40,8 +40,8 @@ def __init__(self, comm, vis_dict):
 
         self.__rank = comm.comm.Get_rank()
         if self.__rank == 0:
-            self.__func_data = []
-            self.__enabled_functionals = []
+            self.__func_data = {}
+            self.__enabled_functionals = {}
 
             self.__time_enabled = self.vis_dict.get("time", False)
             if self.__time_enabled:
@@ -54,7 +54,7 @@ def add_field(self, key, name, callback):
 
     def add_functional(self, key, callback):
         if self.__rank == 0:
-            self.__func_data.append((key, callback))
+            self.__func_data[key] = callback
 
     def start_logging(self, source_id):
         if self.__source_id is not None:
@@ -78,13 +78,12 @@ def start_logging(self, source_id):
             if self.__time_enabled:
                 self.__time = []
 
-            self.__enabled_functionals = []
-            for key, callback in self.__func_data:
+            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.append(Functional(filename, callback))
+                    self.__enabled_functionals[key] = Functional(filename, callback)
 
     def stop_logging(self):
         self.__source_id = None
@@ -93,7 +92,7 @@ def stop_logging(self):
             if self.__time_enabled:
                 np.save(self.__time_filename, self.__time)
 
-            for functional in self.__enabled_functionals:
+            for functional in self.__enabled_functionals.values():
                 functional.save()
 
     def log(self, t):
@@ -104,5 +103,8 @@ def log(self, t):
             if self.__time_enabled:
                 self.__time.append(t)
 
-            for functional in self.__enabled_functionals:
-                functional.sample()
\ No newline at end of file
+            for functional in self.__enabled_functionals.values():
+                functional.sample()
+
+    def get(self, key):
+        return self.__enabled_functionals[key].list[-1]
\ No newline at end of file
diff --git a/test/test_elastic_local_abc.py b/test/test_elastic_local_abc.py
new file mode 100644
index 00000000..7ff1a416
--- /dev/null
+++ b/test/test_elastic_local_abc.py
@@ -0,0 +1,33 @@
+import pytest
+
+from spyro.examples.elastic_local_abc import build_solver
+
+# This value was obtained empirically. It is supposed for backward compatibility
+expected_mechanical_energy = 0.25
+
+
+def test_stacey_abc():
+    wave = build_solver("Stacey", "backward")
+    wave.forward_solve()
+    last_mechanical_energy = wave.field_logger.get("mechanical_energy")
+    assert last_mechanical_energy < expected_mechanical_energy
+
+
+def test_clayton_engquist_abc():
+    wave = build_solver("CE_A1", "backward")
+    wave.forward_solve()
+    last_mechanical_energy = wave.field_logger.get("mechanical_energy")
+    assert last_mechanical_energy < expected_mechanical_energy
+
+
+def test_with_central():
+    wave = build_solver("Stacey", "central")
+    with pytest.raises(AssertionError) as e:
+        wave.forward_solve()
+
+
+def test_with_backward_2nd():
+    wave = build_solver("Stacey", "backward_2nd")
+    wave.forward_solve()
+    last_mechanical_energy = wave.field_logger.get("mechanical_energy")
+    assert last_mechanical_energy < expected_mechanical_energy
\ No newline at end of file

From 83e95c4654d4a7d7bc1e711f6622c31b1f1c1561 Mon Sep 17 00:00:00 2001
From: Eduardo Moscatelli de Souza <5752216+SouzaEM@users.noreply.github.com>
Date: Fri, 1 Nov 2024 14:28:43 -0300
Subject: [PATCH 7/7] Lint modifications from issue 130

---
 spyro/examples/elastic_local_abc.py       | 4 ++--
 spyro/io/field_logger.py                  | 2 +-
 spyro/solvers/elastic_wave/functionals.py | 2 +-
 spyro/solvers/elastic_wave/local_abc.py   | 2 ++
 test/test_elastic_local_abc.py            | 4 ++--
 5 files changed, 8 insertions(+), 6 deletions(-)

diff --git a/spyro/examples/elastic_local_abc.py b/spyro/examples/elastic_local_abc.py
index a403b4c9..379da327 100644
--- a/spyro/examples/elastic_local_abc.py
+++ b/spyro/examples/elastic_local_abc.py
@@ -1,4 +1,3 @@
-import firedrake as fire
 import numpy as np
 import spyro
 
@@ -21,6 +20,7 @@
 final_time = 2.0  # [s]
 out_freq = int(0.01/time_step)
 
+
 def build_solver(local_abc, dt_scheme):
     d = {}
 
@@ -91,4 +91,4 @@ def build_solver(local_abc, dt_scheme):
     wave = spyro.IsotropicWave(d)
     wave.set_mesh(mesh_parameters={'dx': h})
 
-    return wave
\ No newline at end of file
+    return wave
diff --git a/spyro/io/field_logger.py b/spyro/io/field_logger.py
index 4194bdb7..358dbe88 100644
--- a/spyro/io/field_logger.py
+++ b/spyro/io/field_logger.py
@@ -107,4 +107,4 @@ def log(self, t):
                 functional.sample()
 
     def get(self, key):
-        return self.__enabled_functionals[key].list[-1]
\ No newline at end of file
+        return self.__enabled_functionals[key].list[-1]
diff --git a/spyro/solvers/elastic_wave/functionals.py b/spyro/solvers/elastic_wave/functionals.py
index 5f90c6ce..e517b284 100644
--- a/spyro/solvers/elastic_wave/functionals.py
+++ b/spyro/solvers/elastic_wave/functionals.py
@@ -18,4 +18,4 @@ def mechanical_energy_form(wave):
     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
\ No newline at end of file
+    return K + U
diff --git a/spyro/solvers/elastic_wave/local_abc.py b/spyro/solvers/elastic_wave/local_abc.py
index 86bfb98e..de515113 100644
--- a/spyro/solvers/elastic_wave/local_abc.py
+++ b/spyro/solvers/elastic_wave/local_abc.py
@@ -86,6 +86,7 @@ def local_abc_form(wave):
                     uz_dx, ux_dx, uy_dx,
                     uz_dy, ux_dy, uy_dy)
 
+
 def clayton_engquist_A1_terms(ndim, rho, c_p, c_s,
                               v, iz, ix, iy, qr_s,
                               uz_dt, ux_dt, uy_dt,
@@ -150,6 +151,7 @@ def clayton_engquist_A1_terms(ndim, rho, c_p, c_s,
 
     return F_t
 
+
 def stacey_terms(ndim, rho, c_p, c_s,
                  v, iz, ix, iy, qr_s,
                  uz_dt, ux_dt, uy_dt,
diff --git a/test/test_elastic_local_abc.py b/test/test_elastic_local_abc.py
index 7ff1a416..724b7896 100644
--- a/test/test_elastic_local_abc.py
+++ b/test/test_elastic_local_abc.py
@@ -22,7 +22,7 @@ def test_clayton_engquist_abc():
 
 def test_with_central():
     wave = build_solver("Stacey", "central")
-    with pytest.raises(AssertionError) as e:
+    with pytest.raises(AssertionError):
         wave.forward_solve()
 
 
@@ -30,4 +30,4 @@ def test_with_backward_2nd():
     wave = build_solver("Stacey", "backward_2nd")
     wave.forward_solve()
     last_mechanical_energy = wave.field_logger.get("mechanical_energy")
-    assert last_mechanical_energy < expected_mechanical_energy
\ No newline at end of file
+    assert last_mechanical_energy < expected_mechanical_energy