BUCToolkit 1.0b0__py3-none-any.whl

BUCToolkit/BatchGenerate/CheckStructures.py

@@ -0,0 +1,198 @@
+"""
+Check Structures by atomic distance, surface smoothness, and coordination.
+"""
+#  Copyright (c) 2024-2025.7.4, BUCToolkit.
+#  Authors: Pu Pengxin, Song Xin
+#  Version: 0.9a
+#  File: CheckStructures.py
+#  Environment: Python 3.12
+
+import time
+import warnings
+from typing import Literal
+
+from BUCToolkit.BatchStructures.BatchStructuresBase import BatchStructures
+import numpy as np
+
+class CheckStructures:
+    """
+
+    """
+    def __init__(self, batch_structures: BatchStructures):
+        assert isinstance(batch_structures, BatchStructures), (f'`batch_structures` must be a BatchStructures type, '
+                                                               f'but occurred {type(batch_structures)}.')
+        self.BS = batch_structures
+
+    def check_distance(self, radius_threshold: float=1):
+        """
+        check atomic pairwise distances that are too close (less than given `radius_threshold`)
+        Args:
+            radius_threshold: float, the threshold radius.
+        Returns:
+            mask_array: ndarray[(n_batch, ), bool], where False for failed to the check.
+        """
+        self.BS.direct2cartesian()
+        coords = self.BS.Coords
+        mask = np.full(len(coords), True, )
+
+        for i, coo in enumerate(coords):
+            if self.BS.Dist_mat is None:
+                dist_mat = np.linalg.norm(coo[:, None, :] - coo[None, :, :], axis=-1)
+            else:
+                dist_mat = self.BS.Dist_mat[i]
+            # exclude the diag. part
+            diag_indx = np.arange(len(coo))
+            dist_mat[diag_indx, diag_indx] = 20.
+
+            if np.min(dist_mat) < radius_threshold:
+                mask[i] = False
+
+        return mask
+
+    @staticmethod
+    #@jit(nopython=True)
+    def _k_means(x:np.ndarray, c:np.ndarray, maxiter:int=100, thres:float=1e-4):
+        """
+        k-mean clustering algorithm
+        Args:
+            x: (n_atom, ), sample points
+            c: (nc, ), cluster center points
+            maxiter: max iteration numbers
+
+        Returns:
+            variance
+        """
+        n_c = len(c)
+        nc_indx = np.arange(n_c)
+        c_old = c.copy()
+        assert maxiter >= 1, f'Invalid `maxiter` value {maxiter}.'
+        for i in range(maxiter):
+            # calc. dist.
+            dist = np.abs(x[:, None] - c[None, :])  # (n_atom, nc)
+            min_dist_indx = np.argmin(dist, -1)  # (n_atom, )
+            # update center
+            _mask = (min_dist_indx[:, None] == nc_indx)  # (n_atom, nc)
+            n_samp_in_cluster = np.sum(_mask, axis=0)  # (n_c, )
+            drop_mask = n_samp_in_cluster == 0  # drop the non-sample cluster
+            c = np.where(drop_mask, 0., np.sum(x[:, None] * _mask, axis=0)/(n_samp_in_cluster + 1e-20))  # (n_atom, 1) * (n_atom, nc) -mean-> (nc, )
+            # delete the empty cluster
+            if np.any(drop_mask):
+                select_indx = np.where(~drop_mask)[0]
+                c = c[select_indx]
+                c_old = c_old[select_indx]
+                _mask = _mask[:, select_indx]
+                n_c = len(c)
+                nc_indx = np.arange(n_c)
+            # judge threshold
+            if np.linalg.norm(c - c_old) < thres:
+                return c, np.max(np.sum((x[:, None] - c[None, :])**2 * _mask, axis=0)/np.sum(_mask, axis=0))  # return the cluster points and metrics for smoothness
+            else:
+                c_old = c.copy()
+
+        #warnings.warn(f'K-means does not converge in {maxiter} steps.', RuntimeWarning)
+        return c, np.max(np.sum((x[:, None] - c[None, :])**2 * _mask, axis=0)/np.sum(_mask, axis=0))
+
+    @staticmethod
+    def _spectrum_clustering(x: np.ndarray, r_cutoff: float=3.):
+        """
+        Diagonalize laplace matrix of surface and bottom of slab to determined subgraph number.
+        Args:
+            x: (n_atom, 3), coordinates of surf. or bott.
+            r_cutoff: threshold of radius.
+
+        Returns:
+            n: number of connected subgraphs.
+        """
+        dist = np.linalg.norm(x[:, None] - x[None, :], axis=-1) + (r_cutoff + 10.) * np.eye(len(x))  # exclude the diag.
+        adj = np.where(dist < r_cutoff, 1., 0.)
+        degree = np.diag(np.sum(adj, axis=-1))
+        laplace = degree - adj
+        eigs = np.linalg.eigvalsh(laplace)
+        n_sub = np.sum(np.abs(eigs - 0.) < 1e-7)
+        return n_sub
+
+    def _periodic_reformat(self, cell, coords):
+        """
+        Transition the atom moved out of a cell into the periodic position in the cell.
+        Args:
+            cell:
+            coords:
+
+        Returns:
+
+        """
+        # TODO
+        raise NotImplementedError
+
+    def check_surf_smoothness(
+            self,
+            smoothness_threshold: float = 0.5,
+            r_cutoff: float = 3.,
+            surf_check_method: Literal['ConnectedGraph', 'C', 'MinDistance', 'M']|None = 'MinDistance'
+    ):
+        """
+        check whether the surface smooth enough (whether the max difference among surface atoms along z axis < `smoothness_threshold`)
+        Args:
+            smoothness_threshold: smoothness threshold.
+            r_cutoff: the cut-off distance to form the graph or inter-atomic distance checking.
+            surf_check_method: method of surface checking.
+                None: no check will applied;
+                'ConnectedGraph' or 'C': it will check whether the surface atoms form only one connected graph. If not, mask returns False.
+                'MinDistance' or 'M': it will check for atoms whose shortest distance from other atoms exceeds the cut-off distance.
+
+        Returns:
+            mask_array: ndarray[(n_batch, ), bool], where False for failed to the check.
+        """
+        self.BS.direct2cartesian()
+        coords = self.BS.Coords
+        mask = np.full(len(coords), True, )
+        std_val = 1e8
+        for i, coo in enumerate(coords):
+            # check empty coordinates
+            if len(coo) == 0:
+                warnings.warn(f'Occurred empty coordinates in {i}th structure, skipped.', RuntimeWarning)
+                continue
+            # convert to float64
+            coo = coo.astype(np.float64)
+            sorted_coo = np.sort(coo[:, -1])
+            std_val_old = np.inf
+            tol_cont = 0  # tolerance step of std_val not decrease.
+
+            for j in range(len(coo)):
+                _centers = np.array([np.mean(_) for _ in np.array_split(sorted_coo, j+1)])
+                _centers, std_val = self._k_means(coo[:, -1], _centers)
+                # raise for 3 previous steps to break the loop
+                if (std_val > std_val_old) or (len(_centers) < j+1):
+                    tol_cont += 1
+                else:
+                    tol_cont = 0
+                if (tol_cont > 2) or (np.min(np.abs(_centers[:, None] - _centers[None, :]) + 20.*np.eye(len(_centers))) < smoothness_threshold):
+                    break
+                elif std_val < 1e-5:  # completely converged
+                    std_val_old = std_val
+                    centers = _centers
+                    break
+                elif tol_cont == 0:
+                    std_val_old = std_val
+                    centers = _centers
+            if std_val > smoothness_threshold:
+                mask[i] = False
+            else:
+                # check surface sparseness
+                surf_mask = np.abs(coo[:, -1] - np.max(centers)) < smoothness_threshold
+                bott_mask = np.abs(coo[:, -1] - np.min(centers)) < smoothness_threshold
+                surf_coo = coo[surf_mask]
+                bott_coo = coo[bott_mask]
+                if surf_check_method == ('ConnectedGraph' or 'C'):
+                    n_sub_surf = self._spectrum_clustering(surf_coo, r_cutoff=r_cutoff)
+                    n_sub_bott = self._spectrum_clustering(bott_coo, r_cutoff=r_cutoff)
+                    if (n_sub_surf > 1) or (n_sub_bott > 1):
+                        mask[i] = False
+                elif surf_check_method == ('MinDistance'or 'M'):
+                    surf_dist = np.linalg.norm(surf_coo[:, None] - surf_coo[None, :], axis=-1) + (r_cutoff + 10.) * np.eye(len(surf_coo))  # exclude the diag.
+                    bott_dist = np.linalg.norm(bott_coo[:, None] - bott_coo[None, :], axis=-1) + (r_cutoff + 10.) * np.eye(len(bott_coo))
+                    if (not np.all(np.sum(surf_dist < r_cutoff, axis=-1))) or (not np.all(np.sum(bott_dist < r_cutoff, axis=-1))):
+                        mask[i] = False
+        return mask
+
+

