vibe-sketcher/Documentation

Source code for vibesketcher.render.schematic

"""Schematic aids for periodic publication figures (#233).

Every function here reads the 3D model and returns drawing data; none of them
changes a structure. The figure in :mod:`vibesketcher.render.periodic` uses
them for heights, coordination polyhedra and slab layer schematics.
"""

from __future__ import annotations

import itertools
from fractions import Fraction

import numpy as np

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

_GLYPHS = {Fraction(1, 2): "½", Fraction(1, 4): "¼", Fraction(3, 4): "¾", Fraction(1, 3): "⅓",
           Fraction(2, 3): "⅔", Fraction(1, 8): "⅛", Fraction(3, 8): "⅜", Fraction(5, 8): "⅝",
           Fraction(7, 8): "⅞", Fraction(1, 6): "⅙", Fraction(5, 6): "⅚"}  # fmt: skip




def view_period(structure: PeriodicStructure, direction: np.ndarray) -> np.ndarray | None:
    """The shortest lattice vector along the view direction, if one exists (indices up to 6)."""
    rows = crystal.lattice_rows(structure)
    d = direction / np.linalg.norm(direction)
    periodic = [i for i in range(3) if structure.pbc[i]]
    best: np.ndarray | None = None
    for coefficients in itertools.product(range(-6, 7), repeat=len(periodic)):
        if not any(coefficients):
            continue
        vector = sum((c * rows[i] for c, i in zip(coefficients, periodic, strict=True)),
                     np.zeros(3))  # fmt: skip
        length = float(np.linalg.norm(vector))
        if abs(abs(float(vector @ d)) - length) < 1e-6 * max(length, 1.0) and (
            best is None or length < float(np.linalg.norm(best)) - 1e-9
        ):
            best = vector if float(vector @ d) > 0 else -vector
    return best





def fraction_text(value: float) -> str:
    """0, ½, ¼, 0.37: a height in the projection, as crystallographers write it."""
    value %= 1.0
    if value > 1 - 1e-3:
        value = 0.0
    exact = Fraction(value).limit_denominator(8)
    if abs(float(exact) - value) < 2e-3:
        if exact == 0:
            return "0"
        return _GLYPHS.get(exact, f"{exact.numerator}/{exact.denominator}")
    return f"{value:.2f}".lstrip("0")





def heights(
    structure: PeriodicStructure, positions: np.ndarray, toward: np.ndarray
) -> tuple[list[str], str]:
    """Height labels of the atoms toward the viewer, and what they measure.

    Along a lattice direction the heights are fractions of that period;
    otherwise (for example across a slab's vacuum) they are in angstrom above
    the lowest atom.
    """
    period = view_period(structure, toward)
    if period is not None:
        length = float(np.linalg.norm(period))
        unit = period / length
        values = positions @ unit / length
        return [fraction_text(v) for v in values], f"heights in units of {length:.3f} Å"
    depth = positions @ (toward / np.linalg.norm(toward))
    return [f"{v:.1f}" for v in depth - depth.min()], "heights in Å above the lowest atom"





def hull_faces(points: np.ndarray, tolerance: float = 1e-3) -> list[list[int]]:
    """Faces (vertex indices, ordered around the face) of the convex hull of a few points."""
    n = len(points)
    if n < 3:
        return []
    centre = points.mean(axis=0)
    if n == 3 or np.linalg.matrix_rank(points - centre, tol=1e-6) < 3:
        normal = np.cross(points[1] - points[0], points[2] - points[0])
        return [_ordered(list(range(n)), points, normal)] if np.linalg.norm(normal) > 1e-9 else []
    faces: dict[tuple[int, ...], list[int]] = {}
    for i, j, k in itertools.combinations(range(n), 3):
        normal = np.cross(points[j] - points[i], points[k] - points[i])
        if np.linalg.norm(normal) < 1e-9:
            continue
        normal /= np.linalg.norm(normal)
        side = (points - points[i]) @ normal
        if np.all(side <= tolerance) or np.all(side >= -tolerance):
            members = tuple(int(m) for m in np.nonzero(np.abs(side) <= tolerance)[0])
            if members not in faces:
                faces[members] = _ordered(list(members), points, normal)
    return list(faces.values())



def _ordered(members: list[int], points: np.ndarray, normal: np.ndarray) -> list[int]:
    face = points[members]
    centre = face.mean(axis=0)
    e1 = face[0] - centre
    e1 /= np.linalg.norm(e1)
    e2 = np.cross(normal / np.linalg.norm(normal), e1)
    angles = [float(np.arctan2((p - centre) @ e2, (p - centre) @ e1)) for p in face]
    return [members[i] for i in np.argsort(angles)]




def coordination(
    structure: PeriodicStructure,
    centre: np.ndarray,
    centre_symbol: str,
    radius: dict[str, float],
    scale: float = 1.15,
) -> np.ndarray:
    """Cartesian positions of the atoms bonded to an atom at ``centre`` (all images)."""
    rows = crystal.lattice_rows(structure)
    shifts = list(itertools.product(*((-1, 0, 1) if p else (0,) for p in structure.pbc)))
    found = []
    for site in structure.sites:
        if site.symbol == centre_symbol:
            continue
        base = np.asarray(site.fractional, dtype=float) @ rows
        limit = scale * (radius[centre_symbol] + radius[site.symbol])
        for shift in shifts:
            p = base + np.asarray(shift) @ rows
            distance = float(np.linalg.norm(p - centre))
            if 0.4 < distance < limit:
                found.append(p)
    return np.asarray(found) if found else np.zeros((0, 3))





def stacking_letters(structure: PeriodicStructure, tolerance: float = 0.3) -> dict[str, str]:
    """Stacking letter (A, B, C, ...) of each site's atomic plane in a slab.

    Planes are grouped by height along the normal; two planes get the same
    letter when their in-plane fractional positions coincide modulo the
    surface lattice.
    """
    rows = crystal.lattice_rows(structure)
    normal = np.cross(rows[0], rows[1])
    normal /= np.linalg.norm(normal)
    xyz = crystal.cartesian(structure)
    z = xyz @ normal
    order = np.argsort(z)
    planes: list[list[int]] = []
    for index in order:
        if planes and z[index] - z[planes[-1][-1]] <= tolerance:
            planes[-1].append(int(index))
        else:
            planes.append([int(index)])
    in_plane = np.linalg.inv(rows)
    letters: dict[str, str] = {}
    registry: list[tuple[frozenset[tuple[str, int, int]], str]] = []
    for plane in planes:
        key = frozenset(
            (structure.sites[i].symbol,
             int(round(((xyz[i] @ in_plane)[0] % 1.0) * 50)) % 50,
             int(round(((xyz[i] @ in_plane)[1] % 1.0) * 50)) % 50)
            for i in plane
        )  # fmt: skip
        letter = next((name for k, name in registry if k == key), None)
        if letter is None:
            letter = chr(ord("A") + len(registry)) if len(registry) < 26 else "?"
            registry.append((key, letter))
        for i in plane:
            letters[structure.sites[i].id] = letter
    return letters