# Execute this cell to load the notebook's style sheet, then ignore it
from IPython.core.display import HTML
css_file = '../../style/custom.css'
HTML(open(css_file, "r").read())

"""
Explanation: Content under Creative Commons Attribution license CC-BY 4.0, code under BSD 3-Clause License © 2018  by D. Koehn, heterogeneous models are from this Jupyter notebook by Heiner Igel (@heinerigel), Florian Wölfl and Lion Krischer (@krischer) which is a supplemenatry material to the book Computational Seismology: A Practical Introduction, notebook style sheet by L.A. Barba, N.C. Clementi
End of explanation
"""


# Import Libraries 
# ----------------
import numpy as np
from numba import jit
import matplotlib
import matplotlib.pyplot as plt
from pylab import rcParams

# Ignore Warning Messages
# -----------------------
import warnings
warnings.filterwarnings("ignore")

from mpl_toolkits.axes_grid1 import make_axes_locatable

# Definition of initial modelling parameters
# ------------------------------------------
xmax = 2000.0 # maximum spatial extension of the 1D model in x-direction (m)
zmax = xmax   # maximum spatial extension of the 1D model in z-direction (m)
dx   = 10.0   # grid point distance in x-direction (m)
dz   = dx     # grid point distance in z-direction (m)

tmax = 0.75    # maximum recording time of the seismogram (s)
dt   = 0.0010 # time step

vp0  = 3000.  # P-wave speed in medium (m/s)

# acquisition geometry
xsrc = 1000.0 # x-source position (m)
zsrc = xsrc   # z-source position (m)

f0   = 100.0 # dominant frequency of the source (Hz)
t0   = 0.1   # source time shift (s)

isnap = 2  # snapshot interval (timesteps)

"""
Explanation: Simple absorbing boundary for 2D acoustic FD modelling
Realistic FD modelling results for surface seismic acquisition geometries require a further modification of the 2D acoustic FD code. Except for the free surface boundary condition on top of the model, we want to suppress the artifical reflections from the other boundaries.
Such absorbing boundaries can be implemented by different approaches. A comprehensive overview is compiled in 
Gao et al. 2015, Comparison of artificial absorbing boundaries for acoustic wave equation modelling
Before implementing the absorbing boundary frame, we modify some parts of the optimized 2D acoustic FD code:
End of explanation
"""


@jit(nopython=True) # use JIT for C-performance
def update_d2px_d2pz(p, dx, dz, nx, nz, d2px, d2pz):
    
    for i in range(1, nx - 1):
        for j in range(1, nz - 1):
                
            d2px[i,j] = (p[i + 1,j] - 2 * p[i,j] + p[i - 1,j]) / dx**2                
            d2pz[i,j] = (p[i,j + 1] - 2 * p[i,j] + p[i,j - 1]) / dz**2
        
    return d2px, d2pz   

"""
Explanation: In order to modularize the code, we move the 2nd partial derivatives of the wave equation into a function update_d2px_d2pz, so the application of the JIT decorator can be restriced to this function:
End of explanation
"""


