BUCToolkit 1.0b0__py3-none-any.whl

BUCToolkit/api/Losses.py

@@ -0,0 +1,170 @@
+"""
+Build-in Losses function
+"""
+#  Copyright (c) 2024-2025.7.4, BUCToolkit.
+#  Authors: Pu Pengxin, Song Xin
+#  Version: 0.9a
+#  File: Losses.py
+#  Environment: Python 3.12
+
+from typing import Literal, List, Dict, Sequence, Tuple
+import math
+import numpy as np
+import torch as th
+import torch.nn as nn
+import torch.nn.functional as F
+
+class Energy_Force_Loss(nn.Module):
+    """
+    A loss function that evaluates both predicted energy and forces.
+    Both energy & force losses would be averaged over each atom.
+
+    Parameters:
+        loss_E: the loss function of energy.
+        loss_F: the loss function of forces.
+        coeff_E: the coefficient of energy.
+        coeff_F: the coefficient of forces.
+
+    forward:
+        pred: Dict[Literal['energy'], th.Tensor], output of models;
+        label: Dict[Literal['energy'], th.Tensor], labels.
+
+    """
+    def __init__(self, 
+                 loss_E: Literal['MAE', 'MSE']|nn.Module='MAE', 
+                 loss_F: Literal['MAE', 'MSE']|nn.Module='MAE', 
+                 coeff_E: float=1., 
+                 coeff_F: float=1.) -> None:
+        super().__init__()
+        if loss_E == 'MAE':
+            self.loss_E = nn.SmoothL1Loss(reduction='sum')
+        elif loss_E == 'MSE':
+            self.loss_E = nn.MSELoss(reduction='sum')
+        elif loss_E == 'SmoothMAE':
+            self.loss_E = nn.SmoothL1Loss(reduction='sum')
+        else:
+            self.loss_E = loss_E
+
+        if loss_F == 'MAE':
+            self.loss_F = nn.SmoothL1Loss(reduction='sum')
+        elif loss_F == 'MSE':
+            self.loss_F = nn.MSELoss(reduction='sum')
+        else:
+            self.loss_F = loss_F
+        
+        self.coeff_E = coeff_E
+        self.coeff_F = coeff_F
+    
+    def forward(self, pred:Dict[Literal['energy', 'forces'], th.Tensor], label:Dict[Literal['energy', 'forces'], th.Tensor]):
+        n_atom = math.prod(pred['forces'].shape[:-1])
+        loss = (self.coeff_E * self.loss_E(pred['energy'], label['energy']) +
+                self.coeff_F * self.loss_F(pred['forces'], label['forces']))/n_atom
+        return loss
+
+class Energy_Loss(nn.Module):
+    """
+    A loss function that evaluates both predicted energy and forces.
+
+    Parameters:
+        loss_E: the loss function of energy.
+
+    forward:
+        pred: Dict[Literal['energy'], th.Tensor], output of models;
+        label: Dict[Literal['energy'], th.Tensor], labels.
+
+    """
+    def __init__(self, loss_E: Literal['MAE', 'MSE', 'SmoothMAE', 'Huber'] | nn.Module = 'MAE',) -> None:
+        super().__init__()
+        if loss_E == 'MAE':
+            self.loss_E = nn.L1Loss()
+        elif loss_E == 'SmoothMAE':
+            self.loss_E = nn.SmoothL1Loss()
+        elif loss_E == 'MSE':
+            self.loss_E = nn.MSELoss()
+        elif loss_E == 'Huber':
+            self.loss_E = nn.HuberLoss()
+        else:
+            self.loss_E = loss_E
+
+    def forward(self, pred: Dict[Literal['energy'], th.Tensor], label: Dict[Literal['energy'], th.Tensor]):
+        loss = self.loss_E(pred['energy'], label['energy'])
+        return loss
+
+
+class WrapperBoostLoss(nn.Module):
+    def __init__(self, loss_E: Literal['MAE', 'MSE', 'SmoothMAE', 'Huber'] | nn.Module = 'MAE', loss_F = None) -> None:
+        super().__init__()
+        if loss_E == 'MAE':
+            self.loss_E = nn.L1Loss()
+        elif loss_E == 'SmoothMAE':
+            self.loss_E = nn.SmoothL1Loss()
+        elif loss_E == 'MSE':
+            self.loss_E = nn.MSELoss()
+        elif loss_E == 'Huber':
+            self.loss_E = nn.HuberLoss()
+        else:
+            self.loss_E = loss_E
+
+        if loss_F is None:
+            self.has_loss_f = False
+        else:
+            if loss_F == 'MAE':
+                self.loss_F = nn.SmoothL1Loss(reduction='sum')
+            elif loss_F == 'MSE':
+                self.loss_F = nn.MSELoss(reduction='sum')
+            else:
+                self.loss_F = loss_F
+
+    def forward(self, pred: Dict[Literal['energy', 'forces'], List[th.Tensor]], label: Dict[Literal['energy', 'forces'], th.Tensor]):
+        loss = 0.
+        res = label['energy']
+        for pred_ in pred['energy']:
+            loss = loss + self.loss_E(res, pred_)
+            res = res - pred_
+        if self.has_loss_f:
+            res_f = label['forces']
+            for pred_ in pred['forces']:
+                loss = loss + self.loss_F(res_f, pred_)
+                res_f = res - pred_
+
+        return loss
+
+
+class WrapperMeanLoss(nn.Module):
+    def __init__(
+            self,
+            loss_E: Literal['MAE', 'MSE', 'SmoothMAE', 'Huber'] | nn.Module = 'MAE',
+            loss_F: Literal['MAE', 'MSE', 'SmoothMAE', 'Huber'] | nn.Module | None = None
+    ) -> None:
+        super().__init__()
+        if loss_E == 'MAE':
+            self.loss_E = nn.L1Loss()
+        elif loss_E == 'SmoothMAE':
+            self.loss_E = nn.SmoothL1Loss()
+        elif loss_E == 'MSE':
+            self.loss_E = nn.MSELoss()
+        elif loss_E == 'Huber':
+            self.loss_E = nn.HuberLoss()
+        else:
+            self.loss_E = loss_E
+
+        if loss_F is None:
+            self.has_loss_f = False
+        else:
+            self.has_loss_f = True
+            if loss_F == 'MAE':
+                self.loss_F = nn.SmoothL1Loss(reduction='sum')
+            elif loss_F == 'MSE':
+                self.loss_F = nn.MSELoss(reduction='sum')
+            else:
+                self.loss_F = loss_F
+
+    def forward(self, pred: Dict[Literal['energy', 'forces'], List[th.Tensor]], label: Dict[Literal['energy', 'forces'], th.Tensor]):
+        loss = 0.
+        for pred_ in pred['energy']:
+            loss = loss + self.loss_E(label['energy'], pred_)
+        if self.has_loss_f:
+            for pred_ in pred['forces']:
+                loss = loss + self.loss_F(label['forces'], pred_)
+
+        return loss

