BUCToolkit 1.0b0__py3-none-any.whl

BUCToolkit/io.py

@@ -0,0 +1,30 @@
+"""
+A fast API for all i/o operations
+"""
+#  Copyright (c) 2026.4.13, BUCToolkit.
+#  Authors: Pu Pengxin, Song Xin
+#  Version: 1.0b
+#  File: io.py
+#  Environment: Python 3.12
+
+from BUCToolkit.Preprocessing import load_files, write_files
+from BUCToolkit.Preprocessing.load_files import OUTCAR2Feat, Xyz2Feat, ExtXyz2Feat, POSCARs2Feat, ASETraj2Feat, Cif2Feat, load_from_structures
+from BUCToolkit.BatchStructures.StructuresIO import read_opt_structures, read_mc_traj, read_md_traj, ArrayDumpReader, ArrayDumper
+
+__all__ = [
+    'OUTCAR2Feat',
+    'Xyz2Feat',
+    'ExtXyz2Feat',
+    'POSCARs2Feat',
+    'ASETraj2Feat',
+    'Cif2Feat',
+    'read_opt_structures',
+    'read_md_traj',
+    'read_mc_traj',
+    'ArrayDumpReader',
+    'ArrayDumper',
+    'load_from_structures',
+    'load_files',
+    'write_files',
+]
+

BUCToolkit/utils/AtomSelector.py

@@ -0,0 +1,64 @@
+""" Selecting atoms by given conditions """
+#  Copyright (c) 2025.11.26, BUCToolkit.
+#  Authors: Pu Pengxin, Song Xin
+#  Version: 0.9a
+#  File: AtomSelector.py
+#  Environment: Python 3.12
+
+import torch as th
+from typing import Sequence, List, Literal
+from ._Element_info import ATOMIC_SYMBOL
+
+def atom_fix_selector(
+        mask: th.Tensor,
+        elements: th.Tensor,
+        coords: th.Tensor,
+        select_element: Sequence[str | int] | None = None,
+        select_height: float | None = None,
+        atom_index: List[List[float]] | None = None,
+        select_mode: Literal['fix', 'inv_fix', 'free'] = 'fix'
+) -> th.Tensor:
+    """
+    Selecting which atoms are fixed and updating the input mask accordingly.
+    Args:
+        mask:
+        elements:
+        coords:
+        select_element:
+        select_height:
+        atom_index:
+        select_mode:
+            fix: the atoms satisfied given conditions will be set to 0 in given mask, and others do not change.
+            inv_fix: only the atoms satisfied given conditions will be set to 1 in given mask, and others set to 0.
+            free: the atoms satisfied given conditions will be set to 1 in given mask, and others do not change.
+
+    Returns:
+
+    """
+    # check
+    if mask.shape != coords.shape:
+        raise ValueError(f'Expected `mask` has the same shape as `coords` ({coords.shape}), but got {mask.shape}.')
+    if elements.shape != (mask.shape[0], ):
+        raise ValueError(f'Expected `elements` has the same shape as the 1st dim of `mask` ({mask.shape[0]}), but got {elements.shape}.')
+
+    # initialize
+    FIX_DICT = {'fix': 0, 'inv_fix': 1, 'free': 1}
+    fill_elem = FIX_DICT[select_mode]
+    if select_mode == 'inv_fix':
+        _mask = th.zeros_like(mask).to(th.int8)
+    else:
+        _mask = th.where(mask.to(th.int8) == 0, 0, 1).to(th.int8)
+    elements = elements.unsqueeze(-1).broadcast_to(mask.shape)
+
+    # main
+    if select_element is not None:
+        sel_elem = [_ if isinstance(_, int) else ATOMIC_SYMBOL[_] for _ in select_element]
+        for elem in sel_elem:
+            _mask = th.where(elements == elem, fill_elem, _mask)
+    if select_height is not None:
+        _mask = th.where(coords[:, -1].unsqueeze(-1) <= select_height, fill_elem, _mask)
+    if atom_index is not None:
+        atom_index = th.tensor(atom_index, dtype=th.int64, device=_mask.device)
+        _mask[atom_index] = fill_elem
+
+    return _mask

BUCToolkit/utils/AtomicNumber2Properties.py

