Add geometry optimisation (#33)

* Add pytest flag for slow tests * Add geometry optimisation * Add geometry optimisation tests * Rename geometry optimization * Update type hints and tidy code * Refactor tests * Add option to save optimised structure * Fix type hints * Update docs * Make type hints consistent * Add option to save optimization trajectory * Fix docstrings * Clarify default filter * Remove unlinking for trajectory file --------- Co-authored-by: ElliottKasoar <[email protected]>
stfc · Feb 23, 2024 · ead43ec · ead43ec
1 parent 53072d0
commit ead43ec
Show file tree

Hide file tree

Showing 12 changed files with 336 additions and 102 deletions.
diff --git a/.pre-commit-config.yaml b/.pre-commit-config.yaml
@@ -13,7 +13,7 @@ repos:
   rev: v3.15.0
   hooks:
   - id: pyupgrade
-    args: ["--py310-plus"]
+    args: ["--py39-plus"]
 
 - repo: https://github.com/PyCQA/isort
   rev: 5.13.2

diff --git a/conftest.py b/conftest.py
@@ -14,21 +14,31 @@ def pytest_addoption(parser):
         default=False,
         help="Test additional MLIPs",
     )
+    parser.addoption(
+        "--run-slow",
+        action="store_true",
+        default=False,
+        help="Run slow tests",
+    )
 
 
 def pytest_configure(config):
     """Configure pytest to include marker for extra MLIPs."""
     config.addinivalue_line(
         "markers", "extra_mlips: mark test as containing extra MLIPs"
     )
+    config.addinivalue_line("markers", "slow: mark test as slow")
 
 
 def pytest_collection_modifyitems(config, items):
     """Skip tests if marker applied to unit tests."""
-    if config.getoption("--run-extra-mlips"):
-        # --run-extra-mlips given in cli: do not skip tests for extra MLIPs
+    if config.getoption("--run-extra-mlips") or config.getoption("--run-slow"):
+        # --run-extra-mlips or --run-slow given in cli: do not skip tests
         return
     skip_extra_mlips = pytest.mark.skip(reason="need --run-extra-mlips option to run")
+    skip_slow = pytest.mark.skip(reason="need --run-slow option to run")
     for item in items:
         if "extra_mlips" in item.keywords:
             item.add_marker(skip_extra_mlips)
+        if "slow" in item.keywords:
+            item.add_marker(skip_slow)
diff --git a/docs/source/apidoc/janus_core.rst b/docs/source/apidoc/janus_core.rst
@@ -4,20 +4,20 @@ janus\_core package
 Submodules
 ----------
 
-janus\_core.mlip\_calculators module
-------------------------------------
+janus\_core.geom\_opt module
+----------------------------
 
-.. automodule:: janus_core.mlip_calculators
+.. automodule:: janus_core.geom_opt
    :members:
    :special-members:
    :private-members:
    :undoc-members:
    :show-inheritance:
 
-janus\_core.read\_args module
------------------------------
+janus\_core.mlip\_calculators module
+------------------------------------
 
-.. automodule:: janus_core.read_args
+.. automodule:: janus_core.mlip_calculators
    :members:
    :special-members:
    :private-members:

diff --git a/docs/source/conf.py b/docs/source/conf.py
@@ -36,7 +36,7 @@
 intersphinx_mapping = {
     "python": ("https://docs.python.org/3", None),
     "numpy": ("https://numpy.org/doc/stable/", None),
-    "ase": ("https://wiki.fysik.dtu.dk/ase", None),
+    "ase": ("https://wiki.fysik.dtu.dk/ase/", None),
 }
 
 # Add any paths that contain templates here, relative to this directory.

