{
 "cells": [
  {
   "attachments": {},
   "cell_type": "markdown",
   "id": "f4206d20",
   "metadata": {},
   "source": [
    "# IPU RAY LIBRARY TESTING\n",
    "\n",
    "This notebook contains instructions to configure, build, run, and test a Poplar C++ Ray/Path Tracer for Graphcore IPUs. This serves to both test the application and to document how test/debug.\n",
    "\n",
    "## Testing Method\n",
    "\n",
    "There are automated tests and interactive ones. We will run the automated tests first, then the interactive ones. We recommend you click \"run all\" in the notebook to execute everything and then read through to understand whether everything is working as expected.\n",
    "\n",
    "### Build the Code\n",
    "\n",
    "This notebook assumes you are starting from a clean checkout. The following cell configures and builds everything. The build uses CMake:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "5de4cfb7-1c6e-4ad1-890d-dcd211334bfe",
   "metadata": {
    "execution": {
     "iopub.execute_input": "2023-01-10T16:25:04.773610Z",
     "iopub.status.busy": "2023-01-10T16:25:04.773333Z",
     "iopub.status.idle": "2023-01-10T16:25:05.554944Z",
     "shell.execute_reply": "2023-01-10T16:25:05.554395Z",
     "shell.execute_reply.started": "2023-01-10T16:25:04.773583Z"
    }
   },
   "outputs": [],
   "source": [
    "!mkdir -p build\n",
    "%cd build\n",
    "!cmake -Wno-dev -G Ninja ..\n",
    "!ninja -j64"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "id": "d26fd596",
   "metadata": {},
   "source": [
    "## Automated Tests\n",
    "\n",
    "First we can check the automated tests are ok:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "50307127",
   "metadata": {},
   "outputs": [],
   "source": [
    "!./tests --show_progress"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "id": "5e79d8cd-4cca-49d4-9df4-fff6e3baf42d",
   "metadata": {},
   "source": [
    "### Test the \"trace\" Program\n",
    "\n",
    "Now check we can run the application by rendering a path traced image of the Cornell Box:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "8f2b07c6-e572-435d-a9c8-a8aad8beaa67",
   "metadata": {
    "execution": {
     "iopub.execute_input": "2023-01-10T16:25:05.556220Z",
     "iopub.status.busy": "2023-01-10T16:25:05.555930Z",
     "iopub.status.idle": "2023-01-10T16:25:31.152592Z",
     "shell.execute_reply": "2023-01-10T16:25:31.151962Z",
     "shell.execute_reply.started": "2023-01-10T16:25:05.556200Z"
    }
   },
   "outputs": [],
   "source": [
    "!./trace -w 720 -h 720 --render-mode path-trace --visualise rgb --samples 1000 --ipus 4 --ipu-only --scene box-simple"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "id": "7678c0fc-2116-4344-bc5d-bb875284ab00",
   "metadata": {},
   "source": [
    "The output image is high dynamic range (HDR) in EXR format. We can make a function\n",
    "to perform a quick tone-mapping and display the resulting image in Python:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "784f56d3",
   "metadata": {
    "execution": {
     "iopub.execute_input": "2023-01-10T16:25:31.154052Z",
     "iopub.status.busy": "2023-01-10T16:25:31.153872Z",
     "iopub.status.idle": "2023-01-10T16:25:32.669425Z",
     "shell.execute_reply": "2023-01-10T16:25:32.668805Z",
     "shell.execute_reply.started": "2023-01-10T16:25:31.154032Z"
    }
   },
   "outputs": [],
   "source": [
    "import matplotlib.pyplot as plt\n",
    "import cv2\n",
    "import numpy as np\n",
    "\n",
    "# Function to apply simple gamma correction, rescale,\n",
    "# and clip values into range 0-255:\n",
    "def gamma_correct(x, exposure, gamma):\n",
    "  scale = 2.0 ** exposure\n",
    "  y = np.power(x * scale, 1.0 / gamma) * 255.0\n",
    "  return np.clip(y, 0.0, 255.0)\n",
    "\n",
    "# Function to plot an opencv image:\n",
    "def display_image(img):\n",
    "  plt.figure(figsize=(6, 6))\n",
    "  plt.style.use('dark_background')\n",
    "  plt.imshow(cv2.cvtColor(ldr, cv2.COLOR_BGR2RGB), interpolation='bicubic')\n",
    "  plt.show()\n",
    "\n",
    "EXR_FLAGS = cv2.IMREAD_UNCHANGED | cv2.IMREAD_ANYCOLOR | cv2.IMREAD_ANYDEPTH\n",
    "hdr = cv2.imread('out_rgb_ipu.exr', EXR_FLAGS)\n",
    "print(f\"HDR image shape: {hdr.shape} type: {hdr.dtype} min: {np.min(hdr)} max: {np.max(hdr)}\")\n",
    "\n",
    "ldr = gamma_correct(hdr, exposure=1.2, gamma=2.4).astype(np.uint8)\n",
    "cv2.imwrite('tonemapped.png', ldr)\n",
    "display_image(ldr)"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "id": "2cde1d64",
   "metadata": {},
   "source": [
    "If you want to render a CPU reference image remove the option `--ipu-only` but be aware it will take\n",
    "much much longer to render. (For a list of all command options run `./test --help`.)\n",
    "\n",
    "### Comparing AOVs with Embree\n",
    "\n",
    "We can render the same scene using Embree's CPU ray tracing kernels and compare arbitrary output variables (AOVs). You can also compare AOVs between the IPU and CPU implementations (which use almost identical code). These compariasons demonstrate that basic ray-trace functionality is working and is also useful for debugging when things are broken. For this we use a quicker render-mode `shadow-trace`.\n",
    "\n",
    "#### Check Normals\n",
    "\n",
    "For example to compare normals with Embree:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "0d011ec7-708b-42a1-914d-0d17157e5672",
   "metadata": {
    "execution": {
     "iopub.execute_input": "2023-01-10T16:25:32.670659Z",
     "iopub.status.busy": "2023-01-10T16:25:32.670411Z",
     "iopub.status.idle": "2023-01-10T16:25:51.590461Z",
     "shell.execute_reply": "2023-01-10T16:25:51.589697Z",
     "shell.execute_reply.started": "2023-01-10T16:25:32.670631Z"
    }
   },
   "outputs": [],
   "source": [
    "!./trace -w 1440 -h 1440 --render-mode shadow-trace --visualise normal --ipus 4"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "d87524c7-f6be-4712-99d8-712735eb96c1",
   "metadata": {},
   "source": [
    "Once the outputs are ready we can load them into Python to compare:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "1529e6c9",
   "metadata": {
    "execution": {
     "iopub.execute_input": "2023-01-10T16:25:51.591793Z",
     "iopub.status.busy": "2023-01-10T16:25:51.591600Z",
     "iopub.status.idle": "2023-01-10T16:25:52.515071Z",
     "shell.execute_reply": "2023-01-10T16:25:52.514391Z",
     "shell.execute_reply.started": "2023-01-10T16:25:51.591773Z"
    }
   },
   "outputs": [],
   "source": [
    "# Load normal images:\n",
    "ipu_normals = cv2.imread('out_normal_ipu.exr', EXR_FLAGS)\n",
    "cpu_normals = cv2.imread('out_normal_cpu.exr', EXR_FLAGS)\n",
    "embree_normals = cv2.imread('out_normal_embree.exr', EXR_FLAGS)\n",
    "\n",
    "# Function to print simple stats:\n",
    "def compare(ipu, embree):\n",
    "  abs_err = np.abs(ipu - embree)\n",
    "  print(f\"IPU min: {np.min(ipu)} max: {np.max(ipu)}\")\n",
    "  print(f\"Embree min: {np.min(embree)} max: {np.max(embree)}\")\n",
    "  print(f\"ABS Error min: {np.min(abs_err)} max: {np.max(abs_err)} mean: {np.mean(abs_err)}\")\n",
    "  return abs_err\n",
    "\n",
    "# Function to plot images side by side:\n",
    "def plot_compare(ipu, embree, labels=['IPU', 'Embree'], size=(12, 6)):\n",
    "  fig, ax = plt.subplots(nrows=1, ncols=2, sharey=True, figsize=size)\n",
    "  ax[0].imshow(ipu)\n",
    "  ax[0].set_title(labels[0])\n",
    "  ax[1].imshow(embree)\n",
    "  ax[1].set_title(labels[1])\n",
    "  plt.show()\n",
    "\n",
    "vis = ((ipu_normals + 1.0) / 2.0)\n",
    "vis_embree = ((embree_normals + 1.0) / 2.0)\n",
    "plot_compare(vis, vis_embree)\n",
    "abs_err = compare(ipu_normals, embree_normals)"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "id": "59b8e962",
   "metadata": {},
   "source": [
    "We can plot an error histogram (using a log scale because the error counts are small). As you can see most errors are tiny but there are a few outliers - these will be rays that hit alternative (i.e. possibly valid within machine precision) objects due to differences between our intersection test code and Embree's. This scene's dimension range is deliberately large to amplify these differences:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "cc034557",
   "metadata": {},
   "outputs": [],
   "source": [
    "plt.hist(abs_err.flatten(), bins=300, range=[0.0, np.max(abs_err)], log=True)\n",
    "plt.show()"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "id": "b24b5e11",
   "metadata": {},
   "source": [
    "#### Check Material Albedo and Shadow Rays\n",
    "\n",
    "The `--visualise` option allows us to vary the AOV that gets recorded for comparison. For example `rgb` will save albedo and shadow information in `shadow-trace` render mode:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "385ee330",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Render and load albedo images:\n",
    "!./trace -w 1440 -h 1440 --render-mode shadow-trace --visualise rgb --ipus 4\n",
    "ipu_albedo = cv2.cvtColor(cv2.imread('out_rgb_ipu.exr', EXR_FLAGS), cv2.COLOR_BGR2RGB)\n",
    "cpu_albedo = cv2.cvtColor(cv2.imread('out_rgb_cpu.exr', EXR_FLAGS), cv2.COLOR_BGR2RGB)\n",
    "embree_albedo = cv2.cvtColor(cv2.imread('out_rgb_embree.exr', EXR_FLAGS), cv2.COLOR_BGR2RGB)\n",
    "\n",
    "# Plot/compare the albedo:\n",
    "plot_compare(ipu_albedo, embree_albedo)\n",
    "abs_err = compare(ipu_albedo, embree_albedo)\n",
    "plt.hist(abs_err.flatten(), bins=300, range=[0.0, np.max(abs_err)], log=True)\n",
    "plt.show()"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "id": "95d4b1dd",
   "metadata": {},
   "source": [
    "#### Check the Primary Hit-Points\n",
    "\n",
    "Similarly we can compare the world co-ordinates of the first hits for each pixel:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "2d080d78",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Render and load hit images:\n",
    "!./trace -w 1440 -h 1440 --render-mode shadow-trace --visualise hitpoint --ipus 4\n",
    "ipu_hit = cv2.imread('out_hitpoint_ipu.exr', EXR_FLAGS)\n",
    "cpu_hit = cv2.imread('out_hitpoint_cpu.exr', EXR_FLAGS)\n",
    "embree_hit = cv2.imread('out_hitpoint_embree.exr', EXR_FLAGS)\n",
    "\n",
    "# Plot/compare the primary hit points:\n",
    "ipu_vis = ipu_hit\n",
    "ipu_vis = (ipu_vis - np.min(ipu_vis)) / (np.max(ipu_vis) - np.min(ipu_vis))\n",
    "ipu_vis = gamma_correct(ipu_vis, 0, 0.25) / 255\n",
    "embree_vis = (embree_hit - np.min(embree_hit)) / (np.max(embree_hit) - np.min(embree_hit))\n",
    "embree_vis = gamma_correct(embree_vis, 0, 0.25) / 255\n",
    "plot_compare(ipu_vis, embree_vis)\n",
    "\n",
    "abs_err = compare(ipu_hit, embree_hit)\n",
    "plt.hist(abs_err.flatten(), bins=300, range=[0.0, np.max(abs_err)], log=True)\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "d14478fc",
   "metadata": {},
   "source": [
    "### Compare Path-traced Image with CPU Reference\n",
    "\n",
    "Because the CPU is much slower we have to reduce the image size and sample count for this test. With a small number of samples the IPU and CPU images will not be identical (because their random number generators are not identical) but the distribution of their colour histograms should be the same:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "6acf0ee5",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Render, load, and compare images:\n",
    "!time ./trace -w 256 -h 256 --render-mode path-trace --visualise rgb --samples 128 --ipus 4\n",
    "ipu_pt = cv2.cvtColor(cv2.imread('out_rgb_ipu.exr', EXR_FLAGS), cv2.COLOR_BGR2RGB)\n",
    "cpu_pt = cv2.cvtColor(cv2.imread('out_rgb_cpu.exr', EXR_FLAGS), cv2.COLOR_BGR2RGB)\n",
    "ipu_pt = gamma_correct(ipu_pt, 0, 2.2).astype(np.uint8)\n",
    "cpu_pt = gamma_correct(cpu_pt, 0, 2.2).astype(np.uint8)\n",
    "plot_compare(ipu_pt, cpu_pt, labels=['IPU', 'CPU'], size=(8, 4))\n",
    "\n",
    "# Plot colour histograms:\n",
    "fig, ax = plt.subplots(nrows=1, ncols=2, sharey=True, figsize=(8, 4))\n",
    "ax[0].hist(ipu_pt[:,:,0].ravel(), bins=255, range=[1, 255], color='red', alpha=0.5)\n",
    "ax[0].hist(ipu_pt[:,:,1].ravel(), bins=255, range=[1, 255], color='Green', alpha=0.5)\n",
    "ax[0].hist(ipu_pt[:,:,2].ravel(), bins=255, range=[1, 255], color='Blue', alpha=0.5)\n",
    "ax[1].hist(cpu_pt[:,:,0].ravel(), bins=255, range=[1, 255], color='red', alpha=0.5)\n",
    "ax[1].hist(cpu_pt[:,:,1].ravel(), bins=255, range=[1, 255], color='Green', alpha=0.5)\n",
    "ax[1].hist(cpu_pt[:,:,2].ravel(), bins=255, range=[1, 255], color='Blue', alpha=0.5)\n",
    "plt.show()"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "id": "0c087d23",
   "metadata": {},
   "source": [
    "### Render the Test Scene (Collada file)\n",
    "\n",
    "Check loading of external scenes. This is a larger scene that should only just fit in tile memory:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "720daf7b",
   "metadata": {},
   "outputs": [],
   "source": [
    "!./trace -w 720 -h 720 --render-mode path-trace --visualise rgb --samples 2500 --ipus 4 --ipu-only --mesh-file ../assets/test_scene.dae --load-normals"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "c2a96f42",
   "metadata": {},
   "outputs": [],
   "source": [
    "hdr = cv2.imread('out_rgb_ipu.exr', EXR_FLAGS)\n",
    "ldr = gamma_correct(hdr, exposure=-0.5, gamma=2.6).astype(np.uint8)\n",
    "display_image(ldr)"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "id": "87dc9e34",
   "metadata": {},
   "source": [
    "### Test Neural Rendering\n",
    "\n",
    "The most experimental part of the path-tracer is the ability to use neural lighting: an HDRI environment map that has been approximated/compressed by a neural network. This is similar to a neural radiance cache (restricted to the external environment lighting):"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "93007679",
   "metadata": {},
   "outputs": [],
   "source": [
    "!./trace -w 720 -h 720 --render-mode path-trace --visualise rgb --samples 12000 --ipus 4 --ipu-only --rays-per-worker 8 --nif-hdri ../assets/nif/urban_alley_01_4k_fp16_yuv/assets.extra/  --scene spheres --hdri-rotation 110 -o hdri"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "c356b52c",
   "metadata": {},
   "outputs": [],
   "source": [
    "hdr = cv2.imread('hdri_rgb_ipu.exr', EXR_FLAGS)\n",
    "ldr = gamma_correct(hdr, exposure=-0.2, gamma=2.2).astype(np.uint8)\n",
    "display_image(ldr)"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "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.8.10"
  },
  "vscode": {
   "interpreter": {
    "hash": "c9ad4864c9c721810277722a61cd300684ef18a24cb01fc57b2c480710a554dd"
   }
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}