"""Convenience wrapper for MolScript and Raster3D image generation.
`MolScript`_ is a tool to create images for molecular structures with the help of
`Raster3D`_. While the system creates nice "out-of-the-box" images, the workflow
is a bit more involved. So this module exists as a convenience wrapper, producing
(opinionated) images from single function or command line calls.
`MolScript`_ and `Raster3d`_ are not bundled with this module. The source code
can be downloaded `here <https://github.com/pekrau/MolScript/>`_ and
`here <http://skuld.bmsc.washington.edu/raster3d/>`_. Installation instructions
are available on the project web pages and/\u2009or in the source code
distributions.
This module can also be used on the command line as ``ma-make-image``.
"""
# Copyright (c) 2026, SIB - Swiss Institute of Bioinformatics and
# Biozentrum - University of Basel
#
# Permission is hereby granted, free of charge, to any person obtaining a copy
# of this software and associated documentation files (the "Software"), to deal
# in the Software without restriction, including without limitation the rights
# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
# copies of the Software, and to permit persons to whom the Software is
# furnished to do so, subject to the following conditions:
#
# The above copyright notice and this permission notice shall be included in all
# copies or substantial portions of the Software.
#
# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
# SOFTWARE.
from contextlib import contextmanager
from pathlib import Path
import gzip
import os
import re
import subprocess
import tempfile
import gemmi
MOLAUTO_BINARY = os.environ.get("MOLAUTO_BINARY", "molauto")
"""Path to the molauto binary, defaults to ``molauto`` from ``$PATH``.
Can be overridden by setting the ``MOLAUTO_BINARY`` environment variable before
import.
"""
MOLSCRIPT_BINARY = os.environ.get("MOLSCRIPT_BINARY", "molscript")
"""Path to the molscript binary, defaults to ``molscript`` from ``$PATH``.
Can be overridden by setting the ``MOLSCRIPT_BINARY`` environment variable before
import.
"""
RENDER_BINARY = os.environ.get("RENDER_BINARY", "render")
"""Path to the render (Raster3D) binary, defaults to ``render`` from ``$PATH``.
Can be overridden by setting the ``RENDER_BINARY`` environment variable before
import.
"""
[docs]
def run_molscript(script, options=None):
"""Execute molscript without checks or cleanup.
Args:
script (list[str]): Lines of the MolScript input script.
options (list[str], optional): Additional command line options passed
to molscript. Defaults to None.
Returns:
subprocess.CompletedProcess: Result of the molscript run.
Raises:
subprocess.CalledProcessError: If molscript exits with a non-zero
return code.
"""
cmd = [MOLSCRIPT_BINARY]
if options:
cmd.extend(options)
return subprocess.run(
cmd,
input=("\n".join(script) + "\n").encode(),
stdout=subprocess.PIPE,
check=True,
)
[docs]
def run_molauto(input_pdb, options=None):
"""Execute molauto without checks or cleanup.
Args:
input_pdb (~pathlib.Path | str): Path to the input PDB file.
options (list[str], optional): Additional command line options passed
to molauto. Defaults to None.
Returns:
subprocess.CompletedProcess: Result of the MAXIT run.
Raises:
subprocess.CalledProcessError: If molauto exits with a non-zero return
code.
"""
cmd = [MOLAUTO_BINARY]
if options:
cmd.extend(options)
cmd.append(str(input_pdb))
return subprocess.run(cmd, stdout=subprocess.PIPE, check=True)
[docs]
def run_render(script_stdout, png_path, options=None):
"""Execute render without checks or cleanup.
Args:
script_stdout (bytes): stdout of a molscript run.
png_path (~pathlib.Path | str): Path to the output PNG file.
options (list[str], optional): Additional command line options passed
to render. Defaults to None.
Returns:
subprocess.CompletedProcess: Result of the render run.
Raises:
subprocess.CalledProcessError: If render exits with a non-zero return
code.
"""
cmd = [RENDER_BINARY, "-png", png_path]
if options:
cmd.extend(options)
return subprocess.run(
cmd,
input=script_stdout,
capture_output=True,
check=True,
)
def _create_molscript_script(input_pdb, colour_scheme=None):
"""Generate the script lines to be used in molscript.
Args:
input_pdb (~pathlib.Path | str): Path to the input PDB file.
colour_scheme (str, optional): Colour scheme to apply. Currently
supported: ``"chain"``. Defaults to None.
Returns:
list[str]: Lines of the MolScript input script.
"""
input_pdb = Path(input_pdb)
# Create initial script using molauto
if input_pdb.stat().st_size > 1000000:
molauto_process = run_molauto(input_pdb, ["-noligand"])
else:
molauto_process = run_molauto(input_pdb, ["-nice"])
# Edit script
lines = molauto_process.stdout.decode().splitlines()
# reformat floating point numbers
script = []
plot_line = None
for i, line in enumerate(lines):
if line == "plot":
plot_line = i
match = re.match(r"(\s*set planecolour hsb )(\d+e[\-\d]+)(.+)", line)
if match:
line = f"{match.group(1)}{float(match.group(2))}{match.group(3)}"
script.append(line.strip())
if plot_line is None: # pragma: no cover (fallback for molauto failures)
raise RuntimeError("'plot' instruction not found in script")
script.insert(plot_line + 2, "background white;")
# Apply chain colour scheme to script
if colour_scheme and colour_scheme == "chain":
colour_lines = []
made_changes = False
for line in script:
if not made_changes and line.find("read mol") > -1:
colour_lines.append(line)
cols = [
[0.106, 0.62, 0.467],
[0.851, 0.373, 0.008],
[0.459, 0.439, 0.702],
[0.906, 0.161, 0.541],
[0.4, 0.651, 0.118],
[0.902, 0.671, 0.008],
[0.651, 0.463, 0.114],
[0.4, 0.4, 0.4],
]
chns = "ABCDEFGH"
colour_lines.append("set colourparts on, ")
for i, chn in enumerate(chns):
if i > 0:
colour_lines.append(", ")
colour_lines.append(
f'residuecolour chain "{chn}" rgb {cols[i][0]} '
+ f"{cols[i][1]} {cols[i][2]}"
)
colour_lines[-1] += ";"
made_changes = True
continue
if line.find(" planecolour ") == -1 and line.find("rainbow") == -1:
colour_lines.append(line)
script = colour_lines
return script
def _render_image(script, png_path, img_size=400):
"""Render a 2D image from am MolScript script."""
# run molscript
if img_size is None:
img_size = 100
molscript_process = run_molscript(
script, ["-size", str(img_size), str(img_size), "-s", "-r"]
)
# Render the image
png_path = Path(png_path)
run_render(molscript_process.stdout, png_path)
if not png_path.exists() or png_path.stat().st_size < 300:
raise RuntimeError("Image creation failed")
@contextmanager
def _get_pdb_file(input_file):
"""Get a PDB file in case of CIF, compressed files allowed."""
if (
str(input_file)
.lower()
.endswith(
(
".cif.gz",
".cif.gzip",
".cif",
".mmcif.gz",
".mmcif.gzip",
".mmcif",
)
)
):
# CIF route
structure = gemmi.read_structure(
str(input_file), format=gemmi.CoorFormat.Mmcif
)
with tempfile.NamedTemporaryFile(suffix=".pdb") as tfh:
structure.write_pdb(tfh.name)
tfh.flush()
yield Path(tfh.name)
elif str(input_file).lower().endswith((".pdb.gz", ".pdb.gzip")):
# PDB gzip route
with tempfile.NamedTemporaryFile(mode="wb", suffix=".pdb") as tfh:
with gzip.open(input_file, "rb") as pfh:
tfh.write(pfh.read())
tfh.flush()
yield Path(tfh.name)
elif str(input_file).lower().endswith(".pdb"):
# PDB route
yield input_file
else:
# ToDo: make own exception with maxit file extensions
raise ValueError(f"Unsupported file extension: {input_file.suffix}")
[docs]
def coordfile2image(input_file, png_path, colour_scheme=None, img_size=400):
"""Create a 2D image for a CIF/\u2009PDB file (gzip allowed).
Args:
input_file (~pathlib.Path | str): Path to the input PDB or
mmCIF file. Gzip-compressed files are supported (extensions
``.pdb.gz``, ``.pdb.gzip``, ``.cif.gz``, ``.cif.gzip``,
``.mmcif.gz``, ``.mmcif.gzip``).
png_path (~pathlib.Path | str): Path to the output PNG file.
colour_scheme (str, optional): Colour scheme to apply. Currently
supported: ``"chain"``. Defaults to None.
img_size (int, optional): Size of the quadratic image in pixels.
Defaults to 400.
Returns:
None
Raises:
RuntimeError: If image creation fails (PNG missing or too small after
rendering) and if the MolScript script is corrupted.
ValueError: If the file extension of ``input_file`` is not supported.
"""
input_file = Path(input_file)
# Get PDB input
with _get_pdb_file(input_file) as pdb_path:
# Generate script
script = _create_molscript_script(pdb_path, colour_scheme=colour_scheme)
# Generate image
_render_image(script, png_path, img_size=img_size)
[docs]
def main():
"""Entry point for the ``ma-make-image`` command line tool."""
# For main functions we allow bad imports
# pylint: disable=import-outside-toplevel
import sys
from modelarchive import _argparse
from modelarchive import _utils
def _parse_command_line():
"""Get arguments."""
parser = _argparse.MAInOutFileArgumentParser(
"Input coordinate file, CIF/ PDB formatted, gzip supported",
"Path to the output PNG file",
output_metavar="<OUTPUT PNG>",
description="Create a 2D image of a molecular structure via "
+ "molauto, molscript, and render.",
)
parser.add_argument(
"--colour-by-chain",
"-c",
help="Paint each chain in a single colour instead of gradient",
action="store_true",
)
parser.add_argument(
"--image-size",
"-i",
default=400,
help="Size of the image in pixels",
type=int,
)
return parser.parse_args()
def _main():
"""Run as script."""
opts = _parse_command_line()
colour_scheme = None
if opts.colour_by_chain:
colour_scheme = "chain"
try:
coordfile2image(
opts.input,
opts.output,
colour_scheme=colour_scheme,
img_size=opts.image_size,
)
except ValueError:
_utils.abort_msg(
f"Input file '{opts.input}' has an unsupported extension, "
+ "allowed are '.cif', '.mmcif' and '.pdb' plus '.gz' or "
+ "'.gzip'",
2,
)
except RuntimeError:
_utils.abort_msg("Image creation failed", 3)
sys.exit(0)
_main()