"""Relax a species' 3D structure with FIRE on an RDKit force field, constraints held.
A conformer from embedding or a file is pulled to a nearby minimum of
MMFF94 (T. A. Halgren, *J. Comput. Chem.* 17, 490-519 (1996)) or UFF
(A. K. Rappé, C. J. Casewit, K. S. Colwell, W. A. Goddard III and W. M.
Skiff, *J. Am. Chem. Soc.* 114, 10024-10035 (1992)), RDKit supplying the
energy and gradient and :mod:`vibesketch.chem.fire` the minimiser. A job's
constraints hold during the relaxation and are exact, not restraints:
- A fixed atom (a ``position`` constraint) feels no force, through FIRE's
mask.
- A distance, angle or dihedral is held by projection. Before the run and
after every move, the coordinates are put back on the constraint surface
with Newton steps of least norm along the constraint gradients, the SHAKE
idea (J.-P. Ryckaert, G. Ciccotti and H. J. C. Berendsen, *J. Comput.
Phys.* 23, 327-341 (1977)); the force the minimiser sees has its
components along those gradients removed, so it only drives motion that
keeps the constraints. The dihedral's gradient follows A. Blondel and M.
Karplus, *J. Comput. Chem.* 17, 1132-1141 (1996).
The force field is the caller's choice and is not swapped silently: a
species it has no parameters for is refused with the reason. The drawn
stereochemistry is checked on the result as embedding checks it, and a
relaxation that changed it is refused. Everything is deterministic.
Units: angstrom, kcal/mol and kcal/mol/angstrom, degrees in constraints.
"""
from __future__ import annotations
import math
from collections.abc import Sequence
from dataclasses import asdict, dataclass, field
from typing import Any, Literal
import numpy as np
from numpy.typing import NDArray
from vibesketch.chem.fire import minimize
from vibesketch.document.model import Conformer, Constraint, Species
__all__ = ["FORCE_FIELDS", "ConstraintCheck", "RelaxFailed", "Relaxed", "relax_3d"]
_Array = NDArray[np.float64]
ForceField = Literal["mmff94", "uff"]
#: The force fields, in the order they are tried when the caller names none.
FORCE_FIELDS: tuple[ForceField, ...] = ("mmff94", "uff")
#: Convergence: the largest force component, in kcal/mol/angstrom.
DEFAULT_FMAX = 0.05
#: Integration steps before giving up.
DEFAULT_MAX_STEPS = 5000
#: FIRE's time steps, chosen under the stability limit of the stiffest MMFF94 bond.
_DT, _DT_MAX = 0.02, 0.1
#: A constraint is satisfied when its value is this close, in angstrom or radians.
_TOLERANCE = 1e-8
#: Newton steps the projection may take.
_PROJECTION_STEPS = 60
[docs]
class RelaxFailed(ValueError):
"""The species could not be relaxed as asked."""
[docs]
@dataclass(frozen=True)
class ConstraintCheck:
"""A constraint and its value at the end, in angstrom or degrees."""
kind: str
atoms: list[int]
target: float
value: float
[docs]
@property
def deviation(self) -> float:
"""How far the value is from the target, in angstrom or degrees."""
if self.kind == "dihedral":
return abs(_wrap_degrees(self.value - self.target))
return abs(self.value - self.target)
[docs]
def to_dict(self) -> dict[str, Any]:
return {**asdict(self), "deviation": self.deviation}
[docs]
@dataclass(frozen=True)
class Relaxed:
"""A relaxed 3D structure and what the relaxation did."""
#: The new structure, method ``fire+mmff94`` or ``fire+uff``.
conformer: Conformer
force_field: ForceField
#: The force field's energy at the start, in kcal/mol.
energy_before: float
#: The force field's energy at the end, in kcal/mol.
energy: float
#: The largest remaining force component that can still move an atom.
fmax: float
steps: int
converged: bool
constraints: list[ConstraintCheck] = field(default_factory=list)
warnings: list[str] = field(default_factory=list)
[docs]
def to_dict(self) -> dict[str, Any]:
return {
"conformer": self.conformer.model_dump(mode="json"),
"force_field": self.force_field,
"energy_before_kcal_mol": self.energy_before,
"energy_kcal_mol": self.energy,
"fmax_kcal_mol_angstrom": self.fmax,
"steps": self.steps,
"converged": self.converged,
"constraints": [check.to_dict() for check in self.constraints],
"warnings": list(self.warnings),
}
[docs]
def relax_3d(
species: Species,
*,
constraints: Sequence[Constraint] | None = None,
force_field: ForceField = "mmff94",
fmax: float = DEFAULT_FMAX,
max_steps: int = DEFAULT_MAX_STEPS,
) -> Relaxed:
"""Relax ``species``' first conformer, holding ``constraints``.
``force_field`` is ``mmff94`` or ``uff``; a species it has no parameters
for is refused. The run stops once no force component that can still
move an atom exceeds ``fmax`` (kcal/mol/angstrom) or after ``max_steps``
steps, and the result says which. Raises :class:`RelaxFailed` with the
reason when the species has no structure, the conformer does not match
it, a constraint names an atom it lacks or cannot be satisfied, the
force field lacks parameters, or the drawn stereochemistry was lost.
"""
from rdkit import Chem, rdBase
from rdkit.Chem import rdForceFieldHelpers as force_fields
from vibesketch.chem._rdkit import untyped
from vibesketch.chem.embed import _stereo_kept, _target
from vibesketch.chem.species import mol_from_smiles
if force_field not in FORCE_FIELDS:
raise ValueError(f"unknown force field {force_field!r}; use mmff94 or uff")
if fmax <= 0 or max_steps < 0:
raise ValueError("fmax must be positive and max_steps must not be negative")
if species.smiles is None:
raise RelaxFailed("it has no structure to relax")
if not species.conformers:
raise RelaxFailed("it has no 3D structure; embed it first")
conformer = species.conformers[0]
heavy = mol_from_smiles(species.smiles)
with_h = Chem.AddHs(heavy)
symbols = [atom.GetSymbol() for atom in with_h.GetAtoms()]
if conformer.symbols != symbols:
raise RelaxFailed(
f"the conformer's atoms ({' '.join(conformer.symbols)}) are not the "
f"structure's ({' '.join(symbols)}); embed it again"
)
count = len(symbols)
held = list(constraints or [])
for constraint in held:
beyond = [index for index in constraint.atoms if index >= count]
if beyond:
raise RelaxFailed(
f"a {constraint.kind} constraint names atom {beyond[0]}, "
f"but the species has {count} atoms"
)
fields = untyped(force_fields)
blocker = rdBase.BlockLogs()
try:
if force_field == "mmff94":
properties = fields.MMFFGetMoleculeProperties(with_h)
if properties is None or not fields.MMFFHasAllMoleculeParams(with_h):
raise RelaxFailed(
"MMFF94 has no parameters for it (metal centres and "
"unusual elements are outside MMFF94); try force_field='uff'"
)
elif not fields.UFFHasAllMoleculeParams(with_h):
raise RelaxFailed("UFF has no parameters for it")
positions = Chem.Conformer(count)
for index, (x, y, z) in enumerate(conformer.xyz_angstrom):
positions.SetAtomPosition(index, (x, y, z))
conformer_id = int(with_h.AddConformer(positions, assignId=True))
if force_field == "mmff94":
engine = fields.MMFFGetMoleculeForceField(with_h, properties, confId=conformer_id)
else:
engine = fields.UFFGetMoleculeForceField(with_h, confId=conformer_id)
finally:
del blocker
if engine is None:
raise RelaxFailed(f"{force_field} could not be set up for it")
start = np.array(conformer.xyz_angstrom, dtype=np.float64)
fixed = np.zeros(start.shape, dtype=bool)
for constraint in held:
if constraint.kind == "position":
fixed[constraint.atoms[0]] = True
internal = [
_Internal.of(constraint, start) for constraint in held if constraint.kind != "position"
]
surface = _Surface(internal, fixed)
def energy_at(x: _Array) -> float:
return float(engine.CalcEnergy(x.reshape(-1).tolist()))
def objective(x: _Array) -> tuple[float, _Array]:
# The energy first: RDKit's CalcGrad reuses state from the last
# energy evaluation and is only right at those positions (RDKit
# 2026.09, checked against a numerical gradient).
flat = x.reshape(-1).tolist()
energy = float(engine.CalcEnergy(flat))
gradient = np.asarray(engine.CalcGrad(flat), dtype=np.float64).reshape(x.shape)
gradient[fixed] = 0.0
return energy, surface.tangent(gradient, x)
energy_before = energy_at(start)
try:
result = minimize(
objective,
start,
fmax=fmax,
max_steps=max_steps,
dt=_DT,
dt_max=_DT_MAX,
fixed=fixed,
constrain=surface.project,
)
except _Unsatisfiable as error:
raise RelaxFailed(str(error)) from None
except ValueError as error: # a non-finite force field value
raise RelaxFailed(f"the force field failed: {error}") from None
final = Chem.Mol(with_h)
positions = final.GetConformer(conformer_id)
for index, row in enumerate(result.x):
positions.SetAtomPosition(index, row.tolist())
if not _stereo_kept(final, conformer_id, heavy, _target(heavy)):
raise RelaxFailed(
"the relaxation changed the drawn stereochemistry; "
"loosen the constraints or embed it again"
)
checks = [
ConstraintCheck(c.kind, list(c.atoms), c.target_shown, c.value_shown(result.x))
for c in internal
]
warnings = []
if not result.converged:
warnings.append(
f"not converged after {result.steps} steps "
f"(largest force {result.fmax:.3g} kcal/mol/Å, asked {fmax:g})"
)
relaxed = Conformer(
method=f"fire+{force_field}",
symbols=list(conformer.symbols),
xyz_angstrom=[[round(float(v), 6) for v in row] for row in result.x],
)
return Relaxed(
conformer=relaxed,
force_field=force_field,
energy_before=energy_before,
energy=result.energy,
fmax=result.fmax,
steps=result.steps,
converged=result.converged,
constraints=checks,
warnings=warnings,
)
class _Unsatisfiable(ValueError):
"""The constraints cannot all hold at once."""
def _wrap_degrees(value: float) -> float:
"""``value`` brought into (-180, 180]."""
return float(-((-value + 180.0) % 360.0) + 180.0)
def _wrap_radians(value: float) -> float:
return float(-((-value + math.pi) % (2.0 * math.pi)) + math.pi)
@dataclass(frozen=True)
class _Internal:
"""One held internal coordinate: its atoms and target, in angstrom or radians."""
kind: str
atoms: tuple[int, ...]
target: float
@classmethod
def of(cls, constraint: Constraint, positions: _Array) -> _Internal:
atoms = tuple(constraint.atoms)
if constraint.value is None:
target = cls(constraint.kind, atoms, 0.0).value(positions)
elif constraint.kind == "distance":
target = constraint.value
else:
target = math.radians(constraint.value)
return cls(constraint.kind, atoms, target)
@property
def target_shown(self) -> float:
"""The target in angstrom or degrees."""
return self.target if self.kind == "distance" else math.degrees(self.target)
def value_shown(self, positions: _Array) -> float:
value = self.value(positions)
return value if self.kind == "distance" else math.degrees(value)
def value(self, positions: _Array) -> float:
"""The coordinate's value, in angstrom or radians."""
return self.value_and_gradient(positions)[0]
def residual(self, positions: _Array) -> float:
"""How far the value is from the target, wrapped for a dihedral."""
gap = self.value(positions) - self.target
return _wrap_radians(gap) if self.kind == "dihedral" else gap
def value_and_gradient(self, positions: _Array) -> tuple[float, _Array]:
"""The value and its gradient with respect to every coordinate."""
gradient = np.zeros_like(positions)
r = [positions[i] for i in self.atoms]
if self.kind == "distance":
d = r[0] - r[1]
length = float(np.linalg.norm(d))
if length > 0.0:
gradient[self.atoms[0]] = d / length
gradient[self.atoms[1]] = -d / length
return length, gradient
if self.kind == "angle":
u, v = r[0] - r[1], r[2] - r[1]
lu, lv = float(np.linalg.norm(u)), float(np.linalg.norm(v))
cosine = float(np.clip(u @ v / (lu * lv), -1.0, 1.0))
theta = math.acos(cosine)
sine = math.sin(theta)
if sine > 1e-9:
gi = (cosine * u / lu - v / lv) / (lu * sine)
gk = (cosine * v / lv - u / lu) / (lv * sine)
gradient[self.atoms[0]] = gi
gradient[self.atoms[2]] = gk
gradient[self.atoms[1]] = -(gi + gk)
return theta, gradient
# A dihedral, with the IUPAC sign: Blondel and Karplus (1996), eqs. 27-30.
b1, b2, b3 = r[1] - r[0], r[2] - r[1], r[3] - r[2]
n1, n2 = np.cross(b1, b2), np.cross(b2, b3)
lb2 = float(np.linalg.norm(b2))
phi = math.atan2(lb2 * float(b1 @ n2), float(n1 @ n2))
n1sq, n2sq = float(n1 @ n1), float(n2 @ n2)
if n1sq > 1e-18 and n2sq > 1e-18:
gi = -lb2 / n1sq * n1
gl = lb2 / n2sq * n2
s12, s32 = float(b1 @ b2) / lb2**2, float(b3 @ b2) / lb2**2
gradient[self.atoms[0]] = gi
gradient[self.atoms[3]] = gl
gradient[self.atoms[1]] = s32 * gl - (1.0 + s12) * gi
gradient[self.atoms[2]] = s12 * gi - (1.0 + s32) * gl
return phi, gradient
@dataclass(frozen=True)
class _Surface:
"""The constraint surface: where every held internal coordinate has its target."""
internal: list[_Internal]
fixed: NDArray[np.bool_]
def _jacobian(self, positions: _Array) -> tuple[_Array, _Array]:
"""The residuals and the Jacobian over the free coordinates."""
residuals = np.array([c.residual(positions) for c in self.internal])
rows = [c.value_and_gradient(positions)[1] for c in self.internal]
jacobian = np.array([np.where(self.fixed, 0.0, row).reshape(-1) for row in rows])
return residuals, jacobian
def project(self, positions: _Array) -> _Array:
"""``positions`` moved the least to put every held coordinate at its target."""
if not self.internal:
return positions
x = positions.copy()
for _ in range(_PROJECTION_STEPS):
residuals, jacobian = self._jacobian(x)
if float(np.max(np.abs(residuals), initial=0.0)) <= _TOLERANCE:
return x
correction, *_ = np.linalg.lstsq(jacobian, -residuals, rcond=None)
x += correction.reshape(x.shape)
residuals, _ = self._jacobian(x)
worst = max(zip(self.internal, residuals, strict=True), key=lambda pair: abs(pair[1]))
raise _Unsatisfiable(
f"the constraints cannot all hold at once; the {worst[0].kind} of atoms "
f"{', '.join(map(str, worst[0].atoms))} is {abs(worst[1]):.3g} off"
)
def tangent(self, gradient: _Array, positions: _Array) -> _Array:
"""``gradient`` without its components along the held coordinates."""
if not self.internal:
return gradient
_, jacobian = self._jacobian(positions)
flat = gradient.reshape(-1)
multipliers, *_ = np.linalg.lstsq(jacobian.T, flat, rcond=None)
return (flat - jacobian.T @ multipliers).reshape(gradient.shape)