ENH: Added the new ARMC job.lattice option (#226)

`job.lattice` allows one to specify one or more CP2K .cell files for the purpose of performing calculations on periodic systems.

FOX/armc/armc.py

@@ -28,6 +28,7 @@
 from nanoutils import Literal
 
 from .monte_carlo import MonteCarloABC
+from ..classes.multi_mol import MultiMolecule
 from ..type_hints import ArrayLikeOrScalar, ArrayOrScalar
 from ..io.hdf5_utils import (
     create_hdf5, to_hdf5, create_xyz_hdf5, _get_filename_xyz, hdf5_clear_status
@@ -37,9 +38,8 @@
     import functools
     from .phi_updater import PhiUpdater
     from ..io.read_psf import PSFContainer
-    from ..classes import MultiMolecule
 else:
-    from ..type_alias import MultiMolecule, PhiUpdater, PSFContainer
+    from ..type_alias import PhiUpdater, PSFContainer
 
 __all__ = ['ARMC']
 
@@ -135,6 +135,18 @@ def to_yaml_dict(  # type: ignore[override]
         ret['job'] = self.package_manager.to_yaml_dict()
         ret['job']['molecule'] = [m.properties.filename for m in self.molecule]
 
+        if self.molecule[0].lattice is not None:
+            ret['job']['lattice'] = lat_lst = []
+            iterator1 = (
+                (m, MultiMolecule(m.lattice, atoms={"Xx": range(3)})) for m in self.molecule
+            )
+            for m, lat_mol in iterator1:
+                filename = m.properties.filename.replace(".xyz", ".lattice.xyz")
+                lat_mol.as_xyz(filename)
+                lat_lst.append(filename)
+        else:
+            ret['job']['lattice'] = None
+
         ret['monte_carlo'] = {
             'type': f'{cls.__module__}.{cls.__name__}',
             'iter_len': self.iter_len,
@@ -147,8 +159,8 @@ def to_yaml_dict(  # type: ignore[override]
         }
 
         # Parse the `pes` and `pes_validation` blocks
-        iterator = ((k, getattr(self, k)) for k in ('pes', 'pes_validation'))
-        for name, attr in iterator:
+        iterator2 = ((k, getattr(self, k)) for k in ('pes', 'pes_validation'))
+        for name, attr in iterator2:
             pes: Dict[str, Any] = {}
             ret[name] = pes
             i_max = 1 + max((int(k.split('.')[-1]) for k in attr.keys()), default=0)

FOX/armc/armc_pt.py

@@ -283,7 +283,7 @@ def to_yaml_dict(  # type: ignore[override]
 
         # Convert lists of settings into just settings
         for k, v in ret['job'].items():
-            if k in {'molecule', 'type'}:
+            if k in {'molecule', 'type', 'lattice'}:
                 continue
             v['settings'] = v['settings'][0]
         return ret

FOX/armc/monte_carlo.py

@@ -159,6 +159,14 @@ def __eq__(self, value: Any) -> bool:
         if not np.allclose(self.molecule, value.molecule):
             return False
 
+        if self.molecule[0].lattice is not None:
+            lat0 = np.array([m.lattice for m in self.molecule], dtype=np.float64)
+            lat1 = np.array([m.lattice for m in value.molecule], dtype=np.float64)
+            if not np.allclose(lat0, lat1):
+                return False
+        elif value.molecule[0].lattice is not None:
+            return False
+
         iterator1 = ((v, value.pes[k]) for k, v in self.pes.items())
         for p1, p2 in iterator1:  # type: partial, partial  # type: ignore
             if p1.func != p2.func or p1.keywords != p2.keywords:

FOX/armc/sanitization.py

@@ -277,11 +277,16 @@ def get_package(dct: JobMapping, phi: Iterable) -> Tuple[PackageManager, Tuple[M
 
     """
     _sub_pkg_dict: Dict[str, SubJobMapping] = split_dict(  # type: ignore[call-overload]
-        dct, preserve_order=True, keep_keys={'type', 'molecule'}
+        dct, preserve_order=True, keep_keys={'type', 'molecule', 'lattice'}
     )
 
     job_dict = validate_job(dct)
     mol_list = [mol.as_Molecule(mol_subset=0)[0] for mol in job_dict['molecule']]
+    if job_dict['lattice'] is not None:
+        if len(job_dict['molecule']) != len(job_dict['lattice']):
+            raise ValueError("`job.molecule` and `job.lattice` must be of equal length")
+        for m, lat in zip(job_dict['molecule'], job_dict['lattice']):
+            m.lattice = lat
 
     data: Dict[str, List[PkgDict]] = {}
     for k, v in _sub_pkg_dict.items():

FOX/armc/schemas.py

@@ -12,12 +12,15 @@
 
 """
 
+from __future__ import annotations
+
 import os
 from functools import partial
 from collections import abc
 from typing import (
     Any, SupportsInt, Type, Mapping, Collection, Sequence, SupportsFloat, Generator,
-    Callable, Union, Optional, Tuple, FrozenSet, Iterable, Dict, TypeVar, List
+    Callable, Union, Optional, Tuple, FrozenSet, Iterable, Dict, TypeVar, List,
+    cast, TYPE_CHECKING,
 )
 
 import numpy as np
@@ -27,7 +30,7 @@
 from qmflows.packages import Package
 from nanoutils import (
     Default, Formatter, supports_int, supports_float, isinstance_factory,
-    issubclass_factory, import_factory, as_nd_array, TypedDict,
+    issubclass_factory, import_factory, as_nd_array, TypedDict, Literal,
 )
 
 from .armc import ARMC
@@ -38,6 +41,12 @@
 from ..type_hints import ArrayLike, Scalar, ArrayLikeOrScalar
 from ..classes import MultiMolecule
 from ..utils import get_move_range
+from ..io.cp2k import lattice_from_cell
+from ..io.read_xyz import read_multi_xyz
+
+if TYPE_CHECKING:
+    SCT = TypeVar("SCT", bound=np.generic)
+    NDArray = np.ndarray[Any, np.dtype[SCT]]
 
 __all__ = [
     'validate_phi', 'validate_monte_carlo', 'validate_psf', 'validate_pes',
@@ -84,12 +93,30 @@
 A_TARGET.setflags(write=False)
 
 
-def abspath(file: Union[str, bytes, os.PathLike]) -> str:
+def abspath(file: str | bytes | os.PathLike[Any]) -> str:
     """Return the absolute path of **file**."""
     file_str = os.fsdecode(file)
     return os.path.abspath(file_str)
 
 
+CellFunc = Callable[[Union[str, bytes, "os.PathLike[Any]"]], "NDArray[np.float64]"]
+EXT_DICT: Dict[Literal["cell", "xyz"], CellFunc] = {
+    "cell": lattice_from_cell,
+    "xyz": lambda n: read_multi_xyz(n)[0],
+}
+
+
+def parse_lattice(file: str | bytes | os.PathLike[Any]) -> NDArray[np.float64]:
+    file_str = os.fsdecode(file)
+    _, ext = os.path.splitext(file_str)
+    ext = cast(Literal['cell'], ext.strip("."))
+    try:
+        func = EXT_DICT[ext]
+    except KeyError:
+        raise ValueError(f"Unsupported file extensions: {ext!r}") from None
+    return func(file_str)
+
+
 #: Schema for validating the ``"phi"`` block.
 phi_schema = Schema({
     Optional_('type', default=lambda: PhiUpdater): Or(
@@ -358,7 +385,18 @@ class PESDict(TypedDict):
             ))
         ),
         error=Formatter(f"'job.molecule' expected one or more .xyz files{EPILOG}")
-    )
+    ),
+
+    Optional_('lattice', default=None): Or(
+        None,
+        And(Or(str, bytes, os.PathLike), Use(lambda n: (parse_lattice(n),))),
+        And(
+            abc.Sequence,
+            lambda n: all(isinstance(i, (str, bytes, os.PathLike)) for i in n),
+            Use(lambda n: tuple(parse_lattice(i) for i in n)),
+        ),
+        error=Formatter(f"'job.lattice' expected one or more .cell files{EPILOG}")
+    ),
 }, name='job_schema', description='Schema for validating the "job" block.')
 
 
@@ -367,13 +405,21 @@ class JobMapping(TypedDict, total=False):
 
     type: Union[None, str, Type[PackageManagerABC]]
     molecule: Union[MultiMolecule, str, os.PathLike, Sequence[Union[MultiMolecule, str, os.PathLike]]]  # noqa: E501
+    lattice: Union[
+        None,
+        str,
+        bytes,
+        os.PathLike[Any],
+        Sequence[str | bytes | os.PathLike[Any]],
+    ]
 
 
 class JobDict(TypedDict):
     """A :class:`~typing.TypedDict` representing the output of :data:`job_schema`."""
 
     type: Type[PackageManagerABC]
     molecule: Tuple[MultiMolecule, ...]
+    lattice: None | Tuple[NDArray[np.float64], ...]
 
 
 def _parse_settings_sequence(seq: Sequence[Mapping]) -> Generator[QmSettings, None, None]:

docs/4_monte_carlo_args.rst

@@ -72,6 +72,7 @@ Index
     =========================================== ===================================================================================================================
     :attr:`job.type`                            The type of package manager.
     :attr:`job.molecule`                        One or more .xyz files with reference (QM) potential energy surfaces.
+    :attr:`job.lattice`                         One or more CP2K .cell files with the lattice vectors of each mol in :attr:`job.molecule`.
     :attr:`job.block.type`                      An instance of a QMFlows :class:`~qmflows.packages.packages.Package`.
     :attr:`job.block.settings`                  The job settings as used by :class:`job.block.type`.
     :attr:`job.block.template`                  A settings template for updating :class:`job.block.settings`.
@@ -596,6 +597,15 @@ Note that these jobs are executed in the order as provided by the user-input.
             one *must* be provided by the user.
 
 
+    .. attribute:: job.lattice
+
+        :Parameter:     * **Type** - :class:`str` or :class:`list[str] <list>`
+
+        One or more CP2K .cell files with the lattice vectors of each mol in :attr:`job.molecule`.
+
+        This option should be specified is one is performing calculations on periodic systems.
+
+
     .. attribute:: job.block.type
 
         :Parameter:     * **Type** - :class:`str` or :class:`qmflows.packages.Package<qmflows.packages.packages.Package>` instance