Nudged Elastic Band
CO
3d ago+2985+2985+0123456789012345678901234567890123456789-217-217-012345678901234567890123456789
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, andassemble_pathwith climbing-image NEB, FIRE (vv_fireandase_fire) / GD optimizer support,FixAtomsconstraint propagation, per-imageTrajectoryReporterfilenames, and_NEBModelAdapterso all optimizers receive constraint-masked NEB forces.
- Added
- Optimizer Groups (
group_idx) acrossSimState, FIRE, Runners, and Autobatchers:SimState: Addedgroup_idxand_group_attributessupport with invariant validation in__post_init__,__setattr__,pop,from_state, slicing/remapping, and docstrings explaining whygroup_idxexists (n_systemsvsn_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/integraterunners: Supported group-level batching, lazyIterablestreaming without input mutation,estimate_max_memory_scalerper-group aggregation, andrestore_original_orderfor multi-system groups.
- Tutorial & Tests:
- Added
examples/tutorials/nudged_elastic_band.py(linked indocs/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, andtests/test_autobatching.py.
- Added