"""chemfig source for LaTeX: structures as a walk with angles, schemes with arrows.
chemfig (Christian Tellechea, CTAN) draws a molecule from a walk over its
atoms: every bond carries an absolute angle and a length coefficient,
branches are parenthesised, and a bond that closes a ring is a hook,
``?[a]`` at the first atom and ``?[a,{=}]`` at the second. This module
reproduces the drawing rather than idealising it: a spanning tree by
depth-first search from the leftmost terminal atom, the bond angles and
relative lengths read off the 2D coordinates, and the remaining bonds as
hooks. Fused and bridged rings therefore come out exactly as drawn, which is
why chemfig's ring macros (``*6(...)``) are not used.
Atom labels follow chemfig's rule that a new atom starts at every capital
letter: carbon is implicit, a heteroatom carries its hydrogens on the side
away from its bonds (``HO-`` and ``-OH``) and its charge as a superscript,
and every bond names the atom it leaves and arrives at (``-[:30,,2,1]``) so a
label such as ``H_2N`` attaches at the nitrogen. A heteroatom with bonds on
both sides draws its hydrogens as atoms of their own in the free direction.
Stereo bonds are chemfig's Cram bonds, ``>`` and ``>:`` widening towards the
atom in front, with the direction taken from the walk. Aromatic rings are
written in a Kekulé form.
Reactions become ``\\schemestart ... \\schemestop`` with ``\\+`` between
species, reagents above and conditions and the yield below the arrow, and
compound numbers under their structures with ``\\chemname``. The arrow kinds
map to ``->``, ``<=>`` and ``<->``; a retrosynthetic arrow is drawn with
TikZ's double ``-Implies`` tip, which chemfig loads. The result is a
fragment for ``\\input`` in a document that loads the ``chemfig`` package.
"""
from __future__ import annotations
import math
from collections.abc import Iterable
from dataclasses import dataclass, field
from typing import Any
from vibesketch.document.model import Document, Reaction
__all__ = ["ChemfigRefused", "chemfig_molecule", "chemfig_source"]
[docs]
class ChemfigRefused(ValueError):
"""The structure or scheme cannot be expressed in chemfig."""
#: chemfig arrow codes per reaction arrow; the retrosynthetic one is TikZ's double arrow.
_ARROWS = {
"forward": "->",
"equilibrium": "<=>",
"resonance": "<->",
"retrosynthetic": "->[][][][double,-Implies]",
}
_BOND_CHARS = {1: "-", 2: "=", 3: "~"}
_LATEX_SPECIAL = {
"\\": r"\textbackslash{}",
"{": r"\{",
"}": r"\}",
"%": r"\%",
"&": r"\&",
"#": r"\#",
"_": r"\_",
"$": r"\$",
"^": r"\^{}",
"~": r"\~{}",
}
#: A bond this far from the median length, as a fraction, gets a length coefficient.
_LENGTH_TOLERANCE = 0.03
def _escape(text: str) -> str:
"""``text`` safe in LaTeX text mode."""
return "".join(_LATEX_SPECIAL.get(char, char) for char in text)
@dataclass(frozen=True)
class _Label:
"""An atom's chemfig label and where its element sits among chemfig's atoms."""
text: str
#: The element's number among the label's chemfig atoms; 0 for an empty label.
element: int
#: How many chemfig atoms the label holds.
atoms: int
#: Hydrogens to draw as atoms of their own, with their directions in degrees.
explicit_hydrogens: tuple[float, ...] = ()
@property
def attaches(self) -> str:
"""The atom number a bond attaches at, or empty when chemfig can tell."""
return str(self.element) if self.atoms > 1 else ""
def _superscript(charge: int, radicals: int) -> str:
parts = []
if charge:
magnitude = "" if abs(charge) == 1 else str(abs(charge))
parts.append(f"{magnitude}{'+' if charge > 0 else '-'}")
if radicals:
parts.append(r"\bullet" * radicals)
return f"^{{{''.join(parts)}}}" if parts else ""
def _hydrogens(count: int) -> str:
if count <= 0:
return ""
return "H" if count == 1 else f"H_{{{count}}}"
def _free_directions(angles: Iterable[float], count: int) -> tuple[float, ...]:
"""``count`` directions, in degrees, inside the widest gap between ``angles``."""
taken = sorted(angle % 360.0 for angle in angles)
if not taken:
return tuple(90.0 + 180.0 * index for index in range(count))
gaps = [
((taken[(index + 1) % len(taken)] - angle) % 360.0 or 360.0, angle)
for index, angle in enumerate(taken)
]
width, start = max(gaps)
return tuple((start + width * (index + 1) / (count + 1)) % 360.0 for index in range(count))
def _label(atom: Any, angles: list[float]) -> _Label:
"""The chemfig label of ``atom``, whose bonds leave at ``angles`` (degrees)."""
symbol = str(atom.GetSymbol())
hydrogens = int(atom.GetTotalNumHs())
charge = int(atom.GetFormalCharge())
radicals = int(atom.GetNumRadicalElectrons())
isotope = int(atom.GetIsotope())
if symbol == "C" and not (charge or radicals or isotope) and angles:
return _Label("", 0, 0)
superscript = _superscript(charge, radicals)
mass = f"{{}}^{{{isotope}}}" if isotope else ""
rightward = any(math.cos(math.radians(a)) > 0.05 for a in angles)
leftward = any(math.cos(math.radians(a)) < -0.05 for a in angles)
if hydrogens and rightward and leftward:
# Bonds on both sides: the hydrogens are atoms of their own.
return _Label(
f"{mass}{symbol}{superscript}",
2 if mass else 1,
2 if mass else 1,
_free_directions(angles, hydrogens),
)
if hydrogens and (rightward or (not angles and symbol in ("O", "S", "F", "Cl", "Br", "I"))):
# Hydrogens on the left, away from the bonds: HO-, H_2N-.
return _Label(f"{_hydrogens(hydrogens)}{mass}{symbol}{superscript}", 2, 2)
text = f"{mass}{symbol}{_hydrogens(hydrogens)}{superscript}"
element = 2 if mass else 1
return _Label(text, element, element + (1 if hydrogens else 0))
def _angle(origin: tuple[float, float], target: tuple[float, float]) -> float:
return math.degrees(math.atan2(target[1] - origin[1], target[0] - origin[0])) % 360.0
def _number(value: float) -> str:
"""An angle in whole degrees: a tenth of a degree is invisible on paper."""
return str(int(round(value)) % 360)
@dataclass
class _Walk:
"""The spanning tree of a molecule and the hooks that close its rings."""
children: dict[int, list[int]] = field(default_factory=dict)
parent: dict[int, int | None] = field(default_factory=dict)
#: Hook name per ring-closing bond, keyed by the bond index.
hooks: dict[int, str] = field(default_factory=dict)
#: Hooks opened at an atom, in order; closed at an atom, as (name, bond index).
opened: dict[int, list[str]] = field(default_factory=dict)
closed: dict[int, list[tuple[str, int]]] = field(default_factory=dict)
def _hook_name(index: int) -> str:
letters = ""
index += 1
while index:
index, remainder = divmod(index - 1, 26)
letters = chr(ord("a") + remainder) + letters
return letters
def _walk(mol: Any, positions: list[tuple[float, float]], stereo: dict[int, Any]) -> _Walk:
"""Depth-first from the leftmost terminal atom; stereo bonds become tree edges first."""
atoms = range(mol.GetNumAtoms())
start = min(atoms, key=lambda i: (mol.GetAtomWithIdx(i).GetDegree() != 1, positions[i]))
walk = _Walk(parent={start: None})
order = 0
def neighbours(index: int) -> list[tuple[int, int]]:
pairs = []
for bond in mol.GetAtomWithIdx(index).GetBonds():
other = bond.GetOtherAtomIdx(index)
pairs.append((bond.GetIdx(), other))
return sorted(
pairs,
key=lambda pair: (pair[0] not in stereo, _angle(positions[index], positions[pair[1]])),
)
stack: list[tuple[int, list[tuple[int, int]]]] = [(start, neighbours(start))]
seen_bonds: set[int] = set()
while stack:
index, pending = stack[-1]
if not pending:
stack.pop()
continue
bond_index, other = pending.pop(0)
if bond_index in seen_bonds:
continue
seen_bonds.add(bond_index)
if other in walk.parent:
name = _hook_name(order)
order += 1
walk.hooks[bond_index] = name
walk.opened.setdefault(other, []).append(name)
walk.closed.setdefault(index, []).append((name, bond_index))
continue
walk.parent[other] = index
walk.children.setdefault(index, []).append(other)
stack.append((other, neighbours(other)))
return walk
def _bond_char(bond: Any, stereo: dict[int, Any], from_atom: int) -> str:
"""chemfig's character for ``bond`` walked from ``from_atom``."""
from rdkit import Chem
order = bond.GetBondType()
if order == Chem.BondType.AROMATIC:
raise ChemfigRefused("an aromatic bond survived kekulisation")
if bond.GetIdx() in stereo:
hashed = stereo[bond.GetIdx()] == Chem.BondDir.BEGINDASH
forward = bond.GetBeginAtomIdx() == from_atom
return (">" if forward else "<") + (":" if hashed else "")
try:
return _BOND_CHARS[int(order)]
except KeyError:
raise ChemfigRefused(f"chemfig cannot draw a {order} bond") from None
def _prepared(mol: Any) -> tuple[Any, list[tuple[float, float]], dict[int, Any], float]:
"""A Kekulé copy with wedges assigned, its 2D positions, stereo bonds and bond length."""
from rdkit import Chem
copy = Chem.Mol(mol)
if copy.GetNumConformers() == 0:
raise ChemfigRefused("the structure has no 2D coordinates")
try:
Chem.Kekulize(copy, clearAromaticFlags=True)
except Exception as error: # RDKit's KekulizeException
raise ChemfigRefused(f"the structure cannot be written in a Kekulé form: {error}") from None
Chem.WedgeMolBonds(copy, copy.GetConformer())
conformer = copy.GetConformer()
positions = [
(float(conformer.GetAtomPosition(i).x), float(conformer.GetAtomPosition(i).y))
for i in range(copy.GetNumAtoms())
]
stereo = {
bond.GetIdx(): bond.GetBondDir()
for bond in copy.GetBonds()
if bond.GetBondDir() in (Chem.BondDir.BEGINWEDGE, Chem.BondDir.BEGINDASH)
}
lengths = sorted(
math.dist(positions[b.GetBeginAtomIdx()], positions[b.GetEndAtomIdx()])
for b in copy.GetBonds()
)
unit = lengths[len(lengths) // 2] if lengths else 1.0
return copy, positions, stereo, unit if unit > 1e-9 else 1.0
[docs]
def chemfig_molecule(mol: Any) -> str:
"""``\\chemfig{...}`` for an RDKit molecule with 2D coordinates.
Hydrogens are implicit unless the molecule keeps them as atoms. Raises
:class:`ChemfigRefused` for a bond chemfig cannot draw or a structure
with no coordinates.
"""
if mol.GetNumAtoms() == 0:
raise ChemfigRefused("the structure has no atoms")
prepared, positions, stereo, unit = _prepared(mol)
angles = {
i: [
_angle(positions[i], positions[bond.GetOtherAtomIdx(i)])
for bond in prepared.GetAtomWithIdx(i).GetBonds()
]
for i in range(prepared.GetNumAtoms())
}
labels = {i: _label(prepared.GetAtomWithIdx(i), angles[i]) for i in angles}
if prepared.GetNumAtoms() == 1:
return f"\\chemfig{{{labels[0].text}}}"
walk = _walk(prepared, positions, stereo)
def bond_text(bond_index: int, from_atom: int, to_atom: int) -> str:
bond = prepared.GetBondWithIdx(bond_index)
char = _bond_char(bond, stereo, from_atom)
angle = _number(_angle(positions[from_atom], positions[to_atom]))
ratio = math.dist(positions[from_atom], positions[to_atom]) / unit
coefficient = f"{ratio:.2f}" if abs(ratio - 1.0) > _LENGTH_TOLERANCE else ""
fields = [f":{angle}", coefficient, labels[from_atom].attaches, labels[to_atom].attaches]
while fields and not fields[-1]:
fields.pop()
return f"{char}[{','.join(fields)}]"
def hook_text(name: str, bond_index: int, at_atom: int) -> str:
bond = prepared.GetBondWithIdx(bond_index)
char = _bond_char(bond, stereo, at_atom)
return f"?[{name}]" if char == "-" else f"?[{name},{{{char}}}]"
def emit(index: int) -> str:
label = labels[index]
parts = [label.text]
parts.extend(f"?[{name}]" for name in walk.opened.get(index, []))
parts.extend(hook_text(name, bond, index) for name, bond in walk.closed.get(index, []))
for direction in label.explicit_hydrogens:
attach = f",,{label.attaches}" if label.attaches else ""
parts.append(f"(-[:{_number(direction)}{attach}]H)")
children = walk.children.get(index, [])
for position, child in enumerate(children):
bond = prepared.GetBondBetweenAtoms(index, child).GetIdx()
branch = bond_text(bond, index, child) + emit(child)
parts.append(branch if position == len(children) - 1 else f"({branch})")
return "".join(parts)
start = next(atom for atom, parent in walk.parent.items() if parent is None)
return f"\\chemfig{{{emit(start)}}}"
def _named(figure: str, label: str | None) -> str:
return f"\\chemname{{{figure}}}{{\\textbf{{{_escape(label)}}}}}" if label else figure
def _arrow(reaction: Reaction, figures: dict[str, str]) -> str:
"""``\\arrow{...}`` with the reagents above and the conditions and yield below."""
above = ", ".join(figures[sid] for sid in reaction.reagents if sid in figures)
below_parts = [*reaction.conditions]
if reaction.yield_ is not None:
below_parts.append(reaction.yield_)
below = _escape(", ".join(part for part in below_parts if part))
code = _ARROWS[reaction.arrow]
if "[" in code: # the retrosynthetic arrow carries its TikZ options already
return f"\\arrow{{{code.replace('[][]', f'[{above}][{below}]', 1)}}}"
if below:
return f"\\arrow{{{code}[{above}][{below}]}}"
if above:
return f"\\arrow{{{code}[{above}]}}"
return f"\\arrow{{{code}}}"
[docs]
def chemfig_source(document: Document) -> str:
"""The document's structures and reactions as a chemfig fragment.
Every reaction is a scheme; species in no reaction follow as structures
of their own. Raises :class:`ChemfigRefused` when there is nothing to
draw or a species a reaction needs has no structure.
"""
from vibesketch.document.formats import ExportRefused, species_mols
try:
mols = species_mols(document)
except ExportRefused as error:
raise ChemfigRefused(str(error)) from None
if not mols:
raise ChemfigRefused("the document has no species to export")
figures = {sid: chemfig_molecule(mol) for sid, mol in mols.items()}
labels = {species.id: species.label for species in document.species}
lines = [
"% chemfig source written by vibe-sketch, for \\input in a document that",
"% loads \\usepackage{chemfig}. Angles and bond lengths follow the drawing.",
]
in_reactions: set[str] = set()
for reaction in document.reactions:
for sid in (*reaction.reactants, *reaction.reagents, *reaction.products):
if sid not in figures:
raise ChemfigRefused(
f"reaction {reaction.id} needs species {sid}, which has no structure"
)
in_reactions.update(reaction.reactants, reaction.reagents, reaction.products)
lines.append("")
lines.append(f"% reaction {reaction.id}")
lines.append("\\schemestart")
for index, sid in enumerate(reaction.reactants):
if index:
lines.append("\\+")
lines.append(_named(figures[sid], labels.get(sid)))
lines.append(_arrow(reaction, figures))
for index, sid in enumerate(reaction.products):
if index:
lines.append("\\+")
lines.append(_named(figures[sid], labels.get(sid)))
lines.append("\\schemestop")
loose = [species.id for species in document.species if species.id in figures]
loose = [sid for sid in loose if sid not in in_reactions]
if loose:
lines.append("")
lines.append("% structures")
lines.extend(_named(figures[sid], labels.get(sid)) for sid in loose)
return "\n".join(lines) + "\n"