[CompRhys]

This revives @mstapelberg's NEB work from #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.

[mstapelberg]

Hi @CompRhys happy to work on this in 2 weeks! Defending my thesis next Friday so slammed with that but will have time to work on this and a few other things if interested after the fact.

[CompRhys]

Best of luck! the @ was just to try ensure you still got credit for the work that you did here on neb

[CompRhys]

This matches ASE, matching ASE treats all images as a single opt batch. Doing each image with it's own fire state led to ~3x the number of opt steps being needed.

This design is really messy and shouldn't be merged. Results shown from the two examples. Those should be cleaned up also.

[CompRhys]

I am unsure why at the end torchsim in this tutorial example stops being an envelope function to the ASE results.

[CompRhys]

To get this to work a concept of group_idx needs to be added to the simstate and then passed to all the optimizers. This is because the whole path is optimized at once.

[thomasloux]

Hey, I'm not super familiar with NEB. For instance, group_idx seems necessary but adds a lot of complexity. Could you add a short doc explaining this design choice, in this case probably by explaining why/how NEB needs to optimize the entire path directly. This overall looks great. But I would need more time to have a thorough read

[orionarcher]

I agree that we do need a group_idx but I don't think this is the right approach to operationalizing it.

As I understand it, if we split the state before autobatching then we limit the batch size to the group size, which will throttle performance in the case of small groups.

Instead, I think we need to modify the logic of the autobatchers so that they can either split by system_idx or split by group_idx. Then we can re-use the same logic to pack the maximum number of groups into memory that we use to pack the maximum number of systems into memory. This will require a deeper rewrite to the autobatchers though.

I suggest having a autobatcher property on instantiation that selects "group_idx" or "system_idx" as the batching index. Or could be a bool batch_by_group=False.

I spent a fair amount of time thinking about this last year and I think everything should transition over nicely, no blockers on an architectural level.

[CompRhys]

I suggest having a autobatcher property on instantiation that selects "group_idx" or "system_idx" as the batching index. Or could be a bool batch_by_group=False.

I was trying to get it such that group_idx was the thing that autobatchers worked off and that the default behavior would be single system groups. I think that's cleaner than having switching logic?

[cubic-dev-ai]

All reported issues were addressed across 12 files

<sub>Shadow auto-approve: would not auto-approve because issues were found.
Tip: cubic can generate docs of your entire codebase and keep them up to date. Try it here.

View guided diff | Turn on auto-fix | Re-trigger cubic</sub>

[CompRhys]

Merged latest main into feature/neb-workflow, resolved all conflicts, and addressed the review feedback (@orionarcher, @thomasloux, and automated review items):

• Documented group_idx (optimizer groups vs system_idx systems) in torch_sim/state.py, torch_sim/workflows/neb.py, and examples/tutorials/nudged_elastic_band.py, and linked the tutorial in docs/tutorials/index.rst.
• Kept origin/main zero-GPU-sync FIRE and lazy BinningAutoBatcher / InFlightAutoBatcher streaming improvements alongside group_idx group batching (estimate_max_memory_scaler, restore_original_order, runner group-to-system index mapping).
• Fixed SimState.pop, __setattr__ group validation, SimState.from_state ambiguous n_groups matching, O(n) index remapping, NEB _NEBModelAdapter (vv_fire + FixAtoms constraint masking), GD lr alias normalization, endpoint constraint validation/propagation, and per-image NEB trajectory filenames.
• Fixed the tutorial script (plt.show(), matching CPU comparison, log-scale floor, converged-group freezing, local ty: ignore), reverted the global pyproject.toml ty override, added regression tests, and verified pytest, prek run --all-files, tutorial execution, and jupytext notebook execution locally.

context rot

[cubic-dev-ai]

0 issues found across 1 file (changes from recent commits).

Confidence score: 5/5

• Automated review surfaced no issues in the provided summaries.
• No files require special attention.

<sub>Shadow auto-approve: would require human review. New NEB workflow and optimizer-group support thread group_idx through SimState, FIRE, autobatchers, and runners, making this a broad public/core API change with design tradeoffs.
Tip: Get a Slack DM when cubic finishes reviewing your PRs. Connect Slack

View guided diff | Turn on auto-fix | Re-trigger cubic</sub>