{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# LagrangeBench Datasets\n",
    "\n",
    "This is an overview of the introduced datasets. We show:\n",
    "\n",
    "- Dataset animations.\n",
    "- Evaluation of physical properties.\n",
    "\n",
    "The structure of all datasets is the same and consists of four files:\n",
    "- `metadata.json` - the data needed to set up the numerical solver and simulate trajectories\n",
    "- `train.h5` - training set\n",
    "- `valid.h5` - validation set\n",
    "- `test.h5` - test set"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "2023-10-24 20:16:52.116209: 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",
    "\n",
    "import numpy as np\n",
    "from jax import vmap\n",
    "from jax_md import space\n",
    "import matplotlib.pyplot as plt\n",
    "\n",
    "import lagrangebench\n",
    "\n",
    "plt.rcParams.update({'font.size': 10})"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Dataset animations\n",
    "\n",
    "We animate a few consecutive frames from our datasets using [ParaView](https://www.paraview.org/) and color the particles by the magnitude of their velocity.\n",
    "\n",
    "- Reverse Poiseuille Flow\n",
    "\n",
    "    ![rpf2d.gif](https://s11.gifyu.com/images/Sce92.gif)\n",
    "    ![rpf3d.gif](https://s11.gifyu.com/images/Sce3X.gif)\n",
    "\n",
    "- Taylor Green Vortex\n",
    "\n",
    "    ![tgv2d.gif](https://s11.gifyu.com/images/Sce9b.gif)\n",
    "    ![tgv3d.gif](https://s11.gifyu.com/images/Sce9z.gif)\n",
    "\n",
    "- Lid-Driven Cavity\n",
    "\n",
    "    ![ldc2d.gif](https://s11.gifyu.com/images/Sce9S.gif)\n",
    "    ![ldc3d.gif](https://s11.gifyu.com/images/Sce3e.gif)\n",
    "\n",
    "- Dam Break\n",
    "\n",
    "    ![dam2d.gif](https://s11.gifyu.com/images/SceKB.gif)\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Dataset shapes\n",
    "\n",
    "Short script to check the shapes of the raw validation datasets.\n",
    "\n",
    "If you haven't downloaded the datasets yet, you can do so by running this script from the root of the repository:\n",
    "\n",
    "```bash\n",
    "bash download_data.sh all notebooks/datasets/\n",
    "```"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "2D_TGV_2500_10kevery100/valid.h5, traj_shape= (126, 2500, 2)      , split_valid_traj_into_n= 4    , num_trajs: 50   , eval_n_trajs_infer: 200\n",
      "2D_RPF_3200_20kevery100/valid.h5, traj_shape= (10001, 3200, 2)    , split_valid_traj_into_n= 384  , num_trajs: 1    , eval_n_trajs_infer: 384\n",
      "2D_LDC_2708_10kevery100/valid.h5, traj_shape= (5001, 2708, 2)     , split_valid_traj_into_n= 192  , num_trajs: 1    , eval_n_trajs_infer: 192\n",
      "2D_DAM_5740_20kevery100/valid.h5, traj_shape= (401, 5740, 2)      , split_valid_traj_into_n= 15   , num_trajs: 25   , eval_n_trajs_infer: 375\n",
      "3D_TGV_8000_10kevery100/valid.h5, traj_shape= (61, 8000, 3)       , split_valid_traj_into_n= 2    , num_trajs: 100  , eval_n_trajs_infer: 200\n",
      "3D_RPF_8000_10kevery100/valid.h5, traj_shape= (5001, 8000, 3)     , split_valid_traj_into_n= 192  , num_trajs: 1    , eval_n_trajs_infer: 192\n",
      "3D_LDC_8160_10kevery100/valid.h5, traj_shape= (5001, 8160, 3)     , split_valid_traj_into_n= 192  , num_trajs: 1    , eval_n_trajs_infer: 192\n"
     ]
    }
   ],
   "source": [
    "import os.path as osp\n",
    "import h5py\n",
    "\n",
    "# the datasets are stored in the following directories, whose names follow the convention:\n",
    "# {dim}_{case}_{num_particles_max}_{num_steps}every{sampling_rate}\n",
    "# dim - dimension. One of {2D|3D}.\n",
    "# case - simulation case. One of {TGV|RPF|LDC|DAM}.\n",
    "# num_particles_max - maximum number of particles in the dataset.\n",
    "# num_steps - number of position frames in the training split of the dataset.\n",
    "# sampling_rate - sampling_rate=100 means that the position frames correspond to every\n",
    "#     100th SPH simulation step.\n",
    "path = [\n",
    "    \"2D_TGV_2500_10kevery100\",\n",
    "    \"2D_RPF_3200_20kevery100\",\n",
    "    \"2D_LDC_2708_10kevery100\",\n",
    "    \"2D_DAM_5740_20kevery100\",\n",
    "    \"3D_TGV_8000_10kevery100\",        \n",
    "    \"3D_RPF_8000_10kevery100\",\n",
    "    \"3D_LDC_8160_10kevery100\",\n",
    "]\n",
    "\n",
    "# 6 history positions (5 past velocities) + 20-step rollout\n",
    "valid_seq_length = 6 + 20  # during validation and inference\n",
    "dataset_root = \"datasets/\"\n",
    "split = \"valid.h5\"\n",
    "\n",
    "for p in path:\n",
    "    p = osp.join(dataset_root, p, split)\n",
    "    with h5py.File(p, \"r\") as f:\n",
    "        pos_shape =  f[\"00000/position\"].shape\n",
    "        split_into = pos_shape[0]//valid_seq_length\n",
    "        num_trajs = len(list(f.keys()))\n",
    "        print(f\"{p[9:]}, traj_shape= {str(pos_shape):<20}, \"\n",
    "              f\"split_valid_traj_into_n= {split_into:<5}, \"\n",
    "              f\"num_trajs: {num_trajs:<5}, \"\n",
    "              f\"eval_n_trajs_infer: {split_into * num_trajs}\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Evaluation of physical properties\n",
    "\n",
    "Here we focus on the kinetic energy as a measure of physical behaviour. We distinguish two cases:\n",
    "\n",
    "1. Statistically stationary $E_{kin}$ (RPF and LDC) - the kinetic energy oscillates around an equilibrium value, and the oscillations are because of the finite number of particle used for the discretization.\n",
    "2. Variable $E_{kin}$ (TGV, DAM) - In the TGV example it decays due to viscous interactions and in dam break it first increases by transfer of potential to kinetic energy, but then also decays due to viscosity."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [],
   "source": [
    "# utils\n",
    "\n",
    "def compute_kinetic_energy(metadata, traj):\n",
    "    bounds = np.array(metadata[\"bounds\"])\n",
    "    displacement_fn, _ = space.periodic(side=bounds[:,1])\n",
    "    displ_vmap = vmap(displacement_fn, (0, 0))\n",
    "    displ_dvmap = vmap(displ_vmap, (0, 0))\n",
    "\n",
    "    displacement = displ_dvmap(traj[:,1:], traj[:, :-1])\n",
    "    vel = displacement / (metadata[\"dt\"] * metadata[\"write_every\"])\n",
    "    vel = np.asarray(vel) # shape: (num_particles, num_steps, dim)\n",
    "    square_absolute_velocity = (vel**2).sum(axis=2) # shape: (num_particles, num_steps)\n",
    "    e_kin = 0.5 * square_absolute_velocity.mean(axis=(0,)) # shape: (num_steps,)\n",
    "    return e_kin\n",
    "\n",
    "def plt_e_kin(dataset_path, seq_len, e_kin_ref=None):\n",
    "    dim, name, _, _ = dataset_path.split('/')[1].split('_')\n",
    "    \n",
    "    dataset = lagrangebench.data.H5Dataset(\n",
    "        \"valid\", \n",
    "        dataset_path=dataset_path,\n",
    "        input_seq_length=seq_len,\n",
    "        n_rollout_steps=1,\n",
    "    )\n",
    "    traj, _ = dataset[0]\n",
    "    e_kin = compute_kinetic_energy(dataset.metadata, traj)\n",
    "        \n",
    "    plt.figure(figsize=(8,4))\n",
    "    plt.plot(np.arange(len(e_kin)) + 0.5, e_kin, 'xb', label='dataset')\n",
    "    \n",
    "    if e_kin_ref is not None:\n",
    "        plt.plot(e_kin_ref, 'r', label='reference')\n",
    "    \n",
    "    plt.legend() \n",
    "    plt.grid()\n",
    "    plt.title(f\"{dim} {name}\")\n",
    "    plt.xlabel('step')\n",
    "    plt.ylabel('Ekin')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### RPF and LDC"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "image/png": "[encoded data, 51788 characters]",
      "text/plain": [
       "<Figure size 800x400 with 1 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "plt_e_kin(\"datasets/2D_RPF_3200_20kevery100\", 500)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "image/png": "[encoded data, 48888 characters]",
      "text/plain": [
       "<Figure size 800x400 with 1 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "plt_e_kin(\"datasets/3D_RPF_8000_10kevery100\", 500)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "image/png": "[encoded data, 58236 characters]",
      "text/plain": [
       "<Figure size 800x400 with 1 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "plt_e_kin(\"datasets/2D_LDC_2708_10kevery100\", 500)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "image/png": "[encoded data, 50964 characters]",
      "text/plain": [
       "<Figure size 800x400 with 1 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "plt_e_kin(\"datasets/3D_LDC_8160_10kevery100\", 500)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### DAM"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "image/png": "[encoded data, 42796 characters]",
      "text/plain": [
       "<Figure size 800x400 with 1 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "plt_e_kin(\"datasets/2D_DAM_5740_20kevery100\", 400)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### 2D TGV\n",
    "\n",
    "The two dimensional Taylor-Green vortex has an analytical solution as outlined in [1]. Here, we compare this solution of the kinetic energy with the kinetic enery in one of our dataset trajectories.\n",
    "\n",
    "[1] - Adami et al., [A transport-velocity formulation for smoothed particle hydrodynamics](https://www.sciencedirect.com/science/article/abs/pii/S002199911300096X), 2013"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "metadata": {},
   "outputs": [],
   "source": [
    "dataset_path = \"datasets/2D_TGV_2500_10kevery100\"\n",
    "dataset = lagrangebench.data.H5Dataset(\n",
    "    \"valid\", \n",
    "    dataset_path=dataset_path,\n",
    "    n_rollout_steps=1,\n",
    ")\n",
    "\n",
    "dt_dataset = (dataset.metadata['dt'] * dataset.metadata['write_every'])\n",
    "len_traj = 1 + dataset.metadata['t_end'] / dt_dataset\n",
    "t = np.linspace(0, dataset.metadata['t_end'], int(len_traj)) # physical time axis\n",
    "\n",
    "rho_ref = 1. # reference density\n",
    "u_ref = 1. # reference velocity, i.e. u_max at t=0\n",
    "L_ref = 1. # reference length, i.e. domain size\n",
    "eta = dataset.metadata['viscosity'] # viscosity\n",
    "Re = rho_ref * u_ref * L_ref / eta # Reynolds number\n",
    "slope_u_max = - 8 * np.pi ** 2 / Re # analytical decay rate of max velocity\n",
    "# factor 0.25 = int(int((sin(2 pi x) cos(2 pi y))^2, x from 0 to 1), y from 0 to 1)\n",
    "e_kin_theory = 0.25 * np.exp(2 * slope_u_max * t) # analytical kinetic energy evolution"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "In the above we used information from the `dataset.metadata` which is the content of a metadata file associated with each of our datasets. The metadata file for 2D TGV looks like this:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "{'case': 'TGV',\n",
       " 'solver': 'SPH',\n",
       " 'density_evolution': False,\n",
       " 'dim': 2,\n",
       " 'dx': 0.02,\n",
       " 'dt': 0.0004,\n",
       " 't_end': 5.0,\n",
       " 'viscosity': 0.01,\n",
       " 'p_bg_factor': 0.0,\n",
       " 'g_ext_magnitude': 0.0,\n",
       " 'artificial_alpha': 0.0,\n",
       " 'free_slip': False,\n",
       " 'write_every': 100,\n",
       " 'is_bc_trick': False,\n",
       " 'sequence_length_train': 126,\n",
       " 'num_trajs_train': 100,\n",
       " 'sequence_length_test': 126,\n",
       " 'num_trajs_test': 50,\n",
       " 'num_particles_max': 2500,\n",
       " 'periodic_boundary_conditions': [True, True, True],\n",
       " 'bounds': [[0.0, 1.0], [0.0, 1.0]],\n",
       " 'default_connectivity_radius': 0.029,\n",
       " 'vel_mean': [2.9928384037702926e-07, 5.485588872033986e-07],\n",
       " 'vel_std': [0.006356782279908657, 0.006356848403811455],\n",
       " 'acc_mean': [1.1536379524465599e-14, -1.9131170172353995e-13],\n",
       " 'acc_std': [0.0008755344315432012, 0.0008755328599363565]}"
      ]
     },
     "execution_count": 11,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "dataset.metadata"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "image/png": "[encoded data, 37104 characters]",
      "text/plain": [
       "<Figure size 800x400 with 1 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "plt_e_kin(dataset_path, 125, e_kin_theory)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We don't perfectly solve the 2D TGV, but this behaviour is a known limitation of purely Lagrangian SPH, as shown in [1]. We remind the reader that we restrict ourselves to purely Lagrangian schemes so that we are able to recover the velocity of particles by computing the difference between positions, something that is not valid with improved semi-Lagrangian SPH schemes."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### 3D TGV\n",
    "\n",
    "The three dimensional Taylor Green system at Re=50 does not have an analytical solution, but we run a high resolution reference simulation with the [JAX-Fluids](https://github.com/tumaer/JAXFLUIDS) solver and store the values of the kinetic energy at seconds [0, 0.5, 1, 1.5, ... 10]. \n",
    "\n",
    "Here, we again don't perfectly recover the true solution, but the initial rapid drop in kinetic energy on the 3D TGV is also a documented limitatation of purely Lagrangian SPH, see [2]. And yet we are rather close to the reference.\n",
    "\n",
    "[2] - Adami et al., [Simulating three-dimensional turbulence with SPH](https://mediatum.ub.tum.de/doc/1360487/978874575863.pdf), 2012"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "image/png": "[encoded data, 32580 characters]",
      "text/plain": [
       "<Figure size 800x400 with 1 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "e_kin_ref = 0.001 * np.array([\n",
    "    125, 118, 111, 103, 95.7, 88.0, 80.1, 72.4, 65, 58, 51.5, 45.5, 40.2, 35.5, 31.3, 27.6, 24.3, 21.5, 19.1, 16.9, 15.1\n",
    "])\n",
    "plt_e_kin(\"datasets/3D_TGV_8000_10kevery100\", 59, e_kin_ref)"
   ]
  }
 ],
 "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
}
