qeq

torch_admp.qeq

Charge equilibration (QEq) implementation for torch-admp.

This module implements charge equilibration methods for determining atomic charges in molecular systems. It includes various optimization approaches including matrix inversion and projected gradient methods, with support for different constraints and damping functions.

GaussianDampingForceModule(units_dict: Optional[Dict] = None)

Bases: BaseForceModule

Gaussian short-range damping force module.

This module implements the Gaussian damping function used in charge equilibration to account for short-range electrostatic interactions.

PARAMETER	DESCRIPTION
units_dict	Dictionary containing unit conversion factors, by default None TYPE: Optional[Dict] DEFAULT: None

Initialize the GaussianDampingForceModule.

PARAMETER	DESCRIPTION
units_dict	Dictionary containing unit conversion factors, by default None TYPE: Optional[Dict] DEFAULT: None

Source code in torch_admp/qeq.py

def __init__(
    self,
    units_dict: Optional[Dict] = None,
) -> None:
    """
    Initialize the GaussianDampingForceModule.

    Parameters
    ----------
    units_dict : Optional[Dict], optional
        Dictionary containing unit conversion factors, by default None
    """
    BaseForceModule.__init__(self, units_dict)

QEqForceModule(rcut: float, ethresh: float = 1e-05, kspace: bool = True, rspace: bool = True, slab_corr: bool = False, slab_axis: int = 2, max_iter: int = 20, ls_eps: float = 0.0001, eps: float = 0.0001, units_dict: Optional[Dict] = None, damping: bool = True, sel: Optional[list[int]] = None, kappa: Optional[float] = None, spacing: Union[List[float], float, None] = None, kmesh: Union[List[int], int, None] = None)

Bases: BaseForceModule

Charge equilibrium (QEq) model

PARAMETER	DESCRIPTION
rcut	cutoff radius for short-range interactions TYPE: float
ethresh	energy threshold for electrostatic interaction, by default 1e-5 TYPE: float DEFAULT: 1e-05
kspace	whether the reciprocal part is included TYPE: bool DEFAULT: True
rspace	whether the real space part is included TYPE: bool DEFAULT: True
slab_corr	whether the slab correction is applied TYPE: bool DEFAULT: False
slab_axis	axis at which the slab correction is applied TYPE: int DEFAULT: 2
max_iter	maximum number of iterations for optimization, by default 20 only used for projected gradient method TYPE: int DEFAULT: 20
ls_eps	threshold for line search, by default 1e-4 only used for projected gradient method TYPE: float DEFAULT: 0.0001
eps	threshold for convergence, by default 1e-4 only used for projected gradient method TYPE: float DEFAULT: 0.0001
units_dict	dictionary of units, by default None TYPE: Optional[Dict] DEFAULT: None

Initialize the QEqForceModule.

PARAMETER	DESCRIPTION
rcut	cutoff radius for short-range interactions TYPE: float
ethresh	energy threshold for electrostatic interaction, by default 1e-5 TYPE: float DEFAULT: 1e-05
kspace	whether the reciprocal part is included TYPE: bool DEFAULT: True
rspace	whether the real space part is included TYPE: bool DEFAULT: True
slab_corr	whether the slab correction is applied TYPE: bool DEFAULT: False
slab_axis	axis at which the slab correction is applied TYPE: int DEFAULT: 2
max_iter	maximum number of iterations for optimization, by default 20 only used for projected gradient method TYPE: int DEFAULT: 20
ls_eps	threshold for line search, by default 1e-4 only used for projected gradient method TYPE: float DEFAULT: 0.0001
eps	threshold for convergence, by default 1e-4 only used for projected gradient method TYPE: float DEFAULT: 0.0001
units_dict	dictionary of units, by default None TYPE: Dict DEFAULT: None
damping	Whether to include Gaussian damping, by default True TYPE: bool DEFAULT: True
sel	Selection list for neighbor list, by default None TYPE: Optional[list[int]] DEFAULT: None
kappa	Inverse screening length [Å^-1], by default None TYPE: Optional[float] DEFAULT: None
spacing	Grid spacing for reciprocal space, by default None TYPE: Optional[List[float]] DEFAULT: None

Source code in torch_admp/qeq.py

