{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Working with other datasets\n",
    "\n",
    "We demonstrate how to train a GNN on one of the datasets provided by DeepMind along with the paper [1]. We divide this notebook into three parts:\n",
    "\n",
    "- Download and preprocess\n",
    "- Inspect the data\n",
    "- GNN training\n",
    "\n",
    "[1] - Sanchez-Gonzalez et al., [\"Learning to Simulate Complex Physics with Graph Networks\"](https://arxiv.org/abs/2002.09405), ICLR 2020"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "2023-10-24 20:07:32.539930: 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 lagrangebench\n",
    "import haiku as hk\n",
    "import numpy as np\n",
    "import matplotlib.pyplot as plt"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Download and preprocess\n",
    "\n",
    "The easiest was to download and preprocess the datasets from the GNS paper is by following the instructions in [`gns_data/README.md`](../gns_data/README.md). For our demonstration, we choose the 2D WaterDrop dataset. \n",
    "\n",
    "Note: the size of this dataset is around 4.5GB."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "--2023-10-15 01:39:30--  https://storage.googleapis.com/learning-to-simulate-complex-physics/Datasets/WaterDrop/metadata.json\n",
      "Resolving storage.googleapis.com (storage.googleapis.com)... 142.250.186.91, 172.217.18.27, 142.250.186.59, ...\n",
      "Connecting to storage.googleapis.com (storage.googleapis.com)|142.250.186.91|:443... connected.\n",
      "HTTP request sent, awaiting response... 200 OK\n",
      "Length: 361 [application/octet-stream]\n",
      "Saving to: ‘./datasets/WaterDrop/metadata.json’\n",
      "\n",
      "./datasets/WaterDro 100%[===================>]     361  --.-KB/s    in 0s      \n",
      "\n",
      "2023-10-15 01:39:30 (400 MB/s) - ‘./datasets/WaterDrop/metadata.json’ saved [361/361]\n",
      "\n",
      "--2023-10-15 01:39:30--  https://storage.googleapis.com/learning-to-simulate-complex-physics/Datasets/WaterDrop/train.tfrecord\n",
      "Resolving storage.googleapis.com (storage.googleapis.com)... 142.250.186.91, 142.250.186.187, 142.250.185.123, ...\n",
      "Connecting to storage.googleapis.com (storage.googleapis.com)|142.250.186.91|:443... connected.\n",
      "HTTP request sent, awaiting response... 200 OK\n",
      "Length: 4541246980 (4.2G) [application/octet-stream]\n",
      "Saving to: ‘./datasets/WaterDrop/train.tfrecord’\n",
      "\n",
      "./datasets/WaterDro 100%[===================>]   4.23G  30.0MB/s    in 2m 41s  \n",
      "\n",
      "2023-10-15 01:42:11 (26.9 MB/s) - ‘./datasets/WaterDrop/train.tfrecord’ saved [4541246980/4541246980]\n",
      "\n",
      "--2023-10-15 01:42:11--  https://storage.googleapis.com/learning-to-simulate-complex-physics/Datasets/WaterDrop/valid.tfrecord\n",
      "Resolving storage.googleapis.com (storage.googleapis.com)... 142.250.184.251, 142.250.184.219, 216.58.206.59, ...\n",
      "Connecting to storage.googleapis.com (storage.googleapis.com)|142.250.184.251|:443... connected.\n",
      "HTTP request sent, awaiting response... 200 OK\n",
      "Length: 137190408 (131M) [application/octet-stream]\n",
      "Saving to: ‘./datasets/WaterDrop/valid.tfrecord’\n",
      "\n",
      "./datasets/WaterDro 100%[===================>] 130.83M  28.8MB/s    in 5.1s    \n",
      "\n",
      "2023-10-15 01:42:17 (25.7 MB/s) - ‘./datasets/WaterDrop/valid.tfrecord’ saved [137190408/137190408]\n",
      "\n",
      "--2023-10-15 01:42:17--  https://storage.googleapis.com/learning-to-simulate-complex-physics/Datasets/WaterDrop/test.tfrecord\n",
      "Resolving storage.googleapis.com (storage.googleapis.com)... 172.217.16.155, 172.217.23.123, 142.250.74.219, ...\n",
      "Connecting to storage.googleapis.com (storage.googleapis.com)|172.217.16.155|:443... connected.\n",
      "HTTP request sent, awaiting response... 200 OK\n",
      "Length: 136885800 (131M) [application/octet-stream]\n",
      "Saving to: ‘./datasets/WaterDrop/test.tfrecord’\n",
      "\n",
      "./datasets/WaterDro 100%[===================>] 130.54M  30.5MB/s    in 5.0s    \n",
      "\n",
      "2023-10-15 01:42:22 (26.0 MB/s) - ‘./datasets/WaterDrop/test.tfrecord’ saved [136885800/136885800]\n",
      "\n"
     ]
    }
   ],
   "source": [
    "!mkdir -p ./datasets\n",
    "!bash ../gns_data/download_dataset.sh WaterDrop ./datasets"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "To avoid conflicting library dependencies, we recommend installing a second virtual environment only for the preprocessing. It will have the lightweight CPU version of TensorFlow, which is needed to open the original datasets."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Collecting tensorflow\n",
      "  Downloading tensorflow-2.14.0-cp310-cp310-manylinux_2_17_x86_64.manylinux2014_x86_64.whl (489.8 MB)\n",
      "\u001b[2K     \u001b[90m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\u001b[0m \u001b[32m489.8/489.8 MB\u001b[0m \u001b[31m5.1 MB/s\u001b[0m eta \u001b[36m0:00:00\u001b[0m00:01\u001b[0m00:02\u001b[0m\n",
      "\u001b[?25hCollecting tensorflow-datasets\n",
      "  Downloading tensorflow_datasets-4.9.3-py3-none-any.whl (5.0 MB)\n",
      "\u001b[2K     \u001b[90m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\u001b[0m \u001b[32m5.0/5.0 MB\u001b[0m \u001b[31m15.6 MB/s\u001b[0m eta \u001b[36m0:00:00\u001b[0m00:01\u001b[0m00:01\u001b[0m\n",
      "\u001b[?25hCollecting tensorflow-estimator<2.15,>=2.14.0\n",
      "  Downloading tensorflow_estimator-2.14.0-py2.py3-none-any.whl (440 kB)\n",
      "\u001b[2K     \u001b[90m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\u001b[0m \u001b[32m440.7/440.7 KB\u001b[0m \u001b[31m17.3 MB/s\u001b[0m eta \u001b[36m0:00:00\u001b[0m\n",
      "\u001b[?25hCollecting gast!=0.5.0,!=0.5.1,!=0.5.2,>=0.2.1\n",
      "  Downloading gast-0.5.4-py3-none-any.whl (19 kB)\n",
      "Collecting typing-extensions>=3.6.6\n",
      "  Using cached typing_extensions-4.8.0-py3-none-any.whl (31 kB)\n",
      "Collecting tensorflow-io-gcs-filesystem>=0.23.1\n",
      "  Downloading tensorflow_io_gcs_filesystem-0.34.0-cp310-cp310-manylinux_2_12_x86_64.manylinux2010_x86_64.whl (2.4 MB)\n",
      "\u001b[2K     \u001b[90m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\u001b[0m \u001b[32m2.4/2.4 MB\u001b[0m \u001b[31m13.8 MB/s\u001b[0m eta \u001b[36m0:00:00\u001b[0ma \u001b[36m0:00:01\u001b[0m\n",
      "\u001b[?25hCollecting h5py>=2.9.0\n",
      "  Downloading h5py-3.10.0-cp310-cp310-manylinux_2_17_x86_64.manylinux2014_x86_64.whl (4.8 MB)\n",
      "\u001b[2K     \u001b[90m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\u001b[0m \u001b[32m4.8/4.8 MB\u001b[0m \u001b[31m13.7 MB/s\u001b[0m eta \u001b[36m0:00:00\u001b[0ma \u001b[36m0:00:01\u001b[0m\n",
      "\u001b[?25hCollecting tensorboard<2.15,>=2.14\n",
      "  Downloading tensorboard-2.14.1-py3-none-any.whl (5.5 MB)\n",
      "\u001b[2K     \u001b[90m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\u001b[0m \u001b[32m5.5/5.5 MB\u001b[0m \u001b[31m15.8 MB/s\u001b[0m eta \u001b[36m0:00:00\u001b[0ma \u001b[36m0:00:01\u001b[0m\n",
      "\u001b[?25hCollecting six>=1.12.0\n",
      "  Using cached six-1.16.0-py2.py3-none-any.whl (11 kB)\n",
      "Collecting astunparse>=1.6.0\n",
      "  Using cached astunparse-1.6.3-py2.py3-none-any.whl (12 kB)\n",
      "Collecting wrapt<1.15,>=1.11.0\n",
      "  Downloading wrapt-1.14.1-cp310-cp310-manylinux_2_5_x86_64.manylinux1_x86_64.manylinux_2_17_x86_64.manylinux2014_x86_64.whl (77 kB)\n",
      "\u001b[2K     \u001b[90m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\u001b[0m \u001b[32m77.9/77.9 KB\u001b[0m \u001b[31m40.3 MB/s\u001b[0m eta \u001b[36m0:00:00\u001b[0m\n",
      "\u001b[?25hCollecting termcolor>=1.1.0\n",
      "  Using cached termcolor-2.3.0-py3-none-any.whl (6.9 kB)\n",
      "Collecting grpcio<2.0,>=1.24.3\n",
      "  Downloading grpcio-1.59.0-cp310-cp310-manylinux_2_17_x86_64.manylinux2014_x86_64.whl (5.3 MB)\n",
      "\u001b[2K     \u001b[90m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\u001b[0m \u001b[32m5.3/5.3 MB\u001b[0m \u001b[31m18.1 MB/s\u001b[0m eta \u001b[36m0:00:00\u001b[0m00:01\u001b[0m00:01\u001b[0m\n",
      "\u001b[?25hCollecting absl-py>=1.0.0\n",
      "  Downloading absl_py-2.0.0-py3-none-any.whl (130 kB)\n",
      "\u001b[2K     \u001b[90m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\u001b[0m \u001b[32m130.2/130.2 KB\u001b[0m \u001b[31m63.1 MB/s\u001b[0m eta \u001b[36m0:00:00\u001b[0m\n",
      "\u001b[?25hCollecting opt-einsum>=2.3.2\n",
      "  Using cached opt_einsum-3.3.0-py3-none-any.whl (65 kB)\n",
      "Requirement already satisfied: setuptools in ./venv_tf/lib/python3.10/site-packages (from tensorflow) (59.6.0)\n",
      "Collecting numpy>=1.23.5\n",
      "  Downloading numpy-1.26.1-cp310-cp310-manylinux_2_17_x86_64.manylinux2014_x86_64.whl (18.2 MB)\n",
      "\u001b[2K     \u001b[90m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\u001b[0m \u001b[32m18.2/18.2 MB\u001b[0m \u001b[31m15.9 MB/s\u001b[0m eta \u001b[36m0:00:00\u001b[0m00:01\u001b[0m00:01\u001b[0m\n",
      "\u001b[?25hCollecting flatbuffers>=23.5.26\n",
      "  Using cached flatbuffers-23.5.26-py2.py3-none-any.whl (26 kB)\n",
      "Collecting ml-dtypes==0.2.0\n",
      "  Using cached ml_dtypes-0.2.0-cp310-cp310-manylinux_2_17_x86_64.manylinux2014_x86_64.whl (1.0 MB)\n",
      "Collecting libclang>=13.0.0\n",
      "  Using cached libclang-16.0.6-py2.py3-none-manylinux2010_x86_64.whl (22.9 MB)\n",
      "Collecting packaging\n",
      "  Using cached packaging-23.2-py3-none-any.whl (53 kB)\n",
      "Collecting keras<2.15,>=2.14.0\n",
      "  Downloading keras-2.14.0-py3-none-any.whl (1.7 MB)\n",
      "\u001b[2K     \u001b[90m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\u001b[0m \u001b[32m1.7/1.7 MB\u001b[0m \u001b[31m12.0 MB/s\u001b[0m eta \u001b[36m0:00:00\u001b[0m00:01\u001b[0m00:01\u001b[0m\n",
      "\u001b[?25hCollecting google-pasta>=0.1.1\n",
      "  Using cached google_pasta-0.2.0-py3-none-any.whl (57 kB)\n",
      "Collecting protobuf!=4.21.0,!=4.21.1,!=4.21.2,!=4.21.3,!=4.21.4,!=4.21.5,<5.0.0dev,>=3.20.3\n",
      "  Downloading protobuf-4.24.4-cp37-abi3-manylinux2014_x86_64.whl (311 kB)\n",
      "\u001b[2K     \u001b[90m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\u001b[0m \u001b[32m311.6/311.6 KB\u001b[0m \u001b[31m22.2 MB/s\u001b[0m eta \u001b[36m0:00:00\u001b[0m\n",
      "\u001b[?25hCollecting dm-tree\n",
      "  Using cached dm_tree-0.1.8-cp310-cp310-manylinux_2_17_x86_64.manylinux2014_x86_64.whl (152 kB)\n",
      "Collecting requests>=2.19.0\n",
      "  Using cached requests-2.31.0-py3-none-any.whl (62 kB)\n",
      "Collecting array-record\n",
      "  Using cached array_record-0.4.1-py310-none-any.whl (3.0 MB)\n",
      "Collecting promise\n",
      "  Using cached promise-2.3-py3-none-any.whl\n",
      "Collecting etils[enp,epath,etree]>=0.9.0\n",
      "  Downloading etils-1.5.1-py3-none-any.whl (140 kB)\n",
      "\u001b[2K     \u001b[90m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\u001b[0m \u001b[32m140.5/140.5 KB\u001b[0m \u001b[31m44.4 MB/s\u001b[0m eta \u001b[36m0:00:00\u001b[0m\n",
      "\u001b[?25hCollecting tqdm\n",
      "  Using cached tqdm-4.66.1-py3-none-any.whl (78 kB)\n",
      "Collecting click\n",
      "  Using cached click-8.1.7-py3-none-any.whl (97 kB)\n",
      "Collecting toml\n",
      "  Using cached toml-0.10.2-py2.py3-none-any.whl (16 kB)\n",
      "Collecting tensorflow-metadata\n",
      "  Using cached tensorflow_metadata-1.14.0-py3-none-any.whl (28 kB)\n",
      "Collecting psutil\n",
      "  Using cached psutil-5.9.5-cp36-abi3-manylinux_2_12_x86_64.manylinux2010_x86_64.manylinux_2_17_x86_64.manylinux2014_x86_64.whl (282 kB)\n",
      "Collecting wheel<1.0,>=0.23.0\n",
      "  Using cached wheel-0.41.2-py3-none-any.whl (64 kB)\n",
      "Collecting importlib_resources\n",
      "  Downloading importlib_resources-6.1.0-py3-none-any.whl (33 kB)\n",
      "Collecting fsspec\n",
      "  Downloading fsspec-2023.9.2-py3-none-any.whl (173 kB)\n",
      "\u001b[2K     \u001b[90m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\u001b[0m \u001b[32m173.4/173.4 KB\u001b[0m \u001b[31m32.8 MB/s\u001b[0m eta \u001b[36m0:00:00\u001b[0m\n",
      "\u001b[?25hCollecting zipp\n",
      "  Downloading zipp-3.17.0-py3-none-any.whl (7.4 kB)\n",
      "Collecting idna<4,>=2.5\n",
      "  Using cached idna-3.4-py3-none-any.whl (61 kB)\n",
      "Collecting charset-normalizer<4,>=2\n",
      "  Downloading charset_normalizer-3.3.0-cp310-cp310-manylinux_2_17_x86_64.manylinux2014_x86_64.whl (138 kB)\n",
      "\u001b[2K     \u001b[90m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\u001b[0m \u001b[32m138.7/138.7 KB\u001b[0m \u001b[31m32.3 MB/s\u001b[0m eta \u001b[36m0:00:00\u001b[0m\n",
      "\u001b[?25hCollecting certifi>=2017.4.17\n",
      "  Using cached certifi-2023.7.22-py3-none-any.whl (158 kB)\n",
      "Collecting urllib3<3,>=1.21.1\n",
      "  Downloading urllib3-2.0.6-py3-none-any.whl (123 kB)\n",
      "\u001b[2K     \u001b[90m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\u001b[0m \u001b[32m123.8/123.8 KB\u001b[0m \u001b[31m30.4 MB/s\u001b[0m eta \u001b[36m0:00:00\u001b[0m\n",
      "\u001b[?25hCollecting markdown>=2.6.8\n",
      "  Downloading Markdown-3.5-py3-none-any.whl (101 kB)\n",
      "\u001b[2K     \u001b[90m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\u001b[0m \u001b[32m101.7/101.7 KB\u001b[0m \u001b[31m47.6 MB/s\u001b[0m eta \u001b[36m0:00:00\u001b[0m\n",
      "\u001b[?25hCollecting google-auth<3,>=1.6.3\n",
      "  Downloading google_auth-2.23.3-py2.py3-none-any.whl (182 kB)\n",
      "\u001b[2K     \u001b[90m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\u001b[0m \u001b[32m182.3/182.3 KB\u001b[0m \u001b[31m32.5 MB/s\u001b[0m eta \u001b[36m0:00:00\u001b[0m\n",
      "\u001b[?25hCollecting tensorboard-data-server<0.8.0,>=0.7.0\n",
      "  Using cached tensorboard_data_server-0.7.1-py3-none-manylinux2014_x86_64.whl (6.6 MB)\n",
      "Collecting google-auth-oauthlib<1.1,>=0.5\n",
      "  Using cached google_auth_oauthlib-1.0.0-py2.py3-none-any.whl (18 kB)\n",
      "Collecting werkzeug>=1.0.1\n",
      "  Downloading werkzeug-3.0.0-py3-none-any.whl (226 kB)\n",
      "\u001b[2K     \u001b[90m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\u001b[0m \u001b[32m226.6/226.6 KB\u001b[0m \u001b[31m22.6 MB/s\u001b[0m eta \u001b[36m0:00:00\u001b[0m\n",
      "\u001b[?25hCollecting googleapis-common-protos<2,>=1.52.0\n",
      "  Downloading googleapis_common_protos-1.61.0-py2.py3-none-any.whl (230 kB)\n",
      "\u001b[2K     \u001b[90m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\u001b[0m \u001b[32m230.9/230.9 KB\u001b[0m \u001b[31m20.1 MB/s\u001b[0m eta \u001b[36m0:00:00\u001b[0m\n",
      "\u001b[?25hCollecting absl-py>=1.0.0\n",
      "  Using cached absl_py-1.4.0-py3-none-any.whl (126 kB)\n",
      "Collecting protobuf!=4.21.0,!=4.21.1,!=4.21.2,!=4.21.3,!=4.21.4,!=4.21.5,<5.0.0dev,>=3.20.3\n",
      "  Using cached protobuf-3.20.3-cp310-cp310-manylinux_2_12_x86_64.manylinux2010_x86_64.whl (1.1 MB)\n",
      "Collecting cachetools<6.0,>=2.0.0\n",
      "  Using cached cachetools-5.3.1-py3-none-any.whl (9.3 kB)\n",
      "Collecting pyasn1-modules>=0.2.1\n",
      "  Using cached pyasn1_modules-0.3.0-py2.py3-none-any.whl (181 kB)\n",
      "Collecting rsa<5,>=3.1.4\n",
      "  Using cached rsa-4.9-py3-none-any.whl (34 kB)\n",
      "Collecting requests-oauthlib>=0.7.0\n",
      "  Using cached requests_oauthlib-1.3.1-py2.py3-none-any.whl (23 kB)\n",
      "Collecting MarkupSafe>=2.1.1\n",
      "  Using cached MarkupSafe-2.1.3-cp310-cp310-manylinux_2_17_x86_64.manylinux2014_x86_64.whl (25 kB)\n",
      "Collecting pyasn1<0.6.0,>=0.4.6\n",
      "  Using cached pyasn1-0.5.0-py2.py3-none-any.whl (83 kB)\n",
      "Collecting oauthlib>=3.0.0\n",
      "  Using cached oauthlib-3.2.2-py3-none-any.whl (151 kB)\n",
      "Installing collected packages: libclang, flatbuffers, dm-tree, zipp, wrapt, wheel, urllib3, typing-extensions, tqdm, toml, termcolor, tensorflow-io-gcs-filesystem, tensorflow-estimator, tensorboard-data-server, six, pyasn1, psutil, protobuf, packaging, oauthlib, numpy, MarkupSafe, markdown, keras, importlib_resources, idna, grpcio, gast, fsspec, etils, click, charset-normalizer, certifi, cachetools, absl-py, werkzeug, rsa, requests, pyasn1-modules, promise, opt-einsum, ml-dtypes, h5py, googleapis-common-protos, google-pasta, astunparse, tensorflow-metadata, requests-oauthlib, google-auth, google-auth-oauthlib, array-record, tensorboard, tensorflow-datasets, tensorflow\n",
      "Successfully installed MarkupSafe-2.1.3 absl-py-1.4.0 array-record-0.4.1 astunparse-1.6.3 cachetools-5.3.1 certifi-2023.7.22 charset-normalizer-3.3.0 click-8.1.7 dm-tree-0.1.8 etils-1.5.1 flatbuffers-23.5.26 fsspec-2023.9.2 gast-0.5.4 google-auth-2.23.3 google-auth-oauthlib-1.0.0 google-pasta-0.2.0 googleapis-common-protos-1.61.0 grpcio-1.59.0 h5py-3.10.0 idna-3.4 importlib_resources-6.1.0 keras-2.14.0 libclang-16.0.6 markdown-3.5 ml-dtypes-0.2.0 numpy-1.26.1 oauthlib-3.2.2 opt-einsum-3.3.0 packaging-23.2 promise-2.3 protobuf-3.20.3 psutil-5.9.5 pyasn1-0.5.0 pyasn1-modules-0.3.0 requests-2.31.0 requests-oauthlib-1.3.1 rsa-4.9 six-1.16.0 tensorboard-2.14.1 tensorboard-data-server-0.7.1 tensorflow-2.14.0 tensorflow-datasets-4.9.3 tensorflow-estimator-2.14.0 tensorflow-io-gcs-filesystem-0.34.0 tensorflow-metadata-1.14.0 termcolor-2.3.0 toml-0.10.2 tqdm-4.66.1 typing-extensions-4.8.0 urllib3-2.0.6 werkzeug-3.0.0 wheel-0.41.2 wrapt-1.14.1 zipp-3.17.0\n"
     ]
    }
   ],
   "source": [
    "!python3 -m venv venv_tf\n",
    "!venv_tf/bin/pip install tensorflow tensorflow-datasets"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Finally, we transform the `*.tfrecord` files to our `*.h5` format. We automatically add two further features to the `metadata.json` file, namely:\n",
    "\n",
    "- `num_particles_max`: needed for jitability (via padding)\n",
    "- `periodic_boundary_conditions`: specifying the type of boundary conditions per dimension\n",
    "\n",
    "Note: This might take a few minutes and will double the space taken by the dataset."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "2023-10-15 01:44:07.609558: E tensorflow/compiler/xla/stream_executor/cuda/cuda_dnn.cc:9342] Unable to register cuDNN factory: Attempting to register factory for plugin cuDNN when one has already been registered\n",
      "2023-10-15 01:44:07.609617: E tensorflow/compiler/xla/stream_executor/cuda/cuda_fft.cc:609] Unable to register cuFFT factory: Attempting to register factory for plugin cuFFT when one has already been registered\n",
      "2023-10-15 01:44:07.609667: E tensorflow/compiler/xla/stream_executor/cuda/cuda_blas.cc:1518] Unable to register cuBLAS factory: Attempting to register factory for plugin cuBLAS when one has already been registered\n",
      "2023-10-15 01:44:08.294080: W tensorflow/compiler/tf2tensorrt/utils/py_utils.cc:38] TF-TRT Warning: Could not find TensorRT\n",
      "Start conversion of ./datasets/WaterDrop/train.tfrecord to .h5\n",
      "2023-10-15 01:44:09.073610: E tensorflow/compiler/xla/stream_executor/cuda/cuda_driver.cc:268] failed call to cuInit: CUDA_ERROR_NO_DEVICE: no CUDA-capable device is detected\n",
      "Finish conversion to ./datasets/WaterDrop/train.h5\n",
      "Start conversion of ./datasets/WaterDrop/valid.tfrecord to .h5\n",
      "Finish conversion to ./datasets/WaterDrop/valid.h5\n",
      "Start conversion of ./datasets/WaterDrop/test.tfrecord to .h5\n",
      "Finish conversion to ./datasets/WaterDrop/test.h5\n",
      "Max number of particles in valid.h5: 1044\n",
      "Max number of particles in train.h5: 1108\n",
      "Max number of particles in test.h5: 1108\n"
     ]
    }
   ],
   "source": [
    "!venv_tf/bin/python ../gns_data/tfrecord_to_h5.py --dataset-path=./datasets/WaterDrop"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "If everything worked fine and you see the `*.h5` files, you can remove the `*.tfrecord` files and the virtual environment."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [],
   "source": [
    "!rm ./datasets/WaterDrop/*.tfrecord\n",
    "!rm -r venv_tf"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Inspect the data\n",
    "\n",
    "Because the number of particles in this dataset varies, we are forced to use `matscipy` as the neighbors search implementation. Matscipy runs on the CPU and can efficiently handle systems with up to milions of particles. To integrate this function into our jit-able codebase, we use `jax.pure_callback()` and pad the non-existing particle entries up to `num_particles_max`."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "This is a 2D dataset from ./datasets/WaterDrop.\n",
      "Train frames have shape (1108, 7, 2) (n_nodes, seq_len, xy pos).\n",
      "Val frames have shape (1108, 26, 2) (n_nodes, rollout, xy pos).\n",
      "And particle types have shape (1108,) (n_nodes,).\n",
      "Total of 995000 train frames and 1140 val frames.\n",
      "\n"
     ]
    }
   ],
   "source": [
    "data_train = lagrangebench.data.H5Dataset(\n",
    "    split=\"train\", \n",
    "    dataset_path=\"./datasets/WaterDrop\",\n",
    "    name=\"waterdrop2d\",\n",
    "    nl_backend=\"matscipy\"\n",
    ")\n",
    "\n",
    "data_valid = lagrangebench.data.H5Dataset(\n",
    "    split=\"valid\", \n",
    "    dataset_path=\"./datasets/WaterDrop\",\n",
    "    name=\"waterdrop2d\",\n",
    "    n_rollout_steps=20,\n",
    "    nl_backend=\"matscipy\"\n",
    ")\n",
    "# We can compute the `split_valid_traj_into_n=38` by:\n",
    "# 1. from [1], Appendix B.1, each trajectory has a length of 1000 steps\n",
    "# 2. Later, we set `n_rollout_steps=20` for the validation loss\n",
    "# 3. We use 5 historic velocities (or 6 historic positions) \n",
    "# => This leads to splitting each training trajectory into \n",
    "# 1000 // (20 + 6) = 38 validation instances\n",
    "\n",
    "print(\n",
    "    f\"This is a {data_train.metadata['dim']}D dataset from {data_train.dataset_path}.\\n\"\n",
    "    f\"Train frames have shape {data_train[0][0].shape} (n_nodes, seq_len, xy pos).\\n\"\n",
    "    f\"Val frames have shape {data_valid[0][0].shape} (n_nodes, rollout, xy pos).\\n\"\n",
    "    f\"And particle types have shape {data_train[0][1].shape} (n_nodes,).\\n\"\n",
    "    f\"Total of {len(data_train)} train frames and {len(data_valid)} val frames.\\n\"\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Visualize slices from the first trajectory"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "image/png": "[encoded data, 41072 characters]",
      "text/plain": [
       "<Figure size 2000x400 with 5 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "image/png": "[encoded data, 42796 characters]",
      "text/plain": [
       "<Figure size 2000x400 with 5 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "image/png": "[encoded data, 39488 characters]",
      "text/plain": [
       "<Figure size 2000x400 with 5 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "\n",
    "bounds = np.array(data_train.metadata[\"bounds\"])\n",
    "\n",
    "# from [1], Appendix B.1 we know that the trajectory length is 1000\n",
    "# If we take indices < 994, we will see samples from the first trajectory\n",
    "for j in [0, 500, 993]:\n",
    "    sample = data_train[j]\n",
    "    \n",
    "    # visualize 5 consecutive frames\n",
    "    fig, axs = plt.subplots(1, 5, figsize=(20, 4))\n",
    "    for i in range(5):\n",
    "        mask = sample[1] != -1\n",
    "        axs[i].scatter(sample[0][mask, i, 0], sample[0][mask, i, 1], s=1)\n",
    "        axs[i].set_xlim(bounds[0])\n",
    "        axs[i].set_ylim(bounds[1])\n",
    "    plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Here we visualize random frames from the dataset."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "image/png": "[encoded data, 189136 characters]",
      "text/plain": [
       "<Figure size 2000x1200 with 15 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "np.random.seed(42)\n",
    "frame_nums = np.random.randint(0, len(data_train), 15)\n",
    "\n",
    "# visualize 5 consecutive frames\n",
    "fig, axs = plt.subplots(3, 5, figsize=(20, 12))\n",
    "for ax, i in zip(axs.flatten(), frame_nums):\n",
    "    sample = data_train[i]\n",
    "    mask = sample[1] != -1\n",
    "    \n",
    "    ax.scatter(sample[0][mask, 0, 0], sample[0][mask, 0, 1], s=1, \n",
    "               label=f\"seed={i}, N={np.sum(mask)}\")\n",
    "    ax.legend()\n",
    "    ax.set_xlim(bounds[0])\n",
    "    ax.set_ylim(bounds[1])\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## GNN training\n",
    "\n",
    "This works as usual. See [`notebooks/tutorial.ipynb`](./tutorial.ipynb) for more details."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [],
   "source": [
    "def gns(x):\n",
    "    return lagrangebench.GNS(\n",
    "        particle_dimension=data_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)\n",
    "\n",
    "gns = hk.without_apply_rng(hk.transform_with_state(gns))\n",
    "\n",
    "noise_std = 3e-4\n",
    "bounds = np.array(data_train.metadata[\"bounds\"])\n",
    "box = bounds[:, 1] - bounds[:, 0]\n",
    "\n",
    "case = lagrangebench.case_builder(\n",
    "    box=box,  \n",
    "    metadata=data_train.metadata,\n",
    "    input_seq_length=6,\n",
    "    isotropic_norm=False,\n",
    "    noise_std=noise_std,\n",
    ")\n",
    "\n",
    "trainer = lagrangebench.Trainer(\n",
    "    model=gns,\n",
    "    case=case,\n",
    "    data_train=data_train,\n",
    "    data_eval=data_valid,\n",
    "    noise_std=noise_std,\n",
    "    metrics=[\"mse\"],\n",
    "    n_rollout_steps=20,\n",
    "    eval_n_trajs=2,\n",
    "    lr_start=5e-4,\n",
    "    log_steps=10,\n",
    "    eval_steps=50,\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Reallocate neighbors list (2, 170875) at step 0\n",
      "To list (2, 442352)\n",
      "000, train/loss: 4.31286.\n",
      "Reallocate neighbors list (2, 442352) at step 3\n",
      "To list (2, 684860)\n",
      "Reallocate neighbors list (2, 684860) at step 6\n",
      "To list (2, 915872)\n",
      "010, train/loss: 0.12886.\n",
      "020, train/loss: 0.40341.\n",
      "030, train/loss: 0.04893.\n",
      "040, train/loss: 0.05427.\n",
      "050, train/loss: 12.54240.\n",
      "{'val/loss': 1.0891471447393997e-05, 'val/mse1': 1.145147032755034e-09, 'val/mse5': 8.628529357679326e-08, 'val/mse10': 9.226690735886223e-07, 'val/stdloss': 9.959435374184977e-07, 'val/stdmse1': 1.827232254392186e-10, 'val/stdmse5': 1.3568257628548963e-08, 'val/stdmse10': 1.5360808447439922e-07}\n",
      "060, train/loss: 0.09790.\n",
      "070, train/loss: 0.74478.\n",
      "080, train/loss: 0.10917.\n",
      "090, train/loss: 0.21384.\n",
      "100, train/loss: 0.74520.\n",
      "{'val/loss': 1.3067739018879365e-05, 'val/mse1': 1.373287872308282e-09, 'val/mse5': 1.0317071641452458e-07, 'val/mse10': 1.100612792015454e-06, 'val/stdloss': 6.564368959516287e-08, 'val/stdmse1': 8.536110707169087e-11, 'val/stdmse5': 5.950870018978094e-09, 'val/stdmse10': 7.309478178285644e-08}\n"
     ]
    }
   ],
   "source": [
    "params, state, _ = trainer(step_max=100)"
   ]
  }
 ],
 "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
}
