biotite 1.5.0__cp311-cp311-manylinux_2_24_x86_64.manylinux_2_28_x86_64.whl

biotite/interface/pymol/shapes.py

@@ -0,0 +1,178 @@
+__name__ = "biotite.interface.pymol"
+__author__ = "Patrick Kunzmann"
+__all__ = ["draw_arrows", "draw_box"]
+
+import numpy as np
+from biotite.interface.pymol.cgo import (
+    _arrayfy,
+    draw_cgo,
+    get_cone_cgo,
+    get_cylinder_cgo,
+    get_multiline_cgo,
+)
+
+
+def draw_arrows(
+    start,
+    end,
+    radius=0.1,
+    head_radius=0.20,
+    head_length=0.5,
+    color=(0.5, 0.5, 0.5),
+    head_color=None,
+    name=None,
+    pymol_instance=None,
+    delete=True,
+):
+    """
+    Draw three-dimensional arrows using *Compiled Graphics Objects* (CGOs).
+
+    Parameters
+    ----------
+    start, end : array-like, shape=(n,3)
+        The start and end position of each arrow.
+    radius, head_radius : float or array-like, shape=(n,), optional
+        The radius of the tail and head for each arrow.
+        Uniform for all arrows, if a single value is given.
+    head_length : float or array-like, shape=(n,), optional
+        The length of each arrow head.
+        Uniform for all arrows, if a single value is given.
+    color, head_color : array-like, shape=(3,) or shape=(n,3), optional
+        The color of the tail and head for each arrow, given as RGB
+        values in the range *(0, 1)*.
+        Uniform for all arrows, if a single value is given.
+        If no `head_color` is given, the arrows are single-colored.
+    name : str, optional
+        The name of the newly created CGO object.
+        If omitted, a unique name is generated.
+    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.
+    delete : bool, optional
+        If set to true, the underlying *PyMOL* object will be removed from the *PyMOL*
+        session, when the returned :class:`PyMOLObject` is garbage collected.
+
+    Returns
+    -------
+    pymol_object : PyMOLObject
+        The created :class:`PyMOLObject` representing the drawn CGOs.
+    """
+    if head_color is None:
+        head_color = color
+
+    start = np.asarray(start)
+    end = np.asarray(end)
+    if start.ndim != 2 or end.ndim != 2:
+        raise IndexError("Expected 2D array for start and end positions")
+    if len(start) != len(end):
+        raise IndexError(
+            f"Got {len(start)} start positions, but expected {len(end)} end positions"
+        )
+    expected_length = len(start)
+    radius = _arrayfy(radius, expected_length, 1)
+    head_radius = _arrayfy(head_radius, expected_length, 1)
+    head_length = _arrayfy(head_length, expected_length, 1)
+    color = _arrayfy(color, expected_length, 2)
+    head_color = _arrayfy(head_color, expected_length, 2)
+
+    normal = (end - start) / np.linalg.norm(end - start, axis=-1)[:, np.newaxis]
+    middle = end - normal * head_length[:, np.newaxis]
+
+    cgo_list = []
+    for i in range(len(start)):
+        cgo_list.extend(
+            [
+                get_cylinder_cgo(start[i], middle[i], radius[i], color[i], color[i]),
+                get_cone_cgo(
+                    middle[i],
+                    end[i],
+                    head_radius[i],
+                    0.0,
+                    head_color[i],
+                    head_color[i],
+                    True,
+                    False,
+                ),
+            ]
+        )
+
+    return draw_cgo(cgo_list, name, pymol_instance, delete)
+
+
+def draw_box(
+    box,
+    color=(0, 1, 0),
+    width=1.0,
+    origin=None,
+    name=None,
+    pymol_instance=None,
+    delete=True,
+):
+    """
+    Draw a box using *Compiled Graphics Objects* (CGOs).
+
+    This can be used to draw the unit cell or periodic box of an
+    :class:`AtomArray`.
+
+    Parameters
+    ----------
+    box : array-like, shape=(3,3)
+        The three box vectors.
+    color : array-like, shape=(3,), optional
+        The color of the box, given as RGB
+        values in the range *(0, 1)*.
+    width : float, optional
+        The width of the drawn lines.
+    origin : array-like, shape=(3,), optional
+        If given the box origin is drawn at the given position instead of the
+        coordinate origin.
+    name : str, optional
+        The name of the newly created CGO object.
+        If omitted, a unique name is generated.
+    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.
+    delete : bool, optional
+        If set to true, the underlying *PyMOL* object will be removed from the *PyMOL*
+        session, when the returned :class:`PyMOLObject` is garbage collected.
+
+    Returns
+    -------
+    pymol_object : PyMOLObject
+        The created :class:`PyMOLObject` representing the drawn CGOs.
+    """
+    box = np.asarray(box)
+    if origin is None:
+        origin = np.zeros(3)
+    else:
+        origin = np.asarray(origin)
+
+    starts = []
+    ends = []
+    for direction_dim in (0, 1, 2):
+        plane_dim1, plane_dim2 = [dim for dim in (0, 1, 2) if dim != direction_dim]
+        starts.append(origin)
+        ends.append(origin + box[direction_dim])
+
+        starts.append(origin + box[plane_dim1])
+        ends.append(origin + box[plane_dim1] + box[direction_dim])
+
+        starts.append(origin + box[plane_dim2])
+        ends.append(origin + box[plane_dim2] + box[direction_dim])
+
+        starts.append(origin + box[plane_dim1] + box[plane_dim2])
+        ends.append(origin + box[plane_dim1] + box[plane_dim2] + box[direction_dim])
+
+    return draw_cgo(
+        [get_multiline_cgo(starts, ends, color, width)], name, pymol_instance, delete
+    )

biotite/interface/pymol/startup.py

@@ -0,0 +1,169 @@
+__name__ = "biotite.interface.pymol"
+__author__ = "Patrick Kunzmann"
+__all__ = [
+    "get_and_set_pymol_instance",
+    "launch_pymol",
+    "launch_interactive_pymol",
+    "reset",
+    "setup_parameters",
+    "DuplicatePyMOLError",
+]
+
+
+_pymol = None
+
+
+def get_and_set_pymol_instance(pymol_instance=None):
+    """
+    Get the global *PyMOL* instance.
+
+    This function is intended for internal purposes and should only be
+    used for advanced usages.
+
+    Parameters
+    ----------
+    pymol_instance : module or PyMOL, optional
+        If a ``PyMOL`` instance is given here, the global instance is set
+        to this instance.
+        If *PyMOL* is already running and both instances are not the
+        same, an exception is raised.
+        By default *PyMOL* is started in library mode, if no *PyMOL*
+        instance is currently running.
+
+    Returns
+    -------
+    pymol_instance : module or PyMOL
+        The global ``pymol`` instance.
+    """
+    global _pymol
+    if pymol_instance is None:
+        if not is_launched():
+            _pymol = launch_pymol()
+        return _pymol
+    elif _pymol is None:
+        if not hasattr(pymol_instance, "cmd"):
+            raise TypeError("Given object is not a PyMOL instance")
+        _pymol = pymol_instance
+    elif _pymol is not pymol_instance:
+        # Both the global pymol instance and the given instance are not
+        # the same -> duplicate PyMOL instances
+        raise DuplicatePyMOLError("A PyMOL instance is already running")
+    return _pymol
+
+
+def is_launched():
+    """
+    Check whether a *PyMOL* session is already running.
+
+    Returns
+    -------
+    running : bool
+        True, if a *PyMOL* instance is already running, false otherwise.
+    """
+    return _pymol is not None
+
+
+def launch_pymol():
+    """
+    Launch *PyMOL* in object-oriented library mode.
+
+    This is the recommended way to launch *PyMOL* if no GUI is
+    required.
+    This function simply creates a :class:`PymMOL` object,
+    calls its :func:`start()` method and sets up necessary parameters using
+    :func:`setup_parameters()`.
+
+    Returns
+    -------
+    pymol : PyMOL
+        The started *PyMOL* instance.
+        *PyMOL* commands can be invoked by using its :attr:`cmd` attribute.
+    """
+    from pymol2 import PyMOL
+
+    global _pymol
+
+    if is_launched():
+        raise DuplicatePyMOLError("A PyMOL instance is already running")
+    else:
+        _pymol = PyMOL()
+        _pymol.start()
+        setup_parameters(_pymol)
+    return _pymol
+
+
+def launch_interactive_pymol(*args):
+    """
+    Launch a *PyMOL* GUI with the given command line arguments.
+
+    It starts *PyMOL* by calling :func:`pymol.finish_launching()`,
+    reinitializes *PyMOL* to clear the workspace and sets up necessary
+    parameters using :func:`setup_parameters()`.
+
+    Parameters
+    ----------
+    *args : str
+        The command line options given to *PyMOL*.
+
+    Returns
+    -------
+    pymol : module
+        The :mod:`pymol` module.
+        *PyMOL* commands can be invoked by using its :attr:`cmd`
+        attribute.
+    """
+    import pymol
+
+    global _pymol
+
+    if is_launched():
+        if _pymol is not pymol:
+            raise DuplicatePyMOLError("PyMOL is already running in library mode")
+        else:
+            raise DuplicatePyMOLError("A PyMOL instance is already running")
+    else:
+        pymol.finish_launching(["pymol"] + list(args))
+        _pymol = pymol
+        pymol.cmd.reinitialize()
+        setup_parameters(_pymol)
+    return pymol
+
+
+def reset():
+    """
+    Delete all objects in the PyMOL workspace and reset parameters to
+    defaults.
+
+    If *PyMOL* is not yet running, launch *PyMOL* in object-oriented
+    library mode.
+    """
+    global _pymol
+
+    if not is_launched():
+        _pymol = launch_pymol()
+    _pymol.cmd.reinitialize()
+    setup_parameters(_pymol)
+
+
+def setup_parameters(pymol_instance):
+    """
+    Sets *PyMOL* parameters that are necessary for *Biotite* to interact
+    properly with *PyMOL*.
+
+    Parameters
+    ----------
+    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.
+    """
+    # The selections only work properly,
+    # if the order stays the same after adding a model to PyMOL
+    pymol_instance.cmd.set("retain_order", 1)
+    # Mute the PyMOL messages when rendering movies
+    pymol_instance.cmd.feedback("disable", "movie", "everything")
+
+
+class DuplicatePyMOLError(Exception):
+    pass

biotite/interface/rdkit/__init__.py

@@ -0,0 +1,19 @@
+# This source code is part of the Biotite package and is distributed
+# under the 3-Clause BSD License. Please see 'LICENSE.rst' for further
+# information.
+
+"""
+This subpackage provides an interface to the `RDKit <https://www.rdkit.org/>`_
+cheminformatics package.
+It allows conversion between :class:`.AtomArray` and :class:`rdkit.Chem.rdchem.Mol`
+objects.
+"""
+
+__name__ = "biotite.interface.rdkit"
+__author__ = "Patrick Kunzmann"
+
+from biotite.interface.version import require_package
+
+require_package("rdkit")
+
+from .mol import *