Adelta Tutorial — Part 2: Raymarching Primitive

[Adelta: Automatic Differentiation for Discontinuous Programs — Part 1: Math Basics]

[Adelta: Automatic Differentiation for Discontinuous Programs — Part 2: Introduction to the DSL]

[Adelta Tutorial — Part 1: Differentiating a Simple Shader Program]

[Adelta Tutorial — Part 3: Animating the SIGGRAPH logo]

[Adelta Tutorial — Part 4: Animating the Celtic Knot]

In the previous tutorial, we introduced how to author shader programs using the basic DSL syntax from Aδ. While this already gives us the flexibility to represent complex programs, we sometimes face practical issues of lower efficiency that occur when the gradient tape from our Automatic Differentiation (AD) becomes long, especially in programs that involve loops.

For example, in shader programs, a common way to define geometry is to encode it as an implicit surface, or the zero set of some mathematical function called the signed distance function², and iteratively estimate ray-geometry intersections through methods like raymarching³ or sphere tracing⁴.

While these loops themselves are programs expressible in Aδ’s DSL and can therefore be differentiated, this usually results in a long gradient tape because the number of loop iterations can be arbitrarily large. As an alternative, we can bypass the root-finding process and directly approximate the gradient using the implicit function theorem. In Aδ, we derive specialized gradient rules for the ray marching loops and they are wrapped in the RaymarchingWrapper primitive. This post demonstrates how to use that primitive through examples. For mathematical details regarding this specialized gradient rule, please refer to Appendix D of our research paper.

Example: Sphere

The signed distance function for a sphere can be easily expressed as follows. The function takes the input of a point’s 3D location (x, y, z), and returns the distance to a sphere centered at (ox, oy, oz).

Signed distance function for a sphere

In the DSL, we need to define a similar function as well, which we call a SD function. The input to the SD function is the x, y, z position of a point, and besides outputting the signed distance value, it also outputs a label corresponding to which surface is the point closest to. When the signed distance function is composed of CSG operations between multiple primitives, assigning different label values to each primitive allows us to shade each part differently. In the sphere example, because it only consists of one primitive (sphere), the label should be a constant. The SD function for a sphere is shown below:

def sdSphere(x, y, z):
    pos = np.array([x, y, z])
    sphere_diff = pos - np.array([sphere_px, sphere_py, sphere_pz])        
    dist = length(sphere_diff, 2) - sphere_r
    label = 1
    return dist, label

Note there are several constraints when authoring the SD function:

• The output distance should be C0 continuous. Specifically, only minimum() and maximum() is allowed for C1 discontinuities.
• The output label should be piecewise constant. Additionally, its discontinuities should always correspond to the C1 discontinuity in the distance computation.

In the main shader program, we can instantiate a raymarching loop with the RaymarchingWrapper primitive as below. The output of the loop can be accessed through attributes of the RaymarchingWrapper primitive.

# camera_origin
ro = np.array([origin_x, origin_y, origin_z])
# camera direction
rd = np.array([dir_x, dir_y, dir_z])
# initial t
init_t = 0
# number of iterations in the raymarching loop
raymarching_loop = 32
raymarching = RaymarchingWrapper(sdSphere, ro, rd, init_t, raymarching_loop, include_derivs=False)
cond_converge = raymarching.is_converge
final_t = raymarching.t
final_label = raymarching_ans.label

The final shader program can be found here, where we render constant color whenever the ray hits the sphere. Now let’s try to render the program and visualize its gradient map with the following command:

python approx_gradient.py --dir <path_to_store_result> --shader raymarching_sphere --init_values_pool apps/example_init_values/test_finite_diff_raymarching_sphere_init_values_pool.npy --modes visualize_gradient --render_size 640,640

Below we show the rendering of the sphere shader, as well as the gradient map wrt the x, y, z positions of the sphere.

Rendering and gradient map of a simple sphere shader

Clearly, the constant shading is not interesting and does not demonstrate the 3D scene at all. To allow more complex shading, such as adding diffuse and specular lighting, we also need the geometry’s surface normal. In the next example, we will author a box shader with the proper surface normal that allows additional lighting.

Example: Box with Surface Normal

We use the signed distance function described by iq² to represent the box geometry. Unlike the sphere example, we also want to calculate the box’s surface normal for lighting effects.

Mathematically, surface normal corresponds to the gradient direction of the signed distance function. Ideally, this can be obtained by directly applying traditional AD to differentiate the signed distance function wrt x, y, z position. However, this utility is not yet implemented in the Aδ framework, so we therefore, rely on the user to provide the computation to the surface normal as additional output to the SD function. Similar to the label value, the surface normal has the following constraints:

• The output surface normal should be piecewise continuous. Additionally, its discontinuities should always correspond to the C1 discontinuities in the distance computation.

The SD function for the box shader is shown below:

def sdBox(x, y, z):
    f0 = abs(x) - dimx
    f1 = abs(y) - dimy
    f2 = abs(z) - dimz
    q0 = maximum(f0, 0)
    q1 = maximum(f1, 0)
    q2 = maximum(f2, 0)
    max_f1f2 = maximum(f1, f2)
    dist = (q0 ** 2 + q1 ** 2 + q2 ** 2) ** 0.5 + minimum(maximum(f0, max_f1f2), 0)
    choose_f0 = f0 > max_f1f2
    choose_f1 = f1 > f2
    deriv = [select(choose_f0, sign(x), 0.),
             select(choose_f0, 0., select(choose_f1, sign(y), 0.)),
             select(choose_f0, 0., select(choose_f1, 0., sign(z)))]
        
    label = 1
    return dist, 1, deriv[0], deriv[1], deriv[2]

Because the SD function also outputs surface normal, the RaymarchingWrapper primitive is instantiated with include_derivs set to True:

raymarching = RaymarchingWrapper(sdBox, ro, rd, 0, raymarching_loop, include_derivs=True)

The final shader program can be found here with diffuse lighting. We can visualize the gradient map using the following command:

python approx_gradient.py --dir <path_to_store_result> --shader raymarching_box --init_values_pool apps/example_init_values/test_finite_diff_raymarching_half_cube_init_values_pool.npy --modes visualize_gradient

Below we show the rendering of the box shader, as well as the gradient map wrt the three camera angles:

Rendering and gradient map of a simple box shader

Summary

In this post, we introduced the RaymarchingWrapper primitive, a workaround to avoid long gradient tape when differentiating the raymarching loops that have a large number of iterations. We use two simple examples, a sphere, and a box shader to demonstrate how to use the RaymarchingWrapper primitive to construct the 3D scene. In the next post, we will discuss using the compiler to output a GLSL program that allows manual animation and postprocessing.

Reference

[1] Aδ: Autodiff for Discontinuous Program — Applied to Shaders [project page] [Github]

[2] Signed distance functions [webpage]

[3] Hypertexture [paper]

[4] Sphere Tracing: A Geometric Method for the Antialiased Ray Tracing of Implicit Surfaces [paper]