@@ -0,0 +1,39 @@
+""" Convert from atomic number to various properties, e.g., masses, elements."""
+
+#  Copyright (c) 2025.7.4, BUCToolkit.
+#  Authors: Pu Pengxin, Song Xin
+#  Version: 0.9a
+#  File: AtomicNumber2Properties.py
+#  Environment: Python 3.12
+
+from typing import Sequence, List
+from BUCToolkit.utils._Element_info import MASS, N_MASS, ATOMIC_SYMBOL, ATOMIC_NUMBER
+import numpy as np
+
+
+def atomic_numbers_to_masses(Z: Sequence[int]) -> List[float]:
+    """Convert atomic numbers to masses."""
+    mass = [N_MASS[int(_)] for _ in Z]
+    return mass
+
+
+def atomic_numbers_to_elements(Z: Sequence[int]) -> List[int]:
+    elem = [ATOMIC_NUMBER[int(_)] for _ in Z]
+    return elem
+
+
+def elements_to_atomic_numbers(Z: Sequence[str]) -> List[int]:
+    atomic_numbers = [ATOMIC_SYMBOL[str(_)] for _ in Z]
+    return atomic_numbers
+
+
+__all__ = [
+    "atomic_numbers_to_masses",
+    "atomic_numbers_to_elements",
+    "elements_to_atomic_numbers",
+]
+
+if __name__ == '__main__':
+    z = [1, 2, 5, 3, 66, 32]
+    m = atomic_numbers_to_masses(z)
+    pass

BUCToolkit/utils/ElemListReduce.py

@@ -0,0 +1,52 @@
+#  Copyright (c) 2025.7.4, BUCToolkit.
+#  Authors: Pu Pengxin, Song Xin
+#  Version: 0.9a
+#  File: ElemListReduce.py
+#  Environment: Python 3.12
+"""
+Convert a dense element list (atomic number corresponds to coordinates one by one) into VASP-like compressed format.
+"""
+from typing import List, Tuple
+from ._Element_info import ATOMIC_NUMBER, ATOMIC_SYMBOL
+import numpy as np
+
+
+def elem_list_reduce(
+        atom_list: List[str|int]|np.ndarray
+) -> Tuple[List[str], List[int], List[int]]:
+    """
+    Convert a dense element list (atomic number corresponds to coordinates one by one) into VASP-like compressed format.
+
+    Args:
+        atom_list:
+
+    Returns: element symbol list, atomic numbers list, atom number (count) list
+
+    """
+    if isinstance(atom_list, np.ndarray): atom_list = atom_list.tolist()
+    _elem_old = ATOMIC_NUMBER[atom_list[0]] if isinstance(atom_list[0], int) else atom_list[0]
+    elements = [_elem_old, ]
+    atomic_numbers = [ATOMIC_SYMBOL[_elem_old], ]
+    elem_numbers = list()
+    count_i = 0
+    for elem in atom_list:
+        if isinstance(elem, int):
+            elem_symb = ATOMIC_NUMBER[elem]
+            elem_numb = elem
+        elif isinstance(elem, str):
+            elem_symb = elem
+            elem_numb = ATOMIC_SYMBOL[elem]
+        else:
+            raise TypeError(f'Expected str or int in `atom_list`, but got {type(elem)}.')
+
+        if elem_symb == _elem_old:
+            count_i += 1
+        else:
+            _elem_old = elem_symb
+            elem_numbers.append(count_i)
+            elements.append(elem_symb)
+            atomic_numbers.append(elem_numb)
+            count_i = 1
+    elem_numbers.append(count_i)
+
+    return elements, atomic_numbers, elem_numbers

BUCToolkit/utils/Ewald.py