BUCToolkit/api/Metrics.py

@@ -0,0 +1,84 @@
+""" Metrics """
+#  Copyright (c) 2024-2025.7.4, BUCToolkit.
+#  Authors: Pu Pengxin, Song Xin
+#  Version: 0.9a
+#  File: Metrics.py
+#  Environment: Python 3.12
+
+from typing import Literal, Dict, List, Tuple
+import warnings
+import torch as th
+import torch.nn.functional as F
+
+
+def _r2_score(y_pred: th.Tensor, y_true: th.Tensor) -> th.Tensor:
+        """
+        Calculate R^2
+        """
+        # initialize
+        y1 = y_true.float()
+        y2 = y_pred.float()
+        if y1.dim() == 1:
+            y1 = y1.view(-1, 1)
+            y2 = y2.view(-1, 1)
+
+        # SS_res
+        ss_res = th.sum((y1 - y2) ** 2, dim=0)
+
+        # SS_tol
+        y1_mean = th.mean(y1, dim=0, keepdim=True)
+        ss_tot = th.sum((y1 - y1_mean) ** 2, dim=0)
+
+        # Main
+        r2 = th.where(
+            ss_tot != 0,
+            1 - (ss_res / ss_tot),
+            th.where(
+                ss_res != 0,
+                -th.inf,
+                th.ones_like(ss_res)
+            )
+        )
+
+        return r2.squeeze()
+
+def _rmse(x, y, reduction='mean') -> th.Tensor:
+    return th.sqrt(F.mse_loss(x, y, reduction=reduction))
+
+def E_MAE(pred: Dict[Literal['energy', 'forces'], th.Tensor],
+          label: Dict[Literal['energy', 'forces'], th.Tensor],
+          reduction: Literal['mean', 'sum', 'none'] = 'mean'):
+    with th.no_grad():
+        if isinstance(pred['energy'], th.Tensor):
+            mae = F.l1_loss(pred['energy'], label['energy'], reduction=reduction)
+        elif isinstance(pred['energy'], List|Tuple):
+            mae = sum(F.l1_loss(pred_, label['energy'], reduction=reduction) for pred_ in pred['energy'])
+
+    return mae
+
+
+def E_R2(pred: Dict[Literal['energy', 'forces'], th.Tensor],
+         label: Dict[Literal['energy', 'forces'], th.Tensor],):
+    if len(pred['energy']) <= 2:
+        r2 = th.tensor([1.], dtype=th.float32)
+        warnings.warn('Input samples less than 2, r2 was set to 1.', RuntimeWarning)
+    else:
+        with th.no_grad():
+            r2 = _r2_score(pred['energy'], label['energy'])
+
+    return r2
+
+
+def F_MAE(pred: Dict[Literal['energy', 'forces'], th.Tensor],
+          label: Dict[Literal['energy', 'forces'], th.Tensor],
+          reduction: Literal['mean', 'sum', 'none'] = 'mean'):
+    with th.no_grad():
+        mae = F.l1_loss(pred['forces'], label['forces'], reduction=reduction)
+    return mae
+
+
+def F_MaxE(pred: Dict[Literal['energy', 'forces'], th.Tensor],
+           label: Dict[Literal['energy', 'forces'], th.Tensor]):
+    with th.no_grad():
+        max_ = th.max(pred['forces'] - label['forces'])
+    return max_