def __init__(
    self,
    rcut: float,
    ethresh: float = 1e-5,
    kspace: bool = True,
    rspace: bool = True,
    slab_corr: bool = False,
    slab_axis: int = 2,
    max_iter: int = 20,
    ls_eps: float = 1e-4,
    eps: float = 1e-4,
    units_dict: Optional[Dict] = None,
    damping: bool = True,
    sel: Optional[list[int]] = None,
    kappa: Optional[float] = None,
    spacing: Union[List[float], float, None] = None,
    kmesh: Union[List[int], int, None] = None,
) -> None:
    """
    Initialize the QEqForceModule.

    Parameters
    ----------
    rcut : float
        cutoff radius for short-range interactions
    ethresh : float, optional
        energy threshold for electrostatic interaction, by default 1e-5
    kspace : bool
        whether the reciprocal part is included
    rspace : bool
        whether the real space part is included
    slab_corr : bool
        whether the slab correction is applied
    slab_axis : int
        axis at which the slab correction is applied
    max_iter : int, optional
        maximum number of iterations for optimization, by default 20
        only used for projected gradient method
    ls_eps : float, optional
        threshold for line search, by default 1e-4
        only used for projected gradient method
    eps : float, optional
        threshold for convergence, by default 1e-4
        only used for projected gradient method
    units_dict : Dict, optional
        dictionary of units, by default None
    damping : bool, optional
        Whether to include Gaussian damping, by default True
    sel : Optional[list[int]], optional
        Selection list for neighbor list, by default None
    kappa : Optional[float], optional
        Inverse screening length [Å^-1], by default None
    spacing : Optional[List[float]], optional
        Grid spacing for reciprocal space, by default None
    """
    BaseForceModule.__init__(self, units_dict)

    models: Dict[str, BaseForceModule] = {
        "site": SiteForceModule(units_dict=units_dict),
        "coulomb": CoulombForceModule(
            rcut=rcut,
            ethresh=ethresh,
            kspace=kspace,
            rspace=rspace,
            slab_corr=slab_corr,
            slab_axis=slab_axis,
            units_dict=units_dict,
            kappa=kappa,
            spacing=spacing,
            kmesh=kmesh,
        ),
    }
    if damping:
        models["damping"] = GaussianDampingForceModule(units_dict=units_dict)

    self.submodels = torch.nn.ModuleDict(models)

    self.rcut = rcut
    self.max_iter = max_iter
    self.ls_eps = ls_eps
    self.eps = eps
    self.converge_iter: int = -1
    self.sel = sel

    self.slab_axis = slab_axis
    self.slab_corr = slab_corr

calc_hessian(positions: torch.Tensor, box: Optional[torch.Tensor], chi: torch.Tensor, hardness: torch.Tensor, eta: torch.Tensor, pairs: torch.Tensor, ds: torch.Tensor, buffer_scales: torch.Tensor)

Calculate the Hessian matrix of the energy with respect to charges.

PARAMETER	DESCRIPTION
positions	Atomic positions TYPE: Tensor
box	Simulation box vectors TYPE: Optional[Tensor]
chi	Electronegativity in energy/charge unit TYPE: Tensor
hardness	Atomic hardness in energy/charge^2 unit TYPE: Tensor
eta	Gaussian width in length unit TYPE: Tensor
pairs	Tensor of atom pairs TYPE: Tensor
ds	Distance tensor TYPE: Tensor
buffer_scales	Buffer scales for each pair TYPE: Tensor

RETURNS	DESCRIPTION
Tensor	Hessian matrix with shape (n_atoms, n_atoms)

Source code in torch_admp/qeq.py

@torch.jit.ignore
def calc_hessian(
    self,
    positions: torch.Tensor,
    box: Optional[torch.Tensor],
    chi: torch.Tensor,
    hardness: torch.Tensor,
    eta: torch.Tensor,
    pairs: torch.Tensor,
    ds: torch.Tensor,
    buffer_scales: torch.Tensor,
):
    """
    Calculate the Hessian matrix of the energy with respect to charges.

    Parameters
    ----------
    positions : torch.Tensor
        Atomic positions
    box : Optional[torch.Tensor]
        Simulation box vectors
    chi : torch.Tensor
        Electronegativity in energy/charge unit
    hardness : torch.Tensor
        Atomic hardness in energy/charge^2 unit
    eta : torch.Tensor
        Gaussian width in length unit
    pairs : torch.Tensor
        Tensor of atom pairs
    ds : torch.Tensor
        Distance tensor
    buffer_scales : torch.Tensor
        Buffer scales for each pair

    Returns
    -------
    torch.Tensor
        Hessian matrix with shape (n_atoms, n_atoms)
    """
    n_atoms = positions.shape[0]
    q_tmp = torch.zeros(
        n_atoms, device=positions.device, dtype=positions.dtype, requires_grad=True
    )
    # calculate hessian
    # hessian = torch.func.hessian(self.func_energy)(
    #     q_tmp, positions, box, chi, hardness, eta, pairs, ds, buffer_scales
    # )
    y = self.func_energy(
        q_tmp, positions, box, chi, hardness, eta, pairs, ds, buffer_scales
    )
    grad = torch.autograd.grad(y, q_tmp, retain_graph=True, create_graph=True)
    hessian = []
    for anygrad in grad[0]:
        hessian.append(torch.autograd.grad(anygrad, q_tmp, retain_graph=True)[0])
    hessian = torch.stack(hessian)
    return hessian.reshape([n_atoms, n_atoms])

func_energy(charges: torch.Tensor, positions: torch.Tensor, box: Optional[torch.Tensor], chi: torch.Tensor, hardness: torch.Tensor, eta: torch.Tensor, pairs: torch.Tensor, ds: torch.Tensor, buffer_scales: torch.Tensor) -> torch.Tensor

Energy method for QEq model

