{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "2fc737c1-fd94-4533-993e-e15589140332",
   "metadata": {},
   "source": [
    "# Training with the BubbleML Dataset\n",
    "This notebook shows an example of how to setup training with the BubbleML dataset.\n",
    "It uses a downsampled version of PB Subcooled with fewer simulations.\n",
    "\n",
    "For help with loading the data, check the `data_loading` notebook!"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "ae0243ab",
   "metadata": {},
   "outputs": [],
   "source": [
    "import h5py\n",
    "import torch\n",
    "from torch.utils.data import ConcatDataset, Dataset, DataLoader\n",
    "from neuralop.models import FNO\n",
    "import torch.nn.functional as F\n",
    "import matplotlib.pyplot as plt\n",
    "import matplotlib.animation as animation\n",
    "from scipy.signal import savgol_filter\n",
    "import numpy as np"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "e117944a-9a00-48e1-bd3f-53d87611a5ba",
   "metadata": {},
   "outputs": [],
   "source": [
    "DEVICE = 'cpu'"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "69419e66-7c3d-42f9-bc96-4b69dbf3691d",
   "metadata": {},
   "source": [
    "### Create the PyTorch Dataset\n",
    "We create a dataset that reads from a single HDF5 file, corresponding to one simulation. In this example, we only use one timestep for input to predict the following timestep. Thus, the input has three channels: temperature, x-velocity, and the y-velocity. It just predicts temperatures for the next step."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "64445e91-35d4-44be-bc2d-ef268f6a0318",
   "metadata": {},
   "outputs": [],
   "source": [
    "TEMPERATURE = 'temperature'\n",
    "VELX = 'velx'\n",
    "VELY = 'vely'\n",
    "\n",
    "class HDF5Dataset(Dataset):\n",
    "    def __init__(self, filename):\n",
    "        self.filename = filename\n",
    "        self.data = h5py.File(self.filename, 'r')\n",
    "        self.timesteps = self.data[TEMPERATURE][:].shape[0]\n",
    "    \n",
    "    def __len__(self):\n",
    "        return self.timesteps - 1\n",
    "\n",
    "    def _get_input(self, idx):\n",
    "        r\"\"\"\n",
    "        The input is the temperature, x-velocity, and y-velocity at time == idx\n",
    "        \"\"\"\n",
    "        temp = torch.from_numpy(self.data[TEMPERATURE][idx])\n",
    "        velx = torch.from_numpy(self.data[VELX][idx])\n",
    "        vely = torch.from_numpy(self.data[VELY][idx])\n",
    "        # returns a stack with shape [3 x Y x X]\n",
    "        return torch.stack((temp, velx, vely), dim=0)\n",
    "\n",
    "    def _get_label(self, idx):\n",
    "        r\"\"\"\n",
    "        The output is the temperature at time == idx\n",
    "        \"\"\"\n",
    "        return torch.from_numpy(self.data[TEMPERATURE][idx]).unsqueeze(0)\n",
    "    \n",
    "    def __getitem__(self, idx):\n",
    "        r\"\"\"\n",
    "        As input, get temperature and velocities at time == idx.\n",
    "        As the output label, get the temperature at time == idx + 1.\n",
    "        \"\"\"\n",
    "        input = self._get_input(idx)\n",
    "        label = self._get_label(idx+1)\n",
    "        return input, label"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "19469a0b-f6a6-4287-a009-84befe244cd9",
   "metadata": {},
   "source": [
    "### Create a ConcatDataset and Loaders\n",
    "In order to combine multiple simulations into one larger train/validation set, we use PyTorch's `ConcatDataset`. This is as simple as it sounds, you pass in a list of separate datasets and it concatenates them into one larger dataset. From the user perspective, this acts like a typical dataset. The datalaoders can use this `ConcatDataset` in the normal way."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "id": "c466fa4c",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Train batches: 100\n",
      "Val batches: 50\n"
     ]
    }
   ],
   "source": [
    "train_files = ['Twall-100.hdf5', 'Twall-106.hdf5']\n",
    "val_files = ['Twall-103.hdf5']\n",
    "\n",
    "train_dataset = ConcatDataset(HDF5Dataset(file) for file in train_files)\n",
    "val_dataset = ConcatDataset(HDF5Dataset(file) for file in val_files)\n",
    "\n",
    "train_dataloader = DataLoader(train_dataset, batch_size=4, shuffle=True)\n",
    "val_dataloader = DataLoader(val_dataset, batch_size=4, shuffle=False)\n",
    "\n",
    "print(f'Train batches: {len(train_dataloader)}')\n",
    "print(f'Val batches: {len(val_dataloader)}')"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "a5a16a35-c483-44fa-9c3d-9b6103921e92",
   "metadata": {},
   "source": [
    "### Creating a Model\n",
    "We use `neuralop`s implementation of the Fourier Neural Operator (FNO). It has 3 input channels because we input the `temperature`, `velx`, and `vely` for one timestep. It outputs one channels for the temperature at the following timestep. We keep the lowest (16, 16) modes. As the dataset resolution has been reduced, we can't keep many more modes than this. As this is a simpler example, we use only 64 hidden channels and 4 layers."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "id": "edb3ddcf-fdf7-45b0-9276-2a7e39eaf65e",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "FNO uses 4228097 parameters\n"
     ]
    }
   ],
   "source": [
    "def count_parameters(model):\n",
    "    return sum(p.numel() for p in model.parameters() if p.requires_grad)\n",
    "    \n",
    "model = FNO(in_channels=3,    # 3 channels for temp, velx, vely\n",
    "            out_channels=1,   # 1 channel for temp\n",
    "            n_modes=(16, 16), # keep the lowest fourier modes\n",
    "            hidden_channels=64,\n",
    "            n_layers=4)\n",
    "\n",
    "print(f'FNO uses {count_parameters(model)} parameters')\n",
    "\n",
    "optimizer = torch.optim.AdamW(model.parameters(), lr=1e-4)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "d03c3722-2183-4cfd-9b29-ef982aa03801",
   "metadata": {},
   "source": [
    "### Training!\n",
    "We train the model a short number of epochs and optimize using an MSE loss between the predicted temperature and the true temperature according to the simulation. We plot the average validation loss for each epoch. We see that the validation loss decreases."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "8f0e8f7c",
   "metadata": {},
   "outputs": [],
   "source": [
    "EPOCHS = 15\n",
    "\n",
    "val_losses = []\n",
    "\n",
    "for epoch in range(EPOCHS):\n",
    "    model.train()\n",
    "    for iter, (input, label) in enumerate(train_dataloader):\n",
    "        input = input.to(DEVICE).float()\n",
    "        label = label.to(DEVICE).float()\n",
    "        pred = model(input)\n",
    "        optimizer.zero_grad()\n",
    "        loss = F.mse_loss(pred, label)\n",
    "        loss.backward()\n",
    "        optimizer.step()\n",
    "\n",
    "    val_loss = []\n",
    "    model.eval()\n",
    "    for iter, (input, label) in enumerate(val_dataloader):\n",
    "        input = input.to(DEVICE).float()\n",
    "        label = label.to(DEVICE).float()\n",
    "        pred = model(input)\n",
    "        loss = F.mse_loss(pred, label)\n",
    "        val_loss.append(loss.detach().item())\n",
    "    val_losses.append(torch.mean(torch.tensor(val_loss)))\n",
    "    \n",
    "plt.plot(val_losses)\n",
    "plt.ylabel('MSE Loss')\n",
    "plt.xlabel('Epoch')"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "c784dec7-dfad-4c94-8b72-a875d4e453ce",
   "metadata": {},
   "source": [
    "### One-step error\n",
    "We visualize the one-step error for a sample output. We plot the ground-truth temperature, predicted temperature, and the absolute error between the two. We see that the predicted version is a decent approximation of the ground-truth (with more data and higher-resolution, it would of course look better.)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "0ff762ff-414e-4ba8-b838-242dfdb7bd6a",
   "metadata": {},
   "outputs": [],
   "source": [
    "input, label = val_dataset[35]\n",
    "input = input.to(DEVICE).float().unsqueeze(0)\n",
    "label = label.to(DEVICE).float()\n",
    "pred = model(input)\n",
    "\n",
    "label = label.squeeze().cpu().numpy()\n",
    "pred = pred.squeeze().detach().cpu().numpy()\n",
    "abs_err = np.abs(label - pred)\n",
    "\n",
    "fig, ax = plt.subplots(1, 3, figsize=(10, 5))\n",
    "\n",
    "data = {\n",
    "    'Ground Truth': label,\n",
    "    'Predicted': pred,\n",
    "    'Abs. Error': abs_err\n",
    "}\n",
    "\n",
    "for idx, (key, im) in enumerate(data.items()):\n",
    "    im = ax[idx].imshow(np.flipud(im))\n",
    "    fig.colorbar(im, ax=ax[idx], shrink=0.45)\n",
    "    ax[idx].set_title(key)\n",
    "    ax[idx].set_xticks([])\n",
    "    ax[idx].set_yticks([])"
   ]
  }
 ],
 "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.11.4"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}