How to fix atoms when running ForceFieldRelaxMaker

Atomate
1
  • After running a Structure optimization job using the MACE model, the atoms with the fixed_indices label also moved. I have tested with the pymatgen structure and ASE FixAtoms, neither worked.
  • In contrast, the Structure optimization job via ASE calculators and the MACE model succeeded. The fixed atoms (through FixAtoms) kept fixed well.
  • Here are the key codes I used.

Using atomate2

Scheme A - selective_dynamics of pymatgen structure

from atomate2.forcefields.jobs import ForceFieldRelaxMaker
from pymatgen.core import Structure
import numpy as np

structure = Structure.from_file("example.cif")
...
structure.add_site_property("selective_dynamics", selective_dynamics)

job = ForceFieldRelaxMaker(
    force_field_name="MACE_MP_0B3",
    calculator_kwargs={"default_dtype": "float64"},
    relax_cell=False,
    fix_symmetry=True,
    relax_kwargs={"fmax": 0.03},
    optimizer_kwargs={"optimizer": "BFGSLineSearch"},
    steps=500,
).make(structure)

Scheme B - ASE FixAtoms

from atomate2.forcefields.jobs import ForceFieldRelaxMaker
from pymatgen.core import Structure
import numpy as np
from ase.constraints import FixAtoms

structure = Structure.from_file("example.cif")
...
atoms = structure.to_ase_atoms()
atoms.set_constraint(FixAtoms(indices=fixed_indices))

job = ForceFieldRelaxMaker(
    force_field_name="MACE_MP_0B3",
    calculator_kwargs={"default_dtype": "float64"},
    relax_cell=False,
    fix_symmetry=True,
    relax_kwargs={"fmax": 0.03},
    optimizer_kwargs={"optimizer": "BFGSLineSearch"},
    steps=500,
).make(structure)

Using ASE

from pymatgen.core import Structure
from mace.calculators import mace_mp
import numpy as np
from ase.optimize import BFGSLineSearch

structure = Structure.from_file("example.cif")
...
atoms = structure.to_ase_atoms()
atoms.set_constraint(FixAtoms(indices=fixed_indices))

calc = mace_mp(model="medium", default_dtype="float64", device='cuda')
atoms.calc = calc
dyn = BFGSLineSearch(atoms)
dyn.run(fmax=0.03)

To my knowledge, the atomate2 reuses the ASE calculators. I wonder how I should use Atomate2 to run the jobs with fixed atoms. Thanks for any reply!

2

Hey @ramonmi, a couple of things:

  1. The three calculations you sent all do different things, so the results aren’t going to be directly comparable.

    • In your example B, the constraints are set only on atoms and not on the structure itself, so no constraints are applied during the calculation

    • In example C, you switched to a different MACE model (MP-0 rather than MP-0b3) and no longer relax the cell. Take a look at the FrechetCellFilter and how it is used in atomate2 when the cell is relaxed

  2. The underlying issue seems to be that constraints are only approximately satisfied in ASE when the cell itself is not relaxed

Let's set the jobs up: 
from atomate2.forcefields.jobs import ForceFieldRelaxMaker
from jobflow import run_locally, JobStore
from maggma.stores import MemoryStore
from pymatgen.core import Structure

s = Structure.from_str("""Si2
1.0
   4.3335729972004917    0.0000000000000000    1.9246389981090721
   1.1111909992801432    3.1429239987499362    1.9246389992542632
   0.0000000000000000    0.0000000000000000    3.8492780000000000
Si
2
direct
   0.8750000000000000    0.8750000000000000    0.8750000000000000 Si
   0.1250000000000000    0.1    0.15 Si
""", fmt="poscar")
s = s.scale_lattice(1.1*s.volume)
s.add_site_property("selective_dynamics",[[True,True,True],[False,False,True]])

relax_cell_job = ForceFieldRelaxMaker(
    force_field_name = "MACE-MP-0b3",
    optimizer_kwargs = {"optimizer": "LBFGSLineSearch",},
).make(s.copy())
resp = run_locally(relax_cell_job, store = JobStore(MemoryStore(),additional_stores={"data": MemoryStore()}))
output_relax_cell = r[relax_cell_job.uuid][1].output

fix_cell_job = ForceFieldRelaxMaker(
    force_field_name = "MACE-MP-0b3",
    optimizer_kwargs = {"optimizer": "LBFGSLineSearch",},
    relax_cell = False,
).make(s.copy())
resp = run_locally(fix_cell_job, store = JobStore(MemoryStore(),additional_stores={"data": MemoryStore()}))
output_fix_cell = r[fix_cell_job.uuid][1].output

And now look at the output:

for i in range(len(s)):
    print(
        f"Site index {i}: " + ", ".join(
            f"{typ} : {calc.structure[i].frac_coords}"
            for typ, calc in {
                "input": output_relax_cell.input,
                "fixed cell": output_fix_cell.output,
                "relaxed cell": output_relax_cell.output
            }.items()
        )
    )
>>> Site index 0: input : [0.875 0.875 0.875], fixed cell : [0.87442788 0.84934953 0.89422488], relaxed cell : [0.87520869 0.85082895 0.89330742]
>>> Site index 1: input : [0.125 0.1   0.15 ], fixed cell : [0.125      0.1        0.14388642], relaxed cell : [0.12497618 0.099979   0.14370873]

We see that the first site always relaxes coordinates, while the x and y coordinates of the second site are held fixed exactly when the cell isn’t relaxed, and only approximately when the cell is relaxed.

3

Thanks for your reply!

Actually, I have added relax_cell=False, in the ForceFieldRelaxMaker, but fix_symmetry=True, will break it. After I removed this line, both schemes A and B worked well.

Fixed typos of scheme B in the original post

job = ForceFieldRelaxMaker(
    force_field_name="MACE_MP_0B3",
    calculator_kwargs={"default_dtype": "float64"},
    relax_cell=False,
    fix_symmetry=True,
    relax_kwargs={"fmax": 0.03},
    optimizer_kwargs={"optimizer": "BFGSLineSearch"},
    steps=500,
-).make(structure)
+).make(atoms)