PARAMETER	DESCRIPTION
charges	atomic charges TYPE: Tensor
positions	atomic positions TYPE: Tensor
box	simulation box TYPE: Tensor
chi	eletronegativity in energy / charge unit TYPE: Tensor
hardness	atomic hardness in energy / charge^2 unit TYPE: Tensor
eta	Gaussian width in length unit TYPE: Tensor
pairs	n_pairs * 2 tensor of pairs TYPE: Tensor
ds	i-j distance tensor TYPE: Tensor
buffer_scales	buffer scales for each pair, 1 if i < j else 0 TYPE: Tensor

RETURNS	DESCRIPTION
energy	energy tensor TYPE: Tensor

Source code in torch_admp/qeq.py

@torch.jit.export
def func_energy(
    self,
    charges: torch.Tensor,
    positions: torch.Tensor,
    box: Optional[torch.Tensor],
    chi: torch.Tensor,
    hardness: torch.Tensor,
    eta: torch.Tensor,
    pairs: torch.Tensor,
    ds: torch.Tensor,
    buffer_scales: torch.Tensor,
) -> torch.Tensor:
    """Energy method for QEq model

    Parameters
    ----------
    charges : torch.Tensor
        atomic charges
    positions : torch.Tensor
        atomic positions
    box : torch.Tensor
        simulation box
    chi : torch.Tensor
        eletronegativity in energy / charge unit
    hardness : torch.Tensor
        atomic hardness in energy / charge^2 unit
    eta : torch.Tensor
        Gaussian width in length unit
    pairs : torch.Tensor
        n_pairs * 2 tensor of pairs
    ds : torch.Tensor
        i-j distance tensor
    buffer_scales : torch.Tensor
        buffer scales for each pair, 1 if i < j else 0

    Returns
    -------
    energy: torch.Tensor
        energy tensor
    """
    params = {
        "charge": charges,  # (optional) initial guess for atomic charges,
        "chi": chi,  # eletronegativity in energy / charge unit
        "hardness": hardness,  # atomic hardness in energy / charge^2 unit
        "eta": eta,  # Gaussian width in length unit
    }
    energy = torch.zeros(1, device=positions.device)
    for model in self.submodels.values():
        energy = energy + model(positions, box, pairs, ds, buffer_scales, params)
    return energy

get_rcut() -> float

Get the cutoff radius.

RETURNS	DESCRIPTION
float	Cutoff radius

Source code in torch_admp/qeq.py

def get_rcut(self) -> float:
    """
    Get the cutoff radius.

    Returns
    -------
    float
        Cutoff radius
    """
    return self.rcut

get_sel() -> Optional[list[int]]

Get sel list of DP model.

RETURNS	DESCRIPTION
Optional[list[int]]	The number of selected neighbors for each type of atom.

Source code in torch_admp/qeq.py

def get_sel(self) -> Optional[list[int]]:
    """
    Get `sel` list of DP model.

    Returns
    -------
    Optional[list[int]]
        The number of selected neighbors for each type of atom.
    """
    return self.sel

optimality(charges: torch.Tensor, positions: torch.Tensor, box: Optional[torch.Tensor], chi: torch.Tensor, hardness: torch.Tensor, eta: torch.Tensor, pairs: torch.Tensor, ds: torch.Tensor, buffer_scales: torch.Tensor, constraint_matrix: torch.Tensor, coeff_matrix: torch.Tensor) -> torch.Tensor

Optimality function (normalized projected gradient)

PARAMETER	DESCRIPTION
charges	atomic charges TYPE: Tensor
positions	atomic positions TYPE: Tensor
box	simulation box TYPE: Tensor
chi	eletronegativity in energy / charge unit TYPE: Tensor
hardness	atomic hardness in energy / charge^2 unit TYPE: Tensor
eta	Gaussian width in length unit TYPE: Tensor
pairs	n_pairs * 2 tensor of pairs TYPE: Tensor
ds	i-j distance tensor TYPE: Tensor
buffer_scales	buffer scales for each pair, 1 if i < j else 0 TYPE: Tensor
constraint_matrix	n_const * natoms, constraint matrix TYPE: Tensor
coeff_matrix	n_atoms * n_const, coefficient matrix TYPE: Tensor

RETURNS	DESCRIPTION
pgrad_norm	normalized projected gradient, ~zero at for optimal charges TYPE: Tensor

Source code in torch_admp/qeq.py

