BUCToolkit 1.0b2__tar.gz → 1.0b3__tar.gz

buctoolkit-1.0b2/BUCToolkit/BatchMD/constrained_md/_ConstrBaseMD.py → buctoolkit-1.0b3/BUCToolkit/Bases/BaseConstraints.py

@@ -1,38 +1,49 @@
-""" Molecular Dynamics base framework with constrains """
-
-#  Copyright (c) 2024-2025.7.4, BUCToolkit.
+#  Copyright (c) 2026.4.24, BUCToolkit.
 #  Authors: Pu Pengxin, Song Xin
-#  Version: 0.9a
-#  File: NVE.py
+#  Version: 1.0b
+#  File: BaseConstraints.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 .._BaseMD import _BaseMD
 from BUCToolkit.utils._print_formatter import FLOAT_ARRAY_FORMAT, SCIENTIFIC_ARRAY_FORMAT
-from BUCToolkit.utils.grad_functions import bjvp, bhvp
+from BUCToolkit.utils.grad_functions import bhvp
+from BUCToolkit.Bases.BaseMotion import BaseIO
 
 
-class _rConstrBase(_BaseMD):
+class BaseConstr(BaseIO):
     """
-    Constrained Base Dynamics
+    Base Constraints used for constraints MD or optimizations.
+    It can be combined into any normal gradient-based algorithms by applying `project1` to gradient or any other tangent space vectors,
+     and `project2` as a retraction mapping to coordinates or any other vectors on the manifold.
+
+     I strongly suggest that using combination with proxy instead of inheritance when it is applied in other class.
+     e.g., define following method in the class with aligning `self._constr = BaseConstr(...)` in the self.__init__ method.
+     ```python
+        def __getattr__(self, name):
+            if '_constr' in self.__dict__:
+                constr = self._constr
+                if hasattr(constr, name):
+                    return getattr(constr, name)
+            raise AttributeError(f"'{type(self).__name__}' object has no attribute '{name}'")
+    ```
 
     Args:
-        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.
-        constr_func: Callable, a tuple of Python functions as the constraint functions s_k(X) that map R^n -> R^k. It takes one or more arguments, one of which must be a Tensor, and returns one Tensor with shape (k, ). `None` for identity function. see example below.
-        constr_val: Callable[th.Tensor[1], th.Tensor] | th.Tensor, the constraint value of `constr_func`, i.e., constraints are `constr_func(X) = constr_val`.
-        By defining it as a callable constr_val = constr_val(t) where `t` is a scalar Tensor, it can be set to the time-dependent constraints.
+        constr_func: Callable, a tuple of Python functions as the constraint functions s_k(X) that map R^n -> R^k.
+            It takes one or more arguments, one of which must be a Tensor, and returns one Tensor with shape (k, ).
+            `None` for constant function that always return [0., ]. It should support auto-gradient ops. see example below.
+        constr_val: Callable[th.Tensor[1], th.Tensor] | th.Tensor, the constraint value of `constr_func`,
+            i.e., constraints are `constr_func(X) = constr_val`.
+            By defining it as a callable constr_val = constr_val(t) where `t` is a scalar Tensor,
+            it can be set to the time-dependent constraints.
         constr_threshold: float, the threshold of constraint convergence (error of manifold violation)
-        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.
+        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.
 
     Examples for constrains:
         def constr_func(X):
@@ -59,29 +70,22 @@ class _rConstrBase(_BaseMD):
     """
     def __init__(
             self,
-            time_step: float,
-            max_step: int,
-            T_init: float = 298.15,
             constr_func: Callable | None = None,
             constr_val: Callable[[th.Tensor], th.Tensor|Tuple[th.Tensor]] | th.Tensor | None = None,
             constr_threshold: float = 1e-5,
-            output_file: str | None = None,
-            output_structures_per_step: int = 1,
             device: str | th.device = 'cpu',
-            verbose: int = 2
+            verbose: int = 0
     ):
