vibe-sketcher/Documentation

Source code for vibesketcher.document.solid_edit

"""Editing periodic structures: cell parameters, selected sites and slab layers.

Every edit returns a new validated `PeriodicStructure` and appends a record to
``provenance["builder"]["edits"]``; by default the edited structure replaces
its source under the same ID. Cell parameters follow the IUCr convention
(a, b, c in angstrom, alpha = angle(b, c), beta = angle(a, c),
gamma = angle(a, b) in degrees). The cell keeps its orientation: the new rows
are the standard rows of the new parameters carried by the same rotation that
maps the old standard rows onto the old cell.
"""

from __future__ import annotations

import math
from collections import Counter
from typing import Annotated, Any, Literal

import numpy as np
from pydantic import BaseModel, ConfigDict, Field

from vibesketcher.document import crystal
from vibesketcher.document.model import Document, PeriodicStructure

_ID = Annotated[str, Field(pattern=r"^[A-Za-z][A-Za-z0-9_-]{0,63}$")]
Keep = Literal["fractional", "cartesian"]




def cell_parameters(structure: PeriodicStructure) -> dict[str, float]:
    """a, b, c (angstrom), alpha, beta, gamma (degrees) and the volume."""
    rows = crystal.lattice_rows(structure)
    a, b, c = (float(np.linalg.norm(r)) for r in rows)

    def angle(u: np.ndarray, v: np.ndarray) -> float:
        cos = float(u @ v / (np.linalg.norm(u) * np.linalg.norm(v)))
        return math.degrees(math.acos(max(-1.0, min(1.0, cos))))

    return {
        "a": a,
        "b": b,
        "c": c,
        "alpha": angle(rows[1], rows[2]),
        "beta": angle(rows[0], rows[2]),
        "gamma": angle(rows[0], rows[1]),
        "volume": abs(float(np.linalg.det(rows))),
    }



def _record(structure: PeriodicStructure, edit: dict[str, Any]) -> dict[str, Any]:
    builder = crystal._builder(structure)
    edits = list(builder.get("edits") or [])
    edits.append(crystal._json(edit))
    builder["edits"] = edits
    return {**structure.provenance, "builder": crystal._json(builder)}


def _rebuild(
    structure: PeriodicStructure,
    identifier: str,
    lattice: np.ndarray,
    cartesian: np.ndarray,
    edit: dict[str, Any],
    *,
    symbols: list[str] | None = None,
) -> PeriodicStructure:
    inverse = np.linalg.inv(lattice)
    fractional = cartesian @ inverse
    data = structure.model_dump()
    data["id"] = identifier
    data["lattice_angstrom"] = tuple(tuple(float(x) for x in row) for row in lattice)
    for i, site in enumerate(data["sites"]):
        site["fractional"] = tuple(crystal._clean(x) for x in fractional[i])
        if symbols is not None:
            site["symbol"] = symbols[i]
    data["provenance"] = _record(structure, edit)
    data["source_cif"] = None
    data["source_model_sha256"] = None
    edited = PeriodicStructure.model_validate(data)
    _check_contacts(edited)
    return edited


def _check_contacts(structure: PeriodicStructure, limit: float = 0.5) -> None:
    """Refuse an edit that puts two atoms (or an atom and its image) within ``limit``."""
    lattice = crystal.lattice_rows(structure)
    xyz = crystal.cartesian(structure)
    ranges = [(-1, 0, 1) if p else (0,) for p in structure.pbc]
    shifts = np.asarray([(i, j, k) for i in ranges[0] for j in ranges[1] for k in ranges[2]])
    shifts = shifts @ lattice
    if len(xyz) * len(shifts) * len(xyz) > 4e7:
        return  # too large for the all-pairs check; the builders bound such cells
    for n, s in enumerate(shifts):
        delta = xyz[:, None, :] - (xyz[None, :, :] + s)
        distance = np.linalg.norm(delta, axis=-1)
        if not np.any(s):
            np.fill_diagonal(distance, np.inf)
        i, j = np.unravel_index(int(np.argmin(distance)), distance.shape)
        if distance[i, j] < limit:
            raise ValueError(
                f"the edit puts {structure.sites[i].id} within {distance[i, j]:.2f} A of "
                f"{structure.sites[j].id}" + (" (periodic image)" if n and np.any(s) else "")
            )




