Quasi-geostrophic shallow water.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "import numpy as np              # Numpy is the fundamental package for scientific computing in Python.\n",
    "import matplotlib.pyplot as pl  # Matplotlib is a scientific plotting package.\n",
    "\n",
    "# The statement below enforces the plots to be put into this notebook, instead of in their own windows.\n",
    "%matplotlib inline\n",
    "pl.rcParams.update({'font.size': 11})          # Set the standard font size of the plots to 11pt.\n",
    "pl.rcParams.update({'figure.figsize': [13,5]}) # Set the standard figure size."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "nx = 256\n",
    "ny = 128\n",
    "\n",
    "Lx = 1.2e7\n",
    "Ly = 0.6e7\n",
    "\n",
    "u0 = 0.\n",
    "\n",
    "g = 9.81\n",
    "H = 200.\n",
    "a = 6.37e6\n",
    "lat = 45.\n",
    "Omega = 7.292e-5\n",
    "f0 = 2.*Omega*np.sin(np.deg2rad(lat))\n",
    "beta = (2./a)*Omega*np.cos(np.deg2rad(lat))\n",
    "\n",
    "dx = Lx / nx\n",
    "dy = Ly / ny\n",
    "\n",
    "x = np.arange(0, Lx, dx)\n",
    "y = np.arange(0, Ly, dy)\n",
    "\n",
    "xx, yy = np.meshgrid(x, y)\n",
    "\n",
    "Psi0 = u0*yy\n",
    "\n",
    "h_dist = -1.\n",
    "L_dist = 2e6\n",
    "x_sigma = L_dist/6.\n",
    "h0 = np.exp(-( (xx-Lx/2.)**2 + (yy-Ly/2.)**2) / (2.*x_sigma**2))\n",
    "h0 = h_dist*h0/np.max(h0)\n",
    "Psi = Psi0 + g/f0*h0\n",
    "\n",
    "kkk = 2.*np.pi/Lx * np.arange(nx//2+1)\n",
    "lll = 2.*np.pi/Ly * np.arange(ny//2+1)\n",
    "llll = np.zeros(ny)\n",
    "llll[0:ny//2+1] = lll[:]\n",
    "for i in range(1,ny//2+1):\n",
    "    llll[-i] = -llll[i]\n",
    "\n",
    "dPsidx = np.fft.irfft( 1j*kkk[np.newaxis, :]*np.fft.rfft(Psi, axis=1), axis=1 )\n",
    "dPsidy = np.fft.irfft( 1j*lll[:, np.newaxis]*np.fft.rfft(Psi, axis=0), axis=0 )\n",
    "\n",
    "d2Psidx2 = np.fft.irfft( -kkk[np.newaxis, :]**2*np.fft.rfft(Psi, axis=1), axis=1 )\n",
    "d2Psidy2 = np.fft.irfft( -lll[:, np.newaxis]**2*np.fft.rfft(Psi, axis=0), axis=0 )\n",
    "\n",
    "q0 = d2Psidx2 + d2Psidy2 + beta*yy - f0**2/(g*H)*Psi"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "c0 = f0**2/(g*H)\n",
    "\n",
    "def invert_q(q_in):\n",
    "    q_fft = np.fft.rfft2(q_in - beta*yy)\n",
    "    Psi_fft = q_fft / ( - kkk[np.newaxis,:]**2  - llll[:,np.newaxis]**2 - c0 )\n",
    "    return np.fft.irfft2(Psi_fft)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "Psi = invert_q(q0)\n",
    "u = -np.gradient(Psi, dy, axis=0)\n",
    "v =  np.gradient(Psi, dx, axis=1)\n",
    "h = f0/g*Psi\n",
    "\n",
    "h_range   = np.linspace(-abs(h_dist), abs(h_dist), 40)\n",
    "\n",
    "pl.contourf(x/1000, y/1000, h, h_range, cmap=pl.cm.seismic)\n",
    "pl.colorbar()\n",
    "nq = 4\n",
    "pl.quiver(xx[::nq,::nq]/1000, yy[::nq,::nq]/1000, u[::nq,::nq], v[::nq,::nq],\n",
    "          width=0.0012, scale=3., pivot='mid', color='k')\n",
    "pl.xlabel('x (km)')\n",
    "pl.ylabel('y (km)')\n",
    "pl.tight_layout()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": false,
    "scrolled": false
   },
   "outputs": [],
   "source": [
    "dt = 1800.\n",
    "runtime = 6 * 86400.\n",
    "nsteps = round(runtime / dt)\n",
    "\n",
    "q = q0.copy()\n",
    "t = 10.\n",
    "for n in range(nsteps):\n",
    "    Psi = invert_q(q)\n",
    "    \n",
    "    u = -np.fft.irfft( 1j*lll[:, np.newaxis]*np.fft.rfft(Psi, axis=0), axis=0 )\n",
    "    v =  np.fft.irfft( 1j*kkk[np.newaxis, :]*np.fft.rfft(Psi, axis=1), axis=1 )\n",
    "    dqdt = -u*np.fft.irfft( 1j*kkk[np.newaxis, :]*np.fft.rfft(q, axis=1), axis=1 ) \\\n",
    "           -v*np.fft.irfft( 1j*lll[:, np.newaxis]*np.fft.rfft(q, axis=0), axis=0 )\n",
    "    q += dt*dqdt \n",
    "    t += dt"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "Psi = invert_q(q)\n",
    "u = -np.fft.irfft( 1j*lll[:, np.newaxis]*np.fft.rfft(Psi, axis=0), axis=0 )\n",
    "v =  np.fft.irfft( 1j*kkk[np.newaxis, :]*np.fft.rfft(Psi, axis=1), axis=1 )\n",
    "h = f0/g*Psi\n",
    "\n",
    "pl.figure()\n",
    "pl.contourf(x/1000, y/1000, h, h_range, cmap=pl.cm.seismic)\n",
    "pl.colorbar()\n",
    "#pl.contour(x/1000, y/1000, h0, h_range, colors='k', linewidths=2, linestyles='-')\n",
    "nq = 4\n",
    "pl.quiver(xx[::nq,::nq]/1000, yy[::nq,::nq]/1000, u[::nq,::nq], v[::nq,::nq],\n",
    "          width=0.0012, scale=3., pivot='mid', color='k')\n",
    "pl.xlabel('x (km)')\n",
    "pl.ylabel('y (km)')\n",
    "pl.title('h at t = {0:.1f} d, L$_R$ = {1:.0f} km'.format(t/86400., (g*H)**.5/f0 / 1000))\n",
    "pl.tight_layout()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "pl.figure()\n",
    "pl.subplot(121)\n",
    "pl.plot(x/1000, u[ny//2,:], 'b-', label='u')\n",
    "pl.plot(x/1000, v[ny//2,:], 'r-', label='v')\n",
    "pl.legend(loc=0, frameon=False);\n",
    "pl.xlabel('x (km)')\n",
    "pl.ylabel('u,v (m s$^{-1}$)')\n",
    "pl.subplot(122)\n",
    "pl.plot(x/1000, h [ny//2,:], 'r-')\n",
    "pl.plot(x/1000, h0[ny//2,:], 'k:')\n",
    "pl.xlabel('x (km)')\n",
    "pl.ylabel('h (m)')\n",
    "pl.tight_layout()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": []
  }
 ],
 "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.5.2"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 1
}