{
 "cells": [
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Training GNS on the 2D Taylor Green Vortex\n",
    "This tutorial notebook will show how to train a GNS model on the 2D Taylor Green Vortex. It will give a brief overview of LagrangeBench on these topics:\n",
    "- Datasets and case setup\n",
    "- Models\n",
    "- Training tricks and strategies (random-walk noise and pushforward)\n",
    "- Training and inference"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "2023-10-24 20:02:59.973089: E external/xla/xla/stream_executor/cuda/cuda_driver.cc:276] failed call to cuInit: CUDA_ERROR_NO_DEVICE: no CUDA-capable device is detected\n",
      "CUDA backend failed to initialize: FAILED_PRECONDITION: No visible GPU devices. (Set TF_CPP_MIN_LOG_LEVEL=0 and rerun for more info.)\n"
     ]
    }
   ],
   "source": [
    "import os\n",
    "os.environ[\"CUDA_VISIBLE_DEVICES\"] = \"\"\n",
    "os.environ[\"JAX_ENABLE_X64\"] = \"True\"\n",
    "\n",
    "import lagrangebench\n",
    "import haiku as hk\n",
    "import jax.numpy as jnp\n",
    "import numpy as np\n",
    "import matplotlib.pyplot as plt\n",
    "import pickle\n",
    "import matplotlib.animation as animation"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Datasets\n",
    "First thing to do is to load the dataset. The simplest way to do this is by using e.g. the `lagrangebench.data.TGV2D` class for the 2-dimensional Taylor-Green vortex problem. It will automatically download the HDF5 files if they are not found in the respective folder, and it will take care of setting up the dataset. Note that for the validation/test set you need to specify a positive number of rollout steps, e.g. `n_rollout_steps=5`. This means that the dataset will not split the trajectory into subsequences and keep whole rollouts for evaluation."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "This is a 2D dataset called TGV.\n",
      "Train snapshot have shape (2500, 7, 2) (n_nodes, seq_len, xy pos).\n",
      "Val snapshot have shape (2500, 26, 2) (n_nodes, rollout, xy pos).\n",
      "\n"
     ]
    }
   ],
   "source": [
    "tgv2d_train = lagrangebench.data.TGV2D(\"train\")\n",
    "tgv2d_eval = lagrangebench.data.TGV2D(\"valid\", n_rollout_steps=20)\n",
    "\n",
    "print(\n",
    "    f\"This is a {tgv2d_train.metadata['dim']}D dataset \"\n",
    "    f\"called {tgv2d_train.metadata['case']}.\\n\"\n",
    "    f\"Train snapshot have shape {tgv2d_train[0][0].shape} (n_nodes, seq_len, xy pos).\\n\"\n",
    "    f\"Val snapshot have shape {tgv2d_eval[0][0].shape} (n_nodes, rollout, xy pos).\\n\"\n",
    ")"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Similarly, for other datasets one can use the respective subclass, for example\n",
    "```python\n",
    "rpf_3d_data = lagrangebench.data.RPF3D(\"train\") # 3D Reverse Poiseuille flow\n",
    "dam_2d_data = lagrangebench.data.DAM2D(\"train\") # 2D Dam break\n",
    "# etc.\n",
    "# and in general: \n",
    "lagrangebench.data.H5Dataset(\"train\", dataset_path=\"path/to/dataset\")\n",
    "```"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Models\n",
    "All models should inherit from [`models.BaseModel`](/lagrangebench/models/base.py), and generally include a `_transform` function for feature engineering and graph building. \n",
    "\n",
    "Here we use a small GNS model, with latent dimension of 16 and 4 message passing layers and predicting 2D accelerations. Note that we use a function wrapper beause `haiku.Modules` must be initialized inside `haiku.transform`.\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [],
   "source": [
    "def gns(x):\n",
    "    return lagrangebench.GNS(\n",
    "        particle_dimension=tgv2d_train.metadata[\"dim\"],\n",
    "        latent_size=16,\n",
    "        blocks_per_step=2,\n",
    "        num_mp_steps=4,\n",
    "        particle_type_embedding_size=8,\n",
    "    )(x)"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Then apply `haiku.transform` to the model function"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [],
   "source": [
    "gns = hk.without_apply_rng(hk.transform_with_state(gns))"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Training tricks and strategies\n",
    "To improve performance and stabilize training we can use a number of tricks and strategies. In this notebook we will cover random-walk noise and pushforward, which are included by default in LagrangeBench. \n",
    "\n",
    "### Random-walk noise\n",
    "Briefly, random-walk noise adds noise to the velocities in the input sequence and adjusts the target positions accordingly. The standard deviation is passed as a parameter, and each noise step is rescaled so that the cumulated noise has the input standard deviation. It applied to GNNs in physics in the [\"Learning to Simulate Complex Physics with Graph Networks\"](https://arxiv.org/abs/2002.09405) paper (2020).\n",
    "\n",
    "### Pushforward\n",
    "Instead of training always on the next step, the pushforward trick will randomly sample a number of \"unroll\" steps to perform and compute the loss only on the last one. Ideally this should inject the \"right amount\" of noise, since it comes from the physical distribution. Because it adds some overhead, the maximum number of unroll steps should not be too large and the probability of unrolling large sequences should be reasonably small. It comes from the [\"Message Passing Neural PDE Solvers\"](https://arxiv.org/pdf/2202.03376.pdf) paper (2022)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [],
   "source": [
    "noise_std = 3e-4\n",
    "\n",
    "pf_config = lagrangebench.PushforwardConfig(\n",
    "    steps=[-1, 500, 700],  # training steps to unlock the relative stage\n",
    "    unrolls=[0, 2, 5],  # number of unroll steps per stage\n",
    "    probs=[7, 2, 1],  # relative probabilities to unroll to the relative stage\n",
    ")"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "For example, this configuration would apply noise with `std=1e-5` and the pushforward trick with three unroll stages (0, 2 and 5), \"unlocking\" the second stage after 5000 training steps and the third stage after 10000 training steps. After 10000 steps, 0-step unroll (normal training) will happen with a probability of 70%, 2-step unroll with a probability of 20% and finally 5-step unroll with a probability of 10%."
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Case\n",
    "The `CaseSetupFn` is an essential part in LagrangeBench. Most of the operations during training, such as data preprocessing, time integration for the ouput and neighbor computation are managed by `CaseSetupFn`. The easies way to initialize it is through the `lagrangebench.case_builder` function."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [],
   "source": [
    "bounds = np.array(tgv2d_train.metadata[\"bounds\"])\n",
    "box = bounds[:, 1] - bounds[:, 0]\n",
    "\n",
    "tgv2d_case = lagrangebench.case_builder(\n",
    "    box=box,  # (x,y) array with the world size along each axis. (1.0, 1.0) for 2D TGV\n",
    "    metadata=tgv2d_train.metadata,  # metadata dictionary\n",
    "    input_seq_length=6,  # number of consecutive time steps fed to the model\n",
    "    isotropic_norm=False,  # whether to normalize each dimension independently\n",
    "    noise_std=noise_std,  # noise standard deviation used by the random-walk noise\n",
    ")"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Training and inference\n",
    "Finally, to train GNS on Taylor Green (with noise and pushforward) the `lagrangebench.Trainer` methods comes to hand"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "0000, train/loss: 1.68755.\n",
      "0100, train/loss: 0.17377.\n",
      "0200, train/loss: 0.17768.\n",
      "0300, train/loss: 0.16432.\n",
      "0400, train/loss: 3.47414.\n",
      "0500, train/loss: 0.25281.\n",
      "{'val/loss': 0.006483794655650854, 'val/mse1': 3.545727849996183e-06, 'val/mse5': 0.0001387525990139693, 'val/mse10': 0.0009877493139356375, 'val/stdloss': 0.0, 'val/stdmse1': 0.0, 'val/stdmse5': 0.0, 'val/stdmse10': 0.0}\n",
      "0600, train/loss: 0.17975.\n",
      "0700, train/loss: 0.17109.\n",
      "0800, train/loss: 0.21826.\n",
      "0900, train/loss: 0.18034.\n",
      "1000, train/loss: 0.22316.\n",
      "{'val/loss': 0.004053703509271145, 'val/mse1': 2.7963376396655804e-06, 'val/mse5': 9.380515984958038e-05, 'val/mse10': 0.0006204654346220195, 'val/stdloss': 0.0, 'val/stdmse1': 0.0, 'val/stdmse5': 0.0, 'val/stdmse10': 0.0}\n"
     ]
    }
   ],
   "source": [
    "trainer = lagrangebench.Trainer(\n",
    "    model=gns,\n",
    "    case=tgv2d_case,\n",
    "    data_train=tgv2d_train,\n",
    "    data_eval=tgv2d_eval,\n",
    "    pushforward=pf_config,\n",
    "    noise_std=noise_std,\n",
    "    metrics=[\"mse\"],\n",
    "    n_rollout_steps=20,\n",
    "    eval_n_trajs=1,\n",
    "    lr_start=5e-4,\n",
    "    log_steps=100,\n",
    "    eval_steps=500,\n",
    ")\n",
    "\n",
    "params, state, _ = trainer(step_max=1000)"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Now let's see what the trained GNS can do. First, let's get the test data. As for the validation set, `n_rollout_steps=20`."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {},
   "outputs": [],
   "source": [
    "tgv2d_test = lagrangebench.TGV2D(\"test\", n_rollout_steps=20)"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Then let's run then evaluation over 20 timesteps. Here we use three evaluation metrics: __position MSE__ and __Sinkhorn__ distance."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {},
   "outputs": [],
   "source": [
    "metrics = lagrangebench.infer(\n",
    "    gns,\n",
    "    tgv2d_case,\n",
    "    tgv2d_test,\n",
    "    params,\n",
    "    state,\n",
    "    metrics=[\"mse\", \"sinkhorn\"],\n",
    "    eval_n_trajs=1,\n",
    "    n_rollout_steps=20,\n",
    "    rollout_dir=\"rollouts/\",\n",
    "    out_type=\"pkl\",\n",
    ")[\"rollout_0\"]\n",
    "rollout = pickle.load(open(\"rollouts/rollout_0.pkl\", \"rb\"))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "image/png": "[encoded data, 47624 characters]",
      "text/plain": [
       "<Figure size 1200x400 with 2 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "# plot metrics\n",
    "fig, ax = plt.subplots(1, 2, figsize=(12, 4))\n",
    "\n",
    "# plot the metrics\n",
    "for i, m in enumerate([\"mse\", \"sinkhorn\"]):\n",
    "    ax[i].plot(metrics[m])\n",
    "    ax[i].set_title(m)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "rollout of shape (26, 2500, 2) (steps, nodes, xy pos)\n"
     ]
    },
    {
     "data": {
      "text/html": [
       "<img src=\"media/scatter.gif\"/>"
      ],
      "text/plain": [
       "<IPython.core.display.Image object>"
      ]
     },
     "execution_count": 11,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "print(f\"rollout of shape {rollout['predicted_rollout'].shape} (steps, nodes, xy pos)\")\n",
    "\n",
    "fig, ax = plt.subplots(1, 2)\n",
    "ax[0].set_xlim([0, 1.0])\n",
    "ax[0].set_ylim([0, 1.0])\n",
    "ax[1].set_xlim([0, 1.0])\n",
    "ax[1].set_ylim([0, 1.0])\n",
    "fig.set_size_inches(10, 5, forward=True)\n",
    "ax[0].set_title(\"GNS\")\n",
    "ax[1].set_title(\"Ground Truth\")\n",
    "\n",
    "rollout_len = rollout[\"predicted_rollout\"].shape[0] - 1\n",
    "\n",
    "scat0 = ax[0].scatter(\n",
    "    rollout[\"predicted_rollout\"][0, :, 0], rollout[\"predicted_rollout\"][0, :, 1]\n",
    ")\n",
    "scat1 = ax[1].scatter(\n",
    "    rollout[\"ground_truth_rollout\"][0, :, 0], rollout[\"ground_truth_rollout\"][0, :, 1]\n",
    ")\n",
    "\n",
    "\n",
    "def animate(i):\n",
    "    scat0.set_offsets(rollout[\"predicted_rollout\"][i])\n",
    "    scat1.set_offsets(rollout[\"ground_truth_rollout\"][i])\n",
    "    return scat0, scat1\n",
    "\n",
    "\n",
    "ani = animation.FuncAnimation(\n",
    "    fig, animate, repeat=True, frames=rollout_len, interval=50\n",
    ")\n",
    "\n",
    "plt.close(fig)\n",
    "\n",
    "writer = animation.PillowWriter(fps=10, metadata=dict(artist=\"Me\"), bitrate=1800)\n",
    "ani.save(\"media/scatter.gif\", writer=writer)\n",
    "\n",
    "from IPython.display import Image\n",
    "\n",
    "Image(url=\"media/scatter.gif\")"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": ".venv",
   "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
 },
 "nbformat": 4,
 "nbformat_minor": 2
}
