Add fix_symmetry: bool = False option to forcefield relax makers (#789

)

* add fix_symmetry option and tests to force field relaxer

* add fix_symmetry option to one test job

* fix linting of test_jobs

* add missing documentation to GAP/MACE relaxmaker

* restore test_fix_symmetry

* restore fix_symmetry to NequipRelaxMaker doc str

* add fix_symmetry and symprec back into relaxer jobs. Add fix_symmetry/symprec test for MaceRelaxMaker

* add fix_symmetry and symprec to task_doc, add test for chgnet

* separated fix_symmetry from standard relax_maker tests

* fixed test_chgnet_relax_maker_fix_symmetry

* changed fix_symmetry and symprec to optional arguments in from_ase_compatible_results

* test_fix_symmetry set steps=1 in relaxer.relax(), takes test from 6 to 1.5 sec

* assert expected spacegroup in test_fix_symmetry

* cut fix_symmetry and symprec combos in half in test_nequip_relax_maker for speed

* parametrize fix_symmetry=True|False on test_mace_relax_maker

* test_relax_maker_mace check for subgroup instead of group when using fixed symmetry

* fix test_mace_relax_maker

---------

Co-authored-by: Janosh Riebesell <janosh.riebesell@gmail.com>

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)

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,

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):

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)."""