Merge pull request #53 from Olender/issue_0041

Add first part of the code of elastic wave propagation. It is still incomplete, but include some code refactoring that may be used by other branches.

notebook_tutorials/altering_time_integration.ipynb

@@ -338,7 +338,7 @@
    "metadata": {},
    "outputs": [],
    "source": [
-    "Wave_obj._get_initial_velocity_model()\n",
+    "Wave_obj._initialize_model_parameters()\n",
     "Wave_obj.c = Wave_obj.initial_velocity_model\n",
     "Wave_obj.matrix_building()"
    ]
@@ -347,7 +347,7 @@
    "cell_type": "markdown",
    "metadata": {},
    "source": [
-    "The `_get_initial_velocity_model` is not actually necessary for this specific case (where the velocity model is already created in the object) and does not do anything in this case. However, it needs to be called, because in larger cases the velocity model is only loaded, by this method, into our function space right before it is needed (to conserve memory). Spyro will build the solver object based on the velocity model that the `c` attribute points to. This can be diferent than the `initial_velocity_model`, especially in inversion problems."
+    "The `_initialize_model_parameters` is not actually necessary for this specific case (where the velocity model is already created in the object) and does not do anything in this case. However, it needs to be called, because in larger cases the velocity model is only loaded, by this method, into our function space right before it is needed (to conserve memory). Spyro will build the solver object based on the velocity model that the `c` attribute points to. This can be diferent than the `initial_velocity_model`, especially in inversion problems."
    ]
   },
   {

notebook_tutorials/simple_forward_with_overthrust.html

@@ -8051,7 +8051,7 @@ <h1 id="Simple-forward-with-overthrust">Simple forward with overthrust<a class="
 <div class="jp-InputArea jp-Cell-inputArea"><div class="jp-InputPrompt jp-InputArea-prompt">
 </div><div class="jp-RenderedHTMLCommon jp-RenderedMarkdown jp-MarkdownOutput" data-mime-type="text/markdown">
 <p>From the plot above, we can notice two things. First, it is rotated by 90 degrees. In cases like this, the saved PNG image will be in the correct orientation. However, the plotted image in this notebook is rotated because the velocity model we used, commonly with segy files, has the Z-axis before the X-axis in data input.</p>
-<p>Another observation is that, unlike in the <strong>simple forward tutorial</strong>, we looked at the experiment layout after running the forward solve. For memory purposes, a 2D or 3D velocity file is only actually interpolated into our domain when it is necessary for another method of the Wave object class. If you need to force the interpolation sooner, call the _get_initial_velocity_model() method.</p>
+<p>Another observation is that, unlike in the <strong>simple forward tutorial</strong>, we looked at the experiment layout after running the forward solve. For memory purposes, a 2D or 3D velocity file is only actually interpolated into our domain when it is necessary for another method of the Wave object class. If you need to force the interpolation sooner, call the _initialize_model_parameters() method.</p>
 <p>It is also important to note that even though receivers look like a line, they are actually located in points, which can be visible by zooming into the image, not coinciding with nodes.</p>
 </div>
 </div>

notebook_tutorials/simple_forward_with_overthrust.ipynb

@@ -397,7 +397,7 @@
    "source": [
     "From the plot above, we can notice two things. First, it is rotated by 90 degrees. In cases like this, the saved PNG image will be in the correct orientation. However, the plotted image in this notebook is rotated because the velocity model we used, commonly with segy files, has the Z-axis before the X-axis in data input.\n",
     "\n",
-    "Another observation is that, unlike in the **simple forward tutorial**, we looked at the experiment layout after running the forward solve. For memory purposes, a 2D or 3D velocity file is only actually interpolated into our domain when it is necessary for another method of the Wave object class. If you need to force the interpolation sooner, call the _get_initial_velocity_model() method.\n",
+    "Another observation is that, unlike in the **simple forward tutorial**, we looked at the experiment layout after running the forward solve. For memory purposes, a 2D or 3D velocity file is only actually interpolated into our domain when it is necessary for another method of the Wave object class. If you need to force the interpolation sooner, call the _initialize_model_parameters() method.\n",
     "\n",
     "It is also important to note that even though receivers look like a line, they are actually located in points, which can be visible by zooming into the image, not coinciding with nodes."
    ]

spyro/examples/camembert.py

@@ -156,3 +156,7 @@ def _camembert_velocity_model(self):
         )
         self.set_initial_velocity_model(conditional=cond, dg_velocity_model=False)
         return None
+
+if __name__ == "__main__":
+    wave = Camembert_acoustic()
+    wave.forward_solve()

spyro/examples/rectangle.py

@@ -182,25 +182,6 @@ def multiple_layer_velocity_model(self, z_switch, layers):
         # cond = fire.conditional(self.mesh_z > z_switch, layer1, layer2)
         self.set_initial_velocity_model(conditional=cond)
 
