Simple raytracing renderer in Jax

Supports

• Basic deterministic ray casting
• Stochastic path tracing with many rays per pixel (batched or progressive)
• Specular and diffuse reflections
• Emissive materials and diffuse interreflection
• Antialiasing

The renderer is fully jitted and batches render calls along the #rays and viewplane pixels dims, see:

jax-raytracing/raytrace.py

Lines 34 to 37 in 2da55e9

	@jit
	@partial(vmap, in_axes=(1, 1, 1, None, None, None, None, None, None))
	@partial(vmap, in_axes=(0, 0, 0, None, None, None, None, None, None))
	def full_ray_trace(

Samples

Specular and diffuse reflection

Emissive materials

Shadows & occlusion

Video sample here (very few traces per frame -> very noisy).