vimol 0.1.0__py3-none-any.whl

vimol/molecule.py

@@ -0,0 +1,159 @@
+"""The core molecular data model.
+
+A :class:`Molecule` stores atoms as parallel numpy arrays (positions, colors,
+radii) plus a bond list. Keeping everything as numpy arrays lets the renderer
+operate on the whole structure vectorized.
+"""
+from __future__ import annotations
+
+from dataclasses import dataclass, field
+from typing import List, Optional, Sequence, Set, Tuple
+
+import numpy as np
+
+from . import elements
+
+
+@dataclass
+class VectorField:
+    """A set of per-atom vectors (e.g. dipoles, forces) drawn as arrows.
+
+    ``vectors`` is atom-aligned with :attr:`Molecule.positions` -- an atom
+    with a near-zero vector simply has no arrow drawn. Color is the only
+    thing distinguishing one field from another (this layer has no notion
+    of what the vectors represent); magnitudes are never normalized, so
+    arrow length is ``|vector| * scale``.
+    """
+    vectors: np.ndarray                              # (N, 3) angstrom, atom-aligned
+    color: Tuple[float, float, float] = (0.15, 0.95, 0.85)
+    scale: float = 1.0                                # arrow length = |vector| * scale
+    radius: float = 0.05                              # shaft radius, angstrom
+    head_scale: float = 2.5                           # head base radius = radius * head_scale
+    head_length_frac: float = 0.25                    # fraction of arrow length that's the head
+
+
+@dataclass
+class Molecule:
+    """A set of atoms and the bonds between them.
+
+    Attributes
+    ----------
+    symbols:   list of element symbols, length N
+    positions: (N, 3) float array, angstrom
+    bonds:     list of (i, j, order) tuples, i < j
+    name:      free-form label
+    """
+
+    symbols: List[str] = field(default_factory=list)
+    positions: np.ndarray = field(default_factory=lambda: np.zeros((0, 3), np.float64))
+    bonds: List[Tuple[int, int, int]] = field(default_factory=list)
+    name: str = ""
+    vector_fields: List[VectorField] = field(default_factory=list)
+    # Explicit bonds (e.g. drawn by option-drag) that survive re-perception
+    # even beyond the auto distance threshold; (i, j, order) with i < j, like
+    # `bonds`. `editor._reperceive` unions these into `bonds` on every rebuild.
+    manual_bonds: List[Tuple[int, int, int]] = field(default_factory=list)
+    # Indices of atoms created by editing since load (or since the last
+    # `editor.cleanup` accepted the relaxed geometry). `editor.py` marks every
+    # atom it creates; `editor.cleanup` reads this to weight its relaxation
+    # and clears it once the geometry is "accepted".
+    new_atoms: Set[int] = field(default_factory=set)
+
+    # -- construction -----------------------------------------------------
+    def add_atom(self, symbol: str, x: float, y: float, z: float) -> int:
+        idx = len(self.symbols)
+        self.symbols.append(elements.normalize_symbol(symbol))
+        self.positions = np.vstack([self.positions, [x, y, z]]) if len(self.positions) else np.array([[x, y, z]], float)
+        return idx
+
+    def add_bond(self, i: int, j: int, order: int = 1) -> None:
+        if i > j:
+            i, j = j, i
+        self.bonds.append((i, j, order))
+
+    def add_vector_field(self, vectors, color=(0.15, 0.95, 0.85), scale: float = 1.0,
+                         radius: float = 0.05, head_scale: float = 2.5,
+                         head_length_frac: float = 0.25) -> VectorField:
+        """Attach a set of per-atom vectors (e.g. dipoles) to be drawn as arrows."""
+        vectors = np.asarray(vectors, np.float64)
+        if vectors.shape != (self.n_atoms, 3):
+            raise ValueError(
+                f"vectors must have shape ({self.n_atoms}, 3), got {vectors.shape}")
+        vf = VectorField(vectors=vectors, color=tuple(color), scale=scale, radius=radius,
+                         head_scale=head_scale, head_length_frac=head_length_frac)
+        self.vector_fields.append(vf)
+        return vf
+
+    # -- derived quantities ----------------------------------------------
+    @property
+    def n_atoms(self) -> int:
+        return len(self.symbols)
+
+    def centroid(self) -> np.ndarray:
+        if self.n_atoms == 0:
+            return np.zeros(3)
+        return self.positions.mean(axis=0)
+
+    def radius_of_gyration_extent(self) -> float:
+        """A robust size estimate: max distance of any atom from the centroid."""
+        if self.n_atoms == 0:
+            return 1.0
+        d = np.linalg.norm(self.positions - self.centroid(), axis=1)
+        return float(max(d.max(), 1e-3))
+
+    def vector_extent(self) -> float:
+        """Max centroid-to-arrow-tip distance across all vector fields, 0 if none.
+
+        Mirrors :meth:`radius_of_gyration_extent` so the camera's auto-fit can
+        pad for long arrows the same way it pads for atom radii.
+        """
+        if not self.vector_fields or self.n_atoms == 0:
+            return 0.0
+        c = self.centroid()
+        best = 0.0
+        for vf in self.vector_fields:
+            # A field can go stale if atoms are added after it was attached;
+            # skip it here exactly as build_arrow_geometry does (rather than
+            # letting the (N+1,3)+(N,3) broadcast raise mid-render).
+            if np.asarray(vf.vectors).shape != (self.n_atoms, 3):
+                continue
+            tips = self.positions + vf.vectors * vf.scale
+            # pad by the cone-head base radius so a fat arrowhead at the edge
+            # of the scene isn't clipped by auto-fit.
+            d = np.linalg.norm(tips - c, axis=1) + vf.radius * vf.head_scale
+            best = max(best, float(d.max()))
+        return best
+
+    def element_colors(self) -> np.ndarray:
+        return np.array([elements.element_color(s) for s in self.symbols], np.float64)
+
+    def covalent_radii(self) -> np.ndarray:
+        return np.array([elements.covalent_radius(s) for s in self.symbols], np.float64)
+
+    def vdw_radii(self) -> np.ndarray:
+        return np.array([elements.vdw_radius(s) for s in self.symbols], np.float64)
+
+    def atomic_numbers(self) -> np.ndarray:
+        return np.array([elements.symbol_to_z(s) for s in self.symbols], np.int32)
+
+    def recenter(self) -> "Molecule":
+        """Translate so the centroid sits at the origin (in place)."""
+        if self.n_atoms:
+            self.positions = self.positions - self.centroid()
+        return self
+
+    def formula(self) -> str:
+        """Hill-system molecular formula, e.g. 'C8H10N4O2'."""
+        from collections import Counter
+        c = Counter(self.symbols)
+        parts = []
+        for el in ("C", "H"):
+            if c.get(el):
+                parts.append(el + (str(c[el]) if c[el] > 1 else ""))
+                del c[el]
+        for el in sorted(c):
+            parts.append(el + (str(c[el]) if c[el] > 1 else ""))
+        return "".join(parts)
+
+    def __repr__(self) -> str:
+        return f"<Molecule {self.name!r} atoms={self.n_atoms} bonds={len(self.bonds)} formula={self.formula()}>"