BUCToolkit/BatchGenerate/__init__.py

@@ -0,0 +1,6 @@
+#  Copyright (c) 2024-2025.7.4, BUCToolkit.
+#  Authors: Pu Pengxin, Song Xin
+#  Version: 0.9a
+#  File: __init__.py
+#  Environment: Python 3.12
+

BUCToolkit/BatchGenerate/coords_interp.py

@@ -0,0 +1,127 @@
+"""
+Interpolation from 2 structures
+"""
+
+#  Copyright (c) 2024-2025.10.10, BUCToolkit.
+#  Authors: Pu Pengxin, Song Xin
+#  Version: 0.9a
+#  File: coords_interp.py
+#  Environment: Python 3.12
+
+from typing import List
+import numpy as np
+import torch as th
+
+
+def linear_interpolation(
+        coords1: np.ndarray | List,
+        coords2: np.ndarray | List,
+        n_points: int,
+
+):
+    """
+    Linear interpolation between `coords1` and `coords2` to generate a series of coordinates, usually used for NEB-like methods.
+    Args:
+        coords1: the coordinate of 1st structure.
+        coords2: the coordinate of 2nd structure.
+        n_points: the number of interpolation points.
+
+    Returns:
+        np.ndarray[(n_points, n_atoms, 3)], an array of interpolated structures' coordinates.
+    """
+    # check vars
+    if isinstance(coords1, List):
+        coords1 = np.asarray(coords1, dtype=np.float32)
+    elif not isinstance(coords1, np.ndarray):
+        raise TypeError(f'Invalid type of `coord1`: {type(coords1)}')
+    if isinstance(coords2, List):
+        coords2 = np.asarray(coords2, dtype=np.float32)
+    elif not isinstance(coords2, np.ndarray):
+        raise TypeError(f'Invalid type of `coord2`: {type(coords2)}')
+
+    # main
+    points_indx = np.linspace(0, 1, n_points)  # n_pt
+    coord_diff = coords2 - coords1  # n_atom, 3
+    interp_struc = coords1[None, :, :] + points_indx[:, None, None] * coord_diff[None, :, :]  # (1, n_atom, 3) + (n_pt, 1, 1) * (1, n_atom, 3)
+
+    return interp_struc
+
+def linear_interpolation_tens(
+        coords1: th.Tensor,
+        coords2: th.Tensor,
+        n_points: int,
+        atom_mask: th.Tensor|None = None,
+):
+    """
+    The linear interpolation between `coords1` and `coords2` for torch.Tensor inputs. See `linear_interpolation`.
+    Args:
+        coords1: the coordinate of 1st structure.
+        coords2: the coordinate of 2nd structure.
+        n_points: the number of interpolation points.
+        atom_mask: the mask of `coords1` and `coords2` that fixes atoms not to move. wherein, 1 is for free and 0 is for fixation.
+
+    Returns:
+        th.Tensor[(n_points, n_atoms, n_dim)], an array of interpolated structures' coordinates.
+    """
+    # check
+    if atom_mask is None:
+        atom_mask = th.ones_like(coords1)[None, :, :]
+    elif (isinstance(atom_mask, th.Tensor)) or isinstance(atom_mask, np.ndarray):
+        if atom_mask.shape != coords1.shape:
+            raise TypeError(f'Invalid shape of `atom_mask`: {atom_mask.shape}. It should be the same as `coord1`: {coords1.shape}.')
+        atom_mask = atom_mask[None, :, :]
+    else:
+        raise TypeError(f'Invalid type of `atom_mask`: {type(atom_mask)}')
+    # main
+    points_indx = th.linspace(0, 1, n_points, device=coords1.device, dtype=coords1.dtype)  # n_pt
+    coord_diff = coords2 - coords1  # n_atom, 3
+    interp_struc = coords1[None, :, :] + points_indx[:, None, None] * coord_diff[None, :, :] * atom_mask  # (1, n_atom, 3) + (n_pt, 1, 1) * (1, n_atom, 3)
+
+    return interp_struc
+
+def direction_for_finite_diff(
+        coords1: th.Tensor | np.ndarray,
+        coords2: th.Tensor | np.ndarray,
+        atom_mask: th.Tensor | np.ndarray | None = None,
+        interp_position: float = 0.5,
+):
+    """
+    Generate a direction for finite difference. Usually used for Dimer/Lanczos algo.
+    Args:
+        coords1: the initial coordinate.
+        coords2: the final coordinate.
+        atom_mask: the mask of `coords1` and `coords2` that fixes atoms not to move. wherein, 1 is for free and 0 is for fixation.
+        interp_position: the interpolation position of returned interp. coordinate. It must be between 0 and 1,
+            and gives the coordinate of `coords1 + interp_position * (coords2 - coords1)`.
+
+    Returns:
+        interp_coo: th.Tensor[(n_points, n_atoms, n_dim)], an array of interpolated structures' coordinates.
+        diff_coo  : the normalized difference between `coords1` and `coords2`, i.e., a unit vector.
+    """
+    # check
+    if not isinstance(coords1, th.Tensor):
+        dev = 'cpu'
+    else:
+        dev = coords1.device
+    if atom_mask is None:
+        atom_mask = th.ones_like(coords1)[None, :, :]
+    elif isinstance(atom_mask, th.Tensor) or isinstance(atom_mask, np.ndarray):
+        if atom_mask.shape != coords1.shape:
+            raise TypeError(f'Invalid shape of `atom_mask`: {atom_mask.shape}. It should be the same as `coord1`: {coords1.shape}.')
+        atom_mask = atom_mask[None, :, :]
+        atom_mask = th.as_tensor(atom_mask, dtype=atom_mask.dtype, device=dev)
+    else:
+        raise TypeError(f'Invalid type of `atom_mask`: {type(atom_mask)}')
+
+    coords1 = th.as_tensor(coords1, dtype=coords1.dtype, device=dev)
+    coords2 = th.as_tensor(coords2, dtype=coords2.dtype, device=dev)
+
+    interp_position = float(interp_position)
+    if interp_position > 1. or interp_position < 0.:
+        raise ValueError(f"`interp_position` should be between 0 and 1.")
+    diff_coo = (coords2 - coords1) * atom_mask
+    interp_coo = th.add(coords1, diff_coo, alpha=interp_position)
+    th.div(diff_coo, th.linalg.norm(diff_coo, dim=(-2, -1), keepdim=True), out=diff_coo)
+
+    return interp_coo, diff_coo
+

