Fork me on GitHub

# Dex: Multi-step ray tracer

Source file: dex-raytrace.fut

The following is a port of raytrace.dx, which is itself based on this Jax program by Eric Jang. See this blog post for details on the algorithm - this example won’t be getting into those details. We’ve tried to maintain the naming scheme of the original program, but we’ve re-ordered most definitions bring them closer to their first usage. This is mostly necessary due to the somewhat convoluted program structure forced by the use of modules for AD.

``import "dex-prelude"``

The biggest issue is that the Dex program uses the `grad` operator to convert distance functions into normal functions. Futhark does not (yet) support anything like this as a language construct. Instead, we use forward-mode AD with dual numbers, which can be implemented as a library in Futhark. It works, and is reasonably efficient in this case, but it is certainly not as concise.

``````module ad = import "dual-numbers"

module type expanded_field = {
val max : t -> t -> t
val sqrt : t -> t
val cos : t -> t
val sin : t -> t
val abs : t -> t
}

module f64_field : (expanded_field with t = f64) = {
let max = f64.max
let sqrt = f64.sqrt
let cos = f64.cos
let sin = f64.sin
let abs = f64.abs``````

We expose all the definitions from `f64_field` rather than repeating them.

``````  open ad.f64_field
}``````

The field implementation for dual numbers is mostly boilerplate. We cannot reuse the `mk_dual` parametric module from dual-numbers.fut, because that module chose to make the representation abstract.

