Training APIs#
Training optimizes neighbor-dependent local updates so selected rollout frames approach target point clouds.
The top-level waxmorph.train module is a convenience alias for the
default PyTorch training API. The backend-specific entry points remain available
from waxmorph.torch.train and waxmorph.jax.train.
The training loop expects:
an initial particle state:
source_pos,polarities,c(signaling- molecule concentrations), andradii;a
waxmorph.gnn.GNSmodel or its backend-specific equivalent;a PyTorch optimizer, or an Optax transformation and state for JAX;
a backend-compatible shape loss;
a list of
(frame, positions)targets.
Target frames are zero-based indices within a model evaluation. The source is
the trajectory’s initial entry; frame 0 supervises the state after the first
complete learned, mechanics, and diffusion rollout step. See
Supervise multiple target frames.
Post-update selection#
n_epochs counts optimizer updates. Each history contains n_epochs
entries for the resulting models; evaluating them requires one extra rollout.
The returned model, best_loss, best_epoch, and best trajectory describe
the same post-update epoch. The initial evaluation supplies the first gradient;
the n_epochs updated models form the selection candidates.
Checkpoint and backend contracts#
Model checkpoints contain model weights and configuration. During a new PyTorch
run, the supplied optimizer is restored with the selected model. Refining a
trusted, compatible checkpoint starts with fresh optimizer state, performs one
update, and preserves the existing files. JAX validates the saved configuration
and PyTree, replaces the model, reinitializes Optax state, and likewise performs
one refinement update while preserving existing files. TrainResult contains
the selected model and log; callers manage Optax state separately.
JAX training bounds directed edges and neighbor pairs with
max_edges_factor. Its Warp mechanics and diffusion run on CUDA; CPU runs use
mech_steps=0 and diff_steps=0.