#!/usr/bin/env python
# coding: utf-8

# > 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.
# 

# # 5.5. Ray tracing: Cython array buffers

# 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.

# In[ ]:


import numpy as np
import matplotlib.pyplot as plt


# In[ ]:


get_ipython().run_line_magic('matplotlib', 'inline')


# In[ ]:


#%load_ext cythonmagic
get_ipython().run_line_magic('load_ext', 'Cython')


# ## Take 1

# In[ ]:


get_ipython().run_cell_magic('cython', '', 'import numpy as np\ncimport numpy as np\nfrom numpy import dot\nfrom libc.math cimport sqrt\n\nDBL = np.double\nctypedef np.double_t DBL_C\nINT = np.int\nctypedef np.int_t INT_C\ncdef INT_C w, h\n\nw, h = 200, 200  # Size of the screen in pixels.\n\ndef normalize(np.ndarray[DBL_C, ndim=1] x):\n    # This function normalizes a vector.\n    x /= np.linalg.norm(x)\n    return x\n\ndef 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\ndef 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\ndef 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\ncdef DBL_C radius, ambient, diffuse, specular_k, specular_c\n\n# Sphere and light properties.\nradius = 1.\ndiffuse = 1.\nspecular_c = 1.\nspecular_k = 50.\nambient = .05       \n')


# In[ ]:


img = run()
plt.imshow(img);
plt.xticks([]); plt.yticks([]);


# In[ ]:


get_ipython().run_line_magic('timeit', '-n1 -r1 run()')


# ## Take 2

# In this version, we rewrite normalize in pure C.

# In[ ]:


get_ipython().run_cell_magic('cython', '', 'import numpy as np\ncimport numpy as np\nfrom numpy import dot\nfrom libc.math cimport sqrt\n\nDBL = np.double\nctypedef np.double_t DBL_C\nINT = np.int\nctypedef np.int_t INT_C\ncdef INT_C w, h\n\nw, 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\ncdef 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\ndef 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\ndef 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\ndef 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\ncdef DBL_C radius, ambient, diffuse, specular_k, specular_c\n\n# Sphere and light properties.\nradius = 1.\ndiffuse = 1.\nspecular_c = 1.\nspecular_k = 50.\nambient = .05       \n')


# In[ ]:


img = run()
plt.imshow(img);
plt.xticks([]); plt.yticks([]);


# In[ ]:


get_ipython().run_line_magic('timeit', '-n1 -r1 run()')


# > You'll find all the explanations, figures, references, and much more in the book (to be released later this summer).
# 
# > [IPython Cookbook](http://ipython-books.github.io/), by [Cyrille Rossant](http://cyrille.rossant.net), Packt Publishing, 2014 (500 pages).