biotite 1.5.0__cp311-cp311-manylinux_2_24_x86_64.manylinux_2_28_x86_64.whl

biotite/interface/pymol/convert.py

@@ -0,0 +1,185 @@
+__name__ = "biotite.interface.pymol"
+__author__ = "Patrick Kunzmann"
+__all__ = ["to_model", "from_model"]
+
+import warnings
+import numpy as np
+from chempy import Atom, Bond
+from chempy.models import Indexed as IndexedModel
+import biotite.structure as struc
+from biotite.interface import LossyConversionWarning
+from biotite.sequence import ProteinSequence
+
+# Bond orders in PyMOL (https://pymolwiki.org/index.php/Valence):
+# 1: single bond, 2: double bond, 3:triple bond, 4:delocalized
+_FORMAL_BOND_ORDER = {
+    struc.BondType.SINGLE: 1,
+    struc.BondType.DOUBLE: 2,
+    struc.BondType.TRIPLE: 3,
+    struc.BondType.AROMATIC_SINGLE: 1,
+    struc.BondType.AROMATIC_DOUBLE: 2,
+}
+_DELOKALIZED_BOND_ORDER = {
+    struc.BondType.SINGLE: 1,
+    struc.BondType.DOUBLE: 2,
+    struc.BondType.TRIPLE: 3,
+    struc.BondType.AROMATIC_SINGLE: 4,
+    struc.BondType.AROMATIC_DOUBLE: 4,
+    struc.BondType.AROMATIC: 4,
+}
+_BIOTITE_BOND_TYPES = {
+    1: struc.BondType.SINGLE,
+    2: struc.BondType.DOUBLE,
+    3: struc.BondType.TRIPLE,
+    4: struc.BondType.AROMATIC,
+}
+
+
+def to_model(atom_array, delocalize_bonds=False):
+    """
+    Convert an :class:`AtomArray` into a :class:`chempy.models.Indexed`
+    object.
+
+    Parameters
+    ----------
+    atom_array : AtomArray
+        The structure to be converted.
+    delocalize_bonds : bool, optional
+        If set to true, use *PyMOL*'s delocalized bond order for aromatic bonds.
+        Otherwise, always use formal bond orders.
+
+    Returns
+    -------
+    chempy_model : Indexed
+        The converted structure.
+    """
+    if not np.isfinite(atom_array.coord).all():
+        raise ValueError("PyMOL does not support infinite or NaN coordinates")
+
+    model = IndexedModel()
+
+    annot_cat = atom_array.get_annotation_categories()
+    for i in range(atom_array.array_length()):
+        atom = Atom()
+        atom.segi = atom_array.chain_id[i]
+        atom.chain = atom_array.chain_id[i]
+        atom.resi_number = atom_array.res_id[i]
+        atom.ins_code = atom_array.ins_code[i]
+        res_name = atom_array.res_name[i]
+        atom.resn = res_name
+        if len(res_name) == 1:
+            atom.resn_code = res_name
+        else:
+            try:
+                atom.resn_code = ProteinSequence.convert_letter_3to1(res_name)
+            except KeyError:
+                atom.resn_code = "X"
+        atom.hetatm = 1 if atom_array.hetero[i] else 0
+        atom.name = atom_array.atom_name[i]
+        atom.symbol = atom_array.element[i]
+        if "b_factor" in annot_cat:
+            atom.b = atom_array.b_factor[i]
+        if "occupancy" in annot_cat:
+            atom.q = atom_array.occupancy[i]
+        if "charge" in annot_cat:
+            atom.formal_charge = atom_array.charge[i]
+        atom.coord = tuple(atom_array.coord[..., i, :])
+        atom.index = i + 1
+        model.add_atom(atom)
+
+    bond_order_mapping = (
+        _DELOKALIZED_BOND_ORDER if delocalize_bonds else _FORMAL_BOND_ORDER
+    )
+    unmappable_bond_types = []
+    if atom_array.bonds is not None:
+        for i, j, bond_type in atom_array.bonds.as_array():
+            bond = Bond()
+            order = bond_order_mapping.get(bond_type)
+            if order is None:
+                unmappable_bond_types.append(bond_type)
+                order = 1
+            bond.order = order
+            bond.index = [i, j]
+            model.add_bond(bond)
+    else:
+        warnings.warn("The given atom array (stack) has no associated bond information")
+    if unmappable_bond_types:
+        warnings.warn(
+            "The following bond types could not be mapped to PyMOL: "
+            + ", ".join([struc.BondType(bt).name for bt in set(unmappable_bond_types)]),
+            LossyConversionWarning,
+        )
+
+    return model
+
+
+def from_model(chempy_model, include_bonds=False):
+    """
+    Convert a :class:`chempy.models.Indexed`
+    object into an :class:`AtomArray`.
+
+    The returned :class:`AtomArray` contains the optional annotation
+    categories ``b_factor``, ``occupancy``, ``charge`` and
+    ``altloc_id``.
+    No *altloc* ID filtering is performed.
+
+    Parameters
+    ----------
+    chempy_model : Indexed
+        The ``chempy`` model.
+    include_bonds : bool, optional
+        If set to true, an associated :class:`BondList` will be created
+        for the returned atom array.
+
+    Returns
+    -------
+    atom_array : AtomArray
+        The converted structure.
+    """
+    atoms = chempy_model.atom
+
+    bonds = chempy_model.bond
+
+    atom_array = struc.AtomArray(len(atoms))
+
+    # Add annotation arrays
+    atom_array.chain_id = np.array([a.chain for a in atoms], dtype="U3")
+    atom_array.res_id = np.array([a.resi_number for a in atoms], dtype=int)
+    atom_array.ins_code = np.array([a.ins_code for a in atoms], dtype="U1")
+    atom_array.res_name = np.array([a.resn for a in atoms], dtype="U3")
+    atom_array.hetero = np.array([a.hetatm for a in atoms], dtype=bool)
+    atom_array.atom_name = np.array([a.name for a in atoms], dtype="U6")
+    atom_array.element = np.array([a.symbol for a in atoms], dtype="U2")
+
+    atom_array.set_annotation(
+        "b_factor",
+        np.array([a.b if hasattr(a, "b") else 0 for a in atoms], dtype=float),
+    )
+    atom_array.set_annotation(
+        "occupancy",
+        np.array([a.q if hasattr(a, "q") else 1.0 for a in atoms], dtype=float),
+    )
+    atom_array.set_annotation(
+        "charge",
+        np.array(
+            [a.formal_charge if hasattr(a, "formal_charge") else 0 for a in atoms],
+            dtype=int,
+        ),
+    )
+    atom_array.set_annotation(
+        "altloc_id",
+        np.array([a.alt if hasattr(a, "alt") else "" for a in atoms], dtype="U1"),
+    )
+
+    # Set coordinates
+    atom_array.coord = np.array([a.coord for a in atoms], dtype=np.float32)
+
+    # Add bonds
+    if include_bonds:
+        bond_array = np.array(
+            [[b.index[0], b.index[1], _BIOTITE_BOND_TYPES[b.order]] for b in bonds],
+            dtype=np.uint32,
+        )
+        atom_array.bonds = struc.BondList(len(atoms), bond_array)
+
+    return atom_array

biotite/interface/pymol/display.py

@@ -0,0 +1,267 @@
+__name__ = "biotite.interface.pymol"
+__author__ = "Patrick Kunzmann"
+__all__ = ["TimeoutError", "RenderError", "show", "play"]
+
+import datetime
+import shutil
+import subprocess
+import tempfile
+import time
+from os import remove
+from os.path import getsize, join
+from biotite.interface.pymol.startup import get_and_set_pymol_instance
+
+_INTERVAL = 0.1
+
+
+class TimeoutError(Exception):
+    """
+    Exception that is raised after time limit expiry in :func:`show()`.
+    """
+
+    pass
+
+
+class RenderError(Exception):
+    """
+    Exception that is raised when ``imagemagick`` or ``ffmpeg`` fails.
+    """
+
+    pass
+
+
+def show(size=None, use_ray=False, timeout=60.0, pymol_instance=None):
+    """
+    Render an image of the *PyMOL* session and display it in the current
+    *Jupyter* notebook.
+
+    Note that this function works only in a *Jupyter* notebook.
+
+    Parameters
+    ----------
+    size : tuple of (int, int), optional
+        The width and height of the rendered image in pixels.
+        By default, the size of the current *PyMOL* viewport is used.
+    use_ray : bool, optional
+        If set to true, the a ray-traced image is created.
+        This will also increase the rendering time.
+    timeout : float
+        The number of seconds to wait for image output from *PyMOL*.
+    pymol_instance : module or SingletonPyMOL or PyMOL, optional
+        If *PyMOL* is used in library mode, the :class:`PyMOL`
+        or :class:`SingletonPyMOL` object is given here.
+        If otherwise *PyMOL* is used in GUI mode, the :mod:`pymol`
+        module is given.
+        By default the currently active *PyMOL* instance is used.
+        If no *PyMOL* instance is currently running,
+        *PyMOL* is started in library mode.
+
+    Raises
+    ------
+    TimeoutError
+        If no image was created after expiry of the `timeout` limit.
+
+    Notes
+    -----
+    Internally uses the *PyMOL* ``png`` command.
+    """
+    try:
+        from IPython.display import Image
+    except ImportError:
+        raise ImportError("IPython is not installed")
+
+    pymol_instance = get_and_set_pymol_instance(pymol_instance)
+    cmd = pymol_instance.cmd
+
+    if size is None:
+        width = 0
+        height = 0
+    else:
+        width, height = size
+
+    image_file = tempfile.NamedTemporaryFile(
+        delete=False, prefix="biotite_", suffix=".png"
+    )
+    # Close directly and cleanup later
+    # This is necessary, as Windows does not allow writing opened files
+    image_file.close()
+
+    start_time = datetime.datetime.now()
+    cmd.png(image_file.name, width, height, ray=int(use_ray))
+    while True:
+        # After 'timeout' seconds the loop exits with an error
+        if (datetime.datetime.now() - start_time).total_seconds() > timeout:
+            raise TimeoutError("No PNG image was output within the expected time limit")
+        # Check if PyMOL has already written image data to file
+        if getsize(image_file.name) > 0:
+            break
+        time.sleep(_INTERVAL)
+
+    with open(image_file.name, "rb") as f:
+        image_data = f.read()
+    remove(image_file.name)
+    return Image(image_data, embed=True, metadata={"source": "PyMOL"})
+
+
+def play(
+    size=None, fps=30, format="gif", html_attributes="controls", pymol_instance=None
+):
+    """
+    Render an video of the *PyMOL* video frames and display it in the current
+    *Jupyter* notebook.
+
+    Note that this function works only in a *Jupyter* notebook.
+
+    Parameters
+    ----------
+    size : tuple of (int, int), optional
+        The width and height of the rendered video in pixels.
+        By default, the size of the current *PyMOL* viewport is used.
+    fps : int
+        The number of frames per second.
+    format : {"gif", "mp4"}, optional
+        The format of the rendered video.
+        By default, a GIF is created.
+    html_attributes : str, optional
+        The HTML attributes that are passed to the ``<video>`` tag.
+        Only used, if ``format="mp4"``.
+    pymol_instance : module or SingletonPyMOL or PyMOL, optional
+        If *PyMOL* is used in library mode, the :class:`PyMOL`
+        or :class:`SingletonPyMOL` object is given here.
+        If otherwise *PyMOL* is used in GUI mode, the :mod:`pymol`
+        module is given.
+        By default the currently active *PyMOL* instance is used.
+        If no *PyMOL* instance is currently running,
+        *PyMOL* is started in library mode.
+
+    Notes
+    -----
+    Internally uses the *PyMOL* ``mpng`` command.
+    This function requires either the ``ffmpeg`` (``gif`` or ``mp4``)
+    or ``imagemagick`` (``gif``) command line tool to be installed.
+    """
+    try:
+        from IPython.display import Image, Video
+    except ImportError:
+        raise ImportError("IPython is not installed")
+
+    pymol_instance = get_and_set_pymol_instance(pymol_instance)
+    cmd = pymol_instance.cmd
+
+    if size is None:
+        width = 0
+        height = 0
+    else:
+        width, height = size
+
+    with tempfile.TemporaryDirectory(prefix="biotite_") as frame_dir:
+        # Must use ray tracing, as no window is created
+        # Otherwise PyMOL raises 'MoviePNG-Error: Missing rendered image.'
+        cmd.mpng(join(frame_dir, "img_"), mode=2, width=width, height=height)
+        video_data = _create_video(frame_dir, fps, format)
+
+    if format == "mp4":
+        return Video(
+            video_data,
+            embed=True,
+            mimetype="video/mp4",
+            html_attributes=html_attributes,
+        )
+    else:
+        return Image(video_data, embed=True, metadata={"source": "PyMOL"})
+
+
+def _create_video(input_dir, fps, format):
+    """
+    Create a video from the images in the given directory using ``ffmpeg```
+    or ``imagemagick``.
+
+    Parameters
+    ----------
+    input_dir : str
+        The directory containing the frames to be concatenated into a video.
+        The images are consumed in lexographical order.
+    fps : int
+        The number of frames per second.
+    format : {"gif", "mp4"}
+        The format of the video.
+
+    Returns
+    -------
+    video : bytes
+        The video data.
+    """
+    if format == "gif":
+        # GIFs created with 'imagemagick' have less artifacts
+        if _is_installed("magick"):
+            return _create_gif_with_imagemagick(input_dir, fps)
+        elif _is_installed("ffmpeg"):
+            return _create_video_with_ffmpeg(input_dir, fps, format)
+        else:
+            raise RenderError("Neither 'imagemagick' nor 'ffmpeg' is installed")
+    elif format == "mp4":
+        if _is_installed("ffmpeg"):
+            return _create_video_with_ffmpeg(input_dir, fps, format)
+        else:
+            raise RenderError("'ffmpeg' is not installed")
+
+
+def _create_gif_with_imagemagick(input_dir, fps):
+    # See https://usage.imagemagick.org/anim_basics/ for reference
+    completed_process = subprocess.run(
+        [
+            "magick",
+            # GIFs require a multiple of a hundredth of a second to work properly
+            "-delay", str(int(100 / fps)),
+            # Make animation loop infinitely
+            "-loop", "0",
+            # Do not overlay images
+            "-dispose", "Previous",
+            join(input_dir, "*.png"),
+            # Decrease GIF size
+            "-layers", "Optimize",
+            "GIF:-",
+        ],
+        capture_output=True,
+    )  # fmt: skip
+    if completed_process.returncode != 0:
+        raise RenderError(completed_process.stderr.decode())
+    return completed_process.stdout
+
+
+def _create_video_with_ffmpeg(input_dir, fps, format):
+    # 'input_dir' is a temporary directory anyway
+    video_path = join(input_dir, f"video.{format}")
+    completed_process = subprocess.run(
+        [
+            "ffmpeg",
+            "-i", join(input_dir, "img_%04d.png"),
+            "-r", str(fps),
+            # Must be set to obtain a non-corrupted video
+            "-pix_fmt", "yuv420p",
+            video_path,
+        ],
+        capture_output=True,
+    )  # fmt: skip
+    if completed_process.returncode != 0:
+        raise RenderError(completed_process.stderr.decode())
+    with open(video_path, "rb") as f:
+        video_data = f.read()
+    return video_data
+
+
+def _is_installed(program):
+    """
+    Check whether the given program is installed.
+
+    Parameters
+    ----------
+    program : str
+        The name of the program to check.
+
+    Returns
+    -------
+    installed : bool
+        True, if the program is installed, false otherwise.
+    """
+    return shutil.which(program) is not None