# FD_2D_acoustic code with JIT optimization
# -----------------------------------------
def FD_2D_acoustic_JIT(dt,dx,dz,f0):        
    
    # define model discretization
    # ---------------------------

    nx = (int)(xmax/dx) # number of grid points in x-direction
    print('nx = ',nx)

    nz = (int)(zmax/dz) # number of grid points in x-direction
    print('nz = ',nz)

    nt = (int)(tmax/dt) # maximum number of time steps            
    print('nt = ',nt)

    isrc = (int)(xsrc/dx)  # source location in grid in x-direction
    jsrc = (int)(zsrc/dz)  # source location in grid in x-direction

    # Source time function (Gaussian)
    # -------------------------------
    src  = np.zeros(nt + 1)
    time = np.linspace(0 * dt, nt * dt, nt)

    # 1st derivative of Gaussian
    src  = -2. * (time - t0) * (f0 ** 2) * (np.exp(- (f0 ** 2) * (time - t0) ** 2))
    
    # define clip value: 0.1 * absolute maximum value of source wavelet
    clip = 0.1 * max([np.abs(src.min()), np.abs(src.max())]) / (dx*dz) * dt**2    
    
    # Define model
    # ------------    
    vp  = np.zeros((nx,nz))
    vp  = model(nx,nz,vp,dx,dz)
    vp2 = vp**2    
    
    # Initialize empty pressure arrays
    # --------------------------------
    p    = np.zeros((nx,nz)) # p at time n (now)
    pold = np.zeros((nx,nz)) # p at time n-1 (past)
    pnew = np.zeros((nx,nz)) # p at time n+1 (present)
    d2px = np.zeros((nx,nz)) # 2nd spatial x-derivative of p
    d2pz = np.zeros((nx,nz)) # 2nd spatial z-derivative of p        
    
    # Initalize animation of pressure wavefield 
    # -----------------------------------------    
    fig = plt.figure(figsize=(7,3.5))  # define figure size
    plt.tight_layout()
    extent = [0.0,xmax,zmax,0.0]     # define model extension
    
    # Plot pressure wavefield movie
    ax1 = plt.subplot(121)
    image = plt.imshow(p.T, animated=True, cmap="RdBu", extent=extent, 
                          interpolation='nearest', vmin=-clip, vmax=clip)        
    plt.title('Pressure wavefield')
    plt.xlabel('x [m]')
    plt.ylabel('z [m]')
    
    # Plot Vp-model
    ax2 = plt.subplot(122)
    image1 = plt.imshow(vp.T/1000, cmap=plt.cm.viridis, interpolation='nearest', 
                        extent=extent)
    
    plt.title('Vp-model')
    plt.xlabel('x [m]')
    plt.setp(ax2.get_yticklabels(), visible=False) 
    
    divider = make_axes_locatable(ax2)
    cax2 = divider.append_axes("right", size="2%", pad=0.1)
    fig.colorbar(image1, cax=cax2)         
    plt.ion()    
    plt.show(block=False)
    
    # Calculate Partial Derivatives
    # -----------------------------
    for it in range(nt):
    
        # FD approximation of spatial derivative by 3 point operator
        d2px, d2pz = update_d2px_d2pz(p, dx, dz, nx, nz, d2px, d2pz)

        # Time Extrapolation
        # ------------------
        pnew = 2 * p - pold + vp2 * dt**2 * (d2px + d2pz)

        # Add Source Term at isrc
        # -----------------------
        # Absolute pressure w.r.t analytical solution
        pnew[isrc,jsrc] = pnew[isrc,jsrc] + src[it] / (dx * dz) * dt ** 2        
        
        # Remap Time Levels
        # -----------------
        pold, p = p, pnew
    
        # display pressure snapshots 
        if (it % isnap) == 0:
            image.set_data(p.T)
            fig.canvas.draw()       

"""
Explanation: In the FD modelling code FD_2D_acoustic_JIT, a more flexible model definition is introduced by the function model. The block Initalize animation of pressure wavefield before the time loop displays the velocity model and initial pressure wavefield. During the time-loop, the pressure wavefield is updated with             
image.set_data(p.T)
fig.canvas.draw()
at the every isnap timestep:
End of explanation
"""


# Homogeneous model
def model(nx,nz,vp,dx,dz):
    
    vp += vp0 
    
    return vp

"""
Explanation: Homogeneous block model without absorbing boundary frame
As a reference, we first model the homogeneous block model, defined in the function model, without an absorbing boundary frame:
End of explanation
"""


%matplotlib notebook
dx   = 5.0   # grid point distance in x-direction (m)
dz   = dx     # grid point distance in z-direction (m)
f0   = 100.0  # centre frequency of the source wavelet (Hz)

# calculate dt according to the CFL-criterion
dt = dx / (np.sqrt(2.0) * vp0)

FD_2D_acoustic_JIT(dt,dx,dz,f0)

"""
Explanation: After defining the modelling parameters, we can run the modified FD code ...
End of explanation
"""


# Define simple absorbing boundary frame based on wavefield damping 
# according to Cerjan et al., 1985, Geophysics, 50, 705-708
def absorb(nx,nz):
   
    FW =      # thickness of absorbing frame (gridpoints)    
    a =   # damping variation within the frame
    
    coeff = np.zeros(FW)
    
    # define coefficients


    # initialize array of absorbing coefficients
    absorb_coeff = np.ones((nx,nz))

    # compute coefficients for left grid boundaries (x-direction)


    # compute coefficients for right grid boundaries (x-direction)        


    # compute coefficients for bottom grid boundaries (z-direction)        


    return absorb_coeff

"""
Explanation: Notice the strong, artifical boundary reflections in the wavefield movie
Simple absorbing Sponge boundary
The simplest, and unfortunately least efficient, absorbing boundary was developed by Cerjan et al. (1985). It is based on the simple idea to damp the pressure wavefields $p^n_{i,j}$ and $p^{n+1}_{i,j}$ in an absorbing boundary frame by an exponential function:
\begin{equation}
f_{abs} = exp(-a^2(FW-i)^2), \nonumber
\end{equation}
where $FW$ denotes the thickness of the boundary frame in gridpoints, while the factor $a$ defines the damping variation within the frame. It is import to avoid overlaps of the damping profile in the model corners, when defining the absorbing function:
End of explanation
"""


# Plot absorbing damping profile
# ------------------------------
fig = plt.figure(figsize=(6,4))  # define figure size
extent = [0.0,xmax,0.0,zmax]      # define model extension

# calculate absorbing boundary weighting coefficients
nx = 400
nz = 400
absorb_coeff = absorb(nx,nz)

plt.imshow(absorb_coeff.T)
plt.colorbar()
plt.title('Sponge boundary condition')
plt.xlabel('x [m]')
plt.ylabel('z [m]')
plt.show()