-
-# class Rectangle(AcousticWave):
-#     def __init__(self, model_dictionary=None, comm=None):
-#         model_parameters = Rectangle_parameters(
-#             dictionary=model_dictionary, comm=comm
-#         )
-#         super().__init__(
-#             model_parameters=model_parameters, comm=model_parameters.comm
-#         )
-#         comm = self.comm
-#         num_sources = self.number_of_sources
-#         if comm.comm.rank == 0 and comm.ensemble_comm.rank == 0:
-#             print(
-#                 "INFO: Distributing %d shot(s) across %d core(s). \
-#                     Each shot is using %d cores"
-#                 % (
-#                     num_sources,
-#                     fire.COMM_WORLD.size,
-#                     fire.COMM_WORLD.size / comm.ensemble_comm.size,
-#                 ),
-#                 flush=True,
-#             )
+if __name__ == "__main__":
+    wave = Rectangle_acoustic()
+    wave.forward_solve()

spyro/io/model_parameters.py

@@ -383,12 +383,6 @@ def _sanitize_absorving_boundary_condition(self):
         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_active:
-            self._correct_time_integrator_for_abc()
-
-    def _correct_time_integrator_for_abc(self):
-        if self.time_integrator == "central_difference":
-            self.time_integrator = "mixed_space_central_difference"
 
     def _sanitize_output(self):
         #         default_dictionary["visualization"] = {
@@ -608,9 +602,9 @@ def _sanitize_optimization_and_velocity(self):
             if "velocity_conditional" not in dictionary["synthetic_data"]:
                 self.velocity_model_type = None
                 warnings.warn(
-                    "No velocity model set initially. If using \
-                        user defined conditional or expression, please \
-                            input it in the Wave object."
+                    "No velocity model set initially. If using " \
+                        "user defined conditional or expression, please " \
+                            "input it in the Wave object."
                 )
 
         if "velocity_conditional" in dictionary["synthetic_data"]:

spyro/solvers/acoustic_solver_construction_no_pml.py

@@ -35,9 +35,14 @@ def construct_solver_or_matrix_no_pml(Wave_object):
     )
     a = dot(grad(u_n), grad(v)) * dx(scheme=quad_rule)  # explicit
 
+    le = 0
+    q = Wave_object.source_expression
+    if q is not None:
+        le = q * v * dx(scheme=quad_rule)
+
     B = fire.Function(V)
 
-    form = m1 + a
+    form = m1 + a - le
     lhs = fire.lhs(form)
     rhs = fire.rhs(form)
     Wave_object.lhs = lhs

spyro/solvers/acoustic_wave.py

@@ -14,7 +14,7 @@
 from .backward_time_integration import (
     backward_wave_propagator,
 )
-
+from ..utils.typing import override
 
 class AcousticWave(Wave):
     def save_current_velocity_model(self, file_name=None):
@@ -40,7 +40,7 @@ def forward_solve(self):
         if self.function_space is None:
             self.force_rebuild_function_space()
 
-        self._get_initial_velocity_model()
+        self._initialize_model_parameters()
         self.c = self.initial_velocity_model
         self.matrix_building()
         self.wave_propagator()
@@ -68,6 +68,7 @@ def matrix_building(self):
         self.trial_function = None
         self.u_nm1 = None
         self.u_n = None
+        self.u_np1 = fire.Function(self.function_space)
         self.lhs = None
         self.solver = None
         self.rhs = None
@@ -81,6 +82,7 @@ def matrix_building(self):
             self.X = None
             self.X_n = None
             self.X_nm1 = None
+            self.X_np1 = fire.Function(V * Z)
             construct_solver_or_matrix_with_pml(self)
 
     @ensemble_propagator
