{
 "cells": [
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Spherical Fourier Neural Operators\n",
    "\n",
    "A simple notebook to showcase spherical Fourier Neural Operators\n"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Preparation"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "import torch\n",
    "import torch.nn as nn\n",
    "from torch.utils.data import DataLoader\n",
    "from torch.cuda import amp\n",
    "from torch.optim.lr_scheduler import OneCycleLR\n",
    "\n",
    "import numpy as np\n",
    "import matplotlib.pyplot as plt\n",
    "import matplotlib.animation as animation\n",
    "\n",
    "from math import ceil\n",
    "\n",
    "import time\n",
    "\n",
    "cmap='twilight_shifted'"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [],
   "source": [
    "enable_amp = False\n",
    "\n",
    "# set device\n",
    "device = torch.device('cuda:0' if torch.cuda.is_available() else 'cpu')\n",
    "if torch.cuda.is_available():\n",
    "    torch.cuda.set_device(device.index)"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Training data\n",
    "to train our geometric FNOs, we require training data. To this end let us prepare a Dataloader which computes results on the fly:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [],
   "source": [
    "# dataset\n",
    "from torch_harmonics.examples.sfno import PdeDataset\n",
    "\n",
    "# 1 hour prediction steps\n",
    "dt = 1*3600\n",
    "dt_solver = 150\n",
    "nsteps = dt//dt_solver\n",
    "dataset = PdeDataset(dt=dt, nsteps=nsteps, dims=(256, 512), device=device, normalize=True)\n",
    "# There is still an issue with parallel dataloading. Do NOT use it at the moment\n",
    "dataloader = DataLoader(dataset, batch_size=4, shuffle=True, num_workers=0, persistent_workers=False)\n",
    "solver = dataset.solver.to(device)\n",
    "\n",
    "nlat = dataset.nlat\n",
    "nlon = dataset.nlon"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "image/png": "[encoded data, 105620 characters]",
      "text/plain": [
       "<Figure size 432x288 with 2 Axes>"
      ]
     },
     "metadata": {
      "needs_background": "light"
     },
     "output_type": "display_data"
    },
    {
     "data": {
      "image/png": "[encoded data, 103140 characters]",
      "text/plain": [
       "<Figure size 432x288 with 2 Axes>"
      ]
     },
     "metadata": {
      "needs_background": "light"
     },
     "output_type": "display_data"
    }
   ],
   "source": [
    "torch.manual_seed(0)\n",
    "inp, tar = dataset[0]\n",
    "\n",
    "fig = plt.figure()\n",
    "im = solver.plot_griddata(inp[2], fig, vmax=3, vmin=-3)\n",
    "plt.title(\"input\")\n",
    "plt.colorbar(im)\n",
    "plt.show()\n",
    "\n",
    "fig = plt.figure()\n",
    "im = solver.plot_griddata(tar[2], fig, vmax=3, vmin=-3)\n",
    "plt.title(\"target\")\n",
    "plt.colorbar(im)\n",
    "plt.show()"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Defining the geometric Fourier Neural Operator"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [],
   "source": [
    "from torch_harmonics.examples.sfno import SphericalFourierNeuralOperatorNet as SFNO"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [],
   "source": [
    "fno_model = SFNO(spectral_transform='sht', filter_type='linear', operator_type='driscoll-healy', img_size=(nlat, nlon),\n",
    "                 num_layers=4, scale_factor=3, embed_dim=256).to(device)"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Training the model"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [],
   "source": [
    "def l2loss_sphere(solver, prd, tar, relative=False, polar_opt=0):\n",
    "    loss = solver.integrate_grid((prd - tar)**2, dimensionless=True, polar_opt=polar_opt)\n",
    "    if relative:\n",
    "        loss = loss / solver.integrate_grid(tar**2, dimensionless=True, polar_opt=polar_opt)\n",
    "    loss = torch.mean(loss)\n",
    "    return loss\n",
    "\n",
    "def fluct_l2loss_sphere(solver, prd, tar, inp, relative=False, polar_opt=0):\n",
    "    # compute the weighting factor first\n",
    "    fluct = solver.integrate_grid((tar - inp)**2, dimensionless=True, polar_opt=polar_opt)\n",
    "    weight = fluct / torch.sum(fluct, dim=-1, keepdim=True)\n",
    "    # weight = weight.reshape(*weight.shape, 1, 1)\n",
    "    \n",
    "    loss = weight * solver.integrate_grid((prd - tar)**2, dimensionless=True, polar_opt=polar_opt)\n",
    "    if relative:\n",
    "        loss = loss / (weight * solver.integrate_grid(tar**2, dimensionless=True, polar_opt=polar_opt))\n",
    "    loss = torch.mean(loss)\n",
    "    return loss"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {},
   "outputs": [],
   "source": [
    "# training function\n",
    "def train_model(model, dataloader, optimizer, scheduler=None, nepochs=20, nfuture=0, num_examples=256, num_valid=8, loss='abs'):\n",
    "\n",
    "    train_start = time.time()\n",
    "\n",
    "    for epoch in range(nepochs):\n",
    "\n",
    "        # time each epoch\n",
    "        epoch_start = time.time()\n",
    "\n",
    "        dataloader.dataset.set_initial_condition('random')\n",
    "        dataloader.dataset.set_num_examples(num_examples)\n",
    "\n",
    "        optimizer.zero_grad(set_to_none=True)\n",
    "\n",
    "        # do the training\n",
    "        acc_loss = 0\n",
    "        model.train()\n",
    "        for inp, tar in dataloader:\n",
    "            with amp.autocast(enabled=enable_amp):\n",
    "                prd = model(inp)\n",
    "                for _ in range(nfuture):\n",
    "                    prd = model(prd)\n",
    "                if loss == 'fluct':\n",
    "                    loss = fluct_l2loss_sphere(solver, prd, tar, inp, relative=False)\n",
    "                else:\n",
    "                    loss = l2loss_sphere(solver, prd, tar, relative=False)\n",
    "\n",
    "            acc_loss += loss.item() * inp.size(0)\n",
    "\n",
    "            optimizer.zero_grad(set_to_none=True)\n",
    "            # gscaler.scale(loss).backward()\n",
    "            loss.backward()\n",
    "            optimizer.step()\n",
    "            # gscaler.update()\n",
    "\n",
    "        if scheduler is not None:\n",
    "            scheduler.step()\n",
    "\n",
    "        acc_loss = acc_loss / len(dataloader.dataset)\n",
    "\n",
    "        dataloader.dataset.set_initial_condition('random')\n",
    "        dataloader.dataset.set_num_examples(num_valid)\n",
    "\n",
    "        # perform validation\n",
    "        valid_loss = 0\n",
    "        model.eval()\n",
    "        with torch.no_grad():\n",
    "            for inp, tar in dataloader:\n",
    "                prd = model(inp)\n",
    "                for _ in range(nfuture):\n",
    "                    prd = model(prd)\n",
    "                loss = l2loss_sphere(solver, prd, tar, relative=True)\n",
    "\n",
    "                valid_loss += loss.item() * inp.size(0)\n",
    "\n",
    "        valid_loss = valid_loss / len(dataloader.dataset)\n",
    "\n",
    "        epoch_time = time.time() - epoch_start\n",
    "\n",
    "        print(f'--------------------------------------------------------------------------------')\n",
    "        print(f'Epoch {epoch} summary:')\n",
    "        print(f'time taken: {epoch_time}')\n",
    "        print(f'accumulated training loss: {acc_loss}')\n",
    "        print(f'relative validation loss: {valid_loss}')\n",
    "\n",
    "    train_time = time.time() - train_start\n",
    "\n",
    "    print(f'--------------------------------------------------------------------------------')\n",
    "    print(f'done. Training took {train_time}.')\n",
    "    return valid_loss"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "--------------------------------------------------------------------------------\n",
      "Epoch 0 summary:\n",
      "time taken: 30.86558222770691\n",
      "accumulated training loss: 0.6265566123183817\n",
      "relative validation loss: 0.012334701605141163\n",
      "--------------------------------------------------------------------------------\n",
      "Epoch 1 summary:\n",
      "time taken: 30.972100019454956\n",
      "accumulated training loss: 0.19270113436505198\n",
      "relative validation loss: 0.008939797524362803\n",
      "--------------------------------------------------------------------------------\n",
      "Epoch 2 summary:\n",
      "time taken: 31.100343942642212\n",
      "accumulated training loss: 0.16520043299533427\n",
      "relative validation loss: 0.006787700345739722\n",
      "--------------------------------------------------------------------------------\n",
      "Epoch 3 summary:\n",
      "time taken: 31.07342553138733\n",
      "accumulated training loss: 0.13038600096479058\n",
      "relative validation loss: 0.007389088161289692\n",
      "--------------------------------------------------------------------------------\n",
      "Epoch 4 summary:\n",
      "time taken: 31.208359241485596\n",
      "accumulated training loss: 0.13147166837006807\n",
      "relative validation loss: 0.009229443036019802\n",
      "--------------------------------------------------------------------------------\n",
      "Epoch 5 summary:\n",
      "time taken: 31.288078546524048\n",
      "accumulated training loss: 0.12216734082903713\n",
      "relative validation loss: 0.007763240719214082\n",
      "--------------------------------------------------------------------------------\n",
      "Epoch 6 summary:\n",
      "time taken: 31.260853052139282\n",
      "accumulated training loss: 0.10728413611650467\n",
      "relative validation loss: 0.006804670440033078\n",
      "--------------------------------------------------------------------------------\n",
      "Epoch 7 summary:\n",
      "time taken: 31.425776481628418\n",
      "accumulated training loss: 0.09526405163342133\n",
      "relative validation loss: 0.006044752197340131\n",
      "--------------------------------------------------------------------------------\n",
      "Epoch 8 summary:\n",
      "time taken: 31.336848258972168\n",
      "accumulated training loss: 0.10574182099662721\n",
      "relative validation loss: 0.004922760417684913\n",
      "--------------------------------------------------------------------------------\n",
      "Epoch 9 summary:\n",
      "time taken: 31.249579429626465\n",
      "accumulated training loss: 0.09587052289862186\n",
      "relative validation loss: 0.006112596718594432\n",
      "--------------------------------------------------------------------------------\n",
      "done. Training took 311.78155875205994.\n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "0.006112596718594432"
      ]
     },
     "execution_count": 11,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "# set seed\n",
    "torch.manual_seed(333)\n",
    "torch.cuda.manual_seed(333)\n",
    "\n",
    "optimizer = torch.optim.Adam(fno_model.parameters(), lr=3E-3, weight_decay=0.0)\n",
    "gscaler = amp.GradScaler(enabled=enable_amp)\n",
    "train_model(fno_model, dataloader, optimizer, nepochs=10)\n",
    "\n",
    "# multistep training\n",
    "# learning_rate = 5e-4\n",
    "# optimizer = torch.optim.Adam(fno_model.parameters(), lr=learning_rate)\n",
    "# dataloader.dataset.nsteps = 2 * dt//dt_solver\n",
    "# train_model(fno_model, dataloader, optimizer, nepochs=10, nfuture=1)\n",
    "# dataloader.dataset.nsteps = 1 * dt//dt_solver"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "image/png": "[encoded data, 105928 characters]",
      "text/plain": [
       "<Figure size 432x288 with 2 Axes>"
      ]
     },
     "metadata": {
      "needs_background": "light"
     },
     "output_type": "display_data"
    },
    {
     "data": {
      "image/png": "[encoded data, 105544 characters]",
      "text/plain": [
       "<Figure size 432x288 with 2 Axes>"
      ]
     },
     "metadata": {
      "needs_background": "light"
     },
     "output_type": "display_data"
    }
   ],
   "source": [
    "dataloader.dataset.set_initial_condition('random')\n",
    "\n",
    "torch.manual_seed(0)\n",
    "\n",
    "with torch.inference_mode():\n",
    "    inp, tar = next(iter(dataloader))\n",
    "    out = fno_model(inp).detach()\n",
    "\n",
    "fig = plt.figure()\n",
    "im = solver.plot_griddata(out[0, 1], fig, projection='3d', title='prediction')\n",
    "plt.colorbar(im)\n",
    "plt.show()\n",
    "\n",
    "fig = plt.figure()\n",
    "im = solver.plot_griddata(tar[0, 1], fig, projection='3d', title='target')\n",
    "plt.colorbar(im)\n",
    "plt.show()"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "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.12"
  },
  "orig_nbformat": 4,
  "vscode": {
   "interpreter": {
    "hash": "31f2aee4e71d21fbe5cf8b01ff0e069b9275f58929596ceb00d14d90e3e16cd6"
   }
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}
