{
 "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, sqrt\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": 14,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "image/png": "[encoded data, 102568 characters]",
      "text/plain": [
       "<Figure size 432x288 with 2 Axes>"
      ]
     },
     "metadata": {
      "needs_background": "light"
     },
     "output_type": "display_data"
    },
    {
     "data": {
      "image/png": "[encoded data, 99928 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": [
    "model = SFNO(spectral_transform='sht', operator_type='driscoll-healy', img_size=(nlat, nlon), grid=\"equiangular\",\n",
    "                 num_layers=4, scale_factor=3, embed_dim=16, big_skip=True, pos_embed=\"lat\", use_mlp=False, normalization_layer=\"none\").to(device)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [],
   "source": [
    "# pointwise model for sanity checking\n",
    "# class MLP(nn.Module):\n",
    "#     def __init__(self,\n",
    "#                  input_dim = 3,\n",
    "#                  output_dim = 3,\n",
    "#                  num_layers = 2,\n",
    "#                  hidden_dim = 32,\n",
    "#                  activation_function = nn.ReLU,\n",
    "#                  bias = False):\n",
    "#         super().__init__()\n",
    "    \n",
    "#         current_dim = input_dim\n",
    "#         layers = []\n",
    "#         for l in range(num_layers-1):\n",
    "#             fc = nn.Conv2d(current_dim, hidden_dim, 1, bias=True)\n",
    "#             # initialize the weights correctly\n",
    "#             scale = sqrt(2. / current_dim)\n",
    "#             nn.init.normal_(fc.weight, mean=0., std=scale)\n",
    "#             if fc.bias is not None:\n",
    "#                 nn.init.constant_(fc.bias, 0.0)\n",
    "#             layers.append(fc)\n",
    "#             layers.append(activation_function())\n",
    "#             current_dim = hidden_dim\n",
    "#         fc = nn.Conv2d(current_dim, output_dim, 1, bias=False)\n",
    "#         scale = sqrt(1. / current_dim)\n",
    "#         nn.init.normal_(fc.weight, mean=0., std=scale)\n",
    "#         if fc.bias is not None:\n",
    "#             nn.init.constant_(fc.bias, 0.0)\n",
    "#         layers.append(fc)\n",
    "#         self.mlp = nn.Sequential(*layers)\n",
    "\n",
    "#     def forward(self, x):\n",
    "#         return self.mlp(x)\n",
    "\n",
    "# model = MLP(num_layers=10).to(device)"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Training the model"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {},
   "outputs": [],
   "source": [
    "def l2loss_sphere(solver, prd, tar, relative=False, squared=True):\n",
    "    loss = solver.integrate_grid((prd - tar)**2, dimensionless=True).sum(dim=-1)\n",
    "    if relative:\n",
    "        loss = loss / solver.integrate_grid(tar**2, dimensionless=True).sum(dim=-1)\n",
    "    \n",
    "    if not squared:\n",
    "        loss = torch.sqrt(loss)\n",
    "    loss = loss.mean()\n",
    "\n",
    "    return loss\n",
    "\n",
    "def spectral_l2loss_sphere(solver, prd, tar, relative=False, squared=True):\n",
    "    # compute coefficients\n",
    "    coeffs = torch.view_as_real(solver.sht(prd - tar))\n",
    "    coeffs = coeffs[..., 0]**2 + coeffs[..., 1]**2\n",
    "    norm2 = coeffs[..., :, 0] + 2 * torch.sum(coeffs[..., :, 1:], dim=-1)\n",
    "    loss = torch.sum(norm2, dim=(-1,-2))\n",
    "\n",
    "    if relative:\n",
    "        tar_coeffs = torch.view_as_real(solver.sht(tar))\n",
    "        tar_coeffs = tar_coeffs[..., 0]**2 + tar_coeffs[..., 1]**2\n",
    "        tar_norm2 = tar_coeffs[..., :, 0] + 2 * torch.sum(tar_coeffs[..., :, 1:], dim=-1)\n",
    "        tar_norm2 = torch.sum(tar_norm2, dim=(-1,-2))\n",
    "        loss = loss / tar_norm2\n",
    "\n",
    "    if not squared:\n",
    "        loss = torch.sqrt(loss)\n",
    "    loss = loss.mean()\n",
    "\n",
    "    return loss"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "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_fn='l2'):\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_fn == 'l2':\n",
    "                    loss = l2loss_sphere(solver, prd, tar)\n",
    "                elif loss_fn == \"spectral l2\":\n",
    "                    loss = spectral_l2loss_sphere(solver, prd, tar)\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": 10,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "--------------------------------------------------------------------------------\n",
      "Epoch 0 summary:\n",
      "time taken: 15.235767841339111\n",
      "accumulated training loss: 152.01090368628502\n",
      "relative validation loss: 0.48974674940109253\n",
      "--------------------------------------------------------------------------------\n",
      "Epoch 1 summary:\n",
      "time taken: 13.605531454086304\n",
      "accumulated training loss: 18.09713925421238\n",
      "relative validation loss: 0.2691601365804672\n",
      "--------------------------------------------------------------------------------\n",
      "Epoch 2 summary:\n",
      "time taken: 13.389189958572388\n",
      "accumulated training loss: 11.215070217847824\n",
      "relative validation loss: 0.17566757649183273\n",
      "--------------------------------------------------------------------------------\n",
      "Epoch 3 summary:\n",
      "time taken: 13.35526180267334\n",
      "accumulated training loss: 7.873616486787796\n",
      "relative validation loss: 0.1415245682001114\n",
      "--------------------------------------------------------------------------------\n",
      "Epoch 4 summary:\n",
      "time taken: 13.382116079330444\n",
      "accumulated training loss: 6.145691208541393\n",
      "relative validation loss: 0.11577435955405235\n",
      "--------------------------------------------------------------------------------\n",
      "Epoch 5 summary:\n",
      "time taken: 13.4170503616333\n",
      "accumulated training loss: 5.044192872941494\n",
      "relative validation loss: 0.09770810976624489\n",
      "--------------------------------------------------------------------------------\n",
      "Epoch 6 summary:\n",
      "time taken: 13.44334363937378\n",
      "accumulated training loss: 4.255958370864391\n",
      "relative validation loss: 0.08036481589078903\n",
      "--------------------------------------------------------------------------------\n",
      "Epoch 7 summary:\n",
      "time taken: 13.616739749908447\n",
      "accumulated training loss: 3.727783866226673\n",
      "relative validation loss: 0.06929700076580048\n",
      "--------------------------------------------------------------------------------\n",
      "Epoch 8 summary:\n",
      "time taken: 13.629685640335083\n",
      "accumulated training loss: 3.1292022727429867\n",
      "relative validation loss: 0.056656915694475174\n",
      "--------------------------------------------------------------------------------\n",
      "Epoch 9 summary:\n",
      "time taken: 13.507434844970703\n",
      "accumulated training loss: 2.5669259428977966\n",
      "relative validation loss: 0.0476943776011467\n",
      "--------------------------------------------------------------------------------\n",
      "done. Training took 136.58274269104004.\n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "0.0476943776011467"
      ]
     },
     "execution_count": 10,
     "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(model.parameters(), lr=3E-3, weight_decay=0.0)\n",
    "gscaler = amp.GradScaler(enabled=enable_amp)\n",
    "train_model(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": 13,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "image/png": "[encoded data, 100416 characters]",
      "text/plain": [
       "<Figure size 432x288 with 2 Axes>"
      ]
     },
     "metadata": {
      "needs_background": "light"
     },
     "output_type": "display_data"
    },
    {
     "data": {
      "image/png": "[encoded data, 83088 characters]",
      "text/plain": [
       "<Figure size 432x288 with 2 Axes>"
      ]
     },
     "metadata": {
      "needs_background": "light"
     },
     "output_type": "display_data"
    },
    {
     "data": {
      "image/png": "[encoded data, 105408 characters]",
      "text/plain": [
       "<Figure size 432x288 with 2 Axes>"
      ]
     },
     "metadata": {
      "needs_background": "light"
     },
     "output_type": "display_data"
    },
    {
     "data": {
      "image/png": "[encoded data, 102776 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 = model(inp).detach()\n",
    "\n",
    "s = 0; ch = 2\n",
    "\n",
    "fig = plt.figure()\n",
    "im = solver.plot_griddata(inp[s, ch], fig, projection='3d', title='input')\n",
    "plt.colorbar(im)\n",
    "plt.show()\n",
    "\n",
    "fig = plt.figure()\n",
    "im = solver.plot_griddata(out[s, ch], fig, projection='3d', title='prediction')\n",
    "plt.colorbar(im)\n",
    "plt.show()\n",
    "\n",
    "fig = plt.figure()\n",
    "im = solver.plot_griddata(tar[s, ch], fig, projection='3d', title='target')\n",
    "plt.colorbar(im)\n",
    "plt.show()\n",
    "\n",
    "fig = plt.figure()\n",
    "im = solver.plot_griddata((tar-out)[s, ch], fig, projection='3d', title='error')\n",
    "plt.colorbar(im)\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "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
}