BUCToolkit/api/ModelOptims.py

@@ -0,0 +1,296 @@
+#  Copyright (c) 2026.3.29, BUCToolkit.
+#  Authors: Pu Pengxin, Song Xin
+#  Version: 1.0b
+#  File: ModelOptims.py
+#  Environment: Python 3.12
+import math
+from collections import deque
+
+import torch as th
+from torch.optim import Optimizer
+
+class FIRELikeOptimizer(Optimizer):
+    """
+    FIRE-like optimizer based on damped Newton dynamics.
+    When power = -grad @ velocities > 0, which means loss descents, n_count += 1;
+    when power <= 0, apply an infinite pulse that sets velocity and N_min to zero, and lr to initial lr.
+    Once n_count >= N_min, lr *= fi (fi > 1.), i.e., increase learning rate.
+
+    Args:
+        params: model parameters.
+        lr: initial learning rate, i.e., time steplength.
+        mass: hyperparameter of mass.
+        N_min: steps that keep lr. If contiguous loss descent step number >= N_min, increase lr.
+        tolerance: steps that keep lr & veloc. If contiguous loss increment step number >= tolerance, reset lr to the default and veloc. to 0.
+        fi: learning rate increment factor
+        max_lr: maximum learning rate.
+    """
+    def __init__(self, params, lr=1e-2, mass=1.0, N_min=10, tolerance=5, fi=1.1, max_lr=None):
+        if lr <= 0:
+            raise ValueError(f"Invalid lr: {lr}")
+        else:
+            lr = float(lr)
+        if mass <= 0:
+            raise ValueError(f"mass must be positive: {mass}")
+        else:
+            mass = float(mass)
+        if N_min < 1:
+            raise ValueError(f"N_min must be >= 1: {N_min}")
+        else:
+            N_min = int(N_min)
+        if tolerance < 0:
+            raise ValueError(f"tolerance must be positive: {tolerance}")
+        else:
+            tolerance = int(tolerance)
+        if fi <= 1.0:
+            raise ValueError(f"fi must be > 1.0: {fi}")
+        else:
+            fi = float(fi)
+        if max_lr is None:
+            max_lr = lr * 10.
+        elif max_lr < lr:
+            raise ValueError("max_lr must be >= initial lr")
+        else:
+            max_lr = float(max_lr)
+
+        defaults = dict(lr=lr, mass=mass, N_min=N_min, tolerance=tolerance, fi=fi, max_lr=max_lr)
+        super().__init__(params, defaults)
+
+        # store 'lr' as the baseline, and 'lr_now' initialized to current really used lr.
+        for group in self.param_groups:
+            group['lr_now'] = group['lr']
+
+        # decrement counter
+        self.n_count_dec = 0
+        # increment counter
+        self.n_count_inc = 0
+
+    def step(self, closure=None):
+        loss = None
+        if closure is not None:
+            loss = closure()
+
+        # ---------- 1. Power P = -sum(grad · velocity) ----------
+        power = None
+        for group in self.param_groups:
+            for p in group['params']:
+                if p.grad is None:
+                    continue
+                grad = p.grad.data
+                state = self.state[p]
+                # lazy launch initial veloc.
+                if 'velocity' not in state:
+                    state['velocity'] = th.zeros_like(p.data)
+                if power is None: power = th.scalar_tensor(0.0, dtype=p.data.dtype, device=p.device)
+                vel = state['velocity']
+                #power -= th.dot(grad.view(-1), vel.view(-1)).item()
+                power: th.Tensor
+                power -= th.sum(vel * grad)  # DO NOT USE `.item()` which causes global synchronization
+
+        # ---------- 2. infinite damp pulse ----------
+        # scalar Tensor allows directly comparison
+        if power <= 0:  # loss increment
+            # sets velocity and N_min to zero, and lr to initial lr.
+            self.n_count_dec = 0
+            self.n_count_inc += 1
+            for group in self.param_groups:
+                if self.n_count_inc >= group['tolerance']:
+                    group['lr_now'] = group['lr']
+                    for p in group['params']:
+                        state = self.state.get(p)
+                        if (state is not None) and ('velocity' in state):
+                            state['velocity'].zero_()
+        else:
+            self.n_count_dec += 1
+            self.n_count_inc = 0
+            # check if contiguously decreased at least N_min steps
+            for i, group in enumerate(self.param_groups):
+                if self.n_count_dec >= self.param_groups[i]['N_min']:
+                    group['lr_now'] = min(group['lr_now'] * group['fi'], group['max_lr'])
+
+        # ---------- 3. semi-implicit Euler update ----------
+        for group in self.param_groups:
+            lr = group['lr_now']
+            mass = group['mass']
+            for p in group['params']:
+                if p.grad is None:
+                    continue
+                grad = p.grad.data
+                state = self.state[p]
+                vel = state['velocity']
+                # a = -grad / mass
+                neg_acc = grad / mass  # here avoid one O(n) tensor ops
+                # v: v = v + a * dt
+                vel.add_(neg_acc, alpha=-lr)
+                # p: p = p + v * dt
+                p.data.add_(vel, alpha=lr)
+
+        return loss
+
+    # ---------- store/resume states ----------
+    def state_dict(self):
+        state = super().state_dict()
+        state['n'] = self.n_count_dec
+        return state
+
+    def load_state_dict(self, state_dict):
+        self.n_count_dec = state_dict.pop('n')
+        super().load_state_dict(state_dict)
+
+
+class LangevinOptimizer(Optimizer):
+    r"""
+    Langevin optimizer based on Newton dynamics with Langevin thermostat.
+    Use Equipartition theorem or Virial theorem <x_i \partial H / \partial x_i> == k_B T to
+    judge whether equilibrium reached, wherein k_B = 1. If equilibrium reached, temperature decreases.
+
+    Args:
+        params: model parameters.
+        lr: initial learning rate, i.e., time steplength.
+        mass: hyperparameter of mass.
+        alpha: damping coefficient.
+        temperature: initial virtual temperature.
+        anneal_coeff: anneal coefficient to descent temperature.
+        balance_tol: relative tolerance of differences between k_B T and the Virial, which calculates as `(Virial - T)/T`
+        time_window_size: time window size to count the average <x_i \partial H / \partial x_i>.
+        heating_max_steps: the maximum heating steps. Once the number of calling `self.step()` over it, use normal SGD with momenta instead.
+        momenta_in_sgd: momentum decay coefficient after annealing.
+    """
+    def __init__(
+            self,
+            params,
+            lr=1e-3,
+            mass=1.0,
+            alpha=50.,
+            temperature=1.,
+            anneal_coeff=0.9,
+            balance_tol=0.1,
+            time_window_size=500,
+            heating_max_steps=100000,
+            momenta_in_sgd=0.6,
+    ):
+        if lr <= 0:
+            raise ValueError(f"Invalid lr: {lr}")
+        else:
+            lr = float(lr)
+        if mass <= 0:
+            raise ValueError(f"mass must be positive: {mass}")
+        else:
+            mass = th.scalar_tensor(mass, dtype=th.float32)
+        if anneal_coeff > 1. or anneal_coeff < 0.:
+            raise ValueError(f"anneal_coeff must be in [0, 1]: {anneal_coeff}")
+        time_window_size = int(time_window_size)
+        if time_window_size < 1:
+            raise ValueError(f"time_window_size must be positive: {time_window_size}")
+
+        # Virial queue
+        self.virials = deque([0.] * time_window_size, maxlen=time_window_size)
+        self.Virial_now = 0.
+        self.heating_max_steps = heating_max_steps
+        self.heating_step_now = 0
+
+        damp = math.exp(- alpha * lr)
+        defaults = dict(
+            lr=lr,
+            mass=mass,
+            damp=damp,
+            temperature=temperature,
+            anneal_coeff=anneal_coeff,
+            balance_tol=balance_tol,
+            momenta_in_sgd=momenta_in_sgd
+        )
+        super().__init__(params, defaults)
+
+    def step(self, closure=None):
+        loss = None
+        if closure is not None:
+            loss = closure()
+        # ---------- ABOBA update ----------
+        with th.no_grad():
+            for group in self.param_groups:
+                if self.heating_step_now >= self.heating_max_steps:  # enough annealing, use trivial SGD instead
+                    lr = group['lr']
+                    for p in group['params']:
+                        if p.grad is None:
+                            continue
+                        grad = p.grad.data
+                        state = self.state[p]
+                        # lazy launch initial veloc.
+                        if 'velocity' not in state:
+                            state['velocity'] = th.zeros_like(p.data)
+                        vel = state['velocity']
+                        vel.mul_(group['momenta_in_sgd']).add_(grad)
+                        p.add_(vel, alpha=-lr)
+
+                else:  # keep annealing
+                    lr = group['lr']
+                    mass = group['mass']
+                    damp = group['damp']
+                    temp = group['temperature']
+                    anneal_coeff = group['anneal_coeff']
+                    balance_tol = group['balance_tol']
+                    xpHpx = 0.
+                    sizes = 0
+                    _vir = 0.
+                    for p in group['params']:
+                        if p.grad is None:
+                            continue
+                        grad = p.grad
+                        state = self.state[p]
+                        # lazy launch initial veloc.
+                        if 'velocity' not in state:
+                            state['velocity'] = th.randn_like(p) * th.sqrt(temp / mass)
+                        vel = state['velocity']
+                        # Calc. Virial
+                        #xpHpx += th.sum(p * grad)  # x_i * \partial H / \partial x_i
+                        sizes += p.numel()
+                        _vir += (vel.pow(2).sum() * 0.5 * mass).item()
+                        mass = mass.to(p.dtype).to(p.device)
+                        # ABOBA
+                        vel.addcdiv_(grad, mass, value=- 0.5 * lr)
+                        p.add_(vel, alpha=0.5 * lr)
+                        # stochastic update velocity
+                        vel.mul_(damp)
+                        vel.add_(th.sqrt((temp * (1 - damp ** 2)) / mass) * th.randn_like(vel))
+                        # V = damp * V + th.sqrt((self.T_init * (1 - damp ** 2)) / masses) * th.randn_like(V)
+                        # the rest half-step
+                        p.add_(vel, alpha=0.5 * lr)
+                        vel.addcdiv_(grad, mass, value=-0.5 * lr)
+
+                    # Annealing, check the Equipartition theorem
+                    self.Virial_now -= self.virials[0]  # sub the head
+                    #_vir = xpHpx / sizes
+                    _vir /= sizes
+                    self.virials.append(_vir)
+                    self.Virial_now += _vir
+                    #print(f"Virial: {(self.Virial_now/len(self.virials)): <.4e}, T: {temp}")
+                    #avg_grad_norm = sum(p.grad.norm().item() for p in group['params'] if p.grad is not None) / len(group['params'])
+                    #avg_param_norm = sum(p.norm().item() for p in group['params']) / len(group['params'])
+                    #print(f"Step {self.heating_step_now}: Virial={self.Virial_now/(0.5 * len(self.virials)):.4e}, T={temp:.4e}")
+
+                    #if abs((self.Virial_now/len(self.virials) - temp)/temp) <= balance_tol:
+                    #    group['temperature'] *= anneal_coeff
+                    if abs((self.Virial_now/(0.5 * len(self.virials)) - temp) / temp) <= balance_tol:
+                        group['temperature'] *= anneal_coeff
+
+        self.heating_step_now += 1
+
+        return loss
+
+    # ---------- store/resume states ----------
+    def state_dict(self):
+        state = super().state_dict()
+        state.update({
+            'virials': self.virials, 'Virial_now': self.Virial_now,
+            'heating_step_now': self.heating_step_now,
+            'heating_max_steps': self.heating_max_steps,
+        })
+        return state
+
+    def load_state_dict(self, state_dict):
+        self.virials = state_dict.pop('virials')
+        self.virials = deque(self.virials, maxlen=len(self.virials))
+        self.Virial_now = state_dict.pop('Virial_now')
+        self.heating_step_now = state_dict.pop('heating_step_now')
+        self.heating_max_steps = state_dict.pop('heating_max_steps')
+        super().load_state_dict(state_dict)