"""Which reactant atom becomes which product atom: an atom mapping for a reaction (#29).
A reaction path between two sketches needs both endpoints' atoms in the same
order, so every reactant atom must be paired with one product atom of the
same element. The pairing chosen is the one that changes the fewest bonds,
the principle of minimum chemical distance (I. Ugi et al., *Angew. Chem.
Int. Ed.* 18, 111-123 (1979); J. Jauffret, C. Tonnelier, J.-E. Dubois,
*Tetrahedron Comput. Methodol.* 3, 335-349 (1990)):
* the cost of a pairing of the heavy atoms is the number of bonds between
them it forms and breaks, plus two for every hydrogen that has to move to
another atom (one bond broken, one formed);
* every cheapest pairing is found by a depth-first branch and bound over
the heavy atoms, element by element, breadth first through each
molecule, so a partial pairing's cost prunes early; a search that grows
too large stops and asks for pinned atoms instead of guessing;
* pairings that cost the same and differ only by the molecules' own
symmetry describe the same reaction. When the cheapest pairings describe
different reactions (which oxygen leaves in an esterification; a [3,3]
or a [1,3] shift in a Claisen rearrangement, which break and form the
same number of bonds), the mapping is **ambiguous**: no writer guesses,
and the user pins atoms. Of the cheapest, the one changing the fewest
bond orders is shown first; that is a presentation order, not a
chemical argument.
Hydrogens are paired here by the heavy atoms they sit on; where several are
equivalent, the reaction-path endpoints pair them by position
(:mod:`vibesketch.chem.endpoints`). RXNMapper (MIT) maps from learned
attention, but needs torch, too heavy for a core dependency.
"""
from __future__ import annotations
import math
from collections.abc import Sequence
from dataclasses import dataclass, field
from typing import Any
__all__ = ["AtomMapping", "MappingFailed", "Side", "balance_problems", "map_reaction", "side"]
#: The most partial pairings the search visits before it asks for pinned atoms.
_MAX_NODES = 400_000
#: The most equally cheap pairings kept for the comparison by symmetry.
_MAX_SOLUTIONS = 2000
[docs]
class MappingFailed(ValueError):
"""The reaction's atoms cannot be paired; the message says why."""
[docs]
@dataclass(frozen=True)
class Side:
"""One side of a reaction: its species as one molecule, with explicit hydrogens.
Each species contributes its atoms in the order of its 3D structure
(:func:`vibesketch.chem.embed.embed_3d`): its SMILES atoms, then the
hydrogens RDKit adds.
"""
mol: Any
#: Per atom: the species id and the atom's index within that species.
owners: tuple[tuple[str, int], ...]
#: Per species, in order: its id, SMILES and first atom index here.
species: tuple[tuple[str, str, int], ...]
[docs]
def label(self, atom: int) -> str:
"""``m1:C3``: the species, the element and the 1-based atom number within it."""
species, index = self.owners[atom]
return f"{species}:{self.mol.GetAtomWithIdx(atom).GetSymbol()}{index + 1}"
[docs]
def side(species: Sequence[tuple[str, str]]) -> Side:
"""The side made of ``(id, SMILES)`` species, each once, in order."""
from rdkit import Chem
from vibesketch.chem.species import mol_from_smiles
combined: Any = None
owners: list[tuple[str, int]] = []
listed: list[tuple[str, str, int]] = []
for identifier, smiles in species:
mol = Chem.AddHs(mol_from_smiles(smiles))
listed.append((identifier, smiles, len(owners)))
owners.extend((identifier, i) for i in range(mol.GetNumAtoms()))
combined = mol if combined is None else Chem.CombineMols(combined, mol)
if combined is None:
raise MappingFailed("a side of the reaction has no species")
Chem.SanitizeMol(combined) # CombineMols leaves ring information unset
return Side(combined, tuple(owners), tuple(listed))
[docs]
@dataclass(frozen=True)
class AtomMapping:
"""A pairing of every reactant atom with a product atom."""
reactants: Side
products: Side
#: ``pairs[i]`` is the product atom reactant atom ``i`` becomes.
pairs: tuple[int, ...]
#: Bonds formed and broken, and hydrogens moved, as the cost counts them.
cost: int
#: Other cheapest pairings that describe a different reaction, as ``pairs``.
alternatives: tuple[tuple[int, ...], ...] = ()
#: Atoms the user pinned, as (reactant, product) indices.
pinned: tuple[tuple[int, int], ...] = field(default_factory=tuple)
[docs]
@property
def ambiguous(self) -> bool:
return bool(self.alternatives)
[docs]
def changes(self, pairs: Sequence[int] | None = None) -> list[str]:
"""What the pairing does, one line per bond formed, broken or changed."""
return _changes(self.reactants, self.products, pairs or self.pairs)
[docs]
def centre(self, pairs: Sequence[int] | None = None) -> set[int]:
"""The reaction centre: heavy reactant atoms with a bond formed, broken or changed.
A hydrogen that moves puts the heavy atoms it leaves and joins in the centre.
"""
pairs = pairs or self.pairs
r_mol, p_mol = self.reactants.mol, self.products.mol
back = {b: a for a, b in enumerate(pairs)}
r_bonds = {
frozenset((b.GetBeginAtomIdx(), b.GetEndAtomIdx())): b.GetBondTypeAsDouble()
for b in r_mol.GetBonds()
}
p_bonds = {
frozenset((back[b.GetBeginAtomIdx()], back[b.GetEndAtomIdx()])): b.GetBondTypeAsDouble()
for b in p_mol.GetBonds()
}
changed = {bond for bond in r_bonds if p_bonds.get(bond) != r_bonds[bond]}
changed |= {bond for bond in p_bonds if bond not in r_bonds}
found: set[int] = set()
for bond in changed:
for atom in bond:
if r_mol.GetAtomWithIdx(atom).GetAtomicNum() > 1:
found.add(atom)
return found
# ── Costs ───────────────────────────────────────────────────────────────
def _heavy(mol: Any) -> list[int]:
return [atom.GetIdx() for atom in mol.GetAtoms() if atom.GetAtomicNum() > 1]
def _hydrogens(mol: Any, atom: int) -> int:
return sum(1 for n in mol.GetAtomWithIdx(atom).GetNeighbors() if n.GetAtomicNum() == 1)
def _bonds(mol: Any, atoms: set[int]) -> dict[frozenset[int], float]:
return {
frozenset((b.GetBeginAtomIdx(), b.GetEndAtomIdx())): b.GetBondTypeAsDouble()
for b in mol.GetBonds()
if b.GetBeginAtomIdx() in atoms and b.GetEndAtomIdx() in atoms
}
@dataclass(frozen=True)
class _Graphs:
reactant_heavy: list[int]
product_heavy: list[int]
reactant_bonds: dict[frozenset[int], float]
product_bonds: dict[frozenset[int], float]
reactant_h: dict[int, int]
product_h: dict[int, int]
def _graphs(reactants: Side, products: Side) -> _Graphs:
r_heavy, p_heavy = _heavy(reactants.mol), _heavy(products.mol)
return _Graphs(
r_heavy,
p_heavy,
_bonds(reactants.mol, set(r_heavy)),
_bonds(products.mol, set(p_heavy)),
{a: _hydrogens(reactants.mol, a) for a in r_heavy},
{b: _hydrogens(products.mol, b) for b in p_heavy},
)
def _cost(graphs: _Graphs, heavy: dict[int, int]) -> tuple[int, int]:
"""(bonds formed and broken plus two per moved hydrogen, bond orders changed)."""
changed = 0
orders = 0
mapped = set()
for bond, order in graphs.reactant_bonds.items():
a, b = tuple(bond)
image = frozenset((heavy[a], heavy[b]))
mapped.add(image)
product_order = graphs.product_bonds.get(image)
if product_order is None:
changed += 1
elif product_order != order:
orders += 1
changed += sum(1 for bond in graphs.product_bonds if bond not in mapped)
moved = sum(max(0, graphs.reactant_h[a] - graphs.product_h[b]) for a, b in heavy.items())
return changed + 2 * moved, orders
# ── Search ──────────────────────────────────────────────────────────────
class _TooLarge(Exception):
"""The exhaustive search passed its bound."""
def _adjacency(mol: Any, atoms: list[int]) -> dict[int, set[int]]:
keep = set(atoms)
found: dict[int, set[int]] = {a: set() for a in atoms}
for bond in mol.GetBonds():
a, b = bond.GetBeginAtomIdx(), bond.GetEndAtomIdx()
if a in keep and b in keep:
found[a].add(b)
found[b].add(a)
return found
def _order(atoms: list[int], adjacency: dict[int, set[int]], first: list[int]) -> list[int]:
"""The atoms breadth first, pinned ones and then the best-connected first.
Each atom then has assigned neighbours as early as possible, which is
what lets a partial pairing's cost prune the search.
"""
order: list[int] = []
seen: set[int] = set()
roots = [*first, *sorted(atoms, key=lambda a: (-len(adjacency[a]), a))]
for root in roots:
if root in seen:
continue
queue = [root]
seen.add(root)
while queue:
atom = queue.pop(0)
order.append(atom)
for neighbour in sorted(adjacency[atom]):
if neighbour not in seen:
seen.add(neighbour)
queue.append(neighbour)
return order
def _search(
reactants: Side, products: Side, graphs: _Graphs, pinned: dict[int, int]
) -> list[dict[int, int]]:
"""Every cheapest pairing of the heavy atoms, by depth-first branch and bound.
A partial pairing's cost (bonds between paired atoms formed or broken,
hydrogen counts that differ) only grows as atoms are added, so a branch
that already costs more than the best complete pairing is cut. All
pairings at the lowest cost are kept, up to a bound, because the
caller compares them for ambiguity.
"""
r_adj = _adjacency(reactants.mol, graphs.reactant_heavy)
p_adj = _adjacency(products.mol, graphs.product_heavy)
by_element: dict[int, list[int]] = {}
for b in graphs.product_heavy:
by_element.setdefault(products.mol.GetAtomWithIdx(b).GetAtomicNum(), []).append(b)
order = [a for a in _order(graphs.reactant_heavy, r_adj, list(pinned)) if a not in pinned]
image: dict[int, int] = {}
source: dict[int, int] = {}
def increment(a: int, b: int) -> int:
cost = abs(graphs.reactant_h[a] - graphs.product_h[b])
cost += sum(1 for n in r_adj[a] if n in image and image[n] not in p_adj[b])
cost += sum(1 for m in p_adj[b] if m in source and source[m] not in r_adj[a])
return cost
start = 0
for a, b in pinned.items():
start += increment(a, b)
image[a], source[b] = b, a
best = [math.inf]
found: list[dict[int, int]] = []
nodes = [0]
def step(position: int, cost: int) -> None:
nodes[0] += 1
if nodes[0] > _MAX_NODES:
raise _TooLarge
if position == len(order):
if cost < best[0]:
best[0] = cost
found.clear()
if len(found) < _MAX_SOLUTIONS:
found.append(dict(image))
return
a = order[position]
options = sorted(
(increment(a, b), b)
for b in by_element.get(reactants.mol.GetAtomWithIdx(a).GetAtomicNum(), [])
if b not in source
)
for extra, b in options:
if cost + extra > best[0]:
break
image[a], source[b] = b, a
step(position + 1, cost + extra)
del image[a], source[b]
step(0, start)
return found
# ── Hydrogens, symmetry and description ─────────────────────────────────
def _hydrogens_on(mol: Any, atom: int) -> list[int]:
return sorted(
n.GetIdx() for n in mol.GetAtomWithIdx(atom).GetNeighbors() if n.GetAtomicNum() == 1
)
def _with_hydrogens(reactants: Side, products: Side, heavy: dict[int, int]) -> tuple[int, ...]:
"""Every atom's partner: hydrogens follow their heavy atoms, the rest in order."""
r_mol, p_mol = reactants.mol, products.mol
pairs = dict(heavy)
left_r: list[int] = []
taken: set[int] = set()
for a, b in heavy.items():
on_a = _hydrogens_on(r_mol, a)
on_b = _hydrogens_on(p_mol, b)
for h_r, h_p in zip(on_a, on_b, strict=False):
pairs[h_r] = h_p
taken.add(h_p)
left_r.extend(on_a[len(on_b) :])
left_r.extend(
a.GetIdx() for a in r_mol.GetAtoms() if a.GetAtomicNum() == 1 and not a.GetNeighbors()
)
left_p = [
a.GetIdx() for a in p_mol.GetAtoms() if a.GetAtomicNum() == 1 and a.GetIdx() not in taken
]
for h_r, h_p in zip(sorted(left_r), sorted(left_p), strict=True):
pairs[h_r] = h_p
return tuple(pairs[i] for i in range(r_mol.GetNumAtoms()))
def _signature(reactants: Side, products: Side, heavy: dict[int, int]) -> tuple[Any, ...]:
"""What a pairing does, up to the molecules' symmetry: equal for equivalent pairings."""
from rdkit import Chem
rank_r = list(Chem.CanonicalRankAtoms(reactants.mol, breakTies=False))
rank_p = list(Chem.CanonicalRankAtoms(products.mol, breakTies=False))
return tuple(sorted((rank_r[a], rank_p[b]) for a, b in heavy.items()))
def _changes(reactants: Side, products: Side, pairs: Sequence[int]) -> list[str]:
r_mol, p_mol = reactants.mol, products.mol
back = {b: a for a, b in enumerate(pairs)}
names = {1.0: "single", 1.5: "aromatic", 2.0: "double", 3.0: "triple"}
lines: list[str] = []
r_bonds = {
frozenset((b.GetBeginAtomIdx(), b.GetEndAtomIdx())): b.GetBondTypeAsDouble()
for b in r_mol.GetBonds()
}
p_bonds = {
frozenset((back[b.GetBeginAtomIdx()], back[b.GetEndAtomIdx()])): b.GetBondTypeAsDouble()
for b in p_mol.GetBonds()
}
def name(bond: frozenset[int]) -> str:
a, b = sorted(bond)
return f"{reactants.label(a)}–{reactants.label(b)}"
for bond in sorted(r_bonds, key=sorted):
if bond not in p_bonds:
lines.append(f"broken: {name(bond)}")
elif p_bonds[bond] != r_bonds[bond]:
before = names.get(r_bonds[bond], str(r_bonds[bond]))
after = names.get(p_bonds[bond], str(p_bonds[bond]))
lines.append(f"{before} to {after}: {name(bond)}")
for bond in sorted(p_bonds, key=sorted):
if bond not in r_bonds:
lines.append(f"formed: {name(bond)}")
return lines
# ── Entry point ─────────────────────────────────────────────────────────
[docs]
def balance_problems(reactants: Side, products: Side) -> list[str]:
"""Why the two sides cannot be paired atom for atom; empty if they can."""
from collections import Counter
def formula(mol: Any) -> Counter[str]:
return Counter(atom.GetSymbol() for atom in mol.GetAtoms())
r, p = formula(reactants.mol), formula(products.mol)
if r == p:
return []
more = sorted((r - p).items())
fewer = sorted((p - r).items())
parts = []
if more:
parts.append("the reactants have " + ", ".join(f"{n} more {e}" for e, n in more))
if fewer:
parts.append("the products have " + ", ".join(f"{n} more {e}" for e, n in fewer))
return [
"the reaction is not balanced: "
+ " and ".join(parts)
+ "; draw every molecule that takes part, each as often as it reacts"
]
[docs]
def map_reaction(
reactants: Side, products: Side, pinned: Sequence[tuple[int, int]] = ()
) -> AtomMapping:
"""The cheapest pairing of the reactants' atoms with the products'.
``pinned`` holds (reactant atom, product atom) pairs the pairing must
keep: heavy atoms of the same element. Raises :class:`MappingFailed`
when the sides are not balanced or a pin is impossible.
"""
problems = balance_problems(reactants, products)
if problems:
raise MappingFailed(problems[0])
graphs = _graphs(reactants, products)
pins: dict[int, int] = {}
for a, b in pinned:
if a not in graphs.reactant_heavy or b not in graphs.product_heavy:
raise MappingFailed(
f"only heavy atoms can be pinned: {reactants.label(a)} to {products.label(b)}"
if 0 <= a < reactants.mol.GetNumAtoms() and 0 <= b < products.mol.GetNumAtoms()
else "a pinned atom is not on its side of the reaction"
)
if (
reactants.mol.GetAtomWithIdx(a).GetAtomicNum()
!= products.mol.GetAtomWithIdx(b).GetAtomicNum()
):
raise MappingFailed(
f"{reactants.label(a)} cannot become {products.label(b)}: another element"
)
if a in pins or b in pins.values():
raise MappingFailed(f"{reactants.label(a)} or {products.label(b)} is pinned twice")
pins[a] = b
try:
candidates = _search(reactants, products, graphs, pins)
except _TooLarge:
raise MappingFailed(
"the reaction is too large or too symmetric to map exhaustively; pin a few "
"atoms (reactant atom = product atom) and try again"
) from None
if not candidates:
raise MappingFailed("no pairing of the atoms keeps every element")
scored = sorted(
((_cost(graphs, c), index) for index, c in enumerate(candidates)), key=lambda x: x[0]
)
best = scored[0][0]
# Ambiguity is judged on bonds and hydrogens alone; order changes only rank.
cheapest = [candidates[index] for cost, index in scored if cost[0] == best[0]]
chosen = cheapest[0]
seen = {_signature(reactants, products, chosen)}
alternatives: list[tuple[int, ...]] = []
for pairing in cheapest[1:]:
signature = _signature(reactants, products, pairing)
if signature not in seen:
seen.add(signature)
alternatives.append(_with_hydrogens(reactants, products, pairing))
return AtomMapping(
reactants,
products,
_with_hydrogens(reactants, products, chosen),
best[0],
tuple(alternatives),
tuple(sorted(pins.items())),
)