[Bug]: Unable to use BFGS optimizer in BatchRelaxer

[qchempku2017]

Contact Details

yxyhncd@gmail.com

Bug Description

When tying to use mattersim's BatchRelaxer with BFGS optimizer, an error popped up upon hitting relaxer.run():

---

AttributeError Traceback (most recent call last)
Cell In[15], line 2
1 print("Performing relaxations.")
----> 2 trajs = barlx.relax(all_atoms)
3 relaxed_atoms = [t[-1] for t in trajs.values()]
4 total_energies = np.array([a.info["total_energy"] for a in relaxed_atoms])

File ~/miniconda3/envs/mattersim/lib/python3.9/site-packages/mattersim/applications/batch_relax.py:168, in BatchRelaxer.relax(self, atoms_list)
166 self.tqdmcounter.update(1)
167 pointer += 1
--> 168 self.step_batch()
169 self.tqdmcounter.close()
171 return self.trajectories

File ~/miniconda3/envs/mattersim/lib/python3.9/site-packages/mattersim/applications/batch_relax.py:124, in BatchRelaxer.step_batch(self)
119 self.trajectories[opt.atoms.info["structure_index"]] = [
120 opt.atoms.copy()
121 ]
123 opt.step()
--> 124 if opt.converged() or opt.Nsteps >= self.max_n_steps:
125 self.is_active_instance[idx] = False
126 self.total_converged += 1

AttributeError: 'BFGS' object has no attribute 'Nsteps'

MatterSim Version

v1.1.3.dev1 (installed from git main branch)

Python Version

3.9.21

Reproduction Steps

1, Install mattersim with pip install git+https://github.com/microsoft/mattersim.git

2, Create a batch relaxer:

pot = Potential.from_checkpoint(load_path="./mattersim-v1.0.0-5M.pth", device=device, load_training_state=False)
barlx = BatchRelaxer(
potential=pot,
optimizer="BFGS",
filter="EXPCELLFILTER",
fmax=0.02,
max_natoms_per_batch=1024,
max_n_steps=5000,
)

3, trajs = barlx.relax(all_atoms).

Expected Behavior

BatchRelaxer should be able to optimize structures when using any ase supported optimizer.

Actual Behavior

BatchRelaxer only works with FIRE. This is because only FIRE has attribute self.Nsteps.

I recommend that instead of relying on optimizer attributes, we should add an internal attribute to BatchRelaxer to store number of steps elapsed on each optimizer instance, or simply count the length of each optimization trajectory.

Error Logs

Code of Conduct

• I agree to follow the project's Code of Conduct