diff --git a/janus_core/geom_opt.py b/janus_core/geom_opt.py
@@ -0,0 +1,92 @@
+"""Geometry optimization."""
+
+from typing import Any, Optional
+
+from ase import Atoms
+from ase.io import read, write
+
+try:
+    from ase.filters import FrechetCellFilter as DefaultFilter
+except ImportError:
+    from ase.constraints import ExpCellFilter as DefaultFilter
+
+from ase.optimize import LBFGS
+
+
+def optimize(
+    atoms: Atoms,
+    fmax: float = 0.1,
+    dyn_kwargs: Optional[dict[str, Any]] = None,
+    filter_func: Optional[callable] = DefaultFilter,
+    filter_kwargs: Optional[dict[str, Any]] = None,
+    optimizer: callable = LBFGS,
+    opt_kwargs: Optional[dict[str, Any]] = None,
+    struct_kwargs: Optional[dict[str, Any]] = None,
+    traj_kwargs: Optional[dict[str, Any]] = None,
+) -> Atoms:
+    """Optimize geometry of input structure.
+
+    Parameters
+    ----------
+    atoms : Atoms
+        Atoms object to optimize geometry for.
+    fmax : float
+        Set force convergence criteria for optimizer in units eV/Å.
+    dyn_kwargs : Optional[dict[str, Any]]
+        kwargs to pass to dyn.run. Default is {}.
+    filter_func : Optional[callable]
+        Apply constraints to atoms through ASE filter function.
+        Default is `FrechetCellFilter` if available otherwise `ExpCellFilter`.
+    filter_kwargs : Optional[dict[str, Any]]
+        kwargs to pass to filter_func. Default is {}.
+    optimzer : callable
+        ASE optimization function. Default is `LBFGS`.
+    opt_kwargs : Optional[dict[str, Any]]
+        kwargs to pass to optimzer. Default is {}.
+    struct_kwargs : Optional[dict[str, Any]]
+        kwargs to pass to ase.io.write to save optimized structure.
+        Must include "filename" keyword. Default is {}.
+    traj_kwargs : Optional[dict[str, Any]]
+        kwargs to pass to ase.io.write to save optimization trajectory.
+        Must include "filename" keyword. Default is {}.
+
+    Returns
+    -------
+    atoms: Atoms
+        Structure with geometry optimized.
+    """
+    dyn_kwargs = dyn_kwargs if dyn_kwargs else {}
+    filter_kwargs = filter_kwargs if filter_kwargs else {}
+    opt_kwargs = opt_kwargs if opt_kwargs else {}
+    struct_kwargs = struct_kwargs if struct_kwargs else {}
+    traj_kwargs = traj_kwargs if traj_kwargs else {}
+
+    if struct_kwargs and "filename" not in struct_kwargs:
+        raise ValueError("'filename' must be included in struct_kwargs")
+
+    if traj_kwargs and "filename" not in traj_kwargs:
+        raise ValueError("'filename' must be included in traj_kwargs")
+
+    if traj_kwargs and "trajectory" not in opt_kwargs:
+        raise ValueError(
+            "'trajectory' must be a key in opt_kwargs to save the trajectory."
+        )
+
+    if filter_func is not None:
+        filtered_atoms = filter_func(atoms, **filter_kwargs)
+        dyn = optimizer(filtered_atoms, **opt_kwargs)
+    else:
+        dyn = optimizer(atoms, **opt_kwargs)
+
+    dyn.run(fmax=fmax, **dyn_kwargs)
+
+    # Write out optimized structure
+    if struct_kwargs:
+        write(images=atoms, **struct_kwargs)
+
+    # Reformat trajectory file from binary
+    if traj_kwargs:
+        traj = read(opt_kwargs["trajectory"], index=":")
+        write(images=traj, **traj_kwargs)
+
+    return atoms
diff --git a/janus_core/mlip_calculators.py b/janus_core/mlip_calculators.py
@@ -5,8 +5,6 @@
 - https://github.com/Quantum-Accelerators/quacc.git
 """
 
-from __future__ import annotations
-
 from typing import Literal
 
 from ase.calculators.calculator import Calculator

diff --git a/janus_core/single_point.py b/janus_core/single_point.py
@@ -1,8 +1,7 @@
 """Perpare and perform single point calculations."""
 
-from __future__ import annotations
-
 import pathlib
+from typing import Any, Optional, Union
 
 from ase.io import read
 from numpy import ndarray
@@ -13,13 +12,12 @@
 class SinglePoint:
     """Perpare and perform single point calculations."""
 
-    #  pylint: disable=dangerous-default-value
     def __init__(
         self,
         system: str,
         architecture: str = "mace_mp",
         device: str = "cpu",
-        read_kwargs: dict = {},
+        read_kwargs: Optional[dict[str, Any]] = None,
         **kwargs,
     ) -> None:
         """