def edit_cell(
    structure: PeriodicStructure,
    *,
    identifier: str | None = None,
    a: float | None = None,
    b: float | None = None,
    c: float | None = None,
    alpha: float | None = None,
    beta: float | None = None,
    gamma: float | None = None,
    scale: float | None = None,
    vacuum: float | None = None,
    keep: Keep = "fractional",
) -> PeriodicStructure:
    """Change lattice lengths and angles, scale the cell, or set a slab's vacuum.

    ``keep="fractional"`` moves the atoms with the cell (strain);
    ``"cartesian"`` leaves them in place. Along a nonperiodic direction atoms
    always keep their physical distance from the middle of that vector, so a
    layer is never stretched by a vacuum change. ``vacuum`` sets the gap
    between a 2-D slab's top and its image along the normal.
    """
    old = cell_parameters(structure)
    new = {
        k: (v if v is not None else old[k])
        for k, v in {"a": a, "b": b, "c": c, "alpha": alpha, "beta": beta, "gamma": gamma}.items()
    }
    for key in ("a", "b", "c"):
        if not (math.isfinite(new[key]) and new[key] > 0):
            raise ValueError(f"{key} must be a positive length")
    if scale is not None:
        if not (math.isfinite(scale) and scale > 0):
            raise ValueError("scale must be a positive factor")
        for key, axis in (("a", 0), ("b", 1), ("c", 2)):
            if structure.pbc[axis]:
                new[key] *= scale
    rows = crystal.lattice_rows(structure)
    xyz = crystal.cartesian(structure)
    if vacuum is not None:
        if structure.pbc != (True, True, False):
            raise ValueError("vacuum applies to 2-D slabs (PBC T, T, F)")
        normal = np.cross(rows[0], rows[1])
        normal /= np.linalg.norm(normal)
        heights = xyz @ normal
        thickness = float(np.ptp(heights))
        if not (math.isfinite(vacuum) and vacuum > 0):
            raise ValueError("vacuum must be a positive length")
        along = float(rows[2] @ normal)
        new["c"] = (thickness + vacuum) * new["c"] / along
    old_standard = crystal.cell_from_parameters(
        old["a"], old["b"], old["c"], old["alpha"], old["beta"], old["gamma"]
    )
    rotation = np.linalg.solve(old_standard, rows)  # rows = old_standard @ rotation
    lattice = (
        crystal.cell_from_parameters(
            new["a"], new["b"], new["c"], new["alpha"], new["beta"], new["gamma"]
        )


        @ rotation
    )
    if np.linalg.det(lattice) <= 0:
        raise ValueError("the new parameters do not describe a right-handed cell")
    fractional = xyz @ np.linalg.inv(rows)
    cartesian = fractional @ lattice if keep == "fractional" else xyz.copy()
    # Nonperiodic directions: keep the physical offset from the vector's middle.
    for axis in range(3):
        if structure.pbc[axis]:
            continue
        old_vec, new_vec = rows[axis], lattice[axis]
        old_len, new_len = np.linalg.norm(old_vec), np.linalg.norm(new_vec)
        offset = (fractional[:, axis] - 0.5) * old_len
        centre_new = cartesian - np.outer((cartesian @ np.linalg.inv(lattice))[:, axis], new_vec)
        cartesian = centre_new + np.outer(0.5 + offset / new_len, new_vec)
    edit = {
        "edit": "cell",
        "from": {k: round(v, 10) for k, v in old.items() if k != "volume"},
        "to": {k: round(v, 10) for k, v in new.items()},
        "keep": keep,
        **({"scale": scale} if scale is not None else {}),
        **({"vacuum": vacuum} if vacuum is not None else {}),
    }
    bravais = crystal._builder(structure).get("bravais")
    if isinstance(bravais, str) and any(
        abs(new[k] - old[k]) > 1e-9 for k in ("alpha", "beta", "gamma")
    ):
        edit["warning"] = f"angles changed: the cell may no longer be {bravais}"
    return _rebuild(structure, identifier or structure.id, lattice, cartesian, edit)



