#!/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 with structs

# In[ ]:


import numpy as np
import matplotlib.pyplot as plt


# In[ ]:


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


# In[ ]:


import cython


# In[ ]:


get_ipython().run_line_magic('load_ext', 'cythonmagic')


# We now use a pure C structure to represent a 3D vector. We also implement all operations we need by hand in pure C.

# In[ ]:


get_ipython().run_cell_magic('cython', '', 'cimport cython\nimport numpy as np\ncimport numpy as np\nDBL = np.double\nctypedef np.double_t DBL_C\nfrom libc.math cimport sqrt\n\ncdef int w, h\n\ncdef struct Vec3:\n    double x, y, z\n        \ncdef Vec3 vec3(double x, double y, double z):\n    cdef Vec3 v\n    v.x = x\n    v.y = y\n    v.z = z\n    return v\n\ncdef double dot(Vec3 x, Vec3 y):\n    return x.x * y.x + x.y * y.y + x.z * y.z\n\ncdef Vec3 normalize(Vec3 x):\n    cdef double n\n    n = sqrt(x.x * x.x + x.y * x.y + x.z * x.z)\n    return vec3(x.x / n, x.y / n, x.z / n)\n\ncdef double max(double x, double y):\n    return x if x > y else y\n\ncdef double min(double x, double y):\n    return x if x < y else y\n\ncdef double clip_(double x, double m, double M):\n    return min(max(x, m), M)\n\ncdef Vec3 clip(Vec3 x, double m, double M):\n    return vec3(clip_(x.x, m, M), clip_(x.y, m, M), clip_(x.z, m, M),)\n\ncdef Vec3 add(Vec3 x, Vec3 y):\n    return vec3(x.x + y.x, x.y + y.y, x.z + y.z)\n\ncdef Vec3 subtract(Vec3 x, Vec3 y):\n    return vec3(x.x - y.x, x.y - y.y, x.z - y.z)\n\ncdef Vec3 minus(Vec3 x):\n    return vec3(-x.x, -x.y, -x.z)\n\ncdef Vec3 multiply(Vec3 x, Vec3 y):\n    return vec3(x.x * y.x, x.y * y.y, x.z * y.z)\n    \ncdef Vec3 multiply_s(Vec3 x, double c):\n    return vec3(x.x * c, x.y * c, x.z * c)\n    \ncdef double intersect_sphere(Vec3 O, \n                      Vec3 D, \n                      Vec3 S, \n                      double R):\n    # Return the distance from O to the intersection of the ray (O, D) with the \n    # sphere (S, R), or +inf if there is no intersection.\n    # O and S are 3D points, D (direction) is a normalized vector, R is a scalar.\n    cdef double a, b, c, disc, distSqrt, q, t0, t1\n    cdef Vec3 OS\n    \n    a = dot(D, D)\n    OS = subtract(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 = sqrt(disc)\n        q = (-b - distSqrt) / 2.0 if b < 0 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 1000000\n\ncdef Vec3 trace_ray(Vec3 O, Vec3 D,):\n    \n    cdef double t, radius, diffuse, specular_k, specular_c, DF, SP\n    cdef Vec3 M, N, L, toL, toO, col_ray, \\\n        position, color, color_light, ambient\n\n    # Sphere properties.\n    position = vec3(0., 0., 1.)\n    radius = 1.\n    color = vec3(0., 0., 1.)\n    diffuse = 1.\n    specular_c = 1.\n    specular_k = 50.\n    \n    # Light position and color.\n    L = vec3(5., 5., -10.)\n    color_light = vec3(1., 1., 1.)\n    ambient = vec3(.05, .05, .05)\n    \n    # Find first point of intersection with the scene.\n    t = intersect_sphere(O, D, position, radius)\n    # Return None if the ray does not intersect any object.\n    if t == 1000000:\n        col_ray.x = 1000000\n        return col_ray\n    # Find the point of intersection on the object.\n    M = vec3(O.x + D.x * t, O.y + D.y * t, O.z + D.z * t)\n    N = normalize(subtract(M, position))\n    toL = normalize(subtract(L, M))\n    toO = normalize(subtract(O, M))\n    DF = diffuse * max(dot(N, toL), 0)\n    SP = specular_c * max(dot(N, normalize(add(toL, toO))), 0) ** specular_k\n    \n    return add(ambient, add(multiply_s(color, DF), multiply_s(color_light, SP)))\n\ndef run(int w, int h):\n    cdef DBL_C[:,:,:] img = np.zeros((h, w, 3))\n    cdef Vec3 img_\n    cdef int i, j\n    cdef double x, y\n    cdef Vec3 O, Q, D, col_ray\n    cdef double w_ = float(w)\n    cdef double h_ = float(h)\n    \n    col_ray = vec3(0., 0., 0.)\n    \n    # Camera.\n    O = vec3(0., 0., -1.)  # Position.\n        \n    # Loop through all pixels.\n    for i in range(w):\n        Q = vec3(0., 0., 0.)\n        for j in range(h):\n            x = -1. + 2*(i)/w_\n            y = -1. + 2*(j)/h_\n            Q.x = x\n            Q.y = y\n            col_ray = trace_ray(O, normalize(subtract(Q, O)))\n            if col_ray.x == 1000000:\n                continue\n            img_ = clip(col_ray, 0., 1.)\n            img[h - j - 1, i, 0] = img_.x\n            img[h - j - 1, i, 1] = img_.y\n            img[h - j - 1, i, 2] = img_.z\n    return img\n')


# In[ ]:


w, h = 200, 200


# In[ ]:


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


# In[ ]:


get_ipython().run_line_magic('timeit', 'run(w, h)')


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