waxmorph.train.train

waxmorph.train.train(model, optimizer, loss_fn, *, source_pos, polarities, c, radii, targets=None, config=None, save_path=None, device='cuda')

Train through learned, mechanics, and diffusion rollout steps.

Optimizes L_shape + lambda_reg * sum_t ||dt_gns * GNS_X,t||^2. Histories and selection use post-update rollouts; the supplied optimizer is restored to the selected model. Feature gradients remain live; each graph build and Warp substep freezes its own contact or pair topology.

State and target arrays become C-contiguous float32. Positions and polarities have shape [N, 3], concentrations [N, C], and radii [N] with common nonzero N; input polarities are used as supplied for the first graph, whose angle feature assumes unit vectors; later learned updates apply normalization. loss_fn maps predicted [N, 3] and target [M, 3] positions to a scalar. Targets are unique (frame, [M, 3]) pairs with 0 <= frame < t_rollout. The source precedes indexed frames; frame 0 supervises the state after the first complete rollout step.

device must be accepted by Torch and Warp. A new save_path receives the best GNS checkpoint and compressed log. An existing trusted checkpoint must match the supplied model; its weights load into that object, optimizer state starts fresh, and one refinement update preserves both existing files. See TrainResult for the log schema. The JAX counterpart requires CUDA for its Warp custom VJPs and adds a static edge-capacity setting.

Parameters:

• model (GNS)

• optimizer (Optimizer)

• loss_fn (Callable[[Tensor, Tensor], Tensor])

• source_pos (ndarray)

• polarities (ndarray)

• c (ndarray)

• radii (ndarray)

• targets (list[tuple[int, ndarray]] | None)

• config (TrainConfig | None)

• save_path (str | Path | None)

• device (str)

Return type:

TrainResult