Merge pull request #35 from HPCSys-Lab/variable-density-with-higher-space-orders

variable density with multiple space order

Author: jaimesouza

examples/acoustic_2D_density.py

@@ -0,0 +1,87 @@
+from simwave import (
+    SpaceModel, TimeModel, RickerWavelet, Solver, Compiler,
+    Receiver, Source, plot_wavefield, plot_shotrecord, plot_velocity_model
+)
+import numpy as np
+
+
+# set compiler options
+# available language options: c (sequential) or  cpu_openmp (parallel CPU)
+compiler = Compiler(
+    cc='gcc',
+    language='cpu_openmp',
+    cflags='-O3 -fPIC -ffast-math -Wall -std=c99 -shared'
+)
+
+# Velocity model
+vel = np.zeros(shape=(512, 512), dtype=np.float32)
+vel[:] = 1500.0
+vel[100:] = 2000.0
+
+# Density model
+density = np.zeros(shape=(512, 512), dtype=np.float32)
+density[:] = 10
+
+# create the space model
+space_model = SpaceModel(
+    bounding_box=(0, 5120, 0, 5120),
+    grid_spacing=(10, 10),
+    velocity_model=vel,
+    density_model=density,
+    space_order=2,
+    dtype=np.float32
+)
+
+# config boundary conditions
+# (none,  null_dirichlet or null_neumann)
+space_model.config_boundary(
+    damping_length=0,
+    boundary_condition=(
+        "null_neumann", "null_dirichlet",
+        "none", "null_dirichlet"
+    ),
+    damping_polynomial_degree=3,
+    damping_alpha=0.001
+)
+
+# create the time model
+time_model = TimeModel(
+    space_model=space_model,
+    tf=1.0
+)
+
+# create the set of sources
+source = Source(
+    space_model,
+    coordinates=[(2560, 2560)],
+    window_radius=1
+)
+
+# crete the set of receivers
+receiver = Receiver(
+    space_model=space_model,
+    coordinates=[(2560, i) for i in range(0, 5120, 10)],
+    window_radius=1
+)
+
+# create a ricker wavelet with 10hz of peak frequency
+ricker = RickerWavelet(10.0, time_model)
+
+# create the solver
+solver = Solver(
+    space_model=space_model,
+    time_model=time_model,
+    sources=source,
+    receivers=receiver,
+    wavelet=ricker,
+    saving_stride=0,
+    compiler=compiler
+)
+
+# run the forward
+u_full, recv = solver.forward()
+
+print("u_full shape:", u_full.shape)
+plot_velocity_model(space_model.velocity_model)
+plot_wavefield(u_full[-1])
+plot_shotrecord(recv)

simwave/kernel/backend/c_code/forward/variable_density/2d/wave.c

