Skip to content

Commit cd57cff

Browse files
committed
faster rendering with vulcan
1 parent d8c4ab7 commit cd57cff

6 files changed

Lines changed: 216 additions & 94 deletions

File tree

Lines changed: 87 additions & 76 deletions
Original file line numberDiff line numberDiff line change
@@ -1,7 +1,8 @@
11
# Render a molecule to an image file
22

3-
**Goal:** produce a PNG of a molecule with no viewer window, for a script, a
4-
report, or an agent that cannot look at a screen, using {py:meth}`~moleculekit.molecule.Molecule.render`.
3+
**Goal:** save a picture of a molecule to a file, without opening a viewer
4+
window, using {py:meth}`~moleculekit.molecule.Molecule.render`. Useful in a
5+
script that produces figures, or anywhere no screen is available.
56

67
## Example
78

@@ -16,93 +17,103 @@ mol.reps.add(sel="resname BEN", style="Licorice", color=4)
1617
mol.render("trypsin.png", size=(1200, 900), rotate="top", center="resname BEN")
1718
```
1819

19-
Passing no path returns the PNG bytes instead, which suits writing to a stream:
20+
Leave out the file name and you get the image data back instead, which is handy
21+
for writing it somewhere else:
2022

2123
```python
2224
png = mol.render(size=(800, 600))
2325
```
2426

25-
## Parameters
27+
## Options
2628

27-
| Parameter | Meaning |
29+
| Option | What it does |
2830
| --- | --- |
29-
| `output` | Path to write. Omit to get the PNG bytes back. |
30-
| `size` | `(width, height)` in pixels. |
31-
| `quality` | `"fast"` or `"high"`. `"high"` adds ambient occlusion, which is close to free on a GPU and costs roughly three times the time on the software fallback. |
32-
| `center` | Atom selection to frame the camera on. |
33-
| `rotate` | `"front"`, `"back"`, `"left"`, `"right"`, `"top"`, `"bottom"`, or `(rx, ry, rz)` in degrees. |
34-
| `zoom` | Larger values move the camera closer. |
35-
| `clip` | Half-thickness in Angstrom of the slab drawn around what the camera frames, cutting away what is in front of and behind it. Omit to draw the whole structure. |
36-
| `fog` | Depth cueing strength, `0` to `100`. Fog fades distant geometry into the background colour. Omit for Mol\*'s own strength. |
37-
| `background` | Colour name or hex string. Ignored when `transparent=True`. |
38-
| `transparent` | Render onto a transparent background, ignoring `background`. |
39-
| `timeout` | Seconds allowed for one render. |
31+
| `output` | File to write. Leave it out to get the image data back. |
32+
| `size` | Width and height of the image in pixels. |
33+
| `quality` | `"fast"` or `"high"`. `"high"` adds soft shading in the crevices and smoother edges, and takes noticeably longer without a graphics card. |
34+
| `center` | Atom selection to point the camera at, such as `"resname BEN"`. |
35+
| `rotate` | `"front"`, `"back"`, `"left"`, `"right"`, `"top"`, `"bottom"`, or your own `(x, y, z)` rotation in degrees. |
36+
| `zoom` | How close the camera sits. Larger is closer. |
37+
| `clip` | Show only a slab this many Angstrom thick around what the camera is pointed at, cutting away what is in front and behind. Leave it out to show the whole structure. |
38+
| `fog` | How much distant parts fade into the background, from `0` for none to `100` for the most. |
39+
| `background` | Background colour, by name or as `#rrggbb`. |
40+
| `transparent` | Give the image a transparent background instead of a colour. |
41+
| `timeout` | How many seconds to allow for one image before giving up. |
42+
43+
## Choosing what is drawn
44+
45+
Representations come from `mol.reps`, exactly as for
46+
{py:meth}`~moleculekit.molecule.Molecule.view`, so a molecule renders the way
47+
you would see it on screen. Leave `mol.reps` empty and you get a sensible
48+
default picture: a cartoon for the protein or nucleic acid, with ligands, ions
49+
and waters drawn as ball-and-stick.
50+
51+
Adding a representation replaces that default picture rather than adding to it,
52+
so if you only want to change one thing, start from the default with
53+
{py:meth}`~moleculekit.representations.Representations.addDefaults` and edit
54+
it:
4055