``````module f64_dual_field = {
type underlying = f64
type t = (underlying, underlying)

let primal ((x, _) : t) = x
let tangent ((_, x') : t) = x'
let dual0 x : t = (x, 0)
let dual1 x : t = (x, 1)

let max (x,x' : f64) (y,y') =
(f64.max x y,
if x >= y then x' else y')

let sqrt (x,x') =
(f64.sqrt x, x' / (2 * f64.sqrt x))

let cos (x, x') =
(f64.cos x, f64.neg x' * f64.sin x)

let sin (x, x') =
(f64.sin x, x' * f64.cos x)

let abs (x,x') =
(f64.abs x, f64.sgn x * x')

let neg (x,x') : t = (-x, -x')
let recip (x,x') : t = (1/x, -(x'/(x*x)))

let (x,x') + (y,y') : t = f64.((x + y, x' + y'))
let (x,x') * (y,y') : t = f64.((x * y, x' * y + x * y'))
let x - y = x + neg y
let x / y = x * recip y

let (x,_) == (y,_) = f64.(x == y)
let (x,_) < (y,_) = f64.(x < y)
let (x,_) > (y,_) = f64.(x > y)
let (x,_) <= (y,_) = f64.(x <= y)
let (x,_) >= (y,_) = f64.(x >= y)
let (x,_) != (y,_) = f64.(x != y)

let f64 x = dual0 x
let zero = f64 0
let one = f64 1
}``````

## Geometry

Sine we need the basic geometry functions of the ray tracer to work with both normal and dual numbers, we define it as a parametric module.

``````module mk_geometry(Float: expanded_field) = {

type Float = Float.t``````

The original Dex program implements vectors in three-dimensional space with three-element arrays. This is not optimal in Futhark for reasons explained elsewhere, so instead we import the vector module from another example program.

``module threed = import "3d-vectors"``

And instantiate it with whatever `Float` type we have at hand:

``module vspace = threed.mk_vspace_3d Float``

Convenient shorthands for various vector types and operations:

``````type Vec = vspace.vector
let (*>) = vspace.scale
let (<+>) = (vspace.+)
let (<->) = (vspace.-)
let (<*>) = (vspace.*)
let dot = vspace.dot
let length = vspace.length
let normalize = vspace.normalise
let cross = vspace.cross
let rotateY = flip vspace.rot_y``````

These vectors are not arrays and we cannot `map` them, but we can define our own mapping function.

``````let vmap (f: Float -> Float) (v : Vec) =
{x = f v.x, y = f v.y, z = f v.z}``````

Dex calls the scaling operator `.*`, but that name is not valid in Futhark. Note that the `length` function will shadow the one in the Futhark prelude.

Dex’s `Color` type is from the plot.dex file, which we’re not going to implement in its entirety.

Much of the following is just as straighforward in Futhark as in Dex, and so will not have much commentary.

``````type Color = Vec
type Angle = Float -- angle in radians

type Distance = Float

type Position = Vec
type Direction = Vec  -- Should be normalized.

type BlockHalfWidths = Vec

type ObjectGeom
= #Wall Direction Distance
| #Block Position BlockHalfWidths Angle

type Surface
= #Matte Color
| #Mirror

type OrientedSurface = (Direction, Surface)

type Object
= #PassiveObject ObjectGeom Surface
-- position, half-width, intensity (assumed to point down)

### The signed distance function

This function is the reason why the Futhark implementation is so convoluted, as we will need its derivative.

``````let sdObject (pos:Position) (obj:Object) : Distance =
match obj
case #PassiveObject geom _ ->
(match geom
case #Wall nor d -> Float.(d + dot nor pos)
case #Block blockPos halfWidths angle ->
let pos' = rotateY (pos <-> blockPos) angle
in length (vmap (Float.max Float.zero) (vmap Float.abs pos' <-> halfWidths))
case #Sphere spherePos r ->
let pos' = pos <-> spherePos
in Float.(max (length pos' - r) (f64 0)))
case #Light squarePos hw _ ->
let pos' = pos <-> squarePos
let halfWidths = {x=hw, y=Float.f64 0.1, z=hw}
in length (vmap (Float.max Float.zero)
(vmap Float.abs pos' <-> halfWidths))``````

With that done, the rest of the ray tracer can be written in a non-modularised way, using just floating-point numbers.

``}``

## Working with the geometry

We’ll need to instantiate the geometry module twice: once with ordinary floats, and once with dual numbers.

``````module geometry = mk_geometry f64_field
module geometry_dual = mk_geometry f64_dual_field``````

To use `geometry_dual.sdObject`, we’ll need to convert objects to objects-with-dual-numbers. For this, we’ll first define some utility functions.

``````let dual0 = f64_dual_field.dual0
let dual1 = f64_dual_field.dual1
let dualVec {x,y,z} = {x=dual0 x, y=dual0 y, z=dual0 z}``````

And then a bunch of boilerplate definitions.

``````let dualGeom (geom: geometry.ObjectGeom) : geometry_dual.ObjectGeom =
match geom
case #Wall nor d ->
#Wall (dualVec nor) (dual0 d)
case #Block blockPos halfWidths angle ->
#Block (dualVec blockPos) (dualVec halfWidths) (dual0 angle)
case #Sphere spherePos r ->
#Sphere (dualVec spherePos) (dual0 r)

let dualSurface (surf: geometry.Surface) : geometry_dual.Surface =
match surf
case #Matte c -> #Matte (dualVec c)
case #Mirror -> #Mirror

let dualObject (obj: geometry.Object) : geometry_dual.Object =
match obj
case #PassiveObject geom surface ->
#PassiveObject (dualGeom geom) (dualSurface surface)
case #Light squarePos hw r ->
#Light (dualVec squarePos) (dual0 hw) (dualVec r)``````

The non-dual-number geometry module contains various useful type definitions. Instead of replicating their definitions, we make abbreviations.

``````type Object = geometry.Object
type Position = geometry.Position
type Direction = geometry.Direction
type Distance = geometry.Distance
type Color = geometry.Color
type Surface = geometry.Surface
type OrientedSurface = geometry.OrientedSurface
type Vec = geometry.Vec``````

We also need some of the geometry functions:

``````let dot = geometry.dot
let (<+>) = (geometry.<+>)
let (<->) = (geometry.<->)
let (<*>) = (geometry.<*>)
let (*>) = (geometry.*>)
let normalize = geometry.normalize
let length = geometry.length``````

And a few more shorthands:

``let vec x y z : Vec = {x,y,z}``

Finally we can define the function that computes surface normals at a given position. This function is the reason we went to all that trouble with dual numbers.

``````let calcNormal (obj: Object) (pos: Position) : Direction =
geometry.normalize (vjp (\p -> geometry.sdObject p obj) pos 1)``````

## Random sampling

Since dex-prelude.fut defines the same random number generation interface as the Dex program, the sampling functions are straightforward ports.

First, a function for generating a number in a range:

``````let randuniform (lower: f64) (upper: f64) (k: Key) =
let x = rand k
in (lower + x * (upper-lower))``````

Sample within a square with side length 2*hw:

``````let sampleSquare (hw: f64) (k: Key) : Position =
let (kx, kz) = splitKey k
let x = randuniform (- hw) hw kx
let z = randuniform (- hw) hw kz
in {x, y=0.0, z}``````

Sample within a hemisphere in the direction of the normal vector:

``````let sampleCosineWeightedHemisphere (normal: Vec) (k: Key) : Vec =
let (k1, k2) = splitKey k
let u1 = rand k1
let u2 = rand k2
let uu = geometry.normalize (geometry.cross normal (vec 0.0 1.1 1.1))
let vv = geometry.cross uu normal
let ra = f64.sqrt u2
let rx = ra * f64.cos (2 * f64.pi * u1)
let ry = ra * f64.sin (2 * f64.pi * u1)
let rz = f64.sqrt (1.0 - u2)
let rr = (rx *> uu) <+> (ry *> vv) <+> (rz *> normal)
in normalize rr``````

## Ray marching

The essence of ray marching is simple: move along a given vector and find the first object we collide with. We do this the naive way, by invoking the distance function for every single object. This is not efficient, to put it mildly, but it is easy to implement.

First we need a `minimumBy` function. Despite Futhark lacking Dex’s support for ad-hoc polymorphism, which matters here as we end up passing in the less-than operator, it doesn’t end up looking too bad.

``````let minBy 'a 'o (lt: o -> o -> bool) (f: a->o) (x:a) (y:a) : a =
if f x `lt` f y then x else y

let minimumBy [n] 'a 'o (lt: o -> o -> bool) (f: a->o) (xs: [n]a) : a =
reduce (minBy lt f) xs xs``````

A scene is a collection of objects.

``type Scene [n] = [n]Object``

We can now define a function for finding the closest object, given a position.

``````let sdScene (objs: Scene []) (pos: Position) : (Object, Distance) =
let i =
minimumBy (<) (\i -> geometry.sdObject pos objs[i])
(indices objs)
in (objs[i], geometry.sdObject pos objs[i])``````

When we ray march, we either collide with nothing and disappear into the aether, hit a light, or hit an object, in which case we also produce information about the surface, including the surface normal.

``````type Ray = (Position, Direction)

type RayMarchResult
= #HitObj Ray OrientedSurface
| #HitNothing``````

The ray marching function is defined with the `iter` function in Dex. In Futhark, a `while` loop can express it in an equally natural way.

``````let positiveProjection (x: Vec) (y: Vec) =
dot x y > 0

let raymarch [n] (scene:Scene [n]) (ray:Ray) : RayMarchResult =
let max_iters = 100
let tol = 0.01
let (rayOrigin, rayDir) = ray
let (_, _, res) =
loop (i, rayLength, _) = (0, 10 * tol, #HitNothing)
while i < max_iters do
let rayPos = rayOrigin <+> (rayLength *> rayDir)
let (obj, d) = sdScene scene rayPos
let dNew = rayLength + 0.9 * d
in if d >= tol
then (i+1, dNew, #HitNothing)
else
let surfNorm = calcNormal obj rayPos
in if positiveProjection rayDir surfNorm
then (i+1, dNew, #HitNothing)
else (max_iters,
dNew,
match obj
case #PassiveObject _ surf ->
#HitObj (rayPos, rayDir) (surfNorm, surf)
case #Light _ _ radiance ->
in res``````

## Light sampling

These definitions are pretty straightforward, and similar to those in Dex.

To figure out whether a light is shining on us along a given path, we just march along that path and see if we hit a light. If we do, the radiance is that light.

``````let rayDirectRadiance [n] (scene: Scene [n]) (ray: Ray) : Radiance =
match raymarch scene ray
case #HitLight intensity -> intensity
case #HitNothing -> vec 0 0 0
case #HitObj _ _ -> vec 0 0 0``````

Shorthands for vectors along the cardinal axes will become useful in the following.

``````let xHat : Vec = vec 1 0 0
let yHat : Vec = vec 0 1 0
let zHat : Vec = vec 0 0 1

let relu (x: f64) = f64.max x 0

let probReflection ((nor, surf): OrientedSurface) (_:Ray) ((_, outRayDir):Ray) : f64 =
match surf
case #Matte _ -> relu (dot nor outRayDir)
case #Mirror  -> 0

let directionAndLength (x: Vec) =
(normalize x, length x)``````

The following function determines how much light is shining on a given point. It’s very similar to the Dex function. Dex uses an explicit accumulator that is updated via effects, but I don’t think it makes it much more readable in this case.

``````let sampleLightRadiance [n] (scene: Scene [n])
(osurf: OrientedSurface)
(inRay: Ray)
let (surfNor, _) = osurf
let (rayPos, _) = inRay
in loop radiance = vec 0 0 0 for obj in scene do
match obj
case #PassiveObject _ _ -> radiance
case #Light lightPos hw _ ->
let (dirToLight, distToLight) =
directionAndLength (lightPos <+> sampleSquare hw k <-> rayPos)
in if ! (positiveProjection dirToLight surfNor)
then radiance -- light on the far side of current surface
else
let fracSolidAngle = relu (dot dirToLight yHat) *
sq hw / (f64.pi * sq distToLight)
let outRay = (rayPos, dirToLight)
let coeff = fracSolidAngle * probReflection osurf inRay outRay

## Tracing

Almost done. Everything here is very similar to the Dex code.

``````type Filter = Color

let surfaceFilter (filter:Filter) (surf:Surface) : Filter =
match surf
case #Matte color -> filter <*> color
case #Mirror      -> filter

let sampleReflection ((nor, surf): OrientedSurface) ((pos, dir): Ray) (k: Key) : Ray =
let newDir = match surf
case #Matte _ -> sampleCosineWeightedHemisphere nor k
case #Mirror  -> dir <-> 2 * dot dir nor *> nor
in (pos, newDir)``````

We’re excluding Dex’s `shareSeed` field from the tracing parameters, since it seems mostly relevant for showing the impact of poor seeding on the convergence of the algorithm.

``````type Params = { numSamples : i32,
maxBounces : i32 }

let trace [n] (params: Params) (scene: Scene [n]) (init_ray: Ray) (k: Key) : Color =
(.2) <|
loop
(i, filter, radiance, ray) = (0, (vec 1 1 1), (vec 0 0 0), init_ray)
while i < params.maxBounces do
match raymarch scene ray
case #HitNothing ->
case #HitLight intensity ->
if i == 0
then (params.maxBounces, filter, intensity, ray)
case #HitObj incidentRay osurf ->
let (k1, k2) = splitKey (hash k i)
let outRayHemisphere = sampleReflection osurf incidentRay k2
let newFilter = surfaceFilter filter osurf.1

## Camera

The camera controls how the initial rays are sent into the scene.

``````type Camera =
{ numPix     : i64,
pos        : Position,
halfWidth  : f64,
sensorDist : f64 }``````

And now we’re ready to produce the n*n array of initial rays and RNG states.

``````let cameraRays (n: i64) (camera: Camera) : [n][n](Ray, Key) =
let halfWidth = camera.halfWidth
let pixHalfWidth = halfWidth / f64.i64 n
let ys = reverse (linspace n (-halfWidth) halfWidth)
let xs = linspace n (-halfWidth) halfWidth
let kss = tabulate_2d n n (\i j -> newKey (i32.i64 (1+i*n+j)))
let rayForPixel y x k =
let dx = randuniform (-pixHalfWidth) pixHalfWidth k
let dy = randuniform (-pixHalfWidth) pixHalfWidth k
in ((camera.pos,
normalize (vec (x + dx) (y + dy) (-(camera.sensorDist)))),
k)
in map2 (\y ks -> map2 (rayForPixel y) xs ks) ys kss``````

Most ray tracers perform multiple samples per pixel and take the average. The Dex program defines the `sampleAveraged` function to be (potentially) parallel. We know the parallelism would never be exploited anyway, so we define it as a sequential loop, mostly to make the random number state management simpler.

``````let sampleAveraged (sample: Key -> Vec) (n: i32) (k: Key) : Vec =
(loop acc = vec 0 0 0 for i < n do
(acc <+> sample (ixkey k (i64.i32 i))))
|> ((1/f64.i32 n) *>)``````

The Dex implementation uses an unusual relative colorisation strategy, where the final color values are all divided by the average intensity. This is likely so we won’t have to fiddle with the light intensities to avoid very bright or very dark images.

``````let meanIntensity image =
image |> flatten |> map (\{x,y,z} -> (x+y+z)/3) |> mean``````

Smile for the camera!

``````let takePicture [m] (params: Params) (scene: Scene [m]) (camera: Camera) =
let n = camera.numPix
let rays = cameraRays n camera
let sample (r, k) =
sampleAveraged (trace params scene r) params.numSamples k
let image = map (map sample) rays
let mean = meanIntensity image
in map (map ((1/mean)*>)) image``````

## Scene definition

The only odd thing here is that apparently our vectors rotate opposite of the way Dex does it, so I had to give the block a rotation of -0.5 instead of 0.5.

``````let lightColor = vec 0.2 0.2 0.2
let leftWallColor  = 1.5 *> vec 0.611 0.0555 0.062
let rightWallColor = 1.5 *> vec 0.117 0.4125 0.115
let whiteWallColor = (1/255) *> vec 255.0 239.0 196.0
let blockColor     = (1/255) *> vec 200.0 200.0 255.0

let neg = ((-1)*>)

let theScene : Scene [] =
[ #Light (1.9 *> yHat) 0.5 lightColor
, #PassiveObject (#Wall      xHat  2.0) (#Matte leftWallColor)
, #PassiveObject (#Wall (neg xHat) 2.0) (#Matte rightWallColor)
, #PassiveObject (#Wall      yHat  2.0) (#Matte whiteWallColor)
, #PassiveObject (#Wall (neg yHat) 2.0) (#Matte whiteWallColor)
, #PassiveObject (#Wall      zHat  2.0) (#Matte whiteWallColor)
, #PassiveObject (#Block  (vec 1.0 (-1.6) 1.2) (vec 0.6 0.8 0.6) (-0.5)) (#Matte blockColor)
, #PassiveObject (#Sphere (vec (-1.0) (-1.2) 0.2) 0.8) (#Matte (0.7 *> whiteWallColor))
, #PassiveObject (#Sphere (vec 2 2 (-2)) 1.5) #Mirror
]

let defaultParams : Params = { numSamples = 50
, maxBounces = 10 }

let defaultCamera : Camera = { numPix     = 250
, pos        = 10.0 *> zHat
, halfWidth  = 0.3
, sensorDist = 1.0 }``````

## Entry point

This is not part of the Dex program, but is needed if we want Futhark to produce anything for the outside world.

First, a function to convert floating-point colors to RGB-packed colors with 8 bits per channel. Note that we have to cap each channel, since there is no guarantee that the color components produced by the ray tracer cannot exceed 1.0.

``````let pix (c: Color) =
(u32.f64 (f64.min 255 (c.x * 255)) << 16) |
(u32.f64 (f64.min 255 (c.y * 255)) << 8) |
(u32.f64 (f64.min 255 (c.z * 255)))

let main n =
takePicture defaultParams theScene (defaultCamera with numPix = n)
|> map (map pix)``````
``> :img main 500i64`` ``````-- ==
-- input { 100i64 }``````