BUCToolkit/BatchGenerate/initial_ads_structures.py

@@ -0,0 +1,351 @@
+#  Copyright (c) 2024-2025.7.4, BUCToolkit.
+#  Authors: Pu Pengxin, Song Xin
+#  Version: 0.9a
+#  File: initial_ads_structures.py
+#  Environment: Python 3.12
+
+import copy
+from typing import List, Literal, Dict
+import warnings
+
+import torch as th
+from torch import nn
+import numpy as np
+
+from BUCToolkit.BatchStructures.BatchStructuresBase import BatchStructures
+from BUCToolkit.utils.IrregularTensorReformat import IrregularTensorReformat
+from BUCToolkit.BatchOptim.minimize import CG, QN, FIRE
+
+
+class InitAdsStructure:
+    def __init__(
+            self,
+            slabs: BatchStructures,
+            ads: th.Tensor,
+            ads_elements: Dict,
+            vaccuum_direction: Literal['x', 'y', 'z'] = 'z',
+            device: str | th.device='cpu',
+            verbose:int = 0):
+        """
+        Generate initial guesses of surface adsorption structures.
+        Args:
+            slabs: batch structures of slabs.
+            ads: Tensor(n_atom, 3), coordinates of the adsorbate.
+            ads_elements: Dict(element symbol: atom number).
+            vaccuum_direction: direction of vacuum layers.
+            device: the device that program run on.
+            verbose: to control the verboseness of output information.
+        """
+        self.slabs = slabs
+        self.ads = ads.to(device)
+        self.ads_elem = ads_elements
+        self.verbose = verbose
+        self.device = device
+        n_batch = len(slabs)
+        if not isinstance(self.slabs, BatchStructures):
+            raise TypeError(f'`slabs` must be BatchStructures, but occurred {type(slabs)}')
+        if not isinstance(self.ads, th.Tensor):
+            raise TypeError(f'`ads` must be torch.Tensor, but occurred {type(ads)}')
+        if sum(ads_elements.values()) != len(self.ads):
+            raise ValueError(f'Number of atoms in `ads_elements` does not match it in `ads`.')
+
+        # vacuum dirct
+        vac_d: dict = {'x': 0, 'y': 1, 'z': 2}
+        vac_d: int = vac_d[vaccuum_direction]
+
+        # dist. mat. of ads. mol.
+        self.dist_mol = th.linalg.norm(self.ads.unsqueeze(0) - self.ads.unsqueeze(1), dim=-1)  # (n_atom, n_atom)
+        max_atom_len = th.max(self.dist_mol)
+        # coord of slabs
+        # TODO: Convert to Cartesian coordinates.
+        self.slab_coords = [th.from_numpy(_).to(self.device) for _ in self.slabs.Coords]
+        self.slab_cells = [th.from_numpy(_).to(self.device) for _ in self.slabs.Cells]
+        # check vacuum length
+        vacuum_length_pak = [(self.slab_cells[i][vac_d, vac_d] - th.max(self.slab_coords[i][:, vac_d]), th.max(self.slab_coords[i][:, vac_d]))  for i in range(len(slabs))]
+        vacuum_length, highest_slab = zip(*vacuum_length_pak)
+        highest_slab= th.tensor(highest_slab)  # (n_batch, )
+        for i, _vac_len in enumerate(vacuum_length):
+            if _vac_len < 3.:
+                raise RuntimeError(f'Occurred unacceptable vacuum length {_vac_len} in {i}th structure. I HOPE YOU KNOWN WHAT YOU ARE DOING.')
+            elif _vac_len < 9.:
+                warnings.warn(f'The vacuum length of {i}th structure was too short: {vacuum_length}.')
+            elif (_vac_len < (max_atom_len + 5.)) or th.any(th.linalg.norm(self.slab_cells[i], dim=1) - 5. < max_atom_len):
+                warnings.warn(f'The {i}th cell are too small that distance between 2 adsorbates in neighbor lattice might be less than 5 Angstrom.')
+        # set the lowest atom in ads to 3 Angstrom higher than the highest atom in slab
+        lowest_mol = th.min(self.ads[:, vac_d], dim=-1).values  # (n_batch, )
+        self.ads = (self.ads.unsqueeze(0)).expand(n_batch, -1, -1)  # (n_batch, n_mol_atom, n_dim)
+        self.ads = self.ads + (highest_slab + 3. - lowest_mol).unsqueeze(-1).unsqueeze(-1)
+        pass
+
+    def run(
+            self,
+            slab_site: List[List[int]],
+            ads_site: List[int],
+            bond_length: float,
+            bond_coeff: float = 5.,
+            non_bond_coeff: float = 3.,
+            keep_ads_conf_coeff: float = 2.
+    ):
+        """
+
+        Args:
+            slab_site: The indices of adsorption site in slabs.
+            ads_site: The indices of adsorption site in adsorbate.
+            bond_length: The target bond length of adsorbate and slabs.
+            bond_coeff: Strength coefficient of ads-slab bond.
+            non_bond_coeff: Strength coefficient of repulsiveness of non-bond atom between ads. and slab.
+            keep_ads_conf_coeff: Strength coefficient of keeping original adsorbate configuration.
+
+        Returns:
+            BatchStructures: the batch structures of slabs with ads.
+
+        """
+        regularizor = IrregularTensorReformat()
+        slab_pad_coords, pad_mask = regularizor.regularize(self.slab_coords, 0.)
+        n_batch, n_atom, n_dim = slab_pad_coords.shape
+        # slabs
+        if isinstance(slab_site, List):
+            assert len(slab_site) == n_batch, f'number of specified adsorption sites {len(slab_site)} does not match the number of structures {n_batch}.'
+            #assert len(bond_length) == n_batch, f'number of specified adsorption sites {len(slab_site)} does not match the number of structures {n_batch}.'
+            slab_mask = th.full_like(slab_pad_coords, False, device=self.device, dtype=th.bool)  # (n_batch, n_atom, n_dim)
+            #target_bond_length = th.full((n_batch, n_atom), 0., device=self.device)  # (n_batch, n_atom) # TODO auto-adaptive bond length
+            target_bond_length = bond_length
+            for i, _sites in enumerate(slab_site):
+                slab_mask[i, _sites] = True
+                #target_bond_length[i, _sites] = th.tensor(bond_length[i], device=self.device)  # TODO auto-adaptive bond length
+            slab_mask = slab_mask * pad_mask
+        else:
+            raise TypeError('`slab_site` must be a List[List[int]].')
+        # ads molecules
+        if isinstance(ads_site, List):
+            ads_mask = th.full_like(self.ads[0], False, device=self.device, dtype=th.bool)  # (n_mol_atom, n_dim)
+            ads_mask[ads_site] = True
+        else:
+            raise TypeError('`ads_site` must be a List[int].')
+        # calc. force
+        loss_fn = _Loss()
+        #self._random_perturbate()
+
+        # optimize
+        optimizer = FIRE(
+            1e-3,
+            0.1,
+            200,
+            0.01,
+            device=self.device,
+            verbose=self.verbose
+        )
+        min_loss, X_ads = optimizer.run(
+            loss_fn,
+            self.ads,
+            None,
+            (slab_pad_coords, slab_mask, ads_mask, self.dist_mol, target_bond_length, bond_coeff, non_bond_coeff, keep_ads_conf_coeff),
+            None,
+        )
+
+        ads_structures = self._integrate_coords(X_ads)
+        return ads_structures
+
+    def _integrate_coords(self, X_ads: th.Tensor) -> BatchStructures:
+        """
+        Integrate coordinates of adsorption and slabs.
+        Args:
+            X_ads: coordinates of ads.
+
+        Returns: BatchStructure of adsorbed slabs.
+
+        """
+        ads_elem = list(self.ads_elem.keys())
+        ads_number = [self.ads_elem[_] for _ in ads_elem]
+        f = copy.deepcopy(self.slabs)
+        X_ads = X_ads.numpy(force=True)
+        for i in range(len(f)):
+            f.Coords[i] = np.concatenate([f.Coords[i], X_ads[i]])
+            f.Fixed[i] = np.concatenate([f.Fixed[i], np.ones_like(X_ads[i])])
+            f.Elements[i].extend(ads_elem)
+            f.Numbers[i].extend(ads_number)
+        if f.Atom_list is not None:
+            f.generate_atom_list(True)
+        if f.Atomic_number_list is not None:
+            f.generate_atomic_number_list(True)
+        return f
+
+    def _random_perturbate(self, ):
+        """
+        Ramdom perturbate adsorptions by orthogonal and translational transform to search more configuration.
+        Returns:
+
+        """
+        ln = len(self.ads)
+        rx = th.rand(ln) * 2 * th.pi
+        ry = th.rand(ln) * 2 * th.pi
+        rz = th.rand(ln) * 2 * th.pi
+        tx = th.randn(ln) * 0.05
+        ty = th.randn(ln) * 0.05
+        tz = th.randn(ln) * 0.05
+
+        self.ads = self._distance_preserv_transform(rx, ry ,rz, tx, ty, tz, self.ads)
+
+    def _distance_preserv_transform(
+            self,
+            rx: th.Tensor,
+            ry: th.Tensor,
+            rz: th.Tensor,
+            tx: th.Tensor,
+            ty: th.Tensor,
+            tz: th.Tensor,
+            X: th.Tensor
+    ) -> th.Tensor:
+        """
+        Orthogonal and transitional transform on input X.
+        Args:
+            rx: rotation degree around X-axis.
+            ry: rotation degree around Y-axis.
+            rz: rotation degree around Z-axis.
+            tx: transition along X-axis.
+            ty: transition along Y-axis.
+            tz: transition along Z-axis.
+            X: the coordinated to transform.
+
+        Returns:
+
+        """
+        ln = len(X)
+        if not (len(rx) == len(ry) == len(rz) == len(tx) == len(ty) == len(tz) == ln):
+            raise RuntimeError(f'Incompatible shape of rx~rz & tx~tz & X. They must have same batch number.')
+        # mass center
+        center = th.sum(X, dim=1, keepdim=True)  # (n_batch, 1, 3)
+        # rotation matrix
+        rot_x = th.tensor(
+            [
+                [[1.]*ln,             [0.]*ln,              [0.]*ln],
+                [[0.]*ln, th.cos(rx).tolist(), th.sin(-rx).tolist()],
+                [[0.]*ln, th.sin(rx).tolist(),  th.cos(rx).tolist()]
+            ],
+            device=self.device
+        )
+        rot_y = th.tensor(
+            [
+                [th.cos(ry).tolist(), [0.]*ln,  th.sin(ry).tolist()],
+                [[0.]*ln,             [1.]*ln,              [0.]*ln],
+                [th.sin(-ry).tolist(), [0.]*ln, th.cos(ry).tolist()]
+            ],
+            device=self.device
+        )
+        rot_z = th.tensor(
+            [
+                [th.cos(rz).tolist(), th.sin(-rz).tolist(), [0.]*ln],
+                [th.sin(rz).tolist(),  th.cos(rz).tolist(), [0.]*ln],
+                [            [0.]*ln,              [0.]*ln, [1.]*ln]
+            ],
+            device=self.device
+        )
+        rot_x = rot_x.transpose(0, -1)
+        rot_y = rot_y.transpose(0, -1)
+        rot_z = rot_z.transpose(0, -1)
+        trans = th.cat((tx.unsqueeze(-1).unsqueeze(-1), ty.unsqueeze(-1).unsqueeze(-1), tz.unsqueeze(-1).unsqueeze(-1)), dim=-1)  # (n_batch, 1, 3)
+        X_trans = rot_z @ rot_y @ rot_x @ (X - center) + center + trans # (n_batch, n_atom, 3)
+
+        return X_trans
+
+
+def get_highest_atom_indices(structure: BatchStructures, vaccuum_direction: Literal['x', 'y', 'z'] = 'z',):
+    """
+    Get the highest atom index of given batch structures.
+    Returns:
+        List[List[int]], the indices of highest atom in each structure
+    """
+    # vacuum dirct
+    vac_d: dict = {'x': 0, 'y': 1, 'z': 2}
+    vac_d: int = vac_d[vaccuum_direction]
+    # coords
+    converter = IrregularTensorReformat()
+    padded_coords = [th.from_numpy(_) for _ in structure.Coords]
+    padded_coords, pad_mask = converter.regularize(padded_coords, -th.inf)  # (n_batch, n_atom, 3)
+    indices = th.argmax(padded_coords[:, :, vac_d], dim=1)
+
+    return indices.unsqueeze(-1).tolist()
+
+
+class _Loss(nn.Module):
+    def __init__(self):
+        super().__init__()
+        pass
+
+    def forward(self, X_ads, X_slab, slab_mask, ads_mask, mol_origin_dist_mat, target_bond_len, c_b, c_nb, c_k):
+        """
+
+        Args:
+            X_ads: (n_batch, n_mol_atom, n_dim)
+            X_slab: (n_batch, n_atom, n_dim)
+            slab_mask: (n_batch, n_atom, n_dim)
+            ads_mask: (n_mol_atom, n_dim)
+            mol_origin_dist_mat: (n_mol_atom, n_mol_atom)
+            target_bond_len: float
+            c_b:
+            c_nb:
+            c_k:
+
+        Returns:
+
+        """
+        loss = (c_b * self._bond_force_field(X_ads, X_slab, slab_mask, ads_mask, target_bond_len) +
+                c_nb * self._non_bond_loss(X_ads, X_slab, slab_mask, ads_mask) +
+                c_k * self._keep_origin_loss(X_ads, mol_origin_dist_mat))
+        return loss
+
+    @staticmethod
+    def _bond_force_field(X_ads, X_slab, slab_mask: th.Tensor, ads_mask: th.Tensor, target_bond_len: float):
+        """
+
+        Args:
+            X_ads:
+            X_slab:
+            slab_mask:
+            ads_mask:
+
+        Returns:
+
+        """
+        # (n_batch, n_atom, 1, n_dim) - (n_batch, 1, n_mol_atom, n_dim) -> (n_batch, n_atom, n_mol_atom, n_dim)
+        dist_nb = th.linalg.norm(X_slab.unsqueeze(2) - X_ads.unsqueeze(1), dim=-1)
+        # quadratic potential
+        b_loss = (dist_nb - target_bond_len)**2 # (n_batch, n_atom, n_mol_atom)
+        # tol mask
+        tol_mask = (slab_mask.unsqueeze(-2) * ads_mask.unsqueeze(0).unsqueeze(0))[..., 0]  # (n_batch, n_atom, 1, n_dim) * (1, 1, n_mol_atom, n_dim) -> (n_batch, n_atom, n_mol_atom, n_dim)
+        b_loss = th.sum(b_loss * tol_mask, dim=(-1, -2))
+        return b_loss
+
+    @staticmethod
+    def _non_bond_loss(X_ads, X_slab, slab_mask: th.Tensor, ads_mask: th.Tensor):
+        """
+
+        Args:
+            X_ads:
+            X_slab:
+            slab_mask: (n_batch, n_atom, n_dim)
+            ads_mask:  (n_mol_atom, n_dim)
+
+        Returns:
+
+        """
+        # (n_batch, n_atom, 1, n_dim) - (n_batch, 1, n_mol_atom, n_dim) -> (n_batch, n_atom, n_mol_atom, n_dim) -> (n_batch, n_atom, n_mol_atom)
+        dist_nb = th.linalg.norm(X_slab.unsqueeze(2) - X_ads.unsqueeze(1), dim=-1)
+        # exponential repulsive
+        n_b_loss = th.where(dist_nb < 3., th.exp(6./dist_nb) - 1, 0.)  # (n_batch, n_atom, n_mol_atom)
+        # polynomial repulsive
+        #n_b_loss = (th.linalg.norm(dist_nb, dim=-1) - 3.5)**3
+        #n_b_loss = th.where(n_b_loss < 0., 0., n_b_loss) # (n_batch, n_atom, n_mol_atom)
+        # half L-J repulsive
+        #n_b_loss = 5./(th.linalg.norm(dist_nb, dim=-1))**12  # (n_batch, n_atom, n_mol_atom)
+        # tol mask
+        tol_mask = (~slab_mask.unsqueeze(-2) * ~ads_mask.unsqueeze(0).unsqueeze(0))[..., 0]  # (n_batch, n_atom, 1, n_dim) * (1, 1, n_mol_atom, n_dim) -> (n_batch, n_atom, n_mol_atom, n_dim)
+        n_b_loss = th.sum(n_b_loss * tol_mask, dim=(-1, -2))
+        return n_b_loss
+
+    @staticmethod
+    def _keep_origin_loss(X_ads, origin_dist_mat):
+        dist_mat = th.linalg.norm(X_ads.unsqueeze(1) - X_ads.unsqueeze(2), dim=-1)  # (n_batch, n_mol_atom, n_mol_atom)
+        k_o_loss = th.sum((dist_mat - origin_dist_mat)**2)
+        return k_o_loss
+