vimol 0.1.0__py3-none-any.whl

vimol/__init__.py

@@ -0,0 +1,63 @@
+"""vimol — an embeddable terminal molecular viewer using the Kitty graphics protocol.
+
+Quick start
+-----------
+    import vimol
+
+    mol = vimol.load("caffeine.pdb")     # parse a structure
+    scene = vimol.Scene(mol, 640, 480)    # bind camera + renderer
+    scene.camera.orbit(30, 15)              # rotate
+    img = scene.render()                    # (H, W, 3) uint8 numpy array
+
+    # display in a Kitty-capable terminal:
+    import sys, os
+    os.write(1, scene.to_kitty())
+
+    # or run the full interactive viewer:
+    vimol.view(mol)
+
+Everything here is pure Python + numpy, no GPU or windowing system required.
+"""
+from __future__ import annotations
+
+from .molecule import Molecule, VectorField
+from .camera import Camera
+from .render import Renderer, Style
+from .scene import Scene
+from .widget import MoleculeWidget
+from .input import InputDecoder, MouseEvent, KeyEvent, enable_mouse, disable_mouse
+from .parsers import load, load_all, loads, save, SUPPORTED_EXTENSIONS, SAVE_EXTENSIONS
+from .bonds import perceive_bonds, ensure_bonds
+from . import kitty
+from . import elements
+from . import templates
+from . import editor
+
+__version__ = "0.1.0"
+
+__all__ = [
+    "Molecule", "VectorField", "Camera", "Renderer", "Style", "Scene", "MoleculeWidget",
+    "InputDecoder", "MouseEvent", "KeyEvent", "enable_mouse", "disable_mouse",
+    "load", "load_all", "loads", "save", "SUPPORTED_EXTENSIONS", "SAVE_EXTENSIONS",
+    "perceive_bonds", "ensure_bonds", "kitty", "elements", "templates", "editor",
+    "view", "__version__",
+]
+
+
+def view(molecule_or_path, **kwargs):
+    """Launch the interactive terminal viewer.
+
+    Accepts a :class:`Molecule` or a path to a structure file. Extra keyword
+    arguments are forwarded to :class:`vimol.viewer.Viewer`.
+    """
+    from .viewer import Viewer
+
+    if isinstance(molecule_or_path, str):
+        mol = load(molecule_or_path)
+        kwargs.setdefault("source_path", molecule_or_path)
+    else:
+        mol = molecule_or_path
+    ensure_bonds(mol)
+    v = Viewer(mol, **kwargs)
+    v.widget.scene.camera.orbit(20, -15)  # a pleasant default 3/4 view
+    v.run()

vimol/__main__.py

@@ -0,0 +1,4 @@
+from .app import main
+
+if __name__ == "__main__":
+    raise SystemExit(main())

vimol/app.py

