Nudged elastic band implementation

[orionarcher]

Nudged elastic band would be an excellent addition to the project.

[stefanbringuier]

Curious if anyone has started to tackle this?

[orionarcher]

Not that I am aware of, it'd be a welcome addition!

[mstapelberg]

I started working on implementing NEB with a batched MACE model, but the FIRE optimizer seems a bit different from ASE's implementation. (since I imagine it's changed to work better for batches?)

A few things I noticed when comparing torch_sim's FIRE to ASE's FIRE:

1. The torch_sim FIRE implementation seems to calculate the power after the step has occurred. I think this could impact the setting of dt, alpha, and resetting of the velocity.

# torch_sim/optimizers.py ~line 630 (some lines omitted):
# Velocity Verlet first half step (v += 0.5*a*dt)
        atom_wise_dt = state.dt[state.batch].unsqueeze(-1)
        state.velocities += 0.5 * atom_wise_dt * state.forces / state.masses.unsqueeze(-1)

        # Update atomic positions
        atomic_positions_new = atomic_positions + atom_wise_dt * state.velocities

        # Update state with new positions and cell
        state.positions = atomic_positions_new

        # Get new forces, energy, and stress

        # Velocity Verlet first half step (v += 0.5*a*dt)
        state.velocities += 0.5 * atom_wise_dt * state.forces / state.masses.unsqueeze(-1)

        # Calculate power (F·V) for atoms
        atomic_power = (state.forces * state.velocities).sum(dim=1)  # [n_atoms]
       # add atomic_power to each batch

Where as in ase.optimize.fire.FIRE.step, the power is calculated before the updates are applied to the current step:

    # ase/optimize/fire.py (~line 168 onwards, simplified)
    f = optimizable.get_forces()
    # ... downhill check logic ...
    vf = np.vdot(f, self.v) # P calculated HERE, using self.v from previous step
    if vf > 0.0 and not is_uphill:
        # ... adapt based on vf > 0 ...
    else:
        # ... adapt based on vf <= 0 ...
    # ... THEN update velocity and position ...
    self.v += self.dt * f
    dr = self.dt * self.v
    # ... maxstep constraint ...
    r = optimizable.get_positions()
    optimizable.set_positions(r + dr)

2. In torch_sim, if the batch power is <0, the velocity gets reset for the batch, but it seems like whether the velocity is reset or not, the velocity and force directions get mixed.

for batch_idx in range(n_batches):
            # FIRE specific updates
            if batch_power[batch_idx] > 0:  # Power is positive
                state.n_pos[batch_idx] += 1
                if state.n_pos[batch_idx] > n_min:
                    state.dt[batch_idx] = min(state.dt[batch_idx] * f_inc, dt_max)
                    state.alpha[batch_idx] = state.alpha[batch_idx] * f_alpha
            else:  # Power is negative
                state.n_pos[batch_idx] = 0
                state.dt[batch_idx] = state.dt[batch_idx] * f_dec
                state.alpha[batch_idx] = alpha_start[batch_idx]
                # Reset velocities for both atoms and cell
                state.velocities[state.batch == batch_idx] = 0

        # Mix velocity and force direction using FIRE for atoms
        v_norm = torch.norm(state.velocities, dim=1, keepdim=True)
        f_norm = torch.norm(state.forces, dim=1, keepdim=True)
        # Avoid division by zero
        # mask = f_norm > 1e-10
        # state.velocity = torch.where(
        #     mask,
        #     (1.0 - state.alpha) * state.velocity
        #     + state.alpha * state.forces * v_norm / f_norm,
        #     state.velocity,
        # )
        atom_wise_alpha = state.alpha[state.batch].unsqueeze(-1)
        state.velocities = (
            1.0 - atom_wise_alpha
        ) * state.velocities + atom_wise_alpha * state.forces * v_norm / (f_norm + eps)

