NEB: add group-aware dynamics, path and NEB hooks

[ys-teh]

ALCHEMI Toolkit Pull Request

Description

Adds reusable path and nudged elastic band hooks, including group-aware dynamics and convergence behavior for batched NEB workflows.

This PR covers 4 commits starting from (“feat(dynamics): make dynamics group-aware”). It depends on the following PRs being merged first:

• NEB: kernels #3
• NEB: extend batch with group layout #4
• Fix/dynamics hooks and order NVIDIA/nvalchemi-toolkit#182

Intended usage:

Two-stage regular NEB → climbing-image NEB

from nvalchemi.dynamics import ConvergenceHook, FIRE2, FusedStage
from nvalchemi.dynamics.hooks import FreezeAtomsHook
from nvalchemi.dynamics.paths.hooks import PathEnergyStatsHook
from nvalchemi.dynamics.paths.neb.hooks import (
    ClimbingImageSelectionHook,
    NEBForceHook,
)

# Stage 0: regular NEB relaxation.
regular_stats = PathEnergyStatsHook()
regular_neb = FIRE2(
    model=model,
    dt=0.01,
    by_group=True,
    hooks=[
        regular_stats,
        NEBForceHook(
            energy_stats_hook=regular_stats,
            spring=0.1,
            method="improved_tangent",
            endpoint_mode="fixed",
        ),
        FreezeAtomsHook(mask_key="neb_fixed_node_mask"),
    ],
    convergence_hook=ConvergenceHook.from_fmax(threshold=0.5, by_group=True),
)

# Stage 1: climbing-image NEB refinement.
climbing_stats = PathEnergyStatsHook()
climbing_neb = FIRE2(
    model=model,
    dt=0.01,
    by_group=True,
    hooks=[
        climbing_stats,
        ClimbingImageSelectionHook(
            energy_stats_hook=climbing_stats,
            selection="fixed",
        ),
        NEBForceHook(
            energy_stats_hook=climbing_stats,
            spring=0.1,
            method="improved_tangent",
            endpoint_mode="fixed",
        ),
        FreezeAtomsHook(mask_key="neb_fixed_node_mask"),
    ],
    convergence_hook=ConvergenceHook.from_fmax(threshold=0.05, by_group=True),
)

# Each path advances from regular NEB to CI-NEB independently. Stage 1 is
# re-primed on entry because selecting a climbing image changes its forces.
neb = FusedStage(
    sub_stages=[
        (0, regular_neb),
        (1, climbing_neb),
    ],
    by_group=True,
    reprime_on_entry={1},
    n_steps=500,
    hooks=model.make_neighbor_hooks(),
)

relaxed_bands = neb.run(bands)

Type of Change

• Bug fix (non-breaking change that fixes an issue)
• New feature (non-breaking change that adds functionality)
• Breaking change (fix or feature that would cause existing functionality to not work as expected)
• Performance improvement
• Documentation update
• Refactoring (no functional changes)
• CI/CD or infrastructure change

Related Issues

N/A

Changes Made

• Add path energy statistics, climbing-image selection, and NEB force hooks.
• Make dynamics state updates and convergence checks group-aware for batched paths.

Testing

Targeted tests run:

• uv run --frozen --extra cu13 pytest test/dynamics/test_path_hooks.py test/dynamics/test_neb_hooks.py test/dynamics/test_base.py test/dynamics/test_state_management.py test/dynamics/test_fire2.py (passed: 191 tests)

• make lint (passed)

• Unit tests pass locally (make pytest)

• Linting passes (make lint)

• New tests added for new functionality meets coverage expectations?

Checklist

• I have read and understand the Contributing Guidelines
• I have updated the CHANGELOG.md
• I have performed a self-review of my code
• I have added docstrings to new functions/classes
• I have updated the documentation (if applicable)

Additional Notes

Tip

This repository uses Greptile, an AI code review service, to help conduct
pull request reviews. We encourage contributors to read and consider suggestions
made by Greptile, but note that human maintainers will provide the necessary
reviews for merging: Greptile's comments are not a qualitative judgement
of your code, nor is it an indication that the PR will be accepted/rejected.
We encourage the use of emoji reactions to Greptile comments, depending on
their usefulness and accuracy.