@@ -15,7 +15,7 @@
 // forward_2D_variable_density
 double forward(f_type *grid, f_type *velocity, f_type *density, f_type *damp,
                f_type *wavelet, size_t wavelet_size,
-               f_type *coeff, size_t *boundary_conditions,
+               f_type *coeff_order2, f_type *coeff_order1, size_t *boundary_conditions,
                size_t *src_points_interval, size_t src_points_interval_size,
                f_type *src_points_values, size_t src_points_values_size,
                size_t *rec_points_interval, size_t rec_points_interval_size,
@@ -96,20 +96,45 @@ double forward(f_type *grid, f_type *velocity, f_type *density, f_type *damp,
                 // index of the current point in the grid
                 size_t current = i * nx + j;
 
-                //neighbors in the horizontal direction
-                f_type x1 = ((prev_snapshot[current + 1] - prev_snapshot[current]) * (density[current + 1] + density[current])) / density[current + 1];
-                f_type x2 = ((prev_snapshot[current] - prev_snapshot[current - 1]) * (density[current] + density[current - 1])) / density[current - 1];
-                f_type term_x = (x1 - x2) / (2 * dxSquared);
+                // stencil code to update grid
+                f_type value = 0.0;
 
-                //neighbors in the vertical direction
-                f_type z1 = ((prev_snapshot[current + nx] - prev_snapshot[current]) * (density[current + nx] + density[current])) / density[current + nx];
-                f_type z2 = ((prev_snapshot[current] - prev_snapshot[current - nx]) * (density[current] + density[current - nx])) / density[current - nx];
-                f_type term_z = (z1 - z2) / (2 * dzSquared);
+                // second derivative for pressure
+                f_type sd_pressure_x = coeff_order2[0] * prev_snapshot[current];
+                f_type sd_pressure_z = coeff_order2[0] * prev_snapshot[current];
 
-                //nominator with damp coefficient
+                // first derivative for pressure
+                f_type fd_pressure_x = 0.0;
+                f_type fd_pressure_z = 0.0;
+
+                // first derivative for density
+                f_type fd_density_x = 0.0;
+                f_type fd_density_z = 0.0;
+
+                // radius of the stencil
+                for(int ir = 1; ir <= stencil_radius; ir++){
+                    //neighbors in the horizontal direction
+                    sd_pressure_x += coeff_order2[ir] * (prev_snapshot[current + ir] + prev_snapshot[current - ir]);
+                    fd_pressure_x += coeff_order1[ir] * (prev_snapshot[current + ir] - prev_snapshot[current - ir]);
+                    fd_density_x += coeff_order1[ir] * (density[current + ir] - density[current - ir]);
+
+                    //neighbors in the vertical direction
+                    sd_pressure_z += coeff_order2[ir] * (prev_snapshot[current + (ir * nx)] + prev_snapshot[current - (ir * nx)]);
+                    fd_pressure_z += coeff_order1[ir] * (prev_snapshot[current + (ir * nx)] - prev_snapshot[current - (ir * nx)]);
+                    fd_density_z += coeff_order1[ir] * (density[current + (ir * nx)] - density[current - (ir * nx)]);
+                }
+
+                value += sd_pressure_x/dxSquared + sd_pressure_z/dzSquared;
+
+                f_type term_x = (fd_pressure_x * fd_density_x) / (2 * dxSquared);
+                f_type term_z = (fd_pressure_z * fd_density_z) / (2 * dzSquared);
+
+                value -= (term_x + term_z) / density[current];
+
+                //denominator with damp coefficient
                 f_type denominator = (1.0 + damp[current] * dt);
 
-                f_type value = dtSquared * velocity[current] * velocity[current] * (term_z + term_x) / denominator;
+                value *= (dtSquared * velocity[current] * velocity[current]) / denominator;
 
                 next_snapshot[current] = 2.0 / denominator * prev_snapshot[current] - ((1.0 - damp[current] * dt) / denominator) * next_snapshot[current] + value;
             }

simwave/kernel/backend/c_code/forward/variable_density/3d/wave.c

@@ -15,7 +15,7 @@
 // forward_3D_variable_density
 double forward(f_type *grid, f_type *velocity, f_type *density, f_type *damp,
                f_type *wavelet, size_t wavelet_size,
-               f_type *coeff, size_t *boundary_conditions,
+               f_type *coeff_order2, f_type *coeff_order1, size_t *boundary_conditions,
                size_t *src_points_interval, size_t src_points_interval_size,
                f_type *src_points_values, size_t src_points_values_size,
                size_t *rec_points_interval, size_t rec_points_interval_size,
@@ -101,28 +101,56 @@ double forward(f_type *grid, f_type *velocity, f_type *density, f_type *damp,
                     // index of the current point in the grid
                     size_t current = (i * nx + j) * ny + k;
 
-                    //neighbors in the Y direction
-                    f_type y1 = ((prev_snapshot[current + 1] - prev_snapshot[current]) * (density[current + 1] + density[current])) / density[current + 1];
-                    f_type y2 = ((prev_snapshot[current] - prev_snapshot[current - 1]) * (density[current] + density[current - 1])) / density[current - 1];
-                    f_type term_y = (y1 - y2) / (2 * dySquared);
+                    // stencil code to update grid
+                    f_type value = 0.0;
 
-                    //neighbors in the X direction
-                    f_type x1 = ((prev_snapshot[current + ny] - prev_snapshot[current]) * (density[current + ny] + density[current])) / density[current + ny];
-                    f_type x2 = ((prev_snapshot[current] - prev_snapshot[current - ny]) * (density[current] + density[current - ny])) / density[current - ny];
-                    f_type term_x = (x1 - x2) / (2 * dxSquared);
+                    // second derivative for pressure
+                    f_type sd_pressure_y = coeff_order2[0] * prev_snapshot[current];
+                    f_type sd_pressure_x = coeff_order2[0] * prev_snapshot[current];
+                    f_type sd_pressure_z = coeff_order2[0] * prev_snapshot[current];
 
-                    //neighbors in the Z direction
-                    f_type z1 = ((prev_snapshot[current + (nx * ny)] - prev_snapshot[current]) * (density[current + (nx * ny)] + density[current])) / density[current + (nx * ny)];
-                    f_type z2 = ((prev_snapshot[current] - prev_snapshot[current - (nx * ny)]) * (density[current] + density[current - (nx * ny)])) / density[current - (nx * ny)];
-                    f_type term_z = (z1 - z2) / (2 * dzSquared);
+                    // first derivative for pressure
+                    f_type fd_pressure_y = 0.0;
+                    f_type fd_pressure_x = 0.0;
+                    f_type fd_pressure_z = 0.0;
 
-                    //nominator with damp coefficient
+                    // first derivative for density
+                    f_type fd_density_y = 0.0;
+                    f_type fd_density_x = 0.0;
+                    f_type fd_density_z = 0.0;
+
+                    // radius of the stencil
+                    for(int ir = 1; ir <= stencil_radius; ir++){
+                        //neighbors in the Y direction
+                        sd_pressure_y += coeff_order2[ir] * (prev_snapshot[current + ir] + prev_snapshot[current - ir]);
+                        fd_pressure_y += coeff_order1[ir] * (prev_snapshot[current + ir] - prev_snapshot[current - ir]);
+                        fd_density_y += coeff_order1[ir] * (density[current + ir] - density[current - ir]);
+
+                        //neighbors in the X direction
+                        sd_pressure_x += coeff_order2[ir] * (prev_snapshot[current + (ir * ny)] + prev_snapshot[current - (ir * ny)]);
+                        fd_pressure_x += coeff_order1[ir] * (prev_snapshot[current + (ir * nx)] - prev_snapshot[current - (ir * nx)]);
+                        fd_density_x += coeff_order1[ir] * (density[current + (ir * nx)] - density[current - (ir * nx)]);
+
+                        //neighbors in the Z direction
+                        sd_pressure_z += coeff_order2[ir] * (prev_snapshot[current + (ir * nx * ny)] + prev_snapshot[current - (ir * nx * ny)]);
+                        fd_pressure_z += coeff_order1[ir] * (prev_snapshot[current + (ir * nx * ny)] - prev_snapshot[current - (ir * nx * ny)]);
+                        fd_density_z += coeff_order1[ir] * (density[current + (ir * nx * ny)] - density[current - (ir * nx * ny)]);
+                    }
+
+                    value += sd_pressure_y/dySquared + sd_pressure_x/dxSquared + sd_pressure_z/dzSquared;
+
+                    f_type term_y = (fd_pressure_y * fd_density_y) / (2 * dySquared);
+                    f_type term_x = (fd_pressure_x * fd_density_x) / (2 * dxSquared);
+                    f_type term_z = (fd_pressure_z * fd_density_z) / (2 * dzSquared);
+
+                    value -= (term_y + term_x + term_z) / density[current];
+
+                    //denominator with damp coefficient
                     f_type denominator = (1.0 + damp[current] * dt);
 
-                    f_type value = dtSquared * velocity[current] * velocity[current] * (term_z + term_x + term_y) / denominator;
+                    value *= (dtSquared * velocity[current] * velocity[current]) / denominator;
 
                     next_snapshot[current] = 2.0 / denominator * prev_snapshot[current] - ((1.0 - damp[current] * dt) / denominator) * next_snapshot[current] + value;
-
                 }
             }
         }

