BUCToolkit 1.0b0__py3-none-any.whl

BUCToolkit/BatchMD/__init__.py

@@ -0,0 +1,18 @@
+r"""
+Batched Verlet Molecular Dynamics by pytorch models or functions.
+
+2024/7/30 PPX
+"""
+
+#  Copyright (c) 2024-2025.7.4, BUCToolkit.
+#  Authors: Pu Pengxin, Song Xin
+#  Version: 0.9a
+#  File: __init__.py
+#  Environment: Python 3.12
+
+# ruff: noqa: E701, E702, E703
+from .NVE import NVE
+from .NVT import NVT
+from .BiasedMD import BiasedMD
+
+__all__ = ['NVE', 'NVT', 'BiasedMD']

BUCToolkit/BatchMD/constrained_md/BiasedMD.py

@@ -0,0 +1,273 @@
+#  Copyright (c) 2025.7.4, BUCToolkit.
+#  Authors: Pu Pengxin, Song Xin
+#  Version: 0.9a
+#  File: BiasedMD.py
+#  Environment: Python 3.12
+
+
+from typing import Iterable, Dict, Any, List, Literal, Optional, Callable, Sequence, Tuple  # noqa: F401
+import warnings
+
+import numpy as np
+import torch as th
+from torch import nn
+
+from ._BaseMD import _rBaseMD
+from .NVT import NVT
+from .NVE import NVE
+
+
+class BiasedMD(NVT):
+    """
+    Performing MD with given external-biased potentials.
+    """
+    def __init__(
+            self,
+            time_step: float,
+            max_step: int,
+            thermostat: Literal['Langevin', 'VR', 'Nose-Hoover', 'CSVR'],
+            thermostat_config: Dict | None = None,
+            T_init: float = 298.15,
+            output_file: str | None = None,
+            output_structures_per_step: int = 1,
+            device: str | th.device = 'cpu',
+            verbose: int = 2
+    ):
+        super().__init__(
+            time_step,
+            max_step,
+            thermostat,
+            thermostat_config,
+            T_init,
+            output_file,
+            output_structures_per_step,
+            device,
+            verbose
+        )
+        self.ext_pot_grad_kwargs = None
+        self.ext_pot_grad_args = None
+        self.ext_pot_grad = None
+        self.ext_pot_kwargs = None
+        self.ext_pot_args = None
+        self.ext_pot = None
+        self.__is_set_ext_func = False
+
+    @staticmethod
+    def _auto_grad_func(X: th.Tensor, y: th.Tensor | None = None):
+        grad_shape = th.ones_like(y)
+        g = th.autograd.grad(y, X, grad_shape)
+        return g[0]
+
+    def set_external_potential(
+            self,
+            potential_func: Callable,
+            potential_grad_func: Callable | None = None,
+            func_args:Tuple = tuple(),
+            func_kwargs: None | Dict = None,
+            grad_func_args: Tuple = tuple(),
+            grad_func_kwargs: None | Dict = None,
+            is_grad_contain_y: bool = True,
+            #is_dynamic: bool = False
+    ):
+        """
+        Set the external potential function.
+        Args:
+            potential_func: the function of external potential
+            potential_grad_func: the gradient of external potential. if None, torch.autograd is applied.
+            func_args: arguments of `potential_func`
+            func_kwargs: key words arguments of `potential_func`
+            grad_func_args: arguments of `potential_grad_func`
+            grad_func_kwargs: key words arguments of `potential_grad_func`
+            is_grad_contain_y: whether `potential_grad_func` needs dependent variable y as input. if True, potential_grad_func = g(X, y), else g(X)
+            #is_dynamic: whether potential function containing time as independent variables. if True, potential_func = f(X, t), else f(X).
+
+        Returns:
+
+        """
+        if func_kwargs is None: func_kwargs = {}
+        if grad_func_kwargs is None: grad_func_kwargs = {}
+        self.ext_pot = potential_func
+        self.ext_pot_args = func_args
+        self.ext_pot_kwargs = func_kwargs
+        # ext_grad_func
+        if potential_grad_func is None:
+            self.ext_pot_grad = self._auto_grad_func
+        else:
+            self.ext_pot_grad = potential_grad_func
+
+        self.ext_pot_grad_args = grad_func_args
+        self.ext_pot_grad_kwargs = grad_func_kwargs
+        self.is_grad_contain_y = is_grad_contain_y
+        #self.is_dynamic = is_dynamic
+        self.__is_set_ext_func = True
+
+    def run(
+            self,
+            func: Any | nn.Module,
+            X: th.Tensor,
+            Element_list: List[List[str]] | List[List[int]],
+            V_init: th.Tensor | None = None,
+            grad_func: Any | nn.Module = None,
+            func_args: Tuple|List = tuple(),
+            func_kwargs: Dict | None = None,
+            grad_func_args: Tuple|List = tuple(),
+            grad_func_kwargs: Dict | None = None,
+            is_grad_func_contain_y: bool = True,
+            batch_indices: List[int] | Tuple[int, ...] | th.Tensor | np.ndarray | None = None,
+            fixed_atom_tensor: Optional[th.Tensor] = None,
+            is_fix_mass_center: bool = False
+    ) -> None:
+        """
+
+        Args:
+            func:
+            X:
+            Element_list:
+            V_init:
+            grad_func:
+            func_args:
+            func_kwargs:
+            grad_func_args:
+            grad_func_kwargs:
+            is_grad_func_contain_y:
+            batch_indices:
+            fixed_atom_tensor:
+            is_fix_mass_center:
+
+        Returns:
+
+        """
+        if func_kwargs is None: func_kwargs = dict()
+        if grad_func_kwargs is None: grad_func_kwargs = dict()
+        if grad_func is not None:
+            _grad_func = grad_func
+        else:
+            _grad_func = self._auto_grad_func
+            is_grad_func_contain_y = True
+
+        if self.__is_set_ext_func:
+            wrapper = _FuncWrapperWithExtPot(
+                func,
+                self.ext_pot,
+                _grad_func,
+                self.ext_pot_grad,
+                len(func_args),
+                func_kwargs.keys(),
+                len(grad_func_args),
+                grad_func_kwargs.keys()
+            )
+            wrapped_args = func_args + self.ext_pot_args
+            wrapped_kwargs = func_kwargs.update(self.ext_pot_kwargs)
+            wrapped_grad_args = grad_func_args + self.ext_pot_grad_args
+            wrapped_grad_kwargs = grad_func_kwargs.update(self.ext_pot_grad_kwargs)
+
+            wrapped_func = wrapper.Energy
+            if self.is_grad_contain_y:
+                wrapped_grad_func = wrapper.Grad_y
+            else:
+                wrapped_grad_func = wrapper.Grad
+
+            res = super().run(
+                wrapped_func,
+                X,
+                Element_list,
+                V_init,
+                wrapped_grad_func,
+                wrapped_args,
+                wrapped_kwargs,
+                wrapped_grad_args,
+                wrapped_grad_kwargs,
+                is_grad_func_contain_y,
+                batch_indices,
+                fixed_atom_tensor,
+                is_fix_mass_center
+            )
+
+            return res
+        else:
+            raise RuntimeError(f'The external potential function is not set.')
+
+
+class _FuncWrapperWithExtPot:
+    def __init__(
+            self,
+            func,
+            ext_func,
+            grad_func,
+            ext_grad_func,
+            n_func_args = 0,
+            k_func_kwargs = None,
+            n_grad_func_args = 0,
+            k_grad_func_kwargs = None,
+
+    ) -> None:
+        """
+        A wrapper to add external function
+        Wrap the f(X, ...) into f(X, ...) + ext_func(X, ...)
+
+        And the wrapped func_args would be the concatenation of `func_args` and `ext_func_args`.
+        Hence, the number of func_args, i.e., the `n_func_args` is required to determine that
+         the first `n_func_args`-th args belong to `func` and the rest belong to `ext_func`
+
+        Keyword Args is similar, that require the keys of `func_kwargs`, and the rest belongs to `ext_func`.
+
+        `grad_func` is so, too.
+
+        Args:
+            func: the original function
+            ext_func: the external adding function
+            n_func_args: number of args that belongs to `func`
+            k_func_kwargs: keywords list that belongs to `func`
+            n_grad_func_args,
+            k_grad_func_kwargs
+
+        Methods:
+            Energy: input Tensor `X` and PygData `graph`, it will update graph.pos into X and return model(graph)['energy'].
+            Grad: input Tensor `X` and PygData `graph`, it will update graph.pos into X and return model(graph)['forces'].
+
+        """
+        self._model = func
+        self._ext_func = ext_func
+        self._grad_func = grad_func
+        self._ext_grad_func = ext_grad_func
+        self.n_func_args = n_func_args
+        self.k_func_kwargs = set(k_func_kwargs) if k_func_kwargs is not None else set()
+        self.n_grad_func_args = n_grad_func_args
+        self.k_grad_func_kwargs = set(k_grad_func_kwargs) if k_grad_func_kwargs is not None else set()
+
+        self.forces = None
+
+    @staticmethod
+    def _auto_grad_func(X: th.Tensor, y: th.Tensor | None = None):
+        grad_shape = th.ones_like(y)
+        g = th.autograd.grad(y, X, grad_shape)
+        return g[0]
+
+    def Energy(self, X, *func_args, **func_kwargs, ):
+        self.X = X
+        true_func_args = func_args[:self.n_func_args]
+        ext_func_args = func_args[self.n_func_args:]
+        true_func_kwargs = {k:func_kwargs[k] for k in (func_kwargs.keys() & self.k_func_kwargs)}
+        ext_func_kwargs = {k:func_kwargs[k] for k in (func_kwargs.keys() - self.k_func_kwargs)}
+
+        self.y = self._model(self.X, *true_func_args, **true_func_kwargs) + self._ext_func(X, *ext_func_args, **ext_func_kwargs)
+
+        return self.y
+
+    def Grad(self, X, *grad_func_args, **grad_func_kwargs,):
+        true_func_args = grad_func_args[:self.n_func_args]
+        ext_func_args = grad_func_args[self.n_func_args:]
+        true_func_kwargs = {k: grad_func_kwargs[k] for k in (grad_func_kwargs.keys() & self.k_func_kwargs)}
+        ext_func_kwargs = {k: grad_func_kwargs[k] for k in (grad_func_kwargs.keys() - self.k_func_kwargs)}
+        g = self._grad_func(X, *true_func_args, **true_func_kwargs) + self._ext_grad_func(X, *ext_func_args, **ext_func_kwargs)
+
+        return g
+
+    def Grad_y(self, X, y, *grad_func_args, **grad_func_kwargs, ):
+        true_func_args = grad_func_args[:self.n_func_args]
+        ext_func_args = grad_func_args[self.n_func_args:]
+        true_func_kwargs = {k: grad_func_kwargs[k] for k in (grad_func_kwargs.keys() & self.k_func_kwargs)}
+        ext_func_kwargs = {k: grad_func_kwargs[k] for k in (grad_func_kwargs.keys() - self.k_func_kwargs)}
+        g = self._grad_func(X, y, *true_func_args, **true_func_kwargs) + self._ext_grad_func(X, y, *ext_func_args, **ext_func_kwargs)
+
+        return g

