Source code for vibesketch.document.embed

"""3D structures for a document's species, cached in the document.

Each species gets its conformers from :func:`vibesketch.chem.embed.embed_3d`
and keeps them while its structure, charge and multiplicity stay the same, so
a redraw elsewhere on the canvas does not throw them away (see
:mod:`vibesketch.document.derive`). The input writers read the first
conformer.
"""

from __future__ import annotations

from collections.abc import Sequence
from dataclasses import dataclass, field

from vibesketch.document.model import Document

__all__ = ["EmbedReport", "UnknownSpecies", "embed"]




class UnknownSpecies(KeyError):
    """The document has no species with that id."""

    def __str__(self) -> str:
        return f"the document has no species {self.args[0]!r}"





@dataclass
class EmbedReport:
    """What embedding did, per species id."""

    embedded: list[str] = field(default_factory=list)
    #: Species that had conformers already and were left alone.
    kept: list[str] = field(default_factory=list)
    #: Species that could not be embedded, with the reason.
    failed: dict[str, str] = field(default_factory=dict)
    warnings: list[str] = field(default_factory=list)





def embed(
    document: Document,
    species: Sequence[str] | None = None,
    *,
    count: int = 1,
    replace: bool = False,
) -> tuple[Document, EmbedReport]:
    """``document`` with 3D conformers for ``species`` (default: all of them).

    A species that has conformers keeps them unless ``replace`` is set.
    Raises :class:`UnknownSpecies` for an id the document does not have.
    """
    from vibesketch.chem.embed import EmbedFailed, embed_3d

    known = {s.id for s in document.species}
    wanted = list(species) if species is not None else [s.id for s in document.species]
    for identifier in wanted:
        if identifier not in known:
            raise UnknownSpecies(identifier)
    report = EmbedReport()
    updated = []
    for entry in document.species:
        if entry.id not in wanted:
            updated.append(entry)
            continue
        if entry.conformers and not replace:
            report.kept.append(entry.id)
            updated.append(entry)
            continue
        if entry.solid is not None:
            report.failed[entry.id] = (
                f"the solid {entry.solid.text} is given by its formula: its structure comes with "
                "crystal and slab models (roadmap Phase 5)"
            )
            updated.append(entry)
            continue
        if entry.nanostructure is not None:
            report.failed[entry.id] = (
                "a generated nanostructure keeps its generated atoms; generate it anew to "
                "change them"
            )
            updated.append(entry)
            continue
        if entry.smiles is None:
            report.failed[entry.id] = "it has no structure to embed"
            updated.append(entry)
            continue
        try:
            result = embed_3d(entry.smiles, count=count, attempts=max(10, 3 * count))
        except EmbedFailed as error:
            report.failed[entry.id] = str(error)
            updated.append(entry)
            continue
        report.embedded.append(entry.id)
        report.warnings.extend(f"{entry.id}: {warning}" for warning in result.warnings)
        updated.append(entry.model_copy(update={"conformers": result.conformers}))
    rebuilt = document.model_copy(update={"species": updated})
    return Document.model_validate(rebuilt.to_json_data()), report