utils

torch_admp.utils

Utility functions for torch-admp.

This module provides various utility functions used throughout the torch-admp package, including mathematical operations, unit conversions, and helper functions for optimization and calculations.

TorchConstants(units_dict: Optional[Dict] = None)

Bases: Module

Physical constants and unit conversions for torch-admp.

This class provides consistent physical constants and unit conversions compatible with ASE (Atomic Simulation Environment). It handles conversion between different unit systems for energy, length, and other physical quantities used in molecular simulations.

Notes

Electron volts (eV), Ångström (Ang), atomic mass unit and Kelvin are defined as 1.0.

Consistent with ASE: https://wiki.fysik.dtu.dk/ase/ase/units.html Example: units_dict = { "energy": "kJ/mol", "length": "nm", } Electron volts (eV), Ångström (Ang), the atomic mass unit and Kelvin are defined as 1.0.

Source code in torch_admp/utils.py

def __init__(self, units_dict: Optional[Dict] = None):
    """
    Consistent with ASE: https://wiki.fysik.dtu.dk/ase/ase/units.html
    Example:
    units_dict = {
        "energy": "kJ/mol",
        "length": "nm",
    }
    Electron volts (eV), Ångström (Ang), the atomic mass unit and Kelvin are defined as 1.0.
    """
    torch.nn.Module.__init__(self)
    self.pi = np.pi
    self.sqrt_pi = np.sqrt(np.pi)
    if units_dict is None:
        units_dict = {}
    user_energy = units_dict.get("energy", "eV")
    user_length = units_dict.get("length", "Ang")

    # from user-defined units to ASE units (eV, Ang)
    try:
        length_coeff = getattr(units, user_length)
    except AttributeError as exc:
        raise ValueError(f"Unknown length unit: {user_length}") from exc
    try:
        if user_energy == "kJ/mol":
            energy_coeff = units.kJ / units.mol
        else:
            energy_coeff = getattr(units, user_energy)
    except AttributeError as exc:
        raise ValueError(f"Unknown energy unit: {user_energy}") from exc

    self.length_coeff = length_coeff
    self.energy_coeff = energy_coeff
    # self.register_buffer(
    #     "j2ev",
    #     torch.tensor(
    #         constants.physical_constants["joule-electron volt relationship"][0],
    #         device=DEVICE,
    #     ),
    # )
    # # convert energy unit from kJ/mol to eV/particle (DMFF <-> DP)
    # self.register_buffer(
    #     "energy_coeff", self.j2ev * constants.kilo / constants.Avogadro
    # )

    # vacuum electric permittivity in eV^-1 * angstrom^-1
    self.epsilon = (
        constants.epsilon_0 / constants.elementary_charge * constants.angstrom
    )
    # qqrd2e = 1 / (4 * np.pi * EPSILON)
    # eV
    self.dielectric = 1.0 / (4.0 * np.pi * self.epsilon)

calc_grads(t_out: torch.Tensor, t_in: torch.Tensor)

Calculate gradients

PARAMETER	DESCRIPTION
t_out	Outputs of the differentiated function TYPE: Tensor
t_in	Inputs w.r.t. which the gradient will be returned TYPE: Tensor

RETURNS	DESCRIPTION
grad	Gradients TYPE: Tensor

Source code in torch_admp/utils.py

@torch.jit.script
def calc_grads(t_out: torch.Tensor, t_in: torch.Tensor):
    """
    Calculate gradients

    Parameters
    ----------
    t_out : torch.Tensor
        Outputs of the differentiated function
    t_in : torch.Tensor
        Inputs w.r.t. which the gradient will be returned

    Returns
    -------
    grad : torch.Tensor
        Gradients
    """
    assert t_in.requires_grad, "Input tensor requires grad"

    faked_grad = torch.ones_like(t_out)
    lst = torch.jit.annotate(List[Optional[torch.Tensor]], [faked_grad])
    grad = torch.autograd.grad(
        [t_out],
        [t_in],
        grad_outputs=lst,
        retain_graph=True,
    )[0]
    assert grad is not None
    return grad

