{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "> This is one of the 100 recipes of the [IPython Cookbook](http://ipython-books.github.io/), the definitive guide to high-performance scientific computing and data science in Python.\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# 5.5. Ray tracing: Cython array buffers"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "In this example, we will render a sphere with a diffuse and specular material. The principle is to model a scene with a light source and a camera, and use the physical properties of light propagation to calculate the light intensity and color of every pixel of the screen."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "import numpy as np\n",
    "import matplotlib.pyplot as plt"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "%matplotlib inline"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "#%load_ext cythonmagic\n",
    "%load_ext Cython"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Take 1"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "%%cython\n",
    "import numpy as np\n",
    "cimport numpy as np\n",
    "from numpy import dot\n",
    "from libc.math cimport sqrt\n",
    "\n",
    "DBL = np.double\n",
    "ctypedef np.double_t DBL_C\n",
    "INT = np.int\n",
    "ctypedef np.int_t INT_C\n",
    "cdef INT_C w, h\n",
    "\n",
    "w, h = 200, 200  # Size of the screen in pixels.\n",
    "\n",
    "def normalize(np.ndarray[DBL_C, ndim=1] x):\n",
    "    # This function normalizes a vector.\n",
    "    x /= np.linalg.norm(x)\n",
    "    return x\n",
    "\n",
    "def intersect_sphere(np.ndarray[DBL_C, ndim=1] O, np.ndarray[DBL_C, ndim=1] D, \n",
    "                     np.ndarray[DBL_C, ndim=1] S, DBL_C R):\n",
    "    # Return the distance from O to the intersection \n",
    "    # of the ray (O, D) with the sphere (S, R), or \n",
    "    # +inf if there is no intersection.\n",
    "    # O and S are 3D points, D (direction) is a \n",
    "    # normalized vector, R is a scalar.\n",
    "    \n",
    "    cdef DBL_C a, b, c, disc, distSqrt, q, t0, t1\n",
    "    cdef np.ndarray[DBL_C, ndim=1] OS\n",
    "    \n",
    "    a = dot(D, D)\n",
    "    OS = O - S\n",
    "    b = 2 * dot(D, OS)\n",
    "    c = dot(OS, OS) - R*R\n",
    "    disc = b*b - 4*a*c\n",
    "    if disc > 0:\n",
    "        distSqrt = np.sqrt(disc)\n",
    "        q = (-b - distSqrt) / 2.0 if b < 0 \\\n",
    "            else (-b + distSqrt) / 2.0\n",
    "        t0 = q / a\n",
    "        t1 = c / q\n",
    "        t0, t1 = min(t0, t1), max(t0, t1)\n",
    "        if t1 >= 0:\n",
    "            return t1 if t0 < 0 else t0\n",
    "    return np.inf\n",
    "\n",
    "def trace_ray(np.ndarray[DBL_C, ndim=1] O, np.ndarray[DBL_C, ndim=1] D,\n",
    "               np.ndarray[DBL_C, ndim=1] position,\n",
    "               np.ndarray[DBL_C, ndim=1] color,\n",
    "               np.ndarray[DBL_C, ndim=1] L,\n",
    "               np.ndarray[DBL_C, ndim=1] color_light):\n",
    "        \n",
    "    cdef DBL_C t\n",
    "    cdef np.ndarray[DBL_C, ndim=1] M, N, toL, toO, col\n",
    "    \n",
    "    # Find first point of intersection with the scene.\n",
    "    t = intersect_sphere(O, D, position, radius)\n",
    "    # No intersection?\n",
    "    if t == np.inf:\n",
    "        return\n",
    "    # Find the point of intersection on the object.\n",
    "    M = O + D * t\n",
    "    N = normalize(M - position)\n",
    "    toL = normalize(L - M)\n",
    "    toO = normalize(O - M)\n",
    "    # Ambient light.\n",
    "    col = ambient * np.ones(3)\n",
    "    # Lambert shading (diffuse).\n",
    "    col += diffuse * max(dot(N, toL), 0) * color\n",
    "    # Blinn-Phong shading (specular).\n",
    "    col += specular_c * color_light * \\\n",
    "        max(dot(N, normalize(toL + toO)), 0) \\\n",
    "           ** specular_k\n",
    "    return col\n",
    "\n",
    "def run():\n",
    "    cdef np.ndarray[DBL_C, ndim=3] img\n",
    "    img = np.zeros((h, w, 3))\n",
    "    cdef INT_C i, j\n",
    "    cdef DBL_C x, y\n",
    "    cdef np.ndarray[DBL_C, ndim=1] O, Q, D, col, position, color, L, color_light\n",
    "\n",
    "    # Sphere properties.\n",
    "    position = np.array([0., 0., 1.])\n",
    "    color = np.array([0., 0., 1.])\n",
    "    L = np.array([5., 5., -10.])\n",
    "    color_light = np.ones(3)\n",
    "        \n",
    "    # Camera.\n",
    "    O = np.array([0., 0., -1.])  # Position.\n",
    "    Q = np.array([0., 0., 0.])  # Pointing to.\n",
    "        \n",
    "    # Loop through all pixels.\n",
    "    for i, x in enumerate(np.linspace(-1., 1., w)):\n",
    "        for j, y in enumerate(np.linspace(-1., 1., h)):\n",
    "            # Position of the pixel.\n",
    "            Q[0], Q[1] = x, y\n",
    "            # Direction of the ray going through the optical center.\n",
    "            D = normalize(Q - O)\n",
    "            # Launch the ray and get the color of the pixel.\n",
    "            col = trace_ray(O, D, position, color, L, color_light)\n",
    "            if col is None:\n",
    "                continue\n",
    "            img[h - j - 1, i, :] = np.clip(col, 0, 1)\n",
    "    return img\n",
    "\n",
    "cdef DBL_C radius, ambient, diffuse, specular_k, specular_c\n",
    "\n",
    "# Sphere and light properties.\n",
    "radius = 1.\n",
    "diffuse = 1.\n",
    "specular_c = 1.\n",
    "specular_k = 50.\n",
    "ambient = .05       "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "img = run()\n",
    "plt.imshow(img);\n",
    "plt.xticks([]); plt.yticks([]);"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "%timeit -n1 -r1 run()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Take 2"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "In this version, we rewrite normalize in pure C."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "%%cython\n",
    "import numpy as np\n",
    "cimport numpy as np\n",
    "from numpy import dot\n",
    "from libc.math cimport sqrt\n",
    "\n",
    "DBL = np.double\n",
    "ctypedef np.double_t DBL_C\n",
    "INT = np.int\n",
    "ctypedef np.int_t INT_C\n",
    "cdef INT_C w, h\n",
    "\n",
    "w, h = 200, 200  # Size of the screen in pixels.\n",
    "\n",
    "# normalize is now a pure C function that does not make\n",
    "# use NumPy for the computations\n",
    "cdef normalize(np.ndarray[DBL_C, ndim=1] x):\n",
    "    cdef DBL_C n\n",
    "    n = sqrt(x[0] * x[0] + x[1] * x[1] + x[2] * x[2])\n",
    "    x[0] /= n\n",
    "    x[1] /= n\n",
    "    x[2] /= n\n",
    "    return x\n",
    "\n",
    "def intersect_sphere(np.ndarray[DBL_C, ndim=1] O, np.ndarray[DBL_C, ndim=1] D, \n",
    "                     np.ndarray[DBL_C, ndim=1] S, DBL_C R):\n",
    "    # Return the distance from O to the intersection \n",
    "    # of the ray (O, D) with the sphere (S, R), or \n",
    "    # +inf if there is no intersection.\n",
    "    # O and S are 3D points, D (direction) is a \n",
    "    # normalized vector, R is a scalar.\n",
    "    \n",
    "    cdef DBL_C a, b, c, disc, distSqrt, q, t0, t1\n",
    "    cdef np.ndarray[DBL_C, ndim=1] OS\n",
    "    \n",
    "    a = dot(D, D)\n",
    "    OS = O - S\n",
    "    b = 2 * dot(D, OS)\n",
    "    c = dot(OS, OS) - R*R\n",
    "    disc = b*b - 4*a*c\n",
    "    if disc > 0:\n",
    "        distSqrt = np.sqrt(disc)\n",
    "        q = (-b - distSqrt) / 2.0 if b < 0 \\\n",
    "            else (-b + distSqrt) / 2.0\n",
    "        t0 = q / a\n",
    "        t1 = c / q\n",
    "        t0, t1 = min(t0, t1), max(t0, t1)\n",
    "        if t1 >= 0:\n",
    "            return t1 if t0 < 0 else t0\n",
    "    return np.inf\n",
    "\n",
    "def trace_ray(np.ndarray[DBL_C, ndim=1] O, np.ndarray[DBL_C, ndim=1] D,\n",
    "               np.ndarray[DBL_C, ndim=1] position,\n",
    "               np.ndarray[DBL_C, ndim=1] color,\n",
    "               np.ndarray[DBL_C, ndim=1] L,\n",
    "               np.ndarray[DBL_C, ndim=1] color_light):\n",
    "        \n",
    "    cdef DBL_C t\n",
    "    cdef np.ndarray[DBL_C, ndim=1] M, N, toL, toO, col\n",
    "    \n",
    "    # Find first point of intersection with the scene.\n",
    "    t = intersect_sphere(O, D, position, radius)\n",
    "    # No intersection?\n",
    "    if t == np.inf:\n",
    "        return\n",
    "    # Find the point of intersection on the object.\n",
    "    M = O + D * t\n",
    "    N = normalize(M - position)\n",
    "    toL = normalize(L - M)\n",
    "    toO = normalize(O - M)\n",
    "    # Ambient light.\n",
    "    col = ambient * np.ones(3)\n",
    "    # Lambert shading (diffuse).\n",
    "    col += diffuse * max(dot(N, toL), 0) * color\n",
    "    # Blinn-Phong shading (specular).\n",
    "    col += specular_c * color_light * \\\n",
    "        max(dot(N, normalize(toL + toO)), 0) \\\n",
    "           ** specular_k\n",
    "    return col\n",
    "\n",
    "def run():\n",
    "    cdef np.ndarray[DBL_C, ndim=3] img\n",
    "    img = np.zeros((h, w, 3))\n",
    "    cdef INT_C i, j\n",
    "    cdef DBL_C x, y\n",
    "    cdef np.ndarray[DBL_C, ndim=1] O, Q, D, col, position, color, L, color_light\n",
    "\n",
    "    # Sphere properties.\n",
    "    position = np.array([0., 0., 1.])\n",
    "    color = np.array([0., 0., 1.])\n",
    "    L = np.array([5., 5., -10.])\n",
    "    color_light = np.ones(3)\n",
    "        \n",
    "    # Camera.\n",
    "    O = np.array([0., 0., -1.])  # Position.\n",
    "    Q = np.array([0., 0., 0.])  # Pointing to.\n",
    "        \n",
    "    # Loop through all pixels.\n",
    "    for i, x in enumerate(np.linspace(-1., 1., w)):\n",
    "        for j, y in enumerate(np.linspace(-1., 1., h)):\n",
    "            # Position of the pixel.\n",
    "            Q[0], Q[1] = x, y\n",
    "            # Direction of the ray going through the optical center.\n",
    "            D = normalize(Q - O)\n",
    "            # Launch the ray and get the color of the pixel.\n",
    "            col = trace_ray(O, D, position, color, L, color_light)\n",
    "            if col is None:\n",
    "                continue\n",
    "            img[h - j - 1, i, :] = np.clip(col, 0, 1)\n",
    "    return img\n",
    "\n",
    "cdef DBL_C radius, ambient, diffuse, specular_k, specular_c\n",
    "\n",
    "# Sphere and light properties.\n",
    "radius = 1.\n",
    "diffuse = 1.\n",
    "specular_c = 1.\n",
    "specular_k = 50.\n",
    "ambient = .05       "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "img = run()\n",
    "plt.imshow(img);\n",
    "plt.xticks([]); plt.yticks([]);"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "%timeit -n1 -r1 run()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "> You'll find all the explanations, figures, references, and much more in the book (to be released later this summer).\n",
    "\n",
    "> [IPython Cookbook](http://ipython-books.github.io/), by [Cyrille Rossant](http://cyrille.rossant.net), Packt Publishing, 2014 (500 pages)."
   ]
  }
 ],
 "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.4.2"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 0
}