#!/usr/bin/env python
# coding: utf-8

# # 2D Shallow Water Equations 🌊
# 
# This notebook is a Devito reimplementation from this [notebook](https://github.com/daniel-koehn/Differential-equations-earth-system/blob/e291791b6d6b32f776af01532b37a9869c5cf569/10_Shallow_Water_Equation_2D/01_2D_Shallow_Water_Equations.ipynb). We will use an approximation to the Navier-Stokes equations - the 2D Shallow Water equations - in order to model the propagation of Tsunami events.
# 
# I've only made minor modifications to the text, in order to make it more compatible to the new coding format.
# 
# Author: Átila Saraiva Q. Soares.

# ## Governing Equations
# 
# Starting from the continuity and momentum conservation equations, we want to model the following problem:
# 
# <img src="images/shallow_water_sketch.png" style="width: 800px;"/>
# 
# For a given bathymetry model, which can include a complex seafloor topography, we want to model the amplitudes, speed and interaction of waves at the seasurface. At a given point $(x,\; y)$, the thickness of the water column between the seafloor and undisturbed water surface is defined by the variable $h$, while the wave amplitude is $\eta$ and therefore the whole thickness of the water column $D = h + \eta$.
# 
# Using appropriate boundary conditions at the water surface/seafloor, assuming that the horizontal wavelength of the modelled waves are much larger than the water depth and integrating the conservation of mass and momentum equations over the water column, we can derive the following equations to decribe wave propagation 
# 
# \begin{equation}
# \begin{split}
# \frac{\partial \eta}{\partial t} &+ \frac{\partial M}{\partial x} + \frac{\partial N}{\partial y} = 0\\
# \frac{\partial M}{\partial t} &+ \frac{\partial}{\partial x} \biggl(\frac{M^2}{D}\biggr) + \frac{\partial}{\partial y} \biggl(\frac{MN}{D}\biggr) + g D \frac{\partial \eta}{\partial x} + \frac{g \alpha^2}{D^{7/3}} M \sqrt{M^2+N^2} = 0\\
# \frac{\partial N}{\partial t} &+ \frac{\partial}{\partial x} \biggl(\frac{MN}{D}\biggr) + \frac{\partial}{\partial y} \biggl(\frac{N^2}{D}\biggr) + g D \frac{\partial \eta}{\partial y} + \frac{g \alpha^2}{D^{7/3}} N \sqrt{M^2+N^2} = 0
# \end{split}
# \tag{1}
# \end{equation}
# 
# known as **2D Shallow Water Equations (SWE)**. The derivation of these equations is beyond the scope of this notebook. Therefore, I refer to the [Tsunami Modelling Handbook](http://www.tsunami.civil.tohoku.ac.jp/hokusai3/J/projects/manual-ver-3.1.pdf) and the lecture [Shallow Water Derivation and Applications by Christian Kühbacher](http://www.mathematik.tu-dortmund.de/lsiii/cms/papers/Kuehbacher2009.pdf) for further details.
# 
# In Eq. (1) the discharge fluxes $M,\; N$ in x- and y-direction, respectively are given by
# 
# \begin{equation}
# \begin{split}
# M & = \int_{-h}^\eta u dz = u(h+\eta) = uD\\
# N & = \int_{-h}^\eta v dz = v(h+\eta) = vD\\
# \end{split}
# \tag{2}
# \end{equation}
# 
# with the horizontal velocity components $u,\;v$ in x- and y-direction, while $g$ denotes the gravity acceleration. The terms $\frac{g \alpha^2}{D^{7/3}} M \sqrt{M^2+N^2}$ and $\frac{g \alpha^2}{D^{7/3}} N \sqrt{M^2+N^2}$ describe the influence of seafloor friction on the wave amplitude. $\alpha$ denotes the Manning's roughness which can be as small as 0.01 for neat cement or smooth metal up to 0.06 for very poor natural channels (see [Tsunami modelling handbook](http://www.tsunami.civil.tohoku.ac.jp/hokusai3/J/projects/manual-ver-3.1.pdf)).
# 
# The Shallow Water Equations can be applied to 
# 
# - Tsunami prediction
# - Atmospheric flow
# - Storm surges
# - Flows around structures
# - Planetary flows
# 
# and easily extended to incorporate the effects of Coriolis forces, tides or wind stress.

