From 36c3671aa16fc5c4deec8dd0fefe56a0dd4bcad0 Mon Sep 17 00:00:00 2001 From: Sebastian Ehlert Date: Sun, 1 Dec 2024 16:48:12 +0100 Subject: [PATCH] Cleanup Python API for dftd3.library --- python/dftd3/interface.py | 139 ++++++++++--------- python/dftd3/library.py | 273 ++++++++++++++++++++++++++------------ python/dftd3/meson.build | 1 + python/dftd3/py.typed | 0 4 files changed, 265 insertions(+), 148 deletions(-) create mode 100644 python/dftd3/py.typed diff --git a/python/dftd3/interface.py b/python/dftd3/interface.py index a6e8e250..8c606062 100644 --- a/python/dftd3/interface.py +++ b/python/dftd3/interface.py @@ -22,7 +22,7 @@ library in actual workflows than the low-level access provided in the CFFI generated wrappers. """ -from typing import List, Optional, Union +from typing import Optional import numpy as np from . import library @@ -58,7 +58,7 @@ class Structure: on invalid input, like incorrect shape / type of the passed arrays """ - _mol = library.ffi.NULL + _mol = library.StructureHandle.null() def __init__( self, @@ -104,10 +104,10 @@ def __init__( self._mol = library.new_structure( self._natoms, - _cast("int*", _numbers), - _cast("double*", _positions), - _cast("double*", _lattice), - _cast("bool*", _periodic), + _numbers, + _positions, + _lattice, + _periodic, ) def __len__(self): @@ -147,8 +147,8 @@ def update( library.update_structure( self._mol, - _cast("double*", _positions), - _cast("double*", _lattice), + _positions, + _lattice, ) @@ -176,7 +176,7 @@ class DampingParam: of the object should use the second method. """ - _param = library.ffi.NULL + _param = library.ParamHandle.null() def __init__(self, **kwargs): """Create new damping parameter from method name or explicit data""" @@ -190,11 +190,11 @@ def __init__(self, **kwargs): self._param = self.new_param(**kwargs) @staticmethod - def load_param(method, **kwargs): + def load_param(method, **kwargs) -> library.ParamHandle: raise NotImplementedError("Child class has to define parameter loading") @staticmethod - def new_param(**kwargs): + def new_param(**kwargs) -> library.ParamHandle: raise NotImplementedError("Child class has to define parameter construction") @@ -218,15 +218,22 @@ def __init__(self, **kwargs): DampingParam.__init__(self, **kwargs) @staticmethod - def load_param(method, atm=True): - _method = library.ffi.new("char[]", method.encode()) + def load_param(method: str, atm: bool = True) -> library.ParamHandle: return library.load_rational_damping( - _method, + method, atm, ) @staticmethod - def new_param(*, s6=1.0, s8, s9=1.0, a1, a2, alp=14.0): + def new_param( + *, + s6: float = 1.0, + s8: float, + s9: float = 1.0, + a1: float, + a2: float, + alp: float = 14.0, + ) -> library.ParamHandle: return library.new_rational_damping( s6, s8, @@ -254,15 +261,22 @@ def __init__(self, **kwargs): DampingParam.__init__(self, **kwargs) @staticmethod - def load_param(method, atm=True): - _method = library.ffi.new("char[]", method.encode()) + def load_param(method: str, atm: bool = True) -> library.ParamHandle: return library.load_zero_damping( - _method, + method, atm, ) @staticmethod - def new_param(*, s6=1.0, s8, s9=1.0, rs6, rs8=1.0, alp=14.0): + def new_param( + *, + s6: float = 1.0, + s8: float, + s9: float = 1.0, + rs6: float, + rs8: float = 1.0, + alp: float = 14.0, + ) -> library.ParamHandle: return library.new_zero_damping( s6, s8, @@ -289,15 +303,22 @@ def __init__(self, **kwargs): DampingParam.__init__(self, **kwargs) @staticmethod - def load_param(method, atm=True): - _method = library.ffi.new("char[]", method.encode()) + def load_param(method: str, atm: bool = True) -> library.ParamHandle: return library.load_mrational_damping( - _method, + method, atm, ) @staticmethod - def new_param(*, s6=1.0, s8, s9=1.0, a1, a2, alp=14.0): + def new_param( + *, + s6: float = 1.0, + s8: float, + s9: float = 1.0, + a1: float, + a2: float, + alp: float = 14.0, + ) -> library.ParamHandle: return library.new_mrational_damping( s6, s8, @@ -327,15 +348,23 @@ def __init__(self, **kwargs): DampingParam.__init__(self, **kwargs) @staticmethod - def load_param(method, atm=True): - _method = library.ffi.new("char[]", method.encode()) + def load_param(method: str, atm: bool = True) -> library.ParamHandle: return library.load_mzero_damping( - _method, + method, atm, ) @staticmethod - def new_param(*, s6=1.0, s8, s9=1.0, rs6, rs8=1.0, alp=14.0, bet): + def new_param( + *, + s6: float = 1.0, + s8: float, + s9: float = 1.0, + rs6: float, + rs8: float = 1.0, + alp: float = 14.0, + bet: float, + ) -> library.ParamHandle: return library.new_mzero_damping( s6, s8, @@ -364,15 +393,23 @@ def __init__(self, **kwargs): DampingParam.__init__(self, **kwargs) @staticmethod - def load_param(method, atm=True): - _method = library.ffi.new("char[]", method.encode()) + def load_param(method: str, atm: bool = True) -> library.ParamHandle: return library.load_optimizedpower_damping( - _method, + method, atm, ) @staticmethod - def new_param(*, s6=1.0, s8, s9=1.0, a1, a2, alp=14.0, bet): + def new_param( + *, + s6: float = 1.0, + s8: float, + s9: float = 1.0, + a1: float, + a2: float, + alp: float = 14.0, + bet, + ) -> library.ParamHandle: return library.new_optimizedpower_damping( s6, s8, @@ -394,7 +431,7 @@ class DispersionModel(Structure): input. """ - _disp = library.ffi.NULL + _disp = library.ModelHandle.null() def __init__( self, @@ -427,9 +464,9 @@ def get_dispersion(self, param: DampingParam, grad: bool) -> dict: self._mol, self._disp, param._param, - _cast("double*", _energy), - _cast("double*", _gradient), - _cast("double*", _sigma), + _energy, + _gradient, + _sigma, ) results = dict(energy=_energy) @@ -449,8 +486,8 @@ def get_pairwise_dispersion(self, param: DampingParam) -> dict: self._mol, self._disp, param._param, - _cast("double*", _pair_disp2), - _cast("double*", _pair_disp3), + _pair_disp2, + _pair_disp3, ) return { @@ -469,7 +506,7 @@ class GeometricCounterpoise(Structure): superposition error (BSSE). """ - _gcp = library.ffi.NULL + _gcp = library.GCPHandle.null() def __init__( self, @@ -482,14 +519,12 @@ def __init__( ): Structure.__init__(self, numbers, positions, lattice, periodic) - _method = library.ffi.new("char[]", method.encode()) if method else library.ffi.NULL - _basis = library.ffi.new("char[]", basis.encode()) if basis else library.ffi.NULL - self._gcp = library.load_gcp_param(self._mol, _method, _basis) + self._gcp = library.load_gcp_param(self._mol, method, basis) def set_realspace_cutoff(self, bas: float, srb: float): """Set realspace cutoff for evaluation of interactions""" - library.set_gcp_realspace_cutoff(self._disp, bas, srb) + library.set_gcp_realspace_cutoff(self._gcp, bas, srb) def get_counterpoise(self, grad: bool) -> dict: """Evaluate the counterpoise corrected interaction energy""" @@ -502,7 +537,7 @@ def get_counterpoise(self, grad: bool) -> dict: _gradient = None _sigma = None - library.get_counterpoise(self._mol, self._gcp, _cast("double*", _energy), _cast("double*", _gradient), _cast("double*", _sigma)) + library.get_counterpoise(self._mol, self._gcp, _energy, _gradient, _sigma) results = dict(energy=_energy) if _gradient is not None: @@ -512,24 +547,6 @@ def get_counterpoise(self, grad: bool) -> dict: return results -def _cast(ctype, array): - """Cast a numpy array to a FFI pointer""" - return ( - library.ffi.NULL - if array is None - else library.ffi.cast(ctype, array.ctypes.data) - ) - - -def _ref(ctype, value): - """Create a reference to a value""" - if value is None: - return library.ffi.NULL - ref = library.ffi.new(ctype + "*") - ref[0] = value - return ref - - def _rename_kwargs(kwargs, old_name, new_name): if old_name in kwargs and new_name not in kwargs: kwargs[new_name] = kwargs[old_name] diff --git a/python/dftd3/library.py b/python/dftd3/library.py index 8f780726..0f97bbdf 100644 --- a/python/dftd3/library.py +++ b/python/dftd3/library.py @@ -20,6 +20,9 @@ """ import functools +from typing import Optional + +import numpy as np try: from ._libdftd3 import ffi, lib @@ -43,6 +46,59 @@ def get_api_version() -> str: ) +class Handle: + def __init__(self, handle): + self.handle = handle + + @classmethod + def with_gc(cls, handle): + return cls(ffi.gc(handle, cls._delete)) + + @classmethod + def null(cls): + return cls(ffi.NULL) + + @staticmethod + def _delete(handle): + raise NotImplementedError("Delete function not implemented") + + +class StructureHandle(Handle): + @staticmethod + def _delete(handle): + """Delete a DFT-D3 molecular structure data object""" + ptr = ffi.new("dftd3_structure *") + ptr[0] = handle + lib.dftd3_delete_structure(ptr) + + +class ModelHandle(Handle): + @staticmethod + def _delete(handle): + """Delete a DFT-D3 dispersion model object""" + ptr = ffi.new("dftd3_model *") + ptr[0] = handle + lib.dftd3_delete_model(ptr) + + +class ParamHandle(Handle): + @staticmethod + def _delete(handle): + """Delete a DFT-D3 damping parameteter object""" + ptr = ffi.new("dftd3_param *") + ptr[0] = handle + lib.dftd3_delete_param(ptr) + + +class GCPHandle(Handle): + @staticmethod + def _delete(handle): + """Delete a counter-poise parameter object""" + ptr = ffi.new("dftd3_gcp *") + ptr[0] = handle + lib.dftd3_delete_gcp(ptr) + + def _delete_error(mol): """Delete a DFT-D3 error handle object""" ptr = ffi.new("dftd3_error *") @@ -72,148 +128,191 @@ def handle_error(*args, **kwargs): return handle_error -def _delete_structure(mol): - """Delete a DFT-D3 molecular structure data object""" - ptr = ffi.new("dftd3_structure *") - ptr[0] = mol - lib.dftd3_delete_structure(ptr) - - -def new_structure(natoms, numbers, positions, lattice, periodic): +def new_structure( + natoms: int, + numbers: np.ndarray, + positions: np.ndarray, + lattice: Optional[np.ndarray], + periodic: Optional[np.ndarray], +) -> StructureHandle: """Create new molecular structure data""" - return ffi.gc( + return StructureHandle.with_gc( error_check(lib.dftd3_new_structure)( natoms, - numbers, - positions, - lattice, - periodic, - ), - _delete_structure, + _cast("int*", numbers), + _cast("double*", positions), + _cast("double*", lattice), + _cast("bool*", periodic), + ) ) -def _delete_model(disp): - """Delete a DFT-D3 dispersion model object""" - ptr = ffi.new("dftd3_model *") - ptr[0] = disp - lib.dftd3_delete_model(ptr) - - -def new_d3_model(mol): +def new_d3_model(mol: StructureHandle) -> ModelHandle: """Create new D3 dispersion model""" - return ffi.gc(error_check(lib.dftd3_new_d3_model)(mol), _delete_model) - + return ModelHandle.with_gc(error_check(lib.dftd3_new_d3_model)(mol.handle)) -set_model_realspace_cutoff = error_check(lib.dftd3_set_model_realspace_cutoff) - -def _delete_param(param): - """Delete a DFT-D3 damping parameteter object""" - ptr = ffi.new("dftd3_param *") - ptr[0] = param - lib.dftd3_delete_param(ptr) +def set_model_realspace_cutoff( + disp: ModelHandle, disp2: float, disp3: float, cn: float +) -> None: + """Set the realspace cutoff for the dispersion model""" + return error_check(lib.dftd3_set_model_realspace_cutoff)( + disp.handle, disp2, disp3, cn + ) -def new_zero_damping(s6, s8, s9, rs6, rs8, alp): +def new_zero_damping( + s6: float, s8: float, s9: float, rs6: float, rs8: float, alp: float +) -> ParamHandle: """Create new zero damping parameters""" - return ffi.gc( - error_check(lib.dftd3_new_zero_damping)(s6, s8, s9, rs6, rs8, alp), - _delete_param, + return ParamHandle.with_gc( + error_check(lib.dftd3_new_zero_damping)(s6, s8, s9, rs6, rs8, alp) ) -def load_zero_damping(method, atm): +def load_zero_damping(method: str, atm: bool) -> ParamHandle: """Load zero damping parameters from internal storage""" - return ffi.gc(error_check(lib.dftd3_load_zero_damping)(method, atm), _delete_param) + return ParamHandle.with_gc(error_check(lib.dftd3_load_zero_damping)(_char(method), atm)) -def new_rational_damping(s6, s8, s9, a1, a2, alp): +def new_rational_damping( + s6: float, s8: float, s9: float, a1: float, a2: float, alp: float +) -> ParamHandle: """Create new rational damping parameters""" - return ffi.gc( - error_check(lib.dftd3_new_rational_damping)(s6, s8, s9, a1, a2, alp), - _delete_param, + return ParamHandle.with_gc( + error_check(lib.dftd3_new_rational_damping)(s6, s8, s9, a1, a2, alp) ) -def load_rational_damping(method, atm): +def load_rational_damping(method: str, atm: bool) -> ParamHandle: """Load rational damping parameters from internal storage""" - return ffi.gc( - error_check(lib.dftd3_load_rational_damping)(method, atm), _delete_param - ) + return ParamHandle.with_gc(error_check(lib.dftd3_load_rational_damping)(_char(method), atm)) -def new_mzero_damping(s6, s8, s9, rs6, rs8, alp, bet): +def new_mzero_damping( + s6: float, s8: float, s9: float, rs6: float, rs8: float, alp: float, bet: float +) -> ParamHandle: """Create new modified zero damping parameters""" - return ffi.gc( - error_check(lib.dftd3_new_mzero_damping)(s6, s8, s9, rs6, rs8, alp, bet), - _delete_param, + return ParamHandle.with_gc( + error_check(lib.dftd3_new_mzero_damping)(s6, s8, s9, rs6, rs8, alp, bet) ) -def load_mzero_damping(method, atm): +def load_mzero_damping(method: str, atm: bool) -> ParamHandle: """Load modified zero damping parameters from internal storage""" - return ffi.gc(error_check(lib.dftd3_load_mzero_damping)(method, atm), _delete_param) + return ParamHandle.with_gc(error_check(lib.dftd3_load_mzero_damping)(_char(method), atm)) -def new_mrational_damping(s6, s8, s9, a1, a2, alp): +def new_mrational_damping( + s6: float, s8: float, s9: float, a1: float, a2: float, alp: float +) -> ParamHandle: """Create new modified rational damping parameters""" - return ffi.gc( - error_check(lib.dftd3_new_mrational_damping)(s6, s8, s9, a1, a2, alp), - _delete_param, + return ParamHandle.with_gc( + error_check(lib.dftd3_new_mrational_damping)(s6, s8, s9, a1, a2, alp) ) -def load_mrational_damping(method, atm): +def load_mrational_damping(method: str, atm: bool) -> ParamHandle: """Load modified rational damping parameters from internal storage""" - return ffi.gc( - error_check(lib.dftd3_load_mrational_damping)(method, atm), _delete_param + return ParamHandle.with_gc( + error_check(lib.dftd3_load_mrational_damping)(_char(method), atm) ) -def new_optimizedpower_damping(s6, s8, s9, a1, a2, alp, bet): +def new_optimizedpower_damping( + s6: float, s8: float, s9: float, a1: float, a2: float, alp: float, bet: float +) -> ParamHandle: """Create new optimized power damping parameters""" - return ffi.gc( - error_check(lib.dftd3_new_optimizedpower_damping)(s6, s8, s9, a1, a2, alp, bet), - _delete_param, + return ParamHandle.with_gc( + error_check(lib.dftd3_new_optimizedpower_damping)(s6, s8, s9, a1, a2, alp, bet) ) -def load_optimizedpower_damping(method, atm): +def load_optimizedpower_damping(method: str, atm: bool) -> ParamHandle: """Load optimized power damping parameters from internal storage""" - return ffi.gc( - error_check(lib.dftd3_load_optimizedpower_damping)(method, atm), _delete_param + return ParamHandle.with_gc( + error_check(lib.dftd3_load_optimizedpower_damping)(_char(method), atm) ) -update_structure = error_check(lib.dftd3_update_structure) -get_dispersion = error_check(lib.dftd3_get_dispersion) -get_pairwise_dispersion = error_check(lib.dftd3_get_pairwise_dispersion) +def update_structure( + mol: StructureHandle, positions: np.ndarray, lattice: Optional[np.ndarray] +) -> None: + """Update the molecular structure data""" + return error_check(lib.dftd3_update_structure)( + mol.handle, + _cast("double*", positions), + _cast("double*", lattice), + ) -def _delete_gcp(gcp): - """Delete a geometric counter-poise object""" - ptr = ffi.new("dftd3_gcp *") - ptr[0] = gcp - lib.dftd3_delete_gcp(ptr) +def get_dispersion( + mol: StructureHandle, + disp: ModelHandle, + param: ParamHandle, + energy: np.ndarray, + gradient: Optional[np.ndarray], + sigma: Optional[np.ndarray], +) -> None: + return error_check(lib.dftd3_get_dispersion)( + mol.handle, + disp.handle, + param.handle, + _cast("double*", energy), + _cast("double*", gradient), + _cast("double*", sigma), + ) + + +def get_pairwise_dispersion( + mol: StructureHandle, + disp: ModelHandle, + param: ParamHandle, + energy2: np.ndarray, + energy3: np.ndarray, +) -> None: + return error_check(lib.dftd3_get_pairwise_dispersion)( + mol.handle, + disp.handle, + param.handle, + _cast("double*", energy2), + _cast("double*", energy3), + ) -def load_gcp_param(mol, method: str, basis: str): +def load_gcp_param(mol, method: Optional[str], basis: Optional[str]) -> GCPHandle: """Load GCP parameters from internal storage""" - return ffi.gc( - error_check(lib.dftd3_load_gcp_param)(mol, method.encode(), basis.encode()), - _delete_gcp, + return GCPHandle.with_gc( + error_check(lib.dftd3_load_gcp_param)(mol, _char(method), _char(basis)) + ) + + +def set_gcp_realspace_cutoff(gcp: GCPHandle, bas: float, srb: float) -> None: + error_check(lib.dftd3_set_gcp_realspace_cutoff)(gcp.handle, bas, srb) + + +def get_counterpoise( + mol: StructureHandle, + gcp: GCPHandle, + energy: np.ndarray, + gradient: Optional[np.ndarray], + sigma: Optional[np.ndarray], +) -> None: + """Get the counterpoise energy""" + return error_check(lib.dftd3_get_counterpoise_energy)( + mol.handle, + gcp.handle, + _cast("double*", energy), + _cast("double*", gradient), + _cast("double*", sigma), ) -set_gcp_realspace_cutoff = error_check(lib.dftd3_set_gcp_realspace_cutoff) -get_counterpoise = error_check(lib.dftd3_get_counterpoise) +def _char(value: Optional[str]): + """Convert a string to a C char array""" + return ffi.new("char[]", value.encode()) if value is not None else ffi.NULL -def _ref(ctype, value): - """Create a reference to a value""" - if value is None: - return ffi.NULL - ref = ffi.new(ctype + "*") - ref[0] = value - return ref +def _cast(ctype: str, array: Optional[np.ndarray]): + """Cast a numpy array to a FFI pointer""" + return ffi.cast(ctype, array.ctypes.data) if array is not None else ffi.NULL diff --git a/python/dftd3/meson.build b/python/dftd3/meson.build index 5c9334c4..a5be746c 100644 --- a/python/dftd3/meson.build +++ b/python/dftd3/meson.build @@ -82,6 +82,7 @@ pysrcs = files( 'test_parameters.py', 'test_pyscf.py', 'test_qcschema.py', + 'py.typed', ) fs = import('fs') if fs.exists('parameters.toml') diff --git a/python/dftd3/py.typed b/python/dftd3/py.typed new file mode 100644 index 00000000..e69de29b