{
 "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": 14,
   "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",
    "cmap = 'twilight_shifted'\n",
    "\n",
    "from math import ceil, floor\n",
    "\n",
    "import sys\n",
    "sys.path.append(\"../torch_harmonics\")\n",
    "sys.path.append(\"../examples\")\n",
    "\n",
    "from torch_harmonics.sht import *\n",
    "from shallow_water_equations import *"
   ]
  },
  {
   "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(512/3)\n",
    "mmax = lmax\n",
    "# timestepping\n",
    "dt = 50\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()\n",
    "# uspec0 = swe_solver.random_initial_condition()\n",
    "    "
   ]
  },
  {
   "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=0.69 hours\n",
      "t=1.39 hours\n",
      "t=2.08 hours\n",
      "t=2.78 hours\n",
      "t=3.47 hours\n",
      "t=4.17 hours\n",
      "t=4.86 hours\n",
      "t=5.56 hours\n",
      "t=6.25 hours\n",
      "t=6.94 hours\n",
      "t=7.64 hours\n",
      "t=8.33 hours\n",
      "t=9.03 hours\n",
      "t=9.72 hours\n",
      "t=10.42 hours\n",
      "t=11.11 hours\n",
      "t=11.81 hours\n",
      "t=12.50 hours\n",
      "t=13.19 hours\n",
      "t=13.89 hours\n",
      "t=14.58 hours\n",
      "t=15.28 hours\n",
      "t=15.97 hours\n",
      "t=16.67 hours\n",
      "t=17.36 hours\n",
      "t=18.06 hours\n",
      "t=18.75 hours\n",
      "t=19.44 hours\n",
      "t=20.14 hours\n",
      "t=20.83 hours\n",
      "t=21.53 hours\n",
      "t=22.22 hours\n",
      "t=22.92 hours\n",
      "t=23.61 hours\n",
      "t=24.31 hours\n",
      "t=25.00 hours\n",
      "t=25.69 hours\n",
      "t=26.39 hours\n",
      "t=27.08 hours\n",
      "t=27.78 hours\n",
      "t=28.47 hours\n",
      "t=29.17 hours\n",
      "t=29.86 hours\n",
      "t=30.56 hours\n",
      "t=31.25 hours\n",
      "t=31.94 hours\n",
      "t=32.64 hours\n",
      "t=33.33 hours\n",
      "t=34.03 hours\n",
      "t=34.72 hours\n",
      "t=35.42 hours\n",
      "t=36.11 hours\n",
      "t=36.81 hours\n",
      "t=37.50 hours\n",
      "t=38.19 hours\n",
      "t=38.89 hours\n",
      "t=39.58 hours\n",
      "t=40.28 hours\n",
      "t=40.97 hours\n",
      "t=41.67 hours\n",
      "t=42.36 hours\n",
      "t=43.06 hours\n",
      "t=43.75 hours\n",
      "t=44.44 hours\n",
      "t=45.14 hours\n",
      "t=45.83 hours\n",
      "t=46.53 hours\n",
      "t=47.22 hours\n",
      "t=47.92 hours\n",
      "t=48.61 hours\n",
      "t=49.31 hours\n",
      "t=50.00 hours\n",
      "t=50.69 hours\n",
      "t=51.39 hours\n",
      "t=52.08 hours\n",
      "t=52.78 hours\n",
      "t=53.47 hours\n",
      "t=54.17 hours\n",
      "t=54.86 hours\n",
      "t=55.56 hours\n",
      "t=56.25 hours\n",
      "t=56.94 hours\n",
      "t=57.64 hours\n",
      "t=58.33 hours\n",
      "t=59.03 hours\n",
      "t=59.72 hours\n",
      "t=60.42 hours\n",
      "t=61.11 hours\n",
      "t=61.81 hours\n",
      "t=62.50 hours\n",
      "t=63.19 hours\n",
      "t=63.89 hours\n",
      "t=64.58 hours\n",
      "t=65.28 hours\n",
      "t=65.97 hours\n",
      "t=66.67 hours\n",
      "t=67.36 hours\n",
      "t=68.06 hours\n",
      "t=68.75 hours\n",
      "t=69.44 hours\n",
      "t=70.14 hours\n",
      "t=70.83 hours\n",
      "t=71.53 hours\n",
      "t=72.22 hours\n",
      "t=72.92 hours\n",
      "t=73.61 hours\n",
      "t=74.31 hours\n",
      "t=75.00 hours\n",
      "t=75.69 hours\n",
      "t=76.39 hours\n",
      "t=77.08 hours\n",
      "t=77.78 hours\n",
      "t=78.47 hours\n",
      "t=79.17 hours\n",
      "t=79.86 hours\n",
      "t=80.56 hours\n",
      "t=81.25 hours\n",
      "t=81.94 hours\n",
      "t=82.64 hours\n",
      "t=83.33 hours\n",
      "t=84.03 hours\n",
      "t=84.72 hours\n",
      "t=85.42 hours\n",
      "t=86.11 hours\n",
      "t=86.81 hours\n",
      "t=87.50 hours\n",
      "t=88.19 hours\n",
      "t=88.89 hours\n",
      "t=89.58 hours\n",
      "t=90.28 hours\n",
      "t=90.97 hours\n",
      "t=91.67 hours\n",
      "t=92.36 hours\n",
      "t=93.06 hours\n",
      "t=93.75 hours\n",
      "t=94.44 hours\n",
      "t=95.14 hours\n",
      "t=95.83 hours\n",
      "t=96.53 hours\n",
      "t=97.22 hours\n",
      "t=97.92 hours\n",
      "t=98.61 hours\n",
      "t=99.31 hours\n",
      "t=100.00 hours\n",
      "t=100.69 hours\n",
      "t=101.39 hours\n",
      "t=102.08 hours\n",
      "t=102.78 hours\n",
      "t=103.47 hours\n",
      "t=104.17 hours\n",
      "t=104.86 hours\n",
      "t=105.56 hours\n",
      "t=106.25 hours\n",
      "t=106.94 hours\n",
      "t=107.64 hours\n",
      "t=108.33 hours\n",
      "t=109.03 hours\n",
      "t=109.72 hours\n",
      "t=110.42 hours\n",
      "t=111.11 hours\n",
      "t=111.81 hours\n",
      "t=112.50 hours\n",
      "t=113.19 hours\n",
      "t=113.89 hours\n",
      "t=114.58 hours\n",
      "t=115.28 hours\n",
      "t=115.97 hours\n",
      "t=116.67 hours\n",
      "t=117.36 hours\n",
      "t=118.06 hours\n",
      "t=118.75 hours\n",
      "t=119.44 hours\n",
      "t=120.14 hours\n",
      "t=120.83 hours\n",
      "t=121.53 hours\n",
      "t=122.22 hours\n",
      "t=122.92 hours\n",
      "t=123.61 hours\n",
      "t=124.31 hours\n",
      "t=125.00 hours\n",
      "t=125.69 hours\n",
      "t=126.39 hours\n",
      "t=127.08 hours\n",
      "t=127.78 hours\n",
      "t=128.47 hours\n",
      "t=129.17 hours\n",
      "t=129.86 hours\n",
      "t=130.56 hours\n",
      "t=131.25 hours\n",
      "t=131.94 hours\n",
      "t=132.64 hours\n",
      "t=133.33 hours\n",
      "t=134.03 hours\n",
      "t=134.72 hours\n",
      "t=135.42 hours\n",
      "t=136.11 hours\n",
      "t=136.81 hours\n",
      "t=137.50 hours\n",
      "t=138.19 hours\n",
      "t=138.89 hours\n",
      "t=139.58 hours\n",
      "t=140.28 hours\n",
      "t=140.97 hours\n",
      "t=141.67 hours\n",
      "t=142.36 hours\n",
      "t=143.06 hours\n",
      "t=143.75 hours\n",
      "t=144.44 hours\n",
      "t=145.14 hours\n",
      "t=145.83 hours\n",
      "t=146.53 hours\n",
      "t=147.22 hours\n",
      "t=147.92 hours\n",
      "t=148.61 hours\n",
      "t=149.31 hours\n",
      "t=150.00 hours\n",
      "t=150.69 hours\n",
      "t=151.39 hours\n",
      "t=152.08 hours\n",
      "t=152.78 hours\n",
      "t=153.47 hours\n",
      "t=154.17 hours\n",
      "t=154.86 hours\n",
      "t=155.56 hours\n",
      "t=156.25 hours\n",
      "t=156.94 hours\n",
      "t=157.64 hours\n",
      "t=158.33 hours\n",
      "t=159.03 hours\n",
      "t=159.72 hours\n",
      "t=160.42 hours\n",
      "t=161.11 hours\n",
      "t=161.81 hours\n",
      "t=162.50 hours\n",
      "t=163.19 hours\n",
      "t=163.89 hours\n",
      "t=164.58 hours\n",
      "t=165.28 hours\n",
      "t=165.97 hours\n",
      "t=166.67 hours\n",
      "t=167.36 hours\n",
      "t=168.06 hours\n",
      "t=168.75 hours\n",
      "t=169.44 hours\n",
      "t=170.14 hours\n",
      "t=170.83 hours\n",
      "t=171.53 hours\n",
      "t=172.22 hours\n",
      "t=172.92 hours\n",
      "t=173.61 hours\n",
      "t=174.31 hours\n",
      "t=175.00 hours\n",
      "t=175.69 hours\n",
      "t=176.39 hours\n",
      "t=177.08 hours\n",
      "t=177.78 hours\n",
      "t=178.47 hours\n",
      "t=179.17 hours\n",
      "t=179.86 hours\n",
      "t=180.56 hours\n",
      "t=181.25 hours\n",
      "t=181.94 hours\n",
      "t=182.64 hours\n",
      "t=183.33 hours\n",
      "t=184.03 hours\n",
      "t=184.72 hours\n",
      "t=185.42 hours\n",
      "t=186.11 hours\n",
      "t=186.81 hours\n",
      "t=187.50 hours\n",
      "t=188.19 hours\n",
      "t=188.89 hours\n",
      "t=189.58 hours\n",
      "t=190.28 hours\n",
      "t=190.97 hours\n",
      "t=191.67 hours\n",
      "t=192.36 hours\n",
      "t=193.06 hours\n",
      "t=193.75 hours\n",
      "t=194.44 hours\n",
      "t=195.14 hours\n",
      "t=195.83 hours\n",
      "t=196.53 hours\n",
      "t=197.22 hours\n",
      "t=197.92 hours\n",
      "t=198.61 hours\n",
      "t=199.31 hours\n",
      "t=200.00 hours\n",
      "t=200.69 hours\n",
      "t=201.39 hours\n",
      "t=202.08 hours\n",
      "t=202.78 hours\n",
      "t=203.47 hours\n",
      "t=204.17 hours\n",
      "t=204.86 hours\n",
      "t=205.56 hours\n",
      "t=206.25 hours\n",
      "t=206.94 hours\n",
      "t=207.64 hours\n",
      "t=208.33 hours\n",
      "t=209.03 hours\n",
      "t=209.72 hours\n",
      "t=210.42 hours\n",
      "t=211.11 hours\n",
      "t=211.81 hours\n",
      "t=212.50 hours\n",
      "t=213.19 hours\n",
      "t=213.89 hours\n",
      "t=214.58 hours\n",
      "t=215.28 hours\n",
      "t=215.97 hours\n",
      "t=216.67 hours\n",
      "t=217.36 hours\n",
      "t=218.06 hours\n",
      "t=218.75 hours\n",
      "t=219.44 hours\n",
      "t=220.14 hours\n",
      "t=220.83 hours\n",
      "t=221.53 hours\n",
      "t=222.22 hours\n",
      "t=222.92 hours\n",
      "t=223.61 hours\n",
      "t=224.31 hours\n",
      "t=225.00 hours\n",
      "t=225.69 hours\n",
      "t=226.39 hours\n",
      "t=227.08 hours\n",
      "t=227.78 hours\n",
      "t=228.47 hours\n",
      "t=229.17 hours\n",
      "t=229.86 hours\n",
      "t=230.56 hours\n",
      "t=231.25 hours\n",
      "t=231.94 hours\n",
      "t=232.64 hours\n",
      "t=233.33 hours\n",
      "t=234.03 hours\n",
      "t=234.72 hours\n",
      "t=235.42 hours\n",
      "t=236.11 hours\n",
      "t=236.81 hours\n",
      "t=237.50 hours\n",
      "t=238.19 hours\n",
      "t=238.89 hours\n",
      "t=239.58 hours\n",
      "t=240.28 hours\n",
      "t=240.97 hours\n",
      "t=241.67 hours\n",
      "t=242.36 hours\n",
      "t=243.06 hours\n",
      "t=243.75 hours\n",
      "t=244.44 hours\n",
      "t=245.14 hours\n",
      "t=245.83 hours\n",
      "t=246.53 hours\n",
      "t=247.22 hours\n",
      "t=247.92 hours\n",
      "t=248.61 hours\n",
      "t=249.31 hours\n",
      "t=250.00 hours\n",
      "t=250.69 hours\n",
      "t=251.39 hours\n",
      "t=252.08 hours\n",
      "t=252.78 hours\n",
      "t=253.47 hours\n",
      "t=254.17 hours\n",
      "t=254.86 hours\n",
      "t=255.56 hours\n",
      "t=256.25 hours\n",
      "t=256.94 hours\n",
      "t=257.64 hours\n",
      "t=258.33 hours\n",
      "t=259.03 hours\n",
      "t=259.72 hours\n",
      "t=260.42 hours\n",
      "t=261.11 hours\n",
      "t=261.81 hours\n",
      "t=262.50 hours\n",
      "t=263.19 hours\n",
      "t=263.89 hours\n",
      "t=264.58 hours\n",
      "t=265.28 hours\n",
      "t=265.97 hours\n",
      "t=266.67 hours\n",
      "t=267.36 hours\n",
      "t=268.06 hours\n",
      "t=268.75 hours\n",
      "t=269.44 hours\n",
      "t=270.14 hours\n",
      "t=270.83 hours\n",
      "t=271.53 hours\n",
      "t=272.22 hours\n",
      "t=272.92 hours\n",
      "t=273.61 hours\n",
      "t=274.31 hours\n",
      "t=275.00 hours\n",
      "t=275.69 hours\n",
      "t=276.39 hours\n",
      "t=277.08 hours\n",
      "t=277.78 hours\n",
      "t=278.47 hours\n",
      "t=279.17 hours\n",
      "t=279.86 hours\n",
      "t=280.56 hours\n",
      "t=281.25 hours\n",
      "t=281.94 hours\n",
      "t=282.64 hours\n",
      "t=283.33 hours\n",
      "t=284.03 hours\n",
      "t=284.72 hours\n",
      "t=285.42 hours\n",
      "t=286.11 hours\n",
      "t=286.81 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": 15,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "image/png": "[encoded data, 103560 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=cmap, antialiased=False)\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": 31,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "image/png": "[encoded data, 201136 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=cmap, 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": 17,
   "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": 30,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "<matplotlib.legend.Legend at 0x7ffa58e73190>"
      ]
     },
     "execution_count": 30,
     "metadata": {},
     "output_type": "execute_result"
    },
    {
     "data": {
      "image/png": "[encoded data, 19348 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')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "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
}
