From ead43ec56f267b6ab6f361d8d67d1b5f7bfc2c8f Mon Sep 17 00:00:00 2001 From: ElliottKasoar <45317199+ElliottKasoar@users.noreply.github.com> Date: Fri, 23 Feb 2024 14:57:42 +0000 Subject: [PATCH] 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 --- .pre-commit-config.yaml | 2 +- conftest.py | 14 +++- docs/source/apidoc/janus_core.rst | 12 +-- docs/source/conf.py | 2 +- janus_core/geom_opt.py | 92 +++++++++++++++++++++ janus_core/mlip_calculators.py | 2 - janus_core/single_point.py | 38 +++++---- tests/data/H2O.cif | 14 ++++ tests/test_geom_opt.py | 129 ++++++++++++++++++++++++++++++ tests/test_mlip_calculators.py | 40 ++++----- tests/test_single_point.py | 89 ++++++++++----------- tox.ini | 4 +- 12 files changed, 336 insertions(+), 102 deletions(-) create mode 100644 janus_core/geom_opt.py create mode 100644 tests/data/H2O.cif create mode 100644 tests/test_geom_opt.py diff --git a/.pre-commit-config.yaml b/.pre-commit-config.yaml index c69896ad..0f6f6a71 100644 --- 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 index a0d671c4..fc003f79 100644 --- a/conftest.py +++ b/conftest.py @@ -14,6 +14,12 @@ 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): @@ -21,14 +27,18 @@ def pytest_configure(config): 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 index da2b122c..7570e295 100644 --- 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 index 5e662547..c0ba7f27 100755 --- 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 new file mode 100644 index 00000000..607d0ae9 --- /dev/null +++ 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 index 7e5b874a..5107cd87 100644 --- 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 index b948c5b9..5942ae42 100644 --- 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,7 +32,7 @@ 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 @@ -42,8 +40,9 @@ def __init__( 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 new file mode 100644 index 00000000..871b99a8 --- /dev/null +++ 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 diff --git a/tests/test_geom_opt.py b/tests/test_geom_opt.py new file mode 100644 index 00000000..7348478a --- /dev/null +++ b/tests/test_geom_opt.py @@ -0,0 +1,129 @@ +"""Test geometry optimization.""" + +from pathlib import Path + +try: + from ase.filters import UnitCellFilter +except ImportError: + from ase.constraints import UnitCellFilter +from ase.io import read +import pytest + +from janus_core.geom_opt import optimize +from janus_core.single_point import SinglePoint + +DATA_PATH = Path(__file__).parent / "data" +MODEL_PATH = Path(__file__).parent / "models" / "mace_mp_small.model" + +test_data = [ + ("mace", "NaCl.cif", MODEL_PATH, -27.046359959669214, {}), + ("mace", "NaCl.cif", MODEL_PATH, -27.04636199814088, {"fmax": 0.001}), + ( + "mace", + "NaCl.cif", + MODEL_PATH, + -27.0463392211678, + {"filter_func": UnitCellFilter}, + ), + ("mace", "H2O.cif", MODEL_PATH, -14.051389496520015, {"filter_func": None}), + ( + "mace", + "NaCl.cif", + MODEL_PATH, + -27.04631447979008, + {"filter_func": UnitCellFilter, "filter_kwargs": {"scalar_pressure": 0.0001}}, + ), + ( + "mace", + "NaCl.cif", + MODEL_PATH, + -27.046353221978332, + {"opt_kwargs": {"alpha": 100}}, + ), + ( + "mace", + "NaCl.cif", + MODEL_PATH, + -27.03561540212425, + {"filter_func": UnitCellFilter, "dyn_kwargs": {"steps": 1}}, + ), +] + + +@pytest.mark.parametrize( + "architecture, structure, model_path, expected, kwargs", test_data +) +def test_optimize(architecture, structure, model_path, expected, kwargs): + """Test optimizing geometry using MACE.""" + data_path = DATA_PATH / structure + single_point = SinglePoint( + system=data_path, architecture=architecture, model_paths=model_path + ) + + init_energy = single_point.run_single_point("energy")["energy"] + + atoms = optimize(single_point.sys, **kwargs) + + assert atoms.get_potential_energy() < init_energy + assert atoms.get_potential_energy() == pytest.approx(expected) + + +def test_saving_struct(tmp_path): + """Test saving optimized structure.""" + data_path = DATA_PATH / "NaCl.cif" + struct_path = tmp_path / "NaCl.xyz" + single_point = SinglePoint( + system=data_path, architecture="mace", model_paths=MODEL_PATH + ) + + init_energy = single_point.run_single_point("energy")["energy"] + + optimize( + single_point.sys, + struct_kwargs={"filename": struct_path, "format": "extxyz"}, + ) + opt_struct = read(struct_path) + + assert opt_struct.get_potential_energy() < init_energy + + +def test_saving_traj(tmp_path): + """Test saving optimization trajectory output.""" + data_path = DATA_PATH / "NaCl.cif" + single_point = SinglePoint( + system=data_path, architecture="mace", model_paths=MODEL_PATH + ) + optimize(single_point.sys, opt_kwargs={"trajectory": str(tmp_path / "NaCl.traj")}) + traj = read(tmp_path / "NaCl.traj", index=":") + assert len(traj) == 3 + + +def test_traj_reformat(tmp_path): + """Test saving optimization trajectory in different format.""" + data_path = DATA_PATH / "NaCl.cif" + traj_path_binary = tmp_path / "NaCl.traj" + traj_path_xyz = tmp_path / "NaCl-traj.xyz" + + single_point = SinglePoint( + system=data_path, architecture="mace", model_paths=MODEL_PATH + ) + + optimize( + single_point.sys, + opt_kwargs={"trajectory": str(traj_path_binary)}, + traj_kwargs={"filename": traj_path_xyz}, + ) + traj = read(tmp_path / "NaCl-traj.xyz", index=":") + + assert len(traj) == 3 + + +def test_missing_traj_kwarg(tmp_path): + """Test saving optimization trajectory in different format.""" + data_path = DATA_PATH / "NaCl.cif" + traj_path = tmp_path / "NaCl-traj.xyz" + single_point = SinglePoint( + system=data_path, architecture="mace", model_paths=MODEL_PATH + ) + with pytest.raises(ValueError): + optimize(single_point.sys, traj_kwargs={"filename": traj_path}) diff --git a/tests/test_mlip_calculators.py b/tests/test_mlip_calculators.py index 62006b92..46397c34 100644 --- a/tests/test_mlip_calculators.py +++ b/tests/test_mlip_calculators.py @@ -6,31 +6,21 @@ from janus_core.mlip_calculators import choose_calculator - -def test_mace(): - """Test mace calculator can be configured.""" - model_path = Path(Path(__file__).parent, "models", "mace_mp_small.model") - calculator = choose_calculator( - architecture="mace", device="cpu", model_paths=model_path - ) - assert calculator.parameters["version"] is not None - - -def test_mace_off(): - """Test mace_off calculator can be configured.""" - calculator = choose_calculator( - architecture="mace_off", - device="cpu", - ) - assert calculator.parameters["version"] is not None - - -def test_mace_mp(): - """Test mace_mp calculator can be configured.""" - calculator = choose_calculator( - architecture="mace_mp", - device="cpu", - ) +test_data = [ + ( + "mace", + "cpu", + {"model_paths": Path(__file__).parent / "models" / "mace_mp_small.model"}, + ), + ("mace_off", "cpu", {}), + ("mace_mp", "cpu", {}), +] + + +@pytest.mark.parametrize("architecture, device, kwargs", test_data) +def test_mace(architecture, device, kwargs): + """Test mace calculators can be configured.""" + calculator = choose_calculator(architecture=architecture, device=device, **kwargs) assert calculator.parameters["version"] is not None diff --git a/tests/test_single_point.py b/tests/test_single_point.py index 3c899da2..44b77ec4 100644 --- a/tests/test_single_point.py +++ b/tests/test_single_point.py @@ -6,57 +6,54 @@ from janus_core.single_point import SinglePoint - -def test_potential_energy(): - """Test single point energy using MACE calculator.""" - data_path = Path(Path(__file__).parent, "data", "benzene.xyz") - single_point = SinglePoint(system=data_path) - - assert single_point.run_single_point("energy")["energy"] == -76.0605725422795 - - -def test_single_point_kwargs(): - """Test kwargs passed when using MACE calculator for single point energy.""" - - data_path = Path(Path(__file__).parent, "data", "benzene.xyz") - model_path = Path(Path(__file__).parent, "models", "mace_mp_small.model") +DATA_PATH = Path(__file__).parent / "data" +MODEL_PATH = Path(__file__).parent / "models" / "mace_mp_small.model" + +test_data = [ + (DATA_PATH / "benzene.xyz", -76.0605725422795, "energy", "energy", {}), + ( + DATA_PATH / "benzene.xyz", + -76.06057739257812, + ["energy"], + "energy", + {"default_dtype": "float32"}, + ), + ( + DATA_PATH / "benzene.xyz", + -0.0360169762840179, + ["forces"], + "forces", + {"idx": [0, 1]}, + ), + (DATA_PATH / "NaCl.cif", -0.004783275999053424, ["stress"], "stress", {"idx": [0]}), +] + + +@pytest.mark.parametrize("system, expected, properties, prop_key, kwargs", test_data) +def test_potential_energy(system, expected, properties, prop_key, kwargs): + """Test single point energy using MACE calculators.""" single_point = SinglePoint( - system=data_path, architecture="mace", model_paths=model_path + system=system, architecture="mace", model_paths=MODEL_PATH, **kwargs ) + results = single_point.run_single_point(properties)[prop_key] - assert single_point.run_single_point(["energy"])["energy"] == -76.0605725422795 - - -def test_single_point_forces(): - """Test single point forces using MACE calculator.""" - data_path = Path(Path(__file__).parent, "data", "benzene.xyz") - model_path = Path(Path(__file__).parent, "models", "mace_mp_small.model") - single_point = SinglePoint( - system=data_path, architecture="mace", model_paths=model_path - ) - - assert single_point.run_single_point(["forces"])["forces"][0, 1] == pytest.approx( - -0.0360169762840179 - ) - - -def test_single_point_stress(): - """Test single point stress using MACE calculator.""" - data_path = Path(Path(__file__).parent, "data", "NaCl.cif") - model_path = Path(Path(__file__).parent, "models", "mace_mp_small.model") - single_point = SinglePoint( - system=data_path, architecture="mace", model_paths=model_path - ) - - assert single_point.run_single_point(["stress"])["stress"][0] == pytest.approx( - -0.004783275999053424 - ) + # Check correct values returned + idx = kwargs.pop("idx", None) + if idx is not None: + if len(idx) == 1: + assert results[idx[0]] == pytest.approx(expected) + elif len(idx) == 2: + assert results[idx[0], idx[1]] == pytest.approx(expected) + else: + raise ValueError(f"Invalid index: {idx}") + else: + assert results == pytest.approx(expected) def test_single_point_none(): """Test single point stress using MACE calculator.""" - data_path = Path(Path(__file__).parent, "data", "NaCl.cif") - model_path = Path(Path(__file__).parent, "models", "mace_mp_small.model") + data_path = DATA_PATH / "NaCl.cif" + model_path = MODEL_PATH single_point = SinglePoint( system=data_path, architecture="mace", model_paths=model_path ) @@ -68,8 +65,8 @@ def test_single_point_none(): def test_single_point_traj(): """Test single point stress using MACE calculator.""" - data_path = Path(Path(__file__).parent, "data", "benzene-traj.xyz") - model_path = Path(Path(__file__).parent, "models", "mace_mp_small.model") + data_path = DATA_PATH / "benzene-traj.xyz" + model_path = MODEL_PATH single_point = SinglePoint( system=data_path, architecture="mace", diff --git a/tox.ini b/tox.ini index b71fdfa9..254d55d9 100644 --- a/tox.ini +++ b/tox.ini @@ -4,7 +4,7 @@ envlist = py311 [testenv] usedevelop=True -[testenv:py{39,310,311,312}] +[testenv:py{39,310,312}] description = Run the test suite against Python versions allowlist_externals = poetry commands_pre = poetry install --no-root --sync @@ -26,4 +26,4 @@ commands = poetry run sphinx-build -nW --keep-going -b html {posargs} docs/sourc description = Run the additional tests suite against Python versions allowlist_externals = poetry commands_pre = poetry install --no-root --sync --with extra-mlips -commands = poetry run pytest --run-extra-mlips --cov janus_core --import-mode importlib +commands = poetry run pytest --run-extra-mlips --run-slow --cov janus_core --import-mode importlib