TYP: Add basic scm.plams annotations

nanoqm/_data.py

@@ -11,6 +11,7 @@
 import warnings
 from typing import Any, TYPE_CHECKING, Protocol, Literal
 
+import qmflows
 from scm.plams import Settings
 from qmflows.common import AtomXYZ
 from qmflows.warnings_qmflows import QMFlowsDeprecationWarning
@@ -136,8 +137,8 @@ class CP2KGeneralSetting(_DataConfig):
     potential_file_name: None | str
     functional_x: None | str
     functional_c: None | str
-    cp2k_settings_main: Settings
-    cp2k_settings_guess: Settings
+    cp2k_settings_main: qmflows.Settings
+    cp2k_settings_guess: qmflows.Settings
     wfn_restart_file_name: None | str
     file_cell_parameters: None | str
     aux_fit: Literal["low", "medium", "good", "verygood", "excellent"]

pyproject.toml

@@ -14,7 +14,6 @@ files = ["nanoqm", "typings"]
 
 [[tool.mypy.overrides]]
 module = [
-    "scm.plams.*",
     "h5py.*",
     "schema.*",
 ]

typings/scm/plams/__init__.pyi

@@ -0,0 +1,16 @@
+from scm.plams.core.settings import Settings
+from scm.plams.core.functions import add_to_class
+from scm.plams.mol.atom import Atom
+from scm.plams.mol.bond import Bond
+from scm.plams.mol.molecule import Molecule
+from scm.plams.mol.pdbtools import PDBHandler, PDBRecord
+
+__all__ = [
+    "Atom",
+    "Bond",
+    "Molecule",
+    "Settings",
+    "PDBHandler",
+    "PDBRecord",
+    "add_to_class",
+]

typings/scm/plams/core/functions.pyi

@@ -0,0 +1,6 @@
+from collections.abc import Callable
+from typing import Any, TypeVar
+
+_FT = TypeVar("_FT", bound=Callable[..., Any])
+
+def add_to_class(classname: type[Any]) -> Callable[[_FT], _FT]: ...

typings/scm/plams/core/settings.pyi

@@ -0,0 +1,49 @@
+import sys
+import types
+from collections.abc import Callable, Iterable, Mapping, Sequence
+from typing import Any, Generic, TypeVar, overload, Protocol
+
+_KT = TypeVar("_KT")
+_VT = TypeVar("_VT")
+_MT = TypeVar("_MT", bound=_SupportsMissing)
+_ST = TypeVar("_ST", bound=Settings[Any, Any])
+
+class _SupportsMissing(Protocol):
+    def __missing__(self, __key: Any) -> Any: ...
+
+class Settings(dict[_KT, _VT]):
+    def copy(self: _ST) -> _ST: ...
+    def soft_update(self: _ST, other: Mapping[_KT, _VT]) -> _ST: ...
+    def update(self, other: Mapping[_KT, _VT]) -> None: ...  # type: ignore[override]
+    def merge(self: _ST, other: Mapping[_KT, _VT]) -> _ST: ...
+    def find_case(self, key: _KT) -> _KT: ...
+    def as_dict(self) -> dict[_KT, _VT]: ...
+    @classmethod
+    def suppress_missing(cls: type[_MT]) -> SuppressMissing[_MT]: ...
+    def get_nested(self, key_tuple: Iterable[Any], suppress_missing: bool = False) -> Any: ...
+    def set_nested(self, key_tuple: Sequence[Any], value: Any, suppress_missing: bool = False) -> None: ...
+    def flatten(self, flatten_list: bool = ...) -> Settings[tuple[Any, ...], Any]: ...
+    def unflatten(self, unflatten_list: bool = ...) -> Settings[Any, Any]: ...
+    @classmethod  # type: ignore[override]
+    @overload
+    def fromkeys(cls, __iterable: Iterable[_KT]) -> Settings[_KT, Any]: ...
+    @classmethod
+    @overload
+    def fromkeys(cls, __iterable: Iterable[_KT], __value: _VT) -> Settings[_KT, _VT]: ...
+    def __missing__(self, __key: _KT) -> Settings[Any, Any]: ...
+    def __getattr__(self, name: _KT) -> _VT: ...  # type: ignore[misc]
+    def __setattr__(self, name: _KT, value: _VT) -> None: ...  # type: ignore[misc,override]
+    def __str__(self) -> str: ...
+    def __repr__(self) -> str: ...
+    def __add__(self: _ST, other: Mapping[_KT, _VT]) -> _ST: ...
+    def __iadd__(self: _ST, other: Mapping[_KT, _VT]) -> _ST: ...
+    def __copy__(self: _ST) -> _ST: ...
+
+class SuppressMissing(Generic[_MT]):
+    obj: _MT
+    missing: Callable[[Any, _MT, Any], Any]
+    def __init__(self, obj: type[_MT]) -> None: ...
+    def __enter__(self) -> None: ...
+    def __exit__(
+        self, exc_type: None | type[BaseException], exc_value: None | BaseException, traceback: None | types.TracebackType
+    ) -> None: ...

