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

# In this blog post, I'm going to try to show how the numerical tools of 2018 (specifically, Numpy and autograd) allow to efficiently do ray-tracing. The idea for this blog post came from this tweet that nicely correlates with projects currently developing at work:

# <blockquote class="twitter-tweet" data-lang="en"><p lang="en" dir="ltr">Have been enjoying autograd to get derivatives automatically (to plug into scipy.optimize). Think in 2018, there&#39;s no reason to compute gradients by hand anymore: <a href="https://t.co/52P1pvQjpy">https://t.co/52P1pvQjpy</a></p>&mdash; Erik Bernhardsson (@fulhack) <a href="https://twitter.com/fulhack/status/956895871236497408?ref_src=twsrc%5Etfw">January 26, 2018</a></blockquote> <script async src="https://platform.twitter.com/widgets.js" charset="utf-8"></script> 
# 

# # Introduction to ray tracing 

# Ray tracing is a tool used in many physics discipline since high frequency waves can accurately be approximated by rays. Typical applications include 3D rendering (think [povray](http://www.povray.org)), lens design or acoustic wave simulation (which is what I do professionally).
# To trace rays, you usually start with a source and follow reflexions and refractions until some the end of the tracing (exiting the scene, arriving at the camera).
# 
# In this context, the goal for this post is to find the ray that connects the source to the camera through a set of reflexions. Let's define the reference configuration for this post. We'll model points with Numpy arrays of $(x, y)$ coordinates. 
# 
# We define a source and a camera, as well as an interface that separates the two bounding media that have different speeds.

# In[1]:


import numpy as np


# In[2]:


source = np.array([0., 1.]).reshape(-1, 1)
camera = np.array([1., -1.]).reshape(-1, 1)
interface_y = 0.


# Let's plot this.

# In[3]:


import matplotlib.pyplot as plt
get_ipython().run_line_magic('matplotlib', 'notebook')


# In[4]:


fig, ax = plt.subplots()
ax.scatter(*source, label='source')
ax.scatter(*camera, label='camera')
ax.hlines(interface_y, 0, 1, label='interface')
ax.legend()


# So, the question we're asking is: what should the ray from source to camera and refracting across the interface be?
# 
# There are several ways to answer this question. We'll choose one: [Fermat's principle of least time](https://en.wikipedia.org/wiki/Fermat%27s_principle) states that the ray chooses the trajectory taking the least time to travel. Computing a time of flight is easy since rays travel in straight lines and thus $t = d / v$ where $v$ denotes the speed in the medium. Since we're using the example of light, we will use the approximate formula $t = d * n$ where $n$ is the [refractive index of the medium](https://en.wikipedia.org/wiki/Refractive_index).
# 
# Let's see how this works in the case of our example, where we denote by $x$ the location of the refraction point on the horizontal interface.
# 
# We model the top medium to be air and the bottom medium to be water.

# In[5]:


n1, n2 = 1., 1.333


# In[6]:


# euclidian distance written as Numpy function
dist = lambda vector: np.sqrt(np.sum(np.square(vector.ravel()))) 


# In[7]:


def tof(x):
    """Returns time of flight for ray passing through point (x, 0) on the interface."""
    interface_point = np.array([x, 0.]).reshape(-1, 1)
    return dist(source - interface_point) * n1 + dist(interface_point - camera) * n2


# Let's do a sort of line search to determine the minimal time of flight for points of abscissa between 0 and 1.

# In[8]:


x_candidates = np.linspace(0, 1, num=200)
tofs = tof(x_candidates)


# In[9]:


fig, ax = plt.subplots()
ax.plot(x_candidates, tofs, label='time of flights')
ax.legend()
ax.set_xlabel('abscissa of interface refraction point')


# As expected, we find a minimum time of flight, which is the "true" one. Let's plot all possible rays from this search and color them by their times of flight.

# In[10]:


import matplotlib as mpl
norm = mpl.colors.Normalize(vmin=tofs.min(),vmax=tofs.max())
cmap = plt.cm.get_cmap('plasma_r')


# In[11]:


