|
| 1 | +#!/usr/bin/env python3 |
| 2 | +"""CryoCore ChimeraX figure/movie renderer (license-controlled path). |
| 3 | +
|
| 4 | +ChimeraX is a strong cryo-EM map renderer, but it needs an OpenGL context, so |
| 5 | +the headless path is platform-specific: |
| 6 | +
|
| 7 | + * Linux : xvfb-run ChimeraX --offscreen --nogui --script x.cxc --exit |
| 8 | + * macOS : ChimeraX --offscreen HANGS (no OSMesa). Use the GUI + REST |
| 9 | + API (renders via Metal/GPU) -- opt in with --rest. PyMOL |
| 10 | + (scripts/cryocore/render/pymol_render.py) is the reliable |
| 11 | + window-free local alternative. |
| 12 | +
|
| 13 | +This tool generates the .cxc script for a figure/movie and runs it on the |
| 14 | +available path; with no runnable path it writes the .cxc and prints the exact |
| 15 | +Linux command. ChimeraX is not redistributed from this repo. |
| 16 | +
|
| 17 | +Usage: |
| 18 | + python3 chimerax_render.py map-model --map M.mrc --model X.cif --level 0.03 --out fig.png |
| 19 | + python3 chimerax_render.py turntable --map M.mrc --model X.cif --level 0.03 --out spin.mp4 |
| 20 | + python3 chimerax_render.py local-res --map M.mrc --localres LR.mrc --out lr.png |
| 21 | + python3 chimerax_render.py orthoslices --map M.mrc --out ortho.png |
| 22 | + python3 chimerax_render.py density-zone --map M.mrc --model X.cif --range 6 --out zone.png |
| 23 | +Add --rest to drive a GUI ChimeraX on macOS; --execute to actually run (else |
| 24 | +prep-only: writes the .cxc and the command). |
| 25 | +""" |
| 26 | + |
| 27 | +from __future__ import annotations |
| 28 | + |
| 29 | +import argparse |
| 30 | +import glob |
| 31 | +import os |
| 32 | +import platform |
| 33 | +import shutil |
| 34 | +import subprocess |
| 35 | +import sys |
| 36 | +import time |
| 37 | + |
| 38 | + |
| 39 | +def find_chimerax(): |
| 40 | + for c in ("chimerax", "ChimeraX"): |
| 41 | + p = shutil.which(c) |
| 42 | + if p: |
| 43 | + return p |
| 44 | + apps = sorted(glob.glob("/Applications/ChimeraX*.app/Contents/MacOS/ChimeraX")) |
| 45 | + return apps[-1] if apps else None |
| 46 | + |
| 47 | + |
| 48 | +SCENE = [ |
| 49 | + "set bgColor white", |
| 50 | + "graphics silhouettes true width 1.4", |
| 51 | + "graphics quality 3", |
| 52 | + "lighting simple", # simple lighting keeps density colours true |
| 53 | +] |
| 54 | + |
| 55 | + |
| 56 | +def model_style_lines(model_id="#2"): |
| 57 | + # works for protein, RNA, DNA; nucleotides as ladder, ligands as stick |
| 58 | + return [ |
| 59 | + f"cartoon {model_id}", |
| 60 | + f"nucleotides {model_id} ladder", |
| 61 | + f"style {model_id} & ligand stick", |
| 62 | + f"color {model_id} bychain", |
| 63 | + ] |
| 64 | + |
| 65 | + |
| 66 | +def cxc_map_model(a): |
| 67 | + lines = ["close session", f"windowsize {a.width} {a.height}"] + SCENE |
| 68 | + lines += [f"open {a.map}", f"open {a.model}"] |
| 69 | + lvl = f"level {a.level}" if a.level is not None else "sdLevel 4" |
| 70 | + style = "mesh" if a.map_style == "mesh" else "surface" |
| 71 | + lines += [f"volume #1 style {style} {lvl} step 1 color #9aa6b2"] |
| 72 | + if a.map_style == "surface": |
| 73 | + lines += ["transparency #1 55"] |
| 74 | + lines += model_style_lines("#2") |
| 75 | + lines += ["view", f"save {a.out} width {a.width} height {a.height} supersample 3"] |
| 76 | + return lines |
| 77 | + |
| 78 | + |
| 79 | +def cxc_map_surface(a): |
| 80 | + lines = ["close session", f"windowsize {a.width} {a.height}"] + SCENE |
| 81 | + lines += [f"open {a.map}"] |
| 82 | + lvl = f"level {a.level}" if a.level is not None else "sdLevel 4" |
| 83 | + lines += [f"volume #1 style surface {lvl} step 1 color #7fb3d5", |
| 84 | + "view", f"save {a.out} width {a.width} height {a.height} supersample 3"] |
| 85 | + return lines |
| 86 | + |
| 87 | + |
| 88 | +def cxc_turntable(a): |
| 89 | + lines = ["close session", f"windowsize {a.width} {a.height}"] + SCENE |
| 90 | + lines += [f"open {a.map}", f"open {a.model}"] |
| 91 | + lvl = f"level {a.level}" if a.level is not None else "sdLevel 4" |
| 92 | + lines += [f"volume #1 style mesh {lvl} step 1 color #9aa6b2"] |
| 93 | + lines += model_style_lines("#2") |
| 94 | + lines += ["view", "movie record", |
| 95 | + f"turn y {round(360.0/a.frames, 3)} {a.frames}", f"wait {a.frames}", |
| 96 | + f"movie encode {a.out} framerate {a.fps} quality higher"] |
| 97 | + return lines |
| 98 | + |
| 99 | + |
| 100 | +def cxc_local_res(a): |
| 101 | + # color the model/map surface by an accompanying local-resolution map |
| 102 | + lines = ["close session", f"windowsize {a.width} {a.height}"] + SCENE |
| 103 | + lines += [f"open {a.map}", f"open {a.localres}"] |
| 104 | + lvl = f"level {a.level}" if a.level is not None else "sdLevel 4" |
| 105 | + lines += [f"volume #1 style surface {lvl} step 1", |
| 106 | + "color sample #1 map #2 palette 2.5,blue:3.0,cyan:3.5,green:4.0,yellow:5.0,red", |
| 107 | + "key blue:2.5 cyan:3.0 green:3.5 yellow:4.0 red:5.0 pos 0.80,0.08 size 0.18,0.025", |
| 108 | + "2dlabels create lrlabel text 'local resolution (A)' xpos 0.78 ypos 0.12 size 22", |
| 109 | + "view", f"save {a.out} width {a.width} height {a.height} supersample 3"] |
| 110 | + return lines |
| 111 | + |
| 112 | + |
| 113 | +def cxc_orthoslices(a): |
| 114 | + lines = ["close session", f"windowsize {a.width} {a.height}"] + SCENE |
| 115 | + lines += [f"open {a.map}", "volume #1 style image orthoplanes xyz", |
| 116 | + "view", f"save {a.out} width {a.width} height {a.height} supersample 3"] |
| 117 | + return lines |
| 118 | + |
| 119 | + |
| 120 | +def cxc_density_zone(a): |
| 121 | + lines = ["close session", f"windowsize {a.width} {a.height}"] + SCENE |
| 122 | + lines += [f"open {a.map}", f"open {a.model}", |
| 123 | + f"volume zone #1 nearAtoms #2 range {a.range} newMap true"] |
| 124 | + lvl = f"level {a.level}" if a.level is not None else "sdLevel 4" |
| 125 | + lines += [f"volume #3 style mesh {lvl} step 1 color #9aa6b2"] |
| 126 | + lines += model_style_lines("#2") |
| 127 | + lines += ["view", f"save {a.out} width {a.width} height {a.height} supersample 3"] |
| 128 | + return lines |
| 129 | + |
| 130 | + |
| 131 | +BUILDERS = { |
| 132 | + "map-model": cxc_map_model, "map-surface": cxc_map_surface, |
| 133 | + "turntable": cxc_turntable, "local-res": cxc_local_res, |
| 134 | + "orthoslices": cxc_orthoslices, "density-zone": cxc_density_zone, |
| 135 | +} |
| 136 | + |
| 137 | + |
| 138 | +def run_offscreen(chimerax, cxc_path): |
| 139 | + base = [chimerax, "--offscreen", "--nogui", "--script", cxc_path, "--exit"] |
| 140 | + if shutil.which("xvfb-run") and platform.system() == "Linux": |
| 141 | + return subprocess.run(["xvfb-run", "-a", "-s", "-screen 0 1920x1920x24"] + base) |
| 142 | + return subprocess.run(base) |
| 143 | + |
| 144 | + |
| 145 | +def run_rest(chimerax, lines, out): |
| 146 | + """macOS GUI + REST. Pops a ChimeraX window; renders via GPU.""" |
| 147 | + import socket |
| 148 | + import urllib.parse |
| 149 | + import urllib.request |
| 150 | + s = socket.socket(); s.bind(("127.0.0.1", 0)); port = s.getsockname()[1]; s.close() |
| 151 | + proc = subprocess.Popen([chimerax, "--cmd", |
| 152 | + f"remotecontrol rest start port {port} json true log false"]) |
| 153 | + base = f"http://127.0.0.1:{port}/run" |
| 154 | + try: |
| 155 | + for _ in range(40): |
| 156 | + try: |
| 157 | + urllib.request.urlopen(base + "?command=" + urllib.parse.quote("version"), timeout=2) |
| 158 | + break |
| 159 | + except Exception: |
| 160 | + time.sleep(0.5) |
| 161 | + for cmd in lines: |
| 162 | + urllib.request.urlopen(base + "?command=" + urllib.parse.quote(cmd), timeout=240).read() |
| 163 | + # 0-byte PNG race fix: wait + re-poll |
| 164 | + if os.path.exists(out): |
| 165 | + os.unlink(out) |
| 166 | + for _ in range(16): |
| 167 | + if os.path.exists(out) and os.path.getsize(out) > 0: |
| 168 | + break |
| 169 | + time.sleep(0.5) |
| 170 | + finally: |
| 171 | + try: |
| 172 | + urllib.request.urlopen(base + "?command=" + urllib.parse.quote("exit"), timeout=5) |
| 173 | + except Exception: |
| 174 | + pass |
| 175 | + proc.terminate() |
| 176 | + |
| 177 | + |
| 178 | +def main(): |
| 179 | + ap = argparse.ArgumentParser(description="CryoCore ChimeraX renderer") |
| 180 | + ap.add_argument("kind", choices=list(BUILDERS)) |
| 181 | + ap.add_argument("--map"); ap.add_argument("--model"); ap.add_argument("--localres") |
| 182 | + ap.add_argument("--level", type=float, default=None) |
| 183 | + ap.add_argument("--map-style", choices=["mesh", "surface"], default="mesh") |
| 184 | + ap.add_argument("--range", type=float, default=6.0) |
| 185 | + ap.add_argument("--frames", type=int, default=180) |
| 186 | + ap.add_argument("--fps", type=int, default=30) |
| 187 | + ap.add_argument("--width", type=int, default=1600) |
| 188 | + ap.add_argument("--height", type=int, default=1600) |
| 189 | + ap.add_argument("--out", required=True) |
| 190 | + ap.add_argument("--cxc-out", default=None, help="where to write the .cxc (default beside --out)") |
| 191 | + ap.add_argument("--rest", action="store_true", help="macOS GUI+REST render (pops a window)") |
| 192 | + ap.add_argument("--execute", action="store_true", help="actually run ChimeraX (else prep-only)") |
| 193 | + a = ap.parse_args() |
| 194 | + |
| 195 | + lines = BUILDERS[a.kind](a) |
| 196 | + cxc_path = a.cxc_out or (os.path.splitext(a.out)[0] + ".cxc") |
| 197 | + os.makedirs(os.path.dirname(os.path.abspath(cxc_path)) or ".", exist_ok=True) |
| 198 | + with open(cxc_path, "w") as fh: |
| 199 | + fh.write("\n".join(lines) + "\n") |
| 200 | + print(f"[chimerax] wrote scene script {cxc_path}") |
| 201 | + |
| 202 | + chimerax = find_chimerax() |
| 203 | + is_mac = platform.system() == "Darwin" |
| 204 | + linux_cmd = f"xvfb-run -a ChimeraX --offscreen --nogui --script {cxc_path} --exit" |
| 205 | + |
| 206 | + if not a.execute: |
| 207 | + print("[chimerax] prep-only (pass --execute to render).") |
| 208 | + print(f"[chimerax] Linux: {linux_cmd}") |
| 209 | + if is_mac: |
| 210 | + print("[chimerax] macOS: add --rest to render via a GUI instance, or use pymol_render.py.") |
| 211 | + return |
| 212 | + |
| 213 | + if not chimerax: |
| 214 | + sys.exit("[chimerax] ChimeraX not found; renderer output was not produced.") |
| 215 | + if is_mac and not a.rest: |
| 216 | + print("[chimerax] macOS offscreen hangs; re-run with --rest (GUI) or use pymol_render.py.") |
| 217 | + print(f"[chimerax] Linux: {linux_cmd}") |
| 218 | + return |
| 219 | + if a.rest and is_mac: |
| 220 | + run_rest(chimerax, lines, a.out) |
| 221 | + else: |
| 222 | + run_offscreen(chimerax, cxc_path) |
| 223 | + ok = os.path.exists(a.out) and os.path.getsize(a.out) > 0 |
| 224 | + print(f"[chimerax] {'wrote ' + a.out if ok else 'no output produced'}") |
| 225 | + |
| 226 | + |
| 227 | +if __name__ == "__main__": |
| 228 | + main() |
0 commit comments