def select_sites(
    structure: PeriodicStructure,
    *,
    sites: list[str] | None = None,
    elements: list[str] | None = None,
    layers: list[int] | None = None,
    adsorbate: int | str | None = None,
    z_min: float | None = None,
    z_max: float | None = None,
) -> list[str]:
    """Site IDs matching every given criterion (all sites when none is given).

    ``layers`` uses the slab builder's layer numbers (0 is the bottom);
    ``adsorbate`` an index or label of a recorded adsorbate; ``z_min``/``z_max``
    Cartesian heights in angstrom.
    """
    builder = crystal._builder(structure)
    chosen = [s.id for s in structure.sites]
    known = set(chosen)
    if sites is not None:
        unknown = sorted(set(sites) - known)
        if unknown:
            raise ValueError(f"sites not in the structure: {', '.join(unknown)}")
        chosen = [s for s in chosen if s in set(sites)]
    by_id = {s.id: s for s in structure.sites}
    if elements:
        chosen = [s for s in chosen if by_id[s].symbol in set(elements)]
    if layers is not None:
        layer_of = builder.get("layer_of_site")
        if not isinstance(layer_of, dict):
            raise ValueError("layers apply to structures made by the slab builder")
        chosen = [s for s in chosen if layer_of.get(s) in set(layers)]
    if adsorbate is not None:
        records = builder.get("adsorbates")
        if not isinstance(records, list) or not records:
            raise ValueError("this structure records no adsorbate")
        if isinstance(adsorbate, int):
            if not 0 <= adsorbate < len(records):
                raise ValueError("adsorbate index out of range")
            record = records[adsorbate]
        else:
            matches = [r for r in records if isinstance(r, dict) and r.get("label") == adsorbate]
            if not matches:
                raise ValueError(f"no adsorbate labelled {adsorbate!r}")
            record = matches[0]
        members = set(record.get("sites", [])) if isinstance(record, dict) else set()
        chosen = [s for s in chosen if s in members]
    if z_min is not None or z_max is not None:
        xyz = {s.id: structure.cartesian(s)[2] for s in structure.sites}
        lo = -math.inf if z_min is None else z_min
        hi = math.inf if z_max is None else z_max
        chosen = [s for s in chosen if lo <= xyz[s] <= hi]
    return chosen



def _indices(structure: PeriodicStructure, sites: list[str]) -> list[int]:
    if not sites:
        raise ValueError("select at least one site")
    ids = [s.id for s in structure.sites]
    unknown = sorted(set(sites) - set(ids))
    if unknown:
        raise ValueError(f"sites not in the structure: {', '.join(unknown)}")
    wanted = set(sites)
    return [i for i, sid in enumerate(ids) if sid in wanted]




def move_sites(
    structure: PeriodicStructure,
    sites: list[str],
    translation: tuple[float, float, float],
    *,
    frame: Literal["cartesian", "fractional"] = "cartesian",
    identifier: str | None = None,
) -> PeriodicStructure:
    """Translate the selected sites by a vector (angstrom, or fractions of the cell)."""
    vector = np.asarray(translation, dtype=float)
    if vector.shape != (3,) or not np.isfinite(vector).all():
        raise ValueError("a translation is three finite numbers")
    lattice = crystal.lattice_rows(structure)
    shift = vector @ lattice if frame == "fractional" else vector
    xyz = crystal.cartesian(structure)
    chosen = _indices(structure, sites)
    xyz[chosen] += shift
    edit = {"edit": "move", "sites": sites, "translation": list(vector), "frame": frame}
    return _rebuild(structure, identifier or structure.id, lattice, xyz, edit)





