Frolian's blog standing on the shoulders of PyGiants

• About
• Archive

Ray Tracing with Numpy and Autograd

Sun, 28 Jan 2018

Physics Numerical physics Derivatives

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:

Have been enjoying autograd to get derivatives automatically (to plug into scipy.optimize). Think in 2018, there's no reason to compute gradients by hand anymore: https://t.co/52P1pvQjpy

— Erik Bernhardsson (@fulhack) January 26, 2018

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.

import numpy as np

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

Let's plot this.

import matplotlib.pyplot as plt
%matplotlib notebook

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

<matplotlib.legend.Legend at 0x110090b38>

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.

n1, n2 = 1., 1.333

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

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.

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

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

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.

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

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

<matplotlib.colorbar.Colorbar at 0x11033b5c0>

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.

from scipy.optimize import fmin_bfgs

fmin_bfgs(tof, x0=(0.))

Optimization terminated successfully.
         Current function value: 2.601764
         Iterations: 4
         Function evaluations: 15
         Gradient evaluations: 5

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:

x_candidates[np.argmin(tofs)]

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.

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

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

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

tof_gradient = grad(tof_autograd)

Let's test this:

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

array([ 0.70710678])

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

array([ 0.58854217])

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

How much compute time?

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

%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:

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

(2, 1, 5)

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

(5,)

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

dist_parallel(x_parallel)

array([ 2.23606798])

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

tof_parallel(x_parallel)

array([ 3.29936024,  3.08046356,  2.90455287,  2.78823851,  2.74721356])

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

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

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:

tof_parallel_autograd(x_parallel)

array([ 3.29936024,  3.08046356,  2.90455287,  2.78823851,  2.74721356])

And compute the gradient:

from autograd import elementwise_grad

tof_parallel_gradient = elementwise_grad(tof_parallel_autograd)

tof_parallel_gradient(x_parallel)

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

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

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:

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

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

<matplotlib.colorbar.Colorbar at 0x1187f2be0>

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

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

(2, 1, 24)

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

(24,)

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

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

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

<matplotlib.colorbar.Colorbar at 0x1188d5c88>

We can even vary the locations of these points.

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

camera_parallel.shape

(2, 1, 500)

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

(500,)

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

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

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

<matplotlib.colorbar.Colorbar at 0x118e028d0>

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

%timeit gradient_descent(tof_parallel_autograd, x_parallel, tof_parallel_gradient)

37.4 ms ± 6.83 ms per loop (mean ± std. dev. of 7 runs, 10 loops each)

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.

---

This article has been tagged with the following terms:

Physics Numerical physics Derivatives

---

• 
• 
• 
• 
• 

---

Comments