simwave/kernel/backend/middleware.py

@@ -74,8 +74,10 @@ def exec(self, operator, **kwargs):
         kwargs.update({'boundary_condition': bc})
 
         # if density model is None, remove it
+        # and remove first_order_fd_coefficients
         if kwargs.get('density_model') is None:
             kwargs.pop('density_model')
+            kwargs.pop('first_order_fd_coefficients')
 
         # get the grid shape
         grid_shape = kwargs.get('velocity_model').shape
@@ -168,7 +170,8 @@ def _keys_in_order(self):
             'damping_mask',
             'wavelet',
             'wavelet_size',
-            'fd_coefficients',
+            'second_order_fd_coefficients',
+            'first_order_fd_coefficients',
             'boundary_condition',
             'src_points_interval',
             'src_points_interval_size',

simwave/kernel/frontend/fd.py

@@ -2,10 +2,10 @@
 import findiff
 
 
-def half_coefficients(space_order):
+def half_coefficients(derivative_order, space_order):
     """
     Return a list of right side finite differences coefficients
-    according to the space order.
+    according to the space order and derivative order.
 
     Parameters
     ----------
@@ -18,7 +18,7 @@ def half_coefficients(space_order):
         List of FD coefficients.
     """
     # all coefficients
-    coeffs = coefficients(space_order)
+    coeffs = coefficients(derivative_order, space_order)
 
     middle = len(coeffs) // 2
 