fig, ax = plt.subplots()
ax.scatter(*source, label='source')
ax.scatter(*camera, label='camera')
for x in x_candidates:
    interface_point = np.array([x, 0.]).reshape(-1, 1)
    points = np.concatenate((source, interface_point, camera), axis=1).T
    ax.plot(points[:, 0], points[:, 1], color=cmap(norm(tof(x))))

# plotting the minimum time of flight
x = x_candidates[np.argmin(tofs)]
interface_point = np.array([x, 0.]).reshape(-1, 1)
points = np.concatenate((source, interface_point, camera), axis=1).T
ax.plot(points[:, 0], points[:, 1], color='k', label='minimum tof')
ax.legend()


# plotting the colorbar, see https://stackoverflow.com/questions/43805821/matplotlib-add-colorbar-to-non-mappable-object
sm = plt.cm.ScalarMappable(cmap=cmap, norm=norm)
sm.set_array([])
plt.colorbar(sm, 
             ticks=np.linspace(tofs.min(),tofs.max(),21),
             boundaries=np.linspace(tofs.min(),tofs.max(),21),
            label='time of flight')


# # Iterative algorithms 

# Now, the interesting point here is: in the general case, you don't want to compute all possible rays and choose the one with minimum time of flight. Instead, you just change the refraction point location until it is minimal. This is a typical optimization problem. Let's see how we can implement this.

# In[12]:


from scipy.optimize import fmin_bfgs


# In[13]:


fmin_bfgs(tof, x0=(0.))


# The `fmin_bfgs` function is a classical optimization function that minimizes the argument to the function. Let's check that the result is the same with what we had above:

# In[14]:


x_candidates[np.argmin(tofs)]


# Very nice! However, we also see that the optimization call above had to perform 4 iterations and 15 evaluations of the `tof` function. Can we minize this computational cost? One way of doing this is to pass the gradient of the time of flight function to the `fmin_bfgs` function. However, we don't want to compute this by hand. 

# # Here comes autograd 

# Well, luckily for us, an automatic differentiation library comes to the rescue! [autograd](https://github.com/HIPS/autograd) is capable of automatically computing gradients for us by applying the chain rule to Numpy based functions. It turns out our `tof` function satisfies this requirement!
# 
# However, we still have to rewrite it using the autograd Numpy module.

# In[15]:


import autograd.numpy as auto_np # Thinly-wrapped numpy
from autograd import grad


# In[16]:


dist_autograd = lambda vector: auto_np.sqrt(auto_np.sum(auto_np.square(vector.ravel()))) 


# In[17]:


def tof_autograd(x):
    """Returns time of flight for ray passing through point (x, 0) on the interface.
    Autograd version."""
    interface_point = auto_np.array([x, 0.]).reshape(-1, 1)
    return dist_autograd(source - interface_point) * n1 + dist_autograd(interface_point - camera) * n2


# In[18]:


tof_gradient = grad(tof_autograd)


# Let's test this:

# In[19]:


tof_gradient(np.array([1.]))


# In[20]:


fmin_bfgs(tof, x0=np.array([0.]), fprime=tof_gradient)


# Using this method, we are doing less function evaluations during the optimization which saves compute time.
# 
# How much compute time?

# In[21]:


get_ipython().run_line_magic('timeit', 'fmin_bfgs(tof, x0=np.array([0.]), disp=0)')


# In[22]:


get_ipython().run_line_magic('timeit', 'fmin_bfgs(tof, x0=np.array([0.]), fprime=tof_gradient, disp=0)')


# Wait, what? The gradient version of the optimization algorithm is 4.5 times slower than the non-gradient version!
# 
# So, my reasoning here is wrong: this does not save compute time. However, gradient based optimization can be much faster on more complicated problem, so let's ignore the compute time aspect here.

# If you've read this far you must now be wondering: okay, so I get automatic gradients but then again it doesn't help that much if I have to cast thousand of rays, right?
# 
# Well, it turns out that the autograd method can be applied to lots of rays in parallel, resulting in one big optimization problem!

# # Parallel ray casting with autograd 

# Let's assume we now have five camera points that we want to cast rays to. Our time of flight computation should now return 5 times of flights, one for each point. Let's see if this works:

# In[23]:


camera_parallel = np.concatenate((np.linspace(0, 1, num=5).reshape(1, 1, -1), 
                                  -1 * np.ones((5,)).reshape(1, 1, -1)), axis=0)
camera_parallel.shape


# In[24]:


x_parallel = np.ones((5, ), dtype=np.float)
x_parallel.shape


# In[25]:


dist_parallel = lambda vector: np.sqrt(np.sum(np.square(vector), axis=0)).ravel()


# In[26]:


dist_parallel(x_parallel)


# In[27]:


def tof_parallel(x):
    """Returns time of flight for ray passing through point (x, 0) on the interface in parallel."""
    x_parallel = x.reshape(1, 1, -1)
    interface_point = np.concatenate((x_parallel, np.zeros_like(x_parallel)))
    return dist_parallel(source[:, :, np.newaxis] - interface_point) * n1 + \
           dist_parallel(interface_point - camera_parallel) * n2


# In[28]:


tof_parallel(x_parallel)


# Now that we have that, we can rewrite this with autograd.

# In[29]:


dist_parallel_autograd = lambda vector: auto_np.sqrt(auto_np.sum(auto_np.square(vector), axis=0)).ravel()


# In[30]:


def tof_parallel_autograd(x):
    """Returns time of flight for ray passing through point (x, 0) on the interface in parallel.
    Autograd version."""
    x_parallel = x.reshape(1, 1, -1)
    interface_point = auto_np.concatenate((x_parallel, auto_np.zeros_like(x_parallel)))
    return dist_parallel_autograd(source[:, :, np.newaxis] - interface_point) * n1 + \
           dist_parallel_autograd(interface_point - camera_parallel) * n2


# Let's check the result:

# In[31]:


tof_parallel_autograd(x_parallel)


# And compute the gradient:

# In[32]:


from autograd import elementwise_grad


# In[33]:


tof_parallel_gradient = elementwise_grad(tof_parallel_autograd)


# In[34]:


tof_parallel_gradient(x_parallel)


# Which we can now use to solve the problem using gradient descent (unfortunately, I couldn't get the standard scipy optimizers to work with a vector function...).

# In[35]:


def gradient_descent(func, x0, fprime, learning_rate=0.1, iters=50):
    """Gradient descent by hand."""
    x = x0.copy()
    for _ in range(iters):
        x -= learning_rate * fprime(x)
    return x


# In[36]:


x_parallel_sol = gradient_descent(tof_parallel_autograd, x_parallel, tof_parallel_gradient)
tofs_parallel = tof_parallel(x_parallel_sol)


# Let's plot these solutions:

# In[37]:


norm = mpl.colors.Normalize(vmin=tofs_parallel.min(),vmax=tofs_parallel.max())
cmap = plt.cm.get_cmap('plasma_r')


# In[38]:


fig, ax = plt.subplots()
ax.scatter(*source, label='source')
ax.scatter(*camera_parallel, label='camera')
for ind, x in enumerate(x_parallel_sol):
    interface_point = np.array([x, 0.]).reshape(-1, 1)
    camera = camera_parallel[:, :, ind]
    points = np.concatenate((source, interface_point, camera), axis=1).T
    ax.plot(points[:, 0], points[:, 1], color=cmap(norm(tof(x))), alpha=0.3, lw=5)
ax.hlines(interface_y, 0, 1, label='interface')
ax.legend()
sm = plt.cm.ScalarMappable(cmap=cmap, norm=norm)
sm.set_array([])
plt.colorbar(sm, 
             ticks=np.linspace(tofs_parallel.min(),tofs_parallel.max(),21),
             boundaries=np.linspace(tofs_parallel.min(),tofs_parallel.max(),21),
            label='time of flight')


# One of the cool things with this global gradient descent approach is that it generalizes well to a lot of points:

# In[39]:


N = 24
camera_parallel = np.concatenate((np.linspace(0, 1, num=N).reshape(1, 1, -1), 
                                  -1 * np.ones((N,)).reshape(1, 1, -1)), axis=0)
camera_parallel.shape


# In[40]:


x_parallel = np.ones((N, ), dtype=np.float)
x_parallel.shape