@torch.jit.export
def optimality(
    self,
    charges: torch.Tensor,
    positions: torch.Tensor,
    box: Optional[torch.Tensor],
    chi: torch.Tensor,
    hardness: torch.Tensor,
    eta: torch.Tensor,
    pairs: torch.Tensor,
    ds: torch.Tensor,
    buffer_scales: torch.Tensor,
    constraint_matrix: torch.Tensor,
    coeff_matrix: torch.Tensor,
) -> torch.Tensor:
    """Optimality function (normalized projected gradient)

    Parameters
    ----------
    charges : torch.Tensor
        atomic charges
    positions : torch.Tensor
        atomic positions
    box : torch.Tensor
        simulation box
    chi : torch.Tensor
        eletronegativity in energy / charge unit
    hardness : torch.Tensor
        atomic hardness in energy / charge^2 unit
    eta : torch.Tensor
        Gaussian width in length unit
    pairs : torch.Tensor
        n_pairs * 2 tensor of pairs
    ds : torch.Tensor
        i-j distance tensor
    buffer_scales : torch.Tensor
        buffer scales for each pair, 1 if i < j else 0
    constraint_matrix : torch.Tensor
        n_const * natoms, constraint matrix
    coeff_matrix : torch.Tensor
        n_atoms * n_const, coefficient matrix

    Returns
    -------
    pgrad_norm: torch.Tensor
        normalized projected gradient, ~zero at for optimal charges
    """
    energy = self.func_energy(
        charges, positions, box, chi, hardness, eta, pairs, ds, buffer_scales
    )
    pgrads = calc_pgrads(energy, charges, constraint_matrix, coeff_matrix)
    return torch.norm(pgrads) / charges.shape[0]

solve_matrix_inversion(positions: torch.Tensor, box: Optional[torch.Tensor], chi: torch.Tensor, hardness: torch.Tensor, eta: torch.Tensor, pairs: torch.Tensor, ds: torch.Tensor, buffer_scales: torch.Tensor, constraint_matrix: Optional[torch.Tensor], constraint_vals: Optional[torch.Tensor], check_hessian: bool = False)

Solve QEq with matrix inversion method

PARAMETER	DESCRIPTION
positions	atomic positions TYPE: Tensor
box	simulation box TYPE: Tensor
chi	eletronegativity in energy / charge unit TYPE: Tensor
hardness	atomic hardness in energy / charge^2 unit TYPE: Tensor
eta	Gaussian width in length unit TYPE: Tensor
pairs	n_pairs * 2 tensor of pairs TYPE: Tensor
ds	i-j distance tensor TYPE: Tensor
buffer_scales	buffer scales for each pair, 1 if i < j else 0 TYPE: Tensor
constraint_matrix	n_const * natoms, constraint matrix TYPE: Tensor
constraint_vals	n_const, constraint values TYPE: Tensor
check_hessian	(debugger) check whether hessian matrix is positive definite, by default False TYPE: bool DEFAULT: False

RETURNS	DESCRIPTION
energy	energy tensor TYPE: Tensor
q_opt	optimized atomic charges TYPE: Tensor

Source code in torch_admp/qeq.py

@torch.jit.ignore
def solve_matrix_inversion(
    self,
    positions: torch.Tensor,
    box: Optional[torch.Tensor],
    chi: torch.Tensor,
    hardness: torch.Tensor,
    eta: torch.Tensor,
    pairs: torch.Tensor,
    ds: torch.Tensor,
    buffer_scales: torch.Tensor,
    constraint_matrix: Optional[torch.Tensor],
    constraint_vals: Optional[torch.Tensor],
    check_hessian: bool = False,
):
    """Solve QEq with matrix inversion method

    Parameters
    ----------
    positions : torch.Tensor
        atomic positions
    box : torch.Tensor
        simulation box
    chi : torch.Tensor
        eletronegativity in energy / charge unit
    hardness : torch.Tensor
        atomic hardness in energy / charge^2 unit
    eta : torch.Tensor
        Gaussian width in length unit
    pairs : torch.Tensor
        n_pairs * 2 tensor of pairs
    ds : torch.Tensor
        i-j distance tensor
    buffer_scales : torch.Tensor
        buffer scales for each pair, 1 if i < j else 0
    constraint_matrix : torch.Tensor
        n_const * natoms, constraint matrix
    constraint_vals : torch.Tensor
        n_const, constraint values
    check_hessian : bool, optional
        (debugger) check whether hessian matrix is positive definite, by default False

    Returns
    -------
    energy: torch.Tensor
        energy tensor
    q_opt: torch.Tensor
        optimized atomic charges
    """
    # calculate hessian
    # n_atoms * n_atoms
    hessian = self.calc_hessian(
        positions, box, chi, hardness, eta, pairs, ds, buffer_scales
    )

    # coeff matrix as [[hessian, constraint_matrix.T], [constraint_matrix, 0]]
    # (n_atoms + n_const) * (n_atoms + n_const)
    n_atoms = positions.shape[0]
    if constraint_matrix is None:
        coeff_matrix = hessian
        vector = -chi
    else:
        n_const = constraint_matrix.shape[0]
        coeff_matrix = torch.cat(
            [
                torch.cat([hessian, constraint_matrix.T], dim=1),
                torch.cat(
                    [
                        constraint_matrix,
                        torch.zeros(n_const, n_const, device=positions.device),
                    ],
                    dim=1,
                ),
            ],
            dim=0,
        )
        vector = torch.concat([-chi, constraint_vals])

    if check_hessian:
        print(torch.all(torch.diag(hessian) > 0.0))

    _q_opt = torch.linalg.solve(coeff_matrix, vector.reshape(-1, 1)).reshape(-1)

    q_opt = _q_opt[:n_atoms].detach()
    q_opt.requires_grad = True
    energy = self.func_energy(
        q_opt, positions, box, chi, hardness, eta, pairs, ds, buffer_scales
    )
    # forces = -calc_grads(energy, positions)
    fermi = torch.mean(chi + torch.matmul(hessian, _q_opt[:n_atoms]))
    return energy, _q_opt[:n_atoms], torch.diag(hessian), fermi