@@ -34,16 +32,17 @@ def __init__(
             Default is "mace_mp".
         device : str
             Device to run model on. Default is "cpu".
-        read_kwargs : dict
+        read_kwargs : Optional[dict[str, Any]]
             kwargs to pass to ase.io.read. Default is {}.
         """
         self.architecture = architecture
         self.device = device
         self.system = system
 
         # Read system and get calculator
+        read_kwargs = read_kwargs if read_kwargs else {}
         self.read_system(**read_kwargs)
-        self.get_calculator(**kwargs)
+        self.set_calculator(**kwargs)
 
     def read_system(self, **kwargs) -> None:
         """Read system and system name.
@@ -54,12 +53,14 @@ def read_system(self, **kwargs) -> None:
         self.sys = read(self.system, **kwargs)
         self.sysname = pathlib.Path(self.system).stem
 
-    def get_calculator(self, read_kwargs: dict = {}, **kwargs) -> None:
+    def set_calculator(
+        self, read_kwargs: Optional[dict[str, Any]] = None, **kwargs
+    ) -> None:
         """Configure calculator and attach to system.
 
         Parameters
         ----------
-        read_kwargs : dict
+        read_kwargs : Optional[dict[str, Any]]
             kwargs to pass to ase.io.read. Default is {}.
         """
         calculator = choose_calculator(
@@ -68,6 +69,7 @@ def get_calculator(self, read_kwargs: dict = {}, **kwargs) -> None:
             **kwargs,
         )
         if self.sys is None:
+            read_kwargs = read_kwargs if read_kwargs else {}
             self.read_system(**read_kwargs)
 
         if isinstance(self.sys, list):
@@ -76,12 +78,12 @@ def get_calculator(self, read_kwargs: dict = {}, **kwargs) -> None:
         else:
             self.sys.calc = calculator
 
-    def _get_potential_energy(self) -> float | list[float]:
+    def _get_potential_energy(self) -> Union[float, list[float]]:
         """Calculate potential energy using MLIP.
 
         Returns
         -------
-        potential_energy : float | list[float]
+        potential_energy : Union[float, list[float]]
             Potential energy of system(s).
         """
         if isinstance(self.sys, list):
@@ -92,12 +94,12 @@ def _get_potential_energy(self) -> float | list[float]:
 
         return self.sys.get_potential_energy()
 
-    def _get_forces(self) -> ndarray | list[ndarray]:
+    def _get_forces(self) -> Union[ndarray, list[ndarray]]:
         """Calculate forces using MLIP.
 
         Returns
         -------
-        forces : ndarray | list[ndarray]
+        forces : Union[ndarray, list[ndarray]]
             Forces of system(s).
         """
         if isinstance(self.sys, list):
@@ -108,12 +110,12 @@ def _get_forces(self) -> ndarray | list[ndarray]:
 
         return self.sys.get_forces()
 
-    def _get_stress(self) -> ndarray | list[ndarray]:
+    def _get_stress(self) -> Union[ndarray, list[ndarray]]:
         """Calculate stress using MLIP.
 
         Returns
         -------
-        stress : ndarray | list[ndarray]
+        stress : Union[ndarray, list[ndarray]]
             Stress of system(s).
         """
         if isinstance(self.sys, list):
@@ -124,18 +126,20 @@ def _get_stress(self) -> ndarray | list[ndarray]:
 
         return self.sys.get_stress()
 
-    def run_single_point(self, properties: str | list[str] | None = None) -> dict:
+    def run_single_point(
+        self, properties: Optional[Union[str, list[str]]] = None
+    ) -> dict[str, Any]:
         """Run single point calculations.
 
         Parameters
         ----------
-        properties : str | List[str] | None
+        properties : Optional[Union[str, list[str]]]
             Physical properties to calculate. If not specified, "energy",
             "forces", and "stress" will be returned.
 
         Returns
         -------
-        results : dict
+        results : dict[str, Any]
             Dictionary of calculated results.
         """
         results = {}

diff --git a/tests/data/H2O.cif b/tests/data/H2O.cif
@@ -0,0 +1,14 @@
+data_image0
+_chemical_formula_structural       H2O
+_chemical_formula_sum              "H2 O1"
+loop_
+  _atom_site_type_symbol
+  _atom_site_label
+  _atom_site_symmetry_multiplicity
+  _atom_site_Cartn_x
+  _atom_site_Cartn_y
+  _atom_site_Cartn_z
+  _atom_site_occupancy
+  H   H1        1.0  0.95750  0.00000  0.00000  1.0000
+  H   H2        1.0  -0.23990  0.92696  0.00000  1.0000
+  O   O1        1.0  0.00000  0.00000  0.00000  1.0000