molscope 0.6.0__py3-none-any.whl

molscope/molecule.py

@@ -0,0 +1,502 @@
+"""The :class:`Molecule` value type and its geometric operations.
+
+Coordinates are held as an ``(N, 3)`` numpy array. Optional per-atom metadata
+(atom name, residue name, residue id, chain) travels alongside, enabling
+selections such as ``mol.backbone()`` or ``mol.select(chain="A")``.
+
+Transformations return a new ``Molecule`` rather than mutating in place, so
+chains like ``mol.centered().rotate("z", 90)`` read top to bottom and never
+alias state.
+"""
+
+from __future__ import annotations
+
+from dataclasses import dataclass, field, replace
+from typing import Any, Optional
+
+import numpy as np
+
+from . import elements
+
+# Above this size the dense O(n^2) bond search is refused; install scipy for the
+# KD-tree path (pip install 'molscope[fast]') to handle larger structures.
+_DENSE_BOND_LIMIT = 8000
+
+_BACKBONE_ATOMS = ("N", "CA", "C", "O")
+
+
+@dataclass(frozen=True, eq=False)
+class Molecule:
+    coords: np.ndarray
+    elements: list[str] = field(default_factory=list)
+    name: str = ""
+    # Optional per-atom metadata; empty when the source format carries none.
+    atom_names: list[str] = field(default_factory=list)
+    resnames: list[str] = field(default_factory=list)
+    resids: np.ndarray = field(default_factory=lambda: np.empty(0, dtype=int))
+    chains: list[str] = field(default_factory=list)
+    # Optional explicit bonds as an (E, 2) index array. When set, bonds() returns
+    # these instead of inferring from geometry (used by coarse-graining).
+    bond_index: Optional[np.ndarray] = None
+    _mapping_report: Optional[Any] = field(default=None, repr=False, compare=False)
+
+    def __post_init__(self):
+        coords = np.asarray(self.coords, dtype=float).reshape(-1, 3)
+        object.__setattr__(self, "coords", coords)
+        object.__setattr__(self, "resids", np.asarray(self.resids, dtype=int))
+        if self.bond_index is not None:
+            object.__setattr__(
+                self, "bond_index", np.asarray(self.bond_index, dtype=int).reshape(-1, 2)
+            )
+        if not self.elements:
+            object.__setattr__(self, "elements", [""] * len(coords))
+        for name in ("elements", "atom_names", "resnames", "chains"):
+            seq = getattr(self, name)
+            if seq and len(seq) != len(coords):
+                raise ValueError(f"{len(seq)} {name} for {len(coords)} coordinates")
+        if len(self.resids) and len(self.resids) != len(coords):
+            raise ValueError(f"{len(self.resids)} resids for {len(coords)} coordinates")
+
+    def __len__(self) -> int:
+        return len(self.coords)
+
+    def __eq__(self, other) -> bool:
+        # Auto-generated dataclass __eq__ can't compare the numpy field; do it
+        # explicitly. coords are mutable in place, so instances stay unhashable.
+        if not isinstance(other, Molecule):
+            return NotImplemented
+        return (
+            self.name == other.name
+            and self.elements == other.elements
+            and np.array_equal(self.coords, other.coords)
+        )
+
+    __hash__ = None
+
+    def __getitem__(self, selector) -> Molecule:
+        """``mol[mask]`` / ``mol[indices]`` -> a subset molecule (see :meth:`take`)."""
+        return self.take(selector)
+
+    @property
+    def has_topology(self) -> bool:
+        """True if per-atom names/residues/chains were parsed."""
+        return bool(self.atom_names)
+
+    # -- selection ----------------------------------------------------------
+
+    def take(self, selector) -> Molecule:
+        """Return the subset given by a boolean mask or an array of indices."""
+        idx = np.arange(len(self))[np.asarray(selector)]
+
+        def sub(seq):
+            return [seq[i] for i in idx] if seq else []
+
+        return replace(
+            self,
+            coords=self.coords[idx],
+            elements=sub(self.elements),
+            atom_names=sub(self.atom_names),
+            resnames=sub(self.resnames),
+            resids=self.resids[idx] if len(self.resids) else self.resids,
+            chains=sub(self.chains),
+            bond_index=self._subset_bonds(idx),
+            _mapping_report=None,
+        )
+
+    def _subset_bonds(self, idx):
+        """Restrict explicit bonds to a kept-atom index set and renumber them."""
+        if self.bond_index is None:
+            return None
+        remap = {old: new for new, old in enumerate(idx)}
+        kept = [
+            (remap[i], remap[j]) for i, j in self.bond_index
+            if i in remap and j in remap
+        ]
+        return np.array(kept, dtype=int).reshape(-1, 2)
+
+    def select(
+        self,
+        element=None,
+        chain=None,
+        resname=None,
+        atom_name=None,
+        resid=None,
+    ) -> Molecule:
+        """Return the atoms matching every supplied criterion.
+
+        Each of ``element``/``chain``/``resname``/``atom_name`` accepts a single
+        value or a collection. ``resid`` accepts an int, a collection of ints,
+        or a ``(low, high)`` inclusive range. Selecting on metadata the molecule
+        lacks raises ``ValueError``.
+        """
+        mask = np.ones(len(self), dtype=bool)
+        mask &= self._field_mask(self.elements, element, "element", upper=True)
+        mask &= self._field_mask(self.chains, chain, "chain")
+        mask &= self._field_mask(self.resnames, resname, "residue", upper=True)
+        mask &= self._field_mask(self.atom_names, atom_name, "atom name", upper=True)
+        if resid is not None:
+            mask &= self._resid_mask(resid)
+        return self.take(mask)
+
+    def backbone(self) -> Molecule:
+        """Protein backbone atoms (N, CA, C, O)."""
+        return self.select(atom_name=_BACKBONE_ATOMS)
+
+    def alpha_carbons(self) -> Molecule:
+        """Alpha-carbon (CA) atoms, the usual basis for protein RMSD."""
+        return self.select(atom_name="CA")
+
+    def _field_mask(self, values, criterion, label, upper=False):
+        if criterion is None:
+            return np.ones(len(self), dtype=bool)
+        if not values:
+            raise ValueError(f"no {label} information in this molecule")
+        wanted = {criterion} if isinstance(criterion, str) else set(criterion)
+        if upper:
+            wanted = {w.upper() for w in wanted}
+            return np.array([v.upper() in wanted for v in values], dtype=bool)
+        return np.array([v in wanted for v in values], dtype=bool)
+
+    def _resid_mask(self, resid):
+        if len(self.resids) == 0:
+            raise ValueError("no residue-id information in this molecule")
+        if isinstance(resid, tuple) and len(resid) == 2:
+            low, high = resid
+            return (self.resids >= low) & (self.resids <= high)
+        wanted = [resid] if isinstance(resid, int) else list(resid)
+        return np.isin(self.resids, wanted)
+
+    def residue_groups(self):
+        """Yield ``(atom_indices, resname, resid, chain)`` per residue, in order.
+
+        Residues are runs of atoms sharing ``(chain, resid)``. Yields nothing if
+        the molecule has no residue information.
+        """
+        n = len(self)
+        if len(self.resids) == 0 or n == 0:
+            return
+        chains = self.chains or [""] * n
+        resnames = self.resnames or [""] * n
+        resids = self.resids
+        start = 0
+        for i in range(1, n + 1):
+            if i == n or chains[i] != chains[i - 1] or resids[i] != resids[i - 1]:
+                yield (
+                    list(range(start, i)), resnames[start],
+                    int(resids[start]), chains[start],
+                )
+                start = i
+
+    # -- geometry -----------------------------------------------------------
+
+    @property
+    def masses(self) -> np.ndarray:
+        """Per-atom atomic weights (g/mol)."""
+        return np.array([elements.mass(e) for e in self.elements])
+
+    @property
+    def centroid(self) -> np.ndarray:
+        """Geometric centre (mean of all atom positions)."""
+        return self.coords.mean(axis=0)
+
+    @property
+    def center_of_mass(self) -> np.ndarray:
+        """Mass-weighted centre of the molecule."""
+        m = self.masses
+        return (m[:, None] * self.coords).sum(axis=0) / m.sum()
+
+    @property
+    def radius_of_gyration(self) -> float:
+        """Mass-weighted radius of gyration (angstrom)."""
+        m = self.masses
+        d2 = ((self.coords - self.center_of_mass) ** 2).sum(axis=1)
+        return float(np.sqrt((m * d2).sum() / m.sum()))
+
+    @property
+    def dimensions(self) -> np.ndarray:
+        """Axis-aligned bounding-box size (dx, dy, dz) in angstrom."""
+        return self.coords.max(axis=0) - self.coords.min(axis=0)
+
+    @property
+    def formula(self) -> str:
+        """Hill-order molecular formula, e.g. ``"C6 H12 O6"``."""
+        from collections import Counter
+
+        counts = Counter(e.capitalize() for e in self.elements if e)
+        if not counts:
+            return ""
+        ordered = []
+        for sym in ("C", "H"):
+            if sym in counts:
+                ordered.append((sym, counts.pop(sym)))
+        ordered += sorted(counts.items())
+        return " ".join(f"{s}{n}" if n > 1 else s for s, n in ordered)
+
+    # -- transforms (each returns a new Molecule) ---------------------------
+
+    def translate(self, vector) -> Molecule:
+        """Return a copy shifted by ``vector`` (dx, dy, dz)."""
+        return replace(self, coords=self.coords + np.asarray(vector, dtype=float))
+
+    def centered(self, weighted: bool = False) -> Molecule:
+        """Return a copy with its centre at the origin.
+
+        By default the geometric centroid is used; pass ``weighted=True`` to
+        centre on the mass-weighted centre of mass.
+        """
+        origin = self.center_of_mass if weighted else self.centroid
+        return replace(self, coords=self.coords - origin)
+
+    def rotate(self, axis, angle_deg: float) -> Molecule:
+        """Return a copy rotated ``angle_deg`` degrees about ``axis``.
+
+        ``axis`` may be ``"x"``, ``"y"``, ``"z"`` or any 3-vector. Rotation is
+        about the centroid so the molecule spins in place.
+        """
+        vec = {
+            "x": (1.0, 0.0, 0.0),
+            "y": (0.0, 1.0, 0.0),
+            "z": (0.0, 0.0, 1.0),
+        }.get(axis, axis)
+        rot = _rotation_matrix(np.asarray(vec, dtype=float), np.radians(angle_deg))
+        center = self.centroid
+        rotated = (self.coords - center) @ rot.T + center
+        return replace(self, coords=rotated)
+
+    def superpose(self, reference: Molecule) -> Molecule:
+        """Return a copy optimally rotated/translated onto ``reference``.
+
+        Uses the Kabsch algorithm (least-squares rigid-body fit). Requires the
+        same number of atoms, matched by index.
+        """
+        if len(self) != len(reference):
+            raise ValueError(f"atom count mismatch: {len(self)} vs {len(reference)}")
+        p = self.coords - self.centroid
+        q = reference.coords - reference.centroid
+        u, _, vt = np.linalg.svd(p.T @ q)
+        d = np.sign(np.linalg.det(vt.T @ u.T))
+        rot = vt.T @ np.diag([1.0, 1.0, d]) @ u.T
+        aligned = p @ rot.T + reference.centroid
+        return replace(self, coords=aligned)
+
+    # -- measurements & analysis -------------------------------------------
+
+    def distance(self, i: int, j: int) -> float:
+        """Distance between atoms ``i`` and ``j`` (angstrom)."""
+        return float(np.linalg.norm(self.coords[i] - self.coords[j]))
+
+    def angle(self, i: int, j: int, k: int) -> float:
+        """Angle in degrees at atom ``j`` formed by ``i``-``j``-``k``."""
+        a = self.coords[i] - self.coords[j]
+        b = self.coords[k] - self.coords[j]
+        cos = np.dot(a, b) / (np.linalg.norm(a) * np.linalg.norm(b))
+        return float(np.degrees(np.arccos(np.clip(cos, -1.0, 1.0))))
+
+    def dihedral(self, a: int, b: int, c: int, d: int) -> float:
+        """Dihedral (torsion) angle in degrees about the ``b``-``c`` bond."""
+        v0 = self.coords[a] - self.coords[b]
+        v1 = self.coords[c] - self.coords[b]
+        v2 = self.coords[d] - self.coords[c]
+        v1 = v1 / np.linalg.norm(v1)
+        v = v0 - np.dot(v0, v1) * v1
+        w = v2 - np.dot(v2, v1) * v1
+        x = np.dot(v, w)
+        y = np.dot(np.cross(v1, v), w)
+        return float(np.degrees(np.arctan2(y, x)))
+
+    def distance_matrix(self) -> np.ndarray:
+        """Full ``(N, N)`` pairwise distance matrix (angstrom)."""
+        deltas = self.coords[:, None, :] - self.coords[None, :, :]
+        return np.sqrt((deltas ** 2).sum(axis=-1))
+
+    def contacts(self, cutoff: float = 5.0) -> np.ndarray:
+        """Atom index pairs ``(i, j)`` closer than ``cutoff`` angstrom."""
+        n = len(self)
+        if n < 2:
+            return np.empty((0, 2), dtype=int)
+        try:
+            from scipy.spatial import cKDTree
+
+            return cKDTree(self.coords).query_pairs(cutoff, output_type="ndarray")
+        except ImportError:
+            dist = self.distance_matrix()
+            i, j = np.where(np.triu(dist < cutoff, k=1))
+            return np.stack([i, j], axis=1)
+
+    def contact_map(self, cutoff: float = 8.0, level: str = "residue", method: str = "ca"):
+        """Build a contact map. See :func:`molscope.contactmap.contact_map`.
+
+        ``level`` is ``"atom"`` or ``"residue"``; for residue level ``method`` is
+        ``"ca"`` (CA-CA distance), ``"com"`` (centre of mass) or ``"min"``
+        (closest inter-residue atom). Returns a :class:`ContactMap`.
+        """
+        from .contactmap import contact_map
+
+        return contact_map(self, cutoff=cutoff, level=level, method=method)
+
+    def plot_contact_map(self, cutoff: float = 8.0, level: str = "residue",
+                         method: str = "ca", **kwargs):
+        """Shortcut for ``self.contact_map(...).plot()``."""
+        return self.contact_map(cutoff, level, method).plot(**kwargs)
+
+    def rmsd(self, other: Molecule, align: bool = False) -> float:
+        """Root-mean-square deviation from ``other`` (matched by index).
+
+        With ``align=True`` the molecules are Kabsch-superposed first, giving
+        the minimum RMSD over all rigid-body orientations.
+        """
+        if len(self) != len(other):
+            raise ValueError(f"atom count mismatch: {len(self)} vs {len(other)}")
+        a = self.superpose(other).coords if align else self.coords
+        return float(np.sqrt(((a - other.coords) ** 2).sum() / len(self)))
+
+    def bonds(self, tolerance: float = 1.2) -> np.ndarray:
+        """Infer bonds as index pairs ``(i, j)``.
+
+        Two atoms bond when their separation is within ``tolerance`` times the
+        sum of their covalent radii. Returns an ``(M, 2)`` int array.
+
+        Uses ``scipy.spatial.cKDTree`` when available (scales to large
+        structures); otherwise falls back to a dense search that is refused
+        above ``_DENSE_BOND_LIMIT`` atoms. If the molecule carries explicit
+        bonds (e.g. a coarse-grained model), those are returned directly.
+        """
+        if self.bond_index is not None:
+            return self.bond_index
+        n = len(self.coords)
+        if n < 2:
+            return np.empty((0, 2), dtype=int)
+        radii = np.array([elements.covalent_radius(e) for e in self.elements])
+
+        try:
+            from scipy.spatial import cKDTree
+        except ImportError:
+            cKDTree = None
+
+        if cKDTree is not None:
+            tree = cKDTree(self.coords)
+            cand = tree.query_pairs(tolerance * 2 * radii.max(), output_type="ndarray")
+            if len(cand) == 0:
+                return np.empty((0, 2), dtype=int)
+            i, j = cand[:, 0], cand[:, 1]
+        else:
+            if n > _DENSE_BOND_LIMIT:
+                raise ValueError(
+                    f"{n} atoms exceeds the dense bond limit ({_DENSE_BOND_LIMIT}); "
+                    "install scipy (pip install 'molscope[fast]') for large "
+                    "structures."
+                )
+            i, j = np.triu_indices(n, k=1)
+
+        dist = np.linalg.norm(self.coords[i] - self.coords[j], axis=1)
+        cutoff = tolerance * (radii[i] + radii[j])
+        keep = dist < cutoff
+        return np.stack([i[keep], j[keep]], axis=1)
+
+    def summary(self) -> str:
+        """One-line human-readable description of the molecule."""
+        parts = [f"{self.name or 'molecule'}: {len(self)} atoms"]
+        if self.formula:
+            parts.append(f"formula {self.formula}")
+        if self.chains:
+            parts.append(f"chains {','.join(sorted(set(self.chains)))}")
+        dx, dy, dz = self.dimensions
+        parts.append(f"size {dx:.1f}x{dy:.1f}x{dz:.1f} A")
+        return " | ".join(parts)
+
+    # -- coarse-graining ----------------------------------------------------
+
+    def coarse_grain(self, mapping="residue_com", weighted: bool = True,
+                     bonds=None, return_report: bool = False):
+        """Map this structure onto CG beads. See :mod:`molscope.coarsegrain`.
+
+        ``mapping`` is ``"residue_com"``, ``"residue_centroid"``, ``"martini"``,
+        a ``{resname: {bead: [atom_names]}}`` dict (by residue), or a
+        ``{bead: [atom_indices]}`` dict (by index, works on any structure).
+        ``bonds`` optionally defines the bead network as pairs of bead indices,
+        or bead names when those names are unique. Repeated residue bead names
+        such as ``BB``/``SC`` are ambiguous; use indices for those. Returns a
+        new ``Molecule`` of beads with CG bonds attached, or ``(molecule,
+        report)`` when ``return_report=True``.
+        """
+        from .coarsegrain import coarse_grain
+
+        return coarse_grain(
+            self, mapping=mapping, weighted=weighted, bonds=bonds,
+            return_report=return_report,
+        )
+
+    def mapping_report(self) -> str:
+        """Explain how this coarse-grained molecule was mapped."""
+        if self._mapping_report is None:
+            raise ValueError("no coarse-graining report is available for this molecule")
+        return self._mapping_report.format()
+
+    # -- ML descriptors -----------------------------------------------------
+
+    def descriptors(self, **kwargs) -> dict:
+        """Return fixed-size molecular descriptors for quick ML features."""
+        from .descriptors import descriptors
+
+        return descriptors(self, **kwargs)
+
+    # -- graph export -------------------------------------------------------
+
+    def to_graph(self, tolerance: float = 1.2, bonds=None):
+        """Build a :class:`molscope.graph.MolecularGraph` from this molecule.
+
+        Bonds are inferred from covalent radii (see :meth:`bonds`) unless an
+        explicit ``(E, 2)`` array of index pairs is passed. Node and edge
+        attributes (element, residue, chain, distance, ...) are carried along.
+        """
+        from .graph import MolecularGraph
+
+        edges = self.bonds(tolerance) if bonds is None else np.asarray(bonds, dtype=int)
+        edges = edges.reshape(-1, 2)
+        if len(edges):
+            dist = np.linalg.norm(self.coords[edges[:, 0]] - self.coords[edges[:, 1]], axis=1)
+        else:
+            dist = np.empty(0, dtype=float)
+        return MolecularGraph(
+            coords=self.coords, elements=self.elements, edges=edges,
+            edge_distances=dist, edge_types=np.ones(len(edges)),
+            atom_names=self.atom_names, resnames=self.resnames,
+            resids=self.resids, chains=self.chains, name=self.name,
+        )
+
+    def to_networkx(self, **kwargs):
+        """Shortcut for ``self.to_graph(...).to_networkx()``."""
+        return self.to_graph(**kwargs).to_networkx()
+
+    def to_pyg_data(self, **kwargs):
+        """Shortcut for ``self.to_graph(...).to_pyg_data()`` (PyTorch Geometric)."""
+        return self.to_graph(**kwargs).to_pyg_data()
+
+    def to_dgl_graph(self, **kwargs):
+        """Shortcut for ``self.to_graph(...).to_dgl_graph()`` (DGL)."""
+        return self.to_graph(**kwargs).to_dgl_graph()
+
+    def plot(self, **kwargs):
+        """Render the molecule in 3D. See :func:`molscope.plotting.plot`."""
+        from .plotting import plot
+
+        return plot(self, **kwargs)
+
+    def view(self, **kwargs):
+        """Interactive py3Dmol viewer. See :func:`molscope.plotting.view`."""
+        from .plotting import view
+
+        return view(self, **kwargs)
+
+
+def _rotation_matrix(axis: np.ndarray, angle: float) -> np.ndarray:
+    """Rodrigues rotation matrix for ``angle`` radians about ``axis``."""
+    axis = axis / np.linalg.norm(axis)
+    x, y, z = axis
+    c, s = np.cos(angle), np.sin(angle)
+    C = 1 - c
+    return np.array([
+        [c + x * x * C, x * y * C - z * s, x * z * C + y * s],
+        [y * x * C + z * s, c + y * y * C, y * z * C - x * s],
+        [z * x * C - y * s, z * y * C + x * s, c + z * z * C],
+    ])

molscope/plotting.py

@@ -0,0 +1,191 @@
+"""3D visualization of molecules: matplotlib, py3Dmol, and GIF export."""
+
+from __future__ import annotations
+
+import itertools
+import warnings
+from typing import Optional
+
+import numpy as np
+
+from . import elements
+
+
+def plot(
+    molecule,
+    show_bonds: Optional[bool] = None,
+    bond_tolerance: float = 1.2,
+    color_by: str = "element",
+    scale: float = 60.0,
+    ax=None,
+    show: bool = True,
+):
+    """Scatter-plot atoms in 3D with an equal aspect ratio.
+
+    ``color_by`` selects the colouring: ``"element"`` (CPK), ``"chain"`` or
+    ``"residue"`` (categorical palette). Atom sizes scale with covalent radius.
+    Bonds are drawn when ``show_bonds`` is true, or, when ``None``, automatically
+    for molecules small enough to infer bonds cheaply. Returns the ``Axes3D``;
+    pass ``show=False`` to suppress ``plt.show()``.
+    """
+    import matplotlib.pyplot as plt  # imported lazily so the core has no GUI dep
+
+    coords = molecule.coords
+    if ax is None:
+        fig = plt.figure()
+        ax = fig.add_subplot(1, 1, 1, projection="3d")
+
+    colors = _colors(molecule, color_by)
+    sizes = np.array([elements.covalent_radius(e) for e in molecule.elements]) * scale
+    ax.scatter(coords[:, 0], coords[:, 1], coords[:, 2], c=colors, s=sizes, depthshade=True)
+
+    if show_bonds is None:
+        show_bonds = len(molecule) <= 2000
+    if show_bonds and len(molecule) > 1:
+        try:
+            for i, j in molecule.bonds(tolerance=bond_tolerance):
+                seg = coords[[i, j]]
+                ax.plot(seg[:, 0], seg[:, 1], seg[:, 2], color="0.5", linewidth=1.0)
+        except ValueError as exc:
+            warnings.warn(f"skipping bonds: {exc}", stacklevel=2)
+
+    ax.set_xlabel("X")
+    ax.set_ylabel("Y")
+    ax.set_zlabel("Z")
+    if molecule.name:
+        ax.set_title(molecule.name)
+    _set_equal_aspect(ax, coords)
+
+    if show:
+        plt.show()
+    return ax
+
+
+def view(molecule, style: str = "stick", width: int = 480, height: int = 360):
+    """Return an interactive py3Dmol viewer (for Jupyter notebooks).
+
+    Requires py3Dmol (``pip install py3Dmol``). ``style`` is any py3Dmol style
+    name such as ``"stick"``, ``"sphere"``, ``"line"`` or ``"cartoon"``.
+    """
+    try:
+        import py3Dmol
+    except ImportError as exc:  # pragma: no cover - exercised only without py3Dmol
+        raise ImportError(
+            "view() needs py3Dmol; install it with: pip install py3Dmol"
+        ) from exc
+    from .io import _molecule_to_pdb_string
+
+    viewer = py3Dmol.view(width=width, height=height)
+    viewer.addModel(_molecule_to_pdb_string(molecule), "pdb")
+    viewer.setStyle({style: {"colorscheme": "default"}})
+    viewer.zoomTo()
+    return viewer
+
+
+def spin_gif(molecule, path: str, frames: int = 36, fps: int = 15, **plot_kwargs):
+    """Render a spinning 3D view and save it as an animated GIF.
+
+    Rotates a full turn about the vertical axis over ``frames`` steps. Requires
+    Pillow (already a matplotlib dependency).
+    """
+    import matplotlib.pyplot as plt
+    from matplotlib import animation
+
+    fig = plt.figure()
+    ax = fig.add_subplot(1, 1, 1, projection="3d")
+    plot(molecule, ax=ax, show=False, **plot_kwargs)
+
+    def update(i):
+        ax.view_init(elev=20, azim=i * 360 / frames)
+        return ()
+
+    anim = animation.FuncAnimation(fig, update, frames=frames, blit=False)
+    anim.save(path, writer=animation.PillowWriter(fps=fps))
+    plt.close(fig)
+    return path
+
+
+def plot_contact_map(contact_map, ax=None, cmap=None, show: bool = True):
+    """Draw a :class:`~molscope.contactmap.ContactMap` as a heatmap.
+
+    Booleans render as a binary map; ensemble frequencies render with a colour
+    scale and a colourbar. Returns the matplotlib ``Axes``.
+    """
+    import matplotlib.pyplot as plt
+
+    mat = contact_map.matrix
+    freq = contact_map.is_frequency
+    if ax is None:
+        _, ax = plt.subplots(figsize=(5, 4))
+    im = ax.imshow(
+        mat, origin="lower", interpolation="nearest", vmin=0, vmax=1,
+        cmap=cmap or ("viridis" if freq else "Greys"),
+    )
+
+    unit = "residue" if contact_map.level == "residue" else "atom"
+    ax.set_xlabel(f"{unit} index")
+    ax.set_ylabel(f"{unit} index")
+    label = "contact frequency" if freq else f"contact (< {contact_map.cutoff} Å)"
+    ax.figure.colorbar(im, ax=ax, label=label, fraction=0.046, pad=0.04)
+    ax.set_title(f"{unit} contact map ({contact_map.cutoff} Å)")
+
+    if show:
+        plt.show()
+    return ax
+
+
+def plot_rmsd_heatmap(matrix, order=None, ax=None, cmap="viridis", show: bool = True):
+    """Draw a pairwise-RMSD matrix as a heatmap (angstrom).
+
+    Pass ``order`` (e.g. ``clustering.order``) to reorder rows/columns so
+    clusters appear as blocks along the diagonal. Returns the matplotlib ``Axes``.
+    """
+    import matplotlib.pyplot as plt
+
+    matrix = np.asarray(matrix, dtype=float)
+    if order is not None:
+        order = np.asarray(order)
+        matrix = matrix[np.ix_(order, order)]
+    if ax is None:
+        _, ax = plt.subplots(figsize=(5, 4))
+    im = ax.imshow(matrix, origin="lower", interpolation="nearest", cmap=cmap)
+    ax.set_xlabel("model")
+    ax.set_ylabel("model")
+    ax.figure.colorbar(im, ax=ax, label="RMSD (Å)", fraction=0.046, pad=0.04)
+    ax.set_title("pairwise RMSD")
+    if show:
+        plt.show()
+    return ax
+
+
+def _colors(molecule, color_by: str):
+    if color_by == "element":
+        return [elements.color(e) for e in molecule.elements]
+    if color_by == "chain":
+        keys = molecule.chains
+    elif color_by == "residue":
+        keys = [str(r) for r in molecule.resids] if len(molecule.resids) else []
+    else:
+        raise ValueError(f"unknown color_by {color_by!r}")
+    if not keys:
+        raise ValueError(f"no {color_by} information to colour by")
+    return _categorical_colors(keys)
+
+
+def _categorical_colors(keys):
+    import matplotlib.pyplot as plt
+
+    palette = plt.get_cmap("tab20").colors
+    cycle = {}
+    wheel = itertools.cycle(palette)
+    return [cycle.setdefault(k, next(wheel)) for k in keys]
+
+
+def _set_equal_aspect(ax, coords: np.ndarray) -> None:
+    """Force equal scaling on all axes so the molecule isn't distorted."""
+    mins, maxs = coords.min(axis=0), coords.max(axis=0)
+    centers = (maxs + mins) / 2
+    radius = (maxs - mins).max() / 2 or 1.0
+    ax.set_xlim(centers[0] - radius, centers[0] + radius)
+    ax.set_ylim(centers[1] - radius, centers[1] + radius)
+    ax.set_zlim(centers[2] - radius, centers[2] + radius)