{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "%load_ext autoreload\n",
    "%autoreload 2"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [],
   "source": [
    "import sys\n",
    "import os\n",
    "parentdir = os.path.dirname(os.getcwd())\n",
    "sys.path.append(parentdir)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "True"
      ]
     },
     "execution_count": 3,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "import dotenv\n",
    "# load environment variables from `.env` file if it exists\n",
    "# recursively searches for `.env` in all folders starting from work dir\n",
    "dotenv.load_dotenv()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [],
   "source": [
    "# images list return, save\n",
    "import h5py\n",
    "import matplotlib.pyplot as plt\n",
    "from einops import rearrange\n",
    "from phi.flow import *\n",
    "import matplotlib.pyplot as plt\n",
    "import numpy as np\n",
    "from phi.vis import *\n",
    "from src.utils import resolve_path"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "/datasets/work/mlai-fsp-pdeemu/work/PDEBench/working/ns_sim_cmd/2022-06-04/14-20-40/device=GPU,jit=True,label=basic,mode=slurm,n_batch=10,profile=False,save_gif=False,save_images=False,seed=99,sim_name=nsi_2d/98/basic-99.h5\n"
     ]
    }
   ],
   "source": [
    "data_path = resolve_path('${WORKING_DIR}/*/*/*/*/*/*.h5', idx=-1, unique=False)\n",
    "print(data_path)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "<KeysViewHDF5 ['force', 'particles', 't', 'velocity']>"
      ]
     },
     "execution_count": 6,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "data_f = h5py.File(data_path, 'r')\n",
    "data_f.keys()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "(10, 99901, 256, 256, 1)"
      ]
     },
     "execution_count": 7,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "data_f['particles'].shape"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 16,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "image/png": "[encoded data, 217348 characters]",
      "text/plain": [
       "<Figure size 864x172.8 with 1 Axes>"
      ]
     },
     "metadata": {
      "needs_background": "light"
     },
     "output_type": "display_data"
    }
   ],
   "source": [
    "\n",
    "columns = 5\n",
    "fsize = 12\n",
    "\n",
    "arr = data_f['particles'][0, :columns, :, :, 0]\n",
    "\n",
    "fig = plt.figure(figsize=(fsize, fsize/columns))\n",
    "plt.imshow(rearrange(arr, 't x y -> x (t y)'))\n",
    "plt.gca().set_axis_off()\n",
    "plt.tight_layout(pad = 1)\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "image/png": "iVBORw0KGgoAAAANSUhEUgAAAwoAAACnCAYAAACxUquMAAAAOXRFWHRTb2Z0d2FyZQBNYXRwbG90bGliIHZlcnNpb24zLjUuMiwgaHR0cHM6Ly9tYXRwbG90bGliLm9yZy8qNh9FAAAACXBIWXMAAAsTAAALEwEAmpwYAAADaklEQVR4nO3csQ3CUAxAwfyI1RiBKTMCu2FapCdoE6G71o3bJ0teM7MBAAB82s9eAAAAuB6hAAAAhFAAAABCKAAAACEUAACAEAoAAEDcfg3v+8PvVAAA+FPP17G+zVwUAACAEAoAAEAIBQAAIIQCAAAQQgEAAAihAAAAhFAAAABCKAAAACEUAACAEAoAAEAIBQAAIIQCAAAQQgEAAAihAAAAhFAAAABCKAAAACEUAACAEAoAAEAIBQAAIIQCAAAQQgEAAAihAAAAhFAAAABCKAAAACEUAACAEAoAAEAIBQAAIIQCAAAQQgEAAAihAAAAhFAAAABCKAAAACEUAACAEAoAAEAIBQAAIIQCAAAQQgEAAAihAAAAhFAAAABCKAAAACEUAACAEAoAAEAIBQAAIIQCAAAQQgEAAAihAAAAhFAAAABCKAAAACEUAACAEAoAAEAIBQAAIIQCAAAQQgEAAAihAAAAhFAAAABCKAAAACEUAACAEAoAAEAIBQAAIIQCAAAQQgEAAAihAAAAhFAAAABCKAAAACEUAACAEAoAAEAIBQAAIIQCAAAQQgEAAAihAAAAhFAAAABCKAAAACEUAACAEAoAAEAIBQAAIIQCAAAQQgEAAAihAAAAhFAAAABCKAAAACEUAACAEAoAAEAIBQAAIIQCAAAQQgEAAAihAAAAhFAAAABCKAAAACEUAACAEAoAAEAIBQAAIIQCAAAQQgEAAAihAAAAhFAAAABCKAAAACEUAACAEAoAAEAIBQAAIIQCAAAQQgEAAAihAAAAhFAAAABCKAAAACEUAACAEAoAAEAIBQAAIIQCAAAQQgEAAAihAAAAhFAAAABCKAAAACEUAACAEAoAAEAIBQAAIIQCAAAQQgEAAAihAAAAhFAAAABCKAAAACEUAACAEAoAAEAIBQAAIIQCAAAQQgEAAAihAAAAhFAAAABCKAAAACEUAACAEAoAAEAIBQAAIIQCAAAQQgEAAAihAAAAhFAAAABCKAAAACEUAACAEAoAAEAIBQAAIIQCAAAQQgEAAAihAAAAhFAAAABCKAAAACEUAACAEAoAAEAIBQAAIIQCAAAQQgEAAAihAAAAhFAAAABCKAAAALFm5uwdAACAi3FRAAAAQigAAAAhFAAAgBAKAABACAUAACCEAgAAEG9GKAtJmq95tAAAAABJRU5ErkJggg==",
      "text/plain": [
       "<Figure size 864x172.8 with 1 Axes>"
      ]
     },
     "metadata": {
      "needs_background": "light"
     },
     "output_type": "display_data"
    }
   ],
   "source": [
    "\n",
    "columns = 5\n",
    "fsize = 12\n",
    "\n",
    "arr = data_f['particles'][0, 1000:1000+columns, :, :, 0]\n",
    "\n",
    "fig = plt.figure(figsize=(fsize, fsize/columns))\n",
    "plt.imshow(rearrange(arr, 't x y -> x (t y)'))\n",
    "plt.gca().set_axis_off()\n",
    "plt.tight_layout(pad = 1)\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 16,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "image/png": "[encoded data, 286624 characters]",
      "text/plain": [
       "<Figure size 864x172.8 with 1 Axes>"
      ]
     },
     "metadata": {
      "needs_background": "light"
     },
     "output_type": "display_data"
    }
   ],
   "source": [
    "\n",
    "columns = 5\n",
    "fsize = 12\n",
    "\n",
    "arr = data_f['force'][:columns, :, :, 0]\n",
    "\n",
    "fig = plt.figure(figsize=(fsize, fsize/columns))\n",
    "plt.imshow(rearrange(arr, 'b x y -> x (b y)'))\n",
    "plt.gca().set_axis_off()\n",
    "plt.tight_layout(pad = 1)\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "interpreter": {
   "hash": "38db623c9a8c22d26705f5e96929a68995fd63f2ffb05d6c4e502eb3b9fc3e8d"
  },
  "kernelspec": {
   "display_name": "Python 3.9.4 ('venv2': 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.9.4"
  },
  "orig_nbformat": 4,
  "vscode": {
   "interpreter": {
    "hash": "7a0d8eacbb66928d1b3dec5c8bc3662d91007e449cc8cc497b49c1d19e71d254"
   }
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}
