"""Finite, exhaustive R-group enumeration, preserving the generic drawing.
Occurrence and RestH follow MDL CTfile Formats (December 1999), pp. 2-12,
10-26: https://mdli.com/downloads/literature/ctfile.pdf. Unfilled sites with
RestH off are unrestricted chemistry and cannot be exhaustively enumerated.
"""
from __future__ import annotations
import itertools
import re
from collections import Counter
from collections.abc import Iterator
from typing import Any
from vibesketch.document.model import Canvas, Document
MAX_COMBINATIONS = 10_000
MAX_PRODUCTS = 1_000
[docs]
class EnumerationRefused(ValueError):
"""The complete enumeration cannot be represented or fits no bounded operation."""
def _count_allowed(text: str | None, count: int) -> bool:
for term in (text or ">0").replace(" ", "").split(","):
if re.fullmatch(r"\d+", term):
if count == int(term):
return True
elif re.fullmatch(r"\d+-\d+", term):
low, high = map(int, term.split("-"))
if low > high:
raise EnumerationRefused(f"invalid occurrence range {term!r}")
if low <= count <= high:
return True
elif re.fullmatch(r"[<>]\d+", term):
if (count > int(term[1:])) if term[0] == ">" else (count < int(term[1:])):
return True
else:
raise EnumerationRefused(f"invalid occurrence condition {term!r}")
return False
def _parsed(text: str) -> Any:
from rdkit import Chem
mol = Chem.MolFromSmiles(text)
if mol is None:
raise EnumerationRefused("cannot read the generic structure")
if Chem.GetMolSubstanceGroups(mol):
raise EnumerationRefused("enumeration does not invent the length of bracketed polymers")
return mol
def _sites(mol: Any) -> list[tuple[int, int]]:
sites = []
for atom in mol.GetAtoms():
if atom.GetAtomicNum() != 0:
continue
label = atom.GetProp("atomLabel") if atom.HasProp("atomLabel") else ""
if not re.fullmatch(r"_R[1-9]\d*", label):
raise EnumerationRefused("every core attachment must name exactly one R-group")
if atom.GetDegree() not in (1, 2):
raise EnumerationRefused("R-group sites require one or two explicit connecting bonds")
if any(neighbor.GetAtomicNum() == 0 for neighbor in atom.GetNeighbors()):
raise EnumerationRefused("adjacent R-group sites have no fixed attachment connectivity")
sites.append((atom.GetIdx(), int(label[2:])))
return sites
def _attachments(mol: Any) -> list[int]:
from rdkit import Chem
points = []
labels = []
for atom in mol.GetAtoms():
if atom.GetAtomicNum() != 0:
continue
label = atom.GetProp("atomLabel") if atom.HasProp("atomLabel") else ""
if label not in ("_AP1", "_AP2") or atom.GetDegree() != 1:
raise EnumerationRefused("substituents need explicit terminal _AP1/_AP2 attachments")
points.append(atom.GetIdx())
labels.append(label)
if not points or len(set(labels)) != len(labels):
raise EnumerationRefused("substituents need unique numbered attachment points")
if len(Chem.GetMolFrags(mol)) != 1:
raise EnumerationRefused("a substituent must be connected")
return points
def _product(core: Any, sites: list[tuple[int, int]], choices: tuple[Any, ...]) -> Any:
"""Compatibility helper for one substituent per core site."""
return _assembled(
core, [((index,), choice) for (index, _), choice in zip(sites, choices, strict=True)]
)
def _assembled(core: Any, placements: list[tuple[tuple[int, ...], Any]]) -> Any:
from rdkit import Chem
result = Chem.RWMol(core)
remove = [index for indices, choice in placements if choice is not None for index in indices]
stereo = {
a.GetIdx(): [n.GetIdx() for n in a.GetNeighbors()]
for a in core.GetAtoms()
if a.GetChiralTag() != Chem.ChiralType.CHI_UNSPECIFIED
}
replacements: dict[tuple[int, int], int] = {}
for indices, choice in placements:
bonds = [
(index, bond) for index in indices for bond in core.GetAtomWithIdx(index).GetBonds()
]
if choice is None:
if len(bonds) != 1 or bonds[0][1].GetBondType() != Chem.BondType.SINGLE:
raise EnumerationRefused("hydrogen cannot fill a multivalent R-group site")
hydrogen = result.GetAtomWithIdx(indices[0])
hydrogen.SetAtomicNum(1)
hydrogen.SetNoImplicit(True)
continue
sub, points = choice
offset = result.GetNumAtoms()
result.InsertMol(sub)
stereo.update(
{
offset + a.GetIdx(): [offset + n.GetIdx() for n in a.GetNeighbors()]
for a in sub.GetAtoms()
if a.GetChiralTag() != Chem.ChiralType.CHI_UNSPECIFIED
}
)
remove.extend(offset + point for point in points)
for (index, bond), point in zip(bonds, points, strict=True):
sub_atom = sub.GetAtomWithIdx(point)
sub_bond = sub_atom.GetBonds()[0]
if sub_bond.GetBondType() != bond.GetBondType():
raise EnumerationRefused("core and substituent attachment bond orders differ")
left = bond.GetOtherAtomIdx(index)
right = offset + sub_atom.GetNeighbors()[0].GetIdx()
result.AddBond(left, right, bond.GetBondType())
replacements[(left, index)] = right
replacements[(right, offset + point)] = left
if left in stereo:
stereo[left] = [right if n == index else n for n in stereo[left]]
if right in stereo:
stereo[right] = [left if n == offset + point else n for n in stereo[right]]
removed = set(remove)
for index, expected in stereo.items():
actual = [
n.GetIdx()
for n in result.GetAtomWithIdx(index).GetNeighbors()
if n.GetIdx() not in removed
]
order = [actual.index(n) for n in expected]
inversions = sum(a > b for i, a in enumerate(order) for b in order[i + 1 :])
if inversions % 2:
result.GetAtomWithIdx(index).InvertChirality()
stereo_bonds = []
for bond in result.GetBonds():
reference = list(bond.GetStereoAtoms())
if len(reference) == 2:
left, right = bond.GetBeginAtomIdx(), bond.GetEndAtomIdx()
stereo_bonds.append(
(
left,
right,
bond.GetStereo(),
replacements.get((left, reference[0]), reference[0]),
replacements.get((right, reference[1]), reference[1]),
)
)
kept = {
old: new
for new, old in enumerate(i for i in range(result.GetNumAtoms()) if i not in removed)
}
for index in sorted(remove, reverse=True):
result.RemoveAtom(index)
mol = result.GetMol()
for left, right, kind, first, second in stereo_bonds:
bond = mol.GetBondBetweenAtoms(kept[left], kept[right])
bond.SetStereoAtoms(kept[first], kept[second])
bond.SetStereo(kind)
Chem.SetDoubleBondNeighborDirections(mol)
try:
Chem.SanitizeMol(mol)
mol = Chem.RemoveHs(mol)
except ValueError as exc:
raise EnumerationRefused(f"an attachment produces invalid valence: {exc}") from exc
if len(Chem.GetMolFrags(mol)) != 1:
raise EnumerationRefused("enumeration would produce disconnected species")
return mol
def _assignments(
core: Any, sites: list[tuple[int, int]], definitions: dict[int, Any]
) -> Iterator[list[tuple[tuple[int, ...], Any]]]:
"""Partition sites into complete substituent placements, including bridges."""
if len(sites) > 32:
raise EnumerationRefused("more than 32 attachment sites; nothing was added")
members = {}
for _, number in sites:
if number not in definitions:
raise EnumerationRefused(f"R{number} has no definition")
members[number] = [
(sub, _attachments(sub)) for sub in map(_parsed, definitions[number].substituents)
]
def fill(remaining: list[tuple[int, int]]) -> Iterator[list[tuple[tuple[int, ...], Any]]]:
if not remaining:
yield []
return
index, number = remaining[0]
degree = core.GetAtomWithIdx(index).GetDegree()
for rest in fill(remaining[1:]):
yield [((index,), None), *rest]
for sub, points in members[number]:
groups: list[tuple[int, ...]] = [(index,)] if len(points) == degree else []
if degree == 1 and len(points) == 2:
groups += [
(index, other)
for other, label in remaining[1:]
if label == number and core.GetAtomWithIdx(other).GetDegree() == 1
]
for indices in groups:
left = [site for site in remaining if site[0] not in indices]
for order in itertools.permutations(points):
for rest in fill(left):
yield [(indices, (sub, order)), *rest]
yield from fill(sites)
[docs]
def enumerate_rgroups(document: Document, core: str | None = None) -> Document:
"""Append ALL constraint-valid products, or refuse atomically; never truncate.
Existing species, cores, definitions, coordinates and jobs remain intact.
Products are ordered by chemical identity and deduplicated against existing
species. Repeating the operation is idempotent. A divalent site considers
both attachment orientations, deduplicated by resulting chemical identity.
"""
from vibesketch.chem.species import identity
from vibesketch.document.derive import derive
from vibesketch.document.formats import _check_layout, _merged
from vibesketch.document.ket import write_scene
generics = [
g for g in document.generics if g.kind == "rgroup-core" and (core is None or g.id == core)
]
if not generics:
raise EnumerationRefused(f"no R-group core {core or ''!r} to enumerate")
definitions = {r.number: r for r in document.rgroups}
products: dict[str, Any] = {}
budget = 0
for generic in generics:
mol = _parsed(generic.cxsmiles)
sites = _sites(mol)
if not sites:
raise EnumerationRefused("the core has no numbered R-group sites")
numbers = dict(sites)
for placements in _assignments(mol, sites, definitions):
budget += 1
if budget > MAX_COMBINATIONS:
raise EnumerationRefused(
f"more than {MAX_COMBINATIONS:,} candidate combinations; nothing was added"
)
counts = Counter(
numbers[index]
for indices, choice in placements
if choice is not None
for index in indices
)
if any(
not _count_allowed(definitions[n].occurrence, counts[n])
for n in set(numbers.values())
):
continue
if any(
counts[n] and definitions[n].if_then and not counts[definitions[n].if_then or 0]
for n in counts
):
continue
if any(
choice is None and not definitions[numbers[indices[0]]].rest_h
for indices, choice in placements
):
raise EnumerationRefused(
"an unfilled RestH-off site permits arbitrary fragments; "
"restrict occurrence to fill every site or enable RestH"
)
product = _assembled(mol, placements)
products.setdefault(identity(product).key, product)
if len(products) > MAX_PRODUCTS:
raise EnumerationRefused(f"more than {MAX_PRODUCTS:,} products; nothing was added")
if not products:
raise EnumerationRefused("no combination satisfies every occurrence and if-then constraint")
existing = {s.inchikey or s.smiles or s.id for s in document.species}
additions = [products[key] for key in sorted(products) if key not in existing]
if not additions:
return document
_check_layout(additions)
data = _merged(document.canvas.data, write_scene(additions))
changed = derive(
document.model_copy(update={"canvas": Canvas(format="ket", data=data)})
).document
return changed.model_copy(update={"reactions": document.reactions, "jobs": document.jobs})