@@ -0,0 +1,218 @@
+"""Command-line driver for vimol.
+
+    vimol                           # opens the bundled C60 demo (checkout only)
+    vimol file.pdb                 # interactive viewer (opens editable: a=append)
+    vimol file.xyz --spin          # autospinning
+    vimol file.mol --render out.png --size 800x800
+    vimol file.pdb --kitty         # emit one frame to stdout (for pipes/embeds)
+    vimol file.xyz --info          # print structure info and exit
+"""
+from __future__ import annotations
+
+import argparse
+import os
+import sys
+from typing import List, Optional
+
+from .parsers import load_all, SUPPORTED_EXTENSIONS
+from .bonds import ensure_bonds
+from .render import Style
+from .scene import Scene
+from . import kitty
+from .molecule import Molecule
+
+
+def _parse_size(s: str):
+    if "x" in s.lower():
+        w, h = s.lower().split("x")
+        return int(w), int(h)
+    v = int(s)
+    return v, v
+
+
+def build_style(args) -> Style:
+    st = Style(representation=args.style)
+    if args.atom_scale is not None:
+        st.atom_scale = args.atom_scale
+    if args.bond_radius is not None:
+        st.bond_radius = args.bond_radius
+    if args.background is not None:
+        st.background = _parse_color(args.background)
+        st.transparent = False
+    if args.transparent:
+        st.transparent = True
+    if args.opaque:
+        st.transparent = False
+    if args.no_depth_cue:
+        st.depth_cue = 0.0
+    return st
+
+
+def _parse_color(s: str):
+    s = s.strip().lstrip("#")
+    if len(s) == 6:
+        return tuple(int(s[i:i + 2], 16) / 255 for i in (0, 2, 4))
+    parts = s.split(",")
+    if len(parts) == 3:
+        return tuple(float(p) for p in parts)
+    raise argparse.ArgumentTypeError(f"bad color {s!r}")
+
+
+def make_parser() -> argparse.ArgumentParser:
+    p = argparse.ArgumentParser(
+        prog="vimol",
+        description="A terminal molecular viewer using the Kitty graphics protocol.",
+        formatter_class=argparse.RawDescriptionHelpFormatter,
+        epilog="Supported formats: " + ", ".join(SUPPORTED_EXTENSIONS),
+    )
+    p.add_argument("file", nargs="?", help="structure file (xyz/pdb/mol/sdf)")
+    p.add_argument("--style", default="ball_and_stick",
+                   choices=["ball_and_stick", "spacefill", "licorice", "wireframe"])
+    p.add_argument("--backend", default="auto", choices=["auto", "cpu", "gl"],
+                   help="rendering backend: numpy CPU raycaster, GPU (OpenGL, needs "
+                        "vimol[gl]), or auto (GPU if available, else CPU)")
+    p.add_argument("--size", default="0x0", help="pixel size WxH for --render/--kitty (0=auto)")
+    p.add_argument("--supersample", type=int, default=2, help="anti-aliasing factor for stills")
+    p.add_argument("--rotate", nargs=2, type=float, metavar=("YAW", "PITCH"),
+                   default=(20.0, -15.0), help="initial rotation in degrees")
+    p.add_argument("--spin", action="store_true", help="autospin in interactive mode")
+    p.add_argument("--render", metavar="PNG", help="render a still image to a PNG file and exit")
+    p.add_argument("--kitty", action="store_true", help="emit one frame to stdout as Kitty graphics and exit")
+    p.add_argument("--info", action="store_true", help="print structure info and exit")
+    p.add_argument("--frame", type=int, default=0, help="frame/model index for multi-model files")
+    p.add_argument("--atom-scale", type=float, default=None)
+    p.add_argument("--bond-radius", type=float, default=None)
+    p.add_argument("--background", type=str, default=None, help="hex or r,g,b background color (implies --opaque)")
+    p.add_argument("--transparent", action="store_true", help="transparent background (RGBA cutout)")
+    p.add_argument("--opaque", action="store_true", help="solid background (default for --render)")
+    p.add_argument("--no-depth-cue", action="store_true")
+    p.add_argument("--no-bonds", action="store_true", help="do not auto-perceive bonds")
+    p.add_argument("--bond-tolerance", type=float, default=0.45)
+    p.add_argument("--list-formats", action="store_true")
+    p.add_argument("--version", action="store_true")
+    return p
+
+
+def _print_info(mol: Molecule):
+    print(f"name:    {mol.name}")
+    print(f"formula: {mol.formula()}")
+    print(f"atoms:   {mol.n_atoms}")
+    print(f"bonds:   {len(mol.bonds)}")
+    ext = mol.radius_of_gyration_extent()
+    print(f"extent:  {ext:.2f} A (max atom-centroid distance)")
+    from collections import Counter
+    comp = Counter(mol.symbols)
+    print("composition: " + ", ".join(f"{k}:{v}" for k, v in sorted(comp.items())))
+
+
+def _default_demo_path() -> Optional[str]:
+    """Path to the bundled C60 demo, for `vimol` with no file argument.
+
+    Only resolvable from a checkout or editable install -- `examples/` sits
+    next to `src/`, not inside the installed package -- so a plain `vimol`
+    still falls back to --help rather than crashing on a real install.
+    """
+    root = os.path.dirname(os.path.dirname(os.path.dirname(os.path.abspath(__file__))))
+    path = os.path.join(root, "examples", "c60.xyz")
+    return path if os.path.exists(path) else None
+
+
+def main(argv: List[str] | None = None) -> int:
+    args = make_parser().parse_args(argv)
+
+    if args.version:
+        from . import __version__
+        print(f"vimol {__version__}")
+        return 0
+    if args.list_formats:
+        print("Supported formats: " + ", ".join(SUPPORTED_EXTENSIONS))
+        return 0
+    if not args.file:
+        args.file = _default_demo_path()
+        if not args.file:
+            make_parser().print_help()
+            return 1
+    if not os.path.exists(args.file):
+        print(f"error: no such file: {args.file}", file=sys.stderr)
+        return 2
+
+    try:
+        mols = load_all(args.file)
+    except Exception as e:  # noqa: BLE001
+        print(f"error: failed to parse {args.file}: {e}", file=sys.stderr)
+        return 3
+    if not mols:
+        print("error: no molecules parsed", file=sys.stderr)
+        return 3
+
+    for m in mols:
+        if not args.no_bonds:
+            ensure_bonds(m, tolerance=args.bond_tolerance)
+
+    idx = max(0, min(args.frame, len(mols) - 1))
+    mol = mols[idx]
+
+    if args.info:
+        _print_info(mol)
+        if len(mols) > 1:
+            print(f"models:  {len(mols)}")
+        return 0
+
+    style = build_style(args)
+    w, h = _parse_size(args.size)
+
+    # -- still render to PNG ---------------------------------------------
+    if args.render:
+        if not w:
+            w, h = 900, 700
+        scene = Scene(mol, w, h, style=style, supersample=max(1, args.supersample),
+                     backend=args.backend)
+        scene.camera.orbit(args.rotate[0], args.rotate[1])
+        scene.to_png(args.render)
+        print(f"wrote {args.render} ({w}x{h})")
+        return 0
+
+    # -- single kitty frame to stdout ------------------------------------
+    if args.kitty:
+        if not w:
+            cols, rows, xpx, ypx = kitty.terminal_size_px(1)
+            cw, ch = kitty.cell_size_px(1)
+            w = int(min(cols, 60) * cw)
+            h = int(min(rows - 2, 30) * ch)
+            w, h = max(w, 200), max(h, 200)
+        scene = Scene(mol, w, h, style=style, supersample=max(1, args.supersample),
+                     backend=args.backend)
+        scene.camera.orbit(args.rotate[0], args.rotate[1])
+        sys.stdout.write("\n")
+        sys.stdout.flush()
+        os.write(1, scene.to_kitty(move_cursor=True))
+        sys.stdout.write("\n")
+        return 0
+
+    # -- interactive viewer ----------------------------------------------
+    if not sys.stdout.isatty():
+        print("error: interactive mode needs a terminal (use --render or --kitty)", file=sys.stderr)
+        return 4
+    if not kitty.supports_kitty():
+        print("warning: this terminal may not support the Kitty graphics protocol.",
+              file=sys.stderr)
+        print("         Set VIMOL_FORCE_KITTY=1 to try anyway, or use --render out.png.",
+              file=sys.stderr)
+        return 5
+
+    # interactive defaults to a terminal-matching transparent background
+    if not args.opaque and args.background is None:
+        style.transparent = True
+
+    from .viewer import Viewer
+    viewer = Viewer(mol, frames=mols, style=style, autospin=args.spin,
+                    backend=args.backend, source_path=args.file, editable=True)
+    viewer.frame_index = idx
+    # apply initial rotation
+    viewer.widget.scene.camera.orbit(args.rotate[0], args.rotate[1])
+    viewer.run()
+    return 0
+
+
+if __name__ == "__main__":
+    raise SystemExit(main())

