"""Versioned 2D depiction restraints, with analytic gradients for FIRE.
Numeric drawing conventions: J. Brecher, Pure Appl. Chem. 80, 277–410
(2008), doi:10.1351/pac200880020277. The packaged table names each rule's
coverage and section. These are drawing units, never physical geometry.
Force constants, clearance radii and convergence limits are our engineering
choices, not measurements or numbers attributed to IUPAC.
Complex rings take a generated layout's geometry rather than imposing
mutually incompatible regular polygons: the seed's, or a better one among the
layouts the caller offers (:meth:`DepictionField.build`). RDKit supplies
those layouts through its public depictor API; no template data is copied
into this package.
"""
from __future__ import annotations
import json
from collections.abc import Sequence
from dataclasses import dataclass
from functools import lru_cache
from importlib.resources import files
from itertools import combinations
from typing import Any
import numpy as np
from numpy.typing import NDArray
from vibesketch.chem.fire import minimize
_Array = NDArray[np.float64]
_Indices = NDArray[np.intp]
# Deliberately limited to ordinary organic drawing environments. Metals and
# unusual coordination retain the existing depictor layout.
_ORGANIC = frozenset((1, 5, 6, 7, 8, 9, 14, 15, 16, 17, 35, 53))
@lru_cache(maxsize=1)
def _table() -> dict[str, Any]:
data: dict[str, Any] = json.loads(
files("vibesketch").joinpath("chem/data/depiction.json").read_text(encoding="utf-8")
)
if data["version"] != 1:
raise ValueError("unsupported depiction table version")
return data
def _number(rule: str) -> float:
return float(_table()["rules"][rule]["value"])
def _wrap(angles: _Array) -> _Array:
return (angles + np.pi) % (2 * np.pi) - np.pi
def _angles(points: _Array, triples: _Indices) -> _Array:
a, centre, b = triples.T
p, q = points[a] - points[centre], points[b] - points[centre]
return _wrap(np.arctan2(q[:, 1], q[:, 0]) - np.arctan2(p[:, 1], p[:, 0]))
def _scatter(gradient: _Array, index: _Indices, values: _Array) -> None:
"""``gradient[index] += values``, repeated indices summed (faster than ``add.at``)."""
for axis in (0, 1):
gradient[:, axis] += np.bincount(index, weights=values[:, axis], minlength=len(gradient))
def _systems(rings: list[set[int]]) -> list[set[int]]:
"""Rings that share an atom, joined into ring systems."""
systems: list[set[int]] = []
for ring in rings:
joined = [s for s in systems if s & ring]
merged = set(ring).union(*joined)
systems = [s for s in systems if not s & ring] + [merged]
return systems
def _shape_quality(
points: _Array, members: list[int], inside: list[tuple[int, int]]
) -> tuple[int, float]:
"""How well ``points`` draw one ring system: crowded atom pairs, then bond spread."""
lengths = np.array([np.linalg.norm(points[a] - points[b]) for a, b in inside])
if not len(lengths) or lengths.mean() <= 1e-9 or not np.isfinite(points[members]).all():
return len(members) ** 2, np.inf
scaled = points[members] / np.median(lengths)
distances = np.linalg.norm(scaled[:, None] - scaled[None], axis=2)
position = {atom: index for index, atom in enumerate(members)}
for a, b in inside:
distances[position[a], position[b]] = distances[position[b], position[a]] = np.inf
crowded = distances[np.triu_indices(len(members), 1)] < _table()["numerics"]["atom_clearance"]
return int(np.count_nonzero(crowded)), float(lengths.std() / lengths.mean())
def _shape(
shapes: Sequence[_Array], members: list[int], inside: list[tuple[int, int]]
) -> tuple[_Array, float]:
"""The layout a ring system takes, and the factor that gives its bonds a median of one.
The first layout is kept unless a later one crowds the system's atoms less,
or draws its bonds more evenly by more than the table's ``shape_margin``.
"""
margin = _table()["numerics"]["shape_margin"]
chosen, best = shapes[0], _shape_quality(shapes[0], members, inside)
for candidate in shapes[1:]:
quality = _shape_quality(candidate, members, inside)
if quality[0] < best[0] or (quality[0] == best[0] and quality[1] < best[1] - margin):
chosen, best = candidate, quality
lengths = [np.linalg.norm(chosen[a] - chosen[b]) for a, b in inside]
median = float(np.median(lengths)) if lengths else 1.0
return chosen, 1.0 / median if median > 1e-9 else 1.0
def _interior(points: _Array, ring: Sequence[int], atom: int) -> float:
"""The interior angle of ``ring`` at ``atom`` as ``points`` draw it, in radians.
The polygon's orientation (its signed area) tells the inside from the
outside, so a reflex corner of a large ring is measured correctly.
"""
polygon = points[list(ring)]
x, y = polygon[:, 0], polygon[:, 1]
area = float(x @ np.roll(y, -1) - y @ np.roll(x, -1))
i = list(ring).index(atom)
incoming = polygon[i] - polygon[i - 1]
outgoing = polygon[(i + 1) % len(ring)] - polygon[i]
turn = float(
np.arctan2(
incoming[0] * outgoing[1] - incoming[1] * outgoing[0], float(incoming @ outgoing)
)
)
return np.pi - turn if area > 0 else np.pi + turn
def _ring_atom_angles(
centre: int, ring_side: list[int], others: list[int], interior: float, seed: _Array
) -> dict[tuple[int, int], float]:
"""Signed targets at a ring atom whose one or two substituents share the outside.
One substituent bisects the angle outside the ring (GR-4.2.1); two sit the
table's ``ring_geminal`` angle apart, centred on that bisector (GR-4.2.3).
Which way round they go is read from the seed; a seed that puts the
substituents on both sides of the ring gets no targets here.
"""
exterior = 2 * np.pi - interior
if len(others) == 1:
offsets = [exterior * _number("ring_substituent")]
else:
gap = np.deg2rad(_number("ring_geminal"))
if exterior <= gap:
return {}
offsets = [(exterior - gap) / 2, (exterior + gap) / 2]
order = sorted(
[*ring_side, *others],
key=lambda i: np.arctan2(*(seed[i] - seed[centre])[::-1]),
)
first = order.index(ring_side[0])
order = order[first:] + order[:first]
if order[-1] != ring_side[1]:
# The substituents follow the other ring neighbour going round.
first = order.index(ring_side[1])
order = order[first:] + order[:first]
if order[-1] != ring_side[0]:
return {}
positions = dict(zip(order, [0.0, *offsets, exterior], strict=True))
return {
(min(a, b), max(a, b)): float(_wrap(np.array(positions[max(a, b)] - positions[min(a, b)])))
for a, b in combinations(order, 2)
}
def _four_angles(atom: Any, seed: _Array) -> dict[tuple[int, int], float]:
"""Select only the unambiguous GR-4.1.3 styles, preserving cyclic order."""
neighbours = list(atom.GetNeighbors())
terminal = {a.GetIdx() for a in neighbours if a.GetDegree() == 1}
cross = (
len(terminal) == 4
or all(a.GetDegree() >= 3 for a in neighbours)
or sum(b.GetBondTypeAsDouble() == 2 for b in atom.GetBonds()) >= 2
)
if not cross and len(terminal) != 2:
return {}
centre = atom.GetIdx()
order = sorted(
(a.GetIdx() for a in neighbours),
key=lambda i: np.arctan2(*(seed[i] - seed[centre])[::-1]),
)
if not cross:
# Only adjacent terminal neighbours can bound the 60 degree gap
# without exchanging substituents (which might change stereo).
start = next((i for i in range(4) if {order[i], order[(i + 1) % 4]} == terminal), None)
if start is None:
return {}
order = order[start:] + order[:start]
rule = "four_cross" if cross else "four_chain"
gaps = np.deg2rad(_table()["rules"][rule]["value"])
positions = dict(zip(order, np.concatenate(([0.0], np.cumsum(gaps)[:-1])), strict=True))
return {
(min(a, b), max(a, b)): float(_wrap(np.array(positions[max(a, b)] - positions[min(a, b)])))
for a, b in combinations(order, 2)
}
[docs]
@dataclass(frozen=True)
class DepictionField:
"""One molecule's fixed restraints; metrics use the same targets before/after."""
seed: _Array
bonds: _Indices
lengths: _Array
triples: _Indices
angles: _Array
pairs: _Indices
clearance: _Array
ring_pairs: _Indices
ring_distances: _Array
[docs]
@classmethod
def build(cls, mol: Any, seed: _Array, shapes: Sequence[_Array] = ()) -> DepictionField:
"""The restraints of ``mol``; orientation and unruled angles come from ``seed``.
``shapes`` are layouts of ``mol`` (atom order kept) that each fused,
bridged, spiro or large ring system may take its geometry from
(GR-3.3.2 to GR-3.3.4): bond lengths, ring distances and the angles at
its atoms. Per system the first layout is kept unless a later one is
clearly better (:func:`_shape`). Without ``shapes`` each system keeps
the seed's geometry. Callers comparing layouts pass every field the
same ``shapes``, so all of them share the same targets.
"""
from rdkit import Chem
ring_info = mol.GetRingInfo()
atom_rings = [tuple(r) for r in ring_info.AtomRings()]
rings = [set(r) for r in atom_rings]
small = [
r
for r in rings
if len(r) <= _number("ring_max") and all(r is other or not r & other for other in rings)
]
systems = _systems([r for r in rings if r not in small])
bonds = np.array(
[(b.GetBeginAtomIdx(), b.GetEndAtomIdx()) for b in mol.GetBonds()], dtype=np.intp
).reshape(-1, 2)
# Each complex system's reference layout and the factor scaling it.
reference: dict[int, tuple[_Array, float]] = {}
for system in systems:
members = sorted(system)
inside = [(int(a), int(b)) for a, b in bonds if a in system and b in system]
chosen = _shape(shapes, members, inside) if shapes else (seed, 1.0)
reference.update(dict.fromkeys(members, chosen))
lengths = np.full(len(bonds), _number("bond"))
# GR-1.1 permits unequal lengths in complex polycycles. Keep the
# reference layout's exception rather than flattening a bridge into a
# regular ring.
for index, bond in enumerate(mol.GetBonds()):
a, b = bonds[index]
if bond.IsInRing() and a in reference and b in reference:
points, scale = reference[a]
lengths[index] = scale * np.linalg.norm(points[a] - points[b])
triples, targets = [], []
for atom in mol.GetAtoms():
c = atom.GetIdx()
neighbours = sorted(a.GetIdx() for a in atom.GetNeighbors())
degree = len(neighbours)
four = _four_angles(atom, seed) if degree == 4 and not atom.IsInRing() else {}
ring = next((r for r in small if c in r), None)
ruled: dict[tuple[int, int], float] = {}
ring_side = [n for n in neighbours if mol.GetBondBetweenAtoms(c, n).IsInRing()]
others = [n for n in neighbours if n not in ring_side]
if len(ring_side) == 2 and 1 <= len(others) <= 2 and ring_info.NumAtomRings(c) == 1:
own = next(r for r in atom_rings if c in r)
if ring is not None:
interior = np.pi - 2 * np.pi / len(ring)
else:
points, _ = reference[c]
interior = _interior(points, own, c)
ruled = _ring_atom_angles(c, ring_side, others, interior, seed)
for a, b in combinations(neighbours, 2):
triple = np.array([[a, c, b]], dtype=np.intp)
observed = float(_angles(seed, triple)[0])
target = observed
if (a, b) in ruled:
target = ruled[a, b]
elif c in reference:
points, _ = reference[c]
target = np.copysign(abs(float(_angles(points, triple)[0])), observed)
elif ring is not None and degree <= 3:
interior = 180.0 - 360.0 / len(ring)
degrees = (
interior
if a in ring and b in ring
else (360 - interior) * _number("ring_substituent")
)
target = np.copysign(np.deg2rad(degrees), observed)
elif not atom.IsInRing() and atom.GetAtomicNum() in _ORGANIC:
if degree == 2:
rule = "linear" if str(atom.GetHybridization()) == "SP" else "bent"
target = np.copysign(np.deg2rad(_number(rule)), observed)
elif degree == 3:
target = np.copysign(np.deg2rad(_number("trigonal")), observed)
elif (a, b) in four:
target = four[a, b]
triples.append((a, c, b))
targets.append(target)
bonded = {frozenset(pair) for pair in bonds}
pairs = np.array(
[
(a, b)
for a, b in combinations(range(len(seed)), 2)
if frozenset((a, b)) not in bonded
],
dtype=np.intp,
).reshape(-1, 2)
numeric = _table()["numerics"]
# A conservative label proxy, not a font bounding box: charged,
# isotopic and hetero atoms reserve more room than skeletal carbon.
labelled = np.array(
[
a.GetAtomicNum() != 6 or a.GetFormalCharge() != 0 or a.GetIsotope() != 0
for a in mol.GetAtoms()
]
)
# Atoms two bonds apart get their own clearance, unless one is labelled.
topological = Chem.GetDistanceMatrix(mol)
two_bonds = np.array([topological[a][b] == 2 for a, b in pairs], dtype=bool)
clearance = np.where(
labelled[pairs[:, 0]] | labelled[pairs[:, 1]],
numeric["label_clearance"],
np.where(two_bonds, numeric["two_bond_clearance"], numeric["atom_clearance"]),
)
ring_pairs = np.array(
sorted(
{
tuple(sorted(pair))
for r in rings
if r not in small
for pair in combinations(r, 2)
}
),
dtype=np.intp,
).reshape(-1, 2)
distances = np.array(
[
reference[a][1] * np.linalg.norm(reference[a][0][a] - reference[a][0][b])
for a, b in ring_pairs
]
)
return cls(
seed.copy(),
bonds,
lengths,
np.array(triples, dtype=np.intp).reshape(-1, 3),
np.array(targets),
pairs,
clearance,
ring_pairs,
distances,
)
[docs]
def objective(self, points: _Array) -> tuple[float, _Array]:
gradient = np.zeros_like(points)
numeric = _table()["numerics"]
energy = 0.0
def distances(
pairs: _Indices, targets: _Array, stiffness: float, repel: bool = False
) -> None:
nonlocal energy
a, b = pairs.T
delta = points[a] - points[b]
norms = np.linalg.norm(delta, axis=1)
error = norms - targets
if repel:
error = np.minimum(error, 0)
energy += 0.5 * stiffness * float(error @ error)
force = (stiffness * error / np.maximum(norms, 1e-12))[:, None] * delta
_scatter(gradient, a, force)
_scatter(gradient, b, -force)
distances(self.bonds, self.lengths, numeric["bond_stiffness"])
# Repel to just beyond the clearance, so the soft penalty's equilibrium
# leaves no pair inside it for the overlap gate to count.
margin = numeric["repulsion_margin"]
distances(self.pairs, self.clearance + margin, numeric["repulsion_stiffness"], True)
distances(self.ring_pairs, self.ring_distances, numeric["angle_stiffness"])
a, c, b = self.triples.T
p, q = points[a] - points[c], points[b] - points[c]
error = _wrap(_angles(points, self.triples) - self.angles)
stiffness = numeric["angle_stiffness"]
energy += 0.5 * stiffness * float(error @ error)
gp = (stiffness * error / np.maximum((p * p).sum(axis=1), 1e-12))[
:, None
] * np.column_stack((p[:, 1], -p[:, 0]))
gq = (stiffness * error / np.maximum((q * q).sum(axis=1), 1e-12))[
:, None
] * np.column_stack((-q[:, 1], q[:, 0]))
_scatter(gradient, a, gp)
_scatter(gradient, b, gq)
_scatter(gradient, c, -gp - gq)
# Flat-bottom positional restraint: standard local geometry is free
# to settle, while large rearrangements are discouraged. Alignment
# and the substituent-side acceptance check finish preservation.
delta = points - self.seed
norms = np.linalg.norm(delta, axis=1)
error = np.maximum(norms - numeric["position_radius"], 0)
stiffness = numeric["position_stiffness"]
energy += 0.5 * stiffness * float(error @ error)
gradient += (stiffness * error / np.maximum(norms, 1e-12))[:, None] * delta
return energy, gradient
[docs]
def metrics(self, points: _Array) -> tuple[float, float, int, float]:
"""Bond spread, RMS angle error (radians), clashes and clash penalty."""
lengths = np.linalg.norm(points[self.bonds[:, 0]] - points[self.bonds[:, 1]], axis=1)
errors = _wrap(_angles(points, self.triples) - self.angles)
distances = np.linalg.norm(points[self.pairs[:, 0]] - points[self.pairs[:, 1]], axis=1)
overlap = np.maximum(self.clearance - distances, 0)
return (
float(np.ptp(lengths)),
float(np.sqrt(np.mean(errors**2))) if len(errors) else 0.0,
int(np.count_nonzero(overlap > 1e-6)),
float(overlap @ overlap),
)
[docs]
def acceptable(self, candidate: _Array, baseline: _Array) -> bool:
"""No material regression against the legacy layout, including labels."""
if candidate.shape != baseline.shape or not np.isfinite(candidate).all():
return False
a, b = self.bonds.T
lengths = np.linalg.norm(candidate[a] - candidate[b], axis=1)
old_lengths = np.linalg.norm(baseline[a] - baseline[b], axis=1)
if (lengths < 1e-6).any() or np.max(np.abs(lengths - self.lengths)) > np.max(
np.abs(old_lengths - self.lengths)
) + 0.001:
return False
before, after = self.metrics(baseline), self.metrics(candidate)
# Numerical tolerances: 0.1% of one bond and 0.1 degree, below typical
# canvas display precision; never permit an additional clash.
return all(
a <= b + tolerance
for a, b, tolerance in zip(
after, before, (0.001, np.deg2rad(0.1), 0, 1e-6), strict=True
)
)
[docs]
def refine(self) -> _Array | None:
"""A converged FIRE layout, or refusal; the caller keeps its fallback."""
numeric = _table()["numerics"]
if len(self.seed) > numeric["max_atoms"]:
return None
result = minimize(
self.objective,
self.seed,
fmax=numeric["fmax"],
max_steps=numeric["max_steps"],
dt=0.05,
dt_max=0.15,
)
return result.x if result.converged else None