calc_pgrads(t_out: torch.Tensor, t_in: torch.Tensor, constraint_matrix: torch.Tensor, coeff_matrix: torch.Tensor)

Calculate projected gradients for constrained optimization

PARAMETER	DESCRIPTION
t_out	Output tensor TYPE: Tensor
t_in	Input tensor TYPE: Tensor
constraint_matrix	n_const * natoms, constraint matrix TYPE: Tensor
coeff_matrix	natoms * n_const, Coefficient matrix for vector projection TYPE: Tensor

RETURNS	DESCRIPTION
Tensor	Projected gradients

Source code in torch_admp/utils.py

@torch.jit.script
def calc_pgrads(
    t_out: torch.Tensor,
    t_in: torch.Tensor,
    constraint_matrix: torch.Tensor,
    coeff_matrix: torch.Tensor,
):
    """
    Calculate projected gradients for constrained optimization

    Parameters
    ----------
    t_out : torch.Tensor
        Output tensor
    t_in : torch.Tensor
        Input tensor
    constraint_matrix : torch.Tensor
        n_const * natoms, constraint matrix
    coeff_matrix : torch.Tensor
        natoms * n_const, Coefficient matrix for vector projection

    Returns
    -------
    torch.Tensor
        Projected gradients
    """
    raw_grads = calc_grads(t_out, t_in)
    # n_atoms * 1
    raw_grads = raw_grads.reshape(-1, 1)
    # n_const * 1
    residual = -torch.matmul(constraint_matrix, raw_grads)
    # n_atoms * 1
    pgrads = raw_grads + torch.matmul(coeff_matrix, residual)
    return pgrads.reshape(-1)

pair_buffer_scales(pairs: torch.Tensor) -> torch.Tensor

Calculate buffer scales for atom pairs.

PARAMETER	DESCRIPTION
pairs	Tensor of atom pairs TYPE: Tensor

RETURNS	DESCRIPTION
Tensor	Buffer scales for each pair (1 if i < j, else 0)

Source code in torch_admp/utils.py

def pair_buffer_scales(pairs: torch.Tensor) -> torch.Tensor:
    """
    Calculate buffer scales for atom pairs.

    Parameters
    ----------
    pairs : torch.Tensor
        Tensor of atom pairs

    Returns
    -------
    torch.Tensor
        Buffer scales for each pair (1 if i < j, else 0)
    """
    dp = pairs[:, 0] - pairs[:, 1]
    return torch.where(
        dp < 0,
        torch.tensor(1, dtype=torch.long, device=pairs.device),
        torch.tensor(0, dtype=torch.long, device=pairs.device),
    )

regularize_pairs(pairs: torch.Tensor, buffer_scales: torch.Tensor) -> torch.Tensor

Regularize atom pairs based on buffer scales.

PARAMETER	DESCRIPTION
pairs	Tensor of atom pairs TYPE: Tensor
buffer_scales	Buffer scales for each pair TYPE: Tensor

RETURNS	DESCRIPTION
Tensor	Regularized atom pairs

Source code in torch_admp/utils.py

def regularize_pairs(pairs: torch.Tensor, buffer_scales: torch.Tensor) -> torch.Tensor:
    """
    Regularize atom pairs based on buffer scales.

    Parameters
    ----------
    pairs : torch.Tensor
        Tensor of atom pairs
    buffer_scales : torch.Tensor
        Buffer scales for each pair

    Returns
    -------
    torch.Tensor
        Regularized atom pairs
    """
    a = pairs[:, 0] - buffer_scales
    b = pairs[:, 1] - buffer_scales * 2
    return torch.stack((a, b), dim=1)

safe_inverse(x: torch.Tensor, threshold: float = 1e-08) -> torch.Tensor

Safe inverse for numerical stability

PARAMETER	DESCRIPTION
x	Input tensor TYPE: Tensor
threshold	Threshold for numerical stability TYPE: float (default: 1e-8) DEFAULT: 1e-08

