diff --git a/.gitignore b/.gitignore index 64b07fd..f4f79d9 100644 --- a/.gitignore +++ b/.gitignore @@ -107,3 +107,6 @@ ENV/ # In-tree generated files */_version.py + +# Espaloma junk output +**/.model.pt \ No newline at end of file diff --git a/devtools/conda-envs/release-build.yml b/devtools/conda-envs/release-build.yml index 00e4b7a..f50b3da 100644 --- a/devtools/conda-envs/release-build.yml +++ b/devtools/conda-envs/release-build.yml @@ -25,6 +25,7 @@ dependencies: - openmm - lammps - mdtraj + - pint # for units in case OpenFF is not installed # Molecule building - mbuild @@ -42,7 +43,8 @@ dependencies: # OpenFF stack - openff-toolkit ~=0.16 - openff-interchange >=0.3.28 - - openff-nagl + - openff-nagl >= 0.4 + - openff-nagl-models >= 0.3 # Chemical database queries - cirpy diff --git a/devtools/conda-envs/test-env.yml b/devtools/conda-envs/test-env.yml index 447671e..70395f6 100644 --- a/devtools/conda-envs/test-env.yml +++ b/devtools/conda-envs/test-env.yml @@ -3,12 +3,12 @@ channels: - conda-forge - openeye dependencies: - # Basic Python dependencies + # Basic Python dependencies - python - pip - jupyterlab - # Testing and docs + # Testing and docs - pytest - pytest-cov - codecov @@ -25,6 +25,7 @@ dependencies: - openmm - lammps - mdtraj + - pint # for units in case OpenFF is not installed # Molecule building - mbuild @@ -42,7 +43,8 @@ dependencies: # OpenFF stack - openff-toolkit ~=0.16 - openff-interchange >=0.3.28 - - openff-nagl + - openff-nagl >= 0.4 + - openff-nagl-models >= 0.3 # Chemical database queries - cirpy diff --git a/polymerist/genutils/decorators/functional.py b/polymerist/genutils/decorators/functional.py index 5504c93..a5999f7 100644 --- a/polymerist/genutils/decorators/functional.py +++ b/polymerist/genutils/decorators/functional.py @@ -3,7 +3,10 @@ __author__ = 'Timotej Bernat' __email__ = 'timotej.bernat@colorado.edu' -from typing import Callable, Iterable, Optional, Type, Union +from typing import Callable, Concatenate, Iterable, Iterator, Optional, ParamSpec, TypeVar, Union + +T = TypeVar('T') +Params = ParamSpec('Params') from inspect import signature, Parameter from functools import wraps, partial @@ -13,20 +16,18 @@ from .meta import extend_to_methods from . import signatures -from ..typetools.parametric import T, Args, KWArgs -from ..typetools.categorical import ListLike from ..fileutils.pathutils import aspath, asstrpath @extend_to_methods -def optional_in_place(funct : Callable[[object, Args, KWArgs], None]) -> Callable[[object, Args, bool, KWArgs], Optional[object]]: +def optional_in_place(funct : Callable[[Concatenate[object, Params]], None]) -> Callable[[Concatenate[object, Params]], Optional[object]]: '''Decorator function for allowing in-place (writeable) functions which modify object attributes to be not performed in-place (i.e. read-only), specified by a boolean flag''' # TODO : add assertion that the wrapped function has at least one arg AND that the first arg is of the desired (limited) type old_sig = signature(funct) @wraps(funct) # for preserving docstring and type annotations / signatures - def in_place_wrapper(obj : object, *args : Args, in_place : bool=False, **kwargs : KWArgs) -> Optional[object]: # read-only by default + def in_place_wrapper(obj : object, *args : Params.args, in_place : bool=False, **kwargs : Params.kwargs) -> Optional[object]: # read-only by default '''If not in-place, create a clone on which the method is executed''' # NOTE : old_sig.bind screws up arg passing if in_place: funct(obj, *args, **kwargs) # default call to writeable method - implicitly returns None @@ -54,9 +55,9 @@ def in_place_wrapper(obj : object, *args : Args, in_place : bool=False, **kwargs return in_place_wrapper # TODO : implement support for extend_to_methods (current mechanism is broken by additional deocrator parameters) -def flexible_listlike_input(funct : Callable[[ListLike], T]=None, CastType : Type[ListLike]=list, valid_member_types : Union[Type, tuple[Type]]=object) -> Callable[[Iterable], T]: +def flexible_listlike_input(funct : Callable[[Iterator], T]=None, CastType : type[Iterator]=list, valid_member_types : Union[type, tuple[type]]=object) -> Callable[[Iterable], T]: '''Wrapper which allows a function which expects a single list-initializable, Container-like object to accept any Iterable (or even star-unpacked arguments)''' - if not issubclass(CastType, ListLike): + if not issubclass(CastType, Iterator): raise TypeError(f'Cannot wrap listlike input with non-listlike type "{CastType.__name__}"') @wraps(funct) @@ -79,13 +80,13 @@ def wrapper(*args) -> T: # wrapper which accepts an arbitrary number of non-keyw return wrapper @extend_to_methods -def allow_string_paths(funct : Callable[[Path, Args, KWArgs], T]) -> Callable[[Union[Path, str], Args, KWArgs], T]: +def allow_string_paths(funct : Callable[[Concatenate[Path, Params]], T]) -> Callable[[Concatenate[Union[Path, str], Params]], T]: '''Modifies a function which expects a Path as its first argument to also accept string-paths''' # TODO : add assertion that the wrapped function has at least one arg AND that the first arg is of the desired (limited) type old_sig = signature(funct) # lookup old type signature @wraps(funct) # for preserving docstring and type annotations / signatures - def str_path_wrapper(flex_path : Union[str, Path], *args : Args, **kwargs : KWArgs) -> T: + def str_path_wrapper(flex_path : Union[str, Path], *args : Params.args, **kwargs : Params.kwargs) -> T: '''First converts stringy paths into normal Paths, then executes the original function''' return funct(aspath(flex_path), *args, **kwargs) @@ -99,13 +100,13 @@ def str_path_wrapper(flex_path : Union[str, Path], *args : Args, **kwargs : KWAr return str_path_wrapper @extend_to_methods -def allow_pathlib_paths(funct : Callable[[str, Args, KWArgs], T]) -> Callable[[Union[Path, str], Args, KWArgs], T]: +def allow_pathlib_paths(funct : Callable[[Concatenate[str, Params]], T]) -> Callable[[Concatenate[Union[Path, str], Params]], T]: '''Modifies a function which expects a string path as its first argument to also accept canonical pathlib Paths''' # TODO : add assertion that the wrapped function has at least one arg AND that the first arg is of the desired (limited) type old_sig = signature(funct) # lookup old type signature @wraps(funct) # for preserving docstring and type annotations / signatures - def str_path_wrapper(flex_path : Union[str, Path], *args : Args, **kwargs : KWArgs) -> T: + def str_path_wrapper(flex_path : Union[str, Path], *args : Params.args, **kwargs : Params.kwargs) -> T: '''First converts normal Paths into stringy paths, then executes the original function''' return funct(asstrpath(flex_path), *args, **kwargs) diff --git a/polymerist/genutils/decorators/meta.py b/polymerist/genutils/decorators/meta.py index 082dfd6..bedfc0e 100644 --- a/polymerist/genutils/decorators/meta.py +++ b/polymerist/genutils/decorators/meta.py @@ -6,8 +6,9 @@ from typing import Concatenate, Callable, ParamSpec, TypeAlias, TypeVar from functools import update_wrapper, wraps -from ..typetools.parametric import C, O, P, R, Args, KWArgs -Decorator : TypeAlias = Callable[[Callable[P, R]], Callable[P, R]] +Params = ParamSpec('Params') # can also use to typehint *args and **kwargs +ReturnType = TypeVar('ReturnType') +Decorator : TypeAlias = Callable[[Callable[Params, ReturnType]], Callable[Params, ReturnType]] # META DECORATORS @@ -18,16 +19,16 @@ def extend_to_methods(dec : Decorator) -> Decorator: @wraps(dec, updated=()) # transfer decorator signature to decorator adapter class, without updating the __dict__ field class AdaptedDecorator: - def __init__(self, funct : Callable[P, R]) -> None: + def __init__(self, funct : Callable[Params, ReturnType]) -> None: '''Record function''' self.funct = funct update_wrapper(self, funct) # equivalent to functools.wraps, transfers docstring, module, etc. for documentation - def __call__(self, *args : Args, **kwargs : KWArgs) -> ReturnSignature: # TODO : fix this to reflect the decorator's return signature + def __call__(self, *args : Params.args, **kwargs : Params.kwargs) -> ReturnSignature: # TODO : fix this to reflect the decorator's return signature '''Apply decorator to function, then call decorated function''' return dec(self.funct)(*args, **kwargs) - def __get__(self, instance : O, owner : C) -> Callable[[Concatenate[O, P]], R]: + def __get__(self, instance : object, owner : type) -> Callable[[Concatenate[object, Params]], ReturnType]: '''Generate partial application with calling instance as first argument (fills in for "self")''' method = self.funct.__get__(instance, owner) # look up method belonging to owner class return dec(method) # return the decorated method diff --git a/polymerist/genutils/importutils/dependencies.py b/polymerist/genutils/importutils/dependencies.py new file mode 100644 index 0000000..4787afe --- /dev/null +++ b/polymerist/genutils/importutils/dependencies.py @@ -0,0 +1,93 @@ +'''Utilities for checking and enforcing module dependencies within code''' + +__author__ = 'Timotej Bernat' +__email__ = 'timotej.bernat@colorado.edu' + +from typing import Callable, ParamSpec, TypeVar + +Params = ParamSpec('Params') +ReturnType = TypeVar('ReturnType') +TCall = Callable[Params, ReturnType] # generic function of callable class + +# from importlib import import_module +from importlib.util import find_spec +from functools import wraps + + +def module_installed(module_name : str) -> bool: + ''' + Check whether a module of the given name is present on the system + + Parameters + ---------- + module_name : str + The name of the module, as it would occur in an import statement + Do not support direct passing of module objects to avoid circularity + (i.e. no reason to check if a module is present if one has already imported it elsewhere) + + Returns + ------- + module_found : bool + Whether or not the module was found to be installed in the current working environment + ''' + # try: + # package = import_module(module_name) + # except ModuleNotFoundError: + # return False + # else: + # return True + + try: # NOTE: opted for this implementation, as it never actually imports the package in question (faster and fewer side-effects) + return find_spec(module_name) is not None + except (ValueError, AttributeError, ModuleNotFoundError): # these could all be raised by + return False + +def modules_installed(*module_names : list[str]) -> bool: + ''' + Check whether one or more modules are all present + Will only return true if ALL specified modules are found + + Parameters + ---------- + module_names : *str + Any number of module names, passed as a comma-separated sequence of strings + + Returns + ------- + all_modules_found : bool + Whether or not all modules were found to be installed in the current working environment + ''' + return all(module_installed(module_name) for module_name in module_names) + +def requires_modules( + *required_module_names : list[str], + missing_module_error : type[Exception]=ImportError, + ) -> Callable[[TCall[..., ReturnType]], TCall[..., ReturnType]]: + ''' + Decorator which enforces optional module dependencies prior to function execution + + Parameters + ---------- + module_names : *str + Any number of module names, passed as a comma-separated sequence of strings + missing_module_error : type[Exception], default ImportError + The type of Exception to raise if a module is not found installed + Defaults to ImportError + + Raises + ------ + ImportError : Exception + Raised if any of the specified packages is not found to be installed + Exception message will indicate the name of the specific package found missing + ''' + def decorator(func) -> TCall[..., ReturnType]: + @wraps(func) + def req_wrapper(*args : Params.args, **kwargs : Params.kwargs) -> ReturnType: + for module_name in required_module_names: + if not module_installed(module_name): + raise missing_module_error(f'No installation found for module "{module_name}"') + else: + return func(*args, **kwargs) + + return req_wrapper + return decorator \ No newline at end of file diff --git a/polymerist/mdtools/openfftools/__init__.py b/polymerist/mdtools/openfftools/__init__.py index 5e86394..8f1ab65 100644 --- a/polymerist/mdtools/openfftools/__init__.py +++ b/polymerist/mdtools/openfftools/__init__.py @@ -1,60 +1,34 @@ -'''Tools for manipulating and extending OpenFF objects, and for interfacing with other tools and formats''' +'''Extensions, interfaces, and convenience methods built around the functionality in the OpenFF software stack''' __author__ = 'Timotej Bernat' __email__ = 'timotej.bernat@colorado.edu' -from typing import Any -from pathlib import Path - -import openforcefields -from openff.toolkit import ToolkitRegistry -from openff.toolkit import GLOBAL_TOOLKIT_REGISTRY as GTR -from openff.toolkit.utils.base_wrapper import ToolkitWrapper -from openff.toolkit.utils.utils import all_subclasses -from openff.toolkit.utils.exceptions import LicenseError, ToolkitUnavailableException -from openff.toolkit.typing.engines.smirnoff.forcefield import _get_installed_offxml_dir_paths - -from openff.toolkit.utils.openeye_wrapper import OpenEyeToolkitWrapper -from espaloma_charge.openff_wrapper import EspalomaChargeToolkitWrapper -from openff.nagl.toolkits import NAGLRDKitToolkitWrapper, NAGLOpenEyeToolkitWrapper - - -# FORCE FIELD AND ToolkitWrapper REFERENCE -FFDIR = Path(openforcefields.get_forcefield_dirs_paths()[0]) # Locate path where OpenFF forcefields are installed -FF_DIR_REGISTRY : dict[Path, Path] = {} -FF_PATH_REGISTRY : dict[Path, Path] = {} -for ffdir_str in _get_installed_offxml_dir_paths(): - ffdir = Path(ffdir_str) - ffdir_name = ffdir.parent.stem - - FF_DIR_REGISTRY[ ffdir_name] = ffdir - FF_PATH_REGISTRY[ffdir_name] = [path for path in ffdir.glob('*.offxml')] - -# CHECKING FOR OpenEye -ALL_IMPORTABLE_TKWRAPPERS = all_subclasses(ToolkitWrapper) # References to every registered ToolkitWrapper and ToolkitRegistry -try: - _ = OpenEyeToolkitWrapper() - _OE_TKWRAPPER_IS_AVAILABLE = True - OEUnavailableException = None -except (LicenseError, ToolkitUnavailableException) as error: - _OE_TKWRAPPER_IS_AVAILABLE = False - OEUnavailableException = error # catch and record relevant error message for use (rather than trying to replicate it elsewhere) - -# Register OpenFF-compatible GNN ToolkitWrappers -GTR.register_toolkit(EspalomaChargeToolkitWrapper) -GTR.register_toolkit(NAGLRDKitToolkitWrapper) -if _OE_TKWRAPPER_IS_AVAILABLE: - GTR.register_toolkit(NAGLOpenEyeToolkitWrapper) - - -# GENERATE LOOKUP DICTS FOR EVERY REGISTERED ToolkitWrappers and ToolkitRegistry -REGISTERED_TKWRAPPER_TYPES = [type(tkwrapper) for tkwrapper in GTR.registered_toolkits] -TKWRAPPERS = { # NOTE : this must be done AFTER any new registrations to thr GTR (e.g. after registering GNN ToolkitWrappers) - tk_wrap.toolkit_name : tk_wrap - for tk_wrap in GTR.registered_toolkits -} -TKREGS = {} # individually register toolkit wrappers for cases where a registry must be passed -for tk_name, tk_wrap in TKWRAPPERS.items(): - tk_reg = ToolkitRegistry() - tk_reg.register_toolkit(tk_wrap) - TKREGS[tk_name] = tk_reg \ No newline at end of file +# Subpackage-wide precheck to see if OpenFF is even usable in the first place +from ...genutils.importutils.dependencies import modules_installed +if not modules_installed('openff', 'openff.toolkit'): + raise ModuleNotFoundError( + f''' + OpenFF packages which are required to utilitize {__name__} not found in current environment + Please follow installation instructions at https://docs.openforcefield.org/projects/toolkit/en/stable/installation.html, then retry import + ''' + ) + +# Import of toplevel OpenFF object registries +from ._forcefields import ( + FFDIR, + FF_DIR_REGISTRY, + FF_PATH_REGISTRY, +) +from ._toolkits import ( + ## toolkit registries + GLOBAL_TOOLKIT_REGISTRY, GTR, + POLYMERIST_TOOLKIT_REGISTRY, + ## catalogues of available toolkit wrappers + ALL_IMPORTABLE_TKWRAPPERS, + ALL_AVAILABLE_TKWRAPPERS, + TKWRAPPERS, + TKWRAPPER_TYPES, + ## registry of partial charge methods by + CHARGE_METHODS_BY_TOOLKIT, + TOOLKITS_BY_CHARGE_METHOD, +) \ No newline at end of file diff --git a/polymerist/mdtools/openfftools/_forcefields.py b/polymerist/mdtools/openfftools/_forcefields.py new file mode 100644 index 0000000..58199fe --- /dev/null +++ b/polymerist/mdtools/openfftools/_forcefields.py @@ -0,0 +1,29 @@ +'''For dynamically determining and cataloging which SMIRNOFF-copatible force fields are installed (and accompanying functionality) are available''' + +__author__ = 'Timotej Bernat' +__email__ = 'timotej.bernat@colorado.edu' + +from typing import Optional +from pathlib import Path + +from ...genutils.importutils.dependencies import modules_installed + + +# Force field and ToolkitWrapper reference +FFDIR : Optional[Path] = None +if modules_installed('openff.toolkit'): + from openforcefields import get_forcefield_dirs_paths + + FFDIR = Path(get_forcefield_dirs_paths()[0]) # Locate path where OpenFF forcefields are installed + +FF_DIR_REGISTRY : dict[Path, Path] = {} +FF_PATH_REGISTRY : dict[Path, Path] = {} +if modules_installed('openforcefields'): + from openff.toolkit.typing.engines.smirnoff.forcefield import _get_installed_offxml_dir_paths + + for ffdir_str in _get_installed_offxml_dir_paths(): + ffdir = Path(ffdir_str) + ffdir_name = ffdir.parent.stem + + FF_DIR_REGISTRY[ ffdir_name] = ffdir + FF_PATH_REGISTRY[ffdir_name] = [path for path in ffdir.glob('*.offxml')] \ No newline at end of file diff --git a/polymerist/mdtools/openfftools/_toolkits.py b/polymerist/mdtools/openfftools/_toolkits.py new file mode 100644 index 0000000..9e6bb40 --- /dev/null +++ b/polymerist/mdtools/openfftools/_toolkits.py @@ -0,0 +1,118 @@ +'''For dynamically determining and cataloging which ToolkitWrappers (and accompanying functionality) are available''' + +__author__ = 'Timotej Bernat' +__email__ = 'timotej.bernat@colorado.edu' + +# Core OpenFF toolkit component registration +from typing import Union +from collections import defaultdict + +from openff.toolkit.utils.utils import all_subclasses +from openff.toolkit.utils.base_wrapper import ToolkitWrapper +from openff.toolkit.utils.toolkit_registry import ToolkitRegistry +from openff.toolkit.utils.toolkits import ( + OPENEYE_AVAILABLE, + RDKIT_AVAILABLE, + AMBERTOOLS_AVAILABLE, + GLOBAL_TOOLKIT_REGISTRY, +) +GTR = GLOBAL_TOOLKIT_REGISTRY # alias for brevity + +from ...genutils.importutils.dependencies import modules_installed +_REGISTER_TOOLKITS_TO_GLOBAL : bool = True # TODO: find way to avoid setting this config parameter directly in code + + +# Helper functions +def toolkit_wrapper_is_registered(toolkit_wrapper : Union[ToolkitWrapper, type[ToolkitWrapper]], toolkit_registry : ToolkitRegistry=GTR) -> bool: + '''Check whether a ToolkitRegistry instance has already registered a given ToolkitWrapper subclass''' + if not isinstance(toolkit_wrapper, type): # ToolkitWrapper TYPES are needed for this check; any instances therefore... + toolkit_wrapper = type(toolkit_wrapper) # ...will have their respective types extracted + if not issubclass(toolkit_wrapper, ToolkitWrapper): + raise TypeError(f'Expected a ToolkitWrapper instance or subclass, instead received object of type {toolkit_wrapper.__name__}') + + return any(isinstance(tkwrapper, toolkit_wrapper) for tkwrapper in toolkit_registry.registered_toolkits) + +# Setup of initial containers for OpenFF module info +ALL_IMPORTABLE_TKWRAPPERS : list[type[ToolkitWrapper]] = all_subclasses(ToolkitWrapper) # NOTE: just because you can import these, doesn't necessarily mean they can be instantiated +ALL_AVAILABLE_TKWRAPPERS : list[type[ToolkitWrapper]] = [] +CHARGE_METHODS_BY_TOOLKIT : dict[type[ToolkitWrapper], list[str]] = defaultdict(list) + +# Toolkit-specific registrations which depend on available packages +## BuiltIn (not particularly useful in and of itself, but nice to know it's accessible) +if modules_installed('openff.toolkit'): # this check is idempotent to initial OpenFF check, but is nice to have for consistency between all ToolkitWrappers below + from openff.toolkit.utils.builtin_wrapper import BuiltInToolkitWrapper + + ALL_AVAILABLE_TKWRAPPERS.append(BuiltInToolkitWrapper) + CHARGE_METHODS_BY_TOOLKIT[BuiltInToolkitWrapper] = [ + charge_method + for charge_method in BuiltInToolkitWrapper._supported_charge_methods + ] +## RDKit +if modules_installed('rdkit') and RDKIT_AVAILABLE: + from openff.toolkit.utils.rdkit_wrapper import RDKitToolkitWrapper + + ALL_AVAILABLE_TKWRAPPERS.append(RDKitToolkitWrapper) + CHARGE_METHODS_BY_TOOLKIT[RDKitToolkitWrapper] = [ + charge_method + for charge_method in RDKitToolkitWrapper._supported_charge_methods + ] +## Ambertools +if modules_installed('pdb4amber') and AMBERTOOLS_AVAILABLE: # turns out "ambertools" can't actually be imported as a module, need to check for peripheral modules which are better behaved instead + from openff.toolkit.utils.ambertools_wrapper import AmberToolsToolkitWrapper + + ALL_AVAILABLE_TKWRAPPERS.append(AmberToolsToolkitWrapper) + CHARGE_METHODS_BY_TOOLKIT[AmberToolsToolkitWrapper] = [ + charge_method + for charge_method in AmberToolsToolkitWrapper._supported_charge_methods + ] +## OpenEye +if modules_installed('openeye.oechem', 'openeye.oeomega') and OPENEYE_AVAILABLE: + from openff.toolkit.utils.openeye_wrapper import OpenEyeToolkitWrapper + + ALL_AVAILABLE_TKWRAPPERS.append(OpenEyeToolkitWrapper) + CHARGE_METHODS_BY_TOOLKIT[OpenEyeToolkitWrapper] = [ + charge_method + for charge_method in OpenEyeToolkitWrapper._supported_charge_methods + ] +## NAGL - extracting available charge methods is a little different for GNN toolkits +if modules_installed('openff.nagl', 'openff.nagl_models'): + from openff.toolkit.utils.nagl_wrapper import NAGLToolkitWrapper + + ALL_AVAILABLE_TKWRAPPERS.append(NAGLToolkitWrapper) + CHARGE_METHODS_BY_TOOLKIT[NAGLToolkitWrapper] = [ + model_path.name # need to extract dynamically from Paths + for model_path in NAGLToolkitWrapper.list_available_nagl_models() + ] +## Espaloma - extracting available charge methods is a little different for GNN toolkits +if modules_installed('espaloma_charge'): + from espaloma_charge.openff_wrapper import EspalomaChargeToolkitWrapper + + ALL_AVAILABLE_TKWRAPPERS.append(EspalomaChargeToolkitWrapper) + CHARGE_METHODS_BY_TOOLKIT[EspalomaChargeToolkitWrapper] = [ + 'espaloma-am1bcc' + ] # this is, at this of writing, the only available method for EspalomaCharge and unfortunately not accessible dynamically + +# Post-registration info compilation +## Compiling name-based lookups for available ToolkitWrappers and registering all to a local +POLYMERIST_TOOLKIT_REGISTRY = ToolkitRegistry() # retain a local registry separate from GLOBAL_TOOLKIT_REGISTRY +TKWRAPPERS : dict[str, ToolkitWrapper ] = {} +TKWRAPPER_TYPES : dict[str, type[ToolkitWrapper]] = {} + +for tkwrapper_type in ALL_AVAILABLE_TKWRAPPERS: + tkwrapper_instance = tkwrapper_type() # instantiate toolkit wrapper class + POLYMERIST_TOOLKIT_REGISTRY.register_toolkit(tkwrapper_instance) + # if requested, also mirror all found toolkits to the Global ToolkitRegistry + if _REGISTER_TOOLKITS_TO_GLOBAL and not toolkit_wrapper_is_registered(tkwrapper_type, GLOBAL_TOOLKIT_REGISTRY): # make registration idempotent + GLOBAL_TOOLKIT_REGISTRY.register_toolkit(tkwrapper_instance) + + # register to name-based lookup dict + TKWRAPPERS[ tkwrapper_type._toolkit_name] = tkwrapper_instance + TKWRAPPER_TYPES[tkwrapper_type._toolkit_name] = tkwrapper_type + +## Compiling registry of which partial charge methods are supported by which toolkits +TOOLKITS_BY_CHARGE_METHOD : dict[str, list[type[ToolkitWrapper]]] = defaultdict(list) # also compile inverse mapping (compiled once available toolkits are known) +for tkwrapper_type, supported_methods in CHARGE_METHODS_BY_TOOLKIT.items(): + if (tkwrapper_type in ALL_AVAILABLE_TKWRAPPERS): # exclude non-registered toolkits to avoid confusion + for method in supported_methods: + TOOLKITS_BY_CHARGE_METHOD[method].append(tkwrapper_type) +TOOLKITS_BY_CHARGE_METHOD = dict(TOOLKITS_BY_CHARGE_METHOD) # convert to pure dict for typing purposes diff --git a/polymerist/mdtools/openfftools/boxvectors.py b/polymerist/mdtools/openfftools/boxvectors.py index 011a810..49ea028 100644 --- a/polymerist/mdtools/openfftools/boxvectors.py +++ b/polymerist/mdtools/openfftools/boxvectors.py @@ -14,7 +14,7 @@ from openff.toolkit import Topology from openff.interchange.components._packmol import _box_vectors_are_in_reduced_form -from ...unitutils.interop import allow_openmm_units, openff_to_openmm +from .omminter.unitsys import allow_openmm_units, openff_to_openmm # CUSTOM TYPES FOR CLARITY, ESPECIALLY WITH UNITS diff --git a/polymerist/mdtools/openfftools/omminter/__init__.py b/polymerist/mdtools/openfftools/omminter/__init__.py new file mode 100644 index 0000000..db12755 --- /dev/null +++ b/polymerist/mdtools/openfftools/omminter/__init__.py @@ -0,0 +1,12 @@ +'''For interfacing between OpenFF and OpenMM''' + +__author__ = 'Timotej Bernat' +__email__ = 'timotej.bernat@colorado.edu' + +from .mdobjects import forcefield_flexible, openff_topology_to_openmm +from .unitsys import ( + openmm_to_openff, + openff_to_openmm, + allow_openmm_units, + allow_openff_units, +) \ No newline at end of file diff --git a/polymerist/mdtools/openfftools/omminter.py b/polymerist/mdtools/openfftools/omminter/mdobjects.py similarity index 89% rename from polymerist/mdtools/openfftools/omminter.py rename to polymerist/mdtools/openfftools/omminter/mdobjects.py index 199b239..2ed2be9 100644 --- a/polymerist/mdtools/openfftools/omminter.py +++ b/polymerist/mdtools/openfftools/omminter/mdobjects.py @@ -1,4 +1,4 @@ -'''For interfacing between OpenFF and OpenMM representations, along with the file analogues''' +'''For interfacing between OpenFF and OpenMM representations of Topologies and other MD primitives''' __author__ = 'Timotej Bernat' __email__ = 'timotej.bernat@colorado.edu' @@ -6,7 +6,6 @@ from typing import Optional, Union from pathlib import Path -from numpy import ndarray from openff.toolkit import ForceField from openff.interchange import Interchange @@ -16,12 +15,12 @@ from openmm.app import Topology as OMMTopology from openmm.unit import Quantity -from . import FFDIR -from .boxvectors import box_vectors_flexible, VectorQuantity, BoxVectorsQuantity -from ...unitutils.interop import openff_to_openmm +from .unitsys import openff_to_openmm +from .. import FFDIR +from ..boxvectors import box_vectors_flexible, VectorQuantity, BoxVectorsQuantity -def forcefield_flexible(forcefield : Union[ForceField, str, Path]) -> ForceField: +def forcefield_flexible(forcefield : Union[ForceField, str, Path]) -> ForceField: # DEV: consider deprecating '''For making forcefield input to other functions more flexible (can accept a literal ForceField, a string name, or a Path to the forcefield)''' if isinstance(forcefield, ForceField): return forcefield diff --git a/polymerist/unitutils/interop.py b/polymerist/mdtools/openfftools/omminter/unitsys.py similarity index 57% rename from polymerist/unitutils/interop.py rename to polymerist/mdtools/openfftools/omminter/unitsys.py index b36b7d3..ae82d1e 100644 --- a/polymerist/unitutils/interop.py +++ b/polymerist/mdtools/openfftools/omminter/unitsys.py @@ -1,40 +1,40 @@ -'''Decorators for handling interconversion between the OpenMM and OpenFF (Pint) unit engines''' +'''For handling interconversion between the OpenMM and OpenFF (Pint) unit engines''' __author__ = 'Timotej Bernat' __email__ = 'timotej.bernat@colorado.edu' from typing import Callable, TypeVar -R = TypeVar('R') # for representing generic return values -Q = TypeVar('Q') # for representing generic Quantity-like objects -from openmm.unit import Quantity +ReturnType = TypeVar('ReturnType') + +from openmm.unit import Quantity as OpenMMQuantity from pint import Quantity as PintQuantity # this is also the base class for all OpenFF-style units -from openff.units import Quantity as OFFQuantity + from openff.units.openmm import ( from_openmm as openmm_to_openff, to_openmm as openff_to_openmm, ) -def allow_openmm_units(funct : Callable[[Q], R]) -> Callable[[Q], R]: - '''Allow a Callable which expects ALL of its args to be OpenFF Quantities to also accept equivalent OpenMM Quantites''' - def wrapper(*args, **kwargs) -> R: +def allow_openmm_units(funct : Callable[..., ReturnType]) -> Callable[..., ReturnType]: + '''Allow a Callable which expects any of its args to be OpenFF Quantities to also accept equivalent OpenMM Quantites''' + def wrapper(*args, **kwargs) -> ReturnType: new_args = [ - openmm_to_openff(arg) if isinstance(arg, Quantity) else arg + openmm_to_openff(arg) if isinstance(arg, OpenMMQuantity) else arg for arg in args ] new_kwargs = { - key : openmm_to_openff(kwarg) if isinstance(kwarg, Quantity) else kwarg + key : openmm_to_openff(kwarg) if isinstance(kwarg, OpenMMQuantity) else kwarg for key, kwarg in kwargs.items() } return funct(*new_args, **new_kwargs) return wrapper -def allow_openff_units(funct : Callable[[Q], R]) -> Callable[[Q], R]: - '''Allow a Callable which expects ALL of its args to be OpenMM Quantities to also accept equivalent OpenFF Quantites''' - def wrapper(*args, **kwargs) -> R: +def allow_openff_units(funct : Callable[..., ReturnType]) -> Callable[..., ReturnType]: + '''Allow a Callable which expects any of its args to be OpenMM Quantities to also accept equivalent OpenFF Quantites''' + def wrapper(*args, **kwargs) -> ReturnType: new_args = [ openff_to_openmm(arg) if isinstance(arg, PintQuantity) else arg for arg in args diff --git a/polymerist/mdtools/openfftools/partialcharge/__init__.py b/polymerist/mdtools/openfftools/partialcharge/__init__.py index c5612d8..9c9c334 100644 --- a/polymerist/mdtools/openfftools/partialcharge/__init__.py +++ b/polymerist/mdtools/openfftools/partialcharge/__init__.py @@ -2,3 +2,8 @@ __author__ = 'Timotej Bernat' __email__ = 'timotej.bernat@colorado.edu' + +from .._toolkits import ( + CHARGE_METHODS_BY_TOOLKIT, + TOOLKITS_BY_CHARGE_METHOD, +) \ No newline at end of file diff --git a/polymerist/mdtools/openfftools/partialcharge/chargemethods.py b/polymerist/mdtools/openfftools/partialcharge/chargemethods.py deleted file mode 100644 index 81b3cae..0000000 --- a/polymerist/mdtools/openfftools/partialcharge/chargemethods.py +++ /dev/null @@ -1,61 +0,0 @@ -'''Registry module for keeping track of which partial charging toolkit registries and related methods are available''' - -__author__ = 'Timotej Bernat' -__email__ = 'timotej.bernat@colorado.edu' - -from typing import Type -from collections import defaultdict - -from openff.toolkit.utils.base_wrapper import ToolkitWrapper -from openff.toolkit.utils.rdkit_wrapper import RDKitToolkitWrapper -from openff.toolkit.utils.builtin_wrapper import BuiltInToolkitWrapper -from openff.toolkit.utils.openeye_wrapper import OpenEyeToolkitWrapper -from openff.toolkit.utils.ambertools_wrapper import AmberToolsToolkitWrapper - -from openff import nagl_models -from openff.nagl import GNNModel -from espaloma_charge.openff_wrapper import EspalomaChargeToolkitWrapper - -from .. import REGISTERED_TKWRAPPER_TYPES - - -# REFERENCE MAPPING BETWEEN PARTIAL CHARGE METHODS AND SUPPORTING TOOLKIT WRAPPERS -SUPPORTED_PARTIAL_CHARGE_METHODS_BY_TOOLKIT : dict[Type[ToolkitWrapper], list[str]]= { # TOSELF : this unfortunately cannot be accessed dynamically as a attribute of each TollkitWrapper class - BuiltInToolkitWrapper : [ - 'zeros', - 'formal_charge', - ], - RDKitToolkitWrapper : [ - 'mmff94', - 'gasteiger' - ], - AmberToolsToolkitWrapper : [ - 'am1bcc', - 'am1-mulliken', - 'gasteiger', - ], - OpenEyeToolkitWrapper : [ - 'am1bcc', - 'am1-mulliken', - 'gasteiger', - 'mmff94', - 'am1bccnosymspt', - 'am1elf10', - 'am1bccelf10', - ], - EspalomaChargeToolkitWrapper : [ - 'espaloma-am1bcc' - ] -} - -TOOLKITS_BY_CHARGE_METHOD : dict[str, list[Type[ToolkitWrapper]]] = defaultdict(list) -for tkwrapper_type, supported_methods in SUPPORTED_PARTIAL_CHARGE_METHODS_BY_TOOLKIT.items(): - if (tkwrapper_type in REGISTERED_TKWRAPPER_TYPES): # exclude non-registered toolkits to avoid confusion - for method in supported_methods: - TOOLKITS_BY_CHARGE_METHOD[method].append(tkwrapper_type) - - -## NAGL GNN Model -NAGL_MODEL_PATH = nagl_models.list_available_nagl_models()[1] # Path(/home/timber/miniconda3/envs/polymerist-env/lib/python3.11/site-packages/openff/nagl_models/models/openff-gnn-am1bcc-0.1.0-rc.1.pt) -NAGL_MODEL_PATH = nagl_models.validate_nagl_model_path(NAGL_MODEL_PATH) # double check that this model path is still one of the valid entry point -NAGL_MODEL = GNNModel.load(NAGL_MODEL_PATH) \ No newline at end of file diff --git a/polymerist/mdtools/openfftools/partialcharge/molchargers.py b/polymerist/mdtools/openfftools/partialcharge/molchargers.py index cb8d68e..71a8192 100644 --- a/polymerist/mdtools/openfftools/partialcharge/molchargers.py +++ b/polymerist/mdtools/openfftools/partialcharge/molchargers.py @@ -6,18 +6,16 @@ import logging LOGGER = logging.getLogger(__name__) -from typing import Any, ClassVar, Union -from abc import ABC, abstractmethod, abstractproperty +from typing import Union +from abc import ABC, abstractmethod from rdkit import Chem -from openff.units import unit as offunit from openff.toolkit.topology.molecule import Molecule -from openff.toolkit.utils.exceptions import ToolkitUnavailableException # TODO : use chargemethods.TOOLKITS_BY_CHARGE_METHOD to automatically determine whether/which toolkits are available for each method +from openff.toolkit.utils.exceptions import ToolkitUnavailableException -from .. import TKREGS, _OE_TKWRAPPER_IS_AVAILABLE, OEUnavailableException -from .chargemethods import NAGL_MODEL -from ....genutils.decorators.classmod import register_subclasses +from ....genutils.importutils.dependencies import requires_modules from ....genutils.decorators.functional import optional_in_place +from ....genutils.decorators.classmod import register_subclasses, register_abstract_class_attrs def has_partial_charges(mol : Union[Molecule, Chem.Mol]) -> bool: @@ -32,14 +30,9 @@ def has_partial_charges(mol : Union[Molecule, Chem.Mol]) -> bool: # ABSTRACT AND CONCRETE CLASSES FOR CHARGING MOLECULES @register_subclasses(key_attr='CHARGING_METHOD') +@register_abstract_class_attrs('CHARGING_METHOD') class MolCharger(ABC): '''Base interface for defining various methods of generating and storing atomic partial charges''' - @abstractproperty - @classmethod - def CHARGING_METHOD(cls): - '''For setting the name of the method as a class attribute in child classes''' - pass - @abstractmethod @optional_in_place def _charge_molecule(self, uncharged_mol : Molecule) -> None: @@ -55,29 +48,38 @@ def charge_molecule(self, uncharged_mol : Molecule) -> None: LOGGER.info(f'Successfully assigned "{self.CHARGING_METHOD}" charges') # CONCRETE IMPLEMENTATIONS OF DIFFERENT CHARGING METHODS -class ABE10Charger(MolCharger): +class ABE10Charger(MolCharger, CHARGING_METHOD= 'AM1-BCC-ELF10'): '''Charger class for AM1-BCC-ELF10 exact charging''' - CHARGING_METHOD : ClassVar[str] = 'AM1-BCC-ELF10' - + @requires_modules('openeye.oechem', 'openeye.oeomega', missing_module_error=ToolkitUnavailableException) # for whatever weird reason the toplevel openeye package has no module spec, so just checking "openeye" isn't enough @optional_in_place def _charge_molecule(self, uncharged_mol : Molecule) -> None: - if not _OE_TKWRAPPER_IS_AVAILABLE: - raise OEUnavailableException # AM1-BCC-ELF10 is exclusively available thru OpenEye; if it is not present, then, must err - uncharged_mol.assign_partial_charges(partial_charge_method='am1bccelf10', toolkit_registry=TKREGS['OpenEye Toolkit']) # TODO : provide support for AMBER / RDKit if OE license is unavailable + from openff.toolkit.utils.openeye_wrapper import OpenEyeToolkitWrapper + + uncharged_mol.assign_partial_charges( + partial_charge_method='am1bccelf10', + toolkit_registry=OpenEyeToolkitWrapper(), # instance init will raise exception if license or OpenEye packages are missing + ) # TODO : find decent alternative if OpenEye license is missing (AmberTools doesn't do ELF10 and doesn't work on Windows) -class EspalomaCharger(MolCharger): +class EspalomaCharger(MolCharger, CHARGING_METHOD='Espaloma-AM1-BCC'): '''Charger class for EspalomaCharge charging''' - CHARGING_METHOD : ClassVar[str] = 'Espaloma-AM1-BCC' - + @requires_modules('espaloma_charge', missing_module_error=ToolkitUnavailableException) @optional_in_place def _charge_molecule(self, uncharged_mol : Molecule) -> None: - uncharged_mol.assign_partial_charges(partial_charge_method='espaloma-am1bcc', toolkit_registry=TKREGS['Espaloma Charge Toolkit']) + from espaloma_charge.openff_wrapper import EspalomaChargeToolkitWrapper + + uncharged_mol.assign_partial_charges( + partial_charge_method='espaloma-am1bcc', # this is actually the ONLY charge method the EspalomaChargeToolkitWrapper supports + toolkit_registry=EspalomaChargeToolkitWrapper(), + ) -class NAGLCharger(MolCharger): +class NAGLCharger(MolCharger, CHARGING_METHOD='NAGL'): '''Charger class for NAGL charging''' - CHARGING_METHOD : ClassVar[str] = 'NAGL' - + @requires_modules('openff.nagl', 'openff.nagl_models', missing_module_error=ToolkitUnavailableException) @optional_in_place def _charge_molecule(self, uncharged_mol : Molecule) -> None: - nagl_charges = NAGL_MODEL.compute_property(uncharged_mol, check_domains=True, error_if_unsupported=True) - uncharged_mol.partial_charges = nagl_charges * offunit.elementary_charge # need to have OpenFF-style units attached to set "partial_charges" property + from openff.toolkit.utils.nagl_wrapper import NAGLToolkitWrapper + + uncharged_mol.assign_partial_charges( + partial_charge_method='openff-gnn-am1bcc-0.1.0-rc.3.pt', # 'openff-gnn-am1bcc-0.1.0-rc.2.pt', + toolkit_registry=NAGLToolkitWrapper(), + ) diff --git a/polymerist/mdtools/openfftools/partialcharge/rescharge/interface.py b/polymerist/mdtools/openfftools/partialcharge/rescharge/interface.py index dd168e7..2f2810a 100644 --- a/polymerist/mdtools/openfftools/partialcharge/rescharge/interface.py +++ b/polymerist/mdtools/openfftools/partialcharge/rescharge/interface.py @@ -3,7 +3,6 @@ __author__ = 'Timotej Bernat' __email__ = 'timotej.bernat@colorado.edu' -from typing import ClassVar from dataclasses import dataclass from openff.toolkit import Molecule @@ -15,10 +14,9 @@ @dataclass -class LibraryCharger(MolCharger): +class LibraryCharger(MolCharger, CHARGING_METHOD='RCT'): '''Charger class for applying library charges onto residue-mapped Molecules''' charges_by_residue : ChargesByResidue - CHARGING_METHOD : ClassVar[str] = 'RCT' @optional_in_place def _charge_molecule(self, uncharged_mol : Molecule) -> None: diff --git a/polymerist/mdtools/openfftools/solvation/physprops.py b/polymerist/mdtools/openfftools/solvation/physprops.py index 5db9a28..1de70a2 100644 --- a/polymerist/mdtools/openfftools/solvation/physprops.py +++ b/polymerist/mdtools/openfftools/solvation/physprops.py @@ -15,7 +15,7 @@ from openff.units import Quantity as OFFQuantity from ....unitutils.dimensions import is_volume -from ....unitutils.interop import allow_openff_units, openff_to_openmm +from ..omminter.unitsys import allow_openff_units, openff_to_openmm # MASS @@ -50,6 +50,7 @@ def number_density(density : Quantity, MW : Quantity) -> Quantity: return (density / MW) * AVOGADRO_CONSTANT_NA # NUMBER +@allow_openff_units def num_mols_in_box(mol : Union[Mol, Molecule, Topology], box_vol : Quantity, density : Quantity) -> int: '''Return the number of particles/molecules needed to fill a box of given volume to the specified density''' assert(is_volume(box_vol.unit)) diff --git a/polymerist/tests/genutils/importutils/__init__.py b/polymerist/tests/genutils/importutils/__init__.py new file mode 100644 index 0000000..c08e2b6 --- /dev/null +++ b/polymerist/tests/genutils/importutils/__init__.py @@ -0,0 +1,4 @@ +'''Unit tests for `importutils` package''' + +__author__ = 'Timotej Bernat' +__email__ = 'timotej.bernat@colorado.edu' diff --git a/polymerist/tests/genutils/importutils/test_dependencies.py b/polymerist/tests/genutils/importutils/test_dependencies.py new file mode 100644 index 0000000..eea4ee1 --- /dev/null +++ b/polymerist/tests/genutils/importutils/test_dependencies.py @@ -0,0 +1,46 @@ +'''Unit tests for `dependencies` package''' + +__author__ = 'Timotej Bernat' +__email__ = 'timotej.bernat@colorado.edu' + +import pytest + +from typing import Any, Callable +from polymerist.genutils.importutils import dependencies + + +# Testing module finding +@pytest.mark.parametrize( + 'module_names, expected_found', [ + (['polymerist'], True), # we'd better hope the parent module is present if we're running tests on it :P + (['sys'], True), # test stdlib packages which ought to be present if Python is + (['os', 'sys'], True), # test that unpacking also works + (['fake--module'], False), # test an obviously fake module name (don't want to try an actual module in case it becomes an dependency someday) + ([42], False), # test something that isn't even a module to check error handling + ] +) +def test_modules_installed(module_names : list[str], expected_found : bool) -> None: + '''Check that module install checker correctly identifies present and absent modules''' + assert dependencies.modules_installed(*module_names) == expected_found + +# Testing requires_modules decorator +@dependencies.requires_modules('os') +def should_pass() -> str: + '''Dummy function to test requires_modules decorator for dependencies that are present''' + return 'I will run!' + +@dependencies.requires_modules('fake--module') +def should_fail() -> str: + '''Dummy function to test requires_modules decorator for dependencies that are present''' + return 'I will xfail :(' + +@pytest.mark.parametrize( + 'func', + [ + should_pass, + pytest.param(should_fail, marks=pytest.mark.xfail(raises=ImportError, reason='The required module shouldn\'t be found in the environment', strict=True)), + ] +) +def test_requires_modules(func : Callable[..., Any]) -> None: + '''Test that the requires_modules decortor correctly wraps functions''' + _ = func() # no assertion needed, xfail cases should raise Exception while working cases will ternimate without Exception \ No newline at end of file diff --git a/polymerist/tests/genutils/test_pkginspect.py b/polymerist/tests/genutils/importutils/test_pkginspect.py similarity index 99% rename from polymerist/tests/genutils/test_pkginspect.py rename to polymerist/tests/genutils/importutils/test_pkginspect.py index ea07262..81627bf 100644 --- a/polymerist/tests/genutils/test_pkginspect.py +++ b/polymerist/tests/genutils/importutils/test_pkginspect.py @@ -1,4 +1,4 @@ -'''Unit tests for package inspection utilities''' +'''Unit tests for `pkginspect` package`''' __author__ = 'Timotej Bernat' __email__ = 'timotej.bernat@colorado.edu' diff --git a/polymerist/tests/genutils/test_attrs.py b/polymerist/tests/genutils/test_attrs.py index 5f1ac8f..87f84c3 100644 --- a/polymerist/tests/genutils/test_attrs.py +++ b/polymerist/tests/genutils/test_attrs.py @@ -1,4 +1,4 @@ -'''Unit test for attribute inspection''' +'''Unit tests for `attrs` package''' __author__ = 'Timotej Bernat' __email__ = 'timotej.bernat@colorado.edu' diff --git a/polymerist/tests/genutils/trees/test_trees.py b/polymerist/tests/genutils/trees/test_trees.py index 44e408b..73755a9 100644 --- a/polymerist/tests/genutils/trees/test_trees.py +++ b/polymerist/tests/genutils/trees/test_trees.py @@ -1,4 +1,4 @@ -'''Unit tests for trees''' +'''Unit tests for tree-related functionality''' __author__ = 'Timotej Bernat' __email__ = 'timotej.bernat@colorado.edu' diff --git a/polymerist/tests/mdtools/openfftools/omminter/__init__.py b/polymerist/tests/mdtools/openfftools/omminter/__init__.py new file mode 100644 index 0000000..e906b2b --- /dev/null +++ b/polymerist/tests/mdtools/openfftools/omminter/__init__.py @@ -0,0 +1,4 @@ +'''Unit tests for `omminter` package''' + +__author__ = 'Timotej Bernat' +__email__ = 'timotej.bernat@colorado.edu' diff --git a/polymerist/tests/mdtools/openfftools/partialcharge/test_molchargers.py b/polymerist/tests/mdtools/openfftools/partialcharge/test_molchargers.py new file mode 100644 index 0000000..eb291f2 --- /dev/null +++ b/polymerist/tests/mdtools/openfftools/partialcharge/test_molchargers.py @@ -0,0 +1,66 @@ +'''Unit tests for `molchargers` package''' + +__author__ = 'Timotej Bernat' +__email__ = 'timotej.bernat@colorado.edu' + +import pytest + +# NOTE: inverted custom imports here to get polymerist.openfftools import first, +# Done so that if OpenFF is not found, a helpful installation error with be raised prior to attempting direct openff.toolkit imports below +from polymerist.genutils.importutils.dependencies import modules_installed +from polymerist.mdtools.openfftools.partialcharge.molchargers import MolCharger, ABE10Charger, EspalomaCharger, NAGLCharger +from polymerist.mdtools.openfftools.partialcharge.rescharge.interface import LibraryCharger + +from openff.toolkit import Molecule +from openff.toolkit.utils.toolkits import OPENEYE_AVAILABLE + + +# Test MolCharger subclass registration +def test_molcharger_registers_subclasses() -> None: + '''Test that the MolCharger tracks subclasses''' + assert hasattr(MolCharger, 'subclass_registry') + +@pytest.mark.parametrize('expected_charge_method_name, molcharger_subclass', MolCharger.subclass_registry.items()) # NOTE: this will fail if test_molcharger_registers_subclasses() fails +def test_molcharger_subclass_attr_registration(molcharger_subclass : type[MolCharger], expected_charge_method_name : str) -> None: + '''Test that all MolCharger subclasses define and are registered under their "CHARGING_METHOD" class property''' + assert hasattr(molcharger_subclass, 'CHARGING_METHOD') and (getattr(molcharger_subclass, 'CHARGING_METHOD') == expected_charge_method_name) + + +# Test MolCharger subclass implementations +@pytest.fixture +def offmol() -> Molecule: + '''Dummy Molecule object for testing''' + # DEV: worthing double-checking that partial charges are initially empty?? + return Molecule.from_smiles('c1ccccc1C(=O)O') # benzoic acid - nice and small, but with some non-trivial structure + +## selectively register test to avoid failures due to missing optional dependencies +MOLCHARGER_TYPES_TO_TEST : list[type[MolCharger]] = [] +if modules_installed('openeye.oechem', 'openeye.oeomega') and OPENEYE_AVAILABLE: # extra check needed block check when missing license (as is the case for the open-source polymerist repo) + MOLCHARGER_TYPES_TO_TEST.append(ABE10Charger) +if modules_installed('espaloma_charge'): + MOLCHARGER_TYPES_TO_TEST.append(EspalomaCharger) +if modules_installed('openff.nagl', 'openff.nagl_models'): + MOLCHARGER_TYPES_TO_TEST.append(NAGLCharger) +# MOLCHARGER_TYPES_TO_TEST.append( + # LibraryCharger # LibraryCharger behave differently to other MolCharger and are kind of a pain in the ass generally... +# ) # ...intend to deprecate and revamp them eventually, so will just exclude them from testing for now + +@pytest.mark.parametrize('molcharger_subclass', MOLCHARGER_TYPES_TO_TEST) +def test_molchargers_assign_charges(offmol : Molecule, molcharger_subclass : type[MolCharger]) -> None: + charger = molcharger_subclass() + cmol = charger.charge_molecule(offmol) + assert cmol.partial_charges is not None # should assign charges to the new, copied molecule + +@pytest.mark.parametrize('molcharger_subclass', MOLCHARGER_TYPES_TO_TEST) +def test_molchargers_act_readonly(offmol : Molecule, molcharger_subclass : type[MolCharger]) -> None: + charger = molcharger_subclass() + cmol = charger.charge_molecule(offmol) + assert offmol.partial_charges is None # should NOT affect the + +@pytest.mark.parametrize('molcharger_subclass', MOLCHARGER_TYPES_TO_TEST) +def test_molchargers_record_charge_method(offmol : Molecule, molcharger_subclass : type[MolCharger]) -> None: + charger = molcharger_subclass() + cmol = charger.charge_molecule(offmol) + + recorded_charge_method = cmol.properties.get('charge_method', None) + assert (recorded_charge_method is not None) and (recorded_charge_method == getattr(molcharger_subclass, 'CHARGING_METHOD')) diff --git a/polymerist/tests/unitutils/test_dimensions.py b/polymerist/tests/unitutils/test_dimensions.py new file mode 100644 index 0000000..e13720f --- /dev/null +++ b/polymerist/tests/unitutils/test_dimensions.py @@ -0,0 +1,142 @@ +'''Testing that dimensionality checking behaves as expected for both OpenMm and Pint-style unit systems''' + +from typing import Any, Union +from dataclasses import dataclass + +import pytest + +from openmm.unit import ( + Unit as OpenMMUnit, + Quantity as OpenMMQuantity, + centimeter, + second, +) +from pint import ( + Unit as PintUnit, + Quantity as PintQuantity, # this is also the base class for all OpenFF-style units + UnitRegistry, +) +ureg = UnitRegistry() + +from polymerist.unitutils.dimensions import ( + hasunits, + strip_units, + is_volume, +) + + +# Defining test cases and expected outputs +@dataclass +class UnitExample: + '''Internal encapsulation class for indicating expected + properties of unit-like objects in unit tests (no pun intended)''' + value : Any + has_units : bool + is_a_volume : bool # changed name slightly to obviate clash with is_volume() in namespace + +test_cases : list[UnitExample] = [ + # non-units + UnitExample( + value=42, + has_units=False, + is_a_volume=False, + ), + UnitExample( + value=3.1415, + has_units=False, + is_a_volume=False, + ), + UnitExample( + value={1,2,3}, + has_units=False, + is_a_volume=False, + ), + # pure units + UnitExample( + value=second, + has_units=False, + is_a_volume=False, + ), + UnitExample( + value=centimeter**3, + has_units=False, + is_a_volume=True, # despite being a pure unit, this should still count as a volume + ), + UnitExample( + value=ureg.second, + has_units=False, + is_a_volume=False, + ), + UnitExample( + value=ureg.foot**3, + has_units=False, + is_a_volume=True, + ), + # simple quantities + UnitExample( + value=1.1*second, + has_units=True, + is_a_volume=False, + ), + UnitExample( + value=1.2*centimeter, + has_units=True, + is_a_volume=False, + ), + UnitExample( + value=1.3*centimeter**3, + has_units=True, + is_a_volume=True, + ), + UnitExample( + value=2.1*ureg.second, + has_units=True, + is_a_volume=False, + ), + UnitExample( + value=2.2*ureg.centimeter, + has_units=True, + is_a_volume=False, + ), + UnitExample( + value=2.3*ureg.centimeter**3, + has_units=True, + is_a_volume=True, + ), + # mixed quantities + UnitExample( + value=9.8*centimeter*second**-1, + has_units=True, + is_a_volume=False, + ), + UnitExample( + value=9.81*ureg.centimeter*ureg.second**-1, + has_units=True, + is_a_volume=False, + ) +] + +# Unit tests +@pytest.mark.parametrize('unitlike, expected_output', [ + (unit_example.value, unit_example.has_units) + for unit_example in test_cases + ] +) +def test_hasunits(unitlike : Any, expected_output : bool) -> None: + '''Test that objects with (and without) units are correctly identified''' + assert hasunits(unitlike) == expected_output + +@pytest.mark.parametrize('unitlike, expected_output', [ + (unit_example.value, unit_example.is_a_volume) + for unit_example in test_cases + ] +) +def test_hasunits(unitlike : Any, expected_output : bool) -> None: + '''Test that objects which can (and can't) be interpreted as volumes are correctly identified as such''' + assert is_volume(unitlike) == expected_output + +SAMPLE_COORDS : tuple[float] = (1.23, 4.56, 7.89) # these numbers are arbitrary, but need to be consistent across tests +@pytest.mark.parametrize('coordlike', [SAMPLE_COORDS, SAMPLE_COORDS*centimeter, SAMPLE_COORDS*ureg.centimeter]) +def test_strip_units(coordlike : Union[tuple, PintQuantity, OpenMMQuantity]) -> None: + '''Test that removing units works for Pint, OpenMM, and unit-free objects''' + assert tuple(strip_units(coordlike)) == SAMPLE_COORDS # need to re-tuplify to counteract numpy auto-conversion by pint \ No newline at end of file diff --git a/polymerist/unitutils/dimensions.py b/polymerist/unitutils/dimensions.py index 15acd2a..679a749 100644 --- a/polymerist/unitutils/dimensions.py +++ b/polymerist/unitutils/dimensions.py @@ -3,11 +3,25 @@ __author__ = 'Timotej Bernat' __email__ = 'timotej.bernat@colorado.edu' -from typing import Any, Union +from typing import Any, Union, TypeVar +T = TypeVar('T') -from pint import Quantity as PintQuantity # this is also the base class for all OpenFF-style units -from openmm.unit import Quantity, Unit, length_dimension -from .interop import allow_openmm_units, allow_openff_units +from numpy import ndarray +from openmm.unit import ( + Unit as OpenMMUnit, + Quantity as OpenMMQuantity, + length_dimension, +) +OpenMMUnitLike = Union[OpenMMUnit, OpenMMQuantity] # TODO: add union type checkers + +from pint import ( # this is also the base classes for all OpenFF-style units + Unit as PintUnit, + Quantity as PintQuantity, +) +PintUnitLike = Union[PintUnit, PintQuantity] # TODO: add union type checkers + +Unit = Union[PintUnit , OpenMMUnit] +Quantity = Union[PintQuantity, OpenMMQuantity] # CHECKING FOR AND REMOVING UNITS @@ -18,29 +32,45 @@ def hasunits(obj : Any) -> bool: '''Naive but effective way of checking for pint and openmm units''' return any(hasattr(obj, attr) for attr in ('unit', 'units')) -def strip_units(coords : Union[tuple, PintQuantity, Quantity]) -> tuple[float]: +def strip_units(coords : Union[T, PintQuantity, OpenMMQuantity]) -> Union[T, ndarray[Any]]: ''' Sanitize coordinate tuples for cases which require unitless quantities Specifically needed since OpenMM and pint each have their own Quantity and Units classes ''' if isinstance(coords, PintQuantity): - return coords.magnitude - elif isinstance(coords, Quantity): + return coords.magnitude # for container-like values (e.g. tuples), will always return numpy array instead (not type-safe!) + elif isinstance(coords, OpenMMQuantity): return coords._value return coords # CHECKING DIMENSIONALITY -@allow_openff_units -def is_volume(unit_val : Union[Unit, Quantity]) -> bool: - '''Return whether a unit corresponds to a volume''' - if isinstance(unit_val, Quantity): - unit_val = unit_val.unit # extract just the unit component if a Quantity is passed +def _is_volume_openmm(unitlike : OpenMMUnitLike) -> bool: + '''Check whether an OpenMM Unit/Quantity dimensionally corresponds to a volume''' + if isinstance(unitlike, OpenMMQuantity): + unitlike = unitlike.unit # extract just the unit component if a Quantity is passed - for i, (dim, exp) in enumerate(unit_val.iter_base_dimensions()): + for i, (dim, exp) in enumerate(unitlike.iter_base_dimensions()): if i > 0: - return False # immediate rule out if more than just one unit is present + return False # immediate rule out if more than just one dimension is present if (dim == length_dimension) and (exp == 3.0): # if monodimensional, check that the single dimension is L^3 return True - return False \ No newline at end of file + return False + +def _is_volume_pint(unitlike : PintUnitLike) -> bool: + '''Check whether an Pint Unit/Quantity dimensionally corresponds to a volume''' + return unitlike.dimensionality == '[length]**3' # "dimensionality" attr is present on both the Unit and Quantity classes in Pint + +def is_volume(unitlike : Union[Unit, Quantity]) -> bool: + ''' + Check whether a Unit or Quantity dimensionally corresponds to a volume + Accepts both OpenMM-style and Pint-style unit-like objects + ''' + if isinstance(unitlike, OpenMMUnitLike): + return _is_volume_openmm(unitlike) + elif isinstance(unitlike, PintUnitLike): + return _is_volume_pint(unitlike) + else: + # raise TypeError(f'Cannot interpret object of type "{type(unitlike).__name__}" as unit-like') + return False # strictly speaking, anything which has no notion of units cannot be a volume \ No newline at end of file