-        super().__init__(
-            time_step,
-            max_step,
-            T_init,
-            output_file,
-            output_structures_per_step,
-            device,
-            verbose
-        )
         if constr_func is None:
             constr_func = lambda X: th.tensor(0.)
         self._lazy_calc_constr_val = False  # use the value of constr_func(X_init) to determine constr_val (and fixed).
+        self.device = th.device(device)
+
+        self.time_now = th.scalar_tensor(0., device=device)
+        self.is_const_constr = False
+        self.constr_val_func_raw = None
+        self.time_step = 1.
+
         if isinstance(constr_val, th.Tensor):
             self.is_const_constr = True
             self.constr_val_func_raw = None
@@ -110,12 +114,18 @@ class _rConstrBase(_BaseMD):
         self.constr_thres = constr_threshold
         self.lamb = None  # constr force, i.e., the mu of Lagrange multipler
 
+        self._X_qr_cache = None  # check the consistency of X in project1, ensuring it is the new tangent space projection
+
+        super().__init__()
+        self.init_logger('Main.Constraints')
+        self.verbose = int(verbose)
+
     def initialize(
             self,
             func: Any | nn.Module,
             X: th.Tensor,
-            Element_list: List[List[str]] | List[List[int]],
-            masses: th.Tensor,
+            Element_list: List[List[str]] | List[List[int]] | None,
+            masses: th.Tensor | None,
             V_init: th.Tensor | None = None,
             grad_func: Any | nn.Module = None,
             func_args: Sequence = tuple(),
@@ -128,55 +138,28 @@ class _rConstrBase(_BaseMD):
             fixed_atom_tensor: Optional[th.Tensor] = None,
             is_fix_mass_center: bool = False
     ):
+        if masses is None: masses = th.ones_like(X)
         self.sqrtM = th.sqrt(masses)  # M^1/2, (n_batch, n_atoms, n_dim)
         self.negsqrtM = 1 / th.sqrt(masses)  # M^-1/2
         _y_check = th.vmap(self.constr_func)(X)
         if self._lazy_calc_constr_val:
             self.constr_val_now = _y_check
+        if self.verbose:
+            self.logger.info(
+                f'Constraint values are now {np.array2string(self.constr_val_now.squeeze().numpy(force=True), **SCIENTIFIC_ARRAY_FORMAT)}'
+            )
         # check constr_val shape
         if _y_check.shape != self.constr_val_now.shape:
             raise RuntimeError(
                 f'`constr_val` must have the same shape as what constr_func returned {self.constr_val_now.shape}, but got {_y_check.shape}.'
             )
         jac, y = self._jacobian(X)
-        if self.verbose:
-            self.logger.info(
-                f'Constraint values are now {np.array2string(self.constr_val_now.squeeze().numpy(force=True), **SCIENTIFIC_ARRAY_FORMAT)}'
-            )
         if y.ndim != 2:
             raise ValueError(f'`constr_func` must return a 2D tensor of shape (n_batch, n_constr), but got {y.shape}.')
         self._do_qr(jac)
-        ProjV = self._project1(V_init)
-        Ek = th.sum(
-            masses * V_init ** 2,
-            dim=(-2, -1),
-            keepdim=True
-        )
-        Ek_p = th.sum(
-            masses * ProjV ** 2,
-            dim=(-2, -1),
-            keepdim=True
-        )
-        V_init.copy_(th.where(Ek_p < 1e-5, 0., th.sqrt(Ek/Ek_p) * ProjV))
-        self.free_degree -= jac.shape[1]  # reduce the constr. free deg.
-        # recalculate target Ek under constraints
-        _, n_atom, n_dim = X.shape
-        # target kinetic energy for NVT|NPT ensembles
-        if batch_indices:  # Unit: eV/atom. Boltzmann constant kB = 8.6173332621e-5 eV/K
-            self.EK_TARGET = th.tensor(
-                [((self.free_degree / 2.) * 8.617333262145e-5 * self.T_init) for _n_atom in batch_indices],
-                dtype=X.dtype,
-                device=self.device
-            )
-        else:
-            self.EK_TARGET = (self.free_degree / 2.) * 8.617333262145e-5 * self.T_init
         # calc. constr. intensity
         n_batch, n_constr, _ = self.R.shape
         self.lamb = th.zeros((n_batch, n_constr), device=self.device)