# ## Operator implementation
# 
# This is a simple function which returns the operator that solves this equation. The important part is the operator, which contains all the equations expressed above.

# In[1]:


from devito import Eq, TimeFunction, sqrt, Function, Operator, Grid, solve, ConditionalDimension
from matplotlib import pyplot as plt
import numpy as np


def ForwardOperator(etasave, eta, M, N, h, D, g, alpha, grid):
    """
    Operator that solves the equations expressed above.
    It computes and returns the discharge fluxes M, N and wave height eta from
    the 2D Shallow water equation using the FTCS finite difference method.
    
    Parameters
    ----------
    eta : TimeFunction
        The wave height field as a 2D array of floats.
    M : TimeFunction
        The discharge flux field in x-direction as a 2D array of floats.
    N : TimeFunction
        The discharge flux field in y-direction as a 2D array of floats.
    h : Function
        Bathymetry model as a 2D array of floats.
    D : Function
        Total thickness of the water column.
    g : float
        gravity acceleration.
    alpha : float
        Manning's roughness coefficient.
    etasave : TimeFunction
        Function that is sampled in a different interval than the normal propagation
        and is responsible for saving the snapshots required for the following
        animations.
    """
    
    eps = np.finfo(grid.dtype).eps
    
    # Friction term expresses the loss of amplitude from the friction with the seafloor
    frictionTerm = g * alpha**2 * sqrt(M**2 + N**2) / D**(7./3.)

    # System of equations
    pde_eta = Eq(eta.dt + M.dxc + N.dyc)
    pde_M   = Eq(M.dt + (M**2/D).dxc + (M*N/D).dyc + g*D*eta.forward.dxc + frictionTerm*M)
    pde_N   = Eq(N.dt + (M.forward*N/D).dxc + (N**2/D).dyc + g*D*eta.forward.dyc + g * alpha**2 * sqrt(M.forward**2 + N**2) / D**(7./3.)*N)

    stencil_eta = solve(pde_eta, eta.forward)
    stencil_M   = solve(pde_M, M.forward)
    stencil_N   = solve(pde_N, N.forward)

    # Equations with the forward in time term isolated
    update_eta  = Eq(eta.forward, stencil_eta, subdomain=grid.interior)
    update_M    = Eq(M.forward, stencil_M, subdomain=grid.interior)
    update_N    = Eq(N.forward, stencil_N, subdomain=grid.interior)
    eq_D        = Eq(D, eta.forward + h)

    return Operator([update_eta, update_M, update_N, eq_D] + [Eq(etasave, eta)])


# The following function transforms the saved wavefield snapshots and transform them into a video compatible with jupyter notebook.

# In[2]:


from IPython.display import HTML, display
import matplotlib.pyplot as plt
import matplotlib.animation as animation


def snaps2video (eta, title):
    fig, ax = plt.subplots()
    matrice = ax.imshow(eta.data[0, :, :].T, vmin=-1, vmax=1, cmap="seismic")
    plt.colorbar(matrice)

    plt.xlabel('x')
    plt.ylabel('z')
    plt.title(title)    

    def update(i):
        matrice.set_array(eta.data[i, :, :].T)
        return matrice,

    # Animation
    ani = animation.FuncAnimation(fig, update, frames=nsnaps, interval=50, blit=True)

    plt.close(ani._fig)
    display(HTML(ani.to_html5_video()))


# In[3]:


def plotDepthProfile (h, title):
    fig, ax = plt.subplots()
    matrice = ax.imshow(h0)
    plt.colorbar(matrice)
    plt.xlabel('x')
    plt.ylabel('z')
    plt.title(title)
    plt.show()


