vimol 0.1.0__py3-none-any.whl

vimol/scene.py

@@ -0,0 +1,239 @@
+"""A Scene binds a molecule, a camera, a style and a renderer together.
+
+It is the main object embedding apps interact with: set the pixel size, call
+:meth:`render` to get an RGB numpy array, or :meth:`to_kitty` for protocol
+bytes ready to write to a terminal.
+"""
+from __future__ import annotations
+
+from typing import Optional, Tuple
+
+import numpy as np
+
+from .camera import Camera
+from .molecule import Molecule
+from .render import Renderer, Style
+from . import kitty
+
+
+def _block_sums(img: np.ndarray, factor: int) -> np.ndarray:
+    """Per-block uint16 channel sums over factor x factor tiles.
+
+    Strided-slice accumulation instead of ``reshape(...).mean(axis=(1, 3))``:
+    the reshape-mean runs in float64 over two non-contiguous axes and was
+    the single most expensive step of a supersampled frame (~170 ms at
+    3200x2000); f*f strided adds in uint16 do the same reduction ~6x faster.
+    uint16 is safe up to factor 16 (16*16*255 < 65536).
+    """
+    h, w = img.shape[:2]
+    ch = img.shape[2]
+    h2, w2 = h // factor, w // factor
+    img = img[: h2 * factor, : w2 * factor]
+    acc = np.zeros((h2, w2, ch), np.uint16)
+    for dy in range(factor):
+        for dx in range(factor):
+            acc += img[dy::factor, dx::factor]
+    return acc
+
+
+def _downsample(img: np.ndarray, factor: int) -> np.ndarray:
+    if factor <= 1:
+        return img
+    ch = img.shape[2]
+    sums = _block_sums(img, factor)
+    if ch == 4:
+        # premultiplied-alpha average so anti-aliased edges don't fringe toward
+        # black over a transparent background. RGB is already premultiplied
+        # (undrawn pixels are 0,0,0,0), so straight = sum(rgb)/coverage.
+        rgb_sum = sums[..., :3].astype(np.float32)
+        a_sum = sums[..., 3].astype(np.float32)
+        with np.errstate(divide="ignore", invalid="ignore"):
+            straight = np.where(a_sum[..., None] > 0, 255.0 * rgb_sum / a_sum[..., None], 0.0)
+        alpha = a_sum / (factor * factor)
+        out = np.dstack([np.clip(straight, 0, 255), alpha])
+        return out.astype(np.uint8)
+    # integer floor-average == float mean + truncating uint8 cast, so this is
+    # pixel-identical to the old reshape/mean path
+    return (sums // (factor * factor)).astype(np.uint8)
+
+
+class Scene:
+    def __init__(self, molecule: Molecule, width: int = 640, height: int = 480,
+                 style: Optional[Style] = None, supersample: int = 1,
+                 backend: str = "auto"):
+        self.molecule = molecule
+        self.style = style or Style()
+        self.width = width
+        self.height = height
+        self.supersample = max(1, int(supersample))
+        self.render_scale = 1.0        # dynamic-resolution factor, see set_render_scale
+        self.camera = Camera(center=molecule.centroid(), extent=molecule.radius_of_gyration_extent())
+        self._backend_name, self._renderer = self._make_renderer(
+            backend, width * self.supersample, height * self.supersample)
+        self.fit()
+
+    @staticmethod
+    def _make_renderer(backend: str, w: int, h: int):
+        if backend not in ("auto", "cpu", "gl"):
+            raise ValueError(f"unknown backend {backend!r} (expected 'auto', 'cpu', or 'gl')")
+        if backend == "cpu":
+            return "cpu", Renderer(w, h)
+        if backend == "gl":
+            from .gl_render import GLRenderer  # raises if unavailable -- explicit request, no silent fallback
+            return "gl", GLRenderer(w, h)
+        # "auto": prefer GL, silently fall back to the CPU raycaster if no
+        # GL package/context is available (missing moderngl, no driver, no
+        # display -- see gl_render.py's docstring).
+        try:
+            from .gl_render import GLRenderer
+            return "gl", GLRenderer(w, h)
+        except Exception:
+            return "cpu", Renderer(w, h)
+
+    @property
+    def backend(self) -> str:
+        """Which renderer actually got picked: ``"cpu"`` or ``"gl"``."""
+        return self._backend_name
+
+    # -- sizing / framing -------------------------------------------------
+    def set_size(self, width: int, height: int, refit: bool = False) -> None:
+        """Resize the viewport. By default this preserves the current
+        *apparent framing* -- the fraction of the viewport the molecule fills,
+        plus any manual scroll-zoom -- rather than the raw pixels-per-angstrom.
+        It does that by rescaling zoom/pan by the min-dimension ratio, exactly
+        as :meth:`set_supersample` does for a quality switch.
+
+        This matters for two reasons. (1) For a fit-derived zoom the rescale
+        reproduces *precisely* what a fresh :meth:`fit` would compute for the
+        new size (both reduce to ``min(W,H)/2 / (extent*margin)``), so if the
+        host learns the real terminal size in two steps -- an early, slightly
+        wrong report followed by the settled one -- the second, non-refit
+        resize self-heals to the correct framing instead of freezing the
+        molecule at the first size's zoom. (2) For a user's manual zoom it
+        keeps the molecule the same on-screen fraction across a terminal
+        resize, which is the intuitive behaviour.
+
+        Pass ``refit=True`` to snap to a fresh fit regardless (used the very
+        first time a host establishes real geometry from a placeholder size).
+        """
+        old_min = max(min(self._renderer.width, self._renderer.height), 1)
+        self.width = max(1, int(width))
+        self.height = max(1, int(height))
+        self._resize_renderer()
+        if refit:
+            self.fit()
+            return
+        new_min = max(min(self._renderer.width, self._renderer.height), 1)
+        scale = new_min / old_min
+        self.camera.zoom *= scale
+        self.camera.pan = self.camera.pan * scale
+
+    def set_supersample(self, factor: int) -> None:
+        new_ss = max(1, int(factor))
+        if new_ss == self.supersample:
+            return
+        # Rescale zoom/pan (both in supersampled-buffer pixel units) by the
+        # exact buffer-size ratio so the apparent on-screen picture --
+        # including any manual scroll-to-zoom -- is preserved across a
+        # supersample change. This must NOT go through fit()/Camera.fit():
+        # that recomputes zoom purely from the molecule's extent, with no
+        # memory of the user's current zoom, so it silently discarded any
+        # scroll-zoom every time the interactive quality switch (fast while
+        # dragging/scrolling -> crisp ~0.25s after stopping) fired.
+        old_min = max(min(self._renderer.width, self._renderer.height), 1)
+        self.supersample = new_ss
+        self._resize_renderer()
+        new_min = max(min(self._renderer.width, self._renderer.height), 1)
+        scale = new_min / old_min
+        self.camera.zoom *= scale
+        self.camera.pan = self.camera.pan * scale
+
+    def set_render_scale(self, factor: float) -> None:
+        """Render internally at *factor* x the logical resolution (0.1..1).
+
+        Below 1.0 the scene draws fewer pixels and the terminal (or host)
+        stretches the image back up -- slightly soft/grainy, but the frame
+        cost drops with the square of the factor. Interactive hosts use this
+        as dynamic resolution while the camera is moving, snapping back to
+        1.0 (and full supersampling) once idle. Zoom/pan rescale by the exact
+        buffer-size ratio, like :meth:`set_supersample`, so the apparent
+        framing never changes. The 0.1 hard floor sits below the viewer's own
+        operating floor, so it only guards against absurd inputs.
+        """
+        f = min(1.0, max(0.1, float(factor)))
+        if f == self.render_scale:
+            return
+        old_min = max(min(self._renderer.width, self._renderer.height), 1)
+        self.render_scale = f
+        self._resize_renderer()
+        new_min = max(min(self._renderer.width, self._renderer.height), 1)
+        scale = new_min / old_min
+        self.camera.zoom *= scale
+        self.camera.pan = self.camera.pan * scale
+
+    def _resize_renderer(self) -> None:
+        self._renderer.resize(
+            max(1, int(round(self.width * self.supersample * self.render_scale))),
+            max(1, int(round(self.height * self.supersample * self.render_scale))))
+
+    def fit(self, keep_orientation: bool = False, keep_zoom: bool = False) -> None:
+        rot = self.camera.rotation.copy()
+        pan = self.camera.pan.copy()
+        zoom = self.camera.zoom
+        self.camera.center = self.molecule.centroid()
+        ext = max(
+            self.molecule.radius_of_gyration_extent() + self._max_atom_radius(),
+            self.molecule.vector_extent(),
+        )
+        self.camera.fit(self._renderer.width, self._renderer.height, ext)
+        if keep_orientation:
+            self.camera.rotation = rot
+            self.camera.pan = pan
+        if keep_zoom:
+            self.camera.zoom = zoom
+
+    def _max_atom_radius(self) -> float:
+        if self.molecule.n_atoms == 0:
+            return 0.0
+        if self.style.representation == "spacefill":
+            return float(self.molecule.vdw_radii().max())
+        return float(self.molecule.vdw_radii().max()) * self.style.atom_scale
+
+    def set_molecule(self, molecule: Molecule) -> None:
+        self.molecule = molecule
+        self.camera.center = molecule.centroid()
+        self.fit()
+
+    @property
+    def render_size(self):
+        """(width, height) of the internal render buffer, in supersampled px."""
+        return self._renderer.width, self._renderer.height
+
+    # -- rendering --------------------------------------------------------
+    def render(self) -> np.ndarray:
+        if self._backend_name == "gl":
+            from .gl_adapter import molecule_to_gl_inputs
+            w, h = self._renderer.width, self._renderer.height
+            spheres, cylinders, cones, proj, shading = molecule_to_gl_inputs(
+                self.molecule, self.camera, self.style, w, h)
+            # the GL renderer supersamples/downsamples on the GPU, so it returns
+            # a display-size image directly -- no CPU _downsample pass.
+            return self._renderer.render(spheres, cylinders, proj, shading,
+                                         downsample=self.supersample, cones=cones)
+        img = self._renderer.render(self.molecule, self.camera, self.style)
+        return _downsample(img, self.supersample)
+
+    def to_kitty(self, *, image_id: Optional[int] = None, cols: Optional[int] = None,
+                 rows: Optional[int] = None, move_cursor: bool = False) -> bytes:
+        img = self.render()
+        if image_id is None:
+            # per-process id so concurrent Scenes (e.g. two panes of one
+            # kitty process) don't clobber each other's image storage.
+            image_id = kitty.unique_id_base() + 1
+        return kitty.encode_image(img, image_id=image_id, cols=cols, rows=rows,
+                                  move_cursor=move_cursor)
+
+    def to_png(self, path: str) -> None:
+        img = self.render()
+        with open(path, "wb") as fh:
+            fh.write(kitty.png_bytes(img))

vimol/templates.py

@@ -0,0 +1,205 @@
+"""Geometry templates for interactive molecule building.
+
+For each ``(element, valence)`` combination we record the *ideal* local
+arrangement of bonds around a central atom -- tetrahedral for four-coordinate
+carbon, pyramidal for three-coordinate nitrogen, bent for two-coordinate
+oxygen, and so on. Each arrangement is a small set of unit bond directions in a
+canonical local frame.
+
+The editor uses these to grow chemically-sensible fragments: place a central
+atom, orient its first direction toward the atom it attaches to, and cap the
+remaining directions with hydrogens.
+
+Everything here is pure geometry (unit vectors and rotations); actual atom
+placement, bond lengths and molecule mutation live in :mod:`vimol.editor`.
+"""
+from __future__ import annotations
+
+from dataclasses import dataclass
+
+import numpy as np
+
+from .camera import rotation_from_axis_angle
+from .elements import normalize_symbol as _normalize_symbol
+
+
+def _normalize(v: np.ndarray) -> np.ndarray:
+    v = np.asarray(v, float)
+    n = np.linalg.norm(v)
+    return v / n if n > 1e-12 else v
+
+
+def _perpendicular(a: np.ndarray) -> np.ndarray:
+    """Return an arbitrary unit vector perpendicular to *a*."""
+    a = _normalize(a)
+    # cross with whichever cardinal axis is least aligned with a
+    ref = np.array([1.0, 0.0, 0.0]) if abs(a[0]) < 0.9 else np.array([0.0, 1.0, 0.0])
+    return _normalize(np.cross(a, ref))
+
+
+def rotation_between(a: np.ndarray, b: np.ndarray) -> np.ndarray:
+    """Rotation matrix taking unit vector *a* onto unit vector *b*."""
+    a = _normalize(a)
+    b = _normalize(b)
+    v = np.cross(a, b)
+    c = float(np.clip(np.dot(a, b), -1.0, 1.0))
+    s = float(np.linalg.norm(v))
+    if s < 1e-9:                       # parallel or antiparallel
+        if c > 0:
+            return np.eye(3)
+        return rotation_from_axis_angle(_perpendicular(a), np.pi)
+    return rotation_from_axis_angle(v / s, np.arctan2(s, c))
+
+
+# -- canonical direction sets (unit vectors) ------------------------------
+_TETRAHEDRAL = np.array([
+    [1.0, 1.0, 1.0],
+    [1.0, -1.0, -1.0],
+    [-1.0, 1.0, -1.0],
+    [-1.0, -1.0, 1.0],
+]) / np.sqrt(3.0)
+
+_S3 = np.sqrt(3.0) / 2.0
+_TRIGONAL = np.array([
+    [1.0, 0.0, 0.0],
+    [-0.5, _S3, 0.0],
+    [-0.5, -_S3, 0.0],
+])
+
+_LINEAR = np.array([
+    [0.0, 0.0, 1.0],
+    [0.0, 0.0, -1.0],
+])
+
+_TERMINAL = np.array([[0.0, 0.0, 1.0]])
+
+
+@dataclass(frozen=True)
+class AtomTemplate:
+    """The ideal local geometry for one ``(element, valence)`` combination.
+
+    ``directions`` is an ``(valence, 3)`` array of unit bond directions in a
+    canonical frame. ``directions[0]`` is the *attachment slot* -- the site
+    oriented toward the parent atom when the fragment is grown onto an existing
+    molecule; ``directions[1:]`` are the remaining open sites that get capped
+    (with :attr:`cap`, hydrogen by default). When a fragment is born in free
+    space (e.g. a lone methane), every direction is used.
+    """
+
+    element: str
+    valence: int
+    geometry: str
+    directions: np.ndarray
+    cap: str = "H"
+    hybridization: str = ""          # "sp3" / "sp2" / "sp"; "" when N/A (e.g. H, halogens)
+
+    def label(self) -> str:
+        """A short human label, e.g. 'tetrahedral · sp3 · 4 bonds'."""
+        bonds = f"{self.valence} bond" + ("s" if self.valence != 1 else "")
+        parts = [self.geometry]
+        if self.hybridization:
+            parts.append(self.hybridization)
+        parts.append(bonds)
+        return " · ".join(parts)
+
+    def free_directions(self) -> np.ndarray:
+        """All bond directions in the canonical frame (for a free-standing atom)."""
+        return self.directions
+
+    def open_directions(self, attach_dir: np.ndarray) -> np.ndarray:
+        """Open-site directions when the attachment slot points along *attach_dir*.
+
+        *attach_dir* is the world-space direction from the new central atom
+        toward the atom it bonds to. Returns the ``valence - 1`` remaining bond
+        directions, rotated into world space, ready to receive caps.
+        """
+        if self.valence <= 1:
+            return np.zeros((0, 3))
+        R = rotation_between(self.directions[0], _normalize(attach_dir))
+        return self.directions[1:] @ R.T
+
+
+# -- registry: hybridization/valence options per element ------------------
+# Mirrors the set of atom-type choices a builder like GaussView offers for each
+# element: the standard VSEPR hybridization states (sp3 tetrahedral, sp2
+# trigonal, sp linear, and the lone-pair-reduced pyramidal/bent variants). The
+# order here is the order the geometry picker lists them in (most bonds first).
+TEMPLATES = {
+    # Carbon: sp3 / sp2 / sp
+    ("C", 4): AtomTemplate("C", 4, "tetrahedral", _TETRAHEDRAL, hybridization="sp3"),
+    ("C", 3): AtomTemplate("C", 3, "trigonal", _TRIGONAL, hybridization="sp2"),
+    ("C", 2): AtomTemplate("C", 2, "linear", _LINEAR, hybridization="sp"),
+    # Nitrogen: ammonium sp3 / amine sp3 / imine sp2 / nitrile sp
+    ("N", 4): AtomTemplate("N", 4, "tetrahedral", _TETRAHEDRAL, hybridization="sp3"),
+    ("N", 3): AtomTemplate("N", 3, "pyramidal", _TETRAHEDRAL[:3], hybridization="sp3"),
+    ("N", 2): AtomTemplate("N", 2, "bent", _TRIGONAL[:2], hybridization="sp2"),
+    ("N", 1): AtomTemplate("N", 1, "terminal", _TERMINAL, hybridization="sp"),
+    # Oxygen: ether/hydroxyl sp3 / carbonyl-like sp2 / terminal
+    ("O", 2): AtomTemplate("O", 2, "bent", _TETRAHEDRAL[:2], hybridization="sp3"),
+    ("O", 1): AtomTemplate("O", 1, "terminal", _TERMINAL, hybridization="sp2"),
+    # Boron: sp2 / sp3
+    ("B", 3): AtomTemplate("B", 3, "trigonal", _TRIGONAL, hybridization="sp2"),
+    ("B", 4): AtomTemplate("B", 4, "tetrahedral", _TETRAHEDRAL, hybridization="sp3"),
+    # Heavier main-group with a couple of common valences
+    ("P", 3): AtomTemplate("P", 3, "pyramidal", _TETRAHEDRAL[:3], hybridization="sp3"),
+    ("S", 2): AtomTemplate("S", 2, "bent", _TETRAHEDRAL[:2], hybridization="sp3"),
+    # Monovalent (halogens, hydrogen): a single terminal bond
+    ("H", 1): AtomTemplate("H", 1, "terminal", _TERMINAL),
+    ("F", 1): AtomTemplate("F", 1, "terminal", _TERMINAL),
+    ("Cl", 1): AtomTemplate("Cl", 1, "terminal", _TERMINAL),
+    ("Br", 1): AtomTemplate("Br", 1, "terminal", _TERMINAL),
+    ("I", 1): AtomTemplate("I", 1, "terminal", _TERMINAL),
+}
+
+# Typical valence used to pick a default template when none is specified.
+_DEFAULT_VALENCE = {"C": 4, "N": 3, "O": 2, "B": 3, "P": 3, "S": 2,
+                    "H": 1, "F": 1, "Cl": 1, "Br": 1, "I": 1}
+
+
+def default_template(element: str = "C") -> AtomTemplate:
+    """The default building template for *element* (its most common valence).
+
+    Falls back to a tetrahedral, carbon-like template for elements without a
+    dedicated entry, so an unknown element still grows sensibly.
+    """
+    element = _normalize_symbol(element)
+    val = _DEFAULT_VALENCE.get(element, 4)
+    tmpl = TEMPLATES.get((element, val))
+    if tmpl is not None:
+        return tmpl
+    # synthesize a tetrahedral fallback for the requested element
+    return AtomTemplate(element, 4, "tetrahedral", _TETRAHEDRAL, hybridization="sp3")
+
+
+def options_for(element: str):
+    """The geometry/hybridization templates offered for *element*.
+
+    Returns every registered template for the element, most bonds first (the
+    order the picker lists them). Elements with no dedicated entry fall back to
+    a single default template, so the picker always has at least one option.
+    """
+    element = _normalize_symbol(element)
+    opts = [t for (el, _val), t in TEMPLATES.items() if el == element]
+    if not opts:
+        return [default_template(element)]
+    opts.sort(key=lambda t: -t.valence)
+    return opts
+
+
+def free_direction(neighbor_dirs) -> np.ndarray:
+    """A sensible open bond direction on an atom that already has neighbors.
+
+    *neighbor_dirs* is a sequence of unit vectors pointing from the atom toward
+    its existing neighbors. The result points away from their average (VSEPR-ish
+    "as far from everything else as possible"); with no neighbors it defaults to
+    +z, and for a symmetric arrangement (e.g. two opposed bonds) it picks a
+    perpendicular direction rather than collapsing to zero.
+    """
+    dirs = [np.asarray(d, float) for d in neighbor_dirs]
+    if not dirs:
+        return np.array([0.0, 0.0, 1.0])
+    s = np.sum([_normalize(d) for d in dirs], axis=0)
+    n = float(np.linalg.norm(s))
+    if n < 1e-6:
+        return _perpendicular(dirs[0])
+    return -s / n