def rotate_sites(
    structure: PeriodicStructure,
    sites: list[str],
    angle: float,
    *,
    axis: tuple[float, float, float] = (0.0, 0.0, 1.0),
    centre: tuple[float, float, float] | None = None,
    identifier: str | None = None,
) -> PeriodicStructure:
    """Rotate the selected sites by ``angle`` degrees about an axis through a point.

    The default axis is z (the slab normal for builder slabs) through the
    selection's centroid (Rodrigues' rotation formula).
    """
    u = np.asarray(axis, dtype=float)
    if u.shape != (3,) or not np.isfinite(u).all() or np.linalg.norm(u) < 1e-12:
        raise ValueError("the rotation axis is a nonzero vector")
    if not math.isfinite(angle):
        raise ValueError("the rotation angle must be finite")
    u /= np.linalg.norm(u)
    xyz = crystal.cartesian(structure)
    chosen = _indices(structure, sites)
    point = xyz[chosen].mean(axis=0) if centre is None else np.asarray(centre, dtype=float)
    theta = math.radians(angle)
    k = np.asarray([[0, -u[2], u[1]], [u[2], 0, -u[0]], [-u[1], u[0], 0]])
    rot = np.eye(3) + math.sin(theta) * k + (1 - math.cos(theta)) * (k @ k)
    xyz[chosen] = (xyz[chosen] - point) @ rot.T + point
    edit = {
        "edit": "rotate",
        "sites": sites,
        "angle": angle,
        "axis": list(u),
        "centre": list(point),
    }
    return _rebuild(
        structure, identifier or structure.id, crystal.lattice_rows(structure), xyz, edit
    )





def set_sites(
    structure: PeriodicStructure,
    edits: list[dict[str, Any]],
    *,
    identifier: str | None = None,
) -> PeriodicStructure:
    """Set positions (``fractional`` or ``cartesian``) and/or elements of named sites."""
    from vibesketcher.elements import ATOMIC_NUMBER

    lattice = crystal.lattice_rows(structure)
    xyz = crystal.cartesian(structure)
    symbols = [s.symbol for s in structure.sites]
    ids = [s.id for s in structure.sites]
    for change in edits:
        sid = change.get("id")
        if sid not in ids:
            raise ValueError(f"site {sid!r} not found")
        i = ids.index(str(sid))
        if "fractional" in change and "cartesian" in change:
            raise ValueError("give a fractional or a Cartesian position, not both")
        if "fractional" in change:
            xyz[i] = np.asarray(change["fractional"], dtype=float) @ lattice
        if "cartesian" in change:
            xyz[i] = np.asarray(change["cartesian"], dtype=float)
        if "symbol" in change:
            if change["symbol"] not in ATOMIC_NUMBER:
                raise ValueError(f"unknown element {change['symbol']!r}")
            symbols[i] = str(change["symbol"])
    if not np.isfinite(xyz).all():
        raise ValueError("positions must be finite")
    edit = {"edit": "sites", "changes": edits}
    return _rebuild(structure, identifier or structure.id, lattice, xyz, edit, symbols=symbols)





