Project3_tsunami.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<a href=\"https://colab.research.google.com/github/johnryantaylor/QCES/blob/main/Project3_tsunami.ipynb\" target=\"_parent\"><img src=\"https://colab.research.google.com/assets/colab-badge.svg\" alt=\"Open In Colab\"/></a>"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Project 3: Tsunami\n",
    "In this project, we will use Dedalus to model the propagation of Tsunami waves. Tsunami typically have wavelengths larger than 100km and periods of about 1 hour. Since their wavelength is much larger than the depth of the ocean, Tsunami can be modelled using the shallow water equations. \n",
    "\n",
    "Here, we will use Dedalus to solve an initial value problem for the shallow water equations. We will generate a Tsunami using a localized disturbance to the water height. \n",
    "\n",
    "Student tasks are in [Section 1.3]( #section1pt3 ), [Section 2.3]( #section2pt3 ), and [Section 3.3]( #section3pt3 )."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Setup\n",
    "If you are using Google colab, run the script below to install Dedalus"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Set environment variables for best performance\n",
    "%env OMP_NUM_THREADS=1\n",
    "%env NUMEXPR_MAX_THREADS=1\n",
    "\n",
    "# Minimize logging output\n",
    "import logging\n",
    "logging.disable(logging.DEBUG)\n",
    "\n",
    "# Check if running on google colab\n",
    "import os\n",
    "using_google_colab = bool(os.getenv(\"COLAB_RELEASE_TAG\"))\n",
    "\n",
    "# Check for Dedalus\n",
    "try:\n",
    "    import dedalus.public as de\n",
    "    print(\"Dedalus already installed :)\")\n",
    "except:\n",
    "    print(\"Dedalus not installed yet.\")\n",
    "    if using_google_colab:\n",
    "        print(\"Installing for Google Colab.\")\n",
    "        print()\n",
    "        # Step 1: Install FFTW\n",
    "        !apt-get install libfftw3-dev\n",
    "        !apt-get install libfftw3-mpi-dev\n",
    "        # Step 2: Set paths for Dedalus installation\n",
    "        import os\n",
    "        os.environ['MPI_INCLUDE_PATH'] = \"/usr/lib/x86_64-linux-gnu/openmpi/include\"\n",
    "        os.environ['MPI_LIBRARY_PATH'] = \"/usr/lib/x86_64-linux-gnu\"\n",
    "        os.environ['FFTW_INCLUDE_PATH'] = \"/usr/include\"\n",
    "        os.environ['FFTW_LIBRARY_PATH'] = \"/usr/lib/x86_64-linux-gnu\"\n",
    "        # Step 3: Install Dedalus using pip\n",
    "        !pip3 install cython \"mpi4py<4.0\" numpy setuptools wheel\n",
    "        !CC=mpicc pip3 install --no-cache --no-build-isolation http://github.com/dedalusproject/dedalus/zipball/master/\n",
    "        !pip3 install -q ipympl\n",
    "        # Step 4: Check installation\n",
    "        print()\n",
    "        try:\n",
    "            import dedalus.public as de\n",
    "            print(\"Dedalus successfully installed :)\")\n",
    "        except:\n",
    "            print(\"Error installing Dedalus :(\")\n",
    "            raise\n",
    "    else:\n",
    "        print(\"See website for installation instructions:\")\n",
    "        print(\"https://dedalus-project.readthedocs.io/en/latest/pages/installation.html\")\n",
    "\n",
    "# Setup interactive matplotlib\n",
    "if using_google_colab:\n",
    "    from google.colab import output\n",
    "    output.enable_custom_widget_manager()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 1.  1D shallow water equations\n",
    "Let's start by solving the shallow water equations in 1D with periodic boundary conditions in our spatial variable, $x$. The script below solves the following equations:\n",
    "\n",
    "$$ \\partial_t u + g \\partial_x(h) = - u\\partial_x u$$\n",
    "$$ \\partial_t h + \\partial_x(Hu) = -\\partial_x(hu)$$\n",
    "\n",
    "In these equations, $g$ is the constant gravitational acceleration, $u$ is the velocity in the $x$ direction, $H$ is the resting height, and $h$ is the height of the free surface relative to $H$. Note that the terms on the right hand side are nonlinear in our dependent variables, and we will use an explicit timestepping method for these terms. Recall that Dedalus uses explicit timestepping for terms on the right hand side of the inputted equation.\n",
    "\n",
    "Because we are interested in an application to Tsunami, we will solve the equations using dimensional variables. It is convenient to define simulation units at the start of the script so that we can be explicit about the units.\n",
    "\n",
    "We will start with the following initial conditions:\n",
    "$$ h=e^{-(x-x_1)^2/w_1^2} $$\n",
    "$$ u=0 $$\n",
    "where $x_1$ is the center and $w_1=100$ km is the width of the Gaussian perturbation. \n",
    "\n",
    "To see how the waves change as they approach shallow water, we will set the resting water depth to include a localized seamount using the form:\n",
    "$$ H=H_0-(H_0-H_{min})e^{(-(x-x_2)^2/w_2^2)} $$\n",
    "where $x_2$ is the location of the seamount, $H_0=4$ km is the depth of the ocean away from the seamount, and $H_{min}$ is the minimum water depth above the seamount. Note that since $H\\ll w$, the shallow water approximation is appropriate here."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### 1.1  Dedalus script to solve the 1D shallow water equations\n",
    "\n",
    "The script below solves the 1D shallow water equations with the initial conditions given above."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Import the required packages\n",
    "import numpy as np\n",
    "import dedalus.public as dedalus\n",
    "import logging\n",
    "import matplotlib.pyplot as plt\n",
    "\n",
    "logger = logging.getLogger(__name__)\n",
    "\n",
    "# Define the simulation units. This is optional, but will help us keep track of units\n",
    "meter = 1\n",
    "kilometer = meter * 1e3\n",
    "second = 1\n",
    "minute = 60\n",
    "hour = 3600\n",
    "day = hour * 24\n",
    "\n",
    "# Specify the simulation parameters\n",
    "Lx = 1000 * kilometer # The domain size\n",
    "Nx = 512 # The number of gridpoints in x\n",
    "g = 9.81 * meter / second**2 # Gravitational acceleration\n",
    "H0 = 4 * kilometer # Reference ocean depth\n",
    "w1 = 20 * kilometer # Width of the initial perturbation\n",
    "x1 = 200 * kilometer # Center of the initial perturbation\n",
    "x2 = 500 * kilometer # Center of the seamount\n",
    "w2 = 100 * kilometer # Width of the seamount\n",
    "H_min = 100 * meter # minimum depth above the seamount\n",
    "\n",
    "timestep = 60 * second # Specify the timestep\n",
    "stop_sim_time = 100 * minute # Specify the stopping time\n",
    "\n",
    "timestepper = dedalus.RK222\n",
    "max_timestep = 0.125\n",
    "\n",
    "# Create an x-coordinate, a corresponding basis, and a distributor\n",
    "xcoord = dedalus.Coordinate('x')\n",
    "dist = dedalus.Distributor(xcoord, dtype=np.float64) # No mesh for serial / automatic parallelization\n",
    "x_basis = dedalus.RealFourier(xcoord, size=Nx, bounds=(0, Lx))\n",
    "x = dist.local_grid(x_basis)\n",
    "\n",
    "# Specify the Fields and set the initial conditions\n",
    "u = dist.Field(name='u', bases=x_basis)\n",
    "h = dist.Field(name='h', bases=x_basis)\n",
    "H = dist.Field(name='H', bases=x_basis)\n",
    "h.change_scales(1) \n",
    "H.change_scales(1) \n",
    "u.change_scales(1)\n",
    "H['g'] = H0-(H0-H_min)*np.exp(-(x-x2)**2.0/w2**2.0)\n",
    "h['g'] = np.exp(-(x-x1)**2.0/w1**2.0)  # Initialize the height with a Gaussian function with width set above\n",
    "u['g'] = 0\n",
    "\n",
    "# Operators\n",
    "dx = lambda a: dedalus.Differentiate(a, xcoord)\n",
    "\n",
    "# Problem\n",
    "problem = dedalus.IVP([u, h], namespace=locals())\n",
    "problem.add_equation(\"dt(u) + g*dx(h) = - u*dx(u)\")\n",
    "problem.add_equation(\"dt(h) + dx(H*u) = - dx(h*u)\")\n",
    "\n",
    "# Solver\n",
    "solver = problem.build_solver(dedalus.RK222)\n",
    "solver.stop_sim_time = stop_sim_time\n",
    "\n",
    "# Analysis\n",
    "snapshots = solver.evaluator.add_file_handler('snapshots', sim_dt=1*hour, max_writes=10)\n",
    "snapshots.add_task(h, name='h')\n",
    "snapshots.add_task(u, name='u')\n",
    "\n",
    "h.change_scales(1)\n",
    "h_save = [np.copy(h['g'])]; t_save = [solver.sim_time]; # Save the initial condition and the initial time\n",
    "# Main loop\n",
    "try:\n",
    "    logger.info('Starting main loop')\n",
    "    while solver.proceed:\n",
    "        solver.step(timestep)\n",
    "        if (solver.iteration-1) % 1 == 0:\n",
    "            h.change_scales(1)\n",
    "            h_save.append(np.copy(h['g']))\n",
    "            t_save.append(solver.sim_time)\n",
    "except:\n",
    "    logger.error('Exception raised, triggering end of main loop.')\n",
    "    raise\n",
    "finally:\n",
    "    solver.log_stats()\n",
    "\n",
    "h_save=np.array(h_save) # convert h_save to an np.array type to allow data slicing"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### 1.2  Plotting\n",
    "The code below will make a movie of the height."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# create a figure and axes\n",
    "fig = plt.figure(figsize=(12,6))\n",
    "\n",
    "# Specify the axis limits and labels\n",
    "plt.xlim((0, Lx / (1*kilometer)))            \n",
    "plt.ylim((-1, 1))\n",
    "plt.xlabel('x (km)')\n",
    "plt.ylabel('Height (m)')\n",
    "\n",
    "# Create objects that will change during the animation\n",
    "txt_title = plt.title('')\n",
    "height, = plt.plot([], [], 'b', lw=2)    \n",
    "\n",
    "# Define an animation function\n",
    "def drawframe(n):\n",
    "    height.set_data(x/(1*kilometer), h_save[n,:])\n",
    "    txt_title.set_text('Time = {0:.2f} hours'.format(t_save[n]/(1*hour)))\n",
    "    return (height,)\n",
    "\n",
    "from matplotlib import animation\n",
    "# blit=True re-draws only the parts that have changed.\n",
    "anim = animation.FuncAnimation(fig, drawframe, frames=len(t_save), interval=40, blit=True)\n",
    "\n",
    "from IPython.display import HTML\n",
    "HTML(anim.to_html5_video())\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### 1.3 Student investigation\n",
    "<a id='section1pt3'></a>\n",
    "In the cells below, provide an expression for the expected wave speed and state any assumptions that are necessary to arrive at this. Then plot the wave speed from the output above and compare this with the expected value. "
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "ADD EXPRESSION FOR EXPECTED WAVE SPEED AND EXPLANATION HERE"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# CREATE FIGURES HERE"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "COMMENT ON RESULTS"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 2.  2D Shallow water equations\n",
    "\n",
    "Now, let's see how things change if we move to two dimensions. The script below solves the 2D shallow water equations:\n",
    "\n",
    "$$ \\partial_t u + g \\partial_x(h) = - u*\\partial_x u - v*\\partial_y u$$\n",
    "$$ \\partial_t u + g \\partial_y(h) = - u*\\partial_x v - v*\\partial_y v$$\n",
    "$$ \\partial_t h + H (\\partial_x(u) + \\partial_y(v)) = -\\partial_x(h*u) - \\partial_y(h*v)$$\n",
    "\n",
    "We will solve these equations in a square domain with periodic boundary conditions. We will again initialize with a Gaussian height disturbance and with no water motion, but this time in 2D:\n",
    "\n",
    "$$ h=e^{-(x^2+y^2)/w^2} $$\n",
    "$$ u=0 $$\n",
    "$$ v=0 $$"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### 2.1  Dedalus script to solve the 2D shallow water equations"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Import the required packages\n",
    "import numpy as np\n",
    "import dedalus.public as dedalus\n",
    "import logging\n",
    "import matplotlib.pyplot as plt\n",
    "\n",
    "logger = logging.getLogger(__name__)\n",
    "\n",
    "# Define the simulation units. This is optional, but will help us keep track of units\n",
    "meter = 1\n",
    "kilometer = meter * 1e3\n",
    "second = 1\n",
    "minute = 60\n",
    "hour = 3600\n",
    "day = hour * 24\n",
    "\n",
    "# Specify the simulation parameters\n",
    "Lx = 1000 * kilometer # The domain size in x\n",
    "Ly = 1000 * kilometer # The domain size in y\n",
    "Nx = 64 # The number of gridpoints in x\n",
    "Ny = 64 # The number of gridpoints in y\n",
    "g = 9.81 * meter / second**2 # Gravitational acceleration\n",
    "H = 4 * kilometer # Reference ocean depth\n",
    "width = 100 * kilometer # Set the width of the initial Gaussian\n",
    "\n",
    "timestep = 60 * second # Specify the timestep\n",
    "stop_sim_time = 30 * minute # Specify the stopping time\n",
    "\n",
    "timestepper = dedalus.RK222\n",
    "max_timestep = 0.125\n",
    "\n",
    "# Bases\n",
    "coords = dedalus.CartesianCoordinates('x', 'y')\n",
    "dist = dedalus.Distributor(coords, dtype=np.float64)\n",
    "xbasis = dedalus.RealFourier(coords['x'], size=Nx, bounds=(-Lx/2, Lx/2), dealias=3/2)\n",
    "ybasis = dedalus.RealFourier(coords['y'], size=Ny, bounds=(-Ly/2, Ly/2), dealias=3/2)\n",
    "\n",
    "x, y = dist.local_grids(xbasis, ybasis)\n",
    "\n",
    "# Fields\n",
    "u = dist.Field(name='u', bases=(xbasis, ybasis))\n",
    "v = dist.Field(name='v', bases=(xbasis, ybasis))\n",
    "h = dist.Field(name='h', bases=(xbasis, ybasis))\n",
    "\n",
    "# Initial conditions\n",
    "h.change_scales(1) # Set values on regular grid\n",
    "u.change_scales(1)\n",
    "v.change_scales(1)\n",
    "h['g'] = np.exp(-(x**2+y**2)/width**2)  # Initialize the dependent variable c with a Gaussian function (note ** is the exponentiation operator)\n",
    "u['g'] = 0\n",
    "v['g'] = 0\n",
    "\n",
    "# Operators\n",
    "dx = lambda a: dedalus.Differentiate(a, coords['x'])\n",
    "dy = lambda a: dedalus.Differentiate(a, coords['y'])\n",
    "\n",
    "# Problem\n",
    "problem = dedalus.IVP([u, v, h], namespace=locals())\n",
    "problem.add_equation(\"dt(u) + g*dx(h) = - u*dx(u) - v*dy(u)\")\n",
    "problem.add_equation(\"dt(v) + g*dy(h) = - u*dx(v) - v*dy(v)\")\n",
    "problem.add_equation(\"dt(h) + H*(dx(u) + dy(v)) = - dx(h*u) - dy(h*v)\")\n",
    "\n",
    "# Solver\n",
    "solver = problem.build_solver(dedalus.RK222)\n",
    "solver.stop_sim_time = stop_sim_time\n",
    "\n",
    "# Analysis\n",
    "snapshots = solver.evaluator.add_file_handler('snapshots', sim_dt=1*hour, max_writes=10)\n",
    "snapshots.add_task(h, name='h')\n",
    "snapshots.add_task(u, name='u')\n",
    "snapshots.add_task(v, name='v')\n",
    "\n",
    "h.change_scales(1)\n",
    "h_save = [np.copy(h['g'])]; t_save = [solver.sim_time]; # Save the initial condition and the initial time\n",
    "# Main loop\n",
    "try:\n",
    "    logger.info('Starting main loop')\n",
    "    while solver.proceed:\n",
    "        solver.step(timestep)\n",
    "        if (solver.iteration-1) % 1 == 0:\n",
    "            h.change_scales(1)\n",
    "            h_save.append(np.copy(h['g']))\n",
    "            t_save.append(solver.sim_time)\n",
    "except:\n",
    "    logger.error('Exception raised, triggering end of main loop.')\n",
    "    raise\n",
    "finally:\n",
    "    solver.log_stats()\n",
    "\n",
    "h_save=np.array(h_save) # convert h_save to an np.array type to allow data slicing"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### 2.2  Plotting\n",
    "The code below will make a movie of the height as a function of x and y"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# create a figure and axes\n",
    "fig = plt.figure(figsize=(6,6))\n",
    "\n",
    "# Create arrays with the x and y coordinates for plotting using pcolor\n",
    "from dedalus.extras import plot_tools\n",
    "(X, Y) = plot_tools.quad_mesh(np.squeeze(x), np.squeeze(y))  \n",
    "\n",
    "# Define an animation function\n",
    "def drawframe(n):\n",
    "    plt.clf # clear the figure before plotting the next frame\n",
    "    height = plt.pcolor(X/(1*kilometer), Y/(1*kilometer), h_save[n,:,:], vmin=-1, vmax=1)\n",
    "    # Specify the axis limits and labels\n",
    "    plt.xlim((-Lx/2 / (1*kilometer),Lx/2 / (1*kilometer)))   \n",
    "    plt.ylim((-Ly/2 / (1*kilometer),Ly/2 / (1*kilometer)))            \n",
    "    plt.xlabel('x (km)')\n",
    "    plt.ylabel('y (m)')\n",
    "    plt.title('Surface height (m), time = {0:.2f} hours'.format(t_save[n]/(1*hour)))\n",
    "    return (height,)\n",
    "\n",
    "from matplotlib import animation\n",
    "# blit=True re-draws only the parts that have changed.\n",
    "anim = animation.FuncAnimation(fig, drawframe, frames=len(t_save), interval=40, blit=True)\n",
    "\n",
    "from IPython.display import HTML\n",
    "HTML(anim.to_html5_video())"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### 2.3 Student investigation\n",
    "<a id='section2pt3'></a>\n",
    "Notice how the amplitude of the Tsunami wave decreases with distance from the source location. In the cell below, make a prediction for how the wave amplitude will vary as a function of distance from the origin (x=0, y=0) using conservation of energy. Then, make a plot or plots comparing your prediction for the wave amplitude compared with the results from the simulation. \n",
    "\n",
    "Try out several different possible initial conditions for the surface height, $h$. How does the shape of the resulting waves change?"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "ADD PREDICTION OF WAVE AMPLITUDE HERE (a few sentences and/or equations)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# ADD CODE TO MAKE PLOTS HERE"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "ADD DISCUSSION HERE (~1 paragraph)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 3. Shallow water equations on a sphere\n",
    "\n",
    "Of course the Earth isn't a flat plane. A great strength of Dedalus is its flexibility to solve equations in various geometries. The script below solves the 2D shallow water equations on a spherical shell, where the azimuth and colatitude are the independent variables. The colatitude is the complementary angle to the latitude such that the colatitude is $0^\\circ$ at the North Pole, $90^\\circ$ at the Equator, and $180^\\circ$ at the South Pole. The spherical geometry also allows us to take into account the Earth's rotation in a realistic way. \n",
    "\n",
    "Specifically, we will solve the following equations:\n",
    "$$ \\partial_t \\mathbf{u} + g \\nabla h + 2 \\Omega \\hat{\\mathbf{z}} \\times \\mathbf{u} = -\\mathbf{u} \\cdot \\nabla \\mathbf{u} $$\n",
    "$$ \\partial_t + H \\nabla \\cdot \\mathbf{u} = -\\nabla \\cdot (h \\mathbf{u}) $$\n",
    "\n",
    "where $\\mathbf{u}$ is the velocity vector, $\\Omega$ is the Earth's angular velocity (in radians per second), and $\\hat{\\mathbf{z}}$ is a unit vector aligned with the North Pole. Note that we have written the equations above in vector form. Dedalus includes vector calculus operators that make it very natural to solve equations like this without worrying about the particular form of the vector operators in a curvilinear coordinate system."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### 3.1  Dedalus script for 2D shallow water equations on a sphere"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Import the required packages\n",
    "import numpy as np\n",
    "import dedalus.public as dedalus\n",
    "import logging\n",
    "import matplotlib.pyplot as plt\n",
    "\n",
    "logger = logging.getLogger(__name__)\n",
    "\n",
    "# Define the simulation units. This is optional, but will help us keep track of units\n",
    "meter = 1\n",
    "kilometer = meter * 1e3\n",
    "second = 1\n",
    "minute = 60\n",
    "hour = 3600\n",
    "day = hour * 24\n",
    "\n",
    "# Parameters\n",
    "Nphi = 128 # Number of gridpoints in phi (longitude)\n",
    "Ntheta = 64 # Number of gridpoints in theta (colatitude)\n",
    "R = 6.37122e3 * kilometer\n",
    "Omega = 2 * np.pi / (1 * day) # Earth's angular velocity in radians per second\n",
    "g = 9.80616 * meter / second**2 # Gravitational acceleration\n",
    "H = 4 * kilometer # Depth of the ocean\n",
    "width = 2*np.pi*10/360 # Width of the initial Gaussian in radians\n",
    "phi0 = np.pi # Longitude of the initial Gaussian height perturbation in radians\n",
    "theta0 = np.pi/4 # Colatitude of the initial Gaussian height perturbation in radians\n",
    "\n",
    "timestep = 600 * second\n",
    "stop_sim_time = 10 * day\n",
    "\n",
    "# Bases\n",
    "coords = dedalus.S2Coordinates('phi', 'theta')\n",
    "dist = dedalus.Distributor(coords, dtype=np.float64)\n",
    "basis = dedalus.SphereBasis(coords, (Nphi, Ntheta), radius=R, dealias=3/2, dtype=np.float64)\n",
    "\n",
    "# Define the coordinates\n",
    "phi, theta = dist.local_grids(basis) # Longitude and colatitude\n",
    "lat = 360*(np.pi / 2 - theta)/(2*np.pi)  # latitude in degrees\n",
    "\n",
    "# Fields\n",
    "u = dist.VectorField(coords, name='u', bases=basis)\n",
    "h = dist.Field(name='h', bases=basis)\n",
    "\n",
    "# Operators\n",
    "# First, create an operator to perform the cross product: \"z x \", where z is the local vertical unit vector\n",
    "# The \"skew\" operator performs a 90 degree rotation of a vector field,\n",
    "# so that skew([u,v]) = [-v, u].  MulCosine multiplies by the cosine of the latitude\n",
    "zcross = lambda A: dedalus.MulCosine(dedalus.skew(A))\n",
    "\n",
    "# Set the initial conditions\n",
    "h.change_scales(1) # Set values on regular grid\n",
    "u.change_scales(1)\n",
    "h['g'] = np.exp(-((phi-phi0)**2+(theta-theta0)**2)/width**2)  # Initialize the dependent variable c with a Gaussian function (note ** is the exponentiation operator)\n",
    "# Set the initial velocity to zero\n",
    "u['g'][0] = 0 # the x-component of the velocity\n",
    "u['g'][1] = 0 # the y-component of the velocity\n",
    "\n",
    "# Problem\n",
    "problem = dedalus.IVP([u, h], namespace=locals())\n",
    "problem.add_equation(\"dt(u) + g*grad(h) + 2*Omega*zcross(u) = - u@grad(u)\")\n",
    "problem.add_equation(\"dt(h) + H*div(u) = - div(h*u)\")\n",
    "\n",
    "# Solver\n",
    "solver = problem.build_solver(dedalus.RK222)\n",
    "solver.stop_sim_time = stop_sim_time\n",
    "\n",
    "h.change_scales(1)\n",
    "h_save = [np.copy(h['g'])]; t_save = [solver.sim_time]; # Save the initial condition and the initial time\n",
    "# Main loop\n",
    "try:\n",
    "    logger.info('Starting main loop')\n",
    "    while solver.proceed:\n",
    "        solver.step(timestep)\n",
    "        if (solver.iteration-1) % 10 == 0:\n",
    "            logger.info('Iteration=%i, Time=%e, dt=%e' %(solver.iteration, solver.sim_time, timestep))\n",
    "            h.change_scales(1)\n",
    "            h_save.append(np.copy(h['g']))\n",
    "            t_save.append(solver.sim_time)\n",
    "except:\n",
    "    logger.error('Exception raised, triggering end of main loop.')\n",
    "    raise\n",
    "finally:\n",
    "    solver.log_stats()\n",
    "\n",
    "h_save = np.array(h_save)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### 3.2  Plotting\n",
    "The code below will make a movie of the height as a function of x and y"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# create a figure and axes\n",
    "fig = plt.figure(figsize=(12,6))\n",
    "\n",
    "# Create arrays with the latitude and longitude coordinates for plotting using pcolor\n",
    "from dedalus.extras import plot_tools\n",
    "(X, Y) = plot_tools.quad_mesh(np.squeeze(np.pi/2 - theta)*360/2/np.pi, np.squeeze(phi)*360/2/np.pi)  \n",
    "\n",
    "# Define an animation function\n",
    "def drawframe(n):\n",
    "    plt.clf()\n",
    "    height = plt.pcolor(Y, X, h_save[n,:,:], vmin=-0.1, vmax=0.1)\n",
    "    # Specify the axis limits and labels         \n",
    "    plt.xlabel('Longitude')\n",
    "    plt.ylabel('Latitude')\n",
    "    plt.title('Surface height (m), time = {0:.2f} hours'.format(t_save[n]/(1*hour)))\n",
    "    return (height,)\n",
    "\n",
    "from matplotlib import animation\n",
    "# blit=True re-draws only the parts that have changed.\n",
    "anim = animation.FuncAnimation(fig, drawframe, frames=len(t_save), interval=40, blit=True)\n",
    "\n",
    "from IPython.display import HTML\n",
    "HTML(anim.to_html5_video())"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### 3.3 Student investigation\n",
    "<a id='section3pt3'></a>\n",
    "Create a virtual tide gauge by sampling the height at a point at a latitude of 45ºN and longitude of 90ºW and plot the timeseries of the height at this point in the cells below. You should modifty the code to save the height more often so that you can clearly see the shape of the wave. Notice that as the waves propagate around the Earth, their shape changes. Discuss physical mechanisms that might cause the shape of the wave to change as the wave propagates."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "ADD PLOTS AND DISCUSSION HERE"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "In the simulation above, the Tsunami wave travels around the planet many times. In reality, the waves will partially dissipate when they reach shorelines. In the spherical shallow water equation code above (or in a new code block below), add a forcing term to the shallow water equations that damps the waves for longitudes larger than $\\pm 90^\\circ$. Specifically, add a forcing function to the right hand side of the height equation of the following form:\n",
    "\n",
    "$$ - (1/\\tau) f(\\phi) h$$\n",
    "\n",
    "where\n",
    "\n",
    "$$ f(\\phi) = \\frac{1}{2}\\left(1+tanh\\left(\\frac{\\phi-3\\pi/2}{\\pi/10}\\right)\\right) + \\frac{1}{2}\\left(1-tanh\\left(\\frac{\\phi-\\pi/2}{\\pi/10}\\right)\\right) $$\n",
    "\n",
    "and $\\phi$ is the co-latitude in radians and $\\tau$ is a damping timescale. A timescale of $\\tau=0.1$ days should be sufficient to damp the waves.\n",
    "\n",
    "With this damping term active, much more slowly propagating waves should be visible. These are shallow water Rossby waves. Estimate the east/west phase speed of these Rossby waves and compare them with the phase speed that you can infer from the simulation by plotting the height, $h$, as a function of longitude and time. Discuss the spatial structure of the Rossby waves and any discrepancies that you see between your estimated and measured phase speeds."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# ADD CODE HERE TO SOLVE THE SHALLOW WATER EQUATIONS WITH DAMPING.\n",
    "# YOU CAN START BY COPYING THE CODE ABOVE AND PASTING IT HERE"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 4. Optional further investigations\n",
    "\n",
    "* In the example of the shallow water equations on a sphere, introduce bathymetry by making $H$ a function of latitude and longitude. Make sure that you modify the equations in Dedalus appropriately. How do the waves behave when they encounter shallow water (e.g. a localized seamount)?\n",
    "\n",
    "* It is quite difficult to simulate changing sea level (sometimes called wetting and drying), which is why the examples above do not include land. However, we can estimate the distance that a Tsunami will propagate using physical and dimensional arguments. Consider, for example, a flat, sloping beach. In the absence of energy dissipation, estimate the distance that a wave will propagate onshore. How could you modify this estimate to include energy dissipation (or drag)?"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "dedalus3",
   "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.12.2"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}