"""
Explanation: This implementation of the Sponge boundary sets a free-surface boundary condition on top of the model, while inciding waves at the other boundaries are absorbed:
End of explanation
"""


# FD_2D_acoustic code with JIT optimization
# -----------------------------------------
def FD_2D_acoustic_JIT_absorb(dt,dx,dz,f0):        
    
    # define model discretization
    # ---------------------------

    nx = (int)(xmax/dx) # number of grid points in x-direction
    print('nx = ',nx)

    nz = (int)(zmax/dz) # number of grid points in x-direction
    print('nz = ',nz)

    nt = (int)(tmax/dt) # maximum number of time steps            
    print('nt = ',nt)

    isrc = (int)(xsrc/dx)  # source location in grid in x-direction
    jsrc = (int)(zsrc/dz)  # source location in grid in x-direction

    # Source time function (Gaussian)
    # -------------------------------
    src  = np.zeros(nt + 1)
    time = np.linspace(0 * dt, nt * dt, nt)

    # 1st derivative of Gaussian
    src  = -2. * (time - t0) * (f0 ** 2) * (np.exp(- (f0 ** 2) * (time - t0) ** 2))
    
    # define clip value: 0.1 * absolute maximum value of source wavelet
    clip = 0.1 * max([np.abs(src.min()), np.abs(src.max())]) / (dx*dz) * dt**2
    
    # Define absorbing boundary frame
    # ------------------------------- 

    
    # Define model
    # ------------    
    vp  = np.zeros((nx,nz))
    vp  = model(nx,nz,vp,dx,dz)
    vp2 = vp**2    
    
    # Initialize empty pressure arrays
    # --------------------------------
    p    = np.zeros((nx,nz)) # p at time n (now)
    pold = np.zeros((nx,nz)) # p at time n-1 (past)
    pnew = np.zeros((nx,nz)) # p at time n+1 (present)
    d2px = np.zeros((nx,nz)) # 2nd spatial x-derivative of p
    d2pz = np.zeros((nx,nz)) # 2nd spatial z-derivative of p    
    
    # Initalize animation of pressure wavefield 
    # -----------------------------------------    
    fig = plt.figure(figsize=(7,3.5))  # define figure size
    plt.tight_layout()
    extent = [0.0,xmax,zmax,0.0]     # define model extension
    
    # Plot pressure wavefield movie
    ax1 = plt.subplot(121)
    image = plt.imshow(p.T, animated=True, cmap="RdBu", extent=extent, 
                          interpolation='nearest', vmin=-clip, vmax=clip)        
    plt.title('Pressure wavefield')
    plt.xlabel('x [m]')
    plt.ylabel('z [m]')
    
    # Plot Vp-model
    ax2 = plt.subplot(122)
    image1 = plt.imshow(vp.T/1000, cmap=plt.cm.viridis, interpolation='nearest', 
                        extent=extent)
    
    plt.title('Vp-model')
    plt.xlabel('x [m]')    
    plt.setp(ax2.get_yticklabels(), visible=False) 
    
    divider = make_axes_locatable(ax2)
    cax2 = divider.append_axes("right", size="2%", pad=0.1)
    fig.colorbar(image1, cax=cax2)         
    plt.ion()    
    plt.show(block=False)
    
    # Calculate Partial Derivatives
    # -----------------------------
    for it in range(nt):
    
        # FD approximation of spatial derivative by 3 point operator
        d2px, d2pz = update_d2px_d2pz(p, dx, dz, nx, nz, d2px, d2pz)

        # Time Extrapolation
        # ------------------
        pnew = 2 * p - pold + vp2 * dt**2 * (d2px + d2pz)

        # Add Source Term at isrc
        # -----------------------
        # Absolute pressure w.r.t analytical solution
        pnew[isrc,jsrc] = pnew[isrc,jsrc] + src[it] / (dx * dz) * dt ** 2
        
        # Apply absorbing boundary frame to p and pnew

        
        # Remap Time Levels
        # -----------------
        pold, p = p, pnew
    
        # display pressure snapshots 
        if (it % isnap) == 0:
            image.set_data(p.T)
            fig.canvas.draw()       

"""
Explanation: The FD code itself requires only some small modifications, we have to add the absorb function to define the amount of damping in the boundary frame and apply the damping function to the pressure wavefields pnew and p
End of explanation
"""


%matplotlib notebook
dx   = 5.0   # grid point distance in x-direction (m)
dz   = dx     # grid point distance in z-direction (m)
f0   = 100.0  # centre frequency of the source wavelet (Hz)

# calculate dt according to the CFL-criterion
dt = dx / (np.sqrt(2.0) * vp0)

FD_2D_acoustic_JIT_absorb(dt,dx,dz,f0)

"""
Explanation: Let's evaluate the influence of the Sponge boundaries on the artifical boundary reflections:
End of explanation
"""


