{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Custom finite difference coefficients in Devito\n",
    "\n",
    "## Introduction\n",
    "\n",
    "When taking the numerical derivative of a function in Devito, the default behaviour is for 'standard' finite difference weights (obtained via a Taylor series expansion about the point of differentiation) to be applied. Consider the following example for some field $u(\\mathbf{x},t)$, where $\\mathbf{x}=(x,y)$. Let us define a computational domain/grid and differentiate our field with respect to $x$."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "import numpy as np\n",
    "from devito import Grid, TimeFunction\n",
    "\n",
    "# Create our grid (computational domain)\n",
    "Lx = 10\n",
    "Ly = Lx\n",
    "Nx = 11\n",
    "Ny = Nx\n",
    "dx = Lx/(Nx-1)\n",
    "dy = dx\n",
    "grid = Grid(shape=(Nx,Ny), extent=(Lx,Ly))\n",
    "\n",
    "# Define u(x,y,t) on this grid\n",
    "u = TimeFunction(name='u', grid=grid, time_order=2, space_order=2)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Now, lets look at the output of $\\partial u/\\partial x$:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "print(u.dx.evaluate)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "By default the 'standard' Taylor series expansion result, where `h_x` represents the $x$-direction grid spacing, is returned. However, there may be instances when a user wishes to use 'non-standard' weights when, for example, implementing a dispersion-relation-preserving (DRP) scheme. See e.g. \n",
    "\n",
    "[1] Christopher K.W. Tam, Jay C. Webb (1993). ”Dispersion-Relation-Preserving Finite Difference Schemes for Computational Acoustics.” **J. Comput. Phys.**, 107(2), 262--281. https://doi.org/10.1006/jcph.1993.1142\n",
    "\n",
    "for further details. The use of such modified weights is facilitated in Devito via the 'symbolic' finite difference coefficents functionality. Let us start by re-defining the function $u(\\mathbf{x},t)$ in the following manner:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "u = TimeFunction(name='u', grid=grid, time_order=2, space_order=2, coefficients='symbolic')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Note the addition of the `coefficients='symbolic'` keyword. Now, when printing $\\partial u/\\partial x$ we obtain:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "print(u.dx.evaluate)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Owing to the addition of the `coefficients='symbolic'` keyword the weights have been replaced by sympy functions. Now, take for example the weight `W(x - h_x, 1, u(t, x, y), x)`, the notation is as follows:\n",
    "* The first `x - h_x` refers to the spatial location of the weight w.r.t. the evaluation point `x`.\n",
    "* The `1` refers to the order of the derivative.\n",
    "* `u(t, x, y)` refers to the function with which the weight is associated.\n",
    "* Finally, the `x` refers to the dimension along which the derivative is being taken.\n",
    "\n",
    "Symbolic coefficients can then be manipulated using the Devito 'Coefficient' and 'Substitutions' objects. First, let us consider an example where we wish to replace the coefficients with a set of constants throughout the entire computational domain."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from devito import Coefficient, Substitutions # Import the Devito Coefficient and Substitutions objects\n",
    "# Grab the grid spatial dimensions: Note x[0] will correspond to the x-direction and x[1] to y-direction\n",
    "x = grid.dimensions \n",
    "# Form a Coefficient object and then a replacement rules object (to pass to a Devito equation):\n",
    "u_x_coeffs = Coefficient(1, u, x[0], np.array([-0.6, 0.1, 0.6]))\n",
    "coeffs = Substitutions(u_x_coeffs)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Devito Coefficient ojects take arguments in the following order:\n",
    "1. Derivative order (in the above example this is the first derivative)\n",
    "2. Function to which the coefficients 'belong' (in the above example this is the time function `u`)\n",
    "3. Dimension on which coefficients will be applied (in the above example this is the x-direction)\n",
    "4. Coefficient data. Since, in the above example, the coefficients have been applied as a 1-d numpy array replacement will occur at the equation level. (Note that other options are in development and will be the subject of future notebooks)."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Now, lets form a Devito equation, pass it the Substitutions object, and take a look at the output:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from devito import Eq\n",
    "eq = Eq(u.dt+u.dx, coefficients=coeffs)\n",
    "print(eq.evaluate)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We see that in the above equation the standard weights for the first derivative of `u` in the $x$-direction have now been replaced with our user defined weights. Note that since no replacement rules were defined for the time derivative (`u.dt`) standard weights have replaced the symbolic weights.\n",
    "\n",
    "Now, let us consider a more complete example."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Example: Finite difference modeling for a large velocity-contrast acousitc wave model\n",
    "\n",
    "It is advised to read through the 'Introduction to seismic modelling' notebook located in  devito/examples/seismic/tutorials/01_modelling.ipynb before proceeding with this example since much introductory material will be ommited here. The example now considered is based on an example introduced in\n",
    "\n",
    "[2] Yang Liu (2013). ”Globally optimal finite-difference schemes based on least squares.” **GEOPHYSICS**, 78(4), 113--132. https://doi.org/10.1190/geo2012-0480.1.\n",
    "\n",
    "See figure 18 of [2] for further details. Note that here we will simply use Devito to 'reproduce' the simulations leading to two results presented in the aforementioned figure. No analysis of the results will be carried out. The domain under consideration has a sptaial extent of $2km \\times 2km$ and, letting $x$ be the horizontal coordinate and $z$ the depth, a velocity profile such that $v_1(x,z)=1500ms^{-1}$ for $z\\leq1200m$ and $v_2(x,z)=4000ms^{-1}$ for $z>1200m$.\n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "tags": [
     "nbval-ignore-output"
    ]
   },
   "outputs": [],
   "source": [
    "from examples.seismic import Model, plot_velocity\n",
    "%matplotlib inline\n",
    "\n",
    "# Define a physical size\n",
    "Lx = 2000\n",
    "Lz = Lx\n",
    "h = 10\n",
    "Nx = int(Lx/h)+1\n",
    "Nz = Nx\n",
    "\n",
    "shape = (Nx, Nz)  # Number of grid point\n",
    "spacing = (h, h)  # Grid spacing in m. The domain size is now 2km by 2km\n",
    "origin = (0., 0.)\n",
    "\n",
    "# Define a velocity profile. The velocity is in km/s\n",
    "v = np.empty(shape, dtype=np.float32)\n",
    "v[:, :121] = 1.5\n",
    "v[:, 121:] = 4.0\n",
    "\n",
    "# With the velocity and model size defined, we can create the seismic model that\n",
    "# encapsulates these properties. We also define the size of the absorbing layer as 10 grid points\n",
    "nbl = 10\n",
    "model = Model(vp=v, origin=origin, shape=shape, spacing=spacing,\n",
    "              space_order=20, nbl=nbl)\n",
    "\n",
    "plot_velocity(model)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The seismic wave source term will be modelled as a Ricker Wavelet with a peak-frequency of $25$Hz located at $(1000m,800m)$. Before applying the DRP scheme, we begin by generating a 'reference' solution using a spatially high-order standard finite difference scheme and time step well below the model's critical time-step. The scheme will be 2nd order in time."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from examples.seismic import TimeAxis\n",
    "\n",
    "t0 = 0.  # Simulation starts a t=0\n",
    "tn = 500.  # Simulation lasts 0.5 seconds (500 ms)\n",
    "dt = 0.2  # Time step of 0.2ms\n",
    "\n",
    "time_range = TimeAxis(start=t0, stop=tn, step=dt)\n",
    "\n",
    "#NBVAL_IGNORE_OUTPUT\n",
    "from examples.seismic import RickerSource\n",
    "\n",
    "f0 = 0.015  # Source peak frequency is 25Hz (0.025 kHz)\n",
    "src = RickerSource(name='src', grid=model.grid, f0=f0,\n",
    "                   npoint=1, time_range=time_range)\n",
    "\n",
    "# First, position source centrally in all dimensions, then set depth\n",
    "src.coordinates.data[0, :] = np.array(model.domain_size) * .5\n",
    "src.coordinates.data[0, -1] = 800.  # Depth is 800m\n",
    "\n",
    "# We can plot the time signature to see the wavelet\n",
    "src.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Now let us define our wavefield and PDE:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Define the wavefield with the size of the model and the time dimension\n",
    "u = TimeFunction(name=\"u\", grid=model.grid, time_order=2, space_order=20)\n",
    "\n",
    "# We can now write the PDE\n",
    "pde = model.m * u.dt2 - u.laplace + model.damp * u.dt\n",
    "\n",
    "# This discrete PDE can be solved in a time-marching way updating u(t+dt) from the previous time step\n",
    "# Devito as a shortcut for u(t+dt) which is u.forward. We can then rewrite the PDE as \n",
    "# a time marching updating equation known as a stencil using customized SymPy functions\n",
    "from devito import solve\n",
    "\n",
    "stencil = Eq(u.forward, solve(pde, u.forward))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Finally we define the source injection and receiver read function to generate the corresponding code\n",
    "src_term = src.inject(field=u.forward, expr=src * dt**2 / model.m)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Now, lets create the operator and execute the time marching scheme:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from devito import Operator\n",
    "\n",
    "op = Operator([stencil] + src_term, subs=model.spacing_map)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#NBVAL_IGNORE_OUTPUT\n",
    "op(time=time_range.num-1, dt=dt)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "And plot the result:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#import matplotlib\n",
    "import matplotlib.pyplot as plt\n",
    "from matplotlib import cm\n",
    "\n",
    "Lx = 2000\n",
    "Lz = 2000\n",
    "\n",
    "abs_lay = nbl*h\n",
    "\n",
    "dx = h\n",
    "dz = dx\n",
    "X, Z = np.mgrid[-abs_lay: Lx+abs_lay+1e-10: dx, -abs_lay: Lz+abs_lay+1e-10: dz]\n",
    "\n",
    "levels = 100\n",
    "\n",
    "fig = plt.figure(figsize=(14, 7))\n",
    "ax1 = fig.add_subplot(111)\n",
    "cont = ax1.contourf(X,Z,u.data[0,:,:], levels, cmap=cm.binary)\n",
    "fig.colorbar(cont)\n",
    "ax1.axis([0, Lx, 0, Lz])\n",
    "ax1.set_xlabel('$x$')\n",
    "ax1.set_ylabel('$z$')\n",
    "ax1.set_title('$u(x,z,500)$')\n",
    "\n",
    "plt.gca().invert_yaxis()\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We will now reimplement the above model applying the DRP scheme presented in [2].\n",
    "\n",
    "First, since we wish to apply different custom FD coefficients in the upper on lower layers we need to define these two 'subdomains' using the `Devito SubDomain` functionality:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from devito import SubDomain\n",
    "\n",
    "# Define our 'upper' and 'lower' SubDomains:\n",
    "class Upper(SubDomain):\n",
    "    name = 'upper'\n",
    "    def define(self, dimensions):\n",
    "        x, z = dimensions\n",
    "        # We want our upper layer to span the entire x-dimension and all\n",
    "        # but the bottom 80 (+boundary layer) cells in the z-direction, which is achieved via\n",
    "        # the following notation:\n",
    "        return {x: x, z: ('left', 80+nbl)}\n",
    "    \n",
    "class Lower(SubDomain):\n",
    "    name = 'lower'\n",
    "    def define(self, dimensions):\n",
    "        x, z = dimensions\n",
    "        # We want our lower layer to span the entire x-dimension and all\n",
    "        # but the top 121 (+boundary layer) cells in the z-direction.\n",
    "        return {x: x, z: ('right', 121+nbl)}\n",
    "\n",
    "# Create these subdomains:\n",
    "ur = Upper()\n",
    "lr = Lower()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We now create our model incoporating these subdomains:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "tags": [
     "nbval-ignore-output"
    ]
   },
   "outputs": [],
   "source": [
    "# Our scheme will now be 10th order (or less) in space.\n",
    "order = 10\n",
    "\n",
    "# Create our model passing it our 'upper' and 'lower' subdomains: \n",
    "model = Model(vp=v, origin=origin, shape=shape, spacing=spacing,\n",
    "              space_order=order, nbl=nbl, subdomains=(ur,lr))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "And re-define model related objects. Note that now our wave-field will be defined with `coefficients='symbolic'`."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "t0 = 0.  # Simulation starts a t=0\n",
    "tn = 500.  # Simulation last 1 second (500 ms)\n",
    "dt = 1.0  # Time step of 1.0ms\n",
    "\n",
    "time_range = TimeAxis(start=t0, stop=tn, step=dt)\n",
    "\n",
    "f0 = 0.025  # Source peak frequency is 25Hz (0.025 kHz)\n",
    "src = RickerSource(name='src', grid=model.grid, f0=f0,\n",
    "                   npoint=1, time_range=time_range)\n",
    "\n",
    "src.coordinates.data[0, :] = np.array(model.domain_size) * .5\n",
    "src.coordinates.data[0, -1] = 800.  # Depth is 800m\n",
    "\n",
    "# New wave-field\n",
    "u_DRP = TimeFunction(name=\"u_DRP\", grid=model.grid, time_order=2, space_order=order, coefficients='symbolic')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We now create a stencil for each of our 'Upper' and 'Lower' subdomains defining different custom FD weights within each of these subdomains."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# The underlying pde is the same in both subdomains\n",
    "pde_DRP = model.m * u_DRP.dt2 - u_DRP.laplace + model.damp * u_DRP.dt\n",
    "\n",
    "# Define our custom FD coefficients:\n",
    "x, z = model.grid.dimensions\n",
    "# Upper layer\n",
    "weights_u = np.array([ 2.00462e-03, -1.63274e-02,  7.72781e-02, \n",
    "                      -3.15476e-01,  1.77768e+00, -3.05033e+00,  \n",
    "                       1.77768e+00, -3.15476e-01,  7.72781e-02, \n",
    "                      -1.63274e-02,  2.00462e-03])\n",
    "# Lower layer\n",
    "weights_l = np.array([  0.      ,  0.      ,  0.0274017, \n",
    "                       -0.223818,  1.64875 , -2.90467,  \n",
    "                        1.64875 , -0.223818,  0.0274017,  \n",
    "                        0.      ,  0.       ])\n",
    "# Create the Devito Coefficient objects:\n",
    "ux_u_coeffs = Coefficient(2, u_DRP, x, weights_u/x.spacing**2)\n",
    "uz_u_coeffs = Coefficient(2, u_DRP, z, weights_u/z.spacing**2)\n",
    "ux_l_coeffs = Coefficient(2, u_DRP, x, weights_l/x.spacing**2)\n",
    "uz_l_coeffs = Coefficient(2, u_DRP, z, weights_l/z.spacing**2)\n",
    "# And the replacement rules:\n",
    "coeffs_u = Substitutions(ux_u_coeffs,uz_u_coeffs)\n",
    "coeffs_l = Substitutions(ux_l_coeffs,uz_l_coeffs)\n",
    "# Create a stencil for each subdomain:\n",
    "stencil_u = Eq(u_DRP.forward, solve(pde_DRP, u_DRP.forward), subdomain = model.grid.subdomains['upper'],\n",
    "               coefficients=coeffs_u)\n",
    "stencil_l = Eq(u_DRP.forward, solve(pde_DRP, u_DRP.forward), subdomain = model.grid.subdomains['lower'],\n",
    "               coefficients=coeffs_l)\n",
    "\n",
    "# Source term:\n",
    "src_term = src.inject(field=u_DRP.forward, expr=src * dt**2 / model.m)\n",
    "\n",
    "# Create the operator, incoporating both upper and lower stencils:\n",
    "op = Operator([stencil_u, stencil_l] + src_term, subs=model.spacing_map)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "And now execute the operator:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#NBVAL_IGNORE_OUTPUT\n",
    "op(time=time_range.num-1, dt=dt)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "And plot the new results:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "fig = plt.figure(figsize=(14, 7))\n",
    "ax1 = fig.add_subplot(111)\n",
    "cont = ax1.contourf(X,Z,u_DRP.data[0,:,:], levels, cmap=cm.binary)\n",
    "fig.colorbar(cont)\n",
    "ax1.axis([0, Lx, 0, Lz])\n",
    "ax1.set_xlabel('$x$')\n",
    "ax1.set_ylabel('$z$')\n",
    "ax1.set_title('$u_{DRP}(x,z,500)$')\n",
    "\n",
    "plt.gca().invert_yaxis()\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Finally, for comparison, lets plot the difference between the standard 20th order and optimized 10th order models:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "fig = plt.figure(figsize=(14, 7))\n",
    "ax1 = fig.add_subplot(111)\n",
    "cont = ax1.contourf(X,Z,abs(u_DRP.data[0,:,:]-u.data[0,:,:]), levels, cmap=cm.binary)\n",
    "fig.colorbar(cont)\n",
    "ax1.axis([0, Lx, 0, Lz])\n",
    "ax1.set_xlabel('$x$')\n",
    "ax1.set_ylabel('$z$')\n",
    "\n",
    "plt.gca().invert_yaxis()\n",
    "plt.show()"
   ]
  }
 ],
 "metadata": {
  "hide_input": false,
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.7.4"
  },
  "latex_envs": {
   "LaTeX_envs_menu_present": true,
   "autoclose": false,
   "autocomplete": true,
   "bibliofile": "biblio.bib",
   "cite_by": "apalike",
   "current_citInitial": 1,
   "eqLabelWithNumbers": true,
   "eqNumInitial": 1,
   "hotkeys": {
    "equation": "Ctrl-E",
    "itemize": "Ctrl-I"
   },
   "labels_anchors": false,
   "latex_user_defs": false,
   "report_style_numbering": false,
   "user_envs_cfg": false
  },
  "widgets": {
   "state": {},
   "version": "1.1.2"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 1
}