A differentiable Vertex Block Descent (VBD) solver for soft-body simulation, built on JAX. It loads tetrahedral meshes, assembles boundary conditions from selector geometry, and time-integrates a stable Neo-Hookean elastic model (Smith et al. 2018) with optional Chebyshev acceleration and line search.
pip install -e . # CPU
pip install -e ".[cuda]" # CUDA 12 GPU wheels
pip install -e ".[dev]" # + test dependenciesRequires Python 3.10+.
import diff_vbd as dv
# Build a problem from a YAML config (see examples/ for the schema)
problem = dv.load_problem_from_yaml("examples/cantilever_beam.yaml")
state = dv.initial_state(problem)
final_state, history = dv.simulate(problem, state, num_steps=100)
dv.export_simulation_npz("out.npz", problem, history)Everything in the public API is re-exported from the top-level package — see
src/diff_vbd/__init__.py for the full surface
(assemble_problem, step, simulate, boundary-condition builders, exporters, etc.).
Installing the package registers a diff-vbd console script:
diff-vbd --config examples/cantilever_beam.yaml --platform cpu --out out.npzKey flags: --platform {cpu,gpu} (default gpu), --steps N (override the config),
--out PATH (NPZ trajectory output), and GPU memory controls
(--gpu-preallocate, --gpu-mem-fraction).
Trajectories stream to XDMF + HDF5 as they are computed, which ParaView opens directly:
pip install 'diff-vbd[viz]'
diff-vbd --config examples/block_on_plane.yaml --precision float64 \
--paraview-out runs/block --volume-stride 10 --no-npzruns/block/
surface.xdmf open this to scrub the trajectory
volume.xdmf open this to clip into the interior or threshold J < 0
traj.h5 both indices point here; the only large file
colliders.vtp static plane/sphere geometry, load alongside either index
diagnostics.csv min_toi, active pairs, minimum Jacobian per frame
config.yaml copy of the input, so the run can say what produced it
Or from the library, composing sinks yourself:
from diff_vbd.export import open_run_directory
sink = open_run_directory("runs/block", volume_stride=10)
final_state, _ = dv.simulate(problem, state, num_steps=1000, sink=sink)Two streams, two strides. The surface is what you scrub through and it is cheap — boundary vertices grow as V^(2/3), so a million-vertex mesh has roughly forty thousand of them, about 1 MB a frame. The volume carries the tets and per-element fields, is tens of MB a frame, and is rarely needed at full temporal resolution. Writing both at every frame for a million-vertex, thousand-step run costs ~76 GB; the surface every frame plus the volume every fiftieth costs ~2.6 GB and loses nothing you would routinely look at.
Fields are chosen, not assumed — see POINT_FIELDS and CELL_FIELDS in
export/fields.py. The one worth knowing about is
jacobian: stable Neo-Hookean is finite and smooth through inversion by construction, so an
inverted element raises nothing and keeps simulating. Thresholding J < 0 in ParaView is the
only thing that will tell you it happened, and diagnostics.csv carries the same signal as
min_jacobian/inverted_tets per frame.
Frames are written and forgotten, so keep_history defaults to False whenever a sink is
attached — retaining the trajectory is the cost the sink exists to avoid. Pass
keep_history=True for both, which is what the CLI does unless you give it --no-npz.
sink.close() runs in a finally: _audit_step raising mid-run is designed behaviour here,
and the run that stopped at step 37 leaves 36 readable frames plus the diagnostics explaining
why it stopped.
Simulations are described by a YAML config. Minimal cantilever example
(examples/cantilever_beam.yaml):
mesh:
path: beam.msh # Gmsh 2.2 binary tetrahedral mesh
selectors:
clamp:
path: bc_selector.stl # binary STL used to select vertices
material:
mu: 4.638e5 # Lamé parameters + density
lam: 4.174e6
density: 1.04
simulation:
steps: 100
solver:
dt: 0.02
num_iterations: 3
eps: 1.0e-6
acceleration: { enabled: false, rho: 0.95 }
line_search: { enabled: true, num_alphas: 9 } # on by default
body_force: [0.0, 0.0, -9.81]
dirichlet:
- selector: clamp
mode: position
components: ["0.0", "0.0", "0.0"]The cantilever example ships with its mesh assets (examples/beam.msh,
examples/bc_selector.stl), so it runs as-is. For your own simulations, supply tetrahedral
meshes and selector geometry and point the config path fields at them (relative paths
resolve against the config file's directory).
Dirichlet components are sandboxed expressions evaluated per step. Besides time t they
may reference the constrained vertices' reference coordinates x, y, z and the
cylindrical radius r = hypot(x, y), so a single expression can prescribe a field that
varies over the selector — e.g. a cylinder bore driven radially outward by 0.05:
dirichlet:
- selector: bore
mode: position
components: ["0.05 * x / r", "0.05 * y / r", "0.0"]body_force accepts the same expression language: three numbers are a uniform
acceleration (gravity), three strings are a spatially varying acceleration field
evaluated per vertex over the reference coordinates at assembly — e.g.
body_force: ["0.0", "0.0", "-9.81 * r / 10.0"]. Time dependence is rejected: the
field is baked into the problem once, and an expression in t would silently freeze
at its initial value. This is the config face of manufactured-solution loading
(diff_vbd.validation.manufactured derives the exact body force for an analytic
displacement field; see examples/spatial_body_force.yaml for the YAML side).
A Dirichlet entry may also constrain only some axes: axes: [x, y, z] booleans, default
all true. axes: [false, false, true] with zero components is a roller — the
frictionless sliding support that symmetry planes, plane-strain end faces, and uniaxial
platens need, with the unconstrained axes left to the solver:
dirichlet:
- selector: end_faces
mode: position
components: ["0.0", "0.0", "0.0"]
axes: [false, false, true] # pin u_z only; radial slip is freeProgrammatic callers can instead supply an explicit per-vertex displacement array
(mode: "array" plus dirichlet_values={...} in assemble_boundary_conditions) — the
escape hatch for fields easier to state as data, such as manufactured solutions. The load a
prescribed motion transmits is read back with reaction_forces(problem, state): the
constraint reaction at every vertex, summing to the total load over any driven set
(displacement-controlled indentation measures P this way; a roller's reaction is normal
to its sliding plane).
Intersection-free contact via the IPC log barrier: continuous collision detection, lagged Coulomb friction, analytic colliders (plane / sphere), and mesh self-collision. Colliders and mesh-mesh pairs go through the same barrier, the same friction model and the same time-of-impact filter — they are one code path, not two.
contact:
enabled: true
d_hat: 1.0e-3 # activation distance: the barrier turns on inside this gap
# kappa: 1.0e3 # omit to derive a stiffness matched to the inertia term
friction_mu: 0.3 # Coulomb coefficient; 0 for frictionless
eps_v: 1.0e-4 # sliding speed below which a contact is treated as stuck
self_collision: false # also detect the mesh against itself
colliders:
- kind: plane
normal: [0.0, 0.0, 1.0] # points into the free half-space
offset: 0.0
- kind: sphere
center: [0.0, 0.0, -2.0]
radius: 1.0
outside: true # false to be contained *inside* the sphereSee examples/block_on_plane.yaml. A body supported by a
collider needs no Dirichlet constraint — resting on the ground is well-posed on its own,
so that example has no dirichlet block at all.
Contact needs float64. Pass --precision float64. The barrier resolves a gap orders of
magnitude smaller than the mesh coordinates, and it obtains that gap by subtracting them:
in float32, at a coordinate of 100, the absolute resolution is 1.2e-5, so a small positive
gap can round to zero or below — and log(gap / d_hat) is then NaN, losing
intersection-freedom to rounding before any solver logic runs. assemble_problem rejects a
d_hat the float type cannot resolve rather than letting it fail later as a NaN.
The guarantee is enforced by a time-of-impact filter on the whole sweep, not per vertex. A per-vertex bound certifies one vertex moving against frozen geometry, but VBD solves a whole colour in parallel from one snapshot — if a contact pair's endpoints share a colour, both move and neither certificate covers their combined motion. Worse, Chebyshev acceleration extrapolates after every local solve has finished, where no per-vertex check can see it. Scaling the entire update by a single conservative factor bounds where the mesh actually ended up, whatever the interior did.
For analytic colliders the bound is closed-form and holds unconditionally. For mesh-mesh it rests on the pair set being complete, and that needs two things that work only as a pair:
- the detection band is derived from how far the mesh is about to move, not from
d_hat, so every pair that could reach contact this step is in the set; and - the sweep clamps each vertex's cumulative displacement, so none can leave the band it was detected under.
The cost is real and it is the honest limit: the band grows with speed, and past roughly one
body-length of travel per step no once-per-step candidate set can be both bounded and
complete. At that point the pair capacity overflows and says so. That is by design — the
alternative is a band that quietly under-covers and lets a surface tunnel unseen. Reduce
solver.dt, or set self_collision_ccd: false to trade the guarantee for speed (it warns).
Two limits worth stating plainly:
- "No non-adjacent surface primitive pair intersects." Primitives that share a vertex are excluded from detection — a vertex always touches its own triangles — so a triangle inverting through its own neighbour is a tet-inversion problem, not a contact one.
- Grazing motion at speed is over-throttled, not unsafe. The conservative advance is proportional to the gap, so two surfaces sliding past each other fast at a tiny separation exhaust the iteration budget and the step gets shorter than it needed to be. The solver reports this rather than silently freezing.
self_collision also requires a conforming tetrahedralisation, and checks it: a
non-conforming split leaves interior faces unpaired, they get reported as surface faces, and
self-collision then finds phantom contacts deep inside a solid body at rest. Every surface
edge must be shared by exactly two triangles, or assembly fails. (For a structured grid, use
the Kuhn 6-tet split, which conforms unconditionally, rather than the 5-tet one.)
self_collision allocates fixed-capacity pair buffers (capacity, max_per_vertex).
They are fixed because rebuilding them each step at the same shape is a jit cache hit,
whereas a capacity that grew with the contact count would recompile the solver every step.
Overflow is a hard error, never a silent drop of whichever pairs happened to be last.
max_per_vertex is a direct linear multiplier on solve cost, not just a memory bound —
every vertex evaluates all of its slots on every local solve, of which there are
(colours × sweeps) per step, and the whole thing sits inside a Hessian. Measured on
examples/two_blocks.yaml: 128 → 9 s/step, 256 → 19 s/step. Size
it tightly; the overflow error names the exact number you need. capacity is much cheaper (it
only feeds the time-of-impact filter, a shallow kernel), so give that one room.
Flat-on-flat contact between two coincident faces is the case that stresses max_per_vertex:
it generates a large edge-edge fan at a single vertex.
See examples/two_blocks.yaml for the mesh-mesh path end to end —
a block dropped onto a fixed block, with no analytic colliders at all. Its assets are generated
by python examples/make_two_blocks.py.
num_alphas: N builds a linear step-size grid from 1.0 down to 0.0 inclusive; raise it to
refine the search. Refining is cheap — the per-vertex gradient, Hessian and solve dominate, so
9 alphas cost only ~9% more than 4, and every alpha is evaluated in parallel under jax.vmap.
The 0.0 endpoint matters: it lets a vertex decline to move when every positive step would
increase its local objective. Without it the search is forced to return an objective-increasing
step for such vertices, which pumps energy into the mesh and makes the solve diverge as
num_iterations grows. Set alphas: [...] instead to pin the grid explicitly.
The solver never forms a global energy — VBD's premise is a 3×3 solve per vertex against a
frozen block — but the energies it descends are also stated once, whole-mesh, in
solver/potential.py:
import diff_vbd as dv
E = dv.potential_energy(problem, positions, previous_positions) # elastic + contact
G = dv.variational_energy(problem, positions, inertial_target, previous_positions)elastic_potential and inertia_potential are the individual terms; variational_energy
is G(x), the implicit-Euler objective a VBD sweep decreases. One energy definition, two
consumers: the local vertex objective and this global view share the same kernels
(tet_energy, the contact energies), so the forward solve and a future tangent-stiffness
or adjoint path cannot silently disagree about what the energy is. The local objective
sums the elements incident to a vertex, so it counts each tet (and each contact pair)
four times; the tests pin that factor exactly rather than assuming it.
Two limits worth stating: gradients of these potentials do not differentiate through
contact detection (the active set and classifications are frozen inputs, refreshed by
redetect_contacts), and monotone descent of G is only guaranteed in the regime the VBD
paper's argument covers — line search on, Chebyshev acceleration off, no contact filter
rescaling the sweep.
pytestThe test suite is self-contained: it synthesizes small meshes in temp directories, so no external data files are required.
tests/test_validation.py checks the solver against closed-form results, in three tiers that
fail for different reasons and are fixed in different places.
Constitutive and reference verification needs no mesh and always runs. The stable
Neo-Hookean energy is compared against its plain textbook statement — solver/materials.py
rearranges it to avoid cancellation near F = I and claims the result is algebraically
identical, which is checked rather than assumed (it agrees to ~4e-15, including at det F < 0).
The 4/3 and 5/6 parameter remapping is validated by asserting that the small-strain error
against linear elasticity falls at first order in strain — the rate, not a magnitude, which
is what distinguishes a correct remapping from a small constant offset.
Mesh and kinematic verification runs against the author-supplied asset. It checks the
mesh first — orientation, volume, clamp planarity, aspect ratios — then makes exact statements
about the discretisation: the classical patch test (u = A x + b must give every element
exactly F = I + A), rigid translation and rotation costing no energy, uniform scaling giving
J = s³, and lumped mass equalling density × volume.
Analytical validation compares a solved cantilever deflection to Timoshenko beam theory.
This one is xfail(strict=True): the solver does not currently reach static equilibrium, so
it records the gap rather than confirming agreement. Strict means fixing convergence makes the
test XPASS and turn the suite red, forcing the marker off.
Validation geometry lives in assets/validation/: CAD surfaces under
surfaces/, boundary-condition selector volumes under selectors/, and the tetrahedral
.msh files the loader reads at the root. Surfaces are tetrahedralised with
fTetWild, vendored as a submodule, which writes the Gmsh 2.2 binary
the parser requires by default.
REQUIRED_ASSETS in validation/assets.py is the
specification for each mesh — purpose, geometry, meshing, export — and is quoted verbatim into
the error message when one is missing, so the requirements cannot drift from the code that
consumes them. Tests that need geometry skip with that spec rather than substituting a
generated mesh.
The geometry is constrained because beam theory is only valid inside a three-sided envelope —
slenderness L/h ≥ 20, small deflection δ/L < 5%, and Poisson ratio below ~0.40 (linear
tets lock near incompressibility). Fixing one side by hand tends to break another:
examples/beam.msh is 100×20×20, slenderness 5, so it is not usable as a beam-theory
case; thinning it to 100×5×5 at the same stiffness makes it sag 45% of its length, which
is geometrically nonlinear. BeamReference.violations checks all three and reports which one
a case fails:
from diff_vbd.validation import beam_reference
print(beam_reference(length=100., height=5., depth=5.,
mu=1.1772e7, lam=1.7658e7, density=1.04).describe())
# L/h=20.0, delta/L=2.01%, nu=0.300, shear=0.26% -> 2.0052 +-2%src/diff_vbd/ # the installable package
config/ # YAML config loading
setup/ # mesh & selector I/O, topology, boundary conditions
problems/ # ready-made problem builders (e.g. cantilever)
solver/ # VBD kinematics, stable Neo-Hookean materials, time integration
export/ # NPZ archives, and streaming XDMF/HDF5 for ParaView
validation/ # closed-form references, mesh QA, asset loading
cli.py # `diff-vbd` command-line runner
runtime_config.py # JAX/XLA backend configuration
examples/ # example YAML configs
assets/validation/ # author-supplied validation meshes (Gmsh 2.2 binary)
tests/ # unit tests