biotite 1.0.1__cp312-cp312-win_amd64.whl → 1.2.0__cp312-cp312-win_amd64.whl

biotite/interface/pymol/cgo.py

@@ -0,0 +1,346 @@
+__name__ = "biotite.interface.pymol"
+__author__ = "Patrick Kunzmann"
+__all__ = [
+    "draw_cgo",
+    "get_cylinder_cgo",
+    "get_cone_cgo",
+    "get_sphere_cgo",
+    "get_point_cgo",
+    "get_line_cgo",
+    "get_multiline_cgo",
+]
+
+import itertools
+from enum import IntEnum
+import numpy as np
+from biotite.interface.pymol.object import PyMOLObject
+from biotite.interface.pymol.startup import get_and_set_pymol_instance
+
+_object_counter = 0
+
+
+class CGO(IntEnum):
+    # List compiled from uppercase attributes in 'pymol.cgo'
+    ALPHA = 25
+    ALPHA_TRIANGLE = 17
+    BEGIN = 2
+    CHAR = 23
+    COLOR = 6
+    CONE = 27
+    CUSTOM_CYLINDER = 15
+    CYLINDER = 9
+    DISABLE = 13
+    DOTWIDTH = 16
+    ELLIPSOID = 18
+    ENABLE = 12
+    END = 3
+    FONT = 19
+    FONT_AXES = 22
+    FONT_SCALE = 20
+    FONT_VERTEX = 21
+    LINES = 1
+    LINEWIDTH = 10
+    LINE_LOOP = 2
+    LINE_STRIP = 3
+    NORMAL = 5
+    NULL = 1
+    PICK_COLOR = 31
+    POINTS = 0
+    QUADRIC = 26
+    SAUSAGE = 14
+    SPHERE = 7
+    STOP = 0
+    TRIANGLE = 8
+    TRIANGLES = 4
+    TRIANGLE_FAN = 6
+    TRIANGLE_STRIP = 5
+    VERTEX = 4
+    WIDTHSCALE = 11
+
+
+def draw_cgo(cgo_list, name=None, pymol_instance=None, delete=True):
+    """
+    Draw geometric shapes using *Compiled Graphics Objects* (CGOs).
+
+    Each CGO is represented by a list of floats, which can be obtained
+    via the ``get_xxx_cgo()`` functions.
+
+    Parameters
+    ----------
+    cgo_list : list of list of float
+        The CGOs to draw.
+        It is recommended to use a ``get_xxx_cgo()`` function to obtain
+        the elements for this list, if possible.
+        Otherwise, shapes may be drawn incorrectly or omitted entirely,
+        if a CGO is incorrectly formatted.
+    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.
+    """
+    global _object_counter
+    if name is None:
+        name = f"biotite_cgo_{_object_counter}"
+        _object_counter += 1
+    pymol_instance = get_and_set_pymol_instance(pymol_instance)
+    pymol_instance.cmd.load_cgo(
+        # If CGO values are integers instead of floats
+        # the rendering may fail
+        [float(value) for value in list(itertools.chain(*cgo_list))],
+        name,
+    )
+    return PyMOLObject(name, pymol_instance, delete)
+
+
+def get_cylinder_cgo(start, end, radius, start_color, end_color):
+    """
+    Get the CGO for a cylinder.
+
+    Parameters
+    ----------
+    start, end : array-like, shape=(3,)
+        The start and end position of the cylinder.
+    radius : float
+        The radius of the cylinder.
+    start_color, end_color : array-like, shape=(3,)
+        The color at the start and end of the cylinder given as RGB
+        values in the range *(0, 1)*.
+
+    Returns
+    -------
+    cgo : list of float
+        The CGO representation.
+    """
+    _expect_length(start, "start", 3)
+    _expect_length(end, "end", 3)
+    _expect_length(start_color, "start_color", 3)
+    _expect_length(end_color, "end_color", 3)
+    _check_color(start_color)
+    _check_color(end_color)
+    return [CGO.CYLINDER, *start, *end, radius, *start_color, *end_color]
+
+
+def get_cone_cgo(
+    start, end, start_radius, end_radius, start_color, end_color, start_cap, end_cap
+):
+    """
+    Get the CGO for a cone.
+
+    Parameters
+    ----------
+    start, end : array-like, shape=(3,)
+        The start and end position of the cone.
+    start_radius, end_radius : float
+        The radius of the cone at the start and end.
+    start_color, end_color : array-like, shape=(3,)
+        The color at the start and end of the cone given as RGB
+        values in the range *(0, 1)*.
+    start_cap, end_cap : bool
+        If true, a cap is drawn at the start or end of the cone.
+        Otherwise the cone is displayed as *open*.
+
+    Returns
+    -------
+    cgo : list of float
+        The CGO representation.
+    """
+    _expect_length(start, "start", 3)
+    _expect_length(end, "end", 3)
+    _expect_length(start_color, "start_color", 3)
+    _expect_length(end_color, "end_color", 3)
+    _check_color(start_color)
+    _check_color(end_color)
+    return [
+        CGO.CONE,
+        *start,
+        *end,
+        start_radius,
+        end_radius,
+        *start_color,
+        *end_color,
+        start_cap,
+        end_cap,
+    ]
+
+
+def get_sphere_cgo(pos, radius, color):
+    """
+    Get the CGO for a sphere.
+
+    Parameters
+    ----------
+    pos : array-like, shape=(3,)
+        The position of the sphere.
+    radius : float
+        The radius of the sphere.
+    color : array-like, shape=(3,)
+        The color of the sphere given as RGB values in the range
+        *(0, 1)*.
+
+    Returns
+    -------
+    cgo : list of float
+        The CGO representation.
+    """
+    _expect_length(pos, "pos", 3)
+    _expect_length(color, "color", 3)
+    _check_color(color)
+    return [CGO.COLOR, *color, CGO.SPHERE, *pos, radius]
+
+
+def get_point_cgo(pos, color):
+    """
+    Get the CGO for one or multiple points.
+
+    Parameters
+    ----------
+    pos : array-like, shape=(3,), shape=(n,3)
+        The position(s) of the points.
+    color : array-like, shape=(3,) or shape=(n,3)
+        The color of the point(s) given as RGB values in the range
+        *(0, 1)*.
+        Either one color can be given that is used for all points or
+        an individual color for each point can be supplied.
+
+    Returns
+    -------
+    cgo : list of float
+        The CGO representation.
+    """
+    pos = np.atleast_2d(pos)
+    color = _arrayfy(color, len(pos), 2)
+
+    for p in pos:
+        _expect_length(p, "pos", 3)
+    for c in color:
+        _expect_length(c, "color", 3)
+        _check_color(c)
+
+    vertices = []
+    for p, c in zip(pos, color):
+        vertices += [CGO.COLOR, *c, CGO.VERTEX, *p]
+
+    return [CGO.BEGIN, CGO.POINTS] + vertices + [CGO.END]
+
+
+def get_line_cgo(pos, color, width=1.0):
+    """
+    Get the CGO for a line following the given positions.
+
+    Parameters
+    ----------
+    pos : array-like, shape=(n,3)
+        The line follows these positions.
+    color : array-like, shape=(3,) or shape=(n,3)
+        The color of the line given as RGB values in the range
+        *(0, 1)*.
+        Either one color can be given that is used for all positions or
+        an individual color for each position can be supplied.
+    width : float, optional
+        The rendered width of the line.
+        The width is only visible after calling :func:`ray()`.
+
+    Returns
+    -------
+    cgo : list of float
+        The CGO representation.
+    """
+    color = _arrayfy(color, len(pos), 2)
+
+    for p in pos:
+        _expect_length(p, "pos", 3)
+    for c in color:
+        _expect_length(c, "color", 3)
+        _check_color(c)
+
+    vertices = []
+    for p, c in zip(pos, color):
+        vertices += [CGO.COLOR, *c, CGO.VERTEX, *p]
+
+    return [CGO.LINEWIDTH, width, CGO.BEGIN, CGO.LINE_STRIP] + vertices + [CGO.END]
+
+
+def get_multiline_cgo(start, end, color, width=1.0):
+    """
+    Get the CGO for one or multiple straight lines drawn from given
+    start to end positions.
+
+    Parameters
+    ----------
+    start, end : array-like, shape=(3,) or shape=(n,3)
+        The *n* lines are drawn from the `start` to the `end` positions.
+    color : array-like, shape=(3,) or shape=(n,3)
+        The color of the lines given as RGB values in the range
+        *(0, 1)*.
+        Either one color can be given that is used for all lines or
+        an individual color for each line can be supplied.
+    width : float, optional
+        The rendered width of the lines.
+        The width is only visible after calling :func:`ray()`.
+
+    Returns
+    -------
+    cgo : list of float
+        The CGO representation.
+    """
+    start = np.atleast_2d(start)
+    end = np.atleast_2d(end)
+    color = _arrayfy(color, len(start), 2)
+
+    if len(start) != len(end):
+        raise IndexError(
+            f"{len(start)} start positions are given, but {len(end)} end positions"
+        )
+    for p in start:
+        _expect_length(p, "start", 3)
+    for p in end:
+        _expect_length(p, "end", 3)
+    for c in color:
+        _expect_length(c, "color", 3)
+        _check_color(c)
+
+    vertices = []
+    for p1, p2, c in zip(start, end, color):
+        vertices += [CGO.COLOR, *c, CGO.VERTEX, *p1, CGO.VERTEX, *p2]
+
+    return [CGO.LINEWIDTH, width, CGO.BEGIN, CGO.LINES] + vertices + [CGO.END]
+
+
+def _expect_length(values, name, length):
+    if len(values) != length:
+        raise IndexError(
+            f"'{name}' has {len(values)} values, but {length} were expected"
+        )
+
+
+def _check_color(color):
+    if np.any(color) < 0 or np.any(color) > 1:
+        raise ValueError("Colors must be in range (0, 1)")
+
+
+def _arrayfy(value, length, min_dim):
+    """
+    Expand value(s) to the given number of dimensions and repeat value
+    `length` number of times if only a single value is given.
+    """
+    value = np.array(value, ndmin=min_dim)
+    if len(value) == 1 and length > 1:
+        value = np.repeat(value, length, axis=0)
+    elif len(value) != length:
+        raise IndexError(f"Expected {length} values, but got {len(value)}")
+    return value

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