-        # debug
-        #self._debug_X_check = list()
-        #self._debug_V_check = list()
-        # ...
 
     def _constr_func_wrapped(self, X, constr_val_now):
         """
@@ -274,6 +257,12 @@ class _rConstrBase(_BaseMD):
         n_batch, n_atoms, n_dim = X.shape
         sqrtM = self.sqrtM  # M^1/2, (n_batch, n_atoms, n_dim)
         negsqrtM = self.negsqrtM  # M^-1/2
+        # check consistency
+        if (self._X_qr_cache is None) or (not th.allclose(X, self._X_qr_cache)):
+            jac, y = self._jacobian(X)
+            self._do_qr(jac)
+            self._X_qr_cache = X.clone()
+
         # P = M^-1/2 @ (I - Q Q^T) @ M^1/2 = I - M^-1/2 @ Q @ Q^T @ M^1/2
         Q = self.Q  # Q(n_batch, n_atoms * n_dim, n_constr)
         Px = (sqrtM * X).reshape(n_batch, n_atoms * n_dim, 1).contiguous()  # (n_batch, n_atoms * n_dim, 1)
@@ -301,6 +290,12 @@ class _rConstrBase(_BaseMD):
         n_batch, n_atoms, n_dim = X.shape
         sqrtM = self.sqrtM  # M^1/2, (n_batch, n_atoms, n_dim)
         negsqrtM = self.negsqrtM  # M^-1/2
+        # check consistency
+        if (self._X_qr_cache is None) or (not th.allclose(X, self._X_qr_cache)):
+            jac, y = self._jacobian(X)
+            self._do_qr(jac)
+            self._X_qr_cache = X
+
         # P = M^-1/2 @ (I - Q Q^T) @ M^1/2 = I - M^-1/2 @ Q @ Q^T @ M^1/2
         Q = self.Q  # Q(n_batch, n_atoms * n_dim, n_constr)
         Px = (sqrtM * X).reshape(n_batch, n_atoms * n_dim, 1).contiguous()  # (n_batch, n_atoms * n_dim, 1)
@@ -327,15 +322,15 @@ class _rConstrBase(_BaseMD):
             # update Jacobian
             jac, y = self._jacobian(X)
             # print
-            constr_err = th.max(th.abs(y)).item()
+            constr_err = th.max(th.abs(y))
             if self.verbose > 0:
                 self.logger.info(f'{i: < 3d} Constraint errors are now: {constr_err:.4e}')
             # threshold
-            if constr_err < self.constr_thres:
+            if th.all(constr_err < self.constr_thres):
                 # constr. forces lambda
                 Fc = th.linalg.solve_triangular(self.R, _lamb, upper=True)
-                # unit conversion
                 Fc *= 207.28617 / (self.time_step ** 2)  # convert g/mol Angstrom^2 fs^-2 to eV/Atom
+                self._X_qr_cache = X.clone()
                 return Fc
             # QR factor.
             self._do_qr(jac)
@@ -345,75 +340,20 @@ class _rConstrBase(_BaseMD):
             # (n_batch, n_atoms, n_dim) * (n_batch, n_atoms * n_dim, n_constr) @ (n_batch, n_constr, 1)
             corr = self.negsqrtM * (self.Q @ _lamb).reshape(n_batch, n_atoms, n_dim)
             X.add_(corr, alpha=-1.)
-            #V.add_(corr/self.time_step, alpha=-1.)
 
         # if not converged
         self.logger.warning("Projection of X to the manifold is not converged.")
         # constr. forces lambda
         Fc = th.linalg.solve_triangular(self.R, _lamb, upper=True)
-        # unit conversion
+        ######################
+        # FOR MD, the constraint force used for calculate the free energy gradient should be further processed by
+        #   Fc *= 207.28617 / (self.time_step ** 2)  # convert g/mol Angstrom^2 fs^-2 to eV/Atom
+        ######################
         Fc *= 207.28617 / (self.time_step ** 2)  # convert g/mol Angstrom^2 fs^-2 to eV/Atom
-        return Fc
-
-    def _project1_std(self, X:th.Tensor) -> th.Tensor:
-        """
-        Compute the projector for 1st-order quantity (e.g. velocity) of all constrains by QR factorization.
-        for Jacobian matrix J of constraints s_k at X, let J^T M^-1/2 = QR, thus s_k(X + v * dt) = s_k(X) + R^T Q^T M^1/2 v * dt + O(dt^2)
-        define the projector: P = M^1/2 Q_2 Q_2^T M^-1/2 = M^1/2 (I - Q_1 Q_1^T) M^-1/2, where Q = [Q_1, Q_2], Q_1 is from the financial QRF.
-        Args:
-            X: the input tensor of shape (n_batch, n_atom, n_dim). It might be coordinates, velocity, and higher deviations.
-
-        Returns: th.Tensor, the projected X.
-
-        """
-        n_batch, n_atoms, n_dim = X.shape
-        sqrtM = self.sqrtM  # M^1/2, (n_batch, n_atoms, n_dim)
-        negsqrtM = self.negsqrtM  # M^-1/2
-        # P = M^-1/2 @ (I - Q Q^T) @ M^1/2 = I - M^-1/2 @ Q @ Q^T @ M^1/2
-        Q = self.Q  # Q(n_batch, n_atoms * n_dim, n_constr)
-        Px = th.diag_embed(self.sqrtM.flatten(-2, -1), dim1=1, dim2=2) @ X.reshape(n_batch, n_atoms * n_dim, 1).contiguous()  # (n_batch, n_atoms * n_dim, 1)
-        Px = (Q @ (Q.mT.contiguous() @ Px))  # (n_batch, n_atoms, n_dim)
-        Px = th.diag_embed(self.negsqrtM.flatten(-2, -1), dim1=1, dim2=2) @ Px
-        Px = X - Px.reshape(n_batch, n_atoms, n_dim)
-        # manifold veloc. correction
-        if not self.is_const_constr:  # time-dependent constraints
-            th.linalg.solve_triangular(self.R.mT.contiguous(), self.d_constr, upper=False, out=self.q)
-            Px.add_(self.negsqrtM * (self.Q @ self.q.unsqueeze(-1)).reshape(n_batch, n_atoms, n_dim))
+        self._X_qr_cache = X.clone()
 
-        return Px
-
-    def _project2_std(self, X:th.Tensor) -> None:
-        """
-        Compute position corrections by directly solving JM^-1J^T q = s(X)
-        Args:
-            X:
-
-        Returns:
-
-        """
-        n_batch, n_atoms, n_dim = X.shape
-        for i in range(0, self.max_proj_iter):
-            # update Jacobian
-            jac, y = self._jacobian(X)  # jac: (n_batch, n_constr, n_atom, n_dim)
-            # print
-            constr_err = th.max(th.abs(y)).item()
-            lamb = 0.
-            if self.verbose > 0:
-                self.logger.info(f'{i: < 3d} Constraint errors are now: {constr_err:.4e}')
-            # threshold
-            if constr_err < self.constr_thres:
-                return
-            # solve J M^-1 J^T q = s
-            Z = jac*self.negsqrtM  # TODO
-            # Q(n_batch, n_atoms * n_dim, n_constr)
-            corr = th.linalg.solve_triangular(self.R.mT.contiguous(), y.unsqueeze(-1) , upper=False) # (n_batch, n_constr, 1)
-            # (n_batch, n_atoms, n_dim) * (n_batch, n_atoms * n_dim, n_constr) @ (n_batch, n_constr, 1)
-            corr = self.negsqrtM * (self.Q @ corr).reshape(n_batch, n_atoms, n_dim)
-            X.add_(corr, alpha=-1.)
-            #V.add_(corr/self.time_step, alpha=-1.)
+        return Fc
 
-        # if not converged
-        self.logger.warning("Projection of X to the manifold is not converged.")
 
     def _jacobian_derivative(self, X, v: th.Tensor) -> th.Tensor:
         """
