Skip to content

Folders and files

NameName
Last commit message
Last commit date

Latest commit

 

History

9 Commits
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 

Repository files navigation

diff-vbd

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.

Installation

pip install -e .          # CPU
pip install -e ".[cuda]"  # CUDA 12 GPU wheels
pip install -e ".[dev]"   # + test dependencies

Requires Python 3.10+.

Quickstart (library API)

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.).

Quickstart (CLI)

Installing the package registers a diff-vbd console script:

diff-vbd --config examples/cantilever_beam.yaml --platform cpu --out out.npz

Key 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).

Visualization

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-npz
runs/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.

Configuration

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 free

Programmatic 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).

Contact

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 sphere

See 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.

What "intersection-free" actually guarantees

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.

Line search

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.

Energy API

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.

Testing

pytest

The test suite is self-contained: it synthesizes small meshes in temp directories, so no external data files are required.

Validation

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.

Meshes are authored, not generated

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%

Repository layout

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

About

A minimal implementation of vertex block descent in JAX.

Resources

Stars

0 stars

Watchers

0 watching

Forks

Releases

Packages

Contributors

Languages