# ## Example I: Tsunami in ocean with constant depth
# 
# After writing the `Shallow_water_2D` code and all the required functions attached to it, we can define and run our first 2D Tsunami modelling run.
# 
# Let's assume that the ocean model is $ L_x = 100\; m$ in x-direction and $L_y = 100\; m$ in y-direction. The model is discretized with $nx=401$ gridpoints in x-direction and $ny=401$ gridpoints in y-direction, respectively.
# 
# In this first modelling run, we assume a constant bathymetry $h=50\;m$. The initial wave height field $\eta_0$ is defined as a Gaussian at the center of the model, with a half-width of 10 m and an amplitude of 0.5 m. Regarding the initial discharge fluxes, we assume that 
# 
# \begin{equation}
# \begin{split}
# M_0(x,y) &= 100 \eta_0(x,y)\\
# N_0(x,y) &= 0\\
# \end{split}\notag
# \end{equation}
# 
# In order to avoid the occurence of high frequency artifacts, when waves are interacting with the boundaries, boundary conditions will be avoided. Normally Dirichlet conditions are used for the M and N discharge fluxes, and Neumann for the wave height field. Those lead to significant boundary reflections which might be not realistic for a given problem.
# 
# Let's assume the gravity $g = 9.81$ and the Manning's roughness coefficient $\alpha = 0.025$ for the all the remaining examples.

# In[4]:


Lx    = 100.0   # width of the mantle in the x direction []
Ly    = 100.0   # thickness of the mantle in the y direction []
nx    = 401     # number of points in the x direction
ny    = 401     # number of points in the y direction
dx    = Lx / (nx - 1)  # grid spacing in the x direction []
dy    = Ly / (ny - 1)  # grid spacing in the y direction []
g     = 9.81  # gravity acceleration [m/s^2]
alpha = 0.025 # friction coefficient for natural channels in good condition

# Maximum wave propagation time [s]
Tmax  = 3.
dt    = 1/4500.
nt    = (int)(Tmax/dt)
print(dt, nt)

x = np.linspace(0.0, Lx, num=nx)
y = np.linspace(0.0, Ly, num=ny)

# Define initial eta, M, N
X, Y = np.meshgrid(x,y) # coordinates X,Y required to define eta, h, M, N

# Define constant ocean depth profile h = 50 m
h0 = 50. * np.ones_like(X)

# Define initial eta Gaussian distribution [m]
eta0 = 0.5 * np.exp(-((X-50)**2/10)-((Y-50)**2/10))

# Define initial M and N
M0 = 100. * eta0
N0 = 0. * M0
D0 = eta0 + 50.

grid  = Grid(shape=(ny, nx), extent=(Ly, Lx), dtype=np.float32)


# In[5]:


# NBVAL_IGNORE_OUTPUT

nsnaps = 400

# Defining symbolic functions
eta   = TimeFunction(name='eta', grid=grid, space_order=2)
M     = TimeFunction(name='M', grid=grid, space_order=2)
N     = TimeFunction(name='N', grid=grid, space_order=2)
h     = Function(name='h', grid=grid)
D     = Function(name='D', grid=grid)

# Inserting initial conditions
eta.data[0] = eta0.copy()
M.data[0]   = M0.copy()
N.data[0]   = N0.copy()
D.data[:]   = eta0 + h0
h.data[:]   = h0.copy()

# Setting up function to save the snapshots
factor = round(nt / nsnaps)
print(f"The nt/nsnaps factor is {factor}")
time_subsampled = ConditionalDimension('t_sub', parent=grid.time_dim, factor=factor)
etasave = TimeFunction(name='etasave', grid=grid, space_order=2,
                       save=nsnaps, time_dim=time_subsampled)


# Compile the operator
op = ForwardOperator(etasave, eta, M, N, h, D, g, alpha, grid)

# Use the operator
op.apply(eta=eta, etasave=etasave, M=M, N=N, D=D, h=h, time=nt-2, dt=dt)


# In[6]:


# NBVAL_SKIP
snaps2video(etasave, "Modeling a tsunami in ocean with constant depth")


# Let's take a look on what the code looks like

# In[7]:


# To look at the code, uncomment the line below
#print(op.ccode)


# ## Example II: Two Tsunamis in ocean with constant depth
# 
# In example II, we will model the propagation and interaction of two Tsunamis, where the initial conditions for the wave height field $\eta_0$ consists of Gaussian functions with opposite sign located at $(x_1,\;y_1) = (35\; m,\; 35\; m)$ and $(x_2,\;y_2) = (65\; m,\; 65\; m)$. All other modelling parameters remain the same:

# In[8]:


# Define constant ocean depth profile h = 50 m
h0 = 50 * np.ones_like(X)

# Define initial Gaussian eta distribution [m]
eta0 = 0.5 * np.exp(-((X-35)**2/10)-((Y-35)**2/10)) # first Tsunami source
eta0 -= 0.5 * np.exp(-((X-65)**2/10)-((Y-65)**2/10)) # add second Tsunami source

# Define initial M and N
M0 = 100. * eta0
N0 = 0. * M0
D0 = eta0 + h0

# Maximum wave propagation time [s]
Tmax = 3.5
dt = 1/8000.
nt = (int)(Tmax/dt)


# In[9]:


# NBVAL_IGNORE_OUTPUT

# Inserting initial conditions
eta.data[0] = eta0.copy()
M.data[0]   = M0.copy()
N.data[0]   = N0.copy()
D.data[:]   = eta0 + h0
h.data[:]   = h0.copy()

# Setting up function to save the snapshots
factor = round(nt / nsnaps)
print(f"The nt/nsnaps factor is {factor}")
time_subsampled = ConditionalDimension('t_sub', parent=grid.time_dim, factor=factor)
etasave = TimeFunction(name='etasave', grid=grid, space_order=2,
                       save=nsnaps, time_dim=time_subsampled)

# Compile the operator
op = ForwardOperator(etasave, eta, M, N, h, D, g, alpha, grid)

# Use the operator. No need to recompile since the etasave factor is smaller
# than before
op(eta=eta, etasave=etasave, M=M, N=N, D=D, h=h, time=nt-2, dt=dt)


# In[10]:


# NBVAL_SKIP
snaps2video(etasave, "Modeling two tsunamis in ocean with constant depth")


# ## Example III: Tsunami in an ocean with 1D $\tanh$ depth variation
# 
# So far, so good, we can achieve stable and accurate modelling results. However, a constant bathymetry model is a little too simple. In this example, we assume that the bathymetry decreases with a $\tanh$ function from the left to the right boundary in x-direction. Let 's place a Tsunami source at $(x_1,\;y_1) = (30\; m,\; 50\; m)$ and see what 's happening:

# In[11]:


# Define depth profile h [m]
h0 = 50 - 45 * np.tanh((X-70.)/8.)

# Define initial eta [m]
eta0 = 0.5 * np.exp(-((X-30)**2/10)-((Y-50)**2/20))

# Define initial M and N
M0 = 100. * eta0
N0 = 0. * M0
D0 = eta0 + h0

# Maximum wave propagation time [s]
Tmax = 2.
dt = 1/4000.
nt = (int)(Tmax/dt)


# Let's take a look into how this depth profile $h$ looks in the plot below.

# In[12]:


# NBVAL_IGNORE_OUTPUT
plotDepthProfile(h, "Tanh Depth Profile")


# In[13]:


# NBVAL_IGNORE_OUTPUT

# Inserting initial conditions
eta.data[0] = eta0.copy()
M.data[0]   = M0.copy()
N.data[0]   = N0.copy()
D.data[:]   = eta0 + h0
h.data[:]   = h0.copy()

# Setting up function to save the snapshots
factor = round(nt / nsnaps)
print(f"The nt/nsnaps factor is {factor}")
time_subsampled = ConditionalDimension('t_sub', parent=grid.time_dim, factor=factor)
etasave = TimeFunction(name='etasave', grid=grid, space_order=2,
                       save=nsnaps, time_dim=time_subsampled)