vimol/parsers/__init__.py

@@ -0,0 +1,88 @@
+"""Molecular file-format parsers and a format-dispatching loader."""
+from __future__ import annotations
+
+import os
+from typing import List
+
+from ..molecule import Molecule
+from . import xyz as _xyz
+from . import pdb as _pdb
+from . import mol as _mol
+
+_EXT_PARSERS = {
+    ".xyz": _xyz.parse,
+    ".pdb": _pdb.parse,
+    ".ent": _pdb.parse,
+    ".mol": _mol.parse,
+    ".sdf": _mol.parse,
+    ".mdl": _mol.parse,
+}
+
+SUPPORTED_EXTENSIONS = tuple(sorted(_EXT_PARSERS))
+
+
+def load(path: str, fmt: str | None = None) -> Molecule:
+    """Load a single molecule from a file, dispatching on extension (or *fmt*)."""
+    mols = load_all(path, fmt=fmt)
+    if not mols:
+        raise ValueError(f"No molecules found in {path!r}")
+    return mols[0]
+
+
+def load_all(path: str, fmt: str | None = None) -> List[Molecule]:
+    """Load every model/frame from a file (e.g. multi-record SDF, NMR PDB)."""
+    if fmt is None:
+        ext = os.path.splitext(path)[1].lower()
+    else:
+        ext = fmt if fmt.startswith(".") else "." + fmt.lower()
+    parser = _EXT_PARSERS.get(ext)
+    if parser is None:
+        raise ValueError(
+            f"Unsupported format {ext!r}. Supported: {', '.join(SUPPORTED_EXTENSIONS)}"
+        )
+    with open(path, "r", errors="replace") as fh:
+        text = fh.read()
+    mols = parser(text)
+    if isinstance(mols, Molecule):
+        mols = [mols]
+    for m in mols:
+        if not m.name:
+            m.name = os.path.splitext(os.path.basename(path))[0]
+    return mols
+
+
+_EXT_WRITERS = {
+    ".xyz": _xyz.write,
+}
+
+SAVE_EXTENSIONS = tuple(sorted(_EXT_WRITERS))
+
+
+def save(mol: Molecule, path: str, fmt: str | None = None) -> None:
+    """Write a molecule to *path*, dispatching on extension (or *fmt*).
+
+    Only formats in :data:`SAVE_EXTENSIONS` can be written (currently XYZ); a
+    path with any other extension raises :class:`ValueError`.
+    """
+    if fmt is None:
+        ext = os.path.splitext(path)[1].lower()
+    else:
+        ext = fmt if fmt.startswith(".") else "." + fmt.lower()
+    writer = _EXT_WRITERS.get(ext)
+    if writer is None:
+        raise ValueError(
+            f"Cannot write {ext!r}. Writable formats: {', '.join(SAVE_EXTENSIONS)}"
+        )
+    writer(mol, path)
+
+
+def loads(text: str, fmt: str) -> Molecule:
+    """Parse a molecule from an in-memory string given an explicit format."""
+    ext = fmt if fmt.startswith(".") else "." + fmt.lower()
+    parser = _EXT_PARSERS.get(ext)
+    if parser is None:
+        raise ValueError(f"Unsupported format {fmt!r}")
+    mols = parser(text)
+    if isinstance(mols, Molecule):
+        return mols
+    return mols[0]

