BUCToolkit 1.0b0__py3-none-any.whl

BUCToolkit/BatchOptim/minimize/CG.py

@@ -0,0 +1,256 @@
+"""
+Conjugate Gradient Algorithm for optimization.
+"""
+#  Copyright (c) 2024-2025.7.4, BUCToolkit.
+#  Authors: Pu Pengxin, Song Xin
+#  Version: 0.9a
+#  File: CG.py
+#  Environment: Python 3.12
+
+from typing import Literal
+import os
+
+import torch as th
+
+from BUCToolkit.BatchOptim._BaseOpt import _BaseOpt
+from BUCToolkit.utils.index_ops import index_inner_product
+
+os.environ['PYTORCH_CUDA_ALLOC_CONF'] = 'expandable_segments:True'
+
+
+class CG(_BaseOpt):
+    """
+    Conjugate Gradient Algo.
+    Args:
+        iter_scheme: Literal['PR+', 'FR', 'SD'],
+        E_threshold: float = 1e-3,
+        F_threshold: float = 0.05,
+        maxiter: int = 100,
+        linesearch: Literal['Backtrack', 'Wolfe', 'NWolfe', '2PT', '3PT', 'Golden', 'Newton', 'None'] = 'Backtrack',
+        linesearch_maxiter: int = 10,
+        linesearch_thres: float = 0.02,
+        linesearch_factor: float = 0.6,
+        steplength: float = 0.5,
+        use_bb: whether to use Barzilai-Borwein steplength (BB1 or long BB) as initial steplength instead of fixed one.
+        device: str | th.device = 'cpu',
+        verbose: int = 2
+
+    Methods:
+        run: launch the optimization.
+        set_update_batch: setting the method to update the taget function when variables change.
+    """
+    def __init__(
+            self,
+            iter_scheme: Literal['PR+', 'FR', 'SD'],
+            E_threshold: float = 1e-3,
+            F_threshold: float = 0.05,
+            maxiter: int = 100,
+            linesearch: Literal['Backtrack', 'B', 'Wolfe', 'W', 'MT', 'EXACT', 'None', 'N'] = 'Backtrack',
+            linesearch_maxiter: int = 10,
+            linesearch_thres: float = 0.02,
+            linesearch_factor: float = 0.6,
+            steplength: float = 0.5,
+            use_bb: bool = True,
+            output_file: str | None = None,
+            device: str | th.device = 'cpu',
+            verbose: int = 2
+    ):
+        super().__init__(
+            iter_scheme,
+            E_threshold,
+            F_threshold,
+            maxiter,
+            linesearch,
+            linesearch_maxiter,
+            linesearch_thres,
+            linesearch_factor,
+            steplength,
+            use_bb,
+            output_file,
+            device,
+            verbose
+        )
+        self.__update_scheme = None
+        self.MIN_CACHE = th.scalar_tensor(1.e-20, device=self.device)
+
+    @staticmethod
+    def __check_restart(gg, ggo, g, p, beta):
+        """
+        check if restarting is needed.
+
+        """
+        ortho_check = th.abs(ggo) > 0.1 * gg
+        return ortho_check
+
+    def __PRP_irreg(
+            self,
+            g: th.Tensor,
+            g_old: th.Tensor,
+            p: th.Tensor,
+            X: th.Tensor,
+            batch_scatter_indices: th.Tensor | None,
+    ) -> th.Tensor:
+        """
+        POLAK-RIBIERE CG ALGORITHM for irregular batch.
+        """
+        gogo = th.sum(index_inner_product(g_old, g_old, 1, batch_scatter_indices), dim=-1, keepdim=True).clamp_min_(self.MIN_CACHE)  # (1, B, 1)
+        ggo  = th.sum(index_inner_product(g, g_old, 1, batch_scatter_indices), dim=-1, keepdim=True)
+        gg   = th.sum(index_inner_product(g, g, 1, batch_scatter_indices), dim=-1, keepdim=True)
+        beta = (gg - ggo) / gogo  # (1, B, 1)
+        beta.clamp_min_(self.MIN_CACHE)
+        # Restart
+        is_restart = self.__check_restart(gg, ggo, g, p, beta)
+        beta.masked_fill_(is_restart, 0.)
+        beta = beta.index_select(1, batch_scatter_indices)
+        if self.verbose > 1:
+            self.logger.info(f" Restart: {is_restart.flatten().cpu().numpy()}")
+        # update directions
+        #p = -g + beta * p  # (n_batch, n_atom, n_dim)
+        th.addcmul(-g, beta, p, out=p)
+
+        return p
+
+    def __PRP_reg(
+            self ,
+            g: th.Tensor,
+            g_old: th.Tensor,
+            p: th.Tensor,
+            X: th.Tensor,
+            batch_scatter_
+    ) -> th.Tensor:
+        """
+        POLAK-RIBIERE CG ALGORITHM for regular batch.
+        """
+        gogo = th.sum(g_old * g_old, dim=(-2, -1), keepdim=True).clamp_min_(self.MIN_CACHE) # to avoid divide 0
+        ggo = th.sum(g * g_old, dim=(-2, -1), keepdim=True)  # (n_batch, 1, 1)
+        gg = th.sum(g * g, dim=(-2, -1), keepdim=True)  # (n_batch, 1, 1)
+        beta = (gg - ggo) / gogo  # (1, B, 1)
+        beta.clamp_min_(self.MIN_CACHE)
+        # Restart
+        is_restart = self.__check_restart(gg, ggo, g, p, beta)
+        beta.masked_fill_(is_restart, 0.)
+        if self.verbose > 1:
+            self.logger.info(f" Restart: {is_restart.flatten().cpu().numpy()}")
+        # update directions
+        #p = -g + beta * p  # (n_batch, n_atom, n_dim)
+        th.addcmul(-g, beta, p, out=p)
+
+        return p
+
+    def __FR_irreg(
+            self,
+            g: th.Tensor,
+            g_old: th.Tensor,
+            p: th.Tensor,
+            X: th.Tensor,
+            batch_scatter_indices
+    ) -> th.Tensor:
+        """ FLETCHER-REEVES ALGO. for irregular batch. """
+        gogo = th.sum(index_inner_product(g_old, g_old, 1, batch_scatter_indices), dim=-1, keepdim=True).clamp_min_(self.MIN_CACHE)  # (1, B, 1)
+        ggo = th.sum(index_inner_product(g, g_old, 1, batch_scatter_indices), dim=-1, keepdim=True)
+        gg = th.sum(index_inner_product(g, g, 1, batch_scatter_indices), dim=-1, keepdim=True)
+        beta = gg / gogo
+        # Restart
+        is_restart = self.__check_restart(gg, ggo, g, p, beta)
+        beta.masked_fill_(is_restart, 0.)  # (1, B, 1)
+        beta = beta.index_select(1, batch_scatter_indices)  # (1, sumN, 1)
+        if self.verbose > 1:
+            self.logger.info(f" Restart: {is_restart.flatten().cpu().numpy()}")
+        # update directions
+        th.addcmul(-g, beta, p, out=p)
+
+        return p
+
+    def __FR_reg(
+            self,
+            g: th.Tensor,
+            g_old: th.Tensor,
+            p: th.Tensor,
+            X: th.Tensor,
+            batch_scatter_
+    ) -> th.Tensor:
+        """ FLETCHER-REEVES ALGO. for regular batch. """
+        gogo = th.sum(g_old * g_old, dim=(-2, -1), keepdim=True).clamp_min_(self.MIN_CACHE)  # to avoid divide 0
+        ggo = th.sum(g * g_old, dim=(-2, -1), keepdim=True)  # (n_batch, 1, 1)
+        gg = th.sum(g * g, dim=(-2, -1), keepdim=True)  # (n_batch, 1, 1)
+        beta = gg / gogo
+        # Restart
+        is_restart = self.__check_restart(gg, ggo, g, p, beta)
+        beta.masked_fill_(is_restart, 0.)
+        if self.verbose > 1:
+            self.logger.info(f" Restart: {is_restart.flatten().cpu().numpy()}")
+        # update directions
+        th.addcmul(-g, beta, p, out=p)  # (n_batch, n_atom, n_dim)
+
+        return p
+
+    @staticmethod
+    def __SD(g: th.Tensor, g_old: th.Tensor, p: th.Tensor, X: th.Tensor, batch_scatter_) -> th.Tensor:
+        """ Steepest descent """
+        return -g
+
+
+    def initialize_algo_param(self):
+        """
+        Override this method to initialize attribute variables for self._update_direction.
+
+        Returns: None
+        """
+        if self.iterform == "SD":
+            self.__update_scheme = self.__SD
+            return
+        if self.is_concat_X:
+            if self.iterform == "PR+":
+                self.__update_scheme = self.__PRP_irreg
+            elif self.iterform == "FR":
+                self.__update_scheme = self.__FR_irreg
+            else:
+                raise NotImplementedError("Unknown Iterative Scheme.")
+        else:
+            if self.iterform == "PR+":
+                self.__update_scheme = self.__PRP_reg
+            elif self.iterform == "FR":
+                self.__update_scheme = self.__FR_reg
+            else:
+                raise NotImplementedError("Unknown Iterative Scheme.")
+
+    def _update_direction(
+            self,
+            g: th.Tensor,
+            g_old: th.Tensor,
+            p: th.Tensor,
+            X: th.Tensor,
+            batch_scatter_: th.Tensor | None,
+    ) -> th.Tensor:
+        """
+        Override this method to implement X update algorithm.
+        Args:
+            g: (n_batch, n_atom, n_dim), the gradient of X at this step
+            g_old: (n_batch, n_atom, n_dim), the gradient of X at last step
+            p: (n_batch, n_atom, n_dim), the update direction of X at last step
+            X: (n_batch, n_atom, n_dim), the independent vars X.
+
+        Returns:
+            p: th.Tensor, the new update direction of X.
+        """
+        p = self.__update_scheme(g, g_old, p, X, batch_scatter_)
+
+        return p
+
+    def _update_algo_param(
+            self,
+            select_mask: th.Tensor,
+            select_mask_short,
+            batch_scatter_indices,
+            g: th.Tensor,
+            g_old: th.Tensor,
+            p: th.Tensor,
+            displace: th.Tensor
+    ) -> None:
+        """
+        Override this method to update the parameters of X update algorithm i.e., self.iterform.
+
+        Returns: None
+        """
+        pass
+