BUCToolkit/BatchMD/constrained_md/ConstrNVE.py

@@ -0,0 +1,169 @@
+""" Micro canonical ensemble (NVE) Molecular Dynamics via Verlet algo. """
+
+#  Copyright (c) 2024-2025.7.4, BUCToolkit.
+#  Authors: Pu Pengxin, Song Xin
+#  Version: 0.9a
+#  File: NVE.py
+#  Environment: Python 3.12
+
+# ruff: noqa: E701, E702, E703
+from typing import Iterable, Dict, Any, List, Literal, Optional, Callable, Sequence, Tuple  # noqa: F401
+
+import torch as th
+from torch import nn
+import numpy as np
+
+from ._ConstrBaseMD import _rConstrBase
+from BUCToolkit.utils._print_formatter import FLOAT_ARRAY_FORMAT, SCIENTIFIC_ARRAY_FORMAT
+
+
+class ConstrNVE(_rConstrBase):
+    """
+    Constrained micro canonical ensemble (NVE) molecular dynamics implemented via velocity Verlet algo.
+
+    Parameters:
+        time_step: float, time per step (ps).
+        max_step: int, maximum steps.
+        T_init: float, initial temperature, only to generate initial velocities of atoms by Maxwell-Boltzmann distribution. If V_init is given, T_init will be ignored.
+        output_structures_per_step: int, output structures per output_structures_per_step steps.
+        device: str|torch.device, device that program rum on.
+        verbose: int, control the detailed degree of output information. 0 for silence, 1 for output Energy and Forces per step, 2 for output all structures.
+
+    Methods:
+        run: run BatchMD.
+
+    """
+
+    def __init__(
+            self,
+            time_step: float,
+            max_step: int,
+            constr_func: Callable[[th.Tensor], th.Tensor] = None,
+            constr_val: Callable[[th.Tensor], th.Tensor|Tuple[th.Tensor]] | th.Tensor = None,
+            constr_threshold: float = 1e-5,
+            T_init: float = 298.15,
+            output_file: str | None = None,
+            dump_path: str | None = None,
+            output_structures_per_step: int = 1,
+            device: str | th.device = 'cpu',
+            verbose: int = 2
+    ):
+        super().__init__(
+            time_step,
+            max_step,
+            T_init,
+            constr_func,
+            constr_val,
+            constr_threshold,
+            output_file,
+            dump_path,
+            output_structures_per_step,
+            device,
+            verbose
+        )
+        self._X, self._V = None, None
+
+    def initialize(
+            self,
+            func: Any | nn.Module,
+            X: th.Tensor,
+            Element_list: List[List[str]] | List[List[int]],
+            masses: th.Tensor,
+            V_init: th.Tensor | None = None,
+            grad_func: Any | nn.Module = None,
+            func_args: Sequence = tuple(),
+            func_kwargs: Dict | None = None,
+            grad_func_args: Sequence = tuple(),
+            grad_func_kwargs: Dict | None = None,
+            is_grad_func_contain_y: bool = True,
+            require_grad: bool = False,
+            batch_indices: List[int] | Tuple[int, ...] | th.Tensor | np.ndarray | None = None,
+            fixed_atom_tensor: Optional[th.Tensor] = None,
+            is_fix_mass_center: bool = False
+    ):
+        super().initialize(
+            func,
+            X,
+            Element_list,
+            masses,
+            V_init,
+            grad_func,
+            func_args,
+            func_kwargs,
+            grad_func_args,
+            grad_func_kwargs,
+            is_grad_func_contain_y,
+            require_grad,
+            batch_indices,
+            fixed_atom_tensor,
+            is_fix_mass_center
+        )
+        self._X, self._V = X, V_init
+
+    def _updateXV(
+            self, X, V, Force,
+            func, grad_func_, func_args, func_kwargs, grad_func_args, grad_func_kwargs,
+            masses, atom_masks, is_grad_func_contain_y, batch_indices,
+    ) -> Tuple[th.Tensor, th.Tensor, th.Tensor, th.Tensor]:
+        """ Update X, V, Force, and return X, V, Energy, Force. """
+        # X: th.Tensor = X.contiguous()
+        # V: th.Tensor = V.contiguous()
+        # masses: th.Tensor = masses.contiguous()
+        with th.no_grad():
+            # X = X + V * self.time_step + (Force / (2. * masses)) * self.time_step ** 2 * 9.64853329045427e-3
+            V.addcdiv_(Force, masses, value=0.5 * self.time_step * 9.64853329045427e-3)
+            X.add_(V, alpha=self.time_step)
+            Fc = self._project2(X)  # in-place update
+            if self.verbose > 0:
+                self.logger.info(f'Constraint forces \\lambda: {np.array2string(Fc.squeeze().numpy(force=True), **SCIENTIFIC_ARRAY_FORMAT)}')
+            # V = V + (Force / (2. * masses)) * self.time_step * 9.64853329045427e-3  # half-step veloc. update, to avoid saving 2 Forces Tensors.
+            # Update V
+            with th.set_grad_enabled(self.require_grad):
+                X.requires_grad_(self.require_grad)
+                Energy = func(X, *func_args, **func_kwargs)
+                if is_grad_func_contain_y:
+                    Force = - grad_func_(X, Energy, *grad_func_args, **grad_func_kwargs) * atom_masks
+                else:
+                    Force = - grad_func_(X, *grad_func_args, **grad_func_kwargs) * atom_masks
+
+            # V = V + (Force / (2. * masses)) * self.time_step * 9.64853329045427e-3
+            V.addcdiv_(Force, masses, value=0.5 * self.time_step * 9.64853329045427e-3)
+            V.copy_(self._project1(V))
+
+        return X, V, Energy, Force
+
+    def __updateXV(
+            self, X, V, Force,
+            func, grad_func_, func_args, func_kwargs, grad_func_args, grad_func_kwargs,
+            masses, atom_masks, is_grad_func_contain_y,
+    ) -> Tuple[th.Tensor, th.Tensor, th.Tensor]:
+        """
+        TESTING...
+        Here the returned X, V, Force and Energy are all the quantities at the midpoint.
+        And the quantities in the end point are additionally stored in external attr. self._X, self._V,
+        and these attributes would be updated in-place within `self.implicit_midpoint_step`
+        """
+        # predict the new endpoint by last midpoint and endpoint
+        X0 = self._X
+        V0 = self._V
+        f_constr = th.einsum('bijk, bi -> bjk', self.jac, self.lamb)
+        V_mid_pred = V.add((Force + f_constr) / (2. * masses), alpha=self.time_step * 9.64853329045427e-3)
+        X_mid_pred = X.add(V_mid_pred, alpha=self.time_step)
+        self._V = 2 * V_mid_pred - self._V
+        self._X = 2 * X_mid_pred - self._X
+
+        X, V, E, F, lamb = self.implicit_midpoint_step(
+            X0, self._X, V0, self._V, Force, self.lamb,
+            func, grad_func_, func_args, func_kwargs, grad_func_args, grad_func_kwargs,
+            masses, atom_masks, is_grad_func_contain_y,
+        )
+        self.logger.info(
+            f'Constraint forces \\lambda: {np.array2string(lamb.squeeze().numpy(force=True), **SCIENTIFIC_ARRAY_FORMAT)}'
+        )
+        jac, y = self._jacobian(X)
+        self._debug_X_check.append(abs(y[0].item()))
+        self._debug_V_check.append(abs(th.sum(jac * V.unsqueeze(1), dim=(-2, -1))[0].item()))
+        print(abs(th.sum(jac * V.unsqueeze(1), dim=(-2, -1))[0].item()))
+
+        return X, V, E, F
+