vimol/parsers/mol.py

@@ -0,0 +1,57 @@
+"""MDL MOL (V2000) and SDF parser (SDF = one or more MOL records)."""
+from __future__ import annotations
+
+from typing import List
+
+import numpy as np
+
+from ..molecule import Molecule
+from .. import elements
+
+
+def _parse_one(block: str) -> Molecule | None:
+    lines = block.splitlines()
+    if len(lines) < 4:
+        return None
+    name = lines[0].strip()
+    counts = lines[3]
+    try:
+        n_atoms = int(counts[0:3])
+        n_bonds = int(counts[3:6])
+    except ValueError:
+        return None
+    symbols: List[str] = []
+    coords: List[list] = []
+    for k in range(n_atoms):
+        row = lines[4 + k]
+        x = float(row[0:10])
+        y = float(row[10:20])
+        z = float(row[20:30])
+        sym = row[31:34].strip()
+        symbols.append(elements.normalize_symbol(sym))
+        coords.append([x, y, z])
+    bonds: List[tuple] = []
+    for k in range(n_bonds):
+        row = lines[4 + n_atoms + k]
+        a = int(row[0:3]) - 1
+        b = int(row[3:6]) - 1
+        order = int(row[6:9]) if row[6:9].strip() else 1
+        lo, hi = (a, b) if a < b else (b, a)
+        bonds.append((lo, hi, order))
+    return Molecule(
+        symbols=symbols,
+        positions=np.array(coords, float) if coords else np.zeros((0, 3)),
+        bonds=bonds,
+        name=name[:60],
+    )
+
+
+def parse(text: str) -> List[Molecule]:
+    mols: List[Molecule] = []
+    for block in text.split("$$$$"):
+        if not block.strip():
+            continue
+        m = _parse_one(block)
+        if m is not None and m.n_atoms:
+            mols.append(m)
+    return mols

vimol/parsers/pdb.py