BUCToolkit/BatchOptim/minimize/FIRE.py

@@ -0,0 +1,310 @@
+"""
+FIRE Optimization Algorithm. Phys. Rev. Lett., 2006, 97:170201.
+"""
+import copy
+#  Copyright (c) 2024-2025.7.4, BUCToolkit.
+#  Authors: Pu Pengxin, Song Xin
+#  Version: 0.9a
+#  File: FIRE.py
+#  Environment: Python 3.12
+
+import os
+from typing import Iterable, Dict, Any, List, Literal, Optional, Callable, Sequence, Tuple  # noqa: F401
+
+import torch as th
+from torch import nn
+
+from BUCToolkit.utils._Element_info import MASS, N_MASS
+from .._BaseOpt import _BaseOpt
+from BUCToolkit.utils.index_ops import index_inner_product
+
+os.environ['PYTORCH_CUDA_ALLOC_CONF'] = 'expandable_segments:True'
+
+
+class FIRE(_BaseOpt):
+    def __init__(
+            self,
+            E_threshold: float = 1e-3,
+            F_threshold: float = 0.05,
+            maxiter: int = 100,
+            steplength: float = 1.,
+            alpha: float = 0.1,
+            alpha_fac: float = 0.99,
+            fac_inc: float = 1.1,
+            fac_dec: float = 0.5,
+            N_min: int = 5,
+            output_file: str | None = None,
+            device: str | th.device = 'cpu',
+            verbose: int = 2,
+            _hold_samples: bool = False,
+            **kwargs
+    ):
+        r"""
+        FIRE Algorithm for optimization.
+
+        Args:
+            E_threshold: float, threshold of difference of func between 2 iterations.
+            F_threshold: float, threshold of gradient of func.
+            maxiter: int, the maximum iteration steps.
+            steplength: The initial step length, i.e. the BatchMD time step.
+            alpha: FIRE arg. alpha.
+            alpha_fac: FIRE arg. the factor of alpha change.
+            fac_inc: FIRE arg. the increment factor.
+            fac_dec: FIRE arg. the decrement factor.
+            N_min: FIRE arg. the minimum number that keeps steplength static.
+            device: The device that program runs on.
+            verbose: amount of print information.
+            _hold_samples: ONLY FOR SPECIAL USE (e.g., CI-NEB or DEBUG).
+                If True, FIRE optimizer will not remove any sample in a batch even if the sample has converged.
+
+        Method:
+            run: running the main optimization program.
+        """
+        super().__init__(
+            'FIRE',
+            E_threshold,
+            F_threshold,
+            maxiter,
+            'None',
+            1,
+            1.,
+            1.,
+            1.,
+            use_bb = False,
+            output_file=output_file,
+            device = device,
+            verbose = verbose,
+            _hold_samples=_hold_samples
+        )
+        self.MAX_STEPLENGTH = th.scalar_tensor(10 * steplength, device=self.device)
+        self.FIRE_t_init = steplength
+        if not (0. < alpha_fac < 1.):
+            raise ValueError('alpha_fac must between 0 and 1.')
+        if not (0. < fac_dec < 1.):
+            raise ValueError('fac_dec must between 0 and 1.')
+        if fac_inc <= 1.:
+            raise ValueError('fac_inc must be greater than 1.')
+
+        self.alpha = alpha
+        self.alpha_fac = alpha_fac
+        self.fac_inc = fac_inc
+        self.fac_dec = fac_dec
+        self.N_min = N_min
+
+        self.v: th.Tensor | None = None        # shape as X
+        self.v_: th.Tensor | None = None       #
+        self.masses: th.Tensor | None = None   # shape as X
+        self.masses_: th.Tensor | None = None  #
+        self.t: th.Tensor | None = None        # (1, sumN, 1) or (B, 1, 1)
+        self.t_: th.Tensor | None = None
+        self.a: th.Tensor | None = None        # (1, sumN, 1) or (B, 1, 1)
+        self.n_count: th.Tensor | None = None  # (1, sumN, 1) or (B, 1, 1)
+
+    def initialize_algo_param(self):
+        """ """
+        if self.is_concat_X:
+            self.t = th.full((1, self.n_atom, 1), self.FIRE_t_init, device=self.device)
+            self.a = th.full((1, self.n_atom, 1), self.alpha, device=self.device)
+            self.n_count = th.zeros((1, self.n_atom, 1), dtype=th.int, device=self.device)
+        else:
+            self.t = th.full((self.n_true_batch, 1, 1), self.FIRE_t_init, device=self.device)
+            self.a = th.full((self.n_true_batch, 1, 1), self.alpha, device=self.device)
+            self.n_count = th.zeros((self.n_true_batch, 1, 1), dtype=th.int, device=self.device)
+
+    def _update_algo_param(
+            self,
+            select_mask,
+            select_mask_short,
+            batch_scatter_indices,
+            X_grad_,
+            X_grad_old_,
+            p,
+            displace_
+    ):
+        F = - X_grad_
+        if self.is_concat_X:
+            self.masses_ = self.masses[:, select_mask, :]
+            self.t_: th.Tensor = self.t[:, select_mask, :]  # (1, sumN, 1)
+            a_ = self.a[:, select_mask, :]  # (1, sumN, 1)
+            n_count_ = self.n_count[:, select_mask, :]
+            self.v_ = self.v[:, select_mask, :]
+            # (1, sumN, n_dim)
+            F_hat = F / th.sum(index_inner_product(
+                F,
+                F,
+                dim=1,
+                batch_indices=batch_scatter_indices
+            ), dim=-1, keepdim=True).sqrt_().index_select(1, batch_scatter_indices)
+            # (1, sumN, 1)
+            momenta = th.sum(index_inner_product(
+                F,
+                self.v_,
+                dim=1,
+                batch_indices=batch_scatter_indices
+            ), dim=-1, keepdim=True).index_select(1, batch_scatter_indices)
+            # (1, sumN, 1)
+            v_norm = th.sum(index_inner_product(
+                self.v_,
+                self.v_,
+                dim=1,
+                batch_indices=batch_scatter_indices
+            ), dim=-1, keepdim=True).sqrt_().index_select(1, batch_scatter_indices)
+            # update velocity: v = v * (1 - a) + a * |v| * \hat{F}
+            F_hat.mul_(v_norm)
+            self.v_.mul_((1. - a_))
+            self.v_.addcmul_(a_, F_hat)
+            # if P > 0
+            n_count_ += th.where(momenta > 0., 1, -n_count_)  # (1, sumN, 1)
+            is_ncount_gt_Nmin = n_count_ >= self.N_min
+            #
+            new_t_ = (self.t_ * self.fac_inc).clamp_max_(self.MAX_STEPLENGTH)
+            self.t_ = th.where(is_ncount_gt_Nmin, new_t_, self.t_)
+            a_ = th.where(is_ncount_gt_Nmin, (a_ * self.alpha_fac), a_)
+            # if P <= 0.
+            is_p_lt_0 = momenta <= 0.
+            self.t_ = th.where(is_p_lt_0, (self.t_ * self.fac_dec), self.t_)
+            self.v_.masked_fill_(
+                is_p_lt_0,
+                0.
+            )
+            a_.masked_fill_(is_p_lt_0, self.alpha)
+            # re-write
+            select_indices = th.where(select_mask)[0]
+            self.t.index_copy_(1, select_indices, self.t_)
+            self.a.index_copy_(1, select_indices, a_)
+            self.n_count.index_copy_(1, select_indices, n_count_)
+
+        else:
+            self.v_ = self.v[select_mask, ...]
+            self.masses_ = self.masses[select_mask, ...]
+            self.t_ = self.t[select_mask, ...]
+            a_ = self.a[select_mask, ...]
+            n_count_ = self.n_count[select_mask, ...]
+
+            F_hat = F / (th.linalg.norm(F, dim=(-2, -1), keepdim=True) + 1e-20)
+            # (n_batch, n_dim, n_atom) @ (n_batch, n_atom, n_dim) -> (n_batch, 1, 1)
+            momenta = th.sum(F * self.v_, dim=(-1, -2), keepdim=True)
+            # update velocity
+            v_norm = th.linalg.norm(self.v_, dim=(-2, -1), keepdim=True)
+            F_hat.mul_(v_norm)
+            self.v_.mul_((1. - a_))
+            self.v_.addcmul_(a_, F_hat)
+            # if P > 0.
+            n_count_ += th.where(momenta > 0., 1, -n_count_)
+            is_ncount_gt_Nmin = n_count_ >= self.N_min
+            new_t_ = (self.t_ * self.fac_inc).clamp_max_(self.MAX_STEPLENGTH)
+            self.t_ = th.where(is_ncount_gt_Nmin, new_t_, self.t_)
+            a_ = th.where(is_ncount_gt_Nmin, (a_ * self.alpha_fac), a_)
+            # if P <= 0.
+            is_p_lt_0 = momenta <= 0.
+            self.t_ = th.where(is_p_lt_0, (self.t_ * self.fac_dec), self.t_)
+            self.v_.masked_fill_(
+                is_p_lt_0,
+                0.
+            )
+            a_.masked_fill_(is_p_lt_0, self.alpha)
+            # re-write
+            select_indices = th.where(select_mask)[0]
+            self.t.index_copy_(0, select_indices, self.t_)
+            self.a.index_copy_(0, select_indices, a_)
+            self.n_count.index_copy_(0, select_indices, n_count_)
+        pass
+
+    def _update_direction(
+            self,
+            g: th.Tensor,
+            g_old: th.Tensor,
+            p: th.Tensor,
+            X: th.Tensor,
+            batch_scatter_indices: th.Tensor | None,
+    ) -> th.Tensor:
+        self.v_: th.Tensor
+        self.v_.addcdiv_(g * self.t_, self.masses_, value=-9.64853329045427e-3)
+        disp = self.v_ * self.t_
+        return disp
+
+    def _update_algo_batches(
+            self,
+            select_indices: th.Tensor,
+            select_indices_short: th.Tensor | None,
+    ):
+        if self.is_concat_X:
+            self.v.index_copy_(1, select_indices, self.v_)
+        else:
+            self.v.index_copy_(0, select_indices, self.v_)
+
+    def run(
+            self,
+            func: Any | nn.Module,
+            X: th.Tensor,
+            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,
+            output_grad: bool = False,
+            fixed_atom_tensor: Optional[th.Tensor] = None,
+            batch_indices: None | List[int] | Tuple[int, ...] | th.Tensor = None,
+            elements: List[List[str | int]] | None = None,
+    ) -> Tuple[th.Tensor, th.Tensor] | Tuple[th.Tensor, th.Tensor, th.Tensor]:
+        r"""
+        Run the Conjugate gradient
+
+        Args:
+            func: the main function of instantiated torch.nn.Module class.
+            X: Tensor[n_batch, n_atom, 3], the optimization variables (e.g., atom coordinates) that input to func.
+            grad_func: user-defined function that grad_func(X, ...) returns the func's gradient at X. if None, grad_func(X, ...) = th.autograd.grad(func(X, ...), X).
+            func_args: optional, other input of func.
+            func_kwargs: optional, other input of func.
+            grad_func_args: optional, other input of grad_func.
+            grad_func_kwargs: optional, other input of grad_func.
+            is_grad_func_contain_y: bool, if True, grad_func contains output of func followed by X i.e., grad = grad_func(X, y, ...), else grad = grad_func(X, ...)
+            require_grad: bool, if True, autograd will be turned on for func(X, *func_args, **func_kwargs) calculation.
+            output_grad: bool, whether output gradient of last step.
+            fixed_atom_tensor: Optional[th.Tensor], the indices of X that fixed. An integer tensor with the same shape as X where 1 is for free and 0 is for fixation.
+            batch_indices: Sequence | th.Tensor | np.ndarray | None, the split points for given X, Element_list & V_init, must be 1D integer array_like.
+                the format of batch_indices is the same as `split_size_or_sections` in torch.split:
+                batch_indices = (n1, n2, ..., nN) will split X, Element_list & V_init into N parts, and ith parts has ni atoms. sum(n1, ..., nN) = X.shape[1]
+            elements: Optional[List[List[str | int]]], the Element of each given atom in X.
+
+        Returns:
+            min func: Tensor(n_batch, ), the minimum of func.
+            argmin func: Tensor(X.shape), the X corresponds to min func.
+            grad of argmin func: Tensor(X.shape), only output when `output_grad` == True. The gradient of X corresponding to minimum.
+        """
+        # manage Atomic Type & Masses
+        if len(X.shape) == 2:
+            X = X.unsqueeze(0)
+        elif len(X.shape) != 3:
+            raise ValueError(f'`X` must be 2D or 3D, but got shape [{X.shape}]')
+        if elements is None:
+            self.masses = th.ones_like(X)
+        elif isinstance(elements, Sequence):
+            masses = list()
+            for _Elem in elements:
+                masses.append([MASS[__elem] if isinstance(__elem, str) else N_MASS[__elem] for __elem in _Elem])
+            masses = th.tensor(masses, dtype=th.float32, device=self.device)
+            self.masses = masses.unsqueeze(-1).expand_as(X).clone()  # (n_batch, n_atom, n_dim)
+        else:
+            raise TypeError(f'Expected masses is a Sequence[Sequence[...]], but occurred {type(elements)}.')
+        self.v = th.zeros_like(X, device=self.device)
+
+        _results = super().run(
+            func,
+            X,
+            grad_func,
+            func_args,
+            func_kwargs,
+            grad_func_args,
+            grad_func_kwargs,
+            is_grad_func_contain_y=is_grad_func_contain_y,
+            require_grad=require_grad,
+            output_grad=output_grad,
+            fixed_atom_tensor=fixed_atom_tensor,
+            batch_indices=batch_indices,
+        )
+
+        return _results
+