Add single point calculations (#31)

* Change ASE dependency to specific commit

* Allow e as variable name

* Add dependencies for MACE

* Add optional dependencies for additional MLIPs

* Add assertion for example test

* Add tests for other MLIPs

* Add configuring MLIP calculators

* Add single point energy calculation

* Update API docs

* Tidy docstrings

* Add ASE docs

* Allow more arguments with pylint

* Suppress pylint warning

* Add force and stress calculations

* Add single point calculations for trajectories

* Add numpy doc links

* Replace UiO66 with NaCl for tests

* Replace os with pathlib

* Rename mace_mp model file

* Use tagged ASE

* Add suggested installtion for ASE

---------

Co-authored-by: ElliottKasoar <[email protected]>

README.md

@@ -24,6 +24,18 @@ pre-commit install  # install pre-commit hooks
 pytest -v  # discover and run all tests
 ```
 
+Manually updating ASE via https://gitlab.com/ase/ase is strongly recommended, as tags are no longer regularly published. For example:
+
+```shell
+pip install git+https://gitlab.com/ase/ase.git@b31569210d739bd12c8ad2b6ec0290108e049eea
+```
+
+To prevent poetry downgrading ASE when installing in future, add the commit to pyproject.toml:
+
+```shell
+poetry add git+https://gitlab.com:ase/ase.git#b31569210d739bd12c8ad2b6ec0290108e049eea
+```
+
 ## License
 
 [BSD 3-Clause License](LICENSE)

conftest.py

@@ -0,0 +1,34 @@
+"""Configure pytest.
+
+Based on https://docs.pytest.org/en/latest/example/simple.html.
+"""
+
+import pytest
+
+
+def pytest_addoption(parser):
+    """Add flag to run tests for extra MLIPs."""
+    parser.addoption(
+        "--run-extra-mlips",
+        action="store_true",
+        default=False,
+        help="Test additional MLIPs",
+    )
+
+
+def pytest_configure(config):
+    """Configure pytest to include marker for extra MLIPs."""
+    config.addinivalue_line(
+        "markers", "extra_mlips: mark test as containing extra MLIPs"
+    )
+
+
+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
+        return
+    skip_extra_mlips = pytest.mark.skip(reason="need --run-extra-mlips option to run")
+    for item in items:
+        if "extra_mlips" in item.keywords:
+            item.add_marker(skip_extra_mlips)

docs/source/apidoc/janus_core.rst

@@ -1,6 +1,39 @@
 janus\_core package
 ===================
 
+Submodules
+----------
+
+janus\_core.mlip\_calculators module
+------------------------------------
+
+.. automodule:: janus_core.mlip_calculators
+   :members:
+   :special-members:
+   :private-members:
+   :undoc-members:
+   :show-inheritance:
+
+janus\_core.read\_args module
+-----------------------------
+
+.. automodule:: janus_core.read_args
+   :members:
+   :special-members:
+   :private-members:
+   :undoc-members:
+   :show-inheritance:
+
+janus\_core.single\_point module
+--------------------------------
+
+.. automodule:: janus_core.single_point
+   :members:
+   :special-members:
+   :private-members:
+   :undoc-members:
+   :show-inheritance:
+
 Module contents
 ---------------
 

docs/source/conf.py

@@ -35,6 +35,8 @@
 
 intersphinx_mapping = {
     "python": ("https://docs.python.org/3", None),
+    "numpy": ("https://numpy.org/doc/stable/", None),
+    "ase": ("https://wiki.fysik.dtu.dk/ase", None),
 }
 
 # Add any paths that contain templates here, relative to this directory.

janus_core/__init__.py

@@ -1,7 +1,6 @@
-"""
-janus_core
+"""janus_core.
 
-Tools for machine learnt interatomic potentials
+Tools for machine learnt interatomic potentials.
 """
 
 __version__ = "0.1.0a1"

janus_core/mlip_calculators.py

@@ -0,0 +1,101 @@
+"""Configure MLIP calculators.
+
+Similar in spirit with matcalc and quacc approaches
+- https://github.com/materialsvirtuallab/matcalc
+- https://github.com/Quantum-Accelerators/quacc.git
+"""
+
+from __future__ import annotations
+
+from typing import Literal
+
+from ase.calculators.calculator import Calculator
+
+architectures = ["mace", "mace_mp", "mace_off", "m3gnet", "chgnet"]
+
+
+def choose_calculator(
+    architecture: Literal[architectures] = "mace", **kwargs
+) -> Calculator:
+    """Choose MLIP calculator to configure.
+
+    Parameters
+    ----------
+    architecture : Literal[architectures], optional
+        MLIP architecture. Default is "mace".
+
+    Raises
+    ------
+    ModuleNotFoundError
+        MLIP module not correctly been installed.
+    ValueError
+        Invalid architecture specified.
+
+    Returns
+    -------
+    calculator : Calculator
+        Configured MLIP calculator.
+    """
+    # pylint: disable=import-outside-toplevel
+    # pylint: disable=too-many-branches
+    # pylint: disable=import-error
+    # Optional imports handled via `architecture`. We could catch these,
+    # but the error message is clear if imports are missing.
+    if architecture == "mace":
+        from mace import __version__
+        from mace.calculators import MACECalculator
+
+        if "default_dtype" not in kwargs:
+            kwargs["default_dtype"] = "float64"
+        if "device" not in kwargs:
+            kwargs["device"] = "cuda"
+        calculator = MACECalculator(**kwargs)
+
+    elif architecture == "mace_mp":
+        from mace import __version__
+        from mace.calculators import mace_mp
+
+        if "default_dtype" not in kwargs:
+            kwargs["default_dtype"] = "float64"
+        if "device" not in kwargs:
+            kwargs["device"] = "cuda"
+        if "model" not in kwargs:
+            kwargs["model"] = "small"
+        calculator = mace_mp(**kwargs)
+
+    elif architecture == "mace_off":
+        from mace import __version__
+        from mace.calculators import mace_off
+
+        if "default_dtype" not in kwargs:
+            kwargs["default_dtype"] = "float64"
+        if "device" not in kwargs:
+            kwargs["device"] = "cuda"
+        if "model" not in kwargs:
+            kwargs["model"] = "small"
+        calculator = mace_off(**kwargs)
+
+    elif architecture == "m3gnet":
+        from matgl import __version__, load_model
+        from matgl.ext.ase import M3GNetCalculator
+
+        if "model" not in kwargs:
+            model = load_model("M3GNet-MP-2021.2.8-DIRECT-PES")
+        if "stress_weight" not in kwargs:
+            kwargs.setdefault("stress_weight", 1.0 / 160.21766208)
+        calculator = M3GNetCalculator(potential=model, **kwargs)
+
+    elif architecture == "chgnet":
+        from chgnet import __version__
+        from chgnet.model.dynamics import CHGNetCalculator
+
+        calculator = CHGNetCalculator(**kwargs)
+
+    else:
+        raise ValueError(
+            f"Unrecognized {architecture=}. Suported architectures are {architectures}"
+        )
+
+    calculator.parameters["version"] = __version__
+
+    return calculator

janus_core/single_point.py

@@ -0,0 +1,154 @@
+"""Perpare and perform single point calculations."""
+
+from __future__ import annotations
+
+import pathlib
+
+from ase.io import read
+from numpy import ndarray
+
+from janus_core.mlip_calculators import choose_calculator
+
+
+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 = {},
+        **kwargs,
+    ) -> None:
+        """
+        Initialise class.
+
+        Attributes
+        ----------
+        system : str
+            System to simulate.
+        architecture : str
+            MLIP architecture to use for single point calculations.
+            Default is "mace_mp".
+        device : str
+            Device to run model on. Default is "cpu".
+        read_kwargs : dict
+            kwargs to pass to ase.io.read. Default is {}.
+        """
+        self.architecture = architecture
+        self.device = device
+        self.system = system
+
+        # Read system and get calculator
+        self.read_system(**read_kwargs)
+        self.get_calculator(**kwargs)
+
+    def read_system(self, **kwargs) -> None:
+        """Read system and system name.
+
+        If the file contains multiple structures, only the last configuration
+        will be read by default.
+        """
+        self.sys = read(self.system, **kwargs)
+        self.sysname = pathlib.Path(self.system).stem
+
+    def get_calculator(self, read_kwargs: dict = {}, **kwargs) -> None:
+        """Configure calculator and attach to system.
+
+        Parameters
+        ----------
+        read_kwargs : dict
+            kwargs to pass to ase.io.read. Default is {}.
+        """
+        calculator = choose_calculator(
+            architecture=self.architecture,
+            device=self.device,
+            **kwargs,
+        )
+        if self.sys is None:
+            self.read_system(**read_kwargs)
+
+        if isinstance(self.sys, list):
+            for sys in self.sys:
+                sys.calc = calculator
+        else:
+            self.sys.calc = calculator
+
+    def _get_potential_energy(self) -> float | list[float]:
+        """Calculate potential energy using MLIP.
+
+        Returns
+        -------
+        potential_energy : float | list[float]
+            Potential energy of system(s).
+        """
+        if isinstance(self.sys, list):
+            energies = []
+            for sys in self.sys:
+                energies.append(sys.get_potential_energy())
+            return energies
+
+        return self.sys.get_potential_energy()
+
+    def _get_forces(self) -> ndarray | list[ndarray]:
+        """Calculate forces using MLIP.
+
+        Returns
+        -------
+        forces : ndarray | list[ndarray]
+            Forces of system(s).
+        """
+        if isinstance(self.sys, list):
+            forces = []
+            for sys in self.sys:
+                forces.append(sys.get_forces())
+            return forces
+
+        return self.sys.get_forces()
+
+    def _get_stress(self) -> ndarray | list[ndarray]:
+        """Calculate stress using MLIP.
+
+        Returns
+        -------
+        stress : ndarray | list[ndarray]
+            Stress of system(s).
+        """
+        if isinstance(self.sys, list):
+            stress = []
+            for sys in self.sys:
+                stress.append(sys.get_stress())
+            return stress
+
+        return self.sys.get_stress()
+
+    def run_single_point(self, properties: str | list[str] | None = None) -> dict:
+        """Run single point calculations.
+
+        Parameters
+        ----------
+        properties : str | List[str] | None
+            Physical properties to calculate. If not specified, "energy",
+            "forces", and "stress" will be returned.
+
+        Returns
+        -------
+        results : dict
+            Dictionary of calculated results.
+        """
+        results = {}
+        if properties is None:
+            properties = []
+        if isinstance(properties, str):
+            properties = [properties]
+
+        if "energy" in properties or len(properties) == 0:
+            results["energy"] = self._get_potential_energy()
+        if "forces" in properties or len(properties) == 0:
+            results["forces"] = self._get_forces()
+        if "stress" in properties or len(properties) == 0:
+            results["stress"] = self._get_stress()
+
+        return results