Potential JAX-based physics for real-time simulation in Slicer?

Community
1

Hello Slicer community,

I’m a developer interested in physics simulation. My project, brax viewer, visualizes JAX-based physics engines in real time, which is a new open-source package developed as part of Google Summer of Code. Briefly, JAX is an extremely fast, differentiable computation package that runs on GPU.

I am looking for a way for my viewer to be used with Slicer. So I share it with the community and have been browsing some relevant discussions. In threads like the one on biomechanics simulation for pectus excavatum, Blender can animate deformation, and tools like OpenSIM or FEBio can simulate but require patient-specific material properties for precision.

Could JAX engines be used to reverse-estimate these material properties in real-time?

There are two relevant physics simulation engines suited for this problem:

  • MJX for rigid body dynamics
  • JAX-FEM for Finite Element Method

JAX has two advantages:

  • Speed: Because JAX runs natively on the GPU, it’s fast enough for real-time simulation. It can also host many environments in parallel, which is ideal for robot trajectory exploration.
  • Differentiability: The entire simulation is differentiable, making it powerful for optimization problems like material property estimation and learning-based robot training.

Based on this discussion, Slicer’s WebServer and OpenIGTLink seem to be good interfaces.

Thank you for your time and any thoughts you have! I am open to any suggestions.

2

@Bruce_Hu, this sounds interesting, thanks for sharing. We have a 3D Slicer extension integrating the SOFA framework ( GitHub - Slicer/SlicerSOFA: 3D Slicer extension to enable simulations using the SOFA framework ). A common pattern to integrate simulation libraries in 3D Sicer is using numpy.array objects as a common language to transfer data between the foreign library and 3D Slicer. This is what is used in SlicerSOFA and more recently @pieper has succesfully demoed integration with NVIDIA Warp.

@Bruce_Hu, what do you think is the advantage of using Brax viewer with streamed slicer data vs. integrating JAX directly in Slicer? 3D Slicer has already a wealth of functionality for visualization.

3

Welcome @Bruce_Hu - it sounds like you are into some interesting things and I hope you’ll find Slicer useful for your work. Now is an exciting time for simulation and medical imaging and as @RafaelPalomar says, we are actively exploring a lot of options.

You may already know of this, but the Newton Physics project also looks interesting in this space.

In addition to the WebServer and OpenIGTLink, you may want to explore RPyC and shared memory as in the client and server examples here. There’s some description of our electromagnetic simulation examples and visualizations here. Note that we preferred the client-server model for SimNIBS because, like JAX-FEM, it is GPL. But it’s also an option for decoupling simulation from UI processes in general.

As a general rule we also like the option we like being able to have the option of installing all the packages in the Slicer python interpreter for portability, although that’s not always possible. It does seem to work well for Warp in my experience.

Yes, this is a great idea and definitely something that would be valuable. We often have 4D datasets and other data that could be used for this purpose.