"""Species from 3D structures other programs wrote: QVF sections, XYZ frames.
A vibe-qc result, a structure vibe-view saved or exported, or an XYZ file
carries atoms and Cartesian coordinates, sometimes with a charge, a
multiplicity, a SMILES, a compound number and a name, and sometimes the
steps of a reaction scheme. This module turns them into a document:
- each structure's bonds are perceived (:mod:`vibesketch.chem.perceive`),
with the stated charge, or the one its SMILES implies, else neutral;
- the molecules are laid out in 2D, the scheme's steps as reaction rows and
the rest in a row below, and the document is derived from that drawing;
- each species keeps its coordinates as a conformer named after where they
came from (``imported`` by default), with the stated multiplicity, compound
number and name. Frames of one molecule, a trajectory or an optimisation,
are kept as its conformers, the last first and the others named
``<method>:frame-<n>`` (#86).
What cannot be done is said in the warnings: a structure without bonds that
fit, a SMILES the coordinates contradict, a spin the electrons cannot have.
"""
from __future__ import annotations
import contextlib
from dataclasses import dataclass, field
from typing import Any
from vibesketch.document.derive import Derivation, derive
from vibesketch.document.model import Canvas, Conformer, Document, Name, Species
__all__ = ["NothingDrawn", "Step", "Structure", "from_structures"]
[docs]
class NothingDrawn(ValueError):
"""No structure could be drawn."""
[docs]
@dataclass(frozen=True)
class Structure:
"""One structure as another program wrote it."""
#: How the warnings name it: a QVF section id, an XYZ frame.
id: str
symbols: list[str]
positions: list[list[float]]
charge: int | None = None
multiplicity: int | None = None
label: str | None = None
name: str | None = None
smiles: str | None = None
#: The conformer's method: where the coordinates came from.
method: str = "imported"
[docs]
@dataclass(frozen=True)
class Step:
"""One reaction step between structures, by chemical role (QVF § 5.14)."""
id: str
reactants: list[str]
products: list[str]
reagents: list[str] = field(default_factory=list)
arrow: str = "forward"
conditions: list[str] = field(default_factory=list)
yield_text: str | None = None
def _drawable(
structure: Structure, warnings: list[str], bond_orders: bool = False
) -> tuple[Any, Conformer | None] | None:
"""The molecule a structure draws, and its conformer when its bonds are perceived.
When no bonds can be perceived from the coordinates, the SMILES draws
the species without them: a line notation implies no atom order, so the
coordinates cannot be kept.
"""
from rdkit import Chem
from vibesketch.chem.perceive import perceive
from vibesketch.chem.species import identity, mol_from_smiles
named = None
if structure.smiles is not None:
with contextlib.suppress(ValueError):
named = mol_from_smiles(structure.smiles)
charge = structure.charge
if charge is None:
charge = Chem.GetFormalCharge(named) if named is not None else 0
found = perceive(structure.symbols, structure.positions, charge=charge, template=named)
if found is None:
if named is None:
warnings.append(f"{structure.id}: no bonds could be perceived from its coordinates")
return None
warnings.append(
f"{structure.id}: no bonds could be perceived from its coordinates; drawn from "
"its SMILES, without its 3D structure"
)
return named, None
mol = found.mol
crowded = max((a.GetNumRadicalElectrons() for a in mol.GetAtoms()), default=0)
if crowded > 2: # KET and molfiles stop at a triplet
warnings.append(
f"{structure.id}: an atom with {crowded} unpaired electrons cannot be drawn"
)
return None
if found.how == "connectivity" and named is None and bond_orders:
# The file states the charge, so a graph without bond orders is a guess.
warnings.append(
f"{structure.id}: no bond orders fit its coordinates with charge {charge}, "
"so it is left out rather than drawn with guessed bonds"
)
return None
if found.how == "connectivity" and named is None:
warnings.append(
f"{structure.id}: no bond orders fit its coordinates, so it is drawn with single "
f"bonds as {identity(mol).smiles}; check the drawing"
)
if named is not None and identity(named).key != identity(mol).key:
warnings.append(
f"{structure.id}: its coordinates give {identity(mol).smiles}, not the "
f"{structure.smiles} the file names; vibe-sketch drew the coordinates"
)
conformer = Conformer(
method=structure.method, symbols=structure.symbols, xyz_angstrom=structure.positions
)
return mol, conformer
def _ordered(frames: list[Conformer]) -> list[Conformer]:
"""A trajectory's frames as conformers: the last, its result, first; the rest named by frame."""
if len(frames) == 1:
return frames
earlier = [
frame.model_copy(update={"method": f"{frame.method}:frame-{number}"})
for number, frame in enumerate(frames[:-1], start=1)
]
return [frames[-1], *earlier]
def _pieces(mol: Any, conformer: Conformer | None) -> list[tuple[str, Conformer | None]]:
"""Each molecule of a structure, by identity, with its own atoms' coordinates."""
from rdkit import Chem
from vibesketch.chem.species import identity
mapping: list[list[int]] = []
fragments = Chem.GetMolFrags(mol, asMols=True, fragsMolAtomMapping=mapping)
if len(fragments) == 1:
return [(identity(mol).key, conformer)]
pieces: list[tuple[str, Conformer | None]] = []
for fragment, atoms in zip(fragments, mapping, strict=True):
piece = None
if conformer is not None:
piece = Conformer(
method=conformer.method,
symbols=[conformer.symbols[i] for i in atoms],
xyz_angstrom=[conformer.xyz_angstrom[i] for i in atoms],
)
pieces.append((identity(fragment).key, piece))
return pieces
def _fitted(entry: Species, update: dict[str, Any], warnings: list[str]) -> Species:
"""The derived species with what the file says, its spin fitted when that is impossible."""
from pydantic import ValidationError
data = entry.model_dump(by_alias=True) | update
try:
return Species.model_validate(data)
except ValidationError as error:
reason = error.errors()[0]["msg"].removeprefix("Value error, ")
conformers: list[Conformer] = update.get("conformers") or []
electrons = conformers[0].electron_count(entry.charge) if conformers else None
if electrons is None or electrons < 0:
data.pop("conformers", None)
data.pop("multiplicity", None)
warnings.append(f"{entry.id}: {reason}; its 3D structure is left out")
else:
data["multiplicity"] = electrons % 2 + 1
warnings.append(f"{entry.id}: {reason}; multiplicity {data['multiplicity']} instead")
return Species.model_validate(data)
[docs]
def from_structures(
structures: list[Structure],
*,
steps: list[Step] | None = None,
warnings: list[str] | None = None,
source: str = "imported",
bond_orders: bool = False,
) -> Derivation:
"""A document drawn from ``structures``, the ``steps`` as reaction rows.
``warnings`` are the caller's so far, kept first. ``source`` is what a
stated name records as its origin. With ``bond_orders``, for files that
state the charge, a structure whose bond orders do not fit is left out
instead of drawn with single bonds. Raises :class:`NothingDrawn` when no
structure can be drawn.
"""
from rdkit import Chem
from vibesketch.chem.layout import compute_2d
from vibesketch.chem.species import identity
from vibesketch.document.io import new_document
from vibesketch.document.ket import ARROW_MODES, SceneStep, write_scene
said = list(warnings or [])
in_steps = {
sid for step in steps or [] for sid in (*step.reactants, *step.products, *step.reagents)
}
drawn: dict[str, Any] = {}
keys: dict[str, str] = {}
facts: dict[str, dict[str, Any]] = {}
for structure in structures:
found = _drawable(structure, said, bond_orders)
if found is None:
continue
mol, conformer = found
key = identity(mol).key
# A molecule already drawn, a trajectory's next frame, is drawn once.
repeat = key in keys.values() and structure.id not in in_steps
keys[structure.id] = key
if not repeat:
flat = Chem.RemoveHs(mol)
compute_2d(flat)
drawn[structure.id] = flat
pieces = _pieces(mol, conformer)
if len(pieces) > 1:
said.append(
f"{structure.id} holds {len(pieces)} molecules: each keeps its own 3D "
"structure, not where it sits beside the others"
)
for key, piece in pieces:
known = facts.setdefault(key, {})
if piece is not None:
known.setdefault("frames", []).append(piece)
if len(pieces) > 1:
continue # a complex's spin, number and name are not one molecule's
if structure.multiplicity is not None:
known.setdefault("multiplicity", structure.multiplicity)
if structure.label is not None:
known.setdefault("label", structure.label)
if structure.name is not None:
known.setdefault(
"name", Name(value=structure.name, source=source, confidence="high")
)
if not drawn:
raise NothingDrawn("no structure could be drawn: " + "; ".join(said))
for known in facts.values():
frames = known.pop("frames", None)
if frames:
known["conformers"] = _ordered(frames)
rows: list[SceneStep] = []
printed: list[tuple[list[str], list[str], str]] = []
used: set[str] = set()
for step in steps or []:
missing = [r for r in step.reactants + step.products if r not in drawn]
if missing:
said.append(
f"step {step.id}: {', '.join(missing)} could not be drawn, so the step is left out"
)
continue
arrow = step.arrow if step.arrow in ARROW_MODES else "forward"
# vibe-sketch draws by place at the arrow: a retrosynthesis's target is at the tail.
tail, head = (
(step.products, step.reactants)
if arrow == "retrosynthetic"
else (step.reactants, step.products)
)
reagents = [r for r in step.reagents if r in drawn]
rows.append(
SceneStep(
tail=[drawn[r] for r in tail],
head=[drawn[r] for r in head],
reagents=[drawn[r] for r in reagents],
arrow=arrow,
conditions=list(step.conditions),
)
)
if step.yield_text is not None:
printed.append((tail, head, step.yield_text))
used |= {*tail, *head, *reagents}
loose = [mol for sid, mol in drawn.items() if sid not in used]
canvas = Canvas(format="ket", data=write_scene(loose, steps=rows or None))
derived = derive(new_document().model_copy(update={"canvas": canvas}))
if not derived.document.species:
raise NothingDrawn("no structure could be drawn: " + "; ".join(derived.warnings))
species = [
_fitted(entry, facts.get(entry.inchikey or entry.smiles or entry.id, {}), said)
for entry in derived.document.species
]
data = derived.document.model_copy(update={"species": species}).to_json_data()
ids = {entry.inchikey or entry.smiles or entry.id: entry.id for entry in species}
for tail, head, text in printed: # derive orders the steps itself: match them
sides = ({ids.get(keys[r]) for r in tail}, {ids.get(keys[r]) for r in head})
for reaction in data["reactions"]:
if (set(reaction["reactants"]), set(reaction["products"])) == sides:
reaction["yield"] = text
return Derivation(Document.model_validate(data), [*said, *derived.warnings])