@@ -144,10 +146,94 @@ 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")
-
+
+    @override
+    def _initialize_model_parameters(self):
+        if self.initial_velocity_model is not None:
+            return None
+
+        if self.initial_velocity_model_file is None:
+            raise ValueError("No velocity model or velocity file to load.")
+
+        if self.initial_velocity_model_file.endswith(".segy"):
+            vp_filename, vp_filetype = os.path.splitext(
+                self.initial_velocity_model_file
+            )
+            warnings.warn("Converting segy file to hdf5")
+            write_velocity_model(
+                self.initial_velocity_model_file, ofname=vp_filename
+            )
+            self.initial_velocity_model_file = vp_filename + ".hdf5"
+
+        if self.initial_velocity_model_file.endswith((".hdf5", ".h5")):
+            self.initial_velocity_model = interpolate(
+                self,
+                self.initial_velocity_model_file,
+                self.function_space.sub(0),
+            )
+
+        if self.debug_output:
+            fire.File("initial_velocity_model.pvd").write(
+                self.initial_velocity_model, name="velocity"
+            )
+
+    @override
+    def _set_vstate(self, vstate):
+        if self.abc_boundary_layer_type == "PML":
+            self.X_n.assign(vstate)
+        else:
+            self.u_n.assign(vstate)
+
+    @override
+    def _get_vstate(self):
+        if self.abc_boundary_layer_type == "PML":
+            return self.X_n
+        else:
+            return self.u_n
+
+    @override
+    def _set_prev_vstate(self, vstate):
+        if self.abc_boundary_layer_type == "PML":
+            self.X_nm1.assign(vstate)
+        else:
+            self.u_nm1.assign(vstate)
+
+    @override
+    def _get_prev_vstate(self):
+        if self.abc_boundary_layer_type == "PML":
+            return self.X_nm1
+        else:
+            return self.u_nm1
+
+    @override
+    def _set_next_vstate(self, vstate):
+        if self.abc_boundary_layer_type == "PML":
+            self.X_np1.assign(vstate)
+        else:
+            self.u_np1.assign(vstate)
+
+    @override
+    def _get_next_vstate(self):
+        if self.abc_boundary_layer_type == "PML":
+            return self.X_np1
+        else:
+            return self.u_np1
+
+    @override
+    def get_receivers_output(self):
+        if self.abc_boundary_layer_type == "PML":
+            data_with_halos = self.X_n.dat.data_ro_with_halos[0][:]
+        else:
+            data_with_halos = self.u_n.dat.data_ro_with_halos[:]
+        return self.receivers.interpolate(data_with_halos)
+
+    @override
+    def get_function(self):
+        if self.abc_boundary_layer_type == "PML":
+            return self.X_n.sub(0)
+        else:
+            return self.u_n
+
+    @override
+    def get_function_name(self):
+        return "Pressure"

spyro/solvers/elastic_wave/elastic_wave.py

@@ -0,0 +1,29 @@
+from abc import abstractmethod
+
+from ..wave import Wave
+
+class ElasticWave(Wave):
+    '''Base class for elastic wave propagators'''
+    def __init__(self, dictionary, comm=None):
+        super().__init__(dictionary, comm=comm)
+
+    #@override
+    def _initialize_model_parameters(self):
+        d = self.input_dictionary.get("synthetic_data", False)
+        if bool(d) and "type" in d:
+            if d["type"] == "object":
+                self.initialize_model_parameters_from_object(d)
+            elif d["type"] == "file":
+                self.initialize_model_parameters_from_file(d)
+            else:
+                raise Exception(f"Invalid synthetic data type: {d['type']}")
+        else:
+            raise Exception("Input dictionary must contain ['synthetic_data']['type']")
+
+    @abstractmethod
+    def initialize_model_parameters_from_object(self, synthetic_data_dict):
+        pass
+
+    @abstractmethod
+    def initialize_model_parameters_from_file(self, synthetic_data_dict):
+        pass

spyro/solvers/elastic_wave/isotropic_wave.py

@@ -0,0 +1,48 @@
+from .elastic_wave import ElasticWave
+
+class IsotropicWave(ElasticWave):
+    '''Isotropic elastic wave propagator'''
+    def __init__(self, dictionary, comm=None):
+        super().__init__(dictionary, comm=comm)
+
+        self.rho = None   # Density
+        self.lmbda = None # First Lame parameter
+        self.mu = None    # Second Lame parameter
+        self.c_s = None   # Secondary wave velocity
+
+    #@override
+    def initialize_model_parameters_from_object(self, synthetic_data_dict: dict):
+        self.rho = synthetic_data_dict.get("density", None)
+        self.lmbda = synthetic_data_dict.get("lambda", 
+                                            synthetic_data_dict.get("lame_first", None))
+        self.mu = synthetic_data_dict.get("mu", 
+                                          synthetic_data_dict.get("lame_second", None))
+        self.c = synthetic_data_dict.get("p_wave_velocity", None)
+        self.c_s = synthetic_data_dict.get("s_wave_velocity", None)
+
+        # Check if {rho, lambda, mu} is set and {c, c_s} are not
+        option_1 = bool(self.rho) and \
+                   bool(self.lmbda) and \
+                   bool(self.mu) and \
+                   not bool(self.c) and \
+                   not bool(self.c_s)
+        # Check if {rho, c, c_s} is set and {lambda, mu} are not
+        option_2 = bool(self.rho) and \
+                   bool(self.c) and \
+                   bool(self.c_s) and \
+                   not bool(self.lmbda) and \
+                   not bool(self.mu)
+
+        if not option_1 and not option_2:
+            raise Exception(f"Inconsistent selection of isotropic elastic wave parameters:\n" \
+                            f"    Density        : {bool(self.rho)}\n"\
+                            f"    Lame first     : {bool(self.lmbda)}\n"\
+                            f"    Lame second    : {bool(self.mu)}\n"\
+                            f"    P-wave velocity: {bool(self.c)}\n"\
+                            f"    S-wave velocity: {bool(self.c_s)}\n"\
+                            "The valid options are \{Density, Lame first, Lame second\} "\
+                            "or \{Density, P-wave velocity, S-wave velocity\}")
+
+    #@override
+    def initialize_model_parameters_from_file(self, synthetic_data_dict):
+        raise NotImplementedError