@@ -28,13 +28,15 @@ def half_coefficients(space_order):
     return coeffs
 
 
-def coefficients(space_order):
+def coefficients(derivative_order, space_order):
     """
     Return a list of finite differences coefficients according
-    to the space order.
+    to the space order and derivative order.
 
     Parameters
     ----------
+    derivative_order : int
+        Derivative order
     space_order : int
         Spatial order.
 
@@ -45,7 +47,7 @@ def coefficients(space_order):
     """
 
     # fixed second derivative
-    coeffs = findiff.coefficients(deriv=2, acc=space_order)
+    coeffs = findiff.coefficients(deriv=derivative_order, acc=space_order)
 
     return coeffs['center']['coefficients']
 
@@ -77,7 +79,10 @@ def calculate_dt(dimension, space_order, grid_spacing, velocity_model):
     a1 = 4
 
     # FD coeffs to the specific space order
-    fd_coeffs = coefficients(space_order)
+    fd_coeffs = coefficients(
+        derivative_order=2,
+        space_order=space_order
+    )
 
     a2 = dimension * np.sum(np.abs(fd_coeffs))
 

simwave/kernel/frontend/model.py

@@ -38,10 +38,16 @@ def __init__(self, bounding_box, grid_spacing, velocity_model,
 
         self._space_order = space_order
 
-        # space_order are limited to even number and limited up to 20
-        if space_order % 2 != 0 or space_order < 2 or space_order > 20:
+        # space_order are limited to even number
+        if space_order % 2 != 0:
             raise ValueError(
-                "Space order {} not supported".format(space_order)
+                "Odd space order {} not supported".format(space_order)
+            )
+
+        # space_order are limited from 2 to 20
+        if not 2 <= space_order <= 20:
+            raise ValueError(
+                "Space order limited from 2 to 20."
             )
 
         # get the space dimension according to the velocity model
@@ -205,10 +211,23 @@ def damping_alpha(self):
         except AttributeError:
             return 0.001
 
-    @property
-    def fd_coefficients(self):
-        """Central and right side finite differences coefficients."""
-        return self.dtype(fd.half_coefficients(self.space_order))
+    def fd_coefficients(self, derivative_order):
+        """
+        Central and right side finite differences coefficients.
+
+        Parameters
+        ----------
+        derivative_order : int
+            Derivative order.
+
+        Returns
+        ----------
+        ndarray
+            Central and right side FD coefficients.
+        """
+        return self.dtype(
+            fd.half_coefficients(derivative_order, self.space_order)
+        )
 
     def interpolate(self, data):
         """

simwave/kernel/frontend/solver.py

@@ -147,7 +147,8 @@ def forward(self):
             damping_mask=self.space_model.damping_mask,
             wavelet=self.wavelet.values,
             wavelet_size=len(self.wavelet.values),
-            fd_coefficients=self.space_model.fd_coefficients,
+            second_order_fd_coefficients=self.space_model.fd_coefficients(2),
+            first_order_fd_coefficients=self.space_model.fd_coefficients(1),
             boundary_condition=self.space_model.boundary_condition,
             src_points_interval=src_points,
             src_points_interval_size=len(src_points),

tests/test_space_model.py

@@ -153,7 +153,7 @@ def test_fd_coefficients(self, time_order, space_order, coeffs):
         )
 
         assert np.allclose(
-                   space_model.fd_coefficients,
+                   space_model.fd_coefficients(2),
                    np.asarray(coeffs, dtype=space_model.dtype)
                )