def add_slab_layers(
    document: Document,
    structure: PeriodicStructure,
    count: int,
    *,
    identifier: str | None = None,
) -> PeriodicStructure:
    """Add ``count`` layers on top of a builder slab, under any adsorbates.

    The slab is cut again from its recorded bulk source with the same Miller
    index, layer unit, termination and repeat; the existing sites, including
    any moved or edited ones, stay where they are, the new layers' atoms are
    added above the old top layer, adsorbate atoms move up with the surface
    and the normal vector grows by the added thickness.
    """
    from vibesketcher.document.slab import build_slab

    if not 1 <= count <= 100:
        raise ValueError("add between 1 and 100 layers")
    builder = crystal._builder(structure)
    if builder.get("kind") not in ("slab", "adsorbate") or "layers" not in builder:
        raise ValueError("layers can be added to slabs made by the slab builder")
    source_id = builder.get("source")
    bulk = {s.id: s for s in document.periodic_structures}.get(str(source_id))
    template = builder.get("template")
    if bulk is None and source_id is None and isinstance(template, dict):
        # A one-step template slab (#232): rebuild its conventional bulk.
        from vibesketcher.document import crystal_templates

        overrides = template.get("overrides") or {}
        bulk = crystal_templates.build(str(template["name"]), identifier="bulk",
                                       cell="conventional", a=overrides.get("a"),
                                       c=overrides.get("c"))  # fmt: skip
    if bulk is None:
        raise ValueError(f"the slab's bulk source {source_id!r} is not in the document")
    layers = int(builder["layers"])
    repeat = builder.get("repeat")
    h, k, l_ = (int(x) for x in builder["miller"])
    taller = build_slab(
        bulk,
        (h, k, l_),
        identifier="taller",
        layers=layers + count,
        layer_unit=builder.get("layer_unit", "stoichiometric"),
        termination=int(builder.get("termination", 0)),
        vacuum=float(builder.get("vacuum", 15.0)),
        repeat=repeat if isinstance(repeat, list) else (1, 1),
        fixed_layers=0,
        periodicity="3d" if structure.pbc[2] else "2d",
        surface_cell=builder.get("surface_cell", "primitive"),
    )
    tall = crystal._builder(taller)
    tall_layers = tall["layer_of_site"]
    assert isinstance(tall_layers, dict)
    old_layers = builder.get("layer_of_site")
    if not isinstance(old_layers, dict):
        raise ValueError("the slab records no layer map")
    slab_ids = set(old_layers)
    old_count = Counter(int(v) for v in old_layers.values())
    tall_count = Counter(int(v) for v in tall_layers.values())
    if any(tall_count[k] != old_count[k] for k in old_count):
        raise ValueError("the slab no longer matches its recipe; rebuild it with more layers")
    delta = float(tall["thickness"]) - float(builder["thickness"])
    old_rows = crystal.lattice_rows(structure)
    tall_rows = crystal.lattice_rows(taller)
    if not np.allclose(old_rows[:2], tall_rows[:2], atol=1e-6):
        raise ValueError(
            "the in-plane cell was edited since the slab was cut; rebuild the slab with more layers"
        )
    xyz = crystal.cartesian(structure)
    # Align the recut slab on the existing bottom layer (vacuum edits move it).
    bottom_old = [i for i, s in enumerate(structure.sites) if old_layers.get(s.id) == 0]
    bottom_new = [i for i, s in enumerate(taller.sites) if tall_layers.get(s.id) == 0]
    tall_xyz = crystal.cartesian(taller)
    lift = float(xyz[bottom_old, 2].mean() - tall_xyz[bottom_new, 2].mean())
    # Adsorbate atoms (not in the slab's layer map) ride up with the surface.
    for i, site in enumerate(structure.sites):
        if site.id not in slab_ids:
            xyz[i, 2] += delta
    new_atoms = [
        (s.symbol, tall_xyz[i] + np.asarray([0.0, 0.0, lift]), int(tall_layers[s.id]))
        for i, s in enumerate(taller.sites)
        if int(tall_layers[s.id]) >= layers
    ]
    lattice = old_rows.copy()
    lattice[2] = old_rows[2] + delta * old_rows[2] / np.linalg.norm(old_rows[2])
    counts: Counter[str] = Counter(s.symbol for s in structure.sites)
    taken = {s.id for s in structure.sites}
    inverse = np.linalg.inv(lattice)
    data = structure.model_dump()
    data["id"] = identifier or structure.id
    data["lattice_angstrom"] = tuple(tuple(float(x) for x in row) for row in lattice)
    sites = []
    for i, site in enumerate(data["sites"]):
        site["fractional"] = tuple(crystal._clean(x) for x in xyz[i] @ inverse)
        sites.append(site)
    layer_of = dict(old_layers)
    for symbol, position, layer in new_atoms:
        counts[symbol] += 1
        sid = f"{symbol}{counts[symbol]}"
        while sid in taken:
            counts[symbol] += 1
            sid = f"{symbol}{counts[symbol]}"
        taken.add(sid)
        layer_of[sid] = layer
        sites.append(
            {"id": sid, "symbol": symbol,
             "fractional": tuple(crystal._clean(x) for x in position @ inverse)}
        )  # fmt: skip
    data["sites"] = sites
    builder.update(
        {
            "layers": layers + count,
            "thickness": float(tall["thickness"]),
            "layer_of_site": layer_of,
            "top_layer": layers + count - 1,
            "top_composition": tall["top_composition"],
            "warnings": tall["warnings"],
        }
    )
    edits = list(builder.get("edits") or [])
    edits.append({"edit": "add_layers", "count": count, "added_thickness": delta})
    builder["edits"] = edits
    data["provenance"] = {**structure.provenance, "builder": crystal._json(builder)}
    data["source_cif"] = None
    data["source_model_sha256"] = None
    return PeriodicStructure.model_validate(data)