# In[41]:


x_parallel_sol = gradient_descent(tof_parallel_autograd, x_parallel, tof_parallel_gradient)
tofs_parallel = tof_parallel(x_parallel_sol)


# In[42]:


norm = mpl.colors.Normalize(vmin=tofs_parallel.min(),vmax=tofs_parallel.max())
cmap = plt.cm.get_cmap('plasma_r')


# In[43]:


fig, ax = plt.subplots()
ax.scatter(*source, label='source')
ax.scatter(*camera_parallel, label='camera')
for ind, x in enumerate(x_parallel_sol):
    interface_point = np.array([x, 0.]).reshape(-1, 1)
    camera = camera_parallel[:, :, ind]
    points = np.concatenate((source, interface_point, camera), axis=1).T
    ax.plot(points[:, 0], points[:, 1], color=cmap(norm(tof(x))), alpha=0.3, lw=5)
ax.hlines(interface_y, 0, 1, label='interface')
ax.legend()
sm = plt.cm.ScalarMappable(cmap=cmap, norm=norm)
sm.set_array([])
plt.colorbar(sm, 
             ticks=np.linspace(tofs_parallel.min(),tofs_parallel.max(),21),
             boundaries=np.linspace(tofs_parallel.min(),tofs_parallel.max(),21),
            label='time of flight')


# We can even vary the locations of these points.

# In[44]:


X, Y = np.meshgrid(np.linspace(0., 1., num=50), np.linspace(-1., -0.2, num=10))
camera_parallel = np.concatenate((X.reshape(1, -1), Y.reshape(1, -1)))[:, np.newaxis, :]
camera_parallel += np.random.rand(*camera_parallel.shape) / 10


# In[45]:


camera_parallel.shape


# In[46]:


x_parallel = np.ones((camera_parallel.shape[2], ), dtype=np.float)
x_parallel.shape


# In[47]:


x_parallel_sol = gradient_descent(tof_parallel_autograd, x_parallel, tof_parallel_gradient)
tofs_parallel = tof_parallel(x_parallel_sol)


# In[48]:


norm = mpl.colors.Normalize(vmin=tofs_parallel.min(),vmax=tofs_parallel.max())
cmap = plt.cm.get_cmap('plasma_r')


# In[49]:


fig, ax = plt.subplots()
ax.scatter(*source, label='source')
ax.scatter(*camera_parallel, marker='.', label='cameras')
for ind, x in enumerate(x_parallel_sol):
    interface_point = np.array([x, 0.]).reshape(-1, 1)
    camera = camera_parallel[:, :, ind]
    points = np.concatenate((source, interface_point, camera), axis=1).T
    ax.plot(points[:, 0], points[:, 1], color=cmap(norm(tof(x))), alpha=0.3, lw=2)
ax.hlines(interface_y, 0, 1, label='interface')
ax.legend()
sm = plt.cm.ScalarMappable(cmap=cmap, norm=norm)
sm.set_array([])
plt.colorbar(sm, 
             ticks=np.linspace(tofs_parallel.min(),tofs_parallel.max(),21),
             boundaries=np.linspace(tofs_parallel.min(),tofs_parallel.max(),21),
            label='time of flight')


# One of the cool things about this parallel optimization by gradient descent is the computation time:

# In[50]:


get_ipython().run_line_magic('timeit', 'gradient_descent(tof_parallel_autograd, x_parallel, tof_parallel_gradient)')


# Compared to our previous timings of a couple milliseconds for one optimisation, this is pretty nice since the unit time for a single ray goes down to around 100 μs.

# # Conclusions 

# In this post, I've tried to show how you can use the automatic different package autograd to solve a problem as old as computer graphics, namely ray tracing. I hope you've enjoyed the ride and will also try to use this package in your next numerical computing project. 

# *This post was entirely written using the IPython notebook. Its content is BSD-licensed. You can see a static view or download this notebook with the help of nbviewer at [20180127_RayTracingNumpyAutograd.ipynb](http://nbviewer.ipython.org/urls/raw.github.com/flothesof/posts/master/20180127_RayTracingNumpyAutograd.ipynb).*