Source code for ufp.terms.zbl

"""ZBL screened nuclear repulsion term."""

from __future__ import annotations

from collections.abc import Sequence

import torch

from ufp.core.input import UFPInput
from ufp.core.output import UFPOutput
from ufp.terms._base import PairTerm
from ufp.terms._shared import (
    accumulate_pair_energy_forces,
    empty_atomwise_output,
    pair_weight,
)


_COULOMB_EV_ANGSTROM = 14.3996454784255
_SCREENING_LENGTH_FACTOR = 0.46850
_ZBL_COEFFS = (0.1818, 0.5099, 0.2802, 0.02817)
_ZBL_EXPONENTS = (3.2, 0.9423, 0.4029, 0.2016)




class ZBLRepulsionTerm(PairTerm):
    """
    Analytic Ziegler-Biersack-Littmark screened nuclear repulsion.
    """

    def __init__(
        self,
        *,
        cutoff: float,
        atomic_types: Sequence[int] | None = None,
        eps: float = 1.0e-12,
    ) -> None:
        """Store a fast non-fittable pair repulsion term."""
        super().__init__(cutoff=cutoff, atomic_types=atomic_types)
        self.eps = float(eps)

    @property
    def provides_forces(self) -> bool:
        """Report that this term produces analytic pair forces."""
        return True



    def forward(self, inputs: UFPInput) -> UFPOutput:
        """Evaluate ZBL energy and force contributions over covered pairs."""
        if inputs.neighbor_list is None:
            raise RuntimeError("ZBLRepulsionTerm requires a neighbor list")
        pair_mask = None
        if self.atomic_types is not None:
            pair_category = inputs.pair_category_indices(
                self.atomic_types,
                symmetric=True,
            )
            pair_mask = pair_category >= 0
            if not torch.any(pair_mask):
                return empty_atomwise_output(inputs, forces=True)

        first_z, second_z = inputs.pair_atomic_numbers(pair_mask)
        distances = inputs.pair_distances(pair_mask).clamp_min(self.eps)
        dtype = inputs.dtype
        device = inputs.device
        z1 = first_z.to(dtype=dtype, device=device)
        z2 = second_z.to(dtype=dtype, device=device)
        screening = _SCREENING_LENGTH_FACTOR / (
            torch.pow(z1, 0.23) + torch.pow(z2, 0.23)
        )
        scaled = distances / screening
        coeffs = torch.tensor(_ZBL_COEFFS, dtype=dtype, device=device)
        exponents = torch.tensor(_ZBL_EXPONENTS, dtype=dtype, device=device)
        exp_terms = torch.exp(-scaled[:, None] * exponents[None, :])
        phi = torch.sum(coeffs[None, :] * exp_terms, dim=1)
        dphi_dr = (
            -torch.sum(
                coeffs[None, :] * exponents[None, :] * exp_terms,
                dim=1,
            )
            / screening
        )
        prefactor = _COULOMB_EV_ANGSTROM * z1 * z2
        scale = pair_weight(inputs)
        weighted_energy = scale * prefactor * phi / distances
        weighted_grad = (
            scale * prefactor * (dphi_dr / distances - phi / (distances * distances))
        )
        return accumulate_pair_energy_forces(
            inputs,
            pair_mask,
            weighted_pair_energy=weighted_energy,
            weighted_pair_grad=weighted_grad,
            eps=self.eps,
        )




__all__ = [
    "ZBLRepulsionTerm",
]