waxmorph.jax.train.TrainResult

Contents

waxmorph.jax.train.TrainResult#

class waxmorph.jax.train.TrainResult(model, log)[source]#

Bases: object

Best post-update model and its run diagnostics.

losses_total, losses_shape, and unweighted losses_l2 each have length n_epochs and describe post-update rollouts. The best trajectory includes the source frame and has shapes [t_rollout+1, N, 3] for best_traj_pos and best_traj_pol, and [t_rollout+1, N, num_molecules] for best_traj_c. best_epoch is zero-based; best_loss matches that model and trajectory. target_frames is sorted, and every config field appears as config_<field>. JAX therefore adds config_max_edges_factor to the common Torch log schema.

Parameters:
Contents