Nudged Elastic Band

CO
3d ago
+2985-217

This revives @mstapelberg's NEB work from https://github.com/TorchSim/torch-sim/issues/90 and adds optimizer-group (group_idx) support for running batched Nudged Elastic Band calculations cleanly alongside SimState, optimizers, runners, and autobatchers.

Summary of Changes

  • NEB Workflow (torch_sim/workflows/neb.py):
    • Added NEB, neb_init, neb_step, validate_endpoints, interpolate_path, compute_tangents, calculate_neb_forces, and assemble_path with climbing-image NEB, FIRE (vv_fire and ase_fire) / GD optimizer support, FixAtoms constraint propagation, per-image TrajectoryReporter filenames, and _NEBModelAdapter so all optimizers receive constraint-masked NEB forces.
  • Optimizer Groups (group_idx) across SimState, FIRE, Runners, and Autobatchers:
    • SimState: Added group_idx and _group_attributes support with invariant validation in __post_init__, __setattr__, pop, from_state, slicing/remapping, and docstrings explaining why group_idx exists (n_systems vs n_groups).
    • fire_init / fire_step: Adapted reductions (batched_vdot(..., n_systems=n_groups)) and per-group FIRE state (dt, alpha, n_pos) to operate per optimizer group while preserving zero-GPU-sync execution.
    • BinningAutoBatcher & InFlightAutoBatcher + optimize/integrate runners: Supported group-level batching, lazy Iterable streaming without input mutation, estimate_max_memory_scaler per-group aggregation, and restore_original_order for multi-system groups.
  • Tutorial & Tests:
    • Added examples/tutorials/nudged_elastic_band.py (linked in docs/tutorials/index.rst) comparing batched TorchSim NEB paths against sequential ASE NEB.
    • Added unit and regression tests in tests/workflows/test_neb.py, tests/test_state.py, and tests/test_autobatching.py.