class _Edit(BaseModel):
    model_config = ConfigDict(extra="forbid")




class EditCellRecipe(_Edit):
    kind: Literal["edit_cell"] = "edit_cell"
    source: _ID
    identifier: _ID | None = None
    a: float | None = None
    b: float | None = None
    c: float | None = None
    alpha: float | None = None
    beta: float | None = None
    gamma: float | None = None
    scale: float | None = None
    vacuum: float | None = None
    keep: Keep = "fractional"





class Selection(_Edit):
    sites: list[str] | None = None
    elements: list[str] | None = None
    layers: list[int] | None = None
    adsorbate: int | str | None = None
    z_min: float | None = None
    z_max: float | None = None





class MoveSitesRecipe(_Edit):
    kind: Literal["move_sites"] = "move_sites"
    source: _ID
    identifier: _ID | None = None
    select: Selection
    translation: tuple[float, float, float] = (0.0, 0.0, 0.0)
    frame: Literal["cartesian", "fractional"] = "cartesian"
    rotate: float | None = None
    axis: tuple[float, float, float] = (0.0, 0.0, 1.0)





class SetSitesRecipe(_Edit):
    kind: Literal["set_sites"] = "set_sites"
    source: _ID
    identifier: _ID | None = None
    sites: list[dict[str, Any]] = Field(min_length=1, max_length=25000)





class AddLayersRecipe(_Edit):
    kind: Literal["add_layers"] = "add_layers"
    source: _ID
    identifier: _ID | None = None
    count: int = Field(ge=1, le=100)



EditRecipe = EditCellRecipe | MoveSitesRecipe | SetSitesRecipe | AddLayersRecipe
EDIT_KINDS = ("edit_cell", "move_sites", "set_sites", "add_layers")




def apply_edit(document: Document, recipe: EditRecipe) -> PeriodicStructure:
    """Run one edit; without ``identifier`` the result replaces its source."""
    existing = {s.id: s for s in document.periodic_structures}
    if recipe.source not in existing:
        raise ValueError(f"periodic structure {recipe.source!r} not found")
    target = recipe.identifier or recipe.source
    if target != recipe.source and target in existing:
        raise ValueError(f"periodic structure {target!r} already exists")
    source = existing[recipe.source]
    if isinstance(recipe, EditCellRecipe):
        return edit_cell(
            source,
            identifier=target,
            **recipe.model_dump(include={"a", "b", "c", "alpha", "beta", "gamma", "scale",
                                         "vacuum", "keep"}),
        )  # fmt: skip
    if isinstance(recipe, MoveSitesRecipe):
        sites = select_sites(source, **recipe.select.model_dump())
        moved = source
        if recipe.rotate:
            moved = rotate_sites(moved, sites, recipe.rotate, axis=recipe.axis, identifier=target)
        if any(recipe.translation) or not recipe.rotate:
            moved = move_sites(
                moved, sites, recipe.translation, frame=recipe.frame, identifier=target
            )
        return moved
    if isinstance(recipe, SetSitesRecipe):
        return set_sites(source, recipe.sites, identifier=target)
    return add_slab_layers(document, source, recipe.count, identifier=target)





def parse_edit(data: dict[str, Any]) -> EditRecipe:
    from pydantic import TypeAdapter

    adapter: TypeAdapter[EditRecipe] = TypeAdapter(
        Annotated[EditRecipe, Field(discriminator="kind")]
    )
    return adapter.validate_python(data)





def edit(document: Document, recipe: EditRecipe) -> Document:
    """Apply an edit and put its structure into a copy of the document."""
    from vibesketcher.document.periodic import with_periodic

    return with_periodic(document, apply_edit(document, recipe))