elastic_wave_2d_staggered_isotropic_v8_DX4_DT2_with_PML.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Codigo para la propagación de onda elástica 2D PSV en un medio homogeneo mediante el esquema velocidad - esfuerzo y rejilla escalonada (usado en tesis)\n",
    "En este Notebook se presenta un codigo para la propagación de onda elástica 2D PSV en un medio homogeneo mediante el esquema velocidad - esfuerzo y rejilla escalonada propuesto por Virieux (1986), como resultado se muestra en la figura 1 cuatro cuadros donde se muestra el avance de la onda P mediante la divergencia, la onda S mediante el rotor y las velocidades $v_x$ y $v_z$\n",
    "\n",
    "#### Sistema de ecuaciones de onda elástica 2D\n",
    "El sistema de ecuaciones para la propagación de la onda elástica en 2D en el esquema velocidad - esfuerzo donde el movimiento de la particula es en la dirección x (horizontal) y en la dirección z (vertical) son:\n",
    "\n",
    "$$\n",
    "\\frac{\\partial v_x}{\\partial t} = \\frac{1}{\\rho} \\left( \\frac{\\partial \\tau_{xx}}{\\partial x} + \\frac{\\partial \\tau_{xz}}{\\partial z} \\right)\n",
    "$$\n",
    "$$\n",
    "\\frac{\\partial v_z}{\\partial t} = \\frac{1}{\\rho} \\left( \\frac{\\partial \\tau_{xz}}{\\partial x} + \\frac{\\partial \\tau_{zz}}{\\partial z} \\right)\n",
    "$$\n",
    "$$\n",
    "\\frac{\\partial \\tau_{xx}}{\\partial t} = (\\lambda + 2\\mu)\\frac{\\partial v_x}{\\partial x} + \\lambda \\frac{\\partial v_z}{\\partial z}\n",
    "$$\n",
    "$$\n",
    "\\frac{\\partial \\tau_{zz}}{\\partial t} = (\\lambda + 2\\mu)\\frac{\\partial v_z}{\\partial z} + \\lambda \\frac{\\partial v_x}{\\partial x}\n",
    "$$\n",
    "$$\n",
    "\\frac{\\partial \\tau_{xz}}{\\partial t} = \\mu \\left( \\frac{\\partial v_x}{\\partial z} + \\frac{\\partial v_z}{\\partial x} \\right)\n",
    "$$\n",
    "\n",
    "donde $ \\rm {\\rho} $ es la densidad, $ \\rm {\\lambda}$ y $ \\rm {\\mu}$ son los parámetros de Lamé, $ \\rm{(v_x,\\; v_z)}$ es el vector de velocidad de partículas , $ \\rm {\\tau_{xx}} $, $ \\rm{\\tau_{zz}} $, $ \\rm{\\tau_{xz}}$ son los componentes del tensor esfuerzo\n",
    "\n",
    "#### Implementación númerica mediante el metodo de rejilla escalonada\n",
    "Para la solución numérica de las ecuaciones de movimiento elásticas hay que discretizarlas en el tiempo y el espacio en una cuadrícula. Las velocidades de las partículas $\\rm{\\mathbf{v}}$, los esfuerzos $\\rm{\\tau_{ij}}$, los parámetros de Lamé $\\rm{\\lambda}$ y $\\rm{\\mu}$ son calculado y definido en coordenadas cartesianas discretas $\\rm{x = i \\; dx}$, $\\rm{z = j \\; dz}$ y tiempos discretos $\\rm{t = n \\; dt}$. $\\rm{dx}$ y $\\rm{dz}$ denota la distancia espacial entre dos puntos de cuadrícula adyacentes en el eje x y eje z respectivamente y $\\rm{dt}$ es la diferencia entre dos pasos de tiempo sucesivos. Por lo tanto, cada punto de la cuadrícula se encuentra en el intervalo $\\rm{i\\ en N | [1, nx]}$, $\\rm {j\\ en N | [1, nz]}$ y $\\rm{n\\ en N | [1, nt]}$, donde $\\rm{nx}$, $\\rm{nz}$ y $\\rm{nt}$ son el número de puntos discretos de la cuadrícula espacial y el número de pasos de tiempo, respectivamente.\n",
    "\n",
    "Finalmente, las derivadas parciales se reemplazan por operadores **de diferencia finita (FD)**. Se pueden distinguir dos tipos de operadores, operadores hacia adelante y hacia atrás $ \\rm {D^+, \\; D^-}$. La derivada de una función f(x) con respecto a una variable x puede ser aproximada por los siguientes operadores de segundo orden:\n",
    "\n",
    "$$\\rm{D^+_x f = \\frac{f_{i+1}-f_{i}}{dx} \\hspace{1 cm} \\text{operador hacia adelante}} $$\n",
    "\n",
    "$$\\rm{D^-_x f = \\frac{f_i-f_{i-1}}{dx} \\hspace{1 cm} \\text{operador hacia atras}} $$\n",
    "\n",
    "Para calcular las derivadas espaciales de las variables de campo de onda en las posiciones correctas entre sí, las variables no se colocan en los mismos puntos de la cuadrícula, sino que se escalonan por la mitad de la distancia del punto de la cuadrícula espacial (Virieux 1986). La Figura 1 muestra la distribución de los parámetros del material y las variables del campo de onda en la cuadrícula espacial.\n",
    "\n",
    "![Staggered grid](./Staggered_grid.png)\n",
    "\n",
    "En el siguiente paso, discretizamos las ecuaciones de movimiento para el problema 2D PSV utilizando un enfoque de diferencia finita escalonada. Primero, discretizamos la componente x de la ecuación de momento aproximando las derivadas espaciales\n",
    "\n",
    "$$\\rm{\\frac{\\partial \\tau_{xx}}{\\partial x} \\approx \\frac{\\tau_{xx}(j,i+1) - \\tau_{xx}(j,i)}{dx}}, \\rm{\\frac{\\partial \\tau_{xz}}{\\partial z} \\approx \\frac{\\tau_{xz}(j+1/2, i) - \\tau_{xz}(j-1/2,i)}{dz}} $$\n",
    "\n",
    "Y la derivada temporal viene dada por\n",
    "\n",
    "$$\\rho \\rm{\\frac{\\partial v_x}{\\partial t} \\approx \\rho_x(j,i+1/2) \\frac{v_x^{n+1/2}(j,i+1/2) - v_x^{n-1/2}(j,i+1/2)}{dt}} $$\n",
    "\n",
    "Insertandola en la ecuación diferencial parcial\n",
    "\n",
    "$$\\rm{\\rho \\frac{\\partial v_x}{\\partial t} = \\frac{\\partial \\tau_{xx}}{\\partial x} + \\frac{\\partial \\tau_{xz}}{\\partial z}} $$\n",
    "\n",
    "lleva a\n",
    "\n",
    "$$\\rho_x(j,i+1/2) \\frac{v_x^{n+1/2}(j,i+1/2) - v_x^{n-1/2}(j,i+1/2)}{dt} = \\frac{\\tau_{xx}^n(j,i+1) - \\tau_{xx}^n(j,i)}{dx} + \\frac{\\tau_{xz}^n(j+1/2, i) - \\tau_{xz}^n(j-1/2,i)}{dz} $$\n",
    "\n",
    "Despues  reordenando  para $v_x^{n+1/2}(j,i+1/2)$ se obtiene el siguiente esquema de FD explicito para el componente  x del la ecuación de  momentum:\n",
    "\n",
    "$$\\rm{v_x^{n+1/2}(j,i+1/2) = v_x^{n-1/2}(j,i+1/2) + \\frac{dt}{\\rho_x(j,i+1/2)}\\cdot \\biggl(\\frac{\\tau^n_{xx}(j,i+1) - \\tau^n_{xx}(j,i)}{dx} + \\frac{\\tau^n_{xz}(j+1/2, i) - \\tau^n_{xz}(j-1/2,i)}{dz} \\biggr)} $$\n",
    "\n",
    "Usando un enfoque similar, podemos derivar el esquema FD para la componente y de la ecuación de momento ...\n",
    "\n",
    "$$\\rm{v_z^{n+1/2}(j,i+1/2) = v_z^{n-1/2}(j,i+1/2) + \\frac{dt}{\\rho_z(j+1/2,i)}\\cdot \\biggl(\\frac{\\tau^n_{xz}(j, i+1/2) - \\tau^n_{xz}(j,i-1/2)}{dx} + \\frac{\\tau^n_{zz}(j+1,i) - \\tau^n_{zz}(j,i)}{dz} \\biggr)} $$\n",
    "\n",
    "... y los esfuerzos ...\n",
    "\n",
    "$$\n",
    "\\begin{split}   \n",
    "\\rm{\\tau^{n+1}_{xx}(j,i)}\\;&\\rm{= \\tau_{xx}^{n}(j,i) + dt\\cdot\\lambda(j,i)\\cdot \\biggl(v^{n+1/2}_{xx}(j,i) + v^{n+1/2}_{zz}(j,i) \\biggr) + 2 dt\\cdot  \\mu(j,i) \\cdot  v^{n+1/2}_{xx}(j,i)}\\\\ \n",
    "\\rm{\\tau^{n+1}_{zz}(j,i)}\\;&\\rm{= \\tau_{zz}^{n}(j,i) + dt\\cdot\\lambda(j,i)\\cdot \\biggl(v^{n+1/2}_{xx}(j,i) + v^{n+1/2}_{zz}(j,i) \\biggr) + 2 dt\\cdot  \\mu(j,i) \\cdot  v^{n+1/2}_{zz}(j,i)}\\\\ \n",
    "\\rm{\\tau^{n+1}_{xz}(j+1/2,i+1/2)}\\;&\\rm{=\\tau^{n}_{xz}(j+1/2,i+1/2) + dt\\cdot\\mu_{xz}(j+1/2,i+1/2)\\biggl(v^{n+1/2}_{xz}(j+1/2,i+1/2) + v^{n+1/2}_{zx}(j+1/2,i+1/2)\\biggr)}\\\\\n",
    "\\end{split}\n",
    "$$\n",
    "\n",
    "con las derivadas espaciales\n",
    "\n",
    "$$\n",
    "\\begin{split}  \n",
    "\\rm{v_{xx}(j,i)}\\; & \\rm{= \\frac{v_x(j,i+1/2)-v_x(j,i-1/2)}{dx}}\\\\ \n",
    "\\rm{v_{zz}(j,i)}\\; & \\rm{= \\frac{v_z(j+1/2,i)-v_z(j-1/2,i)}{dz}}\\\\ \n",
    "\\rm{v_{zx}(j+1/2,i+1/2)}\\; & \\rm{= \\frac{v_z(j+1/2, i+1)-v_z(j+1/2, i)}{dx}}\\\\ \n",
    "\\rm{v_{xz}(j+1/2,i+1/2)}\\; & \\rm{= \\frac{v_x(j+1,i+1/2)-v_x(j, i+1/2)}{dz}}\\\\\n",
    "\\end{split}\n",
    "$$\n",
    "\n",
    "#### Termino o función fuente\n",
    "\n",
    "En este codigo se uso como función fuente se uso una ondicula  cuya ecuación es:\n",
    "\n",
    "$$\n",
    "s= exp\\left(-f_0^2*((t-t_0)*dt)^2\\right)\n",
    "$$\n",
    "\n",
    "donde $f_0$ es la frecuencia dominante y $t_0$ es el retraso en tiempo"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Import Libraries \n",
    "# ----------------------------------------------\n",
    "%matplotlib notebook\n",
    "import numpy as np\n",
    "import numpy.matlib as npmat\n",
    "import matplotlib.pyplot as plt"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def d(comp,inpu,dire):\n",
    "    if(comp=='x'): idim = 0\n",
    "    elif(comp=='z'): idim = 1\n",
    "    \n",
    "    if(dire == '+'):\n",
    "        return (9/8*(np.roll(inpu, -1, axis = idim)-inpu)\n",
    "                -1/24*(np.roll(inpu, -2, axis = idim)-np.roll(inpu, 1, axis = idim)))/dx\n",
    "    elif(dire == '-'):\n",
    "        return (9/8*(inpu - np.roll(inpu, 1, axis = idim))\n",
    "                -1/24*(np.roll(inpu, -1, axis = idim)-np.roll(inpu, 2, axis = idim)))/dx"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Input Parameter\n",
    "# ----------------------------------------------\n",
    "\n",
    "# Discretization\n",
    "c1=20           # Number of grid points per dominant wavelength\n",
    "c2=0.6          # CFL-Number for 2D CFL< 0.606\n",
    "nx=1400          # number of grid points in x-direction\n",
    "nz=1400          # number of grid points in z-direction\n",
    "T=1.03\n",
    "idisp = 10 # display frequency\n",
    "\n",
    "# The grid number of pml\n",
    "pml = 50                              \n",
    "m = nx + 2*pml\n",
    "n = nz + 2*pml\n",
    " \n",
    "# Source Signal\n",
    "f0= 50      # Center frequency Ricker-wavelet\n",
    "q0= 1      # Maximum amplitude Ricker-Wavelet\n",
    "t0= 1.5/f0\n",
    "xscr = int(m/2)-1;  zscr = int(n/30)+6 # Source position (in grid points)\n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Velocity and density\n",
    "Vpp=2000*np.ones((nx,nz))   # P-wave velocity in m/s\n",
    "Vss=1200*np.ones((nx,nz))   # S-wave velocity in m/s\n",
    "den= 1500*np.ones((nx,nz))       # Density in Kg/m^3\n",
    "\n",
    "Vpp[:,0:200]= 1000\n",
    "Vss[:,0:200]= 600\n",
    "den[:,0:200]= 750\n",
    "\n",
    "for j in range(3*(nz//4)-1,nz):\n",
    "    Vpp[:,j]=Vpp[:,j]*1.5\n",
    "    \n",
    "for j in range(3*(nz//4)-1,nz):\n",
    "    Vss[:,j]=Vss[:,j]*1.5\n",
    "   \n",
    "for j in range(3*(nz//4)-1,nz):\n",
    "    den[:,j]=den[:,j]*1.5"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Plotting velocity and density model\n",
    "plt.ion()\n",
    "fig1 = plt.figure()\n",
    "plt.title('P-wave Velocity Model (m/s)')\n",
    "plt.xlabel('x-grid')\n",
    "plt.ylabel('z-grid')\n",
    "im = plt.imshow(Vpp.transpose(),vmin=Vpp.min(), vmax=Vpp.max(), animated=True,\n",
    "                interpolation=\"nearest\", cmap=plt.cm.RdBu)\n",
    "fig1.colorbar(im)\n",
    "\n",
    "fig2 = plt.figure()\n",
    "plt.title('S-wave Velocity Model (m/s)')\n",
    "plt.xlabel('x-grid')\n",
    "plt.ylabel('z-grid')\n",
    "im = plt.imshow(Vss.transpose(),vmin=Vss.min(), vmax=Vss.max(), animated=True,\n",
    "                interpolation=\"nearest\", cmap=plt.cm.RdBu)\n",
    "fig2.colorbar(im)\n",
    "\n",
    "fig3 = plt.figure()\n",
    "plt.title('Density Model (Kg/$m^3$)')\n",
    "plt.xlabel('x-grid')\n",
    "plt.ylabel('z-grid')\n",
    "im = plt.imshow(den.transpose(),vmin=den.min(), vmax=den.max(), animated=True,\n",
    "                interpolation=\"nearest\", cmap=plt.cm.RdBu)\n",
    "fig3.colorbar(im)\n",
    "\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# PML extra edge\n",
    "\n",
    "Vp=np.zeros((m,n))                        # medium velocity of P wave,m/s,zero fill\n",
    "Vp[pml:pml+nx,pml:pml+nz] = Vpp\n",
    "Vp[pml:pml+nx,0:pml] = np.transpose(npmat.repmat(Vpp[:,0],pml,1))\n",
    "Vp[pml:pml+nx,nz+pml:2*pml+nz] = np.transpose(npmat.repmat(Vpp[:,nz-1],pml,1))\n",
    "Vp[0:pml,:] = npmat.repmat(Vp[pml,:],pml, 1)\n",
    "Vp[pml+nx:2*pml+nx,:] = npmat.repmat(Vp[pml+nx-1,:],pml, 1)\n",
    "del Vpp\n",
    "\n",
    "Vs=np.zeros((m,n))                        # medium velocity of S wave,m/s,zero fill\n",
    "Vs[pml:pml+nx,pml:pml+nz] = Vss\n",
    "Vs[pml:pml+nx,0:pml] = np.transpose(npmat.repmat(Vss[:,0],pml,1))\n",
    "Vs[pml:pml+nx,nz+pml:nz+2*pml] = np.transpose(npmat.repmat(Vss[:,nz-1],pml,1))\n",
    "Vs[0:pml,:] = npmat.repmat(Vs[pml,:],pml, 1)\n",
    "Vs[pml+nx:2*pml+nx,:] = npmat.repmat(Vs[pml+nx-1,:],pml, 1)\n",
    "del Vss;\n",
    "\n",
    "rho=np.zeros((m,n))                        # medium density,kg/m**3,zero fill\n",
    "rho[pml:pml+nx,pml:pml+nz] = den;\n",
    "rho[pml:pml+nx,0:pml] = np.transpose(npmat.repmat(den[:,0],pml,1))\n",
    "rho[pml:pml+nx,nz+pml:nz+2*pml] = np.transpose(npmat.repmat(den[:,nz-1],pml,1))\n",
    "rho[0:pml,:] = npmat.repmat(rho[pml,:],pml, 1)\n",
    "rho[pml+nx:2*pml+nx,:] =npmat.repmat(rho[pml+nx-1,:],pml,1)\n",
    "del den\n",
    "\n",
    "# Calculate Lame  Parameters\n",
    "lam=rho*(Vp**2 - 2*Vs**2)\n",
    "mu=rho*Vs**2\n",
    "del Vs"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "## Preparation \n",
    "\n",
    "# Init wavefields\n",
    "vx=np.zeros((m,n))\n",
    "vz=np.zeros((m,n))\n",
    "sxx=np.zeros((m,n))\n",
    "szz=np.zeros((m,n))\n",
    "sxz=np.zeros((m,n))\n",
    "vx1=np.zeros((m,n))\n",
    "vx2=np.zeros((m,n))\n",
    "vz1=np.zeros((m,n))\n",
    "vz2=np.zeros((m,n))\n",
    "sxx1=np.zeros((m,n))\n",
    "sxx2=np.zeros((m,n))\n",
    "szz1=np.zeros((m,n))\n",
    "szz2=np.zeros((m,n))\n",
    "sxz1=np.zeros((m,n))\n",
    "sxz2=np.zeros((m,n))\n",
    "\n",
    "vol   = np.zeros([m,n])\n",
    "curl  = np.zeros([m,n])\n",
    "\n",
    "cmin=Vp.min()  # Lowest P-wave velocity\n",
    "cmax=Vp.max()  # Highest P-wave velocity\n",
    "fmax=2*f0              # Maximum frequency\n",
    "dx=cmin/(fmax*c1)      # Spatial discretization (in m)\n",
    "dz=dx \n",
    "dt=dx/(cmax)*c2        # Temporal discretization (in s)\n",
    "lampda_min=cmin/fmax   # Smallest wavelength\n",
    "print(dx,dz,dt)\n",
    "del Vp\n",
    "\n",
    "# Create space and time vector\n",
    "x=np.arange(0,dx*nx-dx,dx) # Space vector\n",
    "z=np.arange(0,(dz*nz-dz),dz) # Space vector\n",
    "t=np.arange(0,(T-dt),dt)    # Time vector\n",
    "nt=len(t)       # Number of time steps\n",
    "\n",
    "# Shot gathers\n",
    "seismo_vx=np.zeros((m,nt))\n",
    "seismo_vz=np.zeros((m,nt))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": false
   },
   "outputs": [],
   "source": [
    "# Source signal - Ricker-wavelet\n",
    "\n",
    "# Initialise source\n",
    "tau=np.pi*f0*(t-t0);\n",
    "s=q0*(1-2*tau**2)*np.exp(-tau**2);\n",
    "\n",
    "# Plotting source signal\n",
    "plt.ion()\n",
    "fig4 = plt.figure(figsize=(9,4.5))\n",
    "plt.plot(t-t0,s)\n",
    "plt.title('Source signal Ricker-Wavelet')\n",
    "plt.xlabel('Time in s')\n",
    "plt.ylabel('Amplitude')\n",
    "plt.show() "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#**********************attemuatiom(pnl)************************\n",
    "R=0.5e-0;          # theoretical reflectiom coefficiemt\n",
    "ddx=np.zeros((m,n)) \n",
    "ddz=np.zeros((m,n))\n",
    "plx=pml*dx;\n",
    "plz=pml*dz;\n",
    "for i in range(0,m):\n",
    "     for k in range(0,n):\n",
    "         # zome1\n",
    "         if i>=0 and i<=pml-1 and k>=0 and k<=pml-1:\n",
    "             z=pml-k;x=pml-i;\n",
    "             ddz[i,k]=-np.log(R)*3*cmax*z**2/(2*plz**2);\n",
    "             ddx[i,k]=-np.log(R)*3*cmax*x**2/(2*plx**2);\n",
    "         elif i>=0 and i<=pml-1 and k>n-pml-1 and k<=n:\n",
    "             z=k-(n-pml);x=pml-i;\n",
    "             ddz[i,k]=-np.log(R)*3*cmax*z**2/(2*plz**2);\n",
    "             ddx[i,k]=-np.log(R)*3*cmax*x**2/(2*plx**2);\n",
    "         elif i>m-pml-1 and i<=m and k>=0 and k<=pml-1:\n",
    "             z=pml-k;x=i-(m-pml);\n",
    "             ddz[i,k]=-np.log(R)*3*cmax*z**2/(2*plz**2);\n",
    "             ddx[i,k]=-np.log(R)*3*cmax*x**2/(2*plx**2);\n",
    "         elif i>m-pml-1 and i<=m and k>n-pml-1 and k<=n:\n",
    "             z=k-(n-pml);x=i-(m-pml);\n",
    "             ddz[i,k]=-np.log(R)*3*cmax*z**2/(2*plz**2);\n",
    "             ddx[i,k]=-np.log(R)*3*cmax*x**2/(2*plx**2);\n",
    "  \n",
    "         # xone2\n",
    "         elif i<=pml-1 and k>pml-1 and k<n-pml:\n",
    "             z=0;x=pml-i;\n",
    "             ddz[i,k]=0;ddx[i,k]=-np.log(R)*3*cmax*x**2/(2*plx**2);\n",
    "         elif  i>m-pml-1 and i<=m and k>pml-1 and k<=n-pml-1:\n",
    "             z=0;x=i-(m-pml);\n",
    "             ddz[i,k]=0;ddx[i,k]=-np.log(R)*3*cmax*x**2/(2*plx**2);\n",
    "         # xome3\n",
    "         elif i>pml-1 and i<=m-pml-1 and k<=pml-1:\n",
    "             z=pml-k;x=0;\n",
    "             ddz[i,k]=-np.log(R)*3*cmax*z**2/(2*plz**2);\n",
    "             ddx[i,k]=0;\n",
    "         elif i>pml-1 and i<=m-pml-1 and k>n-pml-1 and k<=n:\n",
    "             z=k-(n-pml);x=0;\n",
    "             ddz[i,k]=-np.log(R)*3*cmax*z**2/(2*plz**2);\n",
    "             ddx[i,k]=0;\n",
    "             \n",
    "#ddx=np.transpose(ddx)\n",
    "#ddz=np.transpose(ddz)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Plot preparation\n",
    "# Initialize animated plot\n",
    "plt.ion()\n",
    "fig5 = plt.figure(figsize=(9,6))\n",
    "f1 = fig5.add_subplot(2, 2, 1)\n",
    "f2 = fig5.add_subplot(2, 2, 2)\n",
    "f3 = fig5.add_subplot(2, 2, 3)\n",
    "f4 = fig5.add_subplot(2, 2, 4)\n",
    "f1.set_title('curl')\n",
    "f2.set_title('div')\n",
    "f3.set_title('vx')\n",
    "f4.set_title('vz')\n",
    "\n",
    "curlp = f1.imshow(curl[pml:m-pml,pml:m-pml].transpose(), interpolation='nearest', animated=True,\n",
    "                   vmin=-2.5e-11, vmax=2.5e-11, cmap=plt.cm.RdBu,extent=[0,(m-2*pml)*dx,(n-2*pml)*dz,0],aspect = 'equal')\n",
    "divp = f2.imshow(vol[pml:m-pml,pml:m-pml].transpose(), interpolation='nearest', animated=True,\n",
    "                   vmin=-2.5e-11, vmax=2.5e-11, cmap=plt.cm.RdBu,extent=[0,(m-2*pml)*dx,(n-2*pml)*dz,0],aspect = 'equal')\n",
    "vxp = f3.imshow(vx[pml:m-pml,pml:m-pml].transpose(), interpolation='nearest', animated=True,\n",
    "                   vmin=-2.5e-10, vmax=2.5e-10, cmap=plt.cm.RdBu,extent=[0,(m-2*pml)*dx,(n-2*pml)*dz,0],aspect = 'equal')\n",
    "vzp = f4.imshow(vz[pml:m-pml,pml:m-pml].transpose(), interpolation='nearest', animated=True,\n",
    "                   vmin=-2.5e-10, vmax=2.5e-10, cmap=plt.cm.RdBu,extent=[0,(m-2*pml)*dx,(n-2*pml)*dz,0],aspect = 'equal')\n",
    "plt.show()\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "for tt in range(1,nt):\n",
    "\n",
    "    # Inject source wavelet (explosion)\n",
    "       \n",
    "    sxx[xscr,zscr]=sxx[xscr,zscr]+s[tt]   \n",
    "    szz[xscr,zscr]=szz[xscr,zscr]+s[tt]\n",
    "    \n",
    "    # Free surface\n",
    "    szz[:,pml-1]=0.\n",
    "    sxz[:,pml-1]=0.\n",
    "               \n",
    "    # Update velocity\n",
    "    vx1=((1-0.5*dt*ddx)*vx1+(dt/rho)*d('x',sxx,'-'))/(1+0.5*dt*ddx)  \n",
    "    vx2=((1-0.5*dt*ddz)*vx2+(dt/rho)*d('z',sxz,'-'))/(1+0.5*dt*ddz)\n",
    "    vx = vx1+vx2\n",
    "    \n",
    "    vz1=((1-0.5*dt*ddx)*vz1+(dt/rho)*d('x',sxz,'+'))/(1+0.5*dt*ddx)\n",
    "    vz2=((1-0.5*dt*ddz)*vz2+(dt/rho)*d('z',szz,'+'))/(1+0.5*dt*ddz)\n",
    "    vz = vz1+vz2 \n",
    "    \n",
    "            \n",
    "    curl = d('z',vx,'+') - d('x',vz,'-')\n",
    "    vol = d('x',vx,'+') + d('z',vz,'-')            \n",
    "                       \n",
    "    # Update pressure\n",
    "    sxx1=((1-0.5*dt*ddx)*sxx1+(lam+2*mu)*dt*d('x',vx,'+'))/(1+0.5*dt*ddx)\n",
    "    sxx2=((1-0.5*dt*ddz)*sxx2+lam*dt*d('z',vz,'-'))/(1+0.5*dt*ddz)\n",
    "    sxx = sxx1 + sxx2    \n",
    "\n",
    "    szz1=((1-0.5*dt*ddx)*szz1+lam*dt*d('x',vx,'+'))/(1+0.5*dt*ddx)\n",
    "    szz2=((1-0.5*dt*ddz)*szz2+(lam+2*mu)*dt*d('z',vz,'-'))/(1+0.5*dt*ddz) \n",
    "    szz = szz1+szz2 \n",
    "\n",
    "    sxz1=((1-0.5*dt*ddx)*sxz1+mu*dt*d('x',vz,'-'))/(1+0.5*dt*ddx)\n",
    "    sxz2=((1-0.5*dt*ddz)*sxz2+mu*dt*d('z',vx,'+'))/(1+0.5*dt*ddz)\n",
    "    sxz = sxz1+sxz2\n",
    "    \n",
    "    seismo_vx[0:m,tt]=vx[0:m,pml]\n",
    "    seismo_vz[0:m,tt]=vz[0:m,pml]\n",
    "    \n",
    "    # Plot every time step (nt)\n",
    "    if (tt % idisp) == 0:\n",
    "        curlp.set_data(curl[pml:m-pml,pml:m-pml].transpose())\n",
    "        f1.set_title(\"Time Step (nt) = %i\" % (tt))\n",
    "        divp.set_data(vol[pml:m-pml,pml:m-pml].transpose())\n",
    "        vxp.set_data(vx[pml:m-pml,pml:m-pml].transpose())\n",
    "        vzp.set_data(vz[pml:m-pml,pml:m-pml].transpose())\n",
    "        plt.gcf().canvas.draw()\n",
    "        \n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": false
   },
   "outputs": [],
   "source": [
    "# Plot Shot gathers 5 pts below free surface\n",
    "    \n",
    "fac=0.1\n",
    "ymx=abs(seismo_vx).max()\n",
    "clim=np.array([-fac*ymx, fac*ymx])\n",
    "plt.ion()\n",
    "fig2 = plt.figure()\n",
    "plt.title('Common Shot Gather (horizontal component phone)')\n",
    "plt.xlabel('x-grid')\n",
    "plt.ylabel('nt')\n",
    "im = plt.imshow(seismo_vx[pml:m-pml,:].transpose(),vmin=clim.min(), vmax=clim.max(), animated=True,\n",
    "                interpolation=\"nearest\", cmap=plt.cm.RdBu,extent=[0,(m-2*pml)*dx,nt*dt - t0,0],aspect = 'auto')\n",
    "fig2.colorbar(im)\n",
    "plt.show()\n",
    "plt.savefig(\"DFDT_het_2D_vx.png\", bbox_inches='tight')\n",
    "\n",
    "fac=0.1\n",
    "ymx=abs(seismo_vz[pml:m-pml,:]).max()\n",
    "clim=np.array([-fac*ymx, fac*ymx])\n",
    "plt.ion()\n",
    "fig3 = plt.figure()\n",
    "plt.title('Common Shot Gather (vertical component phone)')\n",
    "plt.xlabel('x-grid')\n",
    "plt.ylabel('nt')\n",
    "im = plt.imshow(seismo_vz[pml:m-pml,:].transpose(),vmin=clim.min(), vmax=clim.max(), animated=True,\n",
    "                interpolation=\"nearest\", cmap=plt.cm.RdBu,extent=[0,(m-2*pml)*dx,nt*dt - t0,0],aspect = 'auto')\n",
    "fig3.colorbar(im)\n",
    "plt.show()\n",
    "plt.savefig(\"DFDT_het_2D_vz.png\", bbox_inches='tight')"
   ]
  }
 ],
 "metadata": {
  "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"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}