Multax: a warp accelarated multiphase flow engine and handler

[DimitrisDimitropoulos]

Multax is my experiment on a multiphase flow solver that simulates: particle-fluid interactions, drag, gravity, thermodynamics, phase change, and collisions, leveraging an Eulerian-Lagrangian approach. The project started as a fully differentiable engine built entirely in JAX with particle-level parallelism. While JAX handled the physics cleanly, two computational bottlenecks pushed me to incorporate warp: rendering (for the custom visualization engine) and collision detection.

The custom visualization engine initially used Gaussian splatting implemented as a JAX scatter followed by a 2D convolution for smoothing. This proved expensive, since the convolution scales as O(H·W·k²) regardless of particle count. Therefore, I replaced it with a Warp kernel that writes Gaussian splats directly to the output buffer using atomic additions per channel reduced rendering to O(N·r²), efficient enough for the scope of the problem

Collision detection in JAX faced similar issues. Broad-phase neighbor search via Morton code sorting works for moderate counts but degrades with dense systems, and the narrow phase requires irregular memory access patterns poorly suited to JAX's execution model. Here comes warp's spatial hash grid, which achieves O(1) (for the average-case) neighbor lookup. This foundation enabled a full Continuous Collision Detection (CCD) engine, allowing dense particle systems and eliminating tunneling of particles.

The whole incorporation of warp inside the main JAX codebase was done thanks to DLPack. All warp kernels are wrapped with jax_kernel FFI and execute as native nodes inside JAX's XLA graph. Meanwhile, jax and warp interact via dlpack with minimal overhead, thanks to a zero-copy paradigm.

The result is a multiphase engine able to simulate 50M+ particles on consumer hardware.

The project is hosted in this repo

Here are some examples:

cellular_flow.mp4.mp4

cylinder.mp4

wall-thermal.mp4

ccd-lcp-showcase.mp4

[shi-eric]

Congratulations on your project! It's quite encouraging to see that you were able to move the computations you wanted into Warp while allowing the rest of the application to remain in JAX.