# Compile the operator
op = ForwardOperator(etasave, eta, M, N, h, D, g, alpha, grid)

# Use the operator
op(eta=eta, etasave=etasave, M=M, N=N, D=D, h=h, time=nt-2, dt=dt)


# In[14]:


# NBVAL_SKIP
snaps2video(etasave, "Modeling a tsunami in an ocean with 1D Tanh depth variation")


# ## Example IV: Tsunami in an ocean with a seamount
# 
# What happens if we have a constant bathymetry, except for a small seamount reaching 5 m below the water surface? Lets define the model parameters:

# In[15]:


# Define constant ocean depth profile h = 50 m
h0 = 50. * np.ones_like(X)

# Adding seamount to seafloor topography
h0 -= 45. * np.exp(-((X-50)**2/20)-((Y-50)**2/20))

# Define initial eta [m]
eta0 = 0.5 * np.exp(-((X-30)**2/5)-((Y-50)**2/5))

# Define initial M and N
M0 = 100. * eta0
N0 = 0. * M0
D0 = eta0 + h0

# Maximum wave propagation time [s]
Tmax = 2.
dt = 1/8000.
nt = (int)(Tmax/dt)


# Let's take a look into how this depth profile $h$ looks in the plot below.

# In[16]:


# NBVAL_IGNORE_OUTPUT
plotDepthProfile(h, "Seamount Depth Profile")


# In[17]:


# NBVAL_IGNORE_OUTPUT

# Inserting initial conditions
eta.data[0] = eta0.copy()
M.data[0]   = M0.copy()
N.data[0]   = N0.copy()
D.data[:]   = eta0 + h0
h.data[:]   = h0.copy()

# Setting up function to save the snapshots
factor = round(nt / nsnaps)
print(f"The nt/nsnaps factor is {factor}")
time_subsampled = ConditionalDimension('t_sub', parent=grid.time_dim, factor=factor)
etasave = TimeFunction(name='etasave', grid=grid, space_order=2,
                       save=nsnaps, time_dim=time_subsampled)

# Compile the operator
op = ForwardOperator(etasave, eta, M, N, h, D, g, alpha, grid)

# Use the operator
op(eta=eta, etasave=etasave, M=M, N=N, D=D, h=h, time=nt-2, dt=dt)


# In[18]:


# NBVAL_SKIP
snaps2video(etasave, "Modeling a tsunami in an ocean with a seamount")


# ## Example V: Tsunami in an ocean with random seafloor topography variations
# 
# Another modelling idea: what is the influence of the roughness of the seafloor topography. First, we add some random perturbations on the constant bathymetry model $h_0$ by using the `random.rand` function from the `NumPy` library. To smooth the random perturbations, we apply the `gaussian_filter` from the `SciPy`library:

# In[19]:


from scipy.ndimage import gaussian_filter

# Define constant ocean depth profile h = 30 m
h0 = 30. * np.ones_like(X)

# Add random seafloor perturbation of +- 5m
pert = 5.   # perturbation amplitude

np.random.seed(102034)
r = 2.0 * (np.random.rand(ny, nx) - 0.5) * pert # create random number perturbations
r = gaussian_filter(r, sigma=16) # smooth random number perturbation
h0 = h0 * (1 + r) # add perturbations to constant seafloor

# Define initial eta [m]
eta0 = 0.2 * np.exp(-((X-30)**2/5)-((Y-50)**2/5))

# Define initial M and N
M0 = 100. * eta0
N0 = 0. * M0
D0 = eta0 + h0

# Maximum wave propagation time [s]
Tmax = 3.
dt = 1/4000.
nt = (int)(Tmax/dt)


# Let's take a look into how this depth profile $h$ looks in the plot below.

# In[20]:


# NBVAL_IGNORE_OUTPUT
plotDepthProfile(h, "Random seafloor Depth Profile")


