Source code for vibesketch.document.relax

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

Each named species' first conformer goes through
:func:`vibesketch.chem.relax.relax_3d` and is replaced by the result; its
other conformers stay. Constraints name atoms of one species, so a call
with constraints names exactly one species.
"""

from __future__ import annotations

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

from vibesketch.chem.relax import DEFAULT_FMAX, DEFAULT_MAX_STEPS, ForceField, Relaxed
from vibesketch.document.embed import UnknownSpecies
from vibesketch.document.model import Constraint, Document

__all__ = ["RelaxReport", "relax"]




@dataclass
class RelaxReport:
    """What relaxing did, per species id."""

    relaxed: list[str] = field(default_factory=list)
    #: The relaxation of each relaxed species: energies, steps, constraint checks.
    results: dict[str, Relaxed] = field(default_factory=dict)
    #: Species that could not be relaxed, with the reason.
    failed: dict[str, str] = field(default_factory=dict)
    warnings: list[str] = field(default_factory=list)





def relax(
    document: Document,
    species: Sequence[str] | None = None,
    *,
    constraints: Sequence[Constraint] | None = None,
    force_field: ForceField = "mmff94",
    fmax: float | None = None,
    max_steps: int | None = None,
) -> tuple[Document, RelaxReport]:
    """``document`` with the 3D structures of ``species`` (default: all) relaxed.

    ``fmax`` and ``max_steps`` default to the relaxation's own. Raises
    :class:`UnknownSpecies` for an id the document does not have and
    :class:`ValueError` when constraints come with more than one species.
    """
    from vibesketch.chem.relax import RelaxFailed, relax_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)
    if constraints and len(wanted) != 1:
        raise ValueError("constraints name atoms of one species, so name exactly one species")
    report = RelaxReport()
    updated = []
    for entry in document.species:
        if entry.id not in wanted:
            updated.append(entry)
            continue
        try:
            result = relax_3d(
                entry,
                constraints=constraints,
                force_field=force_field,
                fmax=DEFAULT_FMAX if fmax is None else fmax,
                max_steps=DEFAULT_MAX_STEPS if max_steps is None else max_steps,
            )
        except RelaxFailed as error:
            report.failed[entry.id] = str(error)
            updated.append(entry)
            continue
        report.relaxed.append(entry.id)
        report.results[entry.id] = result
        report.warnings.extend(f"{entry.id}: {warning}" for warning in result.warnings)
        conformers = [result.conformer, *entry.conformers[1:]]
        updated.append(entry.model_copy(update={"conformers": conformers}))
    rebuilt = document.model_copy(update={"species": updated})
    return Document.model_validate(rebuilt.to_json_data()), report