@@ -0,0 +1,97 @@
+"""Minimal but practical PDB parser.
+
+Handles ATOM/HETATM records with fixed-column layout, CONECT bonds, TER, and
+MODEL/ENDMDL multi-model files. Element is taken from columns 77-78 when
+present, otherwise guessed from the atom name.
+"""
+from __future__ import annotations
+
+from typing import List
+
+import numpy as np
+
+from ..molecule import Molecule
+from .. import elements
+
+
+def _guess_element(atom_name: str, col_element: str) -> str:
+    col_element = col_element.strip()
+    if col_element:
+        return elements.normalize_symbol(col_element)
+    name = atom_name.strip()
+    # PDB atom names: element is usually the leading alphabetic chunk; names
+    # starting in column 13 with a digit are hydrogens/remoteness-indexed.
+    if not name:
+        return "X"
+    if name[0].isdigit():
+        name = name[1:]
+    # Two-letter elements only for the known set to avoid e.g. 'CA' (calcium)
+    # vs 'CA' (alpha carbon) confusion -> default to single letter for protein.
+    letters = "".join(c for c in name if c.isalpha())
+    if not letters:
+        return "X"
+    return elements.normalize_symbol(letters[0])
+
+
+def parse(text: str) -> List[Molecule]:
+    models: List[Molecule] = []
+    cur = Molecule()
+    serial_to_index: dict = {}
+    have_atoms = False
+
+    def flush():
+        nonlocal cur, serial_to_index, have_atoms
+        if have_atoms:
+            models.append(cur)
+        cur = Molecule()
+        serial_to_index = {}
+        have_atoms = False
+
+    for line in text.splitlines():
+        rec = line[:6].strip()
+        if rec in ("ATOM", "HETATM"):
+            try:
+                serial = int(line[6:11])
+            except ValueError:
+                serial = len(cur.symbols)
+            atom_name = line[12:16]
+            x = float(line[30:38])
+            y = float(line[38:46])
+            z = float(line[46:54])
+            col_el = line[76:78] if len(line) >= 78 else ""
+            el = _guess_element(atom_name, col_el)
+            idx = len(cur.symbols)
+            cur.symbols.append(el)
+            serial_to_index[serial] = idx
+            cur.positions = (
+                np.vstack([cur.positions, [x, y, z]]) if len(cur.positions) else np.array([[x, y, z]], float)
+            )
+            have_atoms = True
+        elif rec == "CONECT":
+            try:
+                a = serial_to_index.get(int(line[6:11]))
+            except ValueError:
+                continue
+            if a is None:
+                continue
+            for c in (11, 16, 21, 26):
+                frag = line[c:c + 5].strip()
+                if not frag:
+                    continue
+                try:
+                    b = serial_to_index.get(int(frag))
+                except ValueError:
+                    continue
+                if b is not None and a != b:
+                    lo, hi = (a, b) if a < b else (b, a)
+                    if (lo, hi, 1) not in cur.bonds:
+                        cur.bonds.append((lo, hi, 1))
+        elif rec == "ENDMDL":
+            flush()
+        elif rec == "END":
+            flush()
+
+    flush()
+    if not models:
+        models.append(cur)
+    return models

vimol/parsers/xyz.py

@@ -0,0 +1,68 @@
+"""XYZ / extended-XYZ parser (supports multi-frame trajectories)."""
+from __future__ import annotations
+
+from typing import List
+
+import numpy as np
+
+from ..molecule import Molecule
+from .. import elements
+
+
+def parse(text: str) -> List[Molecule]:
+    lines = text.splitlines()
+    mols: List[Molecule] = []
+    i = 0
+    n = len(lines)
+    while i < n:
+        # Skip blank lines between frames.
+        while i < n and not lines[i].strip():
+            i += 1
+        if i >= n:
+            break
+        try:
+            count = int(lines[i].strip().split()[0])
+        except (ValueError, IndexError):
+            break
+        comment = lines[i + 1] if i + 1 < n else ""
+        symbols: List[str] = []
+        coords: List[list] = []
+        for k in range(count):
+            row = i + 2 + k
+            if row >= n:
+                break
+            parts = lines[row].split()
+            if len(parts) < 4:
+                continue
+            sym = parts[0]
+            # Extended XYZ may use atomic numbers instead of symbols.
+            if sym.isdigit():
+                sym = elements.SYMBOLS[int(sym)] if 0 < int(sym) < len(elements.SYMBOLS) else "X"
+            symbols.append(elements.normalize_symbol(sym))
+            coords.append([float(parts[1]), float(parts[2]), float(parts[3])])
+        mol = Molecule(
+            symbols=symbols,
+            positions=np.array(coords, float) if coords else np.zeros((0, 3)),
+            name=comment.strip()[:60],
+        )
+        mols.append(mol)
+        i += 2 + count
+    return mols
+
+
+def dumps(mol: Molecule) -> str:
+    """Serialize a molecule to XYZ text (single frame).
+
+    The comment line carries the molecule name. Coordinates are written with
+    six decimals, which round-trips through :func:`parse` without drift.
+    """
+    lines = [str(mol.n_atoms), mol.name or ""]
+    for sym, p in zip(mol.symbols, mol.positions):
+        lines.append(f"{sym:<2s} {p[0]:>14.6f} {p[1]:>14.6f} {p[2]:>14.6f}")
+    return "\n".join(lines) + "\n"
+
+
+def write(mol: Molecule, path: str) -> None:
+    """Write a molecule to *path* as XYZ."""
+    with open(path, "w") as fh:
+        fh.write(dumps(mol))

vimol/periodic_table.py