@@ -0,0 +1,210 @@
+#  Copyright (c) 2026.3.17, BUCToolkit.
+#  Authors: Pu Pengxin, Song Xin
+#  Version: 0.9a
+#  File: Ewald.py
+#  Environment: Python 3.12
+from typing import Callable, Any, List, Tuple, Optional, Dict
+import torch as th
+
+
+class GeneralizedEwald:
+    """
+    A Generalized Ewald Summation Algorithm.
+    It calculates as follows:
+        \sum_n f(r + n R) = \sum_n {f(r + n R) - (a/\pi)^1.5 \int {f(x) exp(-a (x - r - n R)^2) dx}}
+                            + 1/V \sum_k {g(k) exp(-k^2/(4 a) + i k \cdot r)}
+    """
+    def __init__(
+            self,
+            rfunc: Callable[[th.Tensor, Any, ...], th.Tensor],
+            gfunc: Callable[[th.Tensor, Any, ...], th.Tensor],
+            rfunc_gaussian_conv: Callable[[th.Tensor, Any, ...], th.Tensor] | None = None,
+            fft_grids: th.Tensor | List[int] | Tuple[int, int, int] | None = None,
+            alpha: float = 0.1,
+    ):
+        """
+
+        Args:
+            rfunc: the function in the real space.
+                Its first input should be a Tensor X with shape (B, N, 3), and output a Tensor with shape (B, N),
+                where `B` is the batch_size, and `N` is the number of feeler points within a sample.
+                For each N, it returns the function value of f at the given point X[B_i, N_i, :].
+            gfunc: the corresponding function in the reciprocal space. It has the same format as rfunc.
+            rfunc_gaussian_conv: the convolution between rfunc and Gaussian function,
+                i.e., rfunc_gaussian_conv(r) = (a/\pi)^1.5 \int {f(x) exp(-a (x - r)^2) dx}.
+                If None, a numeric FFT would be applied to calculate the reciprocal production instead of real-space convolution.
+            fft_grids: shape of (3, ). If `rfunc_gaussian_conv` is None, this argument is to define the FFT grids,
+                and gives the value at input point `r` by interpolation. Otherwise, this argument is ignored.
+            alpha: the coefficient of Gaussian exponential in Ewald summation.
+        """
+        self.rfunc = rfunc
+        self.gfunc = gfunc
+        self.rgc_func = rfunc_gaussian_conv
+        if rfunc_gaussian_conv is None:
+            fft_grids = th.as_tensor(fft_grids, dtype=th.int64)
+            if fft_grids.shape != (3, ):
+                raise ValueError(f"`fft_grids` must have the shape of (3, ), but got {fft_grids.shape}.")
+            if fft_grids.min() <= 0:
+                raise ValueError(f"FFT grids number must be positive, but got {fft_grids.tolist()}.")
+            self.fft_grids = fft_grids
+        else:
+            self.fft_grids = None
+
+        if alpha <= 1e-7:
+            raise ValueError(f"`alpha` must be positive, but got {alpha}.")
+        self.alpha = alpha
+
+    def run(
+            self,
+            r: th.Tensor,
+            R: th.Tensor,
+            n_r: int = 20,
+            n_g: int | None = None,
+    ):
+        """
+        Run the summation.
+        Args:
+            r: Tensor(B, N, 3), the coordinates of the feeler points.
+            R: Tensor(B, 3, 3), the real-space lattice of each sample in the batch. Each row is a lattice vector.
+            n_r: the number of real-space summation terms.
+            n_g: the number of reciprocal-space summation terms. If None, it will be set to `n_r`.
+
+        Returns: Tensor(B, N), the summations.
+
+        """
+        raise NotImplementedError
+
+    def _run_anal(
+            self,
+            r: th.Tensor,
+            R: th.Tensor,
+            rfunc_args: Tuple[Any, ...] = tuple(),
+            rfunc_kwargs: Dict[str, Any] | None = None,
+            gfunc_args: Tuple[Any, ...] = tuple(),
+            gfunc_kwargs: Dict[str, Any] | None = None,
+            n_r: int = 20,
+            n_g: Optional[int] = None,
+            chunk_size: int = 10000,  # 分块大小，避免内存溢出
+    ) -> th.Tensor:
+        """
+        Run the Ewald summation using analytical convolution function.
+
+        Args:
+            r: (B, N, 3) query points.
+            R: (B, 3, 3) real-space lattice matrices (each row is a lattice vector).
+            n_r: number of real-space cells in each direction (half-width). Total cells = (2*n_r+1)^3.
+            rfunc_args:
+            rfunc_kwargs:
+            gfunc_args:
+            gfunc_kwargs:
+            n_g: number of reciprocal-space cells in each direction. If None, set to n_r.
+            chunk_size: max number of real-space cells processed at once to control memory.
+
+        Returns:
+            (B, N) tensor of summed values.
+        """
+        if rfunc_kwargs is None:
+            rfunc_kwargs = {}
+        if gfunc_kwargs is None:
+            gfunc_kwargs = {}
+
+        device = r.device
+        B, N, _ = r.shape
+        n_g = n_g if n_g is not None else n_r
+
+        # ------------------------------------------------------------------
+        # 1. Real-space part: sum over lattice vectors nR
+        # ------------------------------------------------------------------
+        # Generate integer lattice indices (i,j,k) for real space
+        idx_r = th.arange(-n_r, n_r + 1, device=device)
+        grid_r = th.stack(th.meshgrid(idx_r, idx_r, idx_r, indexing='ij'), dim=-1)  # (Mr,3)
+        Mr = grid_r.shape[0]  # = (2*n_r+1)^3
+        grid_r = grid_r.view(-1, 3)  # (Mr, 3)
+
+        # Precompute offsets = i*R[:,0] + j*R[:,1] + k*R[:,2] for all batch items
+        # offsets shape: (B, Mr, 3)
+        offsets = th.einsum('mi,bij->bmj', grid_r, R)  # sum over i (grid dim), multiply with R rows
+
+        # Prepare real-space sum accumulator
+        real_sum = th.zeros(B, N, device=device)
+
+        # Process real-space cells in chunks to avoid OOM
+        for start in range(0, Mr, chunk_size):
+            end = min(start + chunk_size, Mr)
+            offsets_chunk = offsets[:, start:end, :]  # (B, chunk, 3)
+
+            # Shift query points by each offset: (B, N, chunk, 3)
+            r_shifted = r.unsqueeze(2) + offsets_chunk.unsqueeze(1)
+
+            # Reshape to (B, N*chunk, 3) for function evaluation
+            r_flat = r_shifted.reshape(B, N * (end - start), 3)
+
+            # Evaluate f and its Gaussian convolution
+            f_vals = self.rfunc(r_flat).reshape(B, N, end - start)  # (B, N, chunk)
+            conv_vals = self.rgc_func(r_flat).reshape(B, N, end - start)  # (B, N, chunk)
+
+            # Short-range contribution: f - conv, sum over chunk
+            real_sum += (f_vals - conv_vals).sum(dim=2)
+
+        # ------------------------------------------------------------------
+        # 2. Reciprocal-space part: sum over reciprocal vectors k
+        # ------------------------------------------------------------------
+        # Compute cell volume V = |det(R)|
+        det_R = th.det(R)  # (B,)
+        V = det_R.abs().unsqueeze(-1).unsqueeze(-1)  # (B,1,1) for broadcasting
+
+        # Compute reciprocal lattice vectors (rows of the matrix)
+        # b1 = 2π * (a2 × a3) / V, b2 = 2π * (a3 × a1) / V, b3 = 2π * (a1 × a2) / V
+        a1 = R[:, 0, :]  # (B,3)
+        a2 = R[:, 1, :]  # (B,3)
+        a3 = R[:, 2, :]  # (B,3)
+        cross_a2_a3 = th.cross(a2, a3, dim=-1)  # (B,3)
+        cross_a3_a1 = th.cross(a3, a1, dim=-1)  # (B,3)
+        cross_a1_a2 = th.cross(a1, a2, dim=-1)  # (B,3)
+
+        two_pi = 2 * th.pi
+        b1 = two_pi * cross_a2_a3 / V.squeeze(-1)  # (B,3)
+        b2 = two_pi * cross_a3_a1 / V.squeeze(-1)
+        b3 = two_pi * cross_a1_a2 / V.squeeze(-1)
+
+        # Stack to form reciprocal matrix (B,3,3) where each row is a reciprocal vector
+        recip_matrix = th.stack([b1, b2, b3], dim=1)  # (B,3,3)
+
+        # Generate integer indices for reciprocal space
+        idx_g = th.arange(-n_g, n_g + 1, device=device)
+        grid_g = th.stack(th.meshgrid(idx_g, idx_g, idx_g, indexing='ij'), dim=-1)  # (Mg,3)
+        Mg = grid_g.shape[0]
+        grid_g = grid_g.view(-1, 3)  # (Mg,3)
+
+        # Compute all k vectors: (B, Mg, 3)
+        k_vec = th.einsum('mi,bij->bmj', grid_g, recip_matrix)
+
+        # Compute k^2 = ||k||^2
+        k_sq = (k_vec ** 2).sum(dim=-1)  # (B, Mg)
+
+        # Gaussian attenuation factor exp(-k^2/(4α))
+        atten = th.exp(-k_sq / (4 * self.alpha))  # (B, Mg)
+
+        # Evaluate g(k) for all k points
+        # gfunc expects input (B, N, 3), but here we have (B, Mg, 3). Mg may be large, so we call in one go.
+        g_vals = self.gfunc(k_vec)  # (B, Mg) – assuming gfunc returns real or complex tensor
+
+        # Compute phase factor exp(i k·r) for all query points and all k
+        # k·r: (B, Mg, 3) dot (B, N, 3) -> (B, N, Mg)
+        k_dot_r = th.einsum('bmj,bnj->bnm', k_vec, r)  # (B, N, Mg)
+        phase = th.exp(1j * k_dot_r)  # complex (B, N, Mg)
+
+        # Combine: g(k) * attenuation * phase, sum over k, and divide by V
+        # g_vals may be real or complex; we promote to complex if needed.
+        if not th.is_complex(g_vals):
+            g_vals = g_vals.to(th.complex64)
+        contrib = g_vals.unsqueeze(1) * atten.unsqueeze(1) * phase  # (B, N, Mg)
+        recip_sum = contrib.sum(dim=2) / V  # (B, N) complex
+
+        # Take real part (should be real due to symmetry, but numerically small imaginary may remain)
+        recip_sum = recip_sum.real
+
+        # ------------------------------------------------------------------
+        # 3. Combine and return
+        # ------------------------------------------------------------------
+        return real_sum + recip_sum

