biotite 1.6.0__cp314-cp314-win_amd64.whl

biotite/structure/atoms.py

@@ -0,0 +1,1596 @@
+# 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 module contains the main types of the ``structure`` subpackage:
+:class:`Atom`, :class:`AtomArray` and :class:`AtomArrayStack`.
+"""
+
+__name__ = "biotite.structure"
+__author__ = "Patrick Kunzmann"
+__all__ = [
+    "Atom",
+    "AtomArray",
+    "AtomArrayStack",
+    "concatenate",
+    "array",
+    "stack",
+    "repeat",
+    "from_template",
+    "coord",
+    "set_print_limits",
+]
+
+import abc
+import numbers
+import textwrap
+from collections.abc import Sequence
+import numpy as np
+from biotite.copyable import Copyable
+from biotite.structure.bonds import BondList
+
+
+class _AtomArrayBase(Copyable, metaclass=abc.ABCMeta):
+    """
+    Private base class for :class:`AtomArray` and
+    :class:`AtomArrayStack`.
+    It implements functionality for annotation arrays and also
+    rudimentarily for coordinates.
+
+    Parameters
+    ----------
+    length : int
+        The amount of atoms in the structure.
+    """
+
+    _max_models_printed = 10
+    _max_atoms_printed = 1000
+
+    def __init__(self, length):
+        """
+        Create the annotation arrays
+        """
+        self._annot = {}
+        self._array_length = length
+        self._coord = None
+        self._bonds = None
+        self._box = None
+        self.add_annotation("chain_id", dtype="U4")
+        self.add_annotation("res_id", dtype=int)
+        self.add_annotation("ins_code", dtype="U1")
+        self.add_annotation("res_name", dtype="U5")
+        self.add_annotation("hetero", dtype=bool)
+        self.add_annotation("atom_name", dtype="U6")
+        self.add_annotation("element", dtype="U2")
+
+    def array_length(self):
+        """
+        Get the length of the atom array.
+
+        This value is equivalent to the length of each annotation array.
+        For :class:`AtomArray` it is the same as ``len(array)``.
+
+        Returns
+        -------
+        length : int
+            Length of the array(s).
+        """
+        return self._array_length
+
+    @property
+    @abc.abstractmethod
+    def shape(self):
+        """
+        Tuple of array dimensions.
+
+        This property contains the current shape of the object.
+
+        Returns
+        -------
+        shape : tuple of int
+            Shape of the object.
+        """
+        return
+
+    def add_annotation(self, category, dtype):
+        """
+        Add an annotation category, if not already existing.
+
+        Initially the new annotation is filled with the *zero*
+        representation of the given type.
+
+        Parameters
+        ----------
+        category : str
+            The annotation category to be added.
+        dtype : type or str
+            A type instance or a valid *NumPy* *dtype* string.
+            Defines the type of the annotation.
+
+        See Also
+        --------
+        set_annotation : Assign directly a value to an 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):
+        """
+        Removes an annotation category.
+
+        Parameters
+        ----------
+        category : str
+            The annotation category to be removed.
+        """
+        if category in self._annot:
+            del self._annot[str(category)]
+
+    def get_annotation(self, category):
+        """
+        Return an annotation array.
+
+        Parameters
+        ----------
+        category : str
+            The annotation category to be returned.
+
+        Returns
+        -------
+        array : ndarray
+            The annotation array.
+        """
+        if category not in self._annot:
+            raise ValueError(f"Annotation category '{category}' is not existing")
+        return self._annot[category]
+
+    def set_annotation(self, category, array):
+        """
+        Set an annotation array. If the annotation category does not
+        exist yet, the category is created.
+
+        Parameters
+        ----------
+        category : str
+            The annotation category to be set.
+        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}, but got {len(array)}"
+            )
+        if category in self._annot:
+            # 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] = array
+
+    def get_annotation_categories(self):
+        """
+        Return a list containing all annotation array categories.
+
+        Returns
+        -------
+        categories : list
+            The list containing the names of each annotation array.
+        """
+        return list(self._annot.keys())
+
+    def _subarray(self, index):
+        # Index is one dimensional (boolean mask, index array)
+        new_coord = self._coord[..., index, :]
+        new_length = new_coord.shape[-2]
+        if isinstance(self, AtomArray):
+            # Initialize with length 0 to avoid unnecessary memory allocation
+            # for large arrays, as the individual annotation arrays ar
+            # overwritten later anyway
+            new_object = AtomArray(0)
+        elif isinstance(self, AtomArrayStack):
+            new_depth = new_coord.shape[-3]
+            new_object = AtomArrayStack(new_depth, 0)
+        new_object._coord = new_coord
+        if self._bonds is not None:
+            new_object._bonds = self._bonds[index]
+        if self._box is not None:
+            new_object._box = self._box
+        for annotation in self._annot:
+            new_object._annot[annotation] = self._annot[annotation].__getitem__(index)
+        # Update the array length, since has currently length '0'
+        new_object._array_length = new_length
+        return new_object
+
+    def _set_element(self, index, atom):
+        try:
+            if isinstance(index, (numbers.Integral, np.ndarray)):
+                for name in self._annot:
+                    self._annot[name][index] = atom._annot[name]
+                self._coord[..., index, :] = atom.coord
+            else:
+                raise TypeError(f"Index must be integer, not '{type(index).__name__}'")
+        except KeyError:
+            raise KeyError("The annotations of the 'Atom' are incompatible")
+
+    def _del_element(self, index):
+        if isinstance(index, numbers.Integral):
+            for name in self._annot:
+                self._annot[name] = np.delete(self._annot[name], index, axis=0)
+            self._coord = np.delete(self._coord, index, axis=-2)
+            self._array_length = self._coord.shape[-2]
+            if self._bonds is not None:
+                mask = np.ones(self._bonds.get_atom_count(), dtype=bool)
+                mask[index] = False
+                self._bonds = self._bonds[mask]
+        else:
+            raise TypeError(f"Index must be integer, not '{type(index).__name__}'")
+
+    def equal_annotations(self, item, equal_nan=True):
+        """
+        Check, if this object shares equal annotation arrays with the
+        given :class:`AtomArray` or :class:`AtomArrayStack`.
+
+        Parameters
+        ----------
+        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
+        -------
+        equality : bool
+            True, if the annotation arrays are equal.
+        """
+        if not isinstance(item, _AtomArrayBase):
+            return False
+        if not self.equal_annotation_categories(item):
+            return False
+        for name in self._annot:
+            # ... 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
+
+    def equal_annotation_categories(self, item):
+        """
+        Check, if this object shares equal annotation array categories
+        with the given :class:`AtomArray` or :class:`AtomArrayStack`.
+
+        Parameters
+        ----------
+        item : AtomArray or AtomArrayStack
+            The object to compare the annotation arrays with.
+
+        Returns
+        -------
+        equality : bool
+            True, if the annotation array names are equal.
+        """
+        return sorted(self._annot.keys()) == sorted(item._annot.keys())
+
+    def __getattr__(self, attr):
+        """
+        If the attribute is an annotation, the annotation is returned
+        from the dictionary.
+        Exposes coordinates.
+        """
+        if attr == "coord":
+            return self._coord
+        if attr == "bonds":
+            return self._bonds
+        if attr == "box":
+            return self._box
+        # Call method of 'object' superclass to avoid infinite recursive
+        # calls of '__getattr__()'
+        elif attr in super().__getattribute__("_annot"):
+            return self._annot[attr]
+        else:
+            raise AttributeError(
+                f"'{type(self).__name__}' object has no attribute '{attr}'"
+            )
+
+    def __setattr__(self, attr, value):
+        """
+        If the attribute is an annotation, the :attr:`value` is saved
+        to the annotation in the dictionary.
+        Exposes coordinates.
+        :attr:`value` must have same length as :func:`array_length()`.
+        """
+        if attr == "coord":
+            if not isinstance(value, np.ndarray):
+                raise TypeError("Value must be ndarray of floats")
+            if isinstance(self, AtomArray):
+                if value.ndim != 2:
+                    raise ValueError(
+                        "A 2-dimensional ndarray is expected for an AtomArray"
+                    )
+            elif isinstance(self, AtomArrayStack):
+                if value.ndim != 3:
+                    raise ValueError(
+                        "A 3-dimensional ndarray is expected for an AtomArrayStack"
+                    )
+            if value.shape[-2] != self._array_length:
+                raise ValueError(
+                    f"Expected array length {self._array_length}, but got {len(value)}"
+                )
+            if value.shape[-1] != 3:
+                raise TypeError("Expected 3 coordinates for each atom")
+            super().__setattr__("_coord", value.astype(np.float32, copy=False))
+
+        elif attr == "bonds":
+            if isinstance(value, BondList):
+                if value.get_atom_count() != self._array_length:
+                    raise ValueError(
+                        f"Array length is {self._array_length}, "
+                        f"but bond list has {value.get_atom_count()} atoms"
+                    )
+                super().__setattr__("_bonds", value)
+            elif value is None:
+                # Remove bond list
+                super().__setattr__("_bonds", None)
+            else:
+                raise TypeError("Value must be 'BondList'")
+
+        elif attr == "box":
+            if isinstance(value, np.ndarray):
+                if isinstance(self, AtomArray):
+                    if value.ndim != 2:
+                        raise ValueError(
+                            "A 2-dimensional ndarray is expected for an AtomArray"
+                        )
+                else:  # AtomArrayStack
+                    if value.ndim != 3:
+                        raise ValueError(
+                            "A 3-dimensional ndarray is expected for an AtomArrayStack"
+                        )
+                if value.shape[-2:] != (3, 3):
+                    raise TypeError("Box must be a 3x3 matrix (three vectors)")
+                box = value.astype(np.float32, copy=False)
+                super().__setattr__("_box", box)
+            elif value is None:
+                # Remove box
+                super().__setattr__("_box", None)
+            else:
+                raise TypeError("Box must be ndarray of floats or None")
+
+        elif attr == "_annot":
+            super().__setattr__(attr, value)
+        elif attr in self._annot:
+            self.set_annotation(attr, value)
+        else:
+            super().__setattr__(attr, value)
+
+    def __dir__(self):
+        attr = super().__dir__()
+        attr.append("coord")
+        attr.append("bonds")
+        attr.append("box")
+        for name in self._annot.keys():
+            attr.append(name)
+        return attr
+
+    def __eq__(self, item):
+        """
+        See Also
+        --------
+        equal_annotations
+        """
+        if not self.equal_annotations(item):
+            return False
+        if self._bonds != item._bonds:
+            return False
+        if self._box is None:
+            if item._box is not None:
+                return False
+        else:
+            if not np.array_equal(self._box, item._box):
+                return False
+        return np.array_equal(self._coord, item._coord)
+
+    def __len__(self):
+        """
+        The length of the annotation arrays.
+
+        Returns
+        -------
+        length : int
+            Length of the annotation arrays.
+        """
+        return self._array_length
+
+    def __add__(self, array):
+        return concatenate([self, array])
+
+    def __copy_fill__(self, clone):
+        super().__copy_fill__(clone)
+        self._copy_annotations(clone)
+        clone._coord = np.copy(self._coord)
+
+    def _copy_annotations(self, clone):
+        for name in self._annot:
+            clone._annot[name] = np.copy(self._annot[name])
+        if self._box is not None:
+            clone._box = np.copy(self._box)
+        if self._bonds is not None:
+            clone._bonds = self._bonds.copy()
+
+
+class Atom(Copyable):
+    """
+    A representation of a single atom.
+
+    The coordinates an annotations can be accessed directly.
+    A detailed description of each annotation category can be viewed
+    :doc:`here </apidoc/biotite.structure>`.
+
+    Parameters
+    ----------
+    coord : list or ndarray
+        The x, y and z coordinates.
+    **kwargs
+        Atom annotations as key value pair.
+
+    Attributes
+    ----------
+    {annot} : scalar
+        Annotations for this atom.
+    coord : ndarray, dtype=float
+        ndarray containing the x, y and z coordinate of the atom.
+    shape : tuple of int
+        Shape of the object.
+        In case of an :class:`Atom`, the tuple is empty.
+
+    Examples
+    --------
+
+    >>> atom = Atom([1,2,3], chain_id="A")
+    >>> atom.atom_name = "CA"
+    >>> print(atom.atom_name)
+    CA
+    >>> print(atom.coord)
+    [1. 2. 3.]
+    """
+
+    def __init__(self, coord, **kwargs):
+        self._annot = {}
+        self._annot["chain_id"] = ""
+        self._annot["res_id"] = 0
+        self._annot["ins_code"] = ""
+        self._annot["res_name"] = ""
+        self._annot["hetero"] = False
+        self._annot["atom_name"] = ""
+        self._annot["element"] = ""
+        if "kwargs" in kwargs:
+            # kwargs are given directly as dictionary
+            kwargs = kwargs["kwargs"]
+        for name, annotation in kwargs.items():
+            self._annot[name] = annotation
+        coord = np.array(coord, dtype=np.float32)
+        # Check if coord contains x,y and z coordinates
+        if coord.shape != (3,):
+            raise ValueError("Position must be ndarray with shape (3,)")
+        self.coord = coord
+
+    def __repr__(self):
+        """Represent Atom as a string for debugging."""
+        # print out key-value pairs and format strings in quotation marks
+        annot_parts = [
+            f'{key}="{value}"' if isinstance(value, str) else f"{key}={value}"
+            for key, value in self._annot.items()
+        ]
+
+        annot = ", ".join(annot_parts)
+        return f"Atom(np.{np.array_repr(self.coord)}, {annot})"
+
+    @property
+    def shape(self):
+        return ()
+
+    def __getattr__(self, attr):
+        if attr in super().__getattribute__("_annot"):
+            return self._annot[attr]
+        else:
+            raise AttributeError(
+                f"'{type(self).__name__}' object has no attribute '{attr}'"
+            )
+
+    def __setattr__(self, attr, value):
+        if attr == "_annot":
+            super().__setattr__(attr, value)
+        elif attr == "coord":
+            super().__setattr__(attr, value)
+        else:
+            self._annot[attr] = value
+
+    def __str__(self):
+        hetero = "HET" if self.hetero else ""
+        return (
+            f"{hetero:3} {self.chain_id:3} "
+            f"{self.res_id:5d}{self.ins_code:1} {self.res_name:3} "
+            f"{self.atom_name:6} {self.element:2}     "
+            f"{self.coord[0]:8.3f} "
+            f"{self.coord[1]:8.3f} "
+            f"{self.coord[2]:8.3f}"
+        )
+
+    def __eq__(self, item):
+        if not isinstance(item, Atom):
+            return False
+        if not np.array_equal(self.coord, item.coord):
+            return False
+        if self._annot.keys() != item._annot.keys():
+            return False
+        for name in self._annot:
+            if self._annot[name] != item._annot[name]:
+                return False
+        return True
+
+    def __ne__(self, item):
+        return not self == item
+
+    def __copy_create__(self):
+        return Atom(self.coord, **self._annot)
+
+
+class AtomArray(_AtomArrayBase):
+    """
+    An array representation of a model consisting of multiple atoms.
+
+    An :class:`AtomArray` can be seen as a list of :class:`Atom`
+    instances.
+    Instead of using directly a list, this class uses an *NumPy*
+    :class:`ndarray` for each annotation category and the coordinates.
+    These
+    coordinates can be accessed directly via the :attr:`coord`
+    attribute.
+    The annotations are accessed either via the category as attribute
+    name or the :func:`get_annotation()`, :func:`set_annotation()`
+    method.
+    Usage of custom annotations is achieved via :func:`add_annotation()`
+    or :func:`set_annotation()`.
+    A detailed description of each annotation category can be viewed
+    :doc:`here </apidoc/biotite.structure>`.
+
+    In order to get an an subarray of an :class:`AtomArray`,
+    *NumPy* style indexing is used.
+    This includes slices, boolean arrays, index arrays and even
+    *Ellipsis* notation.
+    Using a single integer as index returns a single :class:`Atom`
+    instance.
+
+    Inserting or appending an :class:`AtomArray` to another
+    :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.
+    It can be accessed using the :attr:`bonds` attribute.
+    If no bond information is available, :attr:`bonds` is ``None``.
+    Consequently the bond information can be removed from the
+    :class:`AtomArray`, by setting :attr:`bonds` to ``None``.
+    When indexing the :class:`AtomArray` the atom indices in the
+    associated :class:`BondList` are updated as well, hence the indices
+    in the :class:`BondList` will always point to the same atoms.
+    If two :class:`AtomArray` instances are concatenated, the resulting
+    :class:`AtomArray` will contain the merged :class:`BondList` if at
+    least one of the operands contains bond information.
+
+    The :attr:`box` attribute contains the box vectors of the unit cell
+    or the MD simulation box, respectively.
+    Hence, it is a *3 x 3* *ndarray* with the vectors in the last
+    dimension.
+    If no box is provided, the attribute is ``None``.
+    Setting the :attr:`box` attribute to ``None`` means removing the
+    box from the atom array.
+
+    Parameters
+    ----------
+    length : int
+        The fixed amount of atoms in the array.
+
+    Attributes
+    ----------
+    {annot} : ndarray
+        Multiple n-length annotation arrays.
+    coord : ndarray, dtype=float, shape=(n,3)
+        ndarray containing the x, y and z coordinate of the
+        atoms.
+    bonds : BondList or None
+        A :class:`BondList`, specifying the indices of atoms
+        that form a chemical bond.
+    box : ndarray, dtype=float, shape=(3,3) or None
+        The surrounding box. May represent a MD simulation box
+        or a crystallographic unit cell.
+    shape : tuple of int
+        Shape of the atom array.
+        The single value in the tuple is
+        the length of the atom array.
+
+    See Also
+    --------
+    AtomArrayStack : Representation of multiple structure models.
+
+    Examples
+    --------
+    Creating an atom array from atoms:
+
+    >>> atom1 = Atom([1,2,3], chain_id="A")
+    >>> atom2 = Atom([2,3,4], chain_id="A")
+    >>> atom3 = Atom([3,4,5], chain_id="B")
+    >>> atom_array = array([atom1, atom2, atom3])
+    >>> print(atom_array.array_length())
+    3
+
+    Accessing an annotation array:
+
+    >>> print(atom_array.chain_id)
+    ['A' 'A' 'B']
+
+    Accessing the coordinates:
+
+    >>> print(atom_array.coord)
+    [[1. 2. 3.]
+     [2. 3. 4.]
+     [3. 4. 5.]]
+
+    *NumPy* style filtering:
+
+    >>> atom_array = atom_array[atom_array.chain_id == "A"]
+    >>> print(atom_array.array_length())
+    2
+
+    Inserting an atom array:
+
+    >>> insert = array([Atom([7,8,9], chain_id="C")])
+    >>> atom_array = atom_array[0:1] + insert + atom_array[1:2]
+    >>> print(atom_array.chain_id)
+    ['A' 'C' 'A']
+    """
+
+    def __init__(self, length):
+        super().__init__(length)
+        if length is None:
+            self._coord = None
+        else:
+            self._coord = np.full((length, 3), np.nan, dtype=np.float32)
+
+    def __repr__(self):
+        """Represent AtomArray as a string for debugging."""
+        atoms = ""
+        for i in range(0, min(self.array_length(), _AtomArrayBase._max_atoms_printed)):
+            atoms = textwrap.indent(self.get_atom(i).__repr__(), "\t") + ",\n"
+        if self.array_length() > _AtomArrayBase._max_atoms_printed:
+            atoms = atoms + "\t...,\n"
+        return f"array([\n{atoms}])"
+
+    @property
+    def shape(self):
+        """
+        Tuple of array dimensions.
+
+        This property contains the current shape of the
+        :class:`AtomArray`.
+
+        Returns
+        -------
+        shape : tuple of int
+            Shape of the array.
+            The single value in the tuple is
+            the :func:`array_length()`.
+        """
+        return (self.array_length(),)
+
+    def get_atom(self, index):
+        """
+        Obtain the atom instance of the array at the specified index.
+
+        The same as ``array[index]``, if `index` is an integer.
+
+        Parameters
+        ----------
+        index : int
+            Index of the atom.
+
+        Returns
+        -------
+        atom : Atom
+            Atom at position `index`.
+        """
+        kwargs = {}
+        for name, annotation in self._annot.items():
+            kwargs[name] = annotation[index]
+        return Atom(coord=self._coord[index], kwargs=kwargs)
+
+    def __iter__(self):
+        """
+        Iterate through the array.
+
+        Yields
+        ------
+        atom : Atom
+        """
+        i = 0
+        while i < len(self):
+            yield self.get_atom(i)
+            i += 1
+
+    def __getitem__(self, index):
+        """
+        Obtain a subarray or the atom instance at the specified index.
+
+        Parameters
+        ----------
+        index : object
+            All index types *NumPy* accepts, are valid.
+
+        Returns
+        -------
+        sub_array : Atom or AtomArray
+            If `index` is an integer an :class:`Atom` instance is
+            returned.
+            Otherwise an :class:`AtomArray` with reduced length is
+            returned.
+        """
+        if isinstance(index, numbers.Integral):
+            return self.get_atom(index)
+        elif isinstance(index, tuple):
+            if len(index) == 2 and index[0] is Ellipsis:
+                # If first index is "...", just ignore the first index
+                return self.__getitem__(index[1])
+            else:
+                raise IndexError("'AtomArray' does not accept multidimensional indices")
+        else:
+            return self._subarray(index)
+
+    def __setitem__(self, index, atom):
+        """
+        Set the atom at the specified array position.
+
+        Parameters
+        ----------
+        index : int
+            The position, where the atom is set.
+        atom : Atom
+            The atom to be set.
+        """
+        self._set_element(index, atom)
+
+    def __delitem__(self, index):
+        """
+        Deletes the atom at the specified array position.
+
+        Parameters
+        ----------
+        index : int
+            The position where the atom should be deleted.
+        """
+        self._del_element(index)
+
+    def __len__(self):
+        """
+        The length of the array.
+
+        Returns
+        -------
+        length : int
+            Length of the array.
+        """
+        return self.array_length()
+
+    def __eq__(self, item):
+        """
+        Check if the array equals another :class:`AtomArray`.
+
+        Parameters
+        ----------
+        item : object
+            Object to campare the array with.
+
+        Returns
+        -------
+        equal : bool
+            True, if `item` is an :class:`AtomArray`
+            and all its attribute arrays equals the ones of this object.
+        """
+        if not super().__eq__(item):
+            return False
+        if not isinstance(item, AtomArray):
+            return False
+        return True
+
+    def __str__(self):
+        """
+        Get a string representation of the array.
+
+        Each line contains the attributes of one atom.
+        """
+        string = "\n".join(
+            [str(atom) for atom in self[: _AtomArrayBase._max_atoms_printed]]
+        )
+        if self.array_length() > _AtomArrayBase._max_atoms_printed:
+            string += "\n\t..."
+        return string
+
+    def __copy_create__(self):
+        return AtomArray(self.array_length())
+
+
+class AtomArrayStack(_AtomArrayBase):
+    """
+    A collection of multiple :class:`AtomArray` instances, where each
+    atom array has equal annotation arrays.
+
+    Effectively, this means that each atom is occurring in every array in
+    the stack at differing coordinates. This situation arises e.g. in
+    NMR-elucidated or simulated structures. Since the annotations are
+    equal for each array, the annotation arrays are 1-D, while the
+    coordinate array is 3-D (m x n x 3).
+    A detailed description of each annotation category can be viewed
+    :doc:`here </apidoc/biotite.structure>`.
+
+    Indexing works similar to :class:`AtomArray`, with the difference,
+    that two index dimensions are possible:
+    The first index dimension specifies the array(s), the second index
+    dimension specifies the atoms in each array (same as the index
+    in :class:`AtomArray`).
+    Using a single integer as first dimension index returns a single
+    :class:`AtomArray` instance.
+
+    Concatenation of atoms for each array in the stack is done using 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
+    might be different for each frame in the atom array stack.
+
+    Parameters
+    ----------
+    depth : int
+        The fixed amount of arrays in the stack. When indexing, this is
+        the length of the first dimension.
+
+    length : int
+        The fixed amount of atoms in each array in the stack. When
+        indexing, this is the length of the second dimension.
+
+    Attributes
+    ----------
+    {annot} : ndarray, shape=(n,)
+        Multiple n-length annotation arrays.
+    coord : ndarray, dtype=float, shape=(m,n,3)
+        ndarray containing the x, y and z coordinate of the
+        atoms.
+    bonds: BondList or None
+        A :class:`BondList`, specifying the indices of atoms
+        that form a chemical bond.
+    box: ndarray, dtype=float, shape=(m,3,3) or None
+        The surrounding box. May represent a MD simulation box
+        or a crystallographic unit cell.
+    shape : tuple of int
+        Shape of the stack.
+        The numbers correspond to the stack depth
+        and array length, respectively.
+
+    See Also
+    --------
+    AtomArray : Representation of a single structure model.
+
+    Examples
+    --------
+    Creating an atom array stack from two arrays:
+
+    >>> atom1 = Atom([1,2,3], chain_id="A")
+    >>> atom2 = Atom([2,3,4], chain_id="A")
+    >>> atom3 = Atom([3,4,5], chain_id="B")
+    >>> atom_array1 = array([atom1, atom2, atom3])
+    >>> print(atom_array1.coord)
+    [[1. 2. 3.]
+     [2. 3. 4.]
+     [3. 4. 5.]]
+    >>> atom_array2 = atom_array1.copy()
+    >>> atom_array2.coord += 3
+    >>> print(atom_array2.coord)
+    [[4. 5. 6.]
+     [5. 6. 7.]
+     [6. 7. 8.]]
+    >>> array_stack = stack([atom_array1, atom_array2])
+    >>> print(array_stack.coord)
+    [[[1. 2. 3.]
+      [2. 3. 4.]
+      [3. 4. 5.]]
+    <BLANKLINE>
+     [[4. 5. 6.]
+      [5. 6. 7.]
+      [6. 7. 8.]]]
+    """
+
+    def __init__(self, depth, length):
+        super().__init__(length)
+        if depth is None or length is None:
+            self._coord = None
+        else:
+            self._coord = np.full((depth, length, 3), np.nan, dtype=np.float32)
+
+    def __repr__(self):
+        """Represent AtomArrayStack as a string for debugging."""
+        arrays = ""
+        for i in range(0, min(self.stack_depth(), _AtomArrayBase._max_models_printed)):
+            arrays = textwrap.indent(self.get_array(i).__repr__(), "\t") + ",\n"
+        if self.stack_depth() > _AtomArrayBase._max_models_printed:
+            arrays = arrays + "\t...,\n"
+        return f"stack([\n{arrays}])"
+
+    def get_array(self, index):
+        """
+        Obtain the atom array instance of the stack at the specified
+        index.
+
+        The same as ``stack[index]``, if `index` is an integer.
+
+        Parameters
+        ----------
+        index : int
+            Index of the atom array.
+
+        Returns
+        -------
+        array : AtomArray
+            AtomArray at position `index`.
+        """
+        array = AtomArray(self.array_length())
+        for name in self._annot:
+            array._annot[name] = self._annot[name]
+        array._coord = self._coord[index]
+        if self._bonds is not None:
+            array._bonds = self._bonds.copy()
+        if self._box is not None:
+            array._box = self._box[index]
+
+        return array
+
+    def stack_depth(self):
+        """
+        Get the depth of the stack.
+
+        This value represents the amount of atom arrays in the stack.
+        It is the same as ``len(array)``.
+
+        Returns
+        -------
+        length : int
+            Length of the array(s).
+        """
+        return len(self)
+
+    @property
+    def shape(self):
+        """
+        Tuple of array dimensions.
+
+        This property contains the current shape of the
+        :class:`AtomArrayStack`.
+
+        Returns
+        -------
+        shape : tuple of int
+            Shape of the stack.
+            The numbers correspond to the :func:`stack_depth()`
+            and :func:`array_length()`, respectively.
+        """
+        return self.stack_depth(), self.array_length()
+
+    def __iter__(self):
+        """
+        Iterate through the array.
+
+        Yields
+        ------
+        array : AtomArray
+        """
+        i = 0
+        while i < len(self):
+            yield self.get_array(i)
+            i += 1
+
+    def __getitem__(self, index):
+        """
+        Obtain the atom array instance or an substack at the specified
+        index.
+
+        Parameters
+        ----------
+        index : object
+            All index types *NumPy* accepts are valid.
+
+        Returns
+        -------
+        sub_array : AtomArray or AtomArrayStack
+            If `index` is an integer an :class:`AtomArray` instance is
+            returned.
+            Otherwise an :class:`AtomArrayStack` with reduced depth and
+            length is returned.
+            In case the index is a tuple(int, int) an :class:`Atom`
+            instance is returned.
+        """
+        if isinstance(index, numbers.Integral):
+            return self.get_array(index)
+        elif isinstance(index, tuple):
+            if len(index) != 2:
+                raise IndexError(
+                    "'AtomArrayStack' does not accept an index "
+                    "with more than two dimensions"
+                )
+            if isinstance(index[0], numbers.Integral):
+                array = self.get_array(index[0])
+                return array.__getitem__(index[1])
+            else:
+                if isinstance(index[1], numbers.Integral):
+                    # Prevent reduction in dimensionality
+                    # in second dimension
+                    new_stack = self._subarray(slice(index[1], index[1] + 1))
+                else:
+                    new_stack = self._subarray(index[1])
+                if index[0] is not Ellipsis:
+                    new_stack._coord = new_stack._coord[index[0]]
+                    if new_stack._box is not None:
+                        new_stack._box = new_stack._box[index[0]]
+                return new_stack
+        else:
+            new_stack = AtomArrayStack(depth=0, length=self.array_length())
+            self._copy_annotations(new_stack)
+            new_stack._coord = self._coord[index]
+            if self._box is not None:
+                new_stack._box = self._box[index]
+            return new_stack
+
+    def __setitem__(self, index, array):
+        """
+        Set the atom array at the specified stack position.
+
+        The array and the stack must have equal annotation arrays.
+
+        Parameters
+        ----------
+        index : int
+            The position, where the array atom is set.
+        array : AtomArray
+            The atom array to be set.
+        """
+        if not self.equal_annotations(array):
+            raise ValueError("The stack and the array have unequal annotations")
+        if self.bonds != array.bonds:
+            raise ValueError("The stack and the array have unequal bonds")
+        if isinstance(index, numbers.Integral):
+            self.coord[index] = array.coord
+            if self.box is not None:
+                self.box[index] = array.box
+        else:
+            raise TypeError(f"Index must be integer, not '{type(index).__name__}'")
+
+    def __delitem__(self, index):
+        """
+        Deletes the atom array at the specified stack position.
+
+        Parameters
+        ----------
+        index : int
+            The position where the atom array should be deleted.
+        """
+        if isinstance(index, numbers.Integral):
+            self._coord = np.delete(self._coord, index, axis=0)
+        else:
+            raise TypeError(f"Index must be integer, not '{type(index).__name__}'")
+
+    def __len__(self):
+        """
+        The depth of the stack, i.e. the amount of models.
+
+        Returns
+        -------
+        depth : int
+            depth of the array.
+        """
+        # length is determined by length of coord attribute
+        return self._coord.shape[0]
+
+    def __eq__(self, item):
+        """
+        Check if the array equals another :class:`AtomArray`
+
+        Parameters
+        ----------
+        item : object
+            Object to campare the array with.
+
+        Returns
+        -------
+        equal : bool
+            True, if `item` is an :class:`AtomArray`
+            and all its attribute arrays equals the ones of this object.
+        """
+        if not super().__eq__(item):
+            return False
+        if not isinstance(item, AtomArrayStack):
+            return False
+        return True
+
+    def __str__(self):
+        """
+        Get a string representation of the stack.
+
+        :class:`AtomArray` strings eparated by blank lines
+        and a line indicating the index.
+        """
+        string = ""
+        for i, array in enumerate(self):
+            if i >= _AtomArrayBase._max_models_printed:
+                string += "..." + "\n" + "\n"
+                break
+            string += "Model " + str(i + 1) + "\n"
+            string += str(array) + "\n" + "\n"
+        return string
+
+    def __copy_create__(self):
+        return AtomArrayStack(self.stack_depth(), self.array_length())
+
+
+def array(atoms):
+    """
+    Create an :class:`AtomArray` from a list of :class:`Atom`.
+
+    Parameters
+    ----------
+    atoms : iterable object of Atom
+        The atoms to be combined in an array.
+        All atoms must share the same annotation categories.
+
+    Returns
+    -------
+    array : AtomArray
+        The listed atoms as array.
+
+    Examples
+    --------
+
+    Creating an atom array from atoms:
+
+    >>> atom1 = Atom([1,2,3], chain_id="A")
+    >>> atom2 = Atom([2,3,4], chain_id="A")
+    >>> atom3 = Atom([3,4,5], chain_id="B")
+    >>> atom_array = array([atom1, atom2, atom3])
+    >>> print(atom_array)
+        A       0                       1.000    2.000    3.000
+        A       0                       2.000    3.000    4.000
+        B       0                       3.000    4.000    5.000
+    """
+    # Check if all atoms have the same annotation names
+    # Equality check requires sorting
+    names = sorted(atoms[0]._annot.keys())
+    for i, atom in enumerate(atoms):
+        if sorted(atom._annot.keys()) != names:
+            raise ValueError(
+                f"The atom at index {i} does not share the same "
+                f"annotation categories as the atom at index 0"
+            )
+    array = AtomArray(len(atoms))
+
+    # Add all (also optional) annotation categories
+    for name in names:
+        value = atoms[0]._annot[name]
+        if isinstance(value, str):
+            # Find maximum string length across all atoms for this annotation
+            max_len = max(len(str(atom._annot[name])) for atom in atoms)
+            dtype = f"<U{max_len}"
+        else:
+            dtype = type(value)
+        array.add_annotation(name, dtype=dtype)
+
+    # Add all atoms to AtomArray
+    for i in range(len(atoms)):
+        for name in names:
+            array._annot[name][i] = atoms[i]._annot[name]
+        array._coord[i] = atoms[i].coord
+    return array
+
+
+def stack(arrays):
+    """
+    Create an :class:`AtomArrayStack` from a list of :class:`AtomArray`.
+
+    Parameters
+    ----------
+    arrays : iterable object of AtomArray
+        The atom arrays to be combined in a stack.
+        All atom arrays must have an equal number of atoms and equal
+        annotation arrays.
+
+    Returns
+    -------
+    stack : AtomArrayStack
+        The stacked atom arrays.
+
+    Examples
+    --------
+    Creating an atom array stack from two arrays:
+
+    >>> atom1 = Atom([1,2,3], chain_id="A")
+    >>> atom2 = Atom([2,3,4], chain_id="A")
+    >>> atom3 = Atom([3,4,5], chain_id="B")
+    >>> atom_array1 = array([atom1, atom2, atom3])
+    >>> print(atom_array1.coord)
+    [[1. 2. 3.]
+     [2. 3. 4.]
+     [3. 4. 5.]]
+    >>> atom_array2 = atom_array1.copy()
+    >>> atom_array2.coord += 3
+    >>> print(atom_array2.coord)
+    [[4. 5. 6.]
+     [5. 6. 7.]
+     [6. 7. 8.]]
+    >>> array_stack = stack([atom_array1, atom_array2])
+    >>> print(array_stack.coord)
+    [[[1. 2. 3.]
+      [2. 3. 4.]
+      [3. 4. 5.]]
+    <BLANKLINE>
+     [[4. 5. 6.]
+      [5. 6. 7.]
+      [6. 7. 8.]]]
+    """
+    array_count = 0
+    ref_array = None
+    for i, array in enumerate(arrays):
+        if ref_array is None:
+            ref_array = array
+        array_count += 1
+        # Check if all arrays share equal annotations
+        if not array.equal_annotations(ref_array):
+            raise ValueError(
+                f"The annotations of the atom array at index {i} are not "
+                f"equal to the annotations of the atom array at index 0"
+            )
+    array_stack = AtomArrayStack(array_count, ref_array.array_length())
+    for name, annotation in ref_array._annot.items():
+        array_stack._annot[name] = annotation
+    coord_list = [array._coord for array in arrays]
+    array_stack._coord = np.stack(coord_list, axis=0)
+    # Take bond list from first array
+    array_stack._bonds = ref_array._bonds
+    # When all atom arrays provide a box, copy the boxes
+    if all([array.box is not None for array in arrays]):
+        array_stack.box = np.array([array.box for array in 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`)
+    multiple times in the same model with different coordinates.
+
+    Parameters
+    ----------
+    atoms : AtomArray, shape=(n,) or AtomArrayStack, shape=(m,n)
+        The atoms to be repeated.
+    coord : ndarray, dtype=float, shape=(k,n,3) or shape=(k,m,n,3)
+        The coordinates to be used for the repeated atoms.
+        The length of first dimension determines the number of repeats.
+        If `atoms` is an :class:`AtomArray` 3 dimensions, otherwise
+        4 dimensions are required.
+
+    Returns
+    -------
+    repeated: AtomArray, shape=(n*k,) or AtomArrayStack, shape=(m,n*k)
+        The repeated atoms.
+        Whether an :class:`AtomArray` or an :class:`AtomArrayStack` is
+        returned depends on the input `atoms`.
+
+    Examples
+    --------
+
+    >>> atoms = 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")
+    ... ])
+    >>> print(atoms)
+                1      N                1.000    2.000    3.000
+                1      CA               4.000    5.000    6.000
+                1      C                7.000    8.000    9.000
+    >>> repeat_coord = np.array([
+    ...     [[0,0,0], [1,1,1], [2,2,2]],
+    ...     [[3,3,3], [4,4,4], [5,5,5]]
+    ... ])
+    >>> print(repeat(atoms, repeat_coord))
+                1      N                0.000    0.000    0.000
+                1      CA               1.000    1.000    1.000
+                1      C                2.000    2.000    2.000
+                1      N                3.000    3.000    3.000
+                1      CA               4.000    4.000    4.000
+                1      C                5.000    5.000    5.000
+    """
+    if isinstance(atoms, AtomArray) and coord.ndim != 3:
+        raise ValueError(
+            f"Expected 3 dimensions for the coordinate array, got {coord.ndim}"
+        )
+    elif isinstance(atoms, AtomArrayStack) and coord.ndim != 4:
+        raise ValueError(
+            f"Expected 4 dimensions for the coordinate array, got {coord.ndim}"
+        )
+
+    repetitions = len(coord)
+    orig_length = atoms.array_length()
+    new_length = orig_length * repetitions
+
+    if isinstance(atoms, AtomArray):
+        if coord.ndim != 3:
+            raise ValueError(
+                f"Expected 3 dimensions for the coordinate array, but got {coord.ndim}"
+            )
+        repeated = AtomArray(new_length)
+        repeated.coord = coord.reshape((new_length, 3))
+
+    elif isinstance(atoms, AtomArrayStack):
+        if coord.ndim != 4:
+            raise ValueError(
+                f"Expected 4 dimensions for the coordinate array, but got {coord.ndim}"
+            )
+        repeated = AtomArrayStack(atoms.stack_depth(), new_length)
+        repeated.coord = coord.reshape((atoms.stack_depth(), new_length, 3))
+
+    else:
+        raise TypeError(
+            f"Expected 'AtomArray' or 'AtomArrayStack', but got {type(atoms).__name__}"
+        )
+
+    for category in atoms.get_annotation_categories():
+        annot = np.tile(atoms.get_annotation(category), repetitions)
+        repeated.set_annotation(category, annot)
+    if atoms.bonds is not None:
+        repeated_bonds = atoms.bonds.copy()
+        for _ in range(repetitions - 1):
+            repeated_bonds += atoms.bonds
+        repeated.bonds = repeated_bonds
+    if atoms.box is not None:
+        repeated.box = atoms.box.copy()
+
+    return repeated
+
+
+def from_template(template, coord, box=None):
+    """
+    Create an :class:`AtomArrayStack` using template atoms and given
+    coordinates.
+
+    Parameters
+    ----------
+    template : AtomArray, shape=(n,) or AtomArrayStack, shape=(m,n)
+        The annotation arrays and bonds of the returned stack are taken
+        from this template.
+    coord : ndarray, dtype=float, shape=(l,n,3)
+        The coordinates for each model of the returned stack.
+    box : ndarray, optional, dtype=float, shape=(l,3,3)
+        The box for each model of the returned stack.
+
+    Returns
+    -------
+    array_stack : AtomArrayStack
+        A stack containing the annotation arrays and bonds from
+        `template` but the coordinates from `coord` and the boxes from
+        `boxes`.
+    """
+    if template.array_length() != coord.shape[-2]:
+        raise ValueError(
+            f"Template has {template.array_length()} atoms, but "
+            f"{coord.shape[-2]} coordinates are given"
+        )
+
+    # Create empty stack with no models
+    new_stack = AtomArrayStack(0, template.array_length())
+
+    for category in template.get_annotation_categories():
+        annot = template.get_annotation(category)
+        new_stack.set_annotation(category, annot)
+    if template.bonds is not None:
+        new_stack.bonds = template.bonds.copy()
+    if box is not None:
+        new_stack.box = box.copy()
+
+    # After setting the coordinates the number of models is the number
+    # of models in the new coordinates
+    new_stack.coord = coord
+
+    return new_stack
+
+
+def coord(item):
+    """
+    Get the atom coordinates of the given array.
+
+    This may be directly and :class:`Atom`, :class:`AtomArray` or
+    :class:`AtomArrayStack` or
+    alternatively an (n x 3) or (m x n x 3)  :class:`ndarray`
+    containing the coordinates.
+
+    Parameters
+    ----------
+    item : Atom or AtomArray or AtomArrayStack or ndarray
+        Returns the :attr:`coord` attribute, if `item` is an
+        :class:`Atom`, :class:`AtomArray` or :class:`AtomArrayStack`.
+        Directly returns the input, if `item` is a :class:`ndarray`.
+
+    Returns
+    -------
+    coord : ndarray
+        Atom coordinates.
+    """
+
+    if isinstance(item, (Atom, _AtomArrayBase)):
+        return item.coord
+    elif isinstance(item, np.ndarray):
+        return item.astype(np.float32, copy=False)
+    else:
+        return np.array(item, dtype=np.float32)
+
+
+def set_print_limits(max_models=10, max_atoms=1000):
+    """
+    Set the maximum number of models and atoms to print in the ``str()`` and ``repr()``
+    representations.
+
+    The remaining models/atoms are abbreviated by ellipses.
+
+    Parameters
+    ----------
+    max_models : int
+        The maximum number of models to print.
+    max_atoms : int
+        The maximum number of atoms to print.
+    """
+    _AtomArrayBase._max_models_printed = max_models
+    _AtomArrayBase._max_atoms_printed = max_atoms