vimol/arrows.py

@@ -0,0 +1,91 @@
+"""Per-atom vector field -> arrow geometry, shared by both render backends.
+
+Each ``Molecule.vector_fields`` entry is a raw ``(N, 3)`` array of
+view-independent directions/magnitudes; this module turns it into view-space
+arrow segments (a cylindrical shaft + a conical head per arrow) that
+``render.Renderer`` (CPU) and ``gl_adapter`` (GPU) both consume, so the arrow
+geometry itself is defined once. Like ``render.py``, this module only depends
+on ``camera.py``/``molecule.py`` -- no GL imports.
+"""
+from __future__ import annotations
+
+from dataclasses import dataclass
+
+import numpy as np
+
+from .camera import Camera
+from .molecule import Molecule
+
+
+@dataclass
+class ArrowGeometry:
+    """View-space arrow segments, split into shaft cylinders and head cones."""
+    shaft_a: np.ndarray        # (M, 3) view space -- base (at the atom)
+    shaft_b: np.ndarray        # (M, 3) view space -- where the head starts
+    shaft_radius: np.ndarray   # (M,)
+    shaft_color: np.ndarray    # (M, 3)
+    head_base: np.ndarray      # (M, 3) view space -- wide end of the cone
+    head_apex: np.ndarray      # (M, 3) view space -- arrow tip
+    head_radius: np.ndarray    # (M,) -- cone base radius
+    head_color: np.ndarray     # (M, 3)
+
+    @staticmethod
+    def empty() -> "ArrowGeometry":
+        z3 = np.zeros((0, 3), np.float64)
+        z1 = np.zeros((0,), np.float64)
+        return ArrowGeometry(
+            shaft_a=z3.copy(), shaft_b=z3.copy(), shaft_radius=z1.copy(), shaft_color=z3.copy(),
+            head_base=z3.copy(), head_apex=z3.copy(), head_radius=z1.copy(), head_color=z3.copy(),
+        )
+
+
+def build_arrow_geometry(mol: Molecule, camera: Camera, view_pos=None) -> ArrowGeometry:
+    """Convert ``mol.vector_fields`` into view-space arrow geometry under *camera*.
+
+    *view_pos* optionally supplies the atoms' already-computed view-space
+    positions (``camera.view_positions(mol.positions)``); the CPU renderer and
+    GL adapter both have it on hand, so passing it avoids re-running that
+    per-frame transform.
+    """
+    if not mol.vector_fields or mol.n_atoms == 0:
+        return ArrowGeometry.empty()
+
+    view_pos = camera.view_positions(mol.positions) if view_pos is None else np.asarray(view_pos, float)
+    shaft_a, shaft_b, shaft_r, shaft_c = [], [], [], []
+    head_base, head_apex, head_r, head_c = [], [], [], []
+
+    for vf in mol.vector_fields:
+        vectors = np.asarray(vf.vectors, np.float64)
+        if vectors.shape != (mol.n_atoms, 3):
+            continue
+        mag = np.linalg.norm(vectors, axis=1)
+        keep = mag > 1e-9
+        if not keep.any():
+            continue
+
+        view_vecs = camera.view_directions(vectors[keep])
+        base = view_pos[keep]
+        tip = base + view_vecs * vf.scale
+        frac = float(np.clip(vf.head_length_frac, 0.0, 1.0))
+        head_start = base + (tip - base) * (1.0 - frac)
+
+        n = base.shape[0]
+        color = np.asarray(vf.color, np.float64)
+        shaft_a.append(base)
+        shaft_b.append(head_start)
+        shaft_r.append(np.full(n, vf.radius))
+        shaft_c.append(np.tile(color, (n, 1)))
+        head_base.append(head_start)
+        head_apex.append(tip)
+        head_r.append(np.full(n, vf.radius * vf.head_scale))
+        head_c.append(np.tile(color, (n, 1)))
+
+    if not shaft_a:
+        return ArrowGeometry.empty()
+
+    return ArrowGeometry(
+        shaft_a=np.concatenate(shaft_a), shaft_b=np.concatenate(shaft_b),
+        shaft_radius=np.concatenate(shaft_r), shaft_color=np.concatenate(shaft_c),
+        head_base=np.concatenate(head_base), head_apex=np.concatenate(head_apex),
+        head_radius=np.concatenate(head_r), head_color=np.concatenate(head_c),
+    )