solve_pgrad(q0: Optional[torch.Tensor], positions: torch.Tensor, box: Optional[torch.Tensor], chi: torch.Tensor, hardness: torch.Tensor, eta: torch.Tensor, pairs: torch.Tensor, ds: torch.Tensor, buffer_scales: torch.Tensor, constraint_matrix: torch.Tensor, constraint_vals: torch.Tensor, coeff_matrix: Optional[torch.Tensor] = None, reinit_q: bool = False, method: str = 'lbfgs')

Solve QEq with projected gradient method

PARAMETER	DESCRIPTION
q0	initial guess for atomic charges, all zeros for None TYPE: Tensor
positions	atomic positions TYPE: Tensor
box	simulation box TYPE: Tensor
chi	eletronegativity in energy / charge unit TYPE: Tensor
hardness	atomic hardness in energy / charge^2 unit TYPE: Tensor
eta	Gaussian width in length unit TYPE: Tensor
pairs	n_pairs * 2 tensor of pairs TYPE: Tensor
ds	i-j distance tensor TYPE: Tensor
buffer_scales	buffer scales for each pair, 1 if i < j else 0 TYPE: Tensor
constraint_matrix	n_const * natoms, constraint matrix TYPE: Tensor
constraint_vals	n_const, constraint values TYPE: Tensor
coeff_matrix	n_atoms * n_const, coefficient matrix TYPE: Tensor DEFAULT: None
reinit_q	if reinitialize the atomic charges based on constraints, by default False TYPE: bool DEFAULT: False
method	optimization method, by default "quadratic" TYPE: str DEFAULT: 'lbfgs'

RETURNS	DESCRIPTION
energy	energy tensor TYPE: Tensor
q_opt	optimized atomic charges TYPE: Tensor

Source code in torch_admp/qeq.py

@torch.jit.ignore
def solve_pgrad(
    self,
    q0: Optional[torch.Tensor],
    positions: torch.Tensor,
    box: Optional[torch.Tensor],
    chi: torch.Tensor,
    hardness: torch.Tensor,
    eta: torch.Tensor,
    pairs: torch.Tensor,
    ds: torch.Tensor,
    buffer_scales: torch.Tensor,
    constraint_matrix: torch.Tensor,
    constraint_vals: torch.Tensor,
    coeff_matrix: Optional[torch.Tensor] = None,
    reinit_q: bool = False,
    method: str = "lbfgs",
):
    """Solve QEq with projected gradient method

    Parameters
    ----------
    q0 : torch.Tensor
        initial guess for atomic charges, all zeros for None
    positions : torch.Tensor
        atomic positions
    box : torch.Tensor
        simulation box
    chi : torch.Tensor
        eletronegativity in energy / charge unit
    hardness : torch.Tensor
        atomic hardness in energy / charge^2 unit
    eta : torch.Tensor
        Gaussian width in length unit
    pairs : torch.Tensor
        n_pairs * 2 tensor of pairs
    ds : torch.Tensor
        i-j distance tensor
    buffer_scales : torch.Tensor
        buffer scales for each pair, 1 if i < j else 0
    constraint_matrix : torch.Tensor
        n_const * natoms, constraint matrix
    constraint_vals : torch.Tensor
        n_const, constraint values
    coeff_matrix : torch.Tensor
        n_atoms * n_const, coefficient matrix
    reinit_q : bool, optional
        if reinitialize the atomic charges based on constraints, by default False
    method : str, optional
        optimization method, by default "quadratic"

    Returns
    -------
    energy: torch.Tensor
        energy tensor
    q_opt: torch.Tensor
        optimized atomic charges
    """
    n_atoms = positions.shape[0]
    # n_const = constraint_matrix.shape[0]

    if q0 is None:
        q0 = torch.rand(n_atoms, device=positions.device, dtype=positions.dtype)
        reinit_q = True
    assert q0.shape[0] == n_atoms
    # make sure the initial guess satisfy constraints
    if reinit_q:
        q0 = vector_projection(q0, constraint_matrix, constraint_vals)

    if coeff_matrix is None:
        coeff_matrix = vector_projection_coeff_matrix(constraint_matrix)

    # choose iterative algorithm
    try:
        solver_fn = getattr(self, f"_optimize_{method}")()
    except KeyError as exc:
        raise ValueError(f"Method {method} is not supported.") from exc

    with torch.device(positions.device):
        _q_opt = custom_root(
            self.optimality,
            argnums=1,
            solve=torchopt.linear_solve.solve_normal_cg(maxiter=5, atol=0),
        )(solver_fn)(
            q0,
            positions,
            box,
            chi,
            hardness,
            eta,
            pairs,
            ds,
            buffer_scales,
            constraint_matrix,
            coeff_matrix,
        )

    q_opt = _q_opt.detach()
    q_opt.requires_grad = True
    energy = self.func_energy(
        q_opt, positions, box, chi, hardness, eta, pairs, ds, buffer_scales
    )
    # forces = -calc_grads(energy, positions)
    return energy, q_opt