BUCToolkit/utils/IrregularTensorReformat.py

@@ -0,0 +1,59 @@
+""" regularize Sequence of Tensors by padding """
+#  Copyright (c) 2024-2025.7.4, BUCToolkit.
+#  Authors: Pu Pengxin, Song Xin
+#  Version: 0.9a
+#  File: IrregularTensorReformat.py
+#  Environment: Python 3.12
+
+from typing import List, Tuple
+
+import torch as th
+
+
+class IrregularTensorReformat:
+    """
+    Regularize List[Tensor] into one Tensor by padding, and then recover them into List[Tensor].
+    The input tensors can only be irregular at 1st dim.
+    """
+    def __init__(self, ):
+        self._length_record = None
+        self.__shape_check = None
+
+    def regularize(self, X: List[th.Tensor], padding_value:float=0.) -> Tuple[th.Tensor, th.Tensor]:
+        """
+        Regularize the List of Tensors.
+        Args:
+            X: the input sequence of Tensors.
+            padding_value: the padding_value.
+
+        Returns:
+            X_r: Tensor, the regularized tensor.
+            mask: Tensor, the mask of paddings. The part of `False` is paddings.
+        """
+        if th.any(th.isnan(th.cat(X, dim=0))):
+            raise ValueError('Occurred NaN in input `X`.')
+        self._length_record = [len(xi) for xi in X]
+        # padding
+        X = th.nn.utils.rnn.pad_sequence(X, True, padding_value=padding_value)
+        pad_mask = ~th.isnan(X)
+        # shape check
+        self.__shape_check = len(X)
+
+        return X, pad_mask
+
+    def recover(self, X:th.Tensor) -> List[th.Tensor]:
+        """
+        Recover the regularized tensor into List[Tensor]
+        Args:
+            X: the Tensor that have same shape of X
+        Returns:
+            List[Tensor], the recovered Tensors.
+        """
+        if self._length_record is None:
+            raise RuntimeError('`X` has not regularized yet.')
+        elif len(X) != self.__shape_check:
+            raise ValueError(f'Wrong shape of `X`. The length of `X` should be {self.__shape_check}, but occurred {len(X)}.')
+
+        X = [X[i, :self._length_record[i]] for i in range(self.__shape_check)]
+
+        return X

BUCToolkit/utils/_CheckModules.py

@@ -0,0 +1,30 @@
+""" check whether imported module exists """
+#  Copyright (c) 2024-2025.7.4, BUCToolkit.
+#  Authors: Pu Pengxin, Song Xin
+#  Version: 0.9a
+#  File: _CheckModules.py
+#  Environment: Python 3.12
+
+import importlib
+import warnings
+
+
+# Check modules func
+def check_module(module_name: str, pkg_name: None | str = None):
+    """
+    Check whether imported module exists.
+    Args:
+        module_name: imported module name.
+        pkg_name: package name. It is required when performing a relative import.
+                  It specifies the package to use as the anchor point from which to resolve the relative import to an absolute import.
+
+    Returns:
+        the imported module if it exists, else return None.
+
+    """
+    try:
+        pkg = importlib.import_module(module_name, pkg_name)
+        return pkg
+    except ImportError:
+        #warnings.warn(f'Package {module_name} was not found, therefore some related methods would be unavailable.')
+        return None