biotite 1.0.0__cp311-cp311-macosx_11_0_arm64.whl → 1.1.0__cp311-cp311-macosx_11_0_arm64.whl

biotite/structure/alphabet/pb.py

@@ -0,0 +1,171 @@
+# 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.
+
+"""
+Conversion of structures into the *Protein Blocks* structural alphabet.
+"""
+
+__name__ = "biotite.structure.alphabet"
+__author__ = "Patrick Kunzmann"
+__all__ = ["ProteinBlocksSequence", "to_protein_blocks"]
+
+import numpy as np
+from biotite.sequence.alphabet import LetterAlphabet
+from biotite.sequence.sequence import Sequence
+from biotite.structure.chains import get_chain_starts
+from biotite.structure.geometry import dihedral_backbone
+
+# PB reference angles, adapted from PBxplore
+PB_ANGLES = np.array(
+    [
+        [41.14,    75.53,   13.92,  -99.80,  131.88,  -96.27, 122.08,  -99.68],
+        [108.24,  -90.12,  119.54,  -92.21,  -18.06, -128.93, 147.04,  -99.90],
+        [-11.61, -105.66,   94.81, -106.09,  133.56, -106.93, 135.97, -100.63],
+        [141.98, -112.79,  132.20, -114.79,  140.11, -111.05, 139.54, -103.16],
+        [133.25, -112.37,  137.64, -108.13,  133.00,  -87.30, 120.54,   77.40],
+        [116.40, -105.53,  129.32,  -96.68,  140.72,  -74.19, -26.65,  -94.51],
+        [0.40,    -81.83,    4.91, -100.59,   85.50,  -71.65, 130.78,   84.98],
+        [119.14, -102.58,  130.83,  -67.91,  121.55,   76.25,  -2.95,  -90.88],
+        [130.68,  -56.92,  119.26,   77.85,   10.42,  -99.43, 141.40,  -98.01],
+        [114.32, -121.47,  118.14,   82.88, -150.05,  -83.81,  23.35,  -85.82],
+        [117.16,  -95.41,  140.40,  -59.35,  -29.23,  -72.39, -25.08,  -76.16],
+        [139.20,  -55.96,  -32.70,  -68.51,  -26.09,  -74.44, -22.60,  -71.74],
+        [-39.62,  -64.73,  -39.52,  -65.54,  -38.88,  -66.89, -37.76,  -70.19],
+        [-35.34,  -65.03,  -38.12,  -66.34,  -29.51,  -89.10,  -2.91,   77.90],
+        [-45.29,  -67.44,  -27.72,  -87.27,    5.13,   77.49,  30.71,  -93.23],
+        [-27.09,  -86.14,    0.30,   59.85,   21.51,  -96.30, 132.67,  -92.91],
+    ]
+)  # fmt: skip
+
+
+class ProteinBlocksSequence(Sequence):
+    """
+    Representation of a structure in the *Protein Blocks* structural alphabet.
+    :footcite:`Brevern2000`
+
+    Parameters
+    ----------
+    sequence : iterable object, optional
+        The *Protein Blocks* sequence.
+        This may either be a list or a string.
+        May take upper or lower case letters.
+        By default the sequence is empty.
+
+    See also
+    --------
+    to_protein_blocks : Create *Protein Blocks* sequences from a structure.
+
+    References
+    ----------
+
+    .. footbibliography::
+
+    """
+
+    alphabet = LetterAlphabet("abcdefghijklmnopz")
+    undefined_symbol = "z"
+
+    def __init__(self, sequence=""):
+        if isinstance(sequence, str):
+            sequence = sequence.lower()
+        else:
+            sequence = [symbol.upper() for symbol in sequence]
+        super().__init__(sequence)
+
+    def get_alphabet(self):
+        return ProteinBlocksSequence.alphabet
+
+    def remove_undefined(self):
+        """
+        Remove undefined symbols from the sequence.
+
+        Returns
+        -------
+        filtered_sequence : ProteinBlocksSequence
+            The sequence without undefined symbols.
+        """
+        undefined_code = ProteinBlocksSequence.alphabet.encode(
+            ProteinBlocksSequence.undefined_symbol
+        )
+        filtered_code = self.code[self.code != undefined_code]
+        filtered_sequence = ProteinBlocksSequence()
+        filtered_sequence.code = filtered_code
+        return filtered_sequence
+
+
+def to_protein_blocks(atoms):
+    """
+    Encode each chain in the given structure to the *Protein Blocks* structural
+    alphabet.
+    :footcite:`Brevern2000`
+
+    Parameters
+    ----------
+    atoms : AtomArray
+        The atom array to encode.
+        May contain multiple chains.
+
+    Returns
+    -------
+    sequences : list of Sequence, length=n
+        The encoded *Protein Blocks* sequence for each peptide chain in the structure.
+    chain_start_indices : ndarray, shape=(n,), dtype=int
+        The atom index where each chain starts.
+
+    References
+    ----------
+
+    .. footbibliography::
+
+    Examples
+    --------
+
+    >>> sequences, chain_starts = to_protein_blocks(atom_array)
+    >>> print(sequences[0])
+    zzmmmmmnopjmnopacdzz
+    """
+    sequences = []
+    chain_start_indices = get_chain_starts(atoms, add_exclusive_stop=True)
+    for i in range(len(chain_start_indices) - 1):
+        start = chain_start_indices[i]
+        stop = chain_start_indices[i + 1]
+        chain = atoms[start:stop]
+        sequences.append(_to_protein_blocks(chain))
+    return sequences, chain_start_indices[:-1]
+
+
+def _to_protein_blocks(chain):
+    undefined_code = ProteinBlocksSequence.alphabet.encode(
+        ProteinBlocksSequence.undefined_symbol
+    )
+
+    phi, psi, _ = dihedral_backbone(chain)
+
+    pb_angles = np.full((len(phi), 8), np.nan)
+    pb_angles[2:-2, 0] = psi[:-4]
+    pb_angles[2:-2, 1] = phi[1:-3]
+    pb_angles[2:-2, 2] = psi[1:-3]
+    pb_angles[2:-2, 3] = phi[2:-2]
+    pb_angles[2:-2, 4] = psi[2:-2]
+    pb_angles[2:-2, 5] = phi[3:-1]
+    pb_angles[2:-2, 6] = psi[3:-1]
+    pb_angles[2:-2, 7] = phi[4:]
+    pb_angles = np.rad2deg(pb_angles)
+
+    # Angle RMSD of all reference angles with all actual angles
+    rmsda = np.sum(
+        ((PB_ANGLES[:, np.newaxis] - pb_angles[np.newaxis, :] + 180) % 360 - 180) ** 2,
+        axis=-1,
+    )
+    # Where RMSDA is NaN, (missing atoms/residues or chain ends) set symbol to unknown
+    pb_seq_code = np.full(len(pb_angles), undefined_code, dtype=np.uint8)
+    pb_available_mask = ~np.isnan(rmsda).any(axis=0)
+    # Chose PB, where the RMSDA to the reference angle is lowest
+    # Due to the definition of Biotite symbol codes
+    # the index of the chosen PB is directly the symbol code
+    pb_seq_code[pb_available_mask] = np.argmin(rmsda[:, pb_available_mask], axis=0)
+    # Put the array of symbol codes into actual sequence objects
+    pb_sequence = ProteinBlocksSequence()
+    pb_sequence.code = pb_seq_code
+    return pb_sequence