typings/scm/plams/mol/atom.pyi

@@ -0,0 +1,76 @@
+import builtins
+from collections.abc import Iterable, Iterator, Mapping, Sequence
+from typing import Any, SupportsFloat
+
+import numpy as np
+from scm.plams import Bond, Molecule, Settings
+
+class Atom:
+    atnum: int
+    mol: Molecule
+    bonds: list[Bond]
+    properties: Settings[Any, Any]
+    coords: tuple[float, float, float]
+    id: int  # Only available after calling `Molecule.set_atoms_id`
+    def __init__(
+        self,
+        atnum: int = ...,
+        symbol: None | builtins.str = ...,
+        coords: None | Iterable[builtins.str | bytes | SupportsFloat] = ...,
+        unit: builtins.str = ...,
+        bonds: None | list[Bond] = ...,
+        mol: None | Molecule = ...,
+        **other: Any,
+    ) -> None: ...
+    def __iter__(self) -> Iterator[float]: ...
+    @property
+    def symbol(self) -> builtins.str: ...
+    @symbol.setter
+    def symbol(self, symbol: builtins.str) -> None: ...
+    @property
+    def x(self) -> float: ...
+    @x.setter
+    def x(self, value: float) -> None: ...
+    @property
+    def y(self) -> float: ...
+    @y.setter
+    def y(self, value: float) -> None: ...
+    @property
+    def z(self) -> float: ...
+    @z.setter
+    def z(self, value: float) -> None: ...
+    @property
+    def mass(self) -> float: ...
+    @property
+    def radius(self) -> float: ...
+    @property
+    def connectors(self) -> float: ...
+    @property
+    def is_metallic(self) -> int: ...
+    @property
+    def is_electronegative(self) -> int: ...
+    def str(
+        self,
+        symbol: bool,
+        suffix: str = ...,
+        suffix_dict: Mapping[str, Any] = ...,
+        unit: str = ...,
+        space: int = ...,
+        decimal: int = ...,
+    ) -> str: ...
+    def translate(self, vector: Iterable[float], unit: builtins.str = ...) -> None: ...
+    def move_to(self, point: Iterable[float], unit: builtins.str = ...) -> None: ...
+    def distance_to(
+        self, point: Iterable[float], unit: builtins.str = ..., result_unit: builtins.str = ...
+    ) -> tuple[float, float, float]: ...
+    def vector_to(self, point: Iterable[float], unit: builtins.str = ..., result_unit: builtins.str = ...) -> float: ...
+    def angle(
+        self,
+        point1: Iterable[float],
+        point2: Iterable[float],
+        point1unit: builtins.str = ...,
+        point2unit: builtins.str = ...,
+        result_unit: builtins.str = ...,
+    ) -> builtins.str: ...
+    def rotate(self, matrix: np.ndarray | Sequence[Sequence[float]]) -> None: ...
+    def neighbors(self) -> list[Atom]: ...