vimol/bonds.py

@@ -0,0 +1,79 @@
+"""Distance-based bond perception.
+
+Two atoms are bonded when their separation is below the sum of covalent radii
+plus a tolerance. A uniform grid (spatial hash) keeps this near-linear so it
+scales to large structures instead of being O(N^2).
+"""
+from __future__ import annotations
+
+from typing import List, Tuple
+
+import numpy as np
+
+from .molecule import Molecule
+
+
+def perceive_bonds(mol: Molecule, tolerance: float = 0.45, max_bonds_per_atom: int = 8) -> List[Tuple[int, int, int]]:
+    """Return a bond list inferred from interatomic distances.
+
+    tolerance is added to the sum of covalent radii (angstrom).
+    """
+    n = mol.n_atoms
+    if n < 2:
+        return []
+    pos = mol.positions
+    cov = mol.covalent_radii()
+    max_r = float(cov.max())
+    cell = max(2 * max_r + tolerance, 1.0)
+
+    # Spatial hash: map each atom to an integer grid cell.
+    mins = pos.min(axis=0)
+    ijk = np.floor((pos - mins) / cell).astype(np.int64)
+    buckets: dict = {}
+    for a in range(n):
+        buckets.setdefault(tuple(ijk[a]), []).append(a)
+
+    bonds: List[Tuple[int, int, int]] = []
+    counts = np.zeros(n, dtype=np.int32)
+    seen = set()
+    neigh = [(-1, 0, 1)] * 3
+    from itertools import product
+    offsets = list(product(*neigh))
+
+    for a in range(n):
+        base = ijk[a]
+        cand: List[int] = []
+        for off in offsets:
+            key = (base[0] + off[0], base[1] + off[1], base[2] + off[2])
+            b = buckets.get(key)
+            if b:
+                cand.extend(b)
+        pa = pos[a]
+        ra = cov[a]
+        for b in cand:
+            if b <= a:
+                continue
+            cutoff = ra + cov[b] + tolerance
+            d = pa - pos[b]
+            if d[0] * d[0] + d[1] * d[1] + d[2] * d[2] <= cutoff * cutoff:
+                # Reject absurdly short contacts (overlapping/duplicate atoms).
+                dist2 = d[0] * d[0] + d[1] * d[1] + d[2] * d[2]
+                if dist2 < 0.16:  # < 0.4 A
+                    continue
+                if counts[a] >= max_bonds_per_atom or counts[b] >= max_bonds_per_atom:
+                    continue
+                key = (a, b)
+                if key in seen:
+                    continue
+                seen.add(key)
+                bonds.append((a, b, 1))
+                counts[a] += 1
+                counts[b] += 1
+    return bonds
+
+
+def ensure_bonds(mol: Molecule, tolerance: float = 0.45) -> Molecule:
+    """Populate mol.bonds if empty. Returns the molecule for chaining."""
+    if not mol.bonds and mol.n_atoms > 1:
+        mol.bonds = perceive_bonds(mol, tolerance=tolerance)
+    return mol