SiteForceModule(units_dict: Optional[Dict] = None)

Bases: BaseForceModule

Chemical site energy force module.

This module implements the chemical site energy term in charge equilibration, accounting for electronegativity and hardness of atomic sites.

PARAMETER	DESCRIPTION
units_dict	Dictionary containing unit conversion factors, by default None TYPE: Optional[Dict] DEFAULT: None

Initialize the SiteForceModule.

PARAMETER	DESCRIPTION
units_dict	Dictionary containing unit conversion factors, by default None TYPE: Optional[Dict] DEFAULT: None

Source code in torch_admp/qeq.py

def __init__(
    self,
    units_dict: Optional[Dict] = None,
) -> None:
    """
    Initialize the SiteForceModule.

    Parameters
    ----------
    units_dict : Optional[Dict], optional
        Dictionary containing unit conversion factors, by default None
    """
    BaseForceModule.__init__(self, units_dict)

calc_hessian(func_energy: Callable, positions: torch.Tensor, box: Optional[torch.Tensor], chi: torch.Tensor, hardness: torch.Tensor, eta: torch.Tensor, pairs: torch.Tensor, ds: torch.Tensor, buffer_scales: torch.Tensor)

Calculate Hessian matrix of the energy with respect to charges.

PARAMETER	DESCRIPTION
func_energy	Energy function TYPE: Callable
positions	Atomic positions TYPE: Tensor
box	Simulation box vectors TYPE: Optional[Tensor]
chi	Electronegativity in energy/charge unit TYPE: Tensor
hardness	Atomic hardness in energy/charge^2 unit TYPE: Tensor
eta	Gaussian width in length unit TYPE: Tensor
pairs	Tensor of atom pairs TYPE: Tensor
ds	Distance tensor TYPE: Tensor
buffer_scales	Buffer scales for each pair TYPE: Tensor

RETURNS	DESCRIPTION
Tensor	Hessian matrix with shape (n_atoms, n_atoms)

Source code in torch_admp/qeq.py

def calc_hessian(
    func_energy: Callable,
    positions: torch.Tensor,
    box: Optional[torch.Tensor],
    chi: torch.Tensor,
    hardness: torch.Tensor,
    eta: torch.Tensor,
    pairs: torch.Tensor,
    ds: torch.Tensor,
    buffer_scales: torch.Tensor,
):
    """
    Calculate Hessian matrix of the energy with respect to charges.

    Parameters
    ----------
    func_energy : Callable
        Energy function
    positions : torch.Tensor
        Atomic positions
    box : Optional[torch.Tensor]
        Simulation box vectors
    chi : torch.Tensor
        Electronegativity in energy/charge unit
    hardness : torch.Tensor
        Atomic hardness in energy/charge^2 unit
    eta : torch.Tensor
        Gaussian width in length unit
    pairs : torch.Tensor
        Tensor of atom pairs
    ds : torch.Tensor
        Distance tensor
    buffer_scales : torch.Tensor
        Buffer scales for each pair

    Returns
    -------
    torch.Tensor
        Hessian matrix with shape (n_atoms, n_atoms)
    """
    n_atoms = positions.shape[0]
    q_tmp = torch.zeros(
        n_atoms, device=positions.device, dtype=positions.dtype, requires_grad=True
    )
    y = func_energy(q_tmp, positions, box, chi, hardness, eta, pairs, ds, buffer_scales)
    grad = torch.autograd.grad(y, q_tmp, retain_graph=True, create_graph=True)
    hessian = []
    for anygrad in grad[0]:
        hessian.append(torch.autograd.grad(anygrad, q_tmp, retain_graph=True)[0])
    hessian = torch.stack(hessian)
    return hessian.reshape([n_atoms, n_atoms])

matinv_optimize(module: QEqForceModule, positions: torch.Tensor, box: Optional[torch.Tensor], chi: torch.Tensor, hardness: torch.Tensor, eta: torch.Tensor, pairs: torch.Tensor, ds: torch.Tensor, buffer_scales: torch.Tensor, constraint_matrix: torch.Tensor, constraint_vals: torch.Tensor, **kwargs)

Function to optimize atomic charges with matrix inversion method.

