diff --git a/src/atomate2/forcefields/jobs.py b/src/atomate2/forcefields/jobs.py index a1d0f4b0ca..67f9425879 100644 --- a/src/atomate2/forcefields/jobs.py +++ b/src/atomate2/forcefields/jobs.py @@ -81,6 +81,11 @@ class ForceFieldRelaxMaker(Maker): The name of the force field. relax_cell : bool = True Whether to allow the cell shape/volume to change during relaxation. + fix_symmetry : bool = False + Whether to fix the symmetry during relaxation. + Refines the symmetry of the initial structure. + symprec : float = 1e-2 + Tolerance for symmetry finding in case of fix_symmetry. steps : int Maximum number of ionic steps allowed during relaxation. relax_kwargs : dict @@ -96,6 +101,8 @@ class ForceFieldRelaxMaker(Maker): name: str = "Force field relax" force_field_name: str = f"{MLFF.Forcefield}" relax_cell: bool = True + fix_symmetry: bool = False + symprec: float = 1e-2 steps: int = 500 relax_kwargs: dict = field(default_factory=dict) optimizer_kwargs: dict = field(default_factory=dict) @@ -126,7 +133,11 @@ def make( with revert_default_dtype(): relaxer = Relaxer( - self._calculator(), relax_cell=self.relax_cell, **self.optimizer_kwargs + self._calculator(), + relax_cell=self.relax_cell, + fix_symmetry=self.fix_symmetry, + symprec=self.symprec, + **self.optimizer_kwargs, ) result = relaxer.relax(structure, steps=self.steps, **self.relax_kwargs) @@ -137,6 +148,8 @@ def make( self.steps, self.relax_kwargs, self.optimizer_kwargs, + self.fix_symmetry, + self.symprec, **self.task_document_kwargs, ) @@ -185,6 +198,11 @@ class CHGNetRelaxMaker(ForceFieldRelaxMaker): The name of the force field. relax_cell : bool = True Whether to allow the cell shape/volume to change during relaxation. + fix_symmetry : bool = False + Whether to fix the symmetry during relaxation. + Refines the symmetry of the initial structure. + symprec : float = 1e-2 + Tolerance for symmetry finding in case of fix_symmetry. steps : int Maximum number of ionic steps allowed during relaxation. relax_kwargs : dict @@ -198,6 +216,8 @@ class CHGNetRelaxMaker(ForceFieldRelaxMaker): name: str = f"{MLFF.CHGNet} relax" force_field_name: str = f"{MLFF.CHGNet}" relax_cell: bool = True + fix_symmetry: bool = False + symprec: float = 1e-2 steps: int = 500 relax_kwargs: dict = field(default_factory=dict) optimizer_kwargs: dict = field(default_factory=dict) @@ -241,6 +261,11 @@ class M3GNetRelaxMaker(ForceFieldRelaxMaker): The name of the force field. relax_cell : bool = True Whether to allow the cell shape/volume to change during relaxation. + fix_symmetry : bool = False + Whether to fix the symmetry during relaxation. + Refines the symmetry of the initial structure. + symprec : float = 1e-2 + Tolerance for symmetry finding in case of fix_symmetry. steps : int Maximum number of ionic steps allowed during relaxation. relax_kwargs : dict @@ -254,6 +279,8 @@ class M3GNetRelaxMaker(ForceFieldRelaxMaker): name: str = f"{MLFF.M3GNet} relax" force_field_name: str = f"{MLFF.M3GNet}" relax_cell: bool = True + fix_symmetry: bool = False + symprec: float = 1e-2 steps: int = 500 relax_kwargs: dict = field(default_factory=dict) optimizer_kwargs: dict = field(default_factory=dict) @@ -276,6 +303,11 @@ class NequipRelaxMaker(ForceFieldRelaxMaker): The name of the force field. relax_cell : bool = True Whether to allow the cell shape/volume to change during relaxation. + fix_symmetry : bool = False + Whether to fix the symmetry during relaxation. + Refines the symmetry of the initial structure. + symprec : float = 1e-2 + Tolerance for symmetry finding in case of fix_symmetry. steps : int Maximum number of ionic steps allowed during relaxation. relax_kwargs : dict @@ -289,6 +321,8 @@ class NequipRelaxMaker(ForceFieldRelaxMaker): name: str = f"{MLFF.Nequip} relax" force_field_name: str = f"{MLFF.Nequip}" relax_cell: bool = True + fix_symmetry: bool = False + symprec: float = 1e-2 steps: int = 500 relax_kwargs: dict = field(default_factory=dict) optimizer_kwargs: dict = field(default_factory=dict) @@ -351,6 +385,11 @@ class MACERelaxMaker(ForceFieldRelaxMaker): The name of the force field. relax_cell : bool = True Whether to allow the cell shape/volume to change during relaxation. + fix_symmetry : bool = False + Whether to fix the symmetry during relaxation. + Refines the symmetry of the initial structure. + symprec : float = 1e-2 + Tolerance for symmetry finding in case of fix_symmetry. steps : int Maximum number of ionic steps allowed during relaxation. relax_kwargs : dict @@ -372,6 +411,8 @@ class MACERelaxMaker(ForceFieldRelaxMaker): name: str = f"{MLFF.MACE} relax" force_field_name: str = f"{MLFF.MACE}" relax_cell: bool = True + fix_symmetry: bool = False + symprec: float = 1e-2 steps: int = 500 relax_kwargs: dict = field(default_factory=dict) optimizer_kwargs: dict = field(default_factory=dict) @@ -419,6 +460,11 @@ class GAPRelaxMaker(ForceFieldRelaxMaker): The name of the force field. relax_cell : bool = True Whether to allow the cell shape/volume to change during relaxation. + fix_symmetry : bool = False + Whether to fix the symmetry during relaxation. + Refines the symmetry of the initial structure. + symprec : float = 1e-2 + Tolerance for symmetry finding in case of fix_symmetry. steps : int Maximum number of ionic steps allowed during relaxation. relax_kwargs : dict @@ -432,6 +478,8 @@ class GAPRelaxMaker(ForceFieldRelaxMaker): name: str = f"{MLFF.GAP} relax" force_field_name: str = f"{MLFF.GAP}" relax_cell: bool = True + fix_symmetry: bool = False + symprec: float = 1e-2 steps: int = 500 relax_kwargs: dict = field(default_factory=dict) optimizer_kwargs: dict = field(default_factory=dict) diff --git a/src/atomate2/forcefields/schemas.py b/src/atomate2/forcefields/schemas.py index ac11bf230d..4c9f3cbfbc 100644 --- a/src/atomate2/forcefields/schemas.py +++ b/src/atomate2/forcefields/schemas.py @@ -55,6 +55,16 @@ class InputDoc(BaseModel): None, description="Whether cell lattice was allowed to change during relaxation.", ) + fix_symmetry: bool = Field( + None, + description=( + "Whether to fix the symmetry of the structure during relaxation. " + "Refines the symmetry of the initial structure." + ), + ) + symprec: float = Field( + None, description="Tolerance for symmetry finding in case of fix_symmetry." + ) steps: int = Field( None, description="Maximum number of steps allowed during relaxation." ) @@ -151,6 +161,8 @@ def from_ase_compatible_result( steps: int, relax_kwargs: dict = None, optimizer_kwargs: dict = None, + fix_symmetry: bool = False, + symprec: float = 1e-2, ionic_step_data: tuple = ("energy", "forces", "magmoms", "stress", "structure"), store_trajectory: StoreTrajectoryOption = StoreTrajectoryOption.NO, **task_document_kwargs, @@ -166,6 +178,10 @@ def from_ase_compatible_result( The outputted results from the task. relax_cell : bool Whether the cell shape/volume was allowed to change during the task. + fix_symmetry : bool + Whether to fix the symmetry of the structure during relaxation. + symprec : float + Tolerance for symmetry finding in case of fix_symmetry. steps : int Maximum number of ionic steps allowed during relaxation. relax_kwargs : dict @@ -199,6 +215,8 @@ def from_ase_compatible_result( input_doc = InputDoc( structure=input_structure, relax_cell=relax_cell, + fix_symmetry=fix_symmetry, + symprec=symprec, steps=steps, relax_kwargs=relax_kwargs, optimizer_kwargs=optimizer_kwargs, diff --git a/src/atomate2/forcefields/utils.py b/src/atomate2/forcefields/utils.py index 4f6d1720fc..e80729a507 100644 --- a/src/atomate2/forcefields/utils.py +++ b/src/atomate2/forcefields/utils.py @@ -20,6 +20,7 @@ import numpy as np from ase.calculators.calculator import PropertyNotImplementedError from ase.calculators.singlepoint import SinglePointCalculator +from ase.constraints import FixSymmetry from ase.io import Trajectory as AseTrajectory from ase.optimize import BFGS, FIRE, LBFGS, BFGSLineSearch, LBFGSLineSearch, MDMin from ase.optimize.sciopt import SciPyFminBFGS, SciPyFminCG @@ -56,7 +57,6 @@ from ase.filters import Filter from ase.optimize.optimize import Optimizer - OPTIMIZERS = { "FIRE": FIRE, "BFGS": BFGS, @@ -290,6 +290,8 @@ def __init__( calculator: Calculator, optimizer: Optimizer | str = "FIRE", relax_cell: bool = True, + fix_symmetry: bool = False, + symprec: float = 1e-2, ) -> None: """ Initialize the Relaxer. @@ -299,6 +301,8 @@ def __init__( calculator (ase Calculator): an ase calculator optimizer (str or ase Optimizer): the optimization algorithm. relax_cell (bool): if True, cell parameters will be optimized. + fix_symmetry (bool): if True, symmetry will be fixed during relaxation. + symprec (float): Tolerance for symmetry finding in case of fix_symmetry. """ self.calculator = calculator @@ -312,6 +316,8 @@ def __init__( self.opt_class: Optimizer = optimizer_obj self.relax_cell = relax_cell self.ase_adaptor = AseAtomsAdaptor() + self.fix_symmetry = fix_symmetry + self.symprec = symprec def relax( self, @@ -350,6 +356,8 @@ def relax( """ if isinstance(atoms, (Structure, Molecule)): atoms = self.ase_adaptor.get_atoms(atoms) + if self.fix_symmetry: + atoms.set_constraint(FixSymmetry(atoms, symprec=self.symprec)) atoms.set_calculator(self.calculator) stream = sys.stdout if verbose else io.StringIO() with contextlib.redirect_stdout(stream): diff --git a/tests/conftest.py b/tests/conftest.py index b9d84e6565..66115a3a27 100644 --- a/tests/conftest.py +++ b/tests/conftest.py @@ -102,6 +102,11 @@ def sr_ti_o3_structure(test_dir): return Structure.from_file(test_dir / "structures" / "SrTiO3.cif") +@pytest.fixture() +def ba_ti_o3_structure(test_dir): + return Structure.from_file(test_dir / "structures" / "BaTiO3.cif") + + @pytest.fixture(autouse=True) def mock_jobflow_settings(memory_jobstore): """Mock the jobflow settings to use our specific jobstore (with data store).""" diff --git a/tests/forcefields/test_jobs.py b/tests/forcefields/test_jobs.py index 76ebd4c401..310a9132ca 100644 --- a/tests/forcefields/test_jobs.py +++ b/tests/forcefields/test_jobs.py @@ -4,6 +4,7 @@ import pytest from jobflow import run_locally from pymatgen.core import Structure +from pymatgen.symmetry.analyzer import SpacegroupAnalyzer from pytest import approx, importorskip from atomate2.forcefields.jobs import ( @@ -40,6 +41,37 @@ def test_chgnet_static_maker(si_structure): assert output1.forcefield_version == get_imported_version("chgnet") +@pytest.mark.parametrize( + "fix_symmetry, symprec", [(True, 1e-2), (False, 1e-2), (True, 1e-1)] +) +def test_chgnet_relax_maker_fix_symmetry( + ba_ti_o3_structure: Structure, + fix_symmetry: bool, + symprec: float, +): + # translate one atom to break symmetry but stay below symprec threshold + ba_ti_o3_structure.translate_sites(1, [symprec / 10.0, 0, 0]) + job = CHGNetRelaxMaker( + relax_kwargs={"fmax": 0.01}, + fix_symmetry=fix_symmetry, + symprec=symprec, + ).make(ba_ti_o3_structure) + # get space group number of input structure + initial_space_group = SpacegroupAnalyzer( + ba_ti_o3_structure, symprec=symprec + ).get_space_group_number() + responses = run_locally(job, ensure_success=True) + output1 = responses[job.uuid][1].output + assert output1.is_force_converged + final_space_group = SpacegroupAnalyzer( + output1.output.structure, symprec=symprec + ).get_space_group_number() + if fix_symmetry: + assert initial_space_group == final_space_group + else: + assert initial_space_group != final_space_group + + @pytest.mark.parametrize("relax_cell", [True, False]) def test_chgnet_relax_maker(si_structure: Structure, relax_cell: bool): # translate one atom to ensure a small number of relaxation steps are taken @@ -138,14 +170,60 @@ def test_mace_static_maker(si_structure: Structure, test_dir: Path, model): assert output1.forcefield_version == get_imported_version("mace-torch") +@pytest.mark.parametrize( + "fix_symmetry, symprec", [(True, 1e-2), (False, 1e-2), (True, 1e-1)] +) +def test_mace_relax_maker_fix_symmetry( + ba_ti_o3_structure: Structure, + fix_symmetry: bool, + symprec: float, +): + # translate one atom to break symmetry but stay below symprec threshold + ba_ti_o3_structure.translate_sites(1, [symprec / 10.0, 0, 0]) + job = MACERelaxMaker( + relax_kwargs={"fmax": 0.02}, + fix_symmetry=fix_symmetry, + symprec=symprec, + ).make(ba_ti_o3_structure) + # get space group number of input structure + initial_space_group = SpacegroupAnalyzer( + ba_ti_o3_structure, symprec=symprec + ).get_space_group_number() + responses = run_locally(job, ensure_success=True) + output1 = responses[job.uuid][1].output + assert output1.is_force_converged + final_space_group = SpacegroupAnalyzer( + output1.output.structure, symprec=symprec + ).get_space_group_number() + if fix_symmetry: + assert initial_space_group == final_space_group + else: + assert initial_space_group != final_space_group + + +@pytest.mark.parametrize( + "fix_symmetry, symprec", [(True, 1e-2), (False, 1e-2), (True, 1e-1)] +) @pytest.mark.parametrize("relax_cell", [True, False]) @mace_paths def test_mace_relax_maker( - si_structure: Structure, test_dir: Path, model, relax_cell: bool + si_structure: Structure, + model, + relax_cell: bool, + fix_symmetry: bool, + symprec: float, ): + from ase.spacegroup.symmetrize import check_symmetry, is_subgroup + + _, init_spg_num = si_structure.get_space_group_info() + assert init_spg_num == 227 + # translate one atom to ensure a small number of relaxation steps are taken si_structure.translate_sites(0, [0, 0, 0.1]) + _, init_spg_num = si_structure.get_space_group_info() + assert init_spg_num == 12 + # generate job # NOTE the test model is not trained on Si, so the energy is not accurate job = MACERelaxMaker( @@ -153,6 +231,9 @@ def test_mace_relax_maker( steps=25, optimizer_kwargs={"optimizer": "BFGSLineSearch"}, relax_cell=relax_cell, + fix_symmetry=fix_symmetry, + symprec=symprec, + relax_kwargs={"fmax": 0.04}, ).make(si_structure) # run the flow or job and ensure that it finished running successfully @@ -162,12 +243,30 @@ def test_mace_relax_maker( output1 = responses[job.uuid][1].output assert output1.is_force_converged assert isinstance(output1, ForceFieldTaskDocument) + + si_atoms = si_structure.to_ase_atoms() + symmetry_ops_init = check_symmetry(si_atoms, symprec=1.0e-3) + si_atoms_final = output1.output.structure.to_ase_atoms() + symmetry_ops_final = check_symmetry(si_atoms_final, symprec=1.0e-3) + + # get space group number of input structure + _, final_spg_num = output1.output.structure.get_space_group_info() if relax_cell: - assert output1.output.energy == approx(-0.0526856, rel=1e-1) + assert final_spg_num == 12 + else: + assert final_spg_num == 74 + + if fix_symmetry: # if symmetry is fixed, the symmetry should be the same or higher + assert is_subgroup(symmetry_ops_init, symmetry_ops_final) + else: # if symmetry is not fixed, it can both increase or decrease + assert not is_subgroup(symmetry_ops_init, symmetry_ops_final) + + if relax_cell: + assert output1.output.energy == approx(-0.071117445, rel=1e-1) assert output1.output.n_steps >= 4 else: - assert output1.output.energy == approx(-0.051912, rel=1e-4) - assert output1.output.n_steps == 4 + assert output1.output.energy == approx(-0.06772976, rel=1e-4) + assert output1.output.n_steps == 7 def test_gap_static_maker(si_structure: Structure, test_dir): @@ -176,7 +275,7 @@ def test_gap_static_maker(si_structure: Structure, test_dir): task_doc_kwargs = {"ionic_step_data": ("structure", "energy")} # generate job - # Test files have been provided by Yuanbin Liu (University of Oxford) + # Test files have been provided by @YuanbinLiu (University of Oxford) job = GAPStaticMaker( calculator_kwargs={ "args_str": "IP GAP", @@ -220,9 +319,15 @@ def test_nequip_static_maker(sr_ti_o3_structure: Structure, test_dir: Path): assert output1.forcefield_version == get_imported_version("nequip") -@pytest.mark.parametrize("relax_cell", [True, False]) +@pytest.mark.parametrize( + ("relax_cell", "fix_symmetry"), + [(True, False), (False, True)], +) def test_nequip_relax_maker( - sr_ti_o3_structure: Structure, test_dir: Path, relax_cell: bool + sr_ti_o3_structure: Structure, + test_dir: Path, + relax_cell: bool, + fix_symmetry: bool, ): importorskip("nequip") # translate one atom to ensure a small number of relaxation steps are taken @@ -232,6 +337,7 @@ def test_nequip_relax_maker( steps=25, optimizer_kwargs={"optimizer": "BFGSLineSearch"}, relax_cell=relax_cell, + fix_symmetry=fix_symmetry, calculator_kwargs={ "model_path": test_dir / "forcefields" / "nequip" / "nequip_ff_sr_ti_o3.pth" }, @@ -250,6 +356,11 @@ def test_nequip_relax_maker( assert output1.output.energy == approx(-44.40015, rel=1e-4) assert output1.output.n_steps == 5 + # fix_symmetry makes no difference for this structure relaxer combo + # just testing that passing fix_symmetry doesn't break + final_spg_num = output1.output.structure.get_space_group_info()[1] + assert final_spg_num == 99 + @pytest.mark.parametrize("relax_cell", [True, False]) def test_gap_relax_maker(si_structure: Structure, test_dir: Path, relax_cell: bool): @@ -259,7 +370,7 @@ def test_gap_relax_maker(si_structure: Structure, test_dir: Path, relax_cell: bo si_structure.translate_sites(0, [0, 0, 0.1]) # generate job - # Test files have been provided by Yuanbin Liu (University of Oxford) + # Test files have been provided by @YuanbinLiu (University of Oxford) max_step = 25 job = GAPRelaxMaker( calculator_kwargs={ diff --git a/tests/forcefields/test_utils.py b/tests/forcefields/test_utils.py index 55f113b926..fe67a88e4d 100644 --- a/tests/forcefields/test_utils.py +++ b/tests/forcefields/test_utils.py @@ -1,10 +1,12 @@ import os import pytest +from ase.build import bulk from ase.calculators.lj import LennardJones from ase.optimize import BFGS +from ase.spacegroup.symmetrize import check_symmetry from numpy.testing import assert_allclose -from pymatgen.io.ase import AseAtomsAdaptor +from pymatgen.core import Structure from atomate2.forcefields import MLFF from atomate2.forcefields.utils import ( @@ -19,8 +21,9 @@ def test_safe_import(): assert FrechetCellFilter is None or FrechetCellFilter.__module__ == "ase.filters" -def test_trajectory_observer(si_structure, test_dir, tmp_dir): - atoms = AseAtomsAdaptor.get_atoms(structure=si_structure, calculator=LennardJones()) +def test_trajectory_observer(si_structure: Structure, test_dir, tmp_dir): + atoms = si_structure.to_ase_atoms() + atoms.set_calculator(LennardJones()) traj = TrajectoryObserver(atoms) @@ -132,11 +135,29 @@ def test_relaxer(si_structure, test_dir, tmp_dir, optimizer, traj_file): def test_ext_load(): - forcefield_to_callable = { + force_field_to_callable = { "CHGNet": {"@module": "chgnet.model.dynamics", "@callable": "CHGNetCalculator"}, "MACE": {"@module": "mace.calculators", "@callable": "mace_mp"}, } - for forcefield in ["CHGNet", "MACE"]: - calc_from_decode = ase_calculator(forcefield_to_callable[forcefield]) - calc_from_preset = ase_calculator(f"{MLFF(forcefield)}") + for force_field in ("CHGNet", "MACE"): + calc_from_decode = ase_calculator(force_field_to_callable[force_field]) + calc_from_preset = ase_calculator(f"{MLFF(force_field)}") assert isinstance(calc_from_decode, type(calc_from_preset)) + + +@pytest.mark.parametrize(("fix_symmetry"), [True, False]) +def test_fix_symmetry(fix_symmetry): + # adapted from the example at https://wiki.fysik.dtu.dk/ase/ase/constraints.html#the-fixsymmetry-class + relaxer = Relaxer( + calculator=LennardJones(), relax_cell=True, fix_symmetry=fix_symmetry + ) + atoms_al = bulk("Al", "bcc", a=2 / 3**0.5, cubic=True) + atoms_al = atoms_al * (2, 2, 2) + atoms_al.positions[0, 0] += 1e-7 + symmetry_init = check_symmetry(atoms_al, 1e-6) + final_struct: Structure = relaxer.relax(atoms=atoms_al, steps=1)["final_structure"] + symmetry_final = check_symmetry(final_struct.to_ase_atoms(), 1e-6) + if fix_symmetry: + assert symmetry_init["number"] == symmetry_final["number"] == 229 + else: + assert symmetry_init["number"] != symmetry_final["number"] == 99 diff --git a/tests/test_data/structures/BaTiO3.cif b/tests/test_data/structures/BaTiO3.cif new file mode 100644 index 0000000000..df0e6eee84 --- /dev/null +++ b/tests/test_data/structures/BaTiO3.cif @@ -0,0 +1,37 @@ +# generated using pymatgen +data_BaTiO3 +_symmetry_space_group_name_H-M 'P 1' +_cell_length_a 4.00768164 +_cell_length_b 4.00768164 +_cell_length_c 4.00768164 +_cell_angle_alpha 90.00000000 +_cell_angle_beta 90.00000000 +_cell_angle_gamma 90.00000000 +_symmetry_Int_Tables_number 1 +_chemical_formula_structural BaTiO3 +_chemical_formula_sum 'Ba1 Ti1 O3' +_cell_volume 64.36942706 +_cell_formula_units_Z 1 +loop_ + _symmetry_equiv_pos_site_id + _symmetry_equiv_pos_as_xyz + 1 'x, y, z' +loop_ + _atom_type_symbol + _atom_type_oxidation_number + Ba2+ 2.0 + Ti4+ 4.0 + O2- -2.0 +loop_ + _atom_site_type_symbol + _atom_site_label + _atom_site_symmetry_multiplicity + _atom_site_fract_x + _atom_site_fract_y + _atom_site_fract_z + _atom_site_occupancy + Ba2+ Ba0 1 0.50000000 0.50000000 0.50000000 1 + Ti4+ Ti1 1 0.00000000 0.00000000 0.00000000 1 + O2- O2 1 0.50000000 0.00000000 0.00000000 1 + O2- O3 1 0.00000000 0.50000000 0.00000000 1 + O2- O4 1 0.00000000 0.00000000 0.50000000 1