Implement lbfgs and bfgs

[abhijeetgangan]

Summary

• Add L-BFGS and BFGS optimization schemes.

Checklist

Before a pull request can be merged, the following items must be checked:

• Doc strings have been added in the Google docstring format.
• Run ruff on your code.
• Tests have been added for any new functionality or bug fixes.

[abhijeetgangan]

See detailed comment here.

[orionarcher]

Going to leave the physics review to someone else, just one API comment

[curtischong]

Just learned the algorithms today so I may have missed things. Great work, especially because the ase implementation isn't very documented.

[curtischong]

consider saying: "scale down step so atoms move at most max_step" or "# Scale step if it exceeds max_step" like you mention in bfgs

[abhijeetgangan]

Thanks for the comments. I will update those over the weekend.

[orionarcher]

@abhijeetgangan is this still on the docket for you?

I know it's a much requested feature so it'd be awesome to get it in. Thanks for your time on it.

[abhijeetgangan]

Will prioritize this on the weekend.

[abhijeetgangan]

Hi @sihoonchoi and @Andrew-S-Rosen the integration is pretty much done. I also tested on 1K structures on WBM with mace comparing FIRE and L-BFGS (with cell filters):

For fmax of 0.02 eV/A

Optimizer    Time (s)     Converged       Max Force (eV/Å)
FIRE         288.55       1000/1000          0.0200         
L-BFGS       101.91       1000/1000          0.0196  

For fmax of 1e-4 eV/A

Optimizer    Time (s)     Converged       Max Force (eV/Å)
FIRE         744.09       995/1000          0.0145         
L-BFGS       192.27       1000/1000          0.0001 

Clearly, here it seem that FIRE can struggle sometimes to relax especially for small fmax. It did not finish all even after 500 steps. BFGS is working but need linesearch to the competitive. That's going to be a separate PR. I am going to merge this as the tests are very comprehensive but would be nice if someone can try it out.

import time

import torch
from mace.calculators.foundations_models import mace_mp
from matbench_discovery.data import DataFiles, ase_atoms_from_zip

import torch_sim as ts
from torch_sim.models.mace import MaceModel


N_STRUCTURES = 1000
MAX_STEPS = 500
FORCE_TOL = 1e-4

device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
dtype = torch.float64
raw_mace = mace_mp(model="small", return_raw_model=True, default_dtype="float64")
mace_model = MaceModel(model=raw_mace, device=device, dtype=dtype, compute_stress=True)

ase_atoms_list = ase_atoms_from_zip(DataFiles.wbm_initial_atoms.path, limit=N_STRUCTURES)
initial_state = ts.io.atoms_to_state(ase_atoms_list, device=device, dtype=dtype)

convergence_fn = ts.generate_force_convergence_fn(
    force_tol=FORCE_TOL, include_cell_forces=True
)

optimizers = [
    ("FIRE", ts.Optimizer.fire, {}),
    ("L-BFGS", ts.Optimizer.lbfgs, {"max_history": 20, "step_size": 0.5}),
]

results = {}

for name, optimizer, init_kwargs in optimizers:
    print(f"Running {name}...")

    state = initial_state.clone()

    autobatcher = ts.autobatching.InFlightAutoBatcher(
        model=mace_model,
        memory_scales_with="n_atoms",
        max_memory_scaler=600.0,
        max_iterations=MAX_STEPS // 2,
    )

    if device.type == "cuda":
        torch.cuda.synchronize()
    start_time = time.perf_counter()

    final_state = ts.optimize(
        system=state,
        model=mace_model,
        optimizer=optimizer,
        max_steps=MAX_STEPS,
        convergence_fn=convergence_fn,
        init_kwargs=dict(cell_filter=ts.CellFilter.frechet, **init_kwargs),
        autobatcher=autobatcher,
    )

    if device.type == "cuda":
        torch.cuda.synchronize()
    elapsed = time.perf_counter() - start_time

    max_forces = torch.norm(final_state.forces, dim=-1).max().item()
    converged = convergence_fn(final_state, None)
    n_converged = int(converged.sum().item())

    results[name] = {
        "time": elapsed,
        "max_force": max_forces,
        "n_converged": n_converged,
    }

    n_total = len(ase_atoms_list)
    print(
        f"Time: {elapsed:.2f}s | Converged: {n_converged}/{n_total} | "
        f"Max force: {max_forces:.4f} eV/Å"
    )