vimol/camera.py

@@ -0,0 +1,77 @@
+"""Orthographic trackball camera.
+
+The camera keeps a 3x3 rotation matrix (world -> view), a look-at center, and a
+zoom expressed as pixels-per-angstrom. View space is right-handed with +x
+right, +y up, +z toward the viewer, so the renderer keeps the fragment with the
+largest z (nearest the camera).
+"""
+from __future__ import annotations
+
+import numpy as np
+
+
+def _normalize(v: np.ndarray) -> np.ndarray:
+    n = np.linalg.norm(v)
+    return v / n if n > 1e-12 else v
+
+
+def rotation_from_axis_angle(axis: np.ndarray, angle: float) -> np.ndarray:
+    axis = _normalize(np.asarray(axis, float))
+    x, y, z = axis
+    c = np.cos(angle)
+    s = np.sin(angle)
+    t = 1.0 - c
+    return np.array([
+        [t * x * x + c,     t * x * y - s * z, t * x * z + s * y],
+        [t * x * y + s * z, t * y * y + c,     t * y * z - s * x],
+        [t * x * z - s * y, t * y * z + s * x, t * z * z + c],
+    ])
+
+
+class Camera:
+    def __init__(self, center=None, extent: float = 5.0):
+        self.center = np.zeros(3) if center is None else np.array(center, float)
+        self.rotation = np.eye(3)
+        # zoom is pixels per angstrom; set by fit()
+        self.zoom = 40.0
+        self.extent = float(extent)
+        # pan offset in screen pixels
+        self.pan = np.zeros(2)
+
+    # -- framing ----------------------------------------------------------
+    def fit(self, width: int, height: int, extent: float, margin: float = 1.15) -> None:
+        """Choose a zoom so a sphere of *extent* angstrom fills the viewport."""
+        self.extent = float(extent)
+        half = max(min(width, height) * 0.5, 1.0)
+        self.zoom = half / (extent * margin)
+        self.pan = np.zeros(2)
+
+    # -- interaction ------------------------------------------------------
+    def orbit(self, dx: float, dy: float, speed: float = 0.01) -> None:
+        """Rotate about screen axes. dx/dy are in pixels."""
+        if dx:
+            self.rotation = rotation_from_axis_angle([0, 1, 0], dx * speed) @ self.rotation
+        if dy:
+            self.rotation = rotation_from_axis_angle([1, 0, 0], dy * speed) @ self.rotation
+
+    def roll(self, angle: float) -> None:
+        self.rotation = rotation_from_axis_angle([0, 0, 1], angle) @ self.rotation
+
+    def zoom_by(self, factor: float) -> None:
+        self.zoom = float(np.clip(self.zoom * factor, 1.0, 1e5))
+
+    def pan_by(self, dx: float, dy: float) -> None:
+        self.pan = self.pan + np.array([dx, dy], float)
+
+    def reset(self) -> None:
+        self.rotation = np.eye(3)
+        self.pan = np.zeros(2)
+
+    # -- transforms -------------------------------------------------------
+    def view_positions(self, positions: np.ndarray) -> np.ndarray:
+        """World coordinates -> view-space coordinates (angstrom)."""
+        return (positions - self.center) @ self.rotation.T
+
+    def view_directions(self, directions: np.ndarray) -> np.ndarray:
+        """World -> view space for free vectors (rotation only, no translation)."""
+        return np.asarray(directions, float) @ self.rotation.T