RETURNS	DESCRIPTION
inv_x	Inverse of x if x.abs() > threshold, otherwise 0 TYPE: Tensor

Source code in torch_admp/utils.py

def safe_inverse(x: torch.Tensor, threshold: float = 1e-8) -> torch.Tensor:
    """Safe inverse for numerical stability

    Parameters
    ----------
    x : torch.Tensor
        Input tensor
    threshold : float (default: 1e-8)
        Threshold for numerical stability

    Returns
    -------
    inv_x : torch.Tensor
        Inverse of x if x.abs() > threshold, otherwise 0
    """
    return torch.where(x.abs() > threshold, 1 / x, torch.zeros_like(x))

vector_projection(vector_in: torch.Tensor, constraint_matrix: torch.Tensor, constraint_vector: Optional[torch.Tensor] = None) -> torch.Tensor

Vector projection subject to linear constraints P(x) = x + A^T(A A^T)^{-1}(b - Ax)

PARAMETER	DESCRIPTION
vector_in	Input vector (n_atoms * 1). TYPE: Tensor
constraint_matrix	Constraint matrix (n_const * natoms). TYPE: Tensor
constraint_vector	Constraint vector (n_const * 1). All zeros when set as None (default). TYPE: Tensor DEFAULT: None

RETURNS	DESCRIPTION
vector_out	n_atoms, projected vector TYPE: Tensor

Source code in torch_admp/utils.py

@torch.jit.script
def vector_projection(
    vector_in: torch.Tensor,
    constraint_matrix: torch.Tensor,
    constraint_vector: Optional[torch.Tensor] = None,
) -> torch.Tensor:
    """Vector projection subject to linear constraints
        P(x) = x + A^T(A A^T)^{-1}(b - Ax)

    Parameters
    ----------
    vector_in : torch.Tensor
        Input vector (n_atoms * 1).
    constraint_matrix : torch.Tensor
        Constraint matrix (n_const * natoms).
    constraint_vector : torch.Tensor, optional
        Constraint vector (n_const * 1).
        All zeros when set as None (default).

    Returns
    -------
    vector_out : torch.Tensor
        n_atoms, projected vector
    """
    if constraint_vector is None:
        constraint_vector = torch.zeros(
            [constraint_matrix.shape[0]], device=vector_in.device
        )

    # n_atoms * n_atoms
    coeff_matrix = vector_projection_coeff_matrix(constraint_matrix)
    # n_atoms * 1
    residual = constraint_vector.reshape(-1, 1) - torch.matmul(
        constraint_matrix, vector_in
    )
    vector_out = vector_in.reshape(-1, 1) + torch.matmul(coeff_matrix, residual)
    return vector_out.reshape(-1)

vector_projection_coeff_matrix(constraint_matrix: torch.Tensor) -> torch.Tensor

Calculate coefficient matrix for vector projection based on constraint matrix

PARAMETER	DESCRIPTION
constraint_matrix	Constraint matrix (n_const * natoms). TYPE: Tensor

RETURNS	DESCRIPTION
coeff_mat	Coefficient matrix (n_atoms * n_const). TYPE: Tensor

Source code in torch_admp/utils.py

@torch.jit.script
def vector_projection_coeff_matrix(constraint_matrix: torch.Tensor) -> torch.Tensor:
    """Calculate coefficient matrix for vector projection based on constraint matrix

    Parameters
    ----------
    constraint_matrix : torch.Tensor
        Constraint matrix (n_const * natoms).

    Returns
    -------
    coeff_mat: torch.Tensor
        Coefficient matrix (n_atoms * n_const).
    """
    constraint_matrix_t = torch.transpose(constraint_matrix, 0, 1)
    # n_atoms * n_const
    coeff_mat = torch.matmul(
        constraint_matrix_t,
        torch.inverse(torch.matmul(constraint_matrix, constraint_matrix_t)),
    )
    return coeff_mat