{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "ba8371b3",
   "metadata": {},
   "source": [
    "# Loading and Viewing BubbleML"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "f60b058a",
   "metadata": {},
   "outputs": [],
   "source": [
    "import h5py\n",
    "import numpy as np\n",
    "import matplotlib.pyplot as plt\n",
    "import numba as nb"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "dab04e01",
   "metadata": {},
   "source": [
    "### Create an h5py file\n",
    "In this case, we are using the simulation run with a wall temperature of 100 degrees Celsius.\n",
    "The simulation file includes a list of keys:\n",
    "  1. `dfun` is a signed distance function from the bubbble interface\n",
    "  2. `pressure` is the pressure gradient\n",
    "  3. `temperature` is the temperature map\n",
    "  4. `velx` is the velocity in the x direction\n",
    "  5. `vely` is the velocity in the y direction\n",
    "  6. `x` and `y `are coordinate grids\n",
    "  7. `int/real-runtime-params` are metadata associated with the simulation run.\n",
    "     This includes things like the Reynold's number, simulation dimensions, etc."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "2c835c16",
   "metadata": {},
   "outputs": [],
   "source": [
    "twall_100 = h5py.File('Twall-100.hdf5', 'r')\n",
    "for idx, key in enumerate(twall_100.keys()):\n",
    "    print(f'{idx + 1}. {key}')"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "7d5a1b53",
   "metadata": {},
   "source": [
    "### Tensor sizes\n",
    "Each tensor is laid out [T x Y x X] (time, y-direction, x-direction). In this case, there are 200 time steps and the domain resolution is 48x48"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "9a0bd89d",
   "metadata": {},
   "outputs": [],
   "source": [
    "[time_res, y_res, x_res] = twall_100['temperature'][:].shape\n",
    "\n",
    "print(f'# Timesteps: {time_res}')\n",
    "print(f'Domain resolution: {y_res} x {x_res}')"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "dbf12cad",
   "metadata": {},
   "source": [
    "### Visualizing the different fields\n",
    "The data can be easily loaded into numpy (or torch, tensorflow, etc) arrays and visualized with matplotlib."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "15334ef1",
   "metadata": {},
   "outputs": [],
   "source": [
    "temp = twall_100['temperature'][:]\n",
    "velx = twall_100['velx'][:]\n",
    "vely = twall_100['vely'][:]\n",
    "pres = twall_100['pressure'][:]\n",
    "dfun = twall_100['dfun'][:]\n",
    "\n",
    "# compute the velocity magnitude\n",
    "mag = np.sqrt(velx**2 + vely**2)\n",
    "\n",
    "# plot the 50-th timestep for each variable\n",
    "fig, ax = plt.subplots(1, 4, figsize=(12, 5))\n",
    "\n",
    "data = {\n",
    "    'Temperature': temp[50],\n",
    "    'Velocity Mag.': mag[50],\n",
    "    'Pressure Grad.': pres[50],\n",
    "    'Distance Func.': dfun[50]\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.5)\n",
    "    ax[idx].set_title(key)\n",
    "    ax[idx].set_xticks([])\n",
    "    ax[idx].set_yticks([])"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "e90f0013-12c9-4d90-8222-a0f9992eb8c3",
   "metadata": {},
   "source": [
    "### Visualizing different timesteps\n",
    "By progressively indexing along the time axis (dimension 0), we are able to see the progression as it leaves the heater surface."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "9fd72bae-590f-4a1e-aaf0-eae31f0e3a79",
   "metadata": {},
   "outputs": [],
   "source": [
    "temp = twall_100['temperature'][:]\n",
    "mag = np.sqrt(twall_100['velx'][:]**2 + twall_100['vely'][:]**2)\n",
    "\n",
    "timesteps = range(40, 52, 2)\n",
    "\n",
    "fig, ax = plt.subplots(2, len(timesteps), figsize=(15, 5))\n",
    "\n",
    "for idx, step in enumerate(timesteps):\n",
    "    ax[0, idx].imshow(np.flipud(temp[step]))\n",
    "    ax[1, idx].imshow(np.flipud(mag[step]))\n",
    "    for row in range(2):\n",
    "        ax[row, idx].set_title(f'Step {step}')\n",
    "        ax[row, idx].set_xticks([])\n",
    "        ax[row, idx].set_yticks([])\n",
    "    ax[0,0].set_ylabel('Temperature')\n",
    "    ax[1,0].set_ylabel('Velocity Mag.')"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "0fd957ba",
   "metadata": {},
   "source": [
    "### Using dfun\n",
    "dfun is a *signed distance function* to the liquid-vapor interfaces. \n",
    "dfun > 0 means the point is in vapor, dfun < 0 means the point is in liquid. It is also a convenient way to identify the bubble interfaces. \n",
    "\n",
    "The function `get_interface_mask` matches the function used in the simulations. [See equation (8)](https://doi.org/10.1016/j.ijmultiphaseflow.2019.103099). We use numba to jit compile this function for performance."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "3c71f0fe",
   "metadata": {},
   "outputs": [],
   "source": [
    "@nb.njit\n",
    "def get_interface_mask(dgrid):\n",
    "    r\"\"\" heavy-side function to determine the bubble interfaces\n",
    "    \"\"\"\n",
    "    interface = np.zeros(dgrid.shape).astype(np.bool_)\n",
    "    [rows, cols] = dgrid.shape\n",
    "    for i in range(rows):\n",
    "        for j in range(cols):\n",
    "            adj = ((i < rows - 1 and dgrid[i][j] * dgrid[i+1, j  ] <= 0) or\n",
    "                   (i > 0 and dgrid[i][j] * dgrid[i-1, j  ] <= 0) or\n",
    "                   (j < cols - 1 and dgrid[i][j] * dgrid[i,   j+1] <= 0) or\n",
    "                   (j > 0 and dgrid[i][j] * dgrid[i,   j-1] <= 0))\n",
    "            interface[i][j] = adj\n",
    "    return interface"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "d6b08304",
   "metadata": {},
   "outputs": [],
   "source": [
    "fig, ax = plt.subplots(1, 5, figsize=(14, 5))\n",
    "\n",
    "bubbles = dfun[50] >= 0  # vapor phase has non-negative distance.\n",
    "liquid = dfun[50] < 0    # Liquid has negative distance.\n",
    "interface = get_interface_mask(dfun[50])\n",
    "\n",
    "data = {\n",
    "    'Temperature': temp[50],\n",
    "    'Velocity Mag.': mag[50],\n",
    "    'Interface': interface,\n",
    "    'Bubbles': bubbles,\n",
    "    'Liquid': liquid\n",
    "}\n",
    "\n",
    "for idx, (key, im) in enumerate(data.items()):\n",
    "    ax[idx].imshow(np.flipud(im))\n",
    "    ax[idx].set_title(key)\n",
    "    ax[idx].set_xticks([])\n",
    "    ax[idx].set_yticks([])"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "2fe398bb-8a38-4eec-972f-56df926012ba",
   "metadata": {},
   "source": [
    "### Accessing the metadata\n",
    "\n",
    "There is a lot of metadata associated with each simulation. Pointers to some of the important metadata is listed in our [dataset documentation](https://github.com/HPCForge/BubbleML/blob/main/bubbleml_data/DOCS.md). The metadata is stored as a numpy array of tuples. Each tuple contains an array of bytes (the key) and a float (the value). The metadata stores critical information for the simulation. Some of these values (like the reynolds nunmber) will be necessary for training physics informed models."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "d5b86131-82ff-43ba-8ab3-47506f7f1db2",
   "metadata": {},
   "outputs": [],
   "source": [
    "real_runtime_params = twall_100['real-runtime-params'][:]\n",
    "key0, val0 = real_runtime_params[0]\n",
    "\n",
    "print(f'Metadata size: {real_runtime_params.shape}')\n",
    "print(f'Key type: {type(key0)}')\n",
    "print(f'Val type: {type(val0)}')\n",
    "\n",
    "def key_to_str(key):\n",
    "    # convert byte string to a standard python utf-8 string.\n",
    "    return key.decode('utf-8').strip()\n",
    "\n",
    "# Convert to a dict of (string, float64)\n",
    "runtime_param_dict = dict([(key_to_str(key), val) for (key, val) in real_runtime_params])\n",
    "\n",
    "# print the reynolds number\n",
    "inv_reynolds = runtime_param_dict['ins_invreynolds']\n",
    "print(f'Reynolds Number: {1 / inv_reynolds}')"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "27c9d2de-fbc4-4774-8d1a-efeba7df5fea",
   "metadata": {},
   "source": [
    "### Getting the Domain Size\n",
    "The simulations in BubbleML are not all the same size. So, it can be beneficial to know how to access the true spatial dimensions. \n",
    "Here, we read the xy-extents from the metadata and set them as axis ticks on an image."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "e7e74bdd-5d4f-44fc-9daa-0b917c6186ef",
   "metadata": {},
   "outputs": [],
   "source": [
    "xmin, xmax = runtime_param_dict['xmin'], runtime_param_dict['xmax']\n",
    "ymin, ymax = runtime_param_dict['ymin'], runtime_param_dict['ymax']\n",
    "\n",
    "print(f'x extents: {xmin}, {xmax}')\n",
    "print(f'y extents: {ymin}, {ymax}')\n",
    "\n",
    "plt.imshow(np.flipud(temp[50]), extent=[xmin,xmax,ymin,ymax])\n",
    "plt.show()"
   ]
  }
 ],
 "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
}