print("SUMMARY")
print(f"{'Optimizer':<12} {'Time (s)':<12} {'Converged':<15} {'Max Force (eV/Å)':<15}")

for name, data in results.items():
    print(
        f"{name:<12} {data['time']:<12.2f} "
        f"{data['n_converged']}/{len(ase_atoms_list):<13} {data['max_force']:<15.4f}"
    )

[sihoonchoi]

@abhijeetgangan I've tried L-BFGS and BFGS along with FIRE with UMA-OMC on 100 structures and it works well. This is great - thanks for your work!

Optimizer    Time (s)     Converged    
FIRE         838.81       100/100          
L-BFGS       530.74       100/100          
BFGS         545.62       100/100          

Below is a comparison with serial relaxations in ASE.

Optimizer    Time (s)     Converged    
FIRE         1670.04     
L-BFGS       1001.56     
BFGS         1133.86     

[abhijeetgangan]

@sihoonchoi Thanks for looking into it. This is very valuable!

There are more improvements that can be made on top of this like introducing line-search and trust regions. Some of which is batching friendly and some of it is not. I am pretty sure the optimization algorithms can be made more robust with it (ASE doesn't use many of these known tricks). I would be curious on what the community thinks but given that there are no benchmarks in matsci it's hard to make comparisons. Maybe guys at Rowan can set this up? Would like to know what others think.

[Andrew-S-Rosen]

I personally find it a bit wild that we don't have good benchmarks on this, but I guess it's historically because DFT is expensive to run. Rowan doesn't have a ton of infra for materials but they have obviously thought about this.

Actually, my group has a funded NSF collaboration with the Argonne folks supporting TorchSim (Ben Blaiszik et al) that might be relevant here. That might be a good avenue to kickstart this. I'd be interested in seeing if we can start something through that and open it up to the community to participate.

[abhijeetgangan]

One thing I always found confusing was when comparing MLIPs vs DFT for relaxation. MLIPs would use FIRE but VASP uses CG and QE uses BFGS both with linesearch and trust regions which makes them very different. Another aspect is the cost. I am sure the devs of DFT codes must have chosen this so that it reduces the DFT calls during relaxation. I think a fair comparison would be ideal but that means going away from exact ASE behavior.

Actually, my group has a funded NSF collaboration with the Argonne folks supporting TorchSim (Ben Blaizsik et al) that might be relevant here. That might be a good avenue to kickstart this. I'd be interested in seeing if we can start something through that and open it up to the community to participate.

I think this is a great idea! Maybe we can start an email thread to discuss this? I can add more optimizers to test.

[Andrew-S-Rosen]

One thing I always found confusing was when comparing MLIPs vs DFT for relaxation. MLIPs would use FIRE but VASP uses CG and QE uses BFGS both with linesearch and trust regions which makes them very different. Another aspect is the cost. I am sure the devs of DFT codes must have chosen this so that it reduces the DFT calls during relaxation. I think a fair comparison would be ideal but that means going away from exact ASE behavior.

I think the reason people have been using MLIPs with FIRE is simply a cultural thing, not based on much other than the original FIRE paper. FIRE isn't in VASP (without external plugins), but yeah, CG is simple enough and works.

I think this is a great idea! Maybe we can start an email thread to discuss this? I can add more optimizers to test.

Sure, happy to kick off an email tomorrow. Can you send me an email from your account? I'm not sure what your email is.

[abhijeetgangan]

Sure, happy to kick off an email tomorrow. Can you send me an email from your account? I'm not sure what your email is.

Sounds good! abhijeetgangan@g.ucla.edu