PARAMETER	DESCRIPTION
module	QEq module TYPE: QEqForceModule
positions	Atomic positions TYPE: Tensor
box	Simulation box vectors TYPE: Optional[Tensor]
chi	Electronegativity in energy/charge unit TYPE: Tensor
hardness	Atomic hardness in energy/charge^2 unit TYPE: Tensor
eta	Gaussian width in length unit TYPE: Tensor
pairs	Tensor of atom pairs TYPE: Tensor
ds	Distance tensor TYPE: Tensor
buffer_scales	Buffer scales for each pair TYPE: Tensor
constraint_matrix	n_const * natoms, constraint matrix TYPE: Tensor
constraint_vals	n_const, constraint values TYPE: Tensor

RETURNS	DESCRIPTION
energy	Energy tensor TYPE: Tensor
q_opt	Optimized atomic charges TYPE: Tensor

Source code in torch_admp/qeq.py

def matinv_optimize(
    module: QEqForceModule,
    positions: torch.Tensor,
    box: Optional[torch.Tensor],
    chi: torch.Tensor,
    hardness: torch.Tensor,
    eta: torch.Tensor,
    pairs: torch.Tensor,
    ds: torch.Tensor,
    buffer_scales: torch.Tensor,
    constraint_matrix: torch.Tensor,
    constraint_vals: torch.Tensor,
    **kwargs,
):
    """
    Function to optimize atomic charges with matrix inversion method.

    Parameters
    ----------
    module : QEqForceModule
        QEq module
    positions : torch.Tensor
        Atomic positions
    box : Optional[torch.Tensor]
        Simulation box vectors
    chi : torch.Tensor
        Electronegativity in energy/charge unit
    hardness : torch.Tensor
        Atomic hardness in energy/charge^2 unit
    eta : torch.Tensor
        Gaussian width in length unit
    pairs : torch.Tensor
        Tensor of atom pairs
    ds : torch.Tensor
        Distance tensor
    buffer_scales : torch.Tensor
        Buffer scales for each pair
    constraint_matrix : torch.Tensor
        n_const * natoms, constraint matrix
    constraint_vals : torch.Tensor
        n_const, constraint values

    Returns
    -------
    energy: torch.Tensor
        Energy tensor
    q_opt: torch.Tensor
        Optimized atomic charges
    """
    device = positions.device
    dtype = positions.dtype
    # calculate hessian
    # n_atoms * n_atoms
    hessian = calc_hessian(
        module.func_energy, positions, box, chi, hardness, eta, pairs, ds, buffer_scales
    )

    # coeff matrix as [[hessian, constraint_matrix.T], [constraint_matrix, 0]]
    # (n_atoms + n_const) * (n_atoms + n_const)
    n_atoms = positions.shape[0]
    if constraint_matrix is None:
        coeff_matrix = hessian
        vector = -chi
    else:
        n_const = constraint_matrix.shape[0]
        coeff_matrix = torch.cat(
            [
                torch.cat([hessian, constraint_matrix.T], dim=1),
                torch.cat(
                    [
                        constraint_matrix,
                        torch.zeros((n_const, n_const), device=device, dtype=dtype),
                    ],
                    dim=1,
                ),
            ],
            dim=0,
        )
        vector = torch.concat([-chi, constraint_vals])

    _q_opt = torch.linalg.solve(coeff_matrix, vector.reshape(-1, 1)).reshape(-1)

    q_opt = _q_opt[:n_atoms].detach()
    q_opt.requires_grad = True
    energy = module.func_energy(
        q_opt, positions, box, chi, hardness, eta, pairs, ds, buffer_scales
    )
    return energy, q_opt

pgrad_optimize(module: QEqForceModule, q0: Optional[torch.Tensor], positions: torch.Tensor, box: Optional[torch.Tensor], chi: torch.Tensor, hardness: torch.Tensor, eta: torch.Tensor, pairs: torch.Tensor, ds: torch.Tensor, buffer_scales: torch.Tensor, constraint_matrix: torch.Tensor, constraint_vals: torch.Tensor, coeff_matrix: Optional[torch.Tensor] = None, reinit_q: bool = False, method: str = 'lbfgs', **kwargs)

Function to optimize atomic charges with projected gradient method.

PARAMETER	DESCRIPTION
module	QEq module TYPE: QEqForceModule
q0	Initial guess for atomic charges, all zeros for None TYPE: Optional[Tensor]
positions	Atomic positions TYPE: Tensor
box	Simulation box vectors TYPE: Optional[Tensor]
chi	Electronegativity in energy/charge unit TYPE: Tensor
hardness	Atomic hardness in energy/charge^2 unit TYPE: Tensor
eta	Gaussian width in length unit TYPE: Tensor
pairs	n_pairs * 2 tensor of pairs TYPE: Tensor
ds	i-j distance tensor TYPE: Tensor
buffer_scales	Buffer scales for each pair, 1 if i < j else 0 TYPE: Tensor
constraint_matrix	n_const * natoms, constraint matrix TYPE: Tensor
constraint_vals	n_const, constraint values TYPE: Tensor
coeff_matrix	n_atoms * n_const, coefficient matrix TYPE: Optional[Tensor] DEFAULT: None
reinit_q	If reinitialize atomic charges based on constraints, by default False TYPE: bool DEFAULT: False
method	Optimization method, by default "lbfgs" TYPE: str DEFAULT: 'lbfgs'

