Ray Tracing with Numpy and Autograd

Physics Numerical physics

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:

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), 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
%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()
Out[4]:

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

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')
Out[9]:
Text(0.5,0,'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')
Out[11]:

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.))
Optimization terminated successfully.
         Current function value: 2.601764
         Iterations: 4
         Function evaluations: 15
         Gradient evaluations: 5
Out[13]:
array([ 0.58854218])

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)]
Out[14]:
0.5879396984924623

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 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.]))
Out[19]:
array([ 0.70710678])
In [20]:
fmin_bfgs(tof, x0=np.array([0.]), fprime=tof_gradient)
Optimization terminated successfully.
         Current function value: 2.601764
         Iterations: 4
         Function evaluations: 5
         Gradient evaluations: 5
Out[20]:
array([ 0.58854217])

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

How much compute time?

In [21]:
%timeit fmin_bfgs(tof, x0=np.array([0.]), disp=0)
870 µs ± 6.84 µs per loop (mean ± std. dev. of 7 runs, 1000 loops each)
In [22]:
%timeit fmin_bfgs(tof, x0=np.array([0.]), fprime=tof_gradient, disp=0)
3.87 ms ± 225 µs per loop (mean ± std. dev. of 7 runs, 100 loops each)

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
Out[23]:
(2, 1, 5)
In [24]:
x_parallel = np.ones((5, ), dtype=np.float)
x_parallel.shape
Out[24]:
(5,)
In [25]:
dist_parallel = lambda vector: np.sqrt(np.sum(np.square(vector), axis=0)).ravel()
In [26]:
dist_parallel(x_parallel)
Out[26]:
array([ 2.23606798])
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)
Out[28]:
array([ 3.29936024,  3.08046356,  2.90455287,  2.78823851,  2.74721356])

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)
Out[31]:
array([ 3.29936024,  3.08046356,  2.90455287,  2.78823851,  2.74721356])

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)
Out[34]:
array([ 1.64968012,  1.50690678,  1.3032425 ,  1.03040677,  0.70710678])

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')