In ASE, if the power is <=0, all the velocities are reset to 0 (self.v[:] *= 0.0 at line 202, but the mixing is not applied afterwards. (note ASE calls power, vf)

            vf = np.vdot(f, self.v)
            if vf > 0.0 and not is_uphill:
                self.v = (1.0 - self.a) * self.v + self.a * f / np.sqrt(
                    np.vdot(f, f)) * np.sqrt(np.vdot(self.v, self.v))
                if self.Nsteps > self.Nmin:
                    self.dt = min(self.dt * self.finc, self.dtmax)
                    self.a *= self.fa
                self.Nsteps += 1
            else:
                self.v[:] *= 0.0
                self.a = self.astart
                self.dt *= self.fdec
                self.Nsteps = 0

I think by mixing regardless of the power, wouldn't you reintroduce velocities based on the current forces and allow for the optimizer to continue uphill? (I noticed this happening when running NEB using FIRE).

3. The torch_sim FIRE implementation also seems to be missing a max_step setting which could help prevent hopping over minima?

Despite these differences, the actual barrier height is 0.01 eV lower compared to my ASE calculated barrier when using FIRE, but the shape of the barrier is different compared to the ASE implementation. Are there performance/reasons to make batching better that informed your implementation of FIRE to be like this? (Never converged)

I switched over to using the GD optimizer provided in torch_sim and get this for my test case (converged in 563 steps):

I also then modified the torch_sim FIRE optimizer to mirror ASE's implementation and got this (converged in about 128 steps):

Additionally I checked my interpolations compared to the ASE interpolations (~0.003 Angstrom differences).

Would it be helpful to submit a pull request on this and figure out what's causing the instability/weird barrier shapes? I'll take another look at my NEB implementation, as that might also be the issue.

Thanks,
Myles

[stefanbringuier]

This looks promising, curious about your implementation details. I've been trying to get something to work following the papers of G. Henkelman and H. Jónsson, but although it runs, the transition states and energy barriers are completely wrong for simple cases so my implementation is off somewhere. Also don't have the batching tested/implemented which is the value.

[mstapelberg]

Thanks! @stefanbringuier , your code is great! Very helpful in finding the errors in my implementation. For context I mostly use NEB for calculating vacancy/interstitial migration energy barriers in alloys; the more complicated use-cases for NEB are above me!

To summarize:

1. I keep all the intermediate images as a single batched SimState, separate from the endpoints, your method seems to hold the entire path (and metadata like tangents) in a dedicated NEBState?
2. I use a single optimizer instance and perform one update step for the entire batch of intermediate images. Whereas your code creates separate optimizer states for each intermediate image, updating them individually in the main loop
3. Your tangent calculation is way more robust; I've implemented it now in my implementation. I was using an energy-difference based logic
4. Climbing Image Implementation: Yours explicitly zeroes the spring force for the climbing image, whereas mine inverts the true force component along the tangent but I don't explicitly zero the spring force
5. I only calculate the endpoint energies/forces initially, whereas your code seems to calculate them at each step?
6. I also think that my code would need some work to better utilize TrajectoryReporter and have more sophisticated convergence criterion

Here's a link to my feature branch; please excuse the mess: https://github.com/mstapelberg/torch-sim/tree/feature/neb-workflow

After implementing your tangent calculation and the better CI logic, I get:

I then realized that I was using VASP Spring constants (i.e SPRING=-5), I saw that ASE uses 0.1 (in eV/Angstrom) by default, I set it to that and re-ran with FIRE:

And with GD (learning rate = 0.01):

I played with adding more images as well, that seemed to perhaps help get the barrier magnitude better, but the actual trajectory with FIRE seems to be off. (GD in Torch_Sim seems to approach the same trajectory as FIRE in ASE).

Looks like with GD this is closer to working? FIRE gets the right barrier, but wrong shape... What changes/suggestions would you recommend?

Thanks,
Myles

P.S: Pretty incredible to run 128 atom NEB simulation with 5 images in a minute on a laptop. Can't imagine how fast this would be with a proper GPU...

[janosh]

@mstapelberg thanks a lot for diagnosing the torch-sim FIRE issue. your proposed fix looks very promising! could you open a PR with those changes and set #146 as the target branch? i'll be happy to review and add tests.

Are there performance/reasons to make batching better that informed your implementation of FIRE to be like this? (Never converged)

no, i think calculating atomic_power after the VV step has nothing to do with batching and was simply done by mistake. @abhijeetgangan correct me if i'm wrong

[orionarcher]

As @janosh said, thank you for the implementations. @stefanbringuier would you mind opening a draft PR with your NEB implementation too? Implementation looks great but I have a couple thoughts on how the batching could work.

[abhijeetgangan]

Thanks for looking into it. The FIRE implementation follows that of the original paper. I will have a batched version that follows the ASE implementation in an upcoming PR. There is an unbatched version available, which follows the ASE implementation. NEB implementation would be great!

[mstapelberg]

@mstapelberg thanks a lot for diagnosing the torch-sim FIRE issue. your proposed fix looks very promising! could you open a PR with those changes and set #146 as the target branch? i'll be happy to review and add tests.

Are there performance/reasons to make batching better that informed your implementation of FIRE to be like this? (Never converged)

no, i think calculating atomic_power after the VV step has nothing to do with batching and was simply done by mistake. @abhijeetgangan correct me if i'm wrong

You bet! @janosh Just to clarify, should I include my changes to the FIRE method in torch-sim/optimizer.py as well? or just the neb.py in torch-sim/workflows?

Thanks,
Myles

[janosh]

yes, i was referring only to the FIRE changes actually. re NEB, maybe coordinate with @stefanbringuier. sounds to me like there's a need to figure out what are the best parts from each of your implementations. either way, that should be a separate PR

[mstapelberg]

yes, i was referring only to the FIRE changes actually. re NEB, maybe coordinate with @stefanbringuier. sounds to me like there's a need to figure out what are the best parts from each of your implementations. either way, that should be a separate PR

Noted, I'll make a PR then for just the FIRE changes. Thanks!

[stefanbringuier]

@mstapelberg, thanks for sharing! You raised some good question for me to review:

1. I keep all the intermediate images as a single batched SimState, separate from the endpoints, your method seems to hold the entire path (and metadata like tangents) in a dedicated NEBState?

Correct, was doing for complete traceability/diagnostics.

2. I use a single optimizer instance and perform one update step for the entire batch of intermediate images. Whereas your code creates separate optimizer states for each intermediate image, updating them individually in the main loop

Hmm, I think your single-optimizer design is cleaner/easier from a batch perspective but one question I have is does it risk "coupling" if the min energy path is highly curved or contains intermediate minima?

4. Climbing Image Implementation: Yours explicitly zeroes the spring force for the climbing image, whereas mine inverts the true force component along the tangent but I don't explicitly zero the spring force

From my understand of the papers, zeroing spring force in addition to inverting is necessary to decouple the climbing image from the image distribution constraints, otherwise it distorts climbing image trajectory causing issues for the saddle convergence? Maybe I'm wrong or its not critical. Curious on your take?

5. I only calculate the endpoint energies/forces initially, whereas your code seems to calculate them at each step?

Yeah, I this should be changed, just doing it for tracking.

I'd like to fully review your code and see where I can make some adjustments. As @janosh and @orionarcher suggested, let's coordinate a PR. If your open to it ,my thinking is:

1. One of us opens a PR for the NEB feature branch.
2. The PR owner gives the other person write access to the fork/branch

This way its a clean single upstream PR. Just need to coordinate and avoid force-pushing.