typings/scm/plams/mol/bond.pyi

@@ -0,0 +1,22 @@
+from collections.abc import Generator, Iterable
+from typing import Any, ClassVar
+
+from scm.plams import Atom, Molecule, Settings
+
+class Bond:
+    AR: ClassVar[float]
+    atom1: Atom
+    atom2: Atom
+    order: float
+    mol: Molecule
+    properties: Settings[Any, Any]
+    def __init__(
+        self, atom1: None | Atom = ..., atom2: None | Atom = ..., order: float = ..., mol: None | Molecule = ..., **other: Any
+    ) -> None: ...
+    def __iter__(self) -> Generator[Atom, None, None]: ...
+    def is_aromatic(self) -> bool: ...
+    def length(self, unit: str = ...) -> float: ...
+    def as_vector(self, start: None | Atom = ..., unit: str = ...) -> tuple[float, float, float]: ...
+    def other_end(self, atom: Atom) -> Atom: ...
+    def resize(self, moving_atom: Atom, length: float, unit: str = ...) -> None: ...
+    def rotate(self, moving_atom: Atom, angle: float, unit: str = ...) -> None: ...

typings/scm/plams/mol/molecule.pyi

@@ -0,0 +1,180 @@
+import builtins
+import os
+from collections.abc import Callable, Collection, Iterable, Iterator, Sequence
+from typing import IO, Any, overload, Literal as L, SupportsIndex, TypedDict
+
+import numpy as np
+from scm.plams import Atom, Bond, PDBHandler, Settings
+
+class _AtomDict(TypedDict):
+    atnum: int
+    bonds: list[int]
+    coords: tuple[float, float, float]
+    properties: Settings[Any, Any]
+
+class _BondDict(TypedDict):
+    atom1: int
+    atom2: int
+    order: float
+    properties: Settings[Any, Any]
+
+class _MolDictBase(TypedDict, total=False):
+    charge: int
+
+class _MolDict(_MolDictBase):
+    atoms: list[_AtomDict]
+    bonds: list[_BondDict]
+    lattice: list[Sequence[float]]
+    properties: Settings[Any, Any]
+
+class _LabelDict(TypedDict, total=False):
+    BO: bool
+    RS: bool
+    EZ: bool
+    DH: bool
+    CO: bool
+    twist_tol: float
+    bend_tol: float
+
+class Molecule:
+    atoms: list[Atom]
+    bonds: list[Bond]
+    lattice: list[Sequence[float]]
+    properties: Settings[Any, Any]
+    def __init__(
+        self,
+        filename: None | builtins.str | os.PathLike[builtins.str] = ...,
+        inputformat: L[None, "xyz", "pdb", "mol", "mol2", "rkf"] = ...,
+        positions: None | Sequence[Sequence[float]] | np.ndarray = ...,
+        numbers: None | Collection[int] = ...,
+        lattice: None | Sequence[Sequence[float]] | np.ndarray = ...,
+        **other: Any,
+    ) -> None: ...
+    def copy(self, atoms: None | Collection[Atom] = ...) -> Molecule: ...
+    def add_molecule(self, other: Molecule, copy: bool = ...) -> None: ...
+    def add_atom(self, atom: Atom, adjacent: None | Iterable[Atom] = ...) -> None: ...
+    def delete_atom(self, atom: Atom) -> None: ...
+    @overload
+    def add_bond(self, arg1: Atom, arg2: Atom, order: float = ...) -> None: ...
+    @overload
+    def add_bond(self, arg1: Bond, arg2: None = ..., order: float = ...) -> None: ...
+    @overload
+    def delete_bond(self, arg1: Atom, arg2: Atom) -> None: ...
+    @overload
+    def delete_bond(self, arg1: Bond, arg2: None = ...) -> None: ...
+    def delete_all_bonds(self) -> None: ...
+    def find_bond(self, atom1: Atom, atom2: Atom) -> None | Bond: ...
+    def set_atoms_id(self, start: int = ...) -> None: ...
+    def unset_atoms_id(self) -> None: ...
+    def neighbors(self, atom: Atom) -> list[Atom]: ...
+    def bond_matrix(self) -> np.ndarray: ...
+    def separate(self) -> list[Molecule]: ...
+    def guess_bonds(self, atom_subset: None | Sequence[Atom] = ..., dmax: float = ...) -> None: ...
+    def in_ring(self, arg: Atom | Bond) -> bool: ...
+    def supercell(self, *args: float | Sequence[float]) -> Molecule: ...
+    def cell_lengths(self, unit: builtins.str) -> list[float]: ...
+    def cell_angles(self, unit: builtins.str) -> list[float]: ...
+    def unit_cell_volume(self, unit: builtins.str = ...) -> float: ...
+    def set_integer_bonds(self, action: builtins.str = ..., tolerance: float = ...) -> None: ...
+    @overload
+    def index(self, value: Atom, start: int = ..., stop: None | int = ...) -> int: ...
+    @overload
+    def index(self, value: Bond, start: int = ..., stop: None | int = ...) -> tuple[int, int]: ...
+    @overload
+    def round_coords(self, decimals: int = ..., *, inplace: L[False] = ...) -> Molecule: ...
+    @overload
+    def round_coords(self, decimals: int = ..., *, inplace: L[True]) -> None: ...
+    def translate(self, vector: tuple[float, float, float], unit: builtins.str = ...) -> None: ...
+    def rotate_lattice(self, matrix: Sequence[Sequence[float]] | np.ndarray) -> None: ...
+    def rotate(self, matrix: Sequence[Sequence[float]] | np.ndarray, lattice: bool = ...) -> None: ...
+    def align_lattice(self, convention: L["AMS", "reax"] = ..., zero: float = ...) -> bool: ...
+    def rotate_bond(self, bond: Bond, moving_atom: Atom, angle: float, unit: builtins.str = ...) -> None: ...
+    def resize_bond(self, bond: Bond, moving_atom: Atom, length: float, unit: builtins.str = ...) -> None: ...
+    def closest_atom(self, point: Atom | tuple[float, float, float], unit: builtins.str = ...) -> Atom: ...
+    def distance_to_point(self, point: Sequence[float], unit: builtins.str = ..., result_unit: builtins.str = ...) -> float: ...
+    @overload
+    def distance_to_mol(self, other: Molecule, result_unit: builtins.str = ..., *, return_atoms: L[False] = ...) -> float: ...
+    @overload
+    def distance_to_mol(
+        self, other: Molecule, result_unit: builtins.str = ..., *, return_atoms: L[True]
+    ) -> tuple[float, Atom, Atom]: ...
+    def wrap(
+        self, length: float, angle: float = ..., length_unit: builtins.str = ..., angle_unit: builtins.str = ...
+    ) -> None: ...
+    def get_center_of_mass(self, unit: builtins.str = ...) -> tuple[float, float, float]: ...
+    def get_mass(self) -> float: ...
+    def get_density(self) -> float: ...
+    @overload
+    def get_formula(self, as_dict: L[False] = ...) -> builtins.str: ...
+    @overload
+    def get_formula(self, as_dict: L[True]) -> dict[builtins.str, int]: ...
+    def apply_strain(self, strain: np.ndarray | Sequence[float], voigt_form: bool = ...) -> None: ...
+    def map_to_central_cell(self, around_origin: bool = ...) -> Molecule: ...
+    def perturb_atoms(
+        self, max_displacement: float = ..., unit: builtins.str = ..., atoms: None | Iterable[Atom] = ...
+    ) -> None: ...
+    def perturb_lattice(self, max_displacement: float = ..., unit: builtins.str = ..., ams_convention: bool = ...) -> None: ...
+    def substitute(
+        self,
+        connector: Bond | tuple[Atom, Atom],
+        ligand: Molecule,
+        ligand_connector: Bond | tuple[Atom, Atom],
+        bond_length: None | float = ...,
+        steps: int = ...,
+        cost_func_mol: None | Callable[[Molecule, Molecule], float] = ...,
+        cost_func_array: None | Callable[[np.ndarray, np.ndarray], float] = ...,
+    ) -> None: ...
+    def __len__(self) -> int: ...
+    def __str__(self) -> builtins.str: ...
+    def str(self, decimal: int = ...) -> builtins.str: ...
+    def __iter__(self) -> Iterator[Atom]: ...
+    @overload
+    def __getitem__(self, key: SupportsIndex) -> Atom: ...
+    @overload
+    def __getitem__(self, key: tuple[SupportsIndex, SupportsIndex]) -> Bond: ...
+    def __add__(self, other: Molecule) -> Molecule: ...
+    def __iadd__(self, other: Molecule) -> Molecule: ...
+    def __copy__(self) -> Molecule: ...
+    def __deepcopy__(self, memo: object) -> Molecule: ...
+    def __round__(self, ndigits: None | int = ...) -> Molecule: ...
+    def __getstate__(self) -> _MolDict: ...
+    def __setstate__(self, state: _MolDict) -> None: ...
+    def as_dict(self) -> _MolDict: ...
+    @classmethod
+    def from_dict(cls, dictionary: _MolDict) -> Molecule: ...
+    @classmethod
+    def from_elements(cls, elements: Iterable[builtins.str]) -> Molecule: ...
+    def as_array(self, atom_subset: None | Iterable[Atom] = ...) -> np.ndarray: ...
+    def from_array(
+        self, xyz_array: Molecule | Iterable[tuple[float, float, float]] | np.ndarray, atom_subset: None | Iterable[Atom] = ...
+    ) -> None: ...
+    def __array__(self, dtype: Any = ...) -> np.ndarray: ...
+    def readxyz(self, f: IO[builtins.str], geometry: int = ..., **other: Any) -> None: ...
+    def writexyz(self, f: IO[builtins.str], **other: Any) -> None: ...
+    def readmol(self, f: IO[builtins.str], **other: Any) -> None: ...
+    def writemol(self, f: IO[builtins.str], **other: Any) -> None: ...
+    def readmol2(self, f: IO[builtins.str], **other: Any) -> None: ...
+    def writemol2(self, f: IO[builtins.str], **other: Any) -> None: ...
+    def readpdb(self, f: IO[builtins.str], geometry: int = ..., **other: Any) -> PDBHandler: ...
+    def writepdb(self, f: IO[builtins.str], **other: Any) -> None: ...
+    def read(
+        self, filename: builtins.str | os.PathLike[builtins.str], inputformat: None | builtins.str = ..., **other: Any
+    ) -> None | PDBHandler: ...
+    def write(
+        self, filename: builtins.str | os.PathLike[builtins.str], outputformat: None | builtins.str = ..., **other: Any
+    ) -> None: ...
+    def add_hatoms(self) -> Molecule: ...
+    @overload
+    @staticmethod
+    def rmsd(mol1: Molecule, mol2: Molecule, *, return_rotmat: L[False] = ..., check: bool = ...) -> np.ndarray: ...
+    @overload
+    @staticmethod
+    def rmsd(mol1: Molecule, mol2: Molecule, *, return_rotmat: L[True], check: bool = ...) -> tuple[np.ndarray, np.ndarray]: ...
+    # `label` and `reorder` are defined in `scm.plams.mol.identity`
+    @overload
+    def label(self, level: int = ..., keep_labels: bool = ..., flags: None | _LabelDict = ...) -> builtins.str: ...
+    @overload
+    def label(
+        self, level: Sequence[int], keep_labels: bool = ..., flags: None | _LabelDict = ...
+    ) -> tuple[builtins.str, ...]: ...
+    def reorder(self, other: Molecule) -> None: ...