"""
amorphgen.utils.classical
--------------------------
Classical pair potential calculators: Lennard-Jones and Buckingham+Coulomb.
Vectorized NumPy implementation with optional PyTorch GPU acceleration.
Uses ASE's Calculator interface. No external dependencies (no LAMMPS, no GULP).
Note: these are rigid-ion potentials (no core-shell model).
For shell model potentials, use LAMMPS or GULP via ASE.
Usage:
from amorphgen.utils.classical import LennardJonesCalculator, BuckinghamCalculator
# LJ (CPU)
calc = LennardJonesCalculator(
params={("Ar", "Ar"): {"epsilon": 0.0104, "sigma": 3.40}},
cutoff=10.0,
)
# Buckingham + Coulomb (GPU)
calc = BuckinghamCalculator(
params={
("Si", "O"): {"A": 13702.905, "rho": 0.193817, "C": 54.681},
("O", "O"): {"A": 2029.2204, "rho": 0.343645, "C": 192.58},
},
charges={"Si": 2.4, "O": -1.2},
cutoff=10.0,
device="cuda", # "cpu", "cuda", or "mps"
)
"""
from __future__ import annotations
from math import erfc as _erfc, pi as _pi, sqrt as _sqrt
import numpy as np
from ase.calculators.calculator import Calculator, all_changes
from ase.neighborlist import neighbor_list
# Lazy torch import — only when GPU is requested
_torch = None
def _get_torch():
global _torch
if _torch is None:
try:
import torch
except ImportError:
raise ImportError(
"The GPU path of the classical calculators requires PyTorch, "
"which is an optional dependency. Install it with "
"`pip install torch`, or use device='cpu' (vectorised NumPy, "
"no torch needed)."
) from None
_torch = torch
return _torch
def _use_gpu(device: str) -> bool:
"""Check if device requests GPU acceleration."""
return device in ("cuda", "mps")
[docs]
class LennardJonesCalculator(Calculator):
"""
Lennard-Jones pair potential calculator (vectorized).
V(r) = 4 * epsilon * [(sigma/r)^12 - (sigma/r)^6]
Parameters
----------
params : dict
Per-pair LJ parameters, e.g.
{("Ar", "Ar"): {"epsilon": 0.0104, "sigma": 3.40}}
Pairs are symmetric: ("A","B") == ("B","A").
cutoff : float
Cutoff distance in Angstrom.
device : str
"cpu" (default, vectorized NumPy) or "cuda"/"mps" (PyTorch GPU).
"""
implemented_properties = ["energy", "forces"]
def __init__(self, params: dict, cutoff: float = 10.0,
device: str = "cpu", **kwargs):
super().__init__(**kwargs)
self.pair_params = {}
for (s1, s2), p in params.items():
self.pair_params[(s1, s2)] = p
self.pair_params[(s2, s1)] = p
self.cutoff = cutoff
self.device = device
[docs]
def calculate(self, atoms=None, properties=["energy"], system_changes=all_changes):
super().calculate(atoms, properties, system_changes)
symbols = self.atoms.get_chemical_symbols()
n = len(self.atoms)
# Neighbour list (always on CPU via ASE)
ii, jj, dd, DD = neighbor_list("ijdD", self.atoms, cutoff=self.cutoff)
# Unique pairs (i < j) plus atom/own-image pairs (i == j, cutoff >
# L/2), which are listed twice and therefore enter with weight 1/2;
# their force contributions cancel by symmetry.
mask = ii <= jj
ii, jj, dd, DD = ii[mask], jj[mask], dd[mask], DD[mask]
pair_w = np.where(ii == jj, 0.5, 1.0)
if len(ii) == 0:
self.results["energy"] = 0.0
self.results["forces"] = np.zeros((n, 3))
return
# Vectorized parameter lookup via integer element indices
unique_syms = sorted(set(symbols))
sym_to_idx = {s: i for i, s in enumerate(unique_syms)}
n_types = len(unique_syms)
atom_type = np.array([sym_to_idx[s] for s in symbols])
eps_table = np.zeros((n_types, n_types))
sig_table = np.zeros((n_types, n_types))
valid_table = np.zeros((n_types, n_types), dtype=bool)
for (s1, s2), p in self.pair_params.items():
if s1 in sym_to_idx and s2 in sym_to_idx:
i1, i2 = sym_to_idx[s1], sym_to_idx[s2]
eps_table[i1, i2] = eps_table[i2, i1] = p["epsilon"]
sig_table[i1, i2] = sig_table[i2, i1] = p["sigma"]
valid_table[i1, i2] = valid_table[i2, i1] = True
ti, tj = atom_type[ii], atom_type[jj]
eps_arr = eps_table[ti, tj] * pair_w
sig_arr = sig_table[ti, tj]
valid = valid_table[ti, tj]
if not np.any(valid):
self.results["energy"] = 0.0
self.results["forces"] = np.zeros((n, 3))
return
# Slice to valid pairs only
ii_v, jj_v = ii[valid], jj[valid]
dd_v, DD_v = dd[valid], DD[valid]
eps_v, sig_v = eps_arr[valid], sig_arr[valid]
if _use_gpu(self.device):
energy, forces = self._calc_gpu(n, ii_v, jj_v, dd_v, DD_v,
eps_v, sig_v)
else:
energy, forces = self._calc_cpu(n, ii_v, jj_v, dd_v, DD_v,
eps_v, sig_v)
self.results["energy"] = energy
self.results["forces"] = forces
@staticmethod
def _calc_cpu(n, ii, jj, dd, DD, eps, sig):
sr6 = (sig / dd) ** 6
sr12 = sr6 ** 2
# Energy
e_pairs = 4.0 * eps * (sr12 - sr6)
energy = float(np.sum(e_pairs))
# Forces: dV/dr * (r_vec / r)
dVdr = 4.0 * eps * (-12.0 * sr12 / dd + 6.0 * sr6 / dd)
f_vecs = dVdr[:, None] * DD / dd[:, None]
forces = np.zeros((n, 3))
np.add.at(forces, ii, f_vecs)
np.add.at(forces, jj, -f_vecs)
return energy, forces
def _calc_gpu(self, n, ii, jj, dd, DD, eps, sig):
torch = _get_torch()
dev = self.device
dtype = torch.float32 if dev == "mps" else torch.float64
dd_t = torch.tensor(dd, dtype=dtype, device=dev)
DD_t = torch.tensor(DD, dtype=dtype, device=dev)
eps_t = torch.tensor(eps, dtype=dtype, device=dev)
sig_t = torch.tensor(sig, dtype=dtype, device=dev)
ii_t = torch.tensor(ii, dtype=torch.long, device=dev)
jj_t = torch.tensor(jj, dtype=torch.long, device=dev)
sr6 = (sig_t / dd_t) ** 6
sr12 = sr6 ** 2
energy = float(torch.sum(4.0 * eps_t * (sr12 - sr6)).cpu())
dVdr = 4.0 * eps_t * (-12.0 * sr12 / dd_t + 6.0 * sr6 / dd_t)
f_vecs = dVdr.unsqueeze(1) * DD_t / dd_t.unsqueeze(1)
forces = torch.zeros(n, 3, dtype=dtype, device=dev)
forces.index_add_(0, ii_t, f_vecs)
forces.index_add_(0, jj_t, -f_vecs)
return energy, forces.cpu().numpy().astype(np.float64)
[docs]
class BuckinghamCalculator(Calculator):
"""
Buckingham + Coulomb pair potential calculator (rigid-ion, vectorized).
V(r) = A * exp(-r/rho) - C/r^6 + q_i * q_j / (4*pi*eps0*r)
The Coulomb term is evaluated by Ewald summation (default; real-space
sum over the neighbour list + reciprocal-space sum + self term) or, if
``coulomb_method="wolf"``, by the damped-shifted Wolf sum.
This is an approximation to full Ewald -- accurate for amorphous
structures but not for crystalline long-range order.
Parameters
----------
params : dict
Per-pair Buckingham parameters, e.g.
{("Si", "O"): {"A": 13702.905, "rho": 0.193817, "C": 54.681}}
charges : dict
Per-element charges in electron units, e.g.
{"Si": 2.4, "O": -1.2}
cutoff : float
Cutoff distance in Angstrom.
alpha : float, optional
Splitting/damping parameter (1/Angstrom). Default ``3.5/cutoff`` for
Ewald (real-space part converged at the cutoff), 0.2 for Wolf.
coulomb_method : {"ewald", "wolf"}, default "ewald"
Ewald reproduces the exact periodic Coulomb energy and forces; the
Wolf sum is faster but only approximate (~10 % force error for an
ionic melt at a 10 A cutoff).
coulomb : bool
If True, include Coulomb interactions. Default True.
device : str
"cpu" (default, vectorized NumPy) or "cuda"/"mps" (PyTorch GPU).
"""
implemented_properties = ["energy", "forces"]
# Coulomb constant: e^2 / (4*pi*eps0) in eV*A
_KE = 14.3996
def __init__(self, params: dict, charges: dict | None = None,
cutoff: float = 10.0, alpha: float | None = None,
coulomb: bool = True, coulomb_method: str = "ewald",
device: str = "cpu", **kwargs):
super().__init__(**kwargs)
self.pair_params = {}
for (s1, s2), p in params.items():
self.pair_params[(s1, s2)] = p
self.pair_params[(s2, s1)] = p
self.charges = charges or {}
self.cutoff = cutoff
coulomb_method = (coulomb_method or "ewald").lower()
if coulomb_method not in ("ewald", "wolf"):
raise ValueError(
f"coulomb_method must be 'ewald' or 'wolf', got {coulomb_method!r}")
self.coulomb_method = coulomb_method
# Ewald: alpha = 3.5/rc makes erfc(alpha*rc) ~ 7e-7, so the real-space
# sum is converged at the pair cutoff. Wolf: 0.2 is the customary
# damping (note the Wolf scheme is only ~10 % accurate in forces for
# ionic melts at rc ~ 10 A; it is kept for speed / comparison).
self.alpha = float(alpha) if alpha is not None else (
3.5 / cutoff if coulomb_method == "ewald" else 0.2)
self.coulomb = coulomb
self.device = device
[docs]
def calculate(self, atoms=None, properties=["energy"], system_changes=all_changes):
super().calculate(atoms, properties, system_changes)
symbols = self.atoms.get_chemical_symbols()
n = len(self.atoms)
alpha = self.alpha
rc = self.cutoff
# Neighbour list (CPU)
ii, jj, dd, DD = neighbor_list("ijdD", self.atoms, cutoff=rc)
# Unique pairs: i < j, plus i == j pairs (an atom with its own
# periodic image, which occur when cutoff > L/2). The i == j image
# pairs are listed twice (+D and -D), so they enter the energy with
# weight 1/2; their force contributions cancel by symmetry.
mask = ii <= jj
ii, jj, dd, DD = ii[mask], jj[mask], dd[mask], DD[mask]
pair_w = np.where(ii == jj, 0.5, 1.0)
n_pairs = len(ii)
# Vectorized parameter lookup via integer element indices
unique_syms = sorted(set(symbols))
sym_to_idx = {s: i for i, s in enumerate(unique_syms)}
n_types = len(unique_syms)
atom_type = np.array([sym_to_idx[s] for s in symbols])
# Build pair-type lookup tables
A_table = np.zeros((n_types, n_types))
rho_table = np.zeros((n_types, n_types))
C_table = np.zeros((n_types, n_types))
valid_table = np.zeros((n_types, n_types), dtype=bool)
q_table = np.zeros(n_types)
for (s1, s2), p in self.pair_params.items():
if s1 in sym_to_idx and s2 in sym_to_idx:
i1, i2 = sym_to_idx[s1], sym_to_idx[s2]
A_table[i1, i2] = A_table[i2, i1] = p["A"]
rho_table[i1, i2] = rho_table[i2, i1] = p["rho"]
C_table[i1, i2] = C_table[i2, i1] = p.get("C", 0.0)
valid_table[i1, i2] = valid_table[i2, i1] = True
if self.coulomb and self.charges:
for s, q in self.charges.items():
if s in sym_to_idx:
q_table[sym_to_idx[s]] = q
# Lookup per-pair parameters via integer indexing (fully vectorized)
ti = atom_type[ii]
tj = atom_type[jj]
A_arr = A_table[ti, tj] * pair_w
rho_arr = rho_table[ti, tj]
C_arr = C_table[ti, tj] * pair_w
buck_valid = valid_table[ti, tj]
qi_arr = q_table[ti] * pair_w
qj_arr = q_table[tj]
# Coulomb bookkeeping. Ewald: the pair kernel below handles the
# real-space erfc part (unshifted); the reciprocal-space and self
# terms are added afterwards. Wolf: shifted pair kernel + Wolf
# self term.
use_ewald = self.coulomb and bool(self.charges) and \
self.coulomb_method == "ewald"
wolf_self = 0.0
if self.coulomb and self.charges:
q_atoms = np.array([self.charges.get(s, 0.0) for s in symbols])
if use_ewald:
wolf_self = -self._KE * np.sum(q_atoms ** 2) * alpha / _sqrt(_pi)
else:
wolf_self = -self._KE * np.sum(q_atoms ** 2) * (
_erfc(alpha * rc) / (2.0 * rc) + alpha / _sqrt(_pi)
)
# For Ewald the pair kernel must not apply the Wolf shift: pass
# rc_shift = inf so erfc(alpha*rc)/rc -> 0 in the kernel.
rc_kernel = np.inf if use_ewald else rc
if _use_gpu(self.device):
energy, forces = self._calc_gpu(
n, ii, jj, dd, DD,
A_arr, rho_arr, C_arr, buck_valid,
qi_arr, qj_arr, alpha, rc_kernel, wolf_self,
)
else:
energy, forces = self._calc_cpu(
n, ii, jj, dd, DD,
A_arr, rho_arr, C_arr, buck_valid,
qi_arr, qj_arr, alpha, rc_kernel, wolf_self,
)
if use_ewald:
if self.atoms.get_volume() <= 0 or not any(self.atoms.pbc):
raise ValueError(
"Ewald Coulomb needs a periodic cell (pbc=True, non-zero "
"volume); use coulomb_method='wolf' with a large cutoff "
"for an isolated cluster.")
e_rec, f_rec = self._ewald_reciprocal(
self.atoms.cell[:], self.atoms.get_positions(), q_atoms, alpha)
energy += e_rec
forces = forces + f_rec
self.results["energy"] = energy
self.results["forces"] = forces
@staticmethod
def _ewald_reciprocal(cell, positions, q, alpha, tol=1e-8):
"""Reciprocal-space Ewald energy and forces (numpy, any cell shape).
E_k = (KE 2pi/V) sum_{k!=0} exp(-k^2/4alpha^2)/k^2 |S(k)|^2,
S(k) = sum_j q_j exp(i k.r_j); F_j = -dE_k/dr_j.
k-vectors are included while exp(-k^2/4alpha^2) > tol.
"""
KE = BuckinghamCalculator._KE
cell = np.asarray(cell, dtype=float)
vol = abs(np.linalg.det(cell))
rec = 2.0 * np.pi * np.linalg.inv(cell).T # rows: b1, b2, b3
k_max = 2.0 * alpha * np.sqrt(-np.log(tol))
# bounds per direction: |n_i| <= k_max / |b_i| is not exact for
# non-orthogonal cells; use the perpendicular widths instead.
widths = 2.0 * np.pi / np.linalg.norm(cell, axis=1) # lower bound on |b_i|
nmax = np.ceil(k_max / widths).astype(int)
rng = [np.arange(-m, m + 1) for m in nmax]
n = np.array(np.meshgrid(*rng, indexing="ij")).reshape(3, -1).T
n = n[np.any(n != 0, axis=1)]
k = n @ rec # (Nk, 3)
k2 = np.einsum("ij,ij->i", k, k)
keep = k2 <= k_max ** 2
k, k2 = k[keep], k2[keep]
A = np.exp(-k2 / (4.0 * alpha ** 2)) / k2 # (Nk,)
phase = k @ positions.T # (Nk, N)
cos_p, sin_p = np.cos(phase), np.sin(phase)
S_re = cos_p @ q
S_im = sin_p @ q
pref = KE * 2.0 * np.pi / vol
e_rec = pref * float(np.sum(A * (S_re ** 2 + S_im ** 2)))
# Uniform neutralising background for a non-neutral cell (removes
# the alpha dependence of the energy); zero for neutral systems.
q_tot = float(np.sum(q))
if abs(q_tot) > 1e-12:
e_rec -= KE * np.pi * q_tot ** 2 / (2.0 * vol * alpha ** 2)
# Im[S* e^{ik.r_j}] = S_re sin(k.r_j) - S_im cos(k.r_j)
im = S_re[:, None] * sin_p - S_im[:, None] * cos_p # (Nk, N)
f = 2.0 * pref * q[None, :] * ((A[:, None] * im).T @ k).T # (3, N)
return e_rec, f.T
@staticmethod
def _calc_cpu(n, ii, jj, dd, DD,
A, rho, C, buck_valid, qi, qj,
alpha, rc, wolf_self):
KE = BuckinghamCalculator._KE
energy = wolf_self
forces = np.zeros((n, 3))
# --- Buckingham ---
if np.any(buck_valid):
bm = buck_valid
r_b = dd[bm]
exp_term = A[bm] * np.exp(-r_b / rho[bm])
r6_term = C[bm] / r_b ** 6
energy += float(np.sum(exp_term - r6_term))
dVdr_buck = -A[bm] / rho[bm] * np.exp(-r_b / rho[bm]) + 6.0 * C[bm] / r_b ** 7
f_buck = dVdr_buck[:, None] * DD[bm] / r_b[:, None]
np.add.at(forces, ii[bm], f_buck)
np.add.at(forces, jj[bm], -f_buck)
# --- Coulomb (Wolf) ---
coul_mask = (qi != 0.0) & (qj != 0.0)
if np.any(coul_mask):
cm = coul_mask
r_c = dd[cm]
qiqj = KE * qi[cm] * qj[cm]
# erfc via scipy for array operation
from scipy.special import erfc as _erfc_arr
erfc_ar = _erfc_arr(alpha * r_c)
erfc_arc = _erfc_arr(alpha * rc)
e_coul = qiqj * (erfc_ar / r_c - erfc_arc / rc)
energy += float(np.sum(e_coul))
dVdr_coul = qiqj * (
-erfc_ar / r_c ** 2
- 2.0 * alpha / _sqrt(_pi) * np.exp(-(alpha * r_c) ** 2) / r_c
)
f_coul = dVdr_coul[:, None] * DD[cm] / r_c[:, None]
np.add.at(forces, ii[cm], f_coul)
np.add.at(forces, jj[cm], -f_coul)
return energy, forces
def _calc_gpu(self, n, ii, jj, dd, DD,
A, rho, C, buck_valid, qi, qj,
alpha, rc, wolf_self):
torch = _get_torch()
dev = self.device
KE = self._KE
dtype = torch.float32 if dev == "mps" else torch.float64
dd_t = torch.tensor(dd, dtype=dtype, device=dev)
DD_t = torch.tensor(DD, dtype=dtype, device=dev)
ii_t = torch.tensor(ii, dtype=torch.long, device=dev)
jj_t = torch.tensor(jj, dtype=torch.long, device=dev)
energy = wolf_self
forces = torch.zeros(n, 3, dtype=dtype, device=dev)
# --- Buckingham ---
bm_t = torch.tensor(buck_valid, dtype=torch.bool, device=dev)
if torch.any(bm_t):
A_t = torch.tensor(A, dtype=dtype, device=dev)[bm_t]
rho_t = torch.tensor(rho, dtype=dtype, device=dev)[bm_t]
C_t = torch.tensor(C, dtype=dtype, device=dev)[bm_t]
r_b = dd_t[bm_t]
D_b = DD_t[bm_t]
ii_b = ii_t[bm_t]
jj_b = jj_t[bm_t]
exp_term = A_t * torch.exp(-r_b / rho_t)
r6_term = C_t / r_b ** 6
energy += float(torch.sum(exp_term - r6_term).cpu())
dVdr = -A_t / rho_t * torch.exp(-r_b / rho_t) + 6.0 * C_t / r_b ** 7
f_b = dVdr.unsqueeze(1) * D_b / r_b.unsqueeze(1)
forces.index_add_(0, ii_b, f_b)
forces.index_add_(0, jj_b, -f_b)
# --- Coulomb (Wolf) ---
qi_t = torch.tensor(qi, dtype=dtype, device=dev)
qj_t = torch.tensor(qj, dtype=dtype, device=dev)
cm_t = (qi_t != 0.0) & (qj_t != 0.0)
if torch.any(cm_t):
r_c = dd_t[cm_t]
D_c = DD_t[cm_t]
ii_c = ii_t[cm_t]
jj_c = jj_t[cm_t]
qiqj = KE * qi_t[cm_t] * qj_t[cm_t]
erfc_ar = torch.erfc(alpha * r_c)
erfc_arc = float(_erfc(alpha * rc))
e_coul = qiqj * (erfc_ar / r_c - erfc_arc / rc)
energy += float(torch.sum(e_coul).cpu())
dVdr_c = qiqj * (
-erfc_ar / r_c ** 2
- 2.0 * alpha / _sqrt(_pi) * torch.exp(-(alpha * r_c) ** 2) / r_c
)
f_c = dVdr_c.unsqueeze(1) * D_c / r_c.unsqueeze(1)
forces.index_add_(0, ii_c, f_c)
forces.index_add_(0, jj_c, -f_c)
return energy, forces.cpu().numpy().astype(np.float64)