SDFs and the Fast sweeping algorithm in Jax

Summary

SDFs and Fast sweeping in JAX This is going to be a fun blog - we'll explore the intuition behind level sets, the Eikonal equation, and implement a speedy algorithm for solving this equation, called the fast sweeping method, in JAX. The level set method and the eikonal equation Fast sweeping method Code numpy jax In action Benchmarks A note on parallel FSM References Fig 1: FSM in action I was recently researching a problem that involved interface evolution over time. Our interface was represented by a set of points. To evolve this interface along its normal direction, we were approximating the normals for these points, extending the points along the normal direction, then resampling them to maintain point density - since, if the shape had expanded, the point density would have decreased. This method of propagating interfaces, involving tracking particles on the front as it evolves, is known as Lagrangian front evolution, and comes with a host of problems, such as resampling issues, handling particles that "grow into" a surface, etc. The level set method and the eikonal equation An alternative view of propagating interfaces was provided by Sethian and Osher, where they developed the level set theory of propagating interfaces. The key difference is that this is an Eulerian approach - the interface is tracked implicitly on a fixed grid, not as particles on the surface as before. In the level set technique, this implicit representation of the surface is the zero level set of a function. What that means, is that you'll have some function defined on your grid, and the grid points where this function is zero will represent your surface. This function (let's call it ϕ\phiϕ) is the level set function, and is a higher order function, with x,y (in 2D) and time as it's inputs. ϕ(xgrid,ygrid,t=k)\phi(x_{grid}, y_{grid}, t=k)ϕ(x​grid​​,y​grid​​,t=k) gives you the zero level set at time kkk, describing your interface. It is this function that we want to learn/approximate. Most of...