# In[21]:


# NBVAL_IGNORE_OUTPUT

# Inserting initial conditions
eta.data[0] = eta0.copy()
M.data[0]   = M0.copy()
N.data[0]   = N0.copy()
D.data[:]   = eta0 + h0
h.data[:]   = h0.copy()

# Setting up function to save the snapshots
factor = round(nt / nsnaps)
print(f"The nt/nsnaps factor is {factor}")
time_subsampled = ConditionalDimension('t_sub', parent=grid.time_dim, factor=factor)
etasave = TimeFunction(name='etasave', grid=grid, space_order=2,
                       save=nsnaps, time_dim=time_subsampled)


# Compile the operator again
op = ForwardOperator(etasave, eta, M, N, h, D, g, alpha, grid)

# Use the operator
op(eta=eta, etasave=etasave, M=M, N=N, D=D, h=h, time=nt-2, dt=dt)


# In[22]:


# NBVAL_SKIP
snaps2video(etasave, "Modeling a tsunami in an ocean with random seafloor topography variations")


# ## Example VI: 2D circular dam break problem
# 
# As a final modelling example, let's take a look at an (academic) engineering problem: a tsunami induced by the collapse of a circular dam in a lake with a constant bathymetry of 30 m. We only need to set the wave height in a circle with radius $r_0 = 5\; m$ to $\eta_0 = 0.5 \; m$ and to zero everywhere else. To avoid the occurence of high frequency artifacts in the wavefield, known as numerical grid dispersion, we apply a Gaussian filter to the initial wave height. To achieve a symmetric dam collapse, the initial discharge fluxes $M_0,N_0$ are set to equal values.

# In[23]:


# Define constant ocean depth profile h = 30 m
h0 = 30. * np.ones_like(X)

# Define initial eta [m]
eta0 = np.zeros_like(X)

# Define mask for circular dam location with radius r0
r0 = 5.
mask = np.where(np.sqrt((X-50)**2 + (Y-50)**2) <= r0)

# Set wave height in dam to 0.5 m
eta0[mask] = 0.5

# Smooth dam boundaries with gaussian filter
eta0 = gaussian_filter(eta0, sigma=8) # smooth random number perturbation

# Define initial M and N
M0 = 1. * eta0
N0 = 1. * M0
D0 = eta0 + h0

# Maximum wave propagation time [s]
Tmax = 3.
dt = 1/4000.
nt = (int)(Tmax/dt)


# In[24]:


# NBVAL_IGNORE_OUTPUT

# Inserting initial conditions
eta.data[0] = eta0.copy()
M.data[0]   = M0.copy()
N.data[0]   = N0.copy()
D.data[:]   = eta0 + h0
h.data[:]   = h0.copy()

# Setting up function to save the snapshots
factor = round(nt / nsnaps)
print(f"The nt/nsnaps factor is {factor}")
time_subsampled = ConditionalDimension('t_sub', parent=grid.time_dim, factor=factor)
etasave = TimeFunction(name='etasave', grid=grid, space_order=2,
                       save=nsnaps, time_dim=time_subsampled)

# Compile the operator again
op = ForwardOperator(etasave, eta, M, N, h, D, g, alpha, grid)

# Use the operator
op(eta=eta, etasave=etasave, M=M, N=N, D=D, h=h, time=nt-2, dt=dt)


# In[25]:


# NBVAL_SKIP
snaps2video(etasave, "Modeling a 2D circular dam break problem")


# ## What we learned:
# 
# * How to solve the 2D Shallow Water Equation using the FTCS finite difference scheme.
# 
# * Propagation of (multiple) Tsunamis in an ocean with constant bathymetry.
# 
# * Tsunamis reaching shallow waters near a coast will slow down, their wavelength decrease, while their wave height increases.
# 
# * The influence of a seamount and roughness of the seafloor topoghraphy on the Tsunami wavefield.
# 
# * Tsunami modelling for engineering problems, like the collapse of dams.