From f3fd254d21b61bf4243d344331e04db6156f98d3 Mon Sep 17 00:00:00 2001 From: jvogan <6239693+jvogan@users.noreply.github.com> Date: Mon, 8 Jun 2026 08:46:20 -0700 Subject: [PATCH] Add rendering, animation, and cryo-EM map tooling - chimerax_rest.py: managed ChimeraX REST renderer (GPU) with 0-byte-save fix - pymol_agent.py: turntable movies, render presets, pLDDT coloring - add_helix_records.py: HELIX records for CA-only backbones - map_info.py: MRC/CCP4 sigma-based contour levels - 6 new gotchas (23 total), reference recipes, ffmpeg detection, hardened .gitignore Co-Authored-By: Claude Opus 4.8 (1M context) --- .gitignore | 36 +++ CHANGELOG.md | 18 ++ README.md | 17 +- SKILL.md | 76 ++++++- references/chimerax.md | 41 ++++ references/pymol.md | 60 +++++ scripts/add_helix_records.py | 138 ++++++++++++ scripts/chimerax_rest.py | 410 +++++++++++++++++++++++++++++++++++ scripts/map_info.py | 114 ++++++++++ scripts/proteus_doctor.py | 7 + scripts/pymol_agent.py | 202 ++++++++++++++--- tests/test_scripts.py | 56 +++++ 12 files changed, 1142 insertions(+), 33 deletions(-) create mode 100755 scripts/add_helix_records.py create mode 100755 scripts/chimerax_rest.py create mode 100755 scripts/map_info.py diff --git a/.gitignore b/.gitignore index 3defedc..3631441 100644 --- a/.gitignore +++ b/.gitignore @@ -1,7 +1,43 @@ __pycache__/ *.py[cod] .DS_Store + +# Fetched structures and maps (keep the repo free of downloaded data) AF-*.pdb AF-*.cif AF-*_pae.json +*.pdb +*.cif +*.mmcif +*.bcif +*.ent +*.map +*.mrc +*.mrcs +*.ccp4 +!tests/fixtures/*.pdb +!tests/fixtures/*.cif + +# Rendered output and movies (host large media externally, not in-repo) +*.mp4 +*.mov +*.gif +*.webm +/*.png +/scripts/*.png banners/ + +# Sequences and model weights +*.fasta +*.fa +*.fastq +*.pt +*.pth +*.npy +*.npz + +# Secrets +*.pem +*.key +.env +.env.* diff --git a/CHANGELOG.md b/CHANGELOG.md index 0849ed9..b5c9457 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -1,5 +1,23 @@ # Changelog +## Unreleased + +- Added `chimerax_rest.py`: a managed ChimeraX REST renderer that launches a GUI + session on an ephemeral port, renders via GPU, defeats the 0-byte-PNG save + race, and guarantees teardown. +- Added headless turntable movies (`pymol_agent.py spin` and `chimerax_rest.py + spin`): ray-traced frames encoded with ffmpeg, degrading gracefully when + ffmpeg is absent. +- Added render presets (`--preset publication|illustration|soft`) and pLDDT + confidence coloring (`--color plddt`) to `pymol_agent.py render`. +- Added `add_helix_records.py`: inject HELIX records into CA-only backbones so + cartoons render for de-novo designs. +- Added `map_info.py`: MRC/CCP4 map inspection with sigma-based contour levels. +- Documented six new gotchas (rendering, animation, and cryo-EM maps) and added + turntable, managed-REST, density-map, and exploded-comparison recipes. +- Added ffmpeg detection to `proteus_doctor.py` and hardened `.gitignore` + against accidentally committing structures, maps, and rendered media. + ## v0.1.0 - Public Launch - Added the core Proteus structural biology agent skill. diff --git a/README.md b/README.md index 311ae00..638a978 100644 --- a/README.md +++ b/README.md @@ -41,10 +41,14 @@ Rosetta-oriented protein design guidance without building a custom plugin. ## What It Provides -- **17 documented gotchas** for PyMOL, ChimeraX, and AlphaFold DB — hard-won from real debugging -- Tool detection for PyMOL and ChimeraX across macOS and Linux installs -- Headless PyMOL rendering for publication-quality structure figures +- **23 documented gotchas** for PyMOL, ChimeraX, AlphaFold DB, rendering, and cryo-EM maps — hard-won from real debugging +- Tool detection for PyMOL, ChimeraX, and ffmpeg across macOS and Linux installs +- Headless PyMOL rendering for publication-quality structure figures, with publication/illustration/soft presets and pLDDT coloring +- Headless turntable movies (PyMOL ray-traced frames + ffmpeg), degrading gracefully when ffmpeg is absent +- Managed ChimeraX REST rendering — launches a GUI session, renders via GPU, defeats the 0-byte-PNG save race, and tears down cleanly - ChimeraX analysis helpers for alignment, SASA, and hydrogen-bond workflows +- HELIX-record injection for CA-only backbones (RFdiffusion / Genie designs) so cartoons render correctly +- MRC/CCP4 map inspection with sigma-based contour-level suggestions - AlphaFold DB fetch with confidence interpretation and pLDDT coloring - RCSB PDB fetch for experimental coordinates, metadata, and biological assembly mmCIF - UniProt lookup for resolving gene/protein names before AlphaFold fetches @@ -160,7 +164,11 @@ python3 scripts/validation_report.py 4HHB --json # wwPDB valid python3 scripts/pocket_report.py 1HSG --json # ligand-pocket contacts python3 scripts/resolve_structure.py TP53 --json # one-command resolver python3 scripts/pymol_agent.py render structure.pdb output.png # headless render +python3 scripts/pymol_agent.py spin structure.pdb spin.mp4 # turntable movie (needs ffmpeg) python3 scripts/chimerax_agent.py align reference.pdb mobile.pdb # structure alignment +python3 scripts/chimerax_rest.py render structure.pdb out.png # GPU render via managed REST +python3 scripts/add_helix_records.py model.pdb --json # fix CA-only backbone cartoons +python3 scripts/map_info.py map.mrc --json # cryo-EM contour levels ``` ## Layout @@ -178,10 +186,13 @@ proteus/ │ ├── pymol.md │ └── rosetta.md └── scripts/ # Agent helper scripts (all stdlib-only) + ├── add_helix_records.py ├── chimerax_agent.py + ├── chimerax_rest.py ├── compare_structures.py ├── fetch_pdb.py ├── fetch_alphafold.py + ├── map_info.py ├── pae_report.py ├── pdb_info.py ├── pocket_report.py diff --git a/SKILL.md b/SKILL.md index 8974068..c34956f 100644 --- a/SKILL.md +++ b/SKILL.md @@ -105,8 +105,11 @@ debugging, patching, or the help text is insufficient for the task. | `scripts/validation_report.py` | Fetch wwPDB/RCSB validation quality metrics | `python3 scripts/validation_report.py 4HHB --json` | | `scripts/pocket_report.py` | Zero-dep ligand pocket contacts from PDB/PDB ID | `python3 scripts/pocket_report.py 1HSG --json` | | `scripts/compare_structures.py` | PyMOL CE alignment + optional per-residue deviations | `python3 scripts/compare_structures.py ref.pdb mobile.pdb --json` | -| `scripts/pymol_agent.py` | Headless PyMOL driver | `python3 scripts/pymol_agent.py info structure.pdb` | -| `scripts/chimerax_agent.py` | Headless ChimeraX driver | `python3 scripts/chimerax_agent.py run "open 1ubq; info chains #1"` | +| `scripts/pymol_agent.py` | Headless PyMOL driver (info, render, **spin movie**) | `python3 scripts/pymol_agent.py render structure.pdb out.png --color plddt` | +| `scripts/chimerax_agent.py` | Headless ChimeraX driver (analysis, `--nogui`) | `python3 scripts/chimerax_agent.py run "open 1ubq; info chains #1"` | +| `scripts/chimerax_rest.py` | Managed ChimeraX REST GUI render (GPU) + turntable | `python3 scripts/chimerax_rest.py render structure.pdb out.png --color plddt` | +| `scripts/add_helix_records.py` | Add HELIX records to CA-only backbones so cartoons render | `python3 scripts/add_helix_records.py model.pdb --json` | +| `scripts/map_info.py` | MRC/CCP4 map stats + sigma-based contour levels | `python3 scripts/map_info.py map.mrc --json` | | `scripts/pdb_info.py` | Legacy zero-dep PDB inspector (PDB only) | `python3 scripts/pdb_info.py structure.pdb` | ## Critical Gotchas (Read This First) @@ -204,6 +207,38 @@ can skip by knowing them upfront. (polyubiquitin-C, 685 residues) instead. Note: this is the full polyubiquitin chain, not the 76-residue monomer. +### Rendering, Animation & Maps + +18. **PyMOL spin loops need `set cache_frames, 0`.** Otherwise PyMOL caches + every ray-traced frame in RAM and the process is OOM-killed partway through + a turntable. Set it before the render loop. (`pymol_agent.py spin` does this + for you.) + +19. **`set auto_zoom, 0` before loading when composing a manual view.** Each + `load`/`show` otherwise re-zooms and fights your `orient`/`zoom`/`turn` + framing. Set it first, frame last. + +20. **Whole-map isosurfaces stall in headless PyMOL.** The headless build lacks + the VTKm accelerator, so contouring a full density map can hang for minutes + even on a small map. Carve the mesh around the model + (`isomesh m, map, level, sele, carve=2.5`) or do the whole-map surface in + ChimeraX. Use `scripts/map_info.py` to pick a sigma-based contour level. + +21. **ChimeraX REST `save` can return before the PNG is flushed** — a 0-byte + file, worse under heavy cartoon recompute. After `save`, issue `wait 1` and + poll the file for non-zero size, retrying a few times. + (`scripts/chimerax_rest.py` handles this.) + +22. **ChimeraX color-name traps.** `gold`/`yellow` atom spheres often render + visibly green — use `orange` for a true gold read. And `color #1 cartoons #...` + silently mis-parses `cartoons` as a color: the command is `cartoon`, then + `color #1 `. + +23. **Deposited coordinates are the asymmetric unit, not necessarily the + biological assembly.** A "dimer" entry may deposit a single chain. Build the + functional oligomer with ChimeraX `sym #1 assembly 1 copies true`, or expand + crystal neighbors in PyMOL with `symexp mate_, obj, sele, 5`. + ## Common Workflows ### Quick Structure Inspection @@ -273,6 +308,39 @@ ray 1200, 900 png output.png ``` +Render presets are also available on the helper: `pymol_agent.py render file.pdb +out.png --preset publication|illustration|soft --color spectrum|chain|bfactor|plddt`. + +### Turntable Movie (Headless) +```bash +# PyMOL ray-traces each frame (works with no display); ffmpeg encodes them. +python3 scripts/pymol_agent.py spin structure.pdb spin.mp4 --frames 60 --color plddt +# Degrades gracefully: with no ffmpeg, the frames are written and returned. +``` +This is the macOS-correct path — ChimeraX needs a GPU/GUI context to render. + +### ChimeraX GPU Rendering (Managed REST) +```bash +# Launches a GUI ChimeraX on an ephemeral port, renders via GPU, tears down. +python3 scripts/chimerax_rest.py render structure.pdb out.png --style surface --color bychain +python3 scripts/chimerax_rest.py spin structure.pdb out.mp4 --frames 72 +``` +Unlike `chimerax_agent.py` (analysis only, `--nogui`), this renders images and +defeats the 0-byte-PNG save race (gotcha 21). + +### CA-only Backbone (de-novo designs) +```bash +# RFdiffusion / Genie backbones render as spaghetti without HELIX records. +python3 scripts/add_helix_records.py model.pdb -o model_ss.pdb --json +# Then render model_ss.pdb; in PyMOL also: set cartoon_trace_atoms, 1 +``` + +### Cryo-EM Contour Level +```bash +# Sigma-based level for `volume`/`isomesh` (absolute level differs per map). +python3 scripts/map_info.py map.mrc --json # -> suggested_level at 1/1.5/2/3 sigma +``` + ## Good Demo Proteins | UniProt / PDB | Protein | Good for | @@ -326,6 +394,10 @@ For multi-step workflows, write a summary JSON report at the end with: | Compare two structures | `python3 scripts/compare_structures.py ref.pdb mobile.pdb --per-residue --json` | | Get structure info via PyMOL | `python3 scripts/pymol_agent.py info file.pdb` | | Render a structure headless | `python3 scripts/pymol_agent.py render file.pdb out.png` | +| Render a turntable movie | `python3 scripts/pymol_agent.py spin file.pdb out.mp4` | +| Render via ChimeraX GPU (REST) | `python3 scripts/chimerax_rest.py render file.pdb out.png` | +| Fix a CA-only backbone for cartoons | `python3 scripts/add_helix_records.py model.pdb` | +| Pick a cryo-EM contour level | `python3 scripts/map_info.py map.mrc --json` | | Fetch an AlphaFold prediction | `python3 scripts/fetch_alphafold.py UNIPROT_ID --pae --json` | | Align two structures (ChimeraX) | `python3 scripts/chimerax_agent.py align ref.pdb mobile.pdb` | | Measure SASA | `python3 scripts/chimerax_agent.py sasa file.pdb` | diff --git a/references/chimerax.md b/references/chimerax.md index 6feb95b..0d52aed 100644 --- a/references/chimerax.md +++ b/references/chimerax.md @@ -128,6 +128,34 @@ except Exception: curl "http://127.0.0.1:50888/run?command=remotecontrol+rest+stop" ``` +### Managed rendering with `chimerax_rest.py` + +`scripts/chimerax_rest.py` packages the full lifecycle so you don't manage the +process by hand: it launches a GUI ChimeraX on an ephemeral port (parallel-safe), +polls `version` until ready, drives commands over HTTP with JSON-envelope error +detection (`json true` surfaces command-level failures even on HTTP 200), and +guarantees teardown. + +```bash +python3 scripts/chimerax_rest.py render structure.pdb out.png --color plddt +python3 scripts/chimerax_rest.py spin structure.pdb out.mp4 --frames 72 +python3 scripts/chimerax_rest.py run "open 1ubq from pdb; cartoon; color bychain" +``` + +**The 0-byte-PNG save race.** REST `save` can return HTTP 200 before the GL +framebuffer is flushed, leaving a 0-byte file (worse under heavy cartoon +recompute). After `save`, issue `wait 1` and poll the output file for non-zero +size, retrying the save a few times: + +```python +cx_run("wait 1") +cx_run(f"save {png} width 1200 height 900 supersample 3") +for _ in range(16): + if os.path.exists(png) and os.path.getsize(png) > 0: + break + time.sleep(0.5) +``` + **CRITICAL: ChimeraX is NOT thread-safe.** Sending REST calls from Python background threads causes `EXC_BAD_ACCESS` crashes on macOS. All REST calls must happen from the main thread. If you need concurrency, use `asyncio` @@ -215,6 +243,14 @@ open 21924 from emdb # Cryo-EM density map directory when launched via subprocess may differ from yours, and unquoted paths with spaces or semicolons can break command parsing. +**Asymmetric unit vs biological assembly.** Deposited coordinates are the +asymmetric unit, which is not always the functional oligomer (a "dimer" entry may +deposit one chain). Build the biological assembly explicitly: + +``` +sym #1 assembly 1 copies true # generate biological assembly 1 +``` + ### Visualization ``` cartoon # Show cartoon ribbon @@ -230,6 +266,11 @@ transparency #1 50 # 50% transparent transparency #1 50 target c # Cartoon only (target: c=cartoon, s=surface, a=atoms) ``` +**Color-name traps.** `gold`/`yellow` atom spheres often render visibly green — +use `orange` for a true gold read. And `color #1 cartoons #...` silently +mis-parses `cartoons` as a (missing) color name: the command is `cartoon`, then +`color #1 ` as separate steps. + ### Presets ``` preset interactive # Quick visualization diff --git a/references/pymol.md b/references/pymol.md index 5ed7c99..eb39291 100644 --- a/references/pymol.md +++ b/references/pymol.md @@ -201,6 +201,20 @@ result = cmd.cealign("target", "mobile") # NOTE: target FIRST, mobile SECOND **Critical:** `cealign` argument order is **(target, mobile)** — the first argument stays fixed, the second gets moved. This is opposite to `align` and `super`. +### Exploded side-by-side comparison + +To show two structures as an unambiguous before/after panel, superpose them, then +move *only one object* aside with `camera=1` (so the shift is in screen space): + +```python +cmd.super("design and name CA", "reference and name CA") # CA-only overlay +cmd.translate([22, 0, 0], object="design", camera=1) # pull one copy aside +cmd.rotate("y", 22, object="design", camera=1) # tilt it for depth +``` + +`camera=1` moves the object relative to the camera, not its own frame, so the +anchored structure stays put while the moved one reads as a separate panel. + ## Measurement & Analysis ``` @@ -256,6 +270,52 @@ center selection # Center on selection turn y, 45 # Rotate 45 degrees around Y ``` +**`set auto_zoom, 0` first when you compose a manual view.** Otherwise each +`load`/`show` re-zooms and fights your `orient`/`zoom`/`turn` framing. Set it +before loading, frame last. + +## Turntable Movies & Animation + +PyMOL ray-traces each frame headlessly; ffmpeg encodes them. `pymol_agent.py spin` +packages this, but the core loop is: + +```python +cmd.set("cache_frames", 0) # CRITICAL: else every frame is cached in RAM -> OOM +cmd.orient() +n = 60 +for i in range(n): + cmd.ray(800, 600) + cmd.png(f"/tmp/frames/frame_{i:04d}.png") + cmd.turn("y", 360.0 / n) +``` + +Then encode (web-playable MP4, even dimensions, faststart): + +```bash +ffmpeg -y -framerate 30 -i /tmp/frames/frame_%04d.png \ + -c:v libx264 -crf 18 -pix_fmt yuv420p -movflags +faststart \ + -vf "scale=trunc(iw/2)*2:trunc(ih/2)*2" spin.mp4 +``` + +`cmd.turn` after `cmd.orient()` rotates about the scene center, so the molecule +spins in place instead of drifting across the frame. + +## Density Maps (PyMOL) + +```python +# Simulate density from a model when you have no experimental map +cmd.map_new("sim_map", "gaussian", 1.0, "struct", 5) +cmd.isomesh("dens", "sim_map", 1.0, "struct", carve=2.5) + +# Show an experimental map carved around the model (fast + local) +cmd.isomesh("dens", "emd_map", level, "struct", carve=2.5) +``` + +**Carve density around the model** (`carve=2.5`). Contouring a *whole* map is slow +and can hang in headless PyMOL (the headless build lacks the VTKm accelerator) — +do whole-map surfaces in ChimeraX instead. Use `scripts/map_info.py` to choose a +sigma-based contour level (absolute levels differ per map). + ## Publication Settings Standard block for high-quality renders: diff --git a/scripts/add_helix_records.py b/scripts/add_helix_records.py new file mode 100755 index 0000000..64de8d5 --- /dev/null +++ b/scripts/add_helix_records.py @@ -0,0 +1,138 @@ +#!/usr/bin/env python3 +"""Add HELIX records to a CA-only backbone so viewers draw helical cartoons. + +CA-only models (RFdiffusion, Genie, and other backbone generators) render as +flat "spaghetti" in PyMOL/ChimeraX: the cartoon engine infers helices from +HELIX/SHEET records that these files don't contain. This detects helices from +backbone geometry alone — the CA(i)-CA(i+3) distance is ~5.0-5.4 A inside an +alpha-helix (vs ~10 A for a beta-strand) — and prepends standard PDB HELIX +records so the cartoon renders correctly. Pairs with PyMOL `set cartoon_trace_atoms, 1`. + +Usage: + python add_helix_records.py model.pdb # writes model_with_ss.pdb + python add_helix_records.py model.pdb -o out.pdb --json + python add_helix_records.py --help +""" + +import argparse +import json +import math +import os +import sys + + +def read_ca_coords(path): + """Return [(resnum, chain, x, y, z), ...] for CA atoms in a PDB file.""" + coords = [] + with open(path) as fh: + for line in fh: + if line.startswith("ATOM") and line[12:16].strip() == "CA": + try: + x = float(line[30:38]) + y = float(line[38:46]) + z = float(line[46:54]) + resnum = int(line[22:26]) + except ValueError: + continue + coords.append((resnum, line[21], x, y, z)) + return coords + + +def _dist(a, b): + return math.sqrt((a[0] - b[0]) ** 2 + (a[1] - b[1]) ** 2 + (a[2] - b[2]) ** 2) + + +def detect_helix_residues(coords, lo=4.7, hi=5.8): + """Flag residues whose CA(i)-CA(i+3) distance matches one alpha-helix turn.""" + is_helix = [False] * len(coords) + for i in range(len(coords) - 3): + if lo <= _dist(coords[i][2:], coords[i + 3][2:]) <= hi: + for k in range(i, i + 4): + is_helix[k] = True + return is_helix + + +def helix_segments(coords, is_helix, min_len=6): + """Coalesce flagged residues into (chain, start, end) runs of >= min_len.""" + segments = [] + i = 0 + while i < len(coords): + if is_helix[i]: + j = i + while j < len(coords) and is_helix[j]: + j += 1 + if j - i >= min_len: + segments.append((coords[i][1], coords[i][0], coords[j - 1][0])) + i = j + else: + i += 1 + return segments + + +def _helix_record(n, chain, start, end): + """Format a standard PDB HELIX record.""" + length = end - start + 1 + return (f"HELIX {n:3d} {n:3d} ALA {chain} {start:4d} " + f"ALA {chain} {end:4d} 1{'':33}{length:5d}") + + +def add_helix_records(input_path, output_path, min_len=6): + coords = read_ca_coords(input_path) + if not coords: + return {"status": "error", "error": f"No CA atoms found in {input_path}"} + is_helix = detect_helix_residues(coords) + segments = helix_segments(coords, is_helix, min_len) + helix_lines = [_helix_record(n, c, s, e) + for n, (c, s, e) in enumerate(segments, 1)] + with open(input_path) as fh: + body = fh.read().splitlines() + with open(output_path, "w") as fh: + fh.write("\n".join(helix_lines + body) + "\n") + helical = sum(is_helix) + return {"status": "ok", "data": { + "input": os.path.abspath(input_path), + "output": os.path.abspath(output_path), + "residues": len(coords), + "helical_residues": helical, + "helical_fraction": round(helical / len(coords), 3), + "helices": [{"chain": c, "start": s, "end": e, "length": e - s + 1} + for c, s, e in segments], + }} + + +def main(): + parser = argparse.ArgumentParser( + description="Add HELIX records to a CA-only PDB from backbone geometry.", + formatter_class=argparse.RawDescriptionHelpFormatter, + epilog="Examples:\n" + " %(prog)s model.pdb\n" + " %(prog)s model.pdb -o model_ss.pdb --json", + ) + parser.add_argument("pdb", help="CA-only PDB file") + parser.add_argument("-o", "--output", help="Output path (default: _with_ss.pdb)") + parser.add_argument("--min-len", type=int, default=6, + help="Minimum helix length in residues (default: 6)") + parser.add_argument("--json", action="store_true", help="Emit machine-readable JSON") + args = parser.parse_args() + + if not os.path.isfile(args.pdb): + result = {"status": "error", "error": f"File not found: {args.pdb}"} + else: + output = args.output or (os.path.splitext(args.pdb)[0] + "_with_ss.pdb") + result = add_helix_records(args.pdb, output, args.min_len) + + if args.json: + print(json.dumps(result, indent=2)) + elif result["status"] == "ok": + d = result["data"] + print(f"{d['input']}: {d['residues']} residues, {d['helical_residues']} helical " + f"({d['helical_fraction'] * 100:.1f}%), {len(d['helices'])} helices -> {d['output']}") + for h in d["helices"]: + print(f" HELIX chain {h['chain']} {h['start']}-{h['end']} (len {h['length']})") + else: + print(result["error"], file=sys.stderr) + sys.exit(1) + + +if __name__ == "__main__": + main() diff --git a/scripts/chimerax_rest.py b/scripts/chimerax_rest.py new file mode 100755 index 0000000..d6d1205 --- /dev/null +++ b/scripts/chimerax_rest.py @@ -0,0 +1,410 @@ +#!/usr/bin/env python3 +"""ChimeraX REST render agent — Proteus skill. + +Drives a *managed* ChimeraX GUI session over its REST API so an agent can render +images (and turntable movies) from the terminal without hand-managing the +process. ChimeraX `--nogui` cannot render on macOS (no OpenGL context); this +launches a real GUI instance on an ephemeral port, talks to it over HTTP, and +guarantees teardown. + +Where chimerax_agent.py does headless analysis (`--nogui`), this does GPU +rendering (REST + GUI). It also defeats the macOS "0-byte PNG" save race. + +Usage: + python chimerax_rest.py render structure.pdb out.png + python chimerax_rest.py render structure.pdb out.png --style surface --color bychain + python chimerax_rest.py render model.pdb out.png --color plddt + python chimerax_rest.py spin model.pdb spin.mp4 --frames 72 # needs ffmpeg + python chimerax_rest.py run "open 1ubq from pdb; cartoon; color bychain" + python chimerax_rest.py --help + +Environment: + CHIMERAX_BIN Override the ChimeraX binary path. +""" + +import argparse +import glob +import http.client +import json +import os +import re +import shutil +import socket +import subprocess +import sys +import tempfile +import time +from urllib.parse import urlencode + + +def _find_chimerax() -> str: + """Auto-detect ChimeraX binary (PATH, then common install locations).""" + found = shutil.which("ChimeraX") or shutil.which("chimerax") + if found: + return found + hits = glob.glob("/Applications/ChimeraX*.app/Contents/bin/ChimeraX") + if hits: + return sorted(hits)[-1] + for p in ["/usr/bin/chimerax", "/usr/local/bin/chimerax", + os.path.expanduser("~/ChimeraX/bin/ChimeraX")]: + if os.path.isfile(p): + return p + return None + + +CHIMERAX = os.environ.get("CHIMERAX_BIN") or _find_chimerax() +DEFAULT_WIDTH = 1200 +DEFAULT_HEIGHT = 900 + + +def find_free_port() -> int: + """Bind an ephemeral port so parallel ChimeraX sessions never collide.""" + with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as sock: + sock.bind(("127.0.0.1", 0)) + return int(sock.getsockname()[1]) + + +def _quote(path: str) -> str: + return '"' + os.path.abspath(path).replace('"', '\\"') + '"' + + +def _validate_color(color: str) -> str: + if color in {"rainbow", "bychain", "bfactor", "plddt"}: + return color + if not re.fullmatch(r"[A-Za-z][A-Za-z0-9_]*", color): + raise ValueError( + "Unsafe ChimeraX color. Use rainbow, bychain, bfactor, plddt, or a simple color name." + ) + return color + + +class ChimeraXRest: + """Launch a local ChimeraX GUI with its REST server and drive it over HTTP. + + Use as a context manager so the process is always torn down: + + with ChimeraXRest() as rest: + rest.run("open 1ubq from pdb") + rest.save_image("/tmp/out.png") + """ + + def __init__(self, chimerax: str = None, port: int = None): + self.chimerax = chimerax or CHIMERAX + self.port = port or find_free_port() + self.process = None + self.history: list = [] + self._log_path = None + + def __enter__(self): + self.start() + return self + + def __exit__(self, *exc): + self.stop() + return False + + def start(self, ready_timeout: int = 60) -> None: + if not self.chimerax: + raise RuntimeError("ChimeraX not found. Install it or set CHIMERAX_BIN.") + log_handle = tempfile.NamedTemporaryFile( + mode="w", suffix=".chimerax.log", delete=False) + self._log_path = log_handle.name + # A GUI (OpenGL) context is required to render; do NOT pass --offscreen on + # macOS (it hangs the Qt event loop). `json true` makes /run return a JSON + # envelope so command-level errors are detectable even on HTTP 200. + self.process = subprocess.Popen( + [self.chimerax, "--cmd", + f"remotecontrol rest start port {self.port} json true log false"], + stdout=log_handle, stderr=subprocess.STDOUT, text=True, + ) + deadline = time.time() + ready_timeout + last_error = "" + while time.time() < deadline: + if self.process.poll() is not None: + raise RuntimeError( + f"ChimeraX exited before REST startup; see {self._log_path}") + try: + self.run("version", timeout=5) + return + except Exception as exc: # startup retry loop + last_error = str(exc) + time.sleep(1) + self.stop() + raise TimeoutError( + f"ChimeraX REST did not start on port {self.port}: {last_error}") + + def run(self, command: str, *, timeout: int = 120, soft: bool = False) -> str: + """Run one ChimeraX command over REST. + + `soft=True` swallows command/HTTP errors so a purely decorative command + never aborts a scene. Errors are detected from the JSON envelope's + `error` key, not just the HTTP status. + """ + conn = http.client.HTTPConnection("127.0.0.1", self.port, timeout=timeout) + path = "/run?" + urlencode({"command": command, "json": "true"}) + started = time.time() + try: + conn.request("GET", path) + response = conn.getresponse() + body = response.read().decode("utf-8", errors="replace") + self.history.append({ + "command": command, + "elapsed_seconds": round(time.time() - started, 3), + "http_status": response.status, + }) + if response.status >= 400: + if soft: + return body + raise RuntimeError(f"ChimeraX REST {response.status}: {body[:300]}") + try: + payload = json.loads(body) + except json.JSONDecodeError: + payload = {} + if payload.get("error"): + if soft: + return body + raise RuntimeError( + f"ChimeraX command failed for {command!r}: {payload['error']}") + return body + finally: + conn.close() + + def run_all(self, commands, *, timeout: int = 120) -> None: + for command in commands: + self.run(command, timeout=timeout) + + def save_image(self, path, width=DEFAULT_WIDTH, height=DEFAULT_HEIGHT, + supersample=3, attempts=4) -> None: + """Save a PNG, defeating the macOS 0-byte race. + + REST `save` can return HTTP 200 before the GL framebuffer is flushed, + leaving a 0-byte PNG (worse under heavy cartoon recompute). Issue + `wait 1` to force a redraw, then poll the file for non-zero size, + retrying the whole save a few times. + """ + path = os.path.abspath(path) + os.makedirs(os.path.dirname(path) or ".", exist_ok=True) + if os.path.exists(path): + os.unlink(path) + last = 0 + for _ in range(attempts): + self.run("wait 1") + self.run(f"save {_quote(path)} width {width} height {height} " + f"supersample {supersample}", timeout=240) + for _ in range(16): + if os.path.exists(path) and os.path.getsize(path) > 0: + return + time.sleep(0.5) + last = os.path.getsize(path) if os.path.exists(path) else 0 + raise RuntimeError( + f"ChimeraX did not write a non-empty image after {attempts} attempts " + f"(last size={last}): {path}") + + def stop(self) -> None: + for command in ("remotecontrol rest stop", "exit"): + try: + self.run(command, timeout=5) + except Exception: + pass + if self.process is not None: + try: + self.process.wait(timeout=10) + except subprocess.TimeoutExpired: + self.process.terminate() + try: + self.process.wait(timeout=5) + except subprocess.TimeoutExpired: + self.process.kill() + self.process.wait(timeout=5) + if self._log_path and os.path.exists(self._log_path): + try: + os.unlink(self._log_path) + except OSError: + pass + + +def _scene_commands(style: str, color: str) -> list: + """A tasteful publication default: white bg, soft light, silhouettes.""" + cmds = [ + "set bgColor white", + "lighting soft", + "graphics silhouettes true", + "set subdivision 3", + ] + styles = { + "cartoon": ["hide #1 atoms", "cartoon #1"], + "surface": ["surface #1"], + "stick": ["hide #1 cartoon", "style #1 stick", "show #1 atoms"], + "sphere": ["hide #1 cartoon", "style #1 sphere", "show #1 atoms"], + } + cmds += styles.get(style, styles["cartoon"]) + color_cmd = { + "rainbow": "rainbow #1", + "bychain": "color bychain #1", + "bfactor": "color bfactor #1", + "plddt": "color bfactor #1 palette alphafold", + }.get(color, f"color #1 {color}") + cmds.append(color_cmd) + return cmds + + +def _encode_movie(frame_dir: str, output: str, fps: int = 30) -> str: + """Encode frame_%04d.png in frame_dir to MP4 (or GIF if output ends .gif). + + yuv420p + even-dimension scaling keeps the MP4 web-playable; GIF uses a + two-pass palette for clean colors. + """ + out = os.path.abspath(output) + pattern = os.path.join(frame_dir, "frame_%04d.png") + even = "scale=trunc(iw/2)*2:trunc(ih/2)*2" + if out.lower().endswith(".gif"): + palette = os.path.join(frame_dir, "palette.png") + subprocess.run(["ffmpeg", "-y", "-framerate", str(fps), "-i", pattern, + "-vf", f"{even},palettegen", palette], + check=True, capture_output=True, text=True) + subprocess.run(["ffmpeg", "-y", "-framerate", str(fps), "-i", pattern, + "-i", palette, "-lavfi", f"{even} [x]; [x][1:v] paletteuse", out], + check=True, capture_output=True, text=True) + else: + subprocess.run(["ffmpeg", "-y", "-framerate", str(fps), "-i", pattern, + "-c:v", "libx264", "-preset", "slow", "-crf", "18", + "-pix_fmt", "yuv420p", "-movflags", "+faststart", + "-vf", even, out], + check=True, capture_output=True, text=True) + return out + + +def rest_run(commands: str) -> dict: + """Open a managed session, run semicolon/newline-separated commands, tear down.""" + if not CHIMERAX: + return {"status": "error", "error": "ChimeraX not found. Install it or set CHIMERAX_BIN."} + cmd_list = [c.strip() for c in commands.replace("\n", ";").split(";") if c.strip()] + try: + with ChimeraXRest() as rest: + rest.run_all(cmd_list) + return {"status": "ok", "data": {"history": rest.history}} + except Exception as exc: + return {"status": "error", "error": str(exc)} + + +def rest_render(structure: str, output: str, style: str = "cartoon", color: str = "rainbow", + width: int = DEFAULT_WIDTH, height: int = DEFAULT_HEIGHT, + supersample: int = 3) -> dict: + """Render a structure to PNG via a managed ChimeraX GUI session (GPU).""" + if not CHIMERAX: + return {"status": "error", "error": "ChimeraX not found. Install it or set CHIMERAX_BIN."} + if not os.path.isfile(structure): + return {"status": "error", "error": f"Structure file not found: {structure}"} + try: + color = _validate_color(color) + except ValueError as exc: + return {"status": "error", "error": str(exc)} + try: + with ChimeraXRest() as rest: + rest.run(f"open {_quote(structure)}") + rest.run_all(_scene_commands(style, color)) + rest.run("view") + rest.save_image(output, width, height, supersample) + return {"status": "ok", "data": { + "rendered": os.path.abspath(output), + "size": f"{width}x{height}", + "supersample": supersample, + }} + except Exception as exc: + return {"status": "error", "error": str(exc)} + + +def rest_spin(structure: str, output: str, frames: int = 72, style: str = "cartoon", + color: str = "rainbow", width: int = 800, height: int = 600, + fps: int = 30) -> dict: + """Render a 360-degree y-spin and encode to a movie (GPU frames + ffmpeg).""" + if not CHIMERAX: + return {"status": "error", "error": "ChimeraX not found. Install it or set CHIMERAX_BIN."} + if not os.path.isfile(structure): + return {"status": "error", "error": f"Structure file not found: {structure}"} + try: + color = _validate_color(color) + except ValueError as exc: + return {"status": "error", "error": str(exc)} + ffmpeg = shutil.which("ffmpeg") + frame_dir = tempfile.mkdtemp(prefix="proteus_cxspin_") + try: + with ChimeraXRest() as rest: + rest.run(f"open {_quote(structure)}") + rest.run_all(_scene_commands(style, color)) + rest.run("view") # centers content so `turn y` spins in place + step = 360.0 / frames + for idx in range(frames): + rest.save_image(os.path.join(frame_dir, f"frame_{idx:04d}.png"), + width, height, supersample=1) + rest.run(f"turn y {step:.4f}") + if not ffmpeg: + return {"status": "ok", "data": { + "movie": None, "frames_dir": frame_dir, "frame_count": frames, + "note": "ffmpeg not found; wrote frames but did not encode a movie.", + }} + _encode_movie(frame_dir, output, fps) + shutil.rmtree(frame_dir, ignore_errors=True) + return {"status": "ok", "data": { + "movie": os.path.abspath(output), "frame_count": frames, "fps": fps, + }} + except subprocess.CalledProcessError as exc: + return {"status": "error", "error": f"ffmpeg failed: {(exc.stderr or '')[:300]}", + "data": {"frames_dir": frame_dir}} + except Exception as exc: + shutil.rmtree(frame_dir, ignore_errors=True) + return {"status": "error", "error": str(exc)} + + +def main(): + parser = argparse.ArgumentParser( + description="ChimeraX REST render agent — GPU rendering and turntable movies via a managed GUI session.", + epilog="Examples:\n" + " %(prog)s render structure.pdb out.png --color plddt\n" + " %(prog)s spin model.pdb spin.mp4 --frames 72\n" + " %(prog)s run 'open 1ubq from pdb; cartoon; color bychain'", + formatter_class=argparse.RawDescriptionHelpFormatter, + ) + sub = parser.add_subparsers(dest="command", required=True) + + p_run = sub.add_parser("run", help="Run commands in a managed REST session") + p_run.add_argument("commands", help="Semicolon-separated ChimeraX commands") + + p_render = sub.add_parser("render", help="Render a structure to PNG (GPU via REST)") + p_render.add_argument("structure", help="Path to structure file") + p_render.add_argument("output", nargs="?", default="/tmp/chimerax_render.png") + p_render.add_argument("--style", default="cartoon", choices=["cartoon", "surface", "stick", "sphere"]) + p_render.add_argument("--color", default="rainbow", help="rainbow, bychain, bfactor, plddt, or a color name") + p_render.add_argument("--width", type=int, default=DEFAULT_WIDTH) + p_render.add_argument("--height", type=int, default=DEFAULT_HEIGHT) + p_render.add_argument("--supersample", type=int, default=3) + + p_spin = sub.add_parser("spin", help="Render a 360-degree turntable movie (needs ffmpeg)") + p_spin.add_argument("structure", help="Path to structure file") + p_spin.add_argument("output", nargs="?", default="/tmp/chimerax_spin.mp4") + p_spin.add_argument("--frames", type=int, default=72) + p_spin.add_argument("--style", default="cartoon", choices=["cartoon", "surface", "stick", "sphere"]) + p_spin.add_argument("--color", default="rainbow", help="rainbow, bychain, bfactor, plddt, or a color name") + p_spin.add_argument("--width", type=int, default=800) + p_spin.add_argument("--height", type=int, default=600) + p_spin.add_argument("--fps", type=int, default=30) + + args = parser.parse_args() + + if args.command == "run": + result = rest_run(args.commands) + elif args.command == "render": + result = rest_render(args.structure, args.output, args.style, args.color, + args.width, args.height, args.supersample) + elif args.command == "spin": + result = rest_spin(args.structure, args.output, args.frames, args.style, + args.color, args.width, args.height, args.fps) + else: + parser.print_help() + sys.exit(1) + + print(json.dumps(result, indent=2)) + + +if __name__ == "__main__": + main() diff --git a/scripts/map_info.py b/scripts/map_info.py new file mode 100755 index 0000000..68960a7 --- /dev/null +++ b/scripts/map_info.py @@ -0,0 +1,114 @@ +#!/usr/bin/env python3 +"""Inspect an MRC/CCP4 density map and suggest contour levels. + +Cryo-EM contour levels are in absolute map units that differ per map, so a +hard-coded `--level` rarely transfers. This reads the MRC/CCP4 header with the +standard library and derives sigma-based levels (mean + N*sigma) from the voxel +data — the protein-agnostic default ChimeraX `volume` / PyMOL `isomesh` want. + +Uses numpy if available for speed; otherwise a pure-stdlib strided sampler keeps +it dependency-free. + +Usage: + python map_info.py map.mrc --json + python map_info.py map.mrc --n-sigma 1.5 + python map_info.py --help +""" + +import argparse +import json +import os +import struct +import sys + + +_MODE_BYTES = {0: 1, 1: 2, 2: 4, 6: 2} +_MODE_STRUCT = {0: "b", 1: "h", 2: "f", 6: "H"} +_MODE_NAME = {0: "int8", 1: "int16", 2: "float32", 6: "uint16"} + + +def read_map_stats(path, sample_target=200000): + """Return (dims, mode, mean, sigma) for an MRC/CCP4 map. + + The MRC header is 1024 bytes: nx/ny/nz/mode are the first four int32s, and + the extended-header byte count (nsymbt) is at offset 92. Voxel data starts + at 1024 + nsymbt. + """ + with open(path, "rb") as fh: + header = fh.read(1024) + if len(header) < 1024: + raise ValueError("File too small to be an MRC/CCP4 map") + nx, ny, nz, mode = struct.unpack("<4i", header[:16]) + nsymbt = struct.unpack(" dict: "pdb_info.py", "pymol_agent.py", "chimerax_agent.py", + "chimerax_rest.py", "pae_report.py", "resolve_structure.py", "validation_report.py", "pocket_report.py", "compare_structures.py", + "add_helix_records.py", + "map_info.py", ] results = {} for script in scripts: @@ -121,6 +124,7 @@ def _network_check(url: str) -> dict: def build_report(include_network: bool) -> dict: pymol = _find_pymol() chimerax = _find_chimerax() + ffmpeg = shutil.which("ffmpeg") data = { "root": str(ROOT), "platform": { @@ -131,6 +135,7 @@ def build_report(include_network: bool) -> dict: "tools": { "pymol": {"ok": bool(pymol), "path": pymol}, "chimerax": {"ok": bool(chimerax), "path": chimerax}, + "ffmpeg": {"ok": bool(ffmpeg), "path": ffmpeg}, }, "scripts": _script_smoke(), "network": None, @@ -141,6 +146,7 @@ def build_report(include_network: bool) -> dict: "alphafold_fetch": include_network, "pymol_rendering": bool(pymol), "chimerax_analysis": bool(chimerax), + "turntable_movies": bool((pymol or chimerax) and ffmpeg), }, } if include_network: @@ -169,6 +175,7 @@ def main(): print(f"Python: {data['python']['version']} ({'ok' if data['python']['ok'] else 'too old'})") print(f"PyMOL: {data['tools']['pymol']['path'] or 'not found'}") print(f"ChimeraX: {data['tools']['chimerax']['path'] or 'not found'}") + print(f"ffmpeg: {data['tools']['ffmpeg']['path'] or 'not found'} (turntable movies)") failed = [name for name, result in data["scripts"].items() if not result["ok"]] print(f"Script help checks: {'ok' if not failed else 'failed: ' + ', '.join(failed)}") if data["network"]: diff --git a/scripts/pymol_agent.py b/scripts/pymol_agent.py index 7a406af..e9821e4 100755 --- a/scripts/pymol_agent.py +++ b/scripts/pymol_agent.py @@ -90,15 +90,101 @@ def _py_literal(value: str) -> str: def _validate_pymol_color(color: str) -> str: - if color in {"spectrum", "bfactor", "chain"}: + if color in {"spectrum", "bfactor", "chain", "plddt"}: return color if not re.fullmatch(r"[A-Za-z][A-Za-z0-9_]*", color): raise ValueError( - "Unsafe PyMOL color name. Use spectrum, bfactor, chain, or a simple PyMOL color identifier." + "Unsafe PyMOL color name. Use spectrum, bfactor, chain, plddt, or a simple PyMOL color identifier." ) return color +def _color_script(color_mode: str, selection: str = "all") -> str: + """PyMOL command lines that apply a color mode to a selection.""" + sel = _py_literal(selection) + if color_mode == "spectrum": + return f'cmd.spectrum("count", "rainbow", {sel})' + if color_mode == "bfactor": + return f'cmd.spectrum("b", "blue_white_red", {sel})' + if color_mode == "chain": + return f'util.cbc({sel})' + if color_mode == "plddt": + # Official AlphaFold bins. Layered broadest-first because PyMOL selection + # algebra has no `<=`: paint everything low, then override upward. + return "\n".join([ + f'cmd.color("orange", {sel})', + f'cmd.color("yellow", {_py_literal(f"({selection}) and b > 50")})', + f'cmd.color("cyan", {_py_literal(f"({selection}) and b > 70")})', + f'cmd.color("blue", {_py_literal(f"({selection}) and b > 90")})', + ]) + return f'cmd.color({_py_literal(color_mode)}, {sel})' + + +def _preset_script(preset: str) -> str: + """PyMOL command lines for a render look. Default: publication.""" + if preset == "illustration": + return "\n".join([ + 'cmd.bg_color("white")', + 'cmd.set("ray_opaque_background", 1)', + 'cmd.set("antialias", 2)', + 'cmd.set("cartoon_fancy_helices", 1)', + 'cmd.set("cartoon_smooth_loops", 1)', + 'cmd.set("cartoon_flat_sheets", 1)', + 'cmd.set("ray_trace_mode", 3)', # quantized colors + black outlines + 'cmd.set("ray_trace_color", "black")', + ]) + if preset == "soft": + return "\n".join([ + 'cmd.bg_color("gray90")', + 'cmd.set("ray_opaque_background", 1)', + 'cmd.set("orthoscopic", 1)', + 'cmd.set("ray_shadows", 0)', + 'cmd.set("antialias", 2)', + 'cmd.set("ambient", 0.4)', + 'cmd.set("specular", 0.15)', + 'cmd.set("cartoon_fancy_helices", 1)', + 'cmd.set("cartoon_smooth_loops", 1)', + 'cmd.set("cartoon_flat_sheets", 1)', + ]) + return "\n".join([ + 'cmd.bg_color("white")', + 'cmd.set("ray_opaque_background", 1)', + 'cmd.set("antialias", 2)', + 'cmd.set("ray_shadows", 1)', + 'cmd.set("specular", 0.25)', + 'cmd.set("ambient", 0.35)', + 'cmd.set("cartoon_fancy_helices", 1)', + 'cmd.set("cartoon_smooth_loops", 1)', + 'cmd.set("cartoon_flat_sheets", 1)', + ]) + + +def _encode_movie(frame_dir: str, output: str, fps: int = 30) -> str: + """Encode frame_%04d.png frames to MP4 (or GIF if output ends .gif). + + yuv420p + even-dimension scaling keeps the MP4 web-playable; GIF uses a + two-pass palette for clean colors. Raises CalledProcessError on failure. + """ + out = os.path.abspath(output) + pattern = os.path.join(frame_dir, "frame_%04d.png") + even = "scale=trunc(iw/2)*2:trunc(ih/2)*2" + if out.lower().endswith(".gif"): + palette = os.path.join(frame_dir, "palette.png") + subprocess.run(["ffmpeg", "-y", "-framerate", str(fps), "-i", pattern, + "-vf", f"{even},palettegen", palette], + check=True, capture_output=True, text=True) + subprocess.run(["ffmpeg", "-y", "-framerate", str(fps), "-i", pattern, + "-i", palette, "-lavfi", f"{even} [x]; [x][1:v] paletteuse", out], + check=True, capture_output=True, text=True) + else: + subprocess.run(["ffmpeg", "-y", "-framerate", str(fps), "-i", pattern, + "-c:v", "libx264", "-preset", "slow", "-crf", "18", + "-pix_fmt", "yuv420p", "-movflags", "+faststart", + "-vf", even, out], + check=True, capture_output=True, text=True) + return out + + def run_pymol_script(script_content: str, timeout: int = 120) -> dict: """Run a PyMOL Python script headlessly and capture output as JSON. @@ -192,14 +278,17 @@ def get_structure_info(pdb_path: str) -> dict: def render_structure(pdb_path: str, output_png: str, width: int = 1200, height: int = 900, - style: str = "cartoon", color: str = "spectrum") -> dict: + style: str = "cartoon", color: str = "spectrum", + preset: str = "publication") -> dict: """Load and render a structure to PNG using PyMOL's software ray tracer. Works fully headless — no display required. Args: style: cartoon, sticks, surface, spheres, lines - color: spectrum (rainbow), bfactor (blue-white-red), chain, or any PyMOL color name + color: spectrum (rainbow), bfactor (blue-white-red), chain, plddt + (AlphaFold confidence bins), or any PyMOL color name + preset: publication, illustration (outlined), or soft (neutral background) """ try: color = _validate_pymol_color(color) @@ -209,36 +298,76 @@ def render_structure(pdb_path: str, output_png: str, width: int = 1200, height: abs_pdb = os.path.abspath(pdb_path) abs_out = os.path.abspath(output_png) script = f''' -structure_path = {_py_literal(abs_pdb)} -output_png = {_py_literal(abs_out)} -style = {_py_literal(style)} -color_mode = {_py_literal(color)} - -cmd.load(structure_path, "struct") +cmd.load({_py_literal(abs_pdb)}, "struct") cmd.hide("everything") -cmd.show(style, "all") -if color_mode == "spectrum": - cmd.spectrum("count", "rainbow", "all") -elif color_mode == "bfactor": - cmd.spectrum("b", "blue_white_red", "all") -elif color_mode == "chain": - util.cbc("all") -else: - cmd.color(color_mode, "all") -cmd.bg_color("white") -cmd.set("ray_opaque_background", 1) -cmd.set("antialias", 2) -cmd.set("cartoon_fancy_helices", 1) -cmd.set("cartoon_smooth_loops", 1) +cmd.show({_py_literal(style)}, "all") +{_color_script(color)} +{_preset_script(preset)} cmd.orient() cmd.ray({width}, {height}) -cmd.png(output_png) -_output["data"]["rendered"] = output_png +cmd.png({_py_literal(abs_out)}) +_output["data"]["rendered"] = {_py_literal(abs_out)} _output["data"]["size"] = "{width}x{height}" ''' return run_pymol_script(script, timeout=300) # Rendering can take longer +def render_spin(pdb_path: str, output: str, frames: int = 60, width: int = 800, + height: int = 600, style: str = "cartoon", color: str = "spectrum", + preset: str = "publication", fps: int = 30) -> dict: + """Render a 360-degree y-spin as ray-traced frames, then encode to a movie. + + This is the headless-correct turntable path: PyMOL ray-traces each frame + (works without a display) and ffmpeg encodes them. Degrades gracefully — + if ffmpeg is missing, the frames are written and their directory returned. + """ + if not PYMOL: + return {"status": "error", "error": "PyMOL not found. Install it or set PYMOL_BIN."} + if not os.path.isfile(pdb_path): + return {"status": "error", "error": f"Structure file not found: {pdb_path}"} + try: + color = _validate_pymol_color(color) + except ValueError as exc: + return {"status": "error", "error": str(exc)} + + ffmpeg = shutil.which("ffmpeg") + frame_dir = tempfile.mkdtemp(prefix="proteus_spin_") + abs_pdb = os.path.abspath(pdb_path) + script = f''' +cmd.load({_py_literal(abs_pdb)}, "struct") +cmd.hide("everything") +cmd.show({_py_literal(style)}, "all") +{_color_script(color)} +{_preset_script(preset)} +cmd.set("cache_frames", 0) # else PyMOL caches every frame in RAM -> OOM on long spins +cmd.orient() +_n = {frames} +for _i in range(_n): + cmd.ray({width}, {height}) + cmd.png(os.path.join({_py_literal(frame_dir)}, "frame_%04d.png" % _i)) + cmd.turn("y", 360.0 / _n) +_output["data"]["frames"] = _n +''' + result = run_pymol_script(script, timeout=max(600, frames * 20)) + if result.get("status") != "ok": + shutil.rmtree(frame_dir, ignore_errors=True) + return result + if not ffmpeg: + return {"status": "ok", "data": { + "movie": None, "frames_dir": frame_dir, "frame_count": frames, + "note": "ffmpeg not found; wrote frames but did not encode a movie.", + }} + try: + _encode_movie(frame_dir, output, fps) + except subprocess.CalledProcessError as exc: + return {"status": "error", "error": f"ffmpeg failed: {(exc.stderr or '')[:300]}", + "data": {"frames_dir": frame_dir}} + shutil.rmtree(frame_dir, ignore_errors=True) + return {"status": "ok", "data": { + "movie": os.path.abspath(output), "frame_count": frames, "fps": fps, + }} + + def main(): parser = argparse.ArgumentParser( description="PyMOL headless agent helper — run commands, inspect structures, render images.", @@ -266,7 +395,20 @@ def main(): p_render.add_argument("--width", type=int, default=1200) p_render.add_argument("--height", type=int, default=900) p_render.add_argument("--style", default="cartoon", choices=["cartoon", "sticks", "surface", "spheres", "lines"]) - p_render.add_argument("--color", default="spectrum", help="spectrum, bfactor, chain, or PyMOL color name") + p_render.add_argument("--color", default="spectrum", help="spectrum, bfactor, chain, plddt, or PyMOL color name") + p_render.add_argument("--preset", default="publication", choices=["publication", "illustration", "soft"]) + + # spin + p_spin = sub.add_parser("spin", help="Render a 360-degree turntable movie (frames -> ffmpeg)") + p_spin.add_argument("pdb", help="Path to structure file") + p_spin.add_argument("output", nargs="?", default="/tmp/pymol_spin.mp4", help="Output movie path (.mp4 or .gif)") + p_spin.add_argument("--frames", type=int, default=60) + p_spin.add_argument("--width", type=int, default=800) + p_spin.add_argument("--height", type=int, default=600) + p_spin.add_argument("--style", default="cartoon", choices=["cartoon", "sticks", "surface", "spheres", "lines"]) + p_spin.add_argument("--color", default="spectrum", help="spectrum, bfactor, chain, plddt, or PyMOL color name") + p_spin.add_argument("--preset", default="publication", choices=["publication", "illustration", "soft"]) + p_spin.add_argument("--fps", type=int, default=30) args = parser.parse_args() @@ -275,7 +417,11 @@ def main(): elif args.command == "info": result = get_structure_info(args.pdb) elif args.command == "render": - result = render_structure(args.pdb, args.output, args.width, args.height, args.style, args.color) + result = render_structure(args.pdb, args.output, args.width, args.height, + args.style, args.color, args.preset) + elif args.command == "spin": + result = render_spin(args.pdb, args.output, args.frames, args.width, args.height, + args.style, args.color, args.preset, args.fps) else: parser.print_help() sys.exit(1) diff --git a/tests/test_scripts.py b/tests/test_scripts.py index 506b7fe..5adee84 100644 --- a/tests/test_scripts.py +++ b/tests/test_scripts.py @@ -1,17 +1,21 @@ #!/usr/bin/env python3 import ast import json +import struct import subprocess import sys +import tempfile import unittest from pathlib import Path ROOT = Path(__file__).resolve().parents[1] sys.path.insert(0, str(ROOT / "scripts")) +import add_helix_records import chimerax_agent import fetch_alphafold import fetch_pdb +import map_info import pymol_agent import uniprot_lookup @@ -97,11 +101,60 @@ def test_pymol_literals_and_color_validation(self): value = 'path with spaces/"quote";still-path' self.assertEqual(ast.literal_eval(pymol_agent._py_literal(value)), value) self.assertEqual(pymol_agent._validate_pymol_color("carbon"), "carbon") + self.assertEqual(pymol_agent._validate_pymol_color("plddt"), "plddt") with self.assertRaises(ValueError): pymol_agent._validate_pymol_color("red; import os") result = pymol_agent.render_structure("tests/fixtures/tiny.pdb", "out.png", color="red; import os") self.assertEqual(result["status"], "error") + def test_pymol_plddt_color_script_is_layered(self): + script = pymol_agent._color_script("plddt") + # Broadest-first layering (no <= in PyMOL selection algebra) + self.assertIn('cmd.color("orange"', script) + self.assertIn("b > 50", script) + self.assertIn("b > 90", script) + self.assertLess(script.index("b > 50"), script.index("b > 90")) + + def test_chimerax_rest_color_validation(self): + import chimerax_rest + self.assertEqual(chimerax_rest._validate_color("plddt"), "plddt") + with self.assertRaises(ValueError): + chimerax_rest._validate_color("red; close session") + result = chimerax_rest.rest_render("tests/fixtures/tiny.pdb", "out.png", color="red; close") + self.assertEqual(result["status"], "error") + + def test_add_helix_detects_ideal_helix(self): + import math + coords = [] + for i in range(1, 15): + ang = math.radians(100 * i) + coords.append((i, "A", 2.3 * math.cos(ang), 2.3 * math.sin(ang), 1.5 * i)) + is_helix = add_helix_records.detect_helix_residues(coords) + self.assertTrue(all(is_helix)) + segments = add_helix_records.helix_segments(coords, is_helix, min_len=6) + self.assertEqual(len(segments), 1) + self.assertEqual(segments[0], ("A", 1, 14)) + + def test_map_info_sigma_from_synthetic_mrc(self): + nx = ny = nz = 4 + vals = [float(i % 7) for i in range(nx * ny * nz)] + header = bytearray(1024) + struct.pack_into("<4i", header, 0, nx, ny, nz, 2) # dims + float32 mode + struct.pack_into("