biotite/structure/alphabet/unkerasify.py

@@ -0,0 +1,122 @@
+# 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.
+
+"""
+Parser for extracting weights from Keras files.
+
+Adapted from `moof2k/kerasify <https://github.com/moof2k/kerasify>`_.
+"""
+
+__name__ = "biotite.structure.alphabet"
+__author__ = "Martin Larralde"
+__all__ = ["load_kerasify"]
+
+import enum
+import functools
+import itertools
+import struct
+import numpy as np
+from biotite.structure.alphabet.layers import DenseLayer, Layer
+
+
+class LayerType(enum.IntEnum):
+    DENSE = 1
+    CONVOLUTION2D = 2
+    FLATTEN = 3
+    ELU = 4
+    ACTIVATION = 5
+    MAXPOOLING2D = 6
+    LSTM = 7
+    EMBEDDING = 8
+
+
+class ActivationType(enum.IntEnum):
+    LINEAR = 1
+    RELU = 2
+    SOFTPLUS = 3
+    SIGMOID = 4
+    TANH = 5
+    HARD_SIGMOID = 6
+
+
+class KerasifyParser:
+    """An incomplete parser for model files serialized with `kerasify`.
+
+    Notes
+    -----
+    Only dense layers are supported, since the ``foldseek`` VQ-VAE model
+    is only using 3 dense layers.
+    """
+
+    def __init__(self, file) -> None:
+        self.file = file
+        self.buffer = bytearray(1024)
+        (self.n_layers,) = self._get("I")
+
+    def read(self):
+        if self.n_layers == 0:
+            return None
+
+        self.n_layers -= 1
+        layer_type = LayerType(self._get("I")[0])
+        if layer_type == LayerType.DENSE:
+            (w0,) = self._get("I")
+            (w1,) = self._get("I")
+            (b0,) = self._get("I")
+            weights = (
+                np.frombuffer(self._read(f"={w0*w1}f"), dtype="f4")
+                .reshape(w0, w1)
+                .copy()
+            )
+            biases = np.frombuffer(self._read(f"={b0}f"), dtype="f4").copy()
+            activation = ActivationType(self._get("I")[0])
+            if activation not in (ActivationType.LINEAR, ActivationType.RELU):
+                raise NotImplementedError(
+                    f"Unsupported activation type: {activation!r}"
+                )
+            return DenseLayer(weights, biases, activation == ActivationType.RELU)
+        else:
+            raise NotImplementedError(f"Unsupported layer type: {layer_type!r}")
+
+    def __iter__(self):
+        return self
+
+    def __next__(self) -> Layer:
+        layer = self.read()
+        if layer is None:
+            raise StopIteration
+        return layer
+
+    def _read(self, format: str) -> memoryview:
+        n = struct.calcsize(format)
+        if len(self.buffer) < n:
+            self.buffer.extend(
+                itertools.islice(itertools.repeat(0), n - len(self.buffer))
+            )
+        v = memoryview(self.buffer)[:n]
+        self.file.readinto(v)  # type: ignore
+        return v
+
+    def _get(self, format: str):
+        v = self._read(format)
+        return struct.unpack(format, v)
+
+
+@functools.cache
+def load_kerasify(file_path):
+    """
+    Load the the model layers from a ``.kerasify`` file.
+
+    Parameters
+    ----------
+    file_path : str
+        The path to the ``.kerasify`` file.
+
+    Returns
+    -------
+    layers : tuple of Layer
+        The model layers.
+    """
+    with open(file_path, "rb") as file:
+        return tuple(KerasifyParser(file))

