{
 "cells": [
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Differentiable Shallow Water Equations\n",
    "\n",
    "We present a differentiable SWE solver, based on `torch-harmonics`. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "import torch\n",
    "import torch.nn as nn\n",
    "import numpy as np\n",
    "import matplotlib.pyplot as plt\n",
    "import matplotlib.animation as animation\n",
    "\n",
    "from math import ceil, floor\n",
    "\n",
    "import sys\n",
    "\n",
    "from torch_harmonics.sht import *\n",
    "from torch_harmonics.examples import ShallowWaterSolver"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [],
   "source": [
    "device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We define a shallow water solver class in `shallow_water_equations.py`"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [],
   "source": [
    "# initialize parameters:\n",
    "nlat = 512\n",
    "nlon = 2*nlat\n",
    "lmax = ceil(128)\n",
    "mmax = lmax\n",
    "# timestepping\n",
    "dt = 75\n",
    "maxiter = 12*int(86400/dt)\n",
    "\n",
    "# initialize solver class\n",
    "swe_solver = ShallowWaterSolver(nlat, nlon, dt, lmax=lmax, mmax=mmax).to(device)\n",
    "\n",
    "lons = swe_solver.lons\n",
    "lats = swe_solver.lats\n",
    "\n",
    "jj, ii = torch.triu_indices(lmax, mmax)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "/home/bbonev/.local/lib/python3.10/site-packages/torch/functional.py:504: UserWarning: torch.meshgrid: in an upcoming release, it will be required to pass the indexing argument. (Triggered internally at ../aten/src/ATen/native/TensorShape.cpp:3190.)\n",
      "  return _VF.meshgrid(tensors, **kwargs)  # type: ignore[attr-defined]\n"
     ]
    }
   ],
   "source": [
    "uspec0 = swe_solver.galewsky_initial_condition()    "
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We are now ready to run the simulation. To perform integration in time, we will use third-order Adams-Bashforth. As we are currently not interested in gradients, we can wrap the function in `torch.inference_mode()`."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "t=0.00 hours\n",
      "t=1.04 hours\n",
      "t=2.08 hours\n",
      "t=3.12 hours\n",
      "t=4.17 hours\n",
      "t=5.21 hours\n",
      "t=6.25 hours\n",
      "t=7.29 hours\n",
      "t=8.33 hours\n",
      "t=9.38 hours\n",
      "t=10.42 hours\n",
      "t=11.46 hours\n",
      "t=12.50 hours\n",
      "t=13.54 hours\n",
      "t=14.58 hours\n",
      "t=15.62 hours\n",
      "t=16.67 hours\n",
      "t=17.71 hours\n",
      "t=18.75 hours\n",
      "t=19.79 hours\n",
      "t=20.83 hours\n",
      "t=21.88 hours\n",
      "t=22.92 hours\n",
      "t=23.96 hours\n",
      "t=25.00 hours\n",
      "t=26.04 hours\n",
      "t=27.08 hours\n",
      "t=28.12 hours\n",
      "t=29.17 hours\n",
      "t=30.21 hours\n",
      "t=31.25 hours\n",
      "t=32.29 hours\n",
      "t=33.33 hours\n",
      "t=34.38 hours\n",
      "t=35.42 hours\n",
      "t=36.46 hours\n",
      "t=37.50 hours\n",
      "t=38.54 hours\n",
      "t=39.58 hours\n",
      "t=40.62 hours\n",
      "t=41.67 hours\n",
      "t=42.71 hours\n",
      "t=43.75 hours\n",
      "t=44.79 hours\n",
      "t=45.83 hours\n",
      "t=46.88 hours\n",
      "t=47.92 hours\n",
      "t=48.96 hours\n",
      "t=50.00 hours\n",
      "t=51.04 hours\n",
      "t=52.08 hours\n",
      "t=53.12 hours\n",
      "t=54.17 hours\n",
      "t=55.21 hours\n",
      "t=56.25 hours\n",
      "t=57.29 hours\n",
      "t=58.33 hours\n",
      "t=59.38 hours\n",
      "t=60.42 hours\n",
      "t=61.46 hours\n",
      "t=62.50 hours\n",
      "t=63.54 hours\n",
      "t=64.58 hours\n",
      "t=65.62 hours\n",
      "t=66.67 hours\n",
      "t=67.71 hours\n",
      "t=68.75 hours\n",
      "t=69.79 hours\n",
      "t=70.83 hours\n",
      "t=71.88 hours\n",
      "t=72.92 hours\n",
      "t=73.96 hours\n",
      "t=75.00 hours\n",
      "t=76.04 hours\n",
      "t=77.08 hours\n",
      "t=78.12 hours\n",
      "t=79.17 hours\n",
      "t=80.21 hours\n",
      "t=81.25 hours\n",
      "t=82.29 hours\n",
      "t=83.33 hours\n",
      "t=84.38 hours\n",
      "t=85.42 hours\n",
      "t=86.46 hours\n",
      "t=87.50 hours\n",
      "t=88.54 hours\n",
      "t=89.58 hours\n",
      "t=90.62 hours\n",
      "t=91.67 hours\n",
      "t=92.71 hours\n",
      "t=93.75 hours\n",
      "t=94.79 hours\n",
      "t=95.83 hours\n",
      "t=96.88 hours\n",
      "t=97.92 hours\n",
      "t=98.96 hours\n",
      "t=100.00 hours\n",
      "t=101.04 hours\n",
      "t=102.08 hours\n",
      "t=103.12 hours\n",
      "t=104.17 hours\n",
      "t=105.21 hours\n",
      "t=106.25 hours\n",
      "t=107.29 hours\n",
      "t=108.33 hours\n",
      "t=109.38 hours\n",
      "t=110.42 hours\n",
      "t=111.46 hours\n",
      "t=112.50 hours\n",
      "t=113.54 hours\n",
      "t=114.58 hours\n",
      "t=115.62 hours\n",
      "t=116.67 hours\n",
      "t=117.71 hours\n",
      "t=118.75 hours\n",
      "t=119.79 hours\n",
      "t=120.83 hours\n",
      "t=121.88 hours\n",
      "t=122.92 hours\n",
      "t=123.96 hours\n",
      "t=125.00 hours\n",
      "t=126.04 hours\n",
      "t=127.08 hours\n",
      "t=128.12 hours\n",
      "t=129.17 hours\n",
      "t=130.21 hours\n",
      "t=131.25 hours\n",
      "t=132.29 hours\n",
      "t=133.33 hours\n",
      "t=134.38 hours\n",
      "t=135.42 hours\n",
      "t=136.46 hours\n",
      "t=137.50 hours\n",
      "t=138.54 hours\n",
      "t=139.58 hours\n",
      "t=140.62 hours\n",
      "t=141.67 hours\n",
      "t=142.71 hours\n",
      "t=143.75 hours\n",
      "t=144.79 hours\n",
      "t=145.83 hours\n",
      "t=146.88 hours\n",
      "t=147.92 hours\n",
      "t=148.96 hours\n",
      "t=150.00 hours\n",
      "t=151.04 hours\n",
      "t=152.08 hours\n",
      "t=153.12 hours\n",
      "t=154.17 hours\n",
      "t=155.21 hours\n",
      "t=156.25 hours\n",
      "t=157.29 hours\n",
      "t=158.33 hours\n",
      "t=159.38 hours\n",
      "t=160.42 hours\n",
      "t=161.46 hours\n",
      "t=162.50 hours\n",
      "t=163.54 hours\n",
      "t=164.58 hours\n",
      "t=165.62 hours\n",
      "t=166.67 hours\n",
      "t=167.71 hours\n",
      "t=168.75 hours\n",
      "t=169.79 hours\n",
      "t=170.83 hours\n",
      "t=171.88 hours\n",
      "t=172.92 hours\n",
      "t=173.96 hours\n",
      "t=175.00 hours\n",
      "t=176.04 hours\n",
      "t=177.08 hours\n",
      "t=178.12 hours\n",
      "t=179.17 hours\n",
      "t=180.21 hours\n",
      "t=181.25 hours\n",
      "t=182.29 hours\n",
      "t=183.33 hours\n",
      "t=184.38 hours\n",
      "t=185.42 hours\n",
      "t=186.46 hours\n",
      "t=187.50 hours\n",
      "t=188.54 hours\n",
      "t=189.58 hours\n",
      "t=190.62 hours\n",
      "t=191.67 hours\n",
      "t=192.71 hours\n",
      "t=193.75 hours\n",
      "t=194.79 hours\n",
      "t=195.83 hours\n",
      "t=196.88 hours\n",
      "t=197.92 hours\n",
      "t=198.96 hours\n",
      "t=200.00 hours\n",
      "t=201.04 hours\n",
      "t=202.08 hours\n",
      "t=203.12 hours\n",
      "t=204.17 hours\n",
      "t=205.21 hours\n",
      "t=206.25 hours\n",
      "t=207.29 hours\n",
      "t=208.33 hours\n",
      "t=209.38 hours\n",
      "t=210.42 hours\n",
      "t=211.46 hours\n",
      "t=212.50 hours\n",
      "t=213.54 hours\n",
      "t=214.58 hours\n",
      "t=215.62 hours\n",
      "t=216.67 hours\n",
      "t=217.71 hours\n",
      "t=218.75 hours\n",
      "t=219.79 hours\n",
      "t=220.83 hours\n",
      "t=221.88 hours\n",
      "t=222.92 hours\n",
      "t=223.96 hours\n",
      "t=225.00 hours\n",
      "t=226.04 hours\n",
      "t=227.08 hours\n",
      "t=228.12 hours\n",
      "t=229.17 hours\n",
      "t=230.21 hours\n",
      "t=231.25 hours\n",
      "t=232.29 hours\n",
      "t=233.33 hours\n",
      "t=234.38 hours\n",
      "t=235.42 hours\n",
      "t=236.46 hours\n",
      "t=237.50 hours\n",
      "t=238.54 hours\n",
      "t=239.58 hours\n",
      "t=240.62 hours\n",
      "t=241.67 hours\n",
      "t=242.71 hours\n",
      "t=243.75 hours\n",
      "t=244.79 hours\n",
      "t=245.83 hours\n",
      "t=246.88 hours\n",
      "t=247.92 hours\n",
      "t=248.96 hours\n",
      "t=250.00 hours\n",
      "t=251.04 hours\n",
      "t=252.08 hours\n",
      "t=253.12 hours\n",
      "t=254.17 hours\n",
      "t=255.21 hours\n",
      "t=256.25 hours\n",
      "t=257.29 hours\n",
      "t=258.33 hours\n",
      "t=259.38 hours\n",
      "t=260.42 hours\n",
      "t=261.46 hours\n",
      "t=262.50 hours\n",
      "t=263.54 hours\n",
      "t=264.58 hours\n",
      "t=265.62 hours\n",
      "t=266.67 hours\n",
      "t=267.71 hours\n",
      "t=268.75 hours\n",
      "t=269.79 hours\n",
      "t=270.83 hours\n",
      "t=271.88 hours\n",
      "t=272.92 hours\n",
      "t=273.96 hours\n",
      "t=275.00 hours\n",
      "t=276.04 hours\n",
      "t=277.08 hours\n",
      "t=278.12 hours\n",
      "t=279.17 hours\n",
      "t=280.21 hours\n",
      "t=281.25 hours\n",
      "t=282.29 hours\n",
      "t=283.33 hours\n",
      "t=284.38 hours\n",
      "t=285.42 hours\n",
      "t=286.46 hours\n",
      "t=287.50 hours\n"
     ]
    }
   ],
   "source": [
    "dudtspec = torch.zeros(3, 3, swe_solver.lmax, swe_solver.mmax, dtype=torch.complex128, device=device)\n",
    "inew = 0\n",
    "inow = 1\n",
    "iold = 2\n",
    "\n",
    "uspec = uspec0.clone().to(device)\n",
    "\n",
    "# save for later:\n",
    "nskip = 50\n",
    "utspec = torch.zeros(floor(maxiter//nskip) + 1, *uspec.shape).cdouble().to(device)\n",
    "\n",
    "with torch.inference_mode():\n",
    "    for iter in range(maxiter+1):\n",
    "        t = iter*dt\n",
    "\n",
    "        if iter % nskip == 0:\n",
    "            utspec[iter//nskip] = uspec\n",
    "            print(f\"t={t/3600:.2f} hours\")\n",
    "\n",
    "        dudtspec[inew] = swe_solver.dudtspec(uspec)\n",
    "        \n",
    "        # update vort,div,phiv with third-order adams-bashforth.\n",
    "        # forward euler, then 2nd-order adams-bashforth time steps to start.\n",
    "        if iter == 0:\n",
    "            dudtspec[inow] = dudtspec[inew]\n",
    "            dudtspec[iold] = dudtspec[inew]\n",
    "        elif iter == 1:\n",
    "            dudtspec[iold] = dudtspec[inew]\n",
    "\n",
    "        uspec = uspec + swe_solver.dt*( (23./12.) * dudtspec[inew] - (16./12.) * dudtspec[inow] + (5./12.) * dudtspec[iold] )\n",
    "\n",
    "        # implicit hyperdiffusion for vort and div.\n",
    "        uspec[1:] = swe_solver.hyperdiff * uspec[1:]\n",
    "        # cycle through the indices\n",
    "        inew = (inew - 1) % 3\n",
    "        inow = (inow - 1) % 3\n",
    "        iold = (iold - 1) % 3\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "image/png": "[encoded data, 90064 characters]",
      "text/plain": [
       "<Figure size 432x288 with 1 Axes>"
      ]
     },
     "metadata": {
      "needs_background": "light"
     },
     "output_type": "display_data"
    }
   ],
   "source": [
    "fig = plt.figure()\n",
    "im = swe_solver.plot_specdata(uspec[1], fig, cmap=\"twilight_shifted\")\n",
    "plt.show()"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Plotting a video\n",
    "\n",
    "let us plot the vorticity for our rollout:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "image/png": "[encoded data, 170188 characters]",
      "text/plain": [
       "<Figure size 576x432 with 1 Axes>"
      ]
     },
     "metadata": {
      "needs_background": "light"
     },
     "output_type": "display_data"
    }
   ],
   "source": [
    "# prepare figure for animation\n",
    "fig = plt.figure(figsize=(8, 6), dpi=72)\n",
    "moviewriter = animation.writers['pillow'](fps=20)\n",
    "moviewriter.setup(fig, './plots/zonal_jet.gif', dpi=72)\n",
    "\n",
    "plot_pvrt = False\n",
    "\n",
    "for i in range(utspec.shape[0]):\n",
    "    t = i*nskip*dt\n",
    "\n",
    "    if plot_pvrt:\n",
    "        variable = swe_solver.potential_vorticity(utspec[i])\n",
    "    else:\n",
    "        variable = swe_solver.spec2grid(utspec[i, 1])\n",
    "\n",
    "    plt.clf()\n",
    "    # swe_solver.plot_griddata(variable, cmap=cmap, vmin=-0.2, vmax=1.8, title=f'zonal jet t={t/3600:.2f} hours')\n",
    "    swe_solver.plot_griddata(variable, fig, cmap=\"twilight_shifted\", antialiased=False)\n",
    "    plt.draw()\n",
    "    moviewriter.grab_frame()\n",
    "\n",
    "\n",
    "moviewriter.finish()"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Conservation of potential vorticity"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {},
   "outputs": [],
   "source": [
    "pvrttspec = torch.zeros(floor(maxiter//nskip) + 1, lmax, mmax).cdouble().to(device)\n",
    "for i in range(utspec.shape[0]):\n",
    "    pvrttspec[i] = swe_solver.grid2spec(swe_solver.potential_vorticity(utspec[i]))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "<matplotlib.legend.Legend at 0x7f4e53406a70>"
      ]
     },
     "execution_count": 9,
     "metadata": {},
     "output_type": "execute_result"
    },
    {
     "data": {
      "image/png": "[encoded data, 9932 characters]",
      "text/plain": [
       "<Figure size 432x288 with 1 Axes>"
      ]
     },
     "metadata": {
      "needs_background": "light"
     },
     "output_type": "display_data"
    }
   ],
   "source": [
    "total_vrt = pvrttspec.abs()**2\n",
    "# total_vrt = utspec[..., 1, :, :].abs()**2\n",
    "total_vrt = torch.sqrt(torch.sum(total_vrt[..., :1], dim=(-1, -2)) + torch.sum(2 * total_vrt[..., 1:], dim=(-1, -2))).cpu()\n",
    "t = nskip*dt * torch.arange(utspec.shape[0])\n",
    "\n",
    "plt.plot(t, total_vrt / total_vrt[0], label='Spectral Solver')\n",
    "plt.title('Total vorticity over time')\n",
    "plt.ylim((0,1))\n",
    "plt.legend(loc='lower left')"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3.10.6 64-bit",
   "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.10.6"
  },
  "orig_nbformat": 4,
  "vscode": {
   "interpreter": {
    "hash": "31f2aee4e71d21fbe5cf8b01ff0e069b9275f58929596ceb00d14d90e3e16cd6"
   }
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}
