Recently I was using JAX to model a few particles moving and interacting in 3D space. Naturally, at some point I wanted to see the behavior of these p

Simple 3D visualization with JAX raycasting

submited by

Style Pass

2022-09-23 13:30:47

Recently I was using JAX to model a few particles moving and interacting in 3D space. Naturally, at some point I wanted to see the behavior of these particles along with some information about their state and interactions with my own eyes. Visualization options available to a researcher are plentiful today, but in the spirit of minimalism I decided to build a visualization with the same tools I was using to make the model. What surprised me was: the effort it took to get a reasonable 3D rendering and the amount of code I had to write was comparable to the amount of code needed to integrate a traditional graphics framework. In principle, I got a working tool for simple 3D scene visualization.

The rendering approach I took is based on raycasting Signed Distance Fields (SDF). Inigo Quilez popularized this approach and made ShaderToy, a thriving creative community of GPU shader developers. Neural Rendering Fields and Differentiable Rendering brought these ideas into machine learning world. There are a number of good SDF tutorials on the internet, including a blog post by Eric Jang, who also used JAX. I recommend having a look at them to get the basic understanding of SDF raycasting before reading further. In this tutorial I'd like to make emphasis on the simplicity and flexibility of using this approach in practice.

The scenes we’ll render consist of a few colored balls. I'm very excited about the underlying experiments which I originally set out to visualise, but that work is not yet ready for the publication, so let's generate some random dummies instead 😉