Differentiable Refractive Ray Tracing

Exploring the GRIN Lens Design Space

Refraction

Ray Tracing, Snell's Law

\frac{\sin{\theta_2}}{\sin{\theta_1}} = \frac{\eta_1}{\eta_2}

Gradient Index Ray Tracing

What about gradient interfaces?

Gradient Refractive Fields

https://www.thorlabs.com

Luneburg Lens

\eta = \sqrt{2 - \left(\frac{r}{R}^2\right)}

Why GRIN Lens?

GRIN Lenses have many more degrees of freedom over traditional lenses.

What sort of effects can we achieve in this space? What are the limitations?

Mostly a theory question

What can GRIN Lenses do that traditional refractive lenses cannot? Vice Versa?

Pretty ill-defined question, so instead...

If we have an objective, what profile achieves it?

We know some (Luneburg, Maxwell, etc.), but what about other objectives?

Outline

  1. Forward Tracing
  2. Inverse Problem
  3. Solve via Autodiff
  4. Adjoint Method
  5. Experiments

Ray Tracing Equations

\begin{aligned} \frac{dx}{ds} &= v \\ \frac{dv}{ds} &= \eta(x) \nabla \eta (x) \end{aligned}
\begin{aligned} \dot{x} &= v \\ \dot{v} &= \eta(x) \nabla \eta (x) \end{aligned}

Simulating GRIN Fields

\begin{aligned} x_i &= v_i \Delta s + x_{i-1} \\ v_i &= \eta(x_{i-1}) \nabla \eta (x_{i-1}) \Delta s + v_{i-1} \end{aligned}
\begin{aligned} \dot{x} &= v \\ \dot{v} &= \eta(x) \nabla \eta (x) \end{aligned}

The Problem Statement

Assuming our desired inputs, what refractive field generates the desired output?

The Problem Statement

\begin{bmatrix} x_o \\ v_o \\ \lambda \end{bmatrix}
\begin{bmatrix} x_f \\ v_f \\ \lambda \end{bmatrix}
\eta(x)

(We will ignore wavelength)

The inverse formulation

\underset{\eta}{\min} ~~ \mathcal{C} \left(x_f, v_f \right)
s.t. \\ \dot{x} = v \\ \dot{v} = \eta \nabla \eta

The cost objective can be anything as long as its a function of \(x\) and \(v\)

Solve it with autodiff!

Initial Conditions

\(x_o, v_o\)

Refractive Field

\(\eta\)

Eikonal Tracer

Cost Objective

Backprop

Gradient

An example problem

\mathcal{C}(x_f, v_f) = \underset{c\in \text{sensors}}{\sum} \lVert \hat{I_c}- I_c(x_f, v_f) \rVert^2

An example is an image formation model, where each ray carries constant irradiance

The goal is to have multiple images based on direction

2 Images, 2 Directions

Collimated Sources

Sensors

\(\eta\)

Autodiff Results

Ground

Truth

Iter 0

Iter 20

Iter 300

Autodiff is expensive...

  • 128x128 pixel target images
  • 10 samples per pixel
  • 128^3 resolution cube
  • ~5s per iteration
  • ~12GB GPU memory usage
    • forward trace + backprop

PyTorch Implementation:

AD Graph

\begin{bmatrix} x_o \\ v_o \\ \eta \end{bmatrix}
\{
\begin{bmatrix} x_f \\ v_f \end{bmatrix}

Every integration step is a "layer" in the graph

\(\Delta s\)

\(\Delta s\)

\(\Delta s\)

Autodiff is step size dependent

Each time step in the simulation adds to the computation graph.

Halving the step size effectively doubles the computation graph!

Higher resolution requires smaller step size

Too low of a step size means the volume itself will be undersampled.

Adjoint State Method

A method for computing the derivative of interest via a set of partial differential equations

\begin{gathered} \dot{\mu} = -\lambda \\ \dot{\lambda} = \left(\nabla \eta \nabla \eta^T + \eta \nabla \nabla \eta \right) \mu \end{gathered}
\dot{x} = v \\ \dot{v} = \eta \nabla \eta

Gradient from Adjoint State

\frac{d\mathcal{C}}{d\eta} = \left(\eta \nabla d\eta + d\eta \nabla \eta \right) \mu

Calculating the adjoint state \(\mu\) gives us the value needed for the derivative of interest

\underset{\eta}{\min} ~~ \mathcal{C} \left(x_f, v_f \right)
s.t. \\ \dot{x} = v \\ \dot{v} = \eta \nabla \eta

With this value, we can use a gradient based optimizer like Adam, SGD, BFGS, etc.

Adjoint via Backtracing

\begin{gathered} \dot{\mu} = -\lambda \\ \dot{\lambda} = \left(\nabla \eta \nabla \eta^T + \eta \nabla \nabla \eta \right) \mu \end{gathered}

\(\mu\) can be thought of as the error vector at the boundary being parallel transported back through the volume

Adjoint via Backtracing

We still need the path for the adjoint state, since its defined on the path.

Naively, this means that we have to save the whole path, which has the same problem as before.

Time-reversibility

The optical paths of the rays are "symplectic" or reversible.

\begin{aligned} x_i &= v_i \Delta s + x_{i-1} \\ v_i &= \eta(x_{i-1}) \nabla \eta (x_{i-1}) \Delta s + v_{i-1} \end{aligned}
\begin{aligned} x_{i-1} &= x_{i} - v_i \Delta s \\ v_{i-1} &= v_{i} - \eta(x_{i-1}) \nabla \eta (x_{i-1}) \Delta s \end{aligned}

Same experiment

Ground

Truth

Autodiff

Adjoint

Performance Numbers

AD Implementation:

  • ~5s per iteration
  • ~18GB GPU mem
  • 128x128 pixel target images
  • 10 samples per pixel
  • 128^3 resolution cube

Adjoint Implementation:

  • ~1s per iteration
  • ~2GB GPU mem

Larger and more accurate simulations

AD can also achieve these results, it would just take much longer and use more memory.

Adjoint Replaces AD

\begin{bmatrix} x_o \\ v_o \\ \eta \end{bmatrix}
\begin{bmatrix} x_f \\ v_f \end{bmatrix}

Adjoint

Tracer

Different Design Tasks

Near and Far Fields

Collimated Source

Near-Field Sensor

\(\eta\)

Far-Field Sensor

Near and Far Fields

Optimization

Ground Truth

Near

Far

Caustic Designs

Same near-far setup, but use a geometric cost objective instead

\mathcal{C}(x, v) = \lVert \text{SDF}_x(x_f) \rVert^2 + \lVert \text{SDF}_v(v_f) \rVert^2

desired image

signed distance field

Caustic Designs

Near Field

Far Field

Ground Truth

Volume

Caustic Designs (problem)

Near Field

Far Field

The energy distribution isn't uniform. That isn't a constraint in the optimization.

As long as the ray reaches the circle, the loss will go to zero.

Luneburg Reconstruction

The Luneburg lens can be defined by a geometric property rather than an index profile.

Luneburg Reconstruction

\mathcal{C}(x, v) = \underset{i \in \mathcal{S}^2}{\int} \lVert \mathcal{T}(x_o) - p_i \rVert^2

Luneburg Reconstruction

It works! We can match the actual Luneburg profile with just the geometric description.

Closing Thoughts

  • Can't handle discrete interfaces.
  • Don't account for real world fabrication constraints
  • Extensions include:
    • polarization
    • wavelength dependence

More Examples?

What sort of things would be cool to see?

What other experiments should I try?