RETURNS	DESCRIPTION
energy	Energy tensor TYPE: Tensor
q_opt	Optimized atomic charges TYPE: Tensor

Source code in torch_admp/qeq.py

def pgrad_optimize(
    module: QEqForceModule,
    q0: Optional[torch.Tensor],
    positions: torch.Tensor,
    box: Optional[torch.Tensor],
    chi: torch.Tensor,
    hardness: torch.Tensor,
    eta: torch.Tensor,
    pairs: torch.Tensor,
    ds: torch.Tensor,
    buffer_scales: torch.Tensor,
    constraint_matrix: torch.Tensor,
    constraint_vals: torch.Tensor,
    coeff_matrix: Optional[torch.Tensor] = None,
    reinit_q: bool = False,
    method: str = "lbfgs",
    **kwargs,
):
    """
    Function to optimize atomic charges with projected gradient method.

    Parameters
    ----------
    module : QEqForceModule
        QEq module
    q0 : Optional[torch.Tensor]
        Initial guess for atomic charges, all zeros for None
    positions : torch.Tensor
        Atomic positions
    box : Optional[torch.Tensor]
        Simulation box vectors
    chi : torch.Tensor
        Electronegativity in energy/charge unit
    hardness : torch.Tensor
        Atomic hardness in energy/charge^2 unit
    eta : torch.Tensor
        Gaussian width in length unit
    pairs : torch.Tensor
        n_pairs * 2 tensor of pairs
    ds : torch.Tensor
        i-j distance tensor
    buffer_scales : torch.Tensor
        Buffer scales for each pair, 1 if i < j else 0
    constraint_matrix : torch.Tensor
        n_const * natoms, constraint matrix
    constraint_vals : torch.Tensor
        n_const, constraint values
    coeff_matrix : Optional[torch.Tensor]
        n_atoms * n_const, coefficient matrix
    reinit_q : bool, optional
        If reinitialize atomic charges based on constraints, by default False
    method : str, optional
        Optimization method, by default "lbfgs"

    Returns
    -------
    energy: torch.Tensor
        Energy tensor
    q_opt: torch.Tensor
        Optimized atomic charges
    """
    """Function to optimize atomic charges with projected gradient method

    Parameters
    ----------
    module : QEqForceModule
        QEq module
    q0 : torch.Tensor
        initial guess for atomic charges, all zeros for None
    positions : torch.Tensor
        atomic positions
    box : torch.Tensor
        simulation box
    chi : torch.Tensor
        eletronegativity in energy / charge unit
    hardness : torch.Tensor
        atomic hardness in energy / charge^2 unit
    eta : torch.Tensor
        Gaussian width in length unit
    pairs : torch.Tensor
        n_pairs * 2 tensor of pairs
    ds : torch.Tensor
        i-j distance tensor
    buffer_scales : torch.Tensor
        buffer scales for each pair, 1 if i < j else 0
    constraint_matrix : torch.Tensor
        n_const * natoms, constraint matrix
    constraint_vals : torch.Tensor
        n_const, constraint values
    coeff_matrix : torch.Tensor
        n_atoms * n_const, coefficient matrix
    reinit_q : bool, optional
        if reinitialize the atomic charges based on constraints, by default False
    method : str, optional
        optimization method, by default "lbfgs"

    Returns
    -------
    energy: torch.Tensor
        energy tensor
    q_opt: torch.Tensor
        optimized atomic charges
    """
    n_atoms = positions.shape[0]
    # n_const = constraint_matrix.shape[0]

    if q0 is None:
        q0 = torch.rand(n_atoms, device=positions.device, dtype=positions.dtype)
        reinit_q = True
    assert q0.shape[0] == n_atoms
    # make sure the initial guess satisfy constraints
    if reinit_q:
        q0 = vector_projection(q0, constraint_matrix, constraint_vals)

    if coeff_matrix is None:
        coeff_matrix = vector_projection_coeff_matrix(constraint_matrix)

    # choose iterative algorithm
    try:
        solver_fn = globals()[f"_pgrad_optimize_{method}"](
            module.func_energy, module.max_iter, module.eps
        )
    except KeyError as exc:
        raise ValueError(f"Method {method} is not supported.") from exc

    with torch.device(positions.device):
        out = custom_root(
            module.optimality,
            argnums=1,
            has_aux=True,
            solve=torchopt.linear_solve.solve_normal_cg(maxiter=5, atol=0),
        )(solver_fn)(
            q0,
            positions,
            box,
            chi,
            hardness,
            eta,
            pairs,
            ds,
            buffer_scales,
            constraint_matrix,
            coeff_matrix,
        )
    if out[1] == -1:
        Warning("Optimization did not converge.")
    module.converge_iter = out[1]
    q_opt = out[0].detach()
    q_opt.requires_grad = True
    energy = module.func_energy(
        q_opt, positions, box, chi, hardness, eta, pairs, ds, buffer_scales
    )
    # forces = -calc_grads(energy, positions)
    return energy, q_opt