@@ -0,0 +1,97 @@
+"""Standard 18-column periodic-table layout, derived from vimol.elements.
+
+Grid coordinates here drive the interactive element picker in the viewer:
+cursor navigation, click hit-testing, and rendering all key off :data:`GRID`.
+The layout follows the common convention of pulling the lanthanides and
+actinides out of the main table into two extra rows below, leaving a
+non-selectable placeholder cell in the gap each leaves behind.
+
+Element placement is computed from atomic-number ranges against
+:data:`vimol.elements.SYMBOLS` rather than typed out by hand, so it can't
+drift out of sync with that table.
+"""
+from __future__ import annotations
+
+from dataclasses import dataclass
+from typing import Dict, List, Optional, Tuple
+
+from . import elements
+
+N_COLS = 18
+
+
+@dataclass(frozen=True)
+class Cell:
+    row: int
+    col: int
+    symbol: Optional[str]           # None for a non-selectable placeholder
+    z: Optional[int] = None
+    label: Optional[str] = None     # display text; defaults to the symbol
+    note: Optional[str] = None      # shown in the info line for a placeholder
+    jump_row: Optional[int] = None  # row Enter/click jumps to, for a placeholder
+
+    @property
+    def text(self) -> str:
+        if self.label is not None:
+            return self.label
+        return self.symbol or ""
+
+
+def _add_row(rows: List[List[Optional[Cell]]], row_idx: int, placements) -> None:
+    """placements: (col, z) for an element, or (col, None, label, note, jump_row)
+    for a non-selectable placeholder."""
+    cells: List[Optional[Cell]] = [None] * N_COLS
+    for item in placements:
+        if len(item) == 2:
+            col, z = item
+            cells[col] = Cell(row_idx, col, elements.SYMBOLS[z], z)
+        else:
+            col, _, label, note, jump_row = item
+            cells[col] = Cell(row_idx, col, None, None, label, note, jump_row)
+    rows.append(cells)
+
+
+def _build() -> List[List[Optional[Cell]]]:
+    rows: List[List[Optional[Cell]]] = []
+    # Period 1: H(1) .. He(2)
+    _add_row(rows, 0, [(0, 1), (17, 2)])
+    # Period 2: Li(3) Be(4) ... B(5)..Ne(10)
+    _add_row(rows, 1, [(0, 3), (1, 4)] + [(12 + i, 5 + i) for i in range(6)])
+    # Period 3: Na(11) Mg(12) ... Al(13)..Ar(18)
+    _add_row(rows, 2, [(0, 11), (1, 12)] + [(12 + i, 13 + i) for i in range(6)])
+    # Period 4: K(19)..Kr(36), 18 consecutive
+    _add_row(rows, 3, [(c, 19 + c) for c in range(18)])
+    # Period 5: Rb(37)..Xe(54), 18 consecutive
+    _add_row(rows, 4, [(c, 37 + c) for c in range(18)])
+    # Period 6: Cs(55) Ba(56) [gap: La-Lu below] Hf(72)..Rn(86)
+    _add_row(rows, 5, [(0, 55), (1, 56), (2, None, "**", "Lanthanides — see the row below", 7)]
+             + [(3 + c, 72 + c) for c in range(15)])
+    # Period 7: Fr(87) Ra(88) [gap: Ac-Lr below] Rf(104)..Og(118)
+    _add_row(rows, 6, [(0, 87), (1, 88), (2, None, "**", "Actinides — see the row below", 8)]
+             + [(3 + c, 104 + c) for c in range(15)])
+    # Lanthanides: La(57)..Lu(71), indented under the gap above
+    _add_row(rows, 7, [(3 + c, 57 + c) for c in range(15)])
+    # Actinides: Ac(89)..Lr(103)
+    _add_row(rows, 8, [(3 + c, 89 + c) for c in range(15)])
+    return rows
+
+
+GRID: List[List[Optional[Cell]]] = _build()
+
+POSITIONS: Dict[str, Tuple[int, int]] = {
+    cell.symbol: (r, c)
+    for r, row in enumerate(GRID)
+    for c, cell in enumerate(row)
+    if cell is not None and cell.symbol is not None
+}
+
+
+def cell_at(row: int, col: int) -> Optional[Cell]:
+    if 0 <= row < len(GRID) and 0 <= col < N_COLS:
+        return GRID[row][col]
+    return None
+
+
+def position_of(symbol: str) -> Tuple[int, int]:
+    """Grid (row, col) for *symbol*, defaulting to Carbon if not placed."""
+    return POSITIONS.get(elements.normalize_symbol(symbol), POSITIONS["C"])