@@ -504,7 +444,6 @@ class _rConstrBase(_BaseMD):
         #    # (n_f, n_f) @ (n_f, n_constr) @ (n_constr, 1) = (n_f, 1)
         #    Gamma = M_sqrt_inv @ Q_mid @ _y
 #
-        #    # 带松弛的更新
         #    v_candidate = v_old - dt * Gamma.reshape(n_atom, n_dim)
         #    v_next = omega * v_candidate + (1 - omega) * v_new
 #
@@ -517,91 +456,3 @@ class _rConstrBase(_BaseMD):
 #
         #print('Not Converged !')
         #return v_new
-
-    def implicit_midpoint_step(
-            self,
-            X0,
-            X,
-            V0,
-            V,
-            F,
-            lamb,
-            func, grad_func_, func_args, func_kwargs, grad_func_args, grad_func_kwargs,
-            masses, atom_masks, is_grad_func_contain_y,
-    ):
-        """
-        Solve manifold dynamics by implicit midpoint method.
-        Here the input X0, V0 are at the endpoint in last step,
-        X, V are the predicted endpoint in this step,
-        input F (the potential forces) and lamb (the constraint forces) are at the midpoint of last step,
-        and returned X_mid, V_mid, Energy, Force and lamb are all at the **midpoint**.
-        The X, V in the end point are updated **in-place** to the input variables.
-        Args:
-            X:
-            X0:
-            V:
-            V0:
-            F:
-            lamb:
-            func:
-            grad_func_:
-            func_args:
-            func_kwargs:
-            grad_func_args:
-            grad_func_kwargs:
-            masses:
-            atom_masks:
-            is_grad_func_contain_y:
-
-        Returns:
-
-        """
-        n_batch, n_atoms, n_dim = X.shape
-        # MAIN LOOP
-        is_converged = False
-        for i in range(self.max_proj_iter):
-            X_mid = 0.5 * (X0 + X)
-            V_mid = 0.5 * (V0 + V)
-            # update Jacobian & QR factor.
-            jac, y = self._jacobian(X_mid)  # (n_batch, n_constr, n_atoms, n_dim)
-            self._do_qr(jac)
-            # 1st forces eval
-            with th.set_grad_enabled(self.require_grad):
-                X_mid.requires_grad_(self.require_grad)
-                E_mid = func(X_mid, *func_args, **func_kwargs)
-                if is_grad_func_contain_y:
-                    F_mid = - grad_func_(X_mid, E_mid, *grad_func_args, **grad_func_kwargs) * atom_masks
-                else:
-                    F_mid = - grad_func_(X_mid, *grad_func_args, **grad_func_kwargs) * atom_masks
-            # residuals
-            f_constr = th.einsum('bijk, bi -> bjk', jac, lamb)
-            RX = X - X0 - V_mid * self.time_step
-            RV = V - V0 - (F_mid + f_constr) / masses * self.time_step * 9.64853329045427e-3
-            # check thres.
-            norm_X, norm_V, norm_s = th.linalg.norm(RX), th.linalg.norm(RV), th.linalg.norm(y)
-            if self.verbose > 0:
-                self.logger.info(
-                    f'{i} residuals X, V, constraints: {norm_X.item():.4e}, {norm_V.item():.4e}, {norm_s.item():.4e}'
-                )
-            if (norm_X < self.constr_thres) and (norm_V < self.constr_thres) and (norm_s < self.constr_thres):
-                is_converged = True
-                break
-
-            # update V & X by quasi-Newton
-            tinvR_y = th.linalg.solve_triangular(self.R.mT.contiguous(), y.unsqueeze(-1) , upper=False)  # R^-T y : (n_batch, n_constr, 1)
-            d_lambda = (
-                    - 4./self.time_step * tinvR_y.squeeze(-1)
-                    + 2./self.time_step * th.einsum('bij, bi -> bj', self.Q, (self.sqrtM * RX).flatten(-2, -1))
-                    + th.einsum('bij, bi -> bj', self.Q, (self.sqrtM * (RV - 2*V_mid)).flatten(-2, -1))
-            )
-            d_lambda = 1./(self.time_step * 9.64853329045427e-3) * th.linalg.solve_triangular(self.R, d_lambda.unsqueeze(-1), upper=True).squeeze(-1)
-            ds = 4./self.time_step * self.negsqrtM * (self.Q @ tinvR_y).reshape(n_batch, n_atoms, n_dim)
-            dV = - ds - self._project1(RV) + 2./self.time_step * self._project_norm(RX)
-            dX = 0.5 * self.time_step * dV - RX
-            lamb.add_(d_lambda, alpha=1.)
-            X.add_(dX, alpha=1.)
-            V.add_(dV, alpha=1.)
-        if not is_converged:
-            self.logger.warning('Implicit midpoint method did not converge!')
-
-        return X_mid, V_mid, E_mid, F_mid, lamb

buctoolkit-1.0b3/BUCToolkit/Bases/BaseMotion.py

@@ -0,0 +1,264 @@
+#  Copyright (c) 2026.4.24, BUCToolkit.
+#  Authors: Pu Pengxin, Song Xin
+#  Version: 1.0b
+#  File: BaseMotion.py
+#  Environment: Python 3.12
+import logging
+import os
+import sys
+import warnings
+from typing import Any, Tuple, Dict, Callable, List
+from itertools import accumulate
+
+import torch as th
+import numpy as np
+
+from BUCToolkit.BatchStructures.StructuresIO import structures_io_dumper
+from BUCToolkit.utils.setup_loggers import has_any_handler, clear_all_handlers
+
+
+class BaseIO:
+    """
+    Base class for I/O operations.
+    """
+    def __init__(self, output_file: str|None = None) -> None:
+        self.logger = None
+        self.log_handler = None
+        self.output_file = str(output_file) if output_file is not None else None
+        self.dumper = structures_io_dumper(
+            path=self.output_file,
+            mode='x',
+        )  # as the default. One can use `reset
+
+    def init_logger(self, logger_name: str):
+        # logging
+        if self.logger is None:
+            # cut off propagation
+            supreme_name = logger_name.split('.')[0]
+            top_logger = logging.getLogger(supreme_name)
+            top_logger.propagate = False
+            # set true logger
+            self.logger = logging.getLogger(logger_name)
+            self.logger.setLevel(logging.INFO)
+            formatter = logging.Formatter('%(message)s')
+            if not has_any_handler(top_logger):
+                self.log_handler = logging.StreamHandler(sys.stdout)
+                self.log_handler.setLevel(logging.INFO)
+                self.log_handler.setFormatter(formatter)
+                top_logger.addHandler(self.log_handler)
+        else:
+            warnings.warn('Logger has already initialized. Nothing will be done.', RuntimeWarning)
+
+    def reset_logger_handler(self, handler: str|logging.StreamHandler|logging.FileHandler, level: int = logging.INFO):
+        """
+        Clear all logging handlers including current logger and its ancestors, and reset one.
+        Args:
+            handler: the new handler.
+            level: the logger handler output level.
+
+        Returns:
+
+        """
+        clear_all_handlers(self.logger)
+        # redirect to supreme logger
+        self.logger.setLevel(level)
+        top_logger = self.logger
+        while top_logger.parent and top_logger.propagate:
+            top_logger = top_logger.parent
+            top_logger.setLevel(level)
+
+        top_logger.setLevel(level)
+
+        formatter = logging.Formatter('%(message)s')
+        if isinstance(handler, logging.Handler):
+            self.log_handler = handler
+        elif isinstance(handler, str):
+            output_path = os.path.dirname(handler)
+            # check whether path exists
+            if not os.path.isdir(output_path): os.makedirs(output_path)
+            # set log handler
+            self.log_handler = logging.FileHandler(handler, 'w', delay=True)
+        else:
+            raise TypeError("handler must be a string path or a logging.Handler instance")
+
+        self.log_handler.setLevel(level)
+        self.log_handler.setFormatter(formatter)
+        top_logger.addHandler(self.log_handler)
+
+    def reset_dumper(self, dumper: Any) -> None:
+        """
+        Reset a new dumper and delete the old one.
+        Args:
+            dumper: A dumper object that satisfied the protocol as `_ArrayDumperPlaceHolder` class.
+
+        Returns:
+
+        """
+        if self.output_file is not None:
+            self.dumper.close()
+            del self.dumper
+            self.dumper = dumper
+        else:
+            self.logger.error(
+                "ERROR: No output file specified. Hence, resetting dumper is meaningless.\n"
+                "'reset_dumper': Operation REFUSED."
+            )
+
+class BaseMotion(BaseIO):
+    """
+    Base class for all processes that atoms move & evolution.
+    """
+
+    def __init__(self, output_file: str|None = None):
+        super().__init__(output_file)
+
+    @staticmethod
+    def handle_motion_mask(
+            X,
+            fixed_atom_tensor,
+    ):
+        """
+        normalize format of atom_masks
+        Returns: the standardized atom_masks
+
+        """
+        if fixed_atom_tensor is None:
+            fixed_atom_tensor = th.ones_like(X, device=X.device)
+        else:
+            fixed_atom_tensor = fixed_atom_tensor.broadcast_to(X.shape)
+        if fixed_atom_tensor.shape == X.shape:
+            atom_masks = fixed_atom_tensor.to(X.device)
+        else:
+            raise RuntimeError(f'The shape of fixed_atom_tensor (shape: {fixed_atom_tensor.shape}) does not match X (shape: {X.shape}).')
+
+        return atom_masks
+
+    @staticmethod
+    def handle_grad_func(
+            grad_func: Callable[[th.Tensor, Any, ...], th.Tensor] | None,
+            is_grad_func_contain_y: bool,
+            require_grad: bool,
+            **kwargs
+    ):
+        """
+
+        Returns: grad_func_, require_grad, is_grad_func_contain_y
+
+        """
+        if grad_func is None:
+            is_grad_func_contain_y = True
+            require_grad = True
+            def grad_func_(x, y, grad_shape=None):
+                if grad_shape is None:
+                    grad_shape = th.ones_like(y)
+                _g = th.autograd.grad(y, x, grad_shape)
+                return _g[0]
+        else:
+            grad_func_ = grad_func
+
+        return grad_func_, require_grad, is_grad_func_contain_y
+
+    @staticmethod
+    def handle_batch_indices(
+            batch_indices,
+            n_batch,
+            device
+    ):
+        r"""
+        Calculating `n_true_batch`, `batch_tensor`, `batch_scatter`, and `batch_slice_indx` from input batch_indices.
+
+        Args:
+            batch_indices: input batch_indices that each element is the atom number of each sample: [n_0, n_1, ..., n_N].
+            n_batch: the length of 1st dimension of X.
+            device: torch device
+
+        Returns: n_true_batch, batch_tensor, batch_scatter, batch_slice_indx;
+            n_true_batch: the true batch size
+            batch_tensor: the batch_indices in torch.Tensor format
+            batch_scatter: the batch indices in the form of Tensor[0, 0, 0, ..., 1, 1, ..., N - 1]
+            batch_slice_indx: the batch indices in the form of ptr List[0, n_0, n_0 + n_1, ..., \sum n_i]
+        """
+        if batch_indices is not None:
+            if n_batch != 1:
+                raise RuntimeError(f'If batch_indices was specified, the 1st dimension of X must be 1 instead of {n_batch}.')
+            if isinstance(batch_indices, (th.Tensor, np.ndarray)):
+                batch_indices = batch_indices.tolist()
+            elif not isinstance(batch_indices, (List, Tuple)):
+                raise TypeError(f'Invalid type of batch_indices {type(batch_indices)}. '
+                                f'It must be List[int] | Tuple[int, ...] | th.Tensor | np.ndarray | None')
+            for i in batch_indices: assert isinstance(i, int), f'All elements in batch_indices must be int, but occurred {type(i)}'
+            batch_slice_indx = [0] + list(accumulate(batch_indices))  # convert n_atom of each batch into split point of each batch
+            batch_tensor = th.as_tensor(batch_indices, device=device)  # the tensor version of batch_indices which is a List.
+            batch_scatter = th.repeat_interleave(
+                th.arange(0, len(batch_indices), dtype=th.int64, device=device),
+                batch_tensor,
+                dim=0
+            )  # scatter mask of the int tensor with the same shape as X.shape[1], which the data in one batch have one index.
+            n_true_batch = len(batch_indices)  # the true batch size for irregular batches
+        else:
+            n_true_batch = n_batch
+            batch_tensor = None
+            batch_scatter = None
+            batch_slice_indx = None
+
+        return n_true_batch, batch_tensor, batch_scatter, batch_slice_indx
+
+    def _calc_EF(
+            self,
+            X: th.Tensor,
+            func: Callable[[th.Tensor, Any, ...], th.Tensor],
+            func_args: Tuple,
+            func_kwargs: Dict,
+            grad_func_: Callable[[th.Tensor, Any, ...], th.Tensor],
+            grad_func_args: Tuple,
+            grad_func_kwargs: Dict,
+            require_grad: bool,
+            is_grad_func_contain_y: bool,
+    ) -> Tuple[th.Tensor, th.Tensor]:
+        """
+        Calculate the energy and forces. In fact, it is function value y and its NEGATIVE gradient -g.
+
+        Returns: Tuple[th.Tensor, th.Tensor], energy and forces.
+
+        """
+        y, g = self._calc_y_grad(
+            X,
+            func,
+            func_args,
+            func_kwargs,
+            grad_func_,
+            grad_func_args,
+            grad_func_kwargs,
+            require_grad,
+            is_grad_func_contain_y,
+        )
+
+        return y, g.neg_()
+
+    @th.compiler.disable
+    def _calc_y_grad(
+            self,
+            X: th.Tensor,
+            func: Callable[[th.Tensor, Any, ...], th.Tensor],
+            func_args: Tuple,
+            func_kwargs: Dict,
+            grad_func_: Callable[[th.Tensor, Any, ...], th.Tensor],
+            grad_func_args: Tuple,
+            grad_func_kwargs: Dict,
+            require_grad: bool,
+            is_grad_func_contain_y: bool,
+    ) -> Tuple[th.Tensor, th.Tensor]:
+        """
+        Calculate the function value y and the corresponding gradient of the y.
+
+        Returns: Tuple[th.Tensor, th.Tensor], y and the gradient of the y
+
+        """
+        with th.set_grad_enabled(require_grad):
+            X.requires_grad_(require_grad)
+            y = func(X, *func_args, **func_kwargs)
+            if is_grad_func_contain_y:
+                g = grad_func_(X, y, *grad_func_args, **grad_func_kwargs)
+            else:
+                g = grad_func_(X, *grad_func_args, **grad_func_kwargs)
+        return y, g

buctoolkit-1.0b3/BUCToolkit/Bases/__init__.py

@@ -0,0 +1,5 @@
+#  Copyright (c) 2026.4.24, BUCToolkit.
+#  Authors: Pu Pengxin, Song Xin
+#  Version: 1.0b
+#  File: __init__.py
+#  Environment: Python 3.12