41-
## Requirements
42-
43-
Rendering needs a chromium binary and the optional `molviewspec` dependency:
44-
45-
```
46-
pip install moleculekit[notebook]
56+
```python
57+
mol.reps.addDefaults()
58+
mol.reps.remove(1) # drop the waters, ligand and ions
59+
mol.reps.add("resname BEN", "VDW", "Name")
4760
```
4861

49-
moleculekit looks for `chromium`, `chromium-browser`, `google-chrome` or
50-
`chrome` on your PATH. To use a specific binary, set `MOLECULEKIT_CHROMIUM` to
51-
its path. No network connection is needed: the Mol* bundle ships with
52-
moleculekit and the structure never leaves the machine.
53-
54-
## Choosing the graphics backend
62+
See [Choose representations and colours](choose-representations.md) for the
63+
available styles, colours, sizes and transparency.
5564

56-
Rendering uses your GPU when one is reachable and falls back to a software
57-
rasteriser when it is not, which is what makes it work unchanged inside
58-
containers that expose no graphics device. The difference is large: a 1200x900
59-
render measured 0.5 seconds on a GPU against 6.9 seconds in software.
65+
## Requirements
6066

61-
The choice is automatic. To pin it, set `MOLECULEKIT_RENDER_GL`:
67+
Rendering uses the Chromium browser behind the scenes, so it needs Chromium
68+
(or Chrome) installed, plus one optional package:
6269

6370
```
64-
MOLECULEKIT_RENDER_GL=hardware # require the GPU, fail if it is unavailable
65-
MOLECULEKIT_RENDER_GL=software # always use the software rasteriser
66-
MOLECULEKIT_RENDER_GL=auto # the default
71+
pip install moleculekit[notebook]
6772
```
6873

69-
Forcing `hardware` is useful when a render unexpectedly falls back and you want
70-
the failure to be loud rather than slow. Forcing `software` is useful for
71-
reproducing what a GPU-less machine will produce.
72-
73-
## Gotchas
74-
75-
- **Representations come from `mol.reps`**, the same ones {py:meth}`~moleculekit.molecule.Molecule.view` uses, and
76-
they replace the automatic scene rather than adding to it. Setting any
77-
representation means you are describing the whole picture, which matches how
78-
the VMD backend behaves. Leave `mol.reps` empty to get the automatic scene of
79-
a cartoon polymer with ball-and-stick ligands, or call {py:meth}`~moleculekit.representations.Representations.addDefaults` to start
80-
from it and edit. See
81-
[Choose representations and colours](choose-representations.md) for the
82-
styles, colour modes, size and transparency.
83-
- **A representation whose selection matches no atoms is dropped with a
84-
warning.** If every representation matches nothing, `render()` raises rather
85-
than writing an empty image.
86-
- **A representation's `frames` argument is ignored.** {py:meth}`~moleculekit.representations.Representations.add`'s `frames`
87-
selects frames for the interactive viewer's trajectory playback; a render is a
88-
single still, so `frames` has nothing to act on.
89-
- **The frame rendered is `mol.frame`.** For a trajectory, set `mol.frame = n`
90-
first. `render()` produces a single still, not a movie.
91-
- **`center` must match at least one atom.** A selection that matches nothing
92-
raises `ValueError`, rather than silently falling back to the default view.
93-
- **`center` is also what `rotate` orbits and what `zoom` scales.** Framing a
94-
ligand and stepping `rotate` through angles gives a turntable around it.
95-
- **The framing follows the requested `size`.** A wider image shows more around
96-
the structure rather than the same picture stretched.
97-
- **The first call is slower** because it starts the browser. Later calls reuse
98-
it, so rendering many images in one process costs the startup once.
99-
- **Speed depends on whether a GPU is available.** With one, a 1200x900 render
100-
takes well under a second. Without one, rendering falls back to a software
101-
rasteriser that fills every pixel on the CPU, which costs several seconds and
102-
scales with pixel count, so a 4K image costs roughly four times a 1080p one.
103-
See "Choosing the graphics backend" below.
104-
- **Call {py:meth}`~moleculekit.molecule.Molecule.render` from one thread at a time.** The headless browser is a
105-
singleton with a single devtools connection, so concurrent calls from
106-
different threads (a multi-threaded service handling several requests at
107-
once, for example) can read each other's responses. Serialise calls to
108-
`render()`, for example behind a lock, if more than one thread might call it.
74+
moleculekit finds `chromium`, `chromium-browser`, `google-chrome` or `chrome`
75+
on your PATH by itself. To point it at a particular one, set the
76+
`MOLECULEKIT_CHROMIUM` environment variable to that program's location.
77+
78+
Nothing is uploaded anywhere: no internet connection is needed and your
79+
structure never leaves your machine.
80+
81+
## If rendering feels slow
82+
83+
With a graphics card, an image takes well under a second. Without one the work
84+
falls to the processor and takes a few seconds, or longer on a machine with few
85+
cores, which is common inside containers and on cloud notebooks.
86+
87+
Three things help, in order of effect:
88+
89+
- **Ask for a smaller image.** The time scales with the number of pixels, so
90+
halving the width and height makes it roughly four times faster.
91+
- **Use `quality="fast"`**, which is the default. `"high"` is worth it for a
92+
final figure and wasteful for a quick look.
93+
- **Render several images in one script.** The first one is slower because the
94+
browser has to start; every image after that reuses it.
95+
96+
## Good to know
97+
98+
- **The picture shows the current frame**, `mol.frame`. For a trajectory, set
99+
`mol.frame = n` first and render again for another frame. Each call makes one
100+
still image, not a movie.
101+
- **`center` is also what `rotate` turns around and what `zoom` scales.**
102+
Pointing at a ligand and stepping `rotate` through several angles gives you a
103+
turntable around it:
104+
105+
```python
106+
for angle in (0, 90, 180, 270):
107+
mol.render(f"ligand_{angle}.png", center="resname BEN", rotate=(0, angle, 0))
108+
```
109+
110+
- **`center` has to match at least one atom**, otherwise you get an error
111+
rather than a picture of something you did not ask for.
112+
- **The framing follows the size you ask for.** A wider image shows more around
113+
the molecule rather than the same picture stretched.
114+
- **A representation matching no atoms is skipped with a warning**, and if none
115+
of them match anything you get an error instead of a blank image.
116+
- **`frames` has no effect here.** It chooses frames for trajectory playback in
117+
the interactive viewer, and a rendered image is a single still.
118+
- **Render one image at a time.** If your program uses threads, do not call
119+
render from several of them at once.

‎moleculekit/viewer/molstar/render.py‎

Lines changed: 68 additions & 14 deletions
Original file line numberDiff line numberDiff line change
@@ -9,8 +9,9 @@
99
Rendering runs on the GPU when one is reachable and falls back to a software
1010
rasteriser when it is not, which is what lets the same code work inside
1111
containers that expose no graphics device. MOLECULEKIT_RENDER_GL pins the
12-
choice. The fallback is not free: a 1200x900 render measured 0.5s on a GPU
13-
against 6.9s in software.
12+
choice. The fallback is not free: a 1200x900 "fast" render measured 0.16s on a
13+
GPU against 1.2s on a 20-core software rasteriser, and software time scales
14+
with pixels and with how many cores it has.
1415
1516
This module is not thread-safe: the singleton browser's devtools session has
1617
one socket and one frame reader, so concurrent render() calls from different
@@ -51,9 +52,13 @@
5152
_PAGE_READY_TIMEOUT = 30.0
5253
_POSIX = os.name == "posix"
5354

55+
# sample_level is the anti-aliasing work: the scene is drawn 2**level times
56+
# and accumulated. Mol*'s own screenshot helper hardcodes 4, so 16 renders,
57+
# which costs nothing on a GPU and dominates the time on the software
58+
# rasteriser. "fast" spends 2 renders on it instead.
5459
QUALITY_PRESETS = {
55-
"fast": {"occlusion": False},
56-
"high": {"occlusion": True},
60+
"fast": {"occlusion": False, "sample_level": 1},
61+
"high": {"occlusion": True, "sample_level": 4},
5762
}
5863

5964
_CHROMIUM_FLAGS = (
@@ -85,12 +90,25 @@
8590
"--mute-audio",
8691
)
8792

88-
# How chromium is told to reach OpenGL. The choice is worth roughly 15x on
89-
# render time: a 1200x900 "fast" render measured 0.43s on an NVIDIA GPU and
90-
# 6.9s on the software rasteriser, because SwiftShader fills every pixel on
91-
# the CPU.
93+
# How chromium is told to reach OpenGL. The choice is worth roughly 7x on
94+
# render time: a 1200x900 "fast" render measured 0.16s on an NVIDIA GPU and
95+
# 1.2s on the software rasteriser, because SwiftShader fills every pixel on
96+
# the CPU, across as many cores as it has.
9297
_GL_BACKENDS = {
9398
"hardware": ("--use-gl=angle", "--use-angle=gl"),
99+
# Mesa's software Vulkan (lavapipe). Measured about 30% faster than
100+
# SwiftShader on the same GPU-less container: 1M63 at 1000x750 took a
101+
# median 1.7s against 2.4s. Only the Vulkan path works headless, because
102+
# ANGLE's GL backend wants an X display and says so ("Could not open the
103+
# default X display"), so llvmpipe cannot be reached through OpenGL here.
104+
"software-vulkan": (
105+
"--use-gl=angle",
106+
"--use-angle=vulkan",
107+
# Chromium refuses a software Vulkan device without these, and then
108+
# reports no WebGL context at all rather than falling back.
109+
"--ignore-gpu-blocklist",
110+
"--enable-features=Vulkan",
111+
),
94112
"software": (
95113
"--use-gl=angle",
96114
"--use-angle=swiftshader",
@@ -102,6 +120,8 @@
102120
_GL_ENV_VAR = "MOLECULEKIT_RENDER_GL"
103121
# A GPU chromium can reach appears here as a DRM render node.
104122
_DRM_DIR = Path("/dev/dri")
123+
# Where Linux advertises Vulkan drivers; Mesa's software one is lvp.
124+
_VULKAN_ICD_DIR = Path("/usr/share/vulkan/icd.d")
105125

106126

107127
@dataclass
@@ -136,12 +156,35 @@ def _hardware_gl_present() -> bool:
136156
return any(node.name.startswith("renderD") for node in _DRM_DIR.iterdir())
137157

138158

159+
def _lavapipe_icd() -> str | None:
160+
"""Mesa's software Vulkan driver, when the machine has one installed.
161+
162+
Linux only, which is where this matters: the GPU-less containers that pay
163+
for software rendering are Linux, while Windows registers Vulkan drivers in
164+
the registry and macOS puts them elsewhere. On those the lookup simply
165+
finds nothing and rendering falls back to SwiftShader, which needs no
166+
driver installed.
167+
168+
Returns
169+
-------
170+
icd : str or None
171+
Path to the lavapipe ICD manifest, or None when Mesa's Vulkan drivers
172+
are not installed (Debian and Ubuntu ship them in mesa-vulkan-drivers).
173+
"""
174+
if not _VULKAN_ICD_DIR.is_dir():
175+
return None
176+
for icd in sorted(_VULKAN_ICD_DIR.glob("lvp_icd*.json")):
177+
return str(icd)
178+
return None
179+
180+
139181
def _gl_backend_order() -> list[str]:
140182
"""The GL backends to try, in order, until one yields a WebGL context.
141183
142-
``MOLECULEKIT_RENDER_GL`` pins the choice to ``hardware`` or ``software``.
143-
Unset (or ``auto``) picks hardware first when a GPU looks reachable, and
144-
always keeps software as the fallback.
184+
``MOLECULEKIT_RENDER_GL`` pins the choice to one of ``_GL_BACKENDS``.
185+
Unset (or ``auto``) picks hardware first when a GPU looks reachable, then
186+
Mesa's software Vulkan where it is installed, and always keeps SwiftShader
187+
as the fallback since it needs nothing installed.
145188
146189
Returns
147190
-------
@@ -161,9 +204,12 @@ def _gl_backend_order() -> list[str]:
161204
f"{_GL_ENV_VAR}={requested!r} is not a known GL backend. Use one of "
162205
f"{sorted(_GL_BACKENDS)}, or 'auto'."
163206
)
164-
if _hardware_gl_present():
165-
return ["hardware", "software"]
166-
return ["software"]
207+
order = ["hardware"] if _hardware_gl_present() else []
208+
if _lavapipe_icd() is not None:
209+
order.append("software-vulkan")
210+
# SwiftShader last and always: it needs nothing installed.
211+
order.append("software")
212+
return order
167213

168214

169215
def _usable_gl(gl_renderer: str | None) -> bool:
@@ -402,6 +448,12 @@ def _start_with_backend(
402448
"""Start the browser on one GL backend. None when it yields no context."""
403449
binary = find_chromium()
404450
profile_dir = tempfile.mkdtemp(prefix="moleculekit-render-")
451+
env = None
452+
if backend == "software-vulkan":
453+
# Pin the driver: with hardware ICDs also present, ANGLE would
454+
# otherwise be free to pick one, and this backend is only ever chosen
455+
# to render in software.
456+
env = {**os.environ, "VK_ICD_FILENAMES": _lavapipe_icd() or ""}
405457
process = subprocess.Popen(
406458
[
407459
binary,
@@ -415,6 +467,7 @@ def _start_with_backend(
415467
],
416468
stdout=subprocess.DEVNULL,
417469
stderr=subprocess.DEVNULL,
470+
env=env,
418471
# New session (POSIX only) so this process is its own process group
419472
# leader: _stop() kills the whole group, taking chromium's helper
420473
# processes down with it instead of leaving them to linger.
@@ -677,6 +730,7 @@ def render(
677730
"width": width,
678731
"height": height,
679732
"occlusion": QUALITY_PRESETS[quality]["occlusion"],
733+
"sampleLevel": QUALITY_PRESETS[quality]["sample_level"],
680734
"transparent": bool(transparent),
681735
}
682736
data_uri = state.ws.evaluate(

0 commit comments

Comments
 (0)