biotite/structure/atoms.py

@@ -13,6 +13,7 @@ __all__ = [
     "Atom",
     "AtomArray",
     "AtomArrayStack",
+    "concatenate",
     "array",
     "stack",
     "repeat",
@@ -22,6 +23,7 @@ __all__ = [
 
 import abc
 import numbers
+from collections.abc import Sequence
 import numpy as np
 from biotite.copyable import Copyable
 from biotite.structure.bonds import BondList
@@ -99,9 +101,24 @@ class _AtomArrayBase(Copyable, metaclass=abc.ABCMeta):
         See Also
         --------
         set_annotation
+
+        Notes
+        -----
+        If the annotation category already exists, a compatible dtype is chosen,
+        that is also able to represent the old values.
         """
         if category not in self._annot:
             self._annot[str(category)] = np.zeros(self._array_length, dtype=dtype)
+        elif np.can_cast(self._annot[str(category)].dtype, dtype):
+            self._annot[str(category)] = self._annot[str(category)].astype(dtype)
+        elif np.can_cast(dtype, self._annot[str(category)].dtype):
+            # The existing dtype is more general
+            pass
+        else:
+            raise ValueError(
+                f"Cannot cast '{str(category)}' "
+                f"with dtype '{self._annot[str(category)].dtype}' into '{dtype}'"
+            )
 
     def del_annotation(self, category):
         """
@@ -142,19 +159,28 @@ class _AtomArrayBase(Copyable, metaclass=abc.ABCMeta):
         ----------
         category : str
             The annotation category to be set.
-        array : ndarray or None
+        array : ndarray
             The new value of the annotation category. The size of the
             array must be the same as the array length.
+
+        Notes
+        -----
+        If the annotation category already exists, a compatible dtype is chosen,
+        that is able to represent the old and new array values.
         """
+        array = np.asarray(array)
         if len(array) != self._array_length:
             raise IndexError(
                 f"Expected array length {self._array_length}, " f"but got {len(array)}"
             )
         if category in self._annot:
-            # Keep the dtype if the annotation already exists
-            self._annot[category] = np.asarray(array, dtype=self._annot[category].dtype)
+            # If the annotation already exists, find the compatible dtype
+            self._annot[category] = array.astype(
+                dtype=np.promote_types(self._annot[category].dtype, array.dtype),
+                copy=False,
+            )
         else:
-            self._annot[category] = np.asarray(array)
+            self._annot[category] = array
 
     def get_annotation_categories(self):
         """
@@ -209,7 +235,7 @@ class _AtomArrayBase(Copyable, metaclass=abc.ABCMeta):
         else:
             raise TypeError(f"Index must be integer, not '{type(index).__name__}'")
 
-    def equal_annotations(self, item):
+    def equal_annotations(self, item, equal_nan=True):
         """
         Check, if this object shares equal annotation arrays with the
         given :class:`AtomArray` or :class:`AtomArrayStack`.
@@ -218,6 +244,8 @@ class _AtomArrayBase(Copyable, metaclass=abc.ABCMeta):
         ----------
         item : AtomArray or AtomArrayStack
             The object to compare the annotation arrays with.
+        equal_nan: bool
+            Whether to count `nan` values as equal. Default: True.
 
         Returns
         -------
@@ -229,7 +257,18 @@ class _AtomArrayBase(Copyable, metaclass=abc.ABCMeta):
         if not self.equal_annotation_categories(item):
             return False
         for name in self._annot:
-            if not np.array_equal(self._annot[name], item._annot[name]):
+            # ... allowing `nan` values causes type-casting, which is
+            #     only possible for floating-point arrays
+            allow_nan = (
+                equal_nan
+                if np.issubdtype(self._annot[name].dtype, np.floating)
+                else False
+            )
+            if not np.array_equal(
+                self._annot[name],
+                item._annot[name],
+                equal_nan=allow_nan,
+            ):
                 return False
         return True
 
@@ -383,42 +422,7 @@ class _AtomArrayBase(Copyable, metaclass=abc.ABCMeta):
         return self._array_length
 
     def __add__(self, array):
-        if not isinstance(self, type(array)):
-            raise TypeError("Can only concatenate two arrays or two stacks")
-        # Create either new array or stack, depending of the own type
-        if isinstance(self, AtomArray):
-            concat = AtomArray(length=self._array_length + array._array_length)
-        if isinstance(self, AtomArrayStack):
-            concat = AtomArrayStack(
-                self.stack_depth(), self._array_length + array._array_length
-            )
-
-        concat._coord = np.concatenate((self._coord, array.coord), axis=-2)
-
-        # Transfer only annotations,
-        # which are existent in both operands
-        arr_categories = list(array._annot.keys())
-        for category in self._annot.keys():
-            if category in arr_categories:
-                annot = self._annot[category]
-                arr_annot = array._annot[category]
-                concat._annot[category] = np.concatenate((annot, arr_annot))
-
-        # Concatenate bonds lists,
-        # if at least one of them contains bond information
-        if self._bonds is not None or array._bonds is not None:
-            bonds1 = self._bonds
-            bonds2 = array._bonds
-            if bonds1 is None:
-                bonds1 = BondList(self._array_length)
-            if bonds2 is None:
-                bonds2 = BondList(array._array_length)
-            concat._bonds = bonds1 + bonds2
-
-        # Copy box
-        if self._box is not None:
-            concat._box = np.copy(self._box)
-        return concat
+        return concatenate([self, array])
 
     def __copy_fill__(self, clone):
         super().__copy_fill__(clone)
@@ -582,6 +586,7 @@ class AtomArray(_AtomArrayBase):
     :class:`AtomArray` is done with the '+' operator.
     Only the annotation categories, which are existing in both arrays,
     are transferred to the new array.
+    For a list of :class:`AtomArray` objects, use :func:`concatenate()`.
 
     Optionally, an :class:`AtomArray` can store chemical bond
     information via a :class:`BondList` object.
@@ -854,7 +859,9 @@ class AtomArrayStack(_AtomArrayBase):
     :class:`AtomArray` instance.
 
     Concatenation of atoms for each array in the stack is done using the
-    '+' operator. For addition of atom arrays onto the stack use the
+    '+' operator.
+    For a list of :class:`AtomArray` objects, use :func:`concatenate()`.
+    For addition of atom arrays onto the stack use the
     :func:`stack()` method.
 
     The :attr:`box` attribute has the shape *m x 3 x 3*, as the cell
@@ -1268,6 +1275,112 @@ def stack(arrays):
     return array_stack
 
 
+def concatenate(atoms):
+    """
+    Concatenate multiple :class:`AtomArray` or :class:`AtomArrayStack` objects into
+    a single :class:`AtomArray` or :class:`AtomArrayStack`, respectively.
+
+    Parameters
+    ----------
+    atoms : iterable object of AtomArray or AtomArrayStack
+        The atoms to be concatenated.
+        :class:`AtomArray` cannot be mixed with :class:`AtomArrayStack`.
+
+    Returns
+    -------
+    concatenated_atoms : AtomArray or AtomArrayStack
+        The concatenated atoms, i.e. its ``array_length()`` is the sum of the
+        ``array_length()`` of the input ``atoms``.
+
+    Notes
+    -----
+    The following rules apply:
+
+    - Only the annotation categories that exist in all elements are transferred.
+    - The box of the first element that has a box is transferred, if any.
+    - The bonds of all elements are concatenated, if any element has associated bonds.
+      For elements without a :class:`BondList` an empty :class:`BondList` is assumed.
+
+    Examples
+    --------
+
+    >>> atoms1 = array([
+    ...     Atom([1,2,3], res_id=1, atom_name="N"),
+    ...     Atom([4,5,6], res_id=1, atom_name="CA"),
+    ...     Atom([7,8,9], res_id=1, atom_name="C")
+    ... ])
+    >>> atoms2 = array([
+    ...     Atom([1,2,3], res_id=2, atom_name="N"),
+    ...     Atom([4,5,6], res_id=2, atom_name="CA"),
+    ...     Atom([7,8,9], res_id=2, atom_name="C")
+    ... ])
+    >>> print(concatenate([atoms1, atoms2]))
+                1      N                1.000    2.000    3.000
+                1      CA               4.000    5.000    6.000
+                1      C                7.000    8.000    9.000
+                2      N                1.000    2.000    3.000
+                2      CA               4.000    5.000    6.000
+                2      C                7.000    8.000    9.000
+    """
+    # Ensure that the atoms can be iterated over multiple times
+    if not isinstance(atoms, Sequence):
+        atoms = list(atoms)
+
+    length = 0
+    depth = None
+    element_type = None
+    common_categories = set(atoms[0].get_annotation_categories())
+    box = None
+    has_bonds = False
+    for element in atoms:
+        if element_type is None:
+            element_type = type(element)
+        else:
+            if not isinstance(element, element_type):
+                raise TypeError(
+                    f"Cannot concatenate '{type(element).__name__}' "
+                    f"with '{element_type.__name__}'"
+                )
+        length += element.array_length()
+        if isinstance(element, AtomArrayStack):
+            if depth is None:
+                depth = element.stack_depth()
+            else:
+                if element.stack_depth() != depth:
+                    raise IndexError("The stack depths are not equal")
+        common_categories &= set(element.get_annotation_categories())
+        if element.box is not None and box is None:
+            box = element.box
+        if element.bonds is not None:
+            has_bonds = True
+
+    if element_type == AtomArray:
+        concat_atoms = AtomArray(length)
+    elif element_type == AtomArrayStack:
+        concat_atoms = AtomArrayStack(depth, length)
+    concat_atoms.coord = np.concatenate([element.coord for element in atoms], axis=-2)
+    for category in common_categories:
+        concat_atoms.set_annotation(
+            category,
+            np.concatenate(
+                [element.get_annotation(category) for element in atoms], axis=0
+            ),
+        )
+    concat_atoms.box = box
+    if has_bonds:
+        # Concatenate bonds of all elements
+        concat_atoms.bonds = BondList.concatenate(
+            [
+                element.bonds
+                if element.bonds is not None
+                else BondList(element.array_length())
+                for element in atoms
+            ]
+        )
+
+    return concat_atoms
+
+
 def repeat(atoms, coord):
     """
     Repeat atoms (:class:`AtomArray` or :class:`AtomArrayStack`)