Installation

Install the base package with pip:

pip install waxmorph

The base install includes NumPy 1.24+, SciPy 1.11+, PyTorch 2+, tqdm 4.66+, and Warp 1.10+. It provides the forward simulator and the default PyTorch graph, GNS, losses, and training APIs. Extras add optional loss, backend, mesh, and rendering dependencies.

Add mesh and rendering support:

pip install "waxmorph[simulation]"

Add GeomLoss and PyKeOps; this extra currently includes simulation:

pip install "waxmorph[learning]"

Add the explicit JAX/Equinox backend:

pip install "waxmorph[jax]"

Use JAX with CUDA 12:

pip install "waxmorph[jax-cuda]"

Set up a development checkout with every optional dependency group:

git clone https://github.com/Computational-Morphogenomics-Group/waxmorph.git
cd waxmorph
pip install -e ".[all]"
pre-commit install

all contains the CPU JAX stack. GPU development adds jax-cuda:

pip install -e ".[all,jax-cuda]"

What each extra pulls in

simulation

Mesh processing and movie rendering: PyVista, trimesh, VTK, OpenGL support, and usd-core.

learning

GeomLoss and PyKeOps for optional PyTorch point-cloud losses, plus simulation.

jax

JAX, Equinox, Optax, and ott-jax, plus simulation.

jax-cuda

CUDA 12 JAX support plus jax.

all

Development, documentation, JAX, learning, and simulation dependencies. jax-cuda supplies CUDA JAX support.

Hardware notes

The forward simulator and PyTorch/Warp training bridge accept supported CPU or CUDA devices; examples default to CUDA. JAX training with mechanics or diffusion uses custom VJPs on CUDA. JAX CPU training uses mech_steps=0 and diff_steps=0. WarpMovieRenderer defaults to CUDA; its OpenGL backend also requires a working graphics and video stack.

Note

USD export uses Warp’s USD renderer and the usd-core package installed by the simulation extra.