biotite 0.41.2__cp310-cp310-win_amd64.whl → 1.0.1__cp310-cp310-win_amd64.whl

biotite/structure/atoms.py

@@ -4,19 +4,27 @@
 
 """
 This module contains the main types of the ``structure`` subpackage:
-:class:`Atom`, :class:`AtomArray` and :class:`AtomArrayStack`. 
+:class:`Atom`, :class:`AtomArray` and :class:`AtomArrayStack`.
 """
 
 __name__ = "biotite.structure"
 __author__ = "Patrick Kunzmann"
-__all__ = ["Atom", "AtomArray", "AtomArrayStack",
-           "array", "stack", "repeat", "from_template", "coord"]
+__all__ = [
+    "Atom",
+    "AtomArray",
+    "AtomArrayStack",
+    "array",
+    "stack",
+    "repeat",
+    "from_template",
+    "coord",
+]
 
-import numbers
 import abc
+import numbers
 import numpy as np
-from .bonds import BondList
-from ..copyable import Copyable
+from biotite.copyable import Copyable
+from biotite.structure.bonds import BondList
 
 
 class _AtomArrayBase(Copyable, metaclass=abc.ABCMeta):
@@ -26,7 +34,7 @@ class _AtomArrayBase(Copyable, metaclass=abc.ABCMeta):
     It implements functionality for annotation arrays and also
     rudimentarily for coordinates.
     """
-    
+
     def __init__(self, length):
         """
         Create the annotation arrays
@@ -43,14 +51,14 @@ class _AtomArrayBase(Copyable, metaclass=abc.ABCMeta):
         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
@@ -71,15 +79,15 @@ class _AtomArrayBase(Copyable, metaclass=abc.ABCMeta):
         shape : tuple of int
             Shape of the object.
         """
-        return 
-        
+        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
@@ -87,19 +95,33 @@ class _AtomArrayBase(Copyable, metaclass=abc.ABCMeta):
         dtype : type or str
             A type instance or a valid *NumPy* *dtype* string.
             Defines the type of the annotation
-        
+
         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)
-            
+            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
@@ -107,32 +129,30 @@ class _AtomArrayBase(Copyable, metaclass=abc.ABCMeta):
         """
         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"
-            )
+            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
@@ -140,31 +160,37 @@ class _AtomArrayBase(Copyable, metaclass=abc.ABCMeta):
         array : ndarray or None
             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)}"
+                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):
         """
         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, :]
@@ -180,10 +206,9 @@ class _AtomArrayBase(Copyable, metaclass=abc.ABCMeta):
         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))
+            new_object._annot[annotation] = self._annot[annotation].__getitem__(index)
         return new_object
-        
+
     def _set_element(self, index, atom):
         try:
             if isinstance(index, (numbers.Integral, np.ndarray)):
@@ -191,12 +216,10 @@ class _AtomArrayBase(Copyable, metaclass=abc.ABCMeta):
                     self._annot[name][index] = atom._annot[name]
                 self._coord[..., index, :] = atom.coord
             else:
-                raise TypeError(
-                    f"Index must be integer, not '{type(index).__name__}'"
-                )
+                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:
@@ -208,20 +231,18 @@ class _AtomArrayBase(Copyable, metaclass=abc.ABCMeta):
                 mask[index] = False
                 self._bonds = self._bonds[mask]
         else:
-            raise TypeError(
-                    f"Index must be integer, not '{type(index).__name__}'"
-                )
-    
+            raise TypeError(f"Index must be integer, not '{type(index).__name__}'")
+
     def equal_annotations(self, item):
         """
         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.
-        
+
         Returns
         -------
         equality : bool
@@ -235,24 +256,24 @@ class _AtomArrayBase(Copyable, metaclass=abc.ABCMeta):
             if not np.array_equal(self._annot[name], item._annot[name]):
                 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
@@ -273,7 +294,7 @@ class _AtomArrayBase(Copyable, metaclass=abc.ABCMeta):
             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
@@ -287,15 +308,13 @@ class _AtomArrayBase(Copyable, metaclass=abc.ABCMeta):
             if isinstance(self, AtomArray):
                 if value.ndim != 2:
                     raise ValueError(
-                        "A 2-dimensional ndarray is expected "
-                        "for an AtomArray"
-                )
+                        "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"
-                )
+                        "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}, "
@@ -304,7 +323,7 @@ class _AtomArrayBase(Copyable, metaclass=abc.ABCMeta):
             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:
@@ -318,22 +337,21 @@ class _AtomArrayBase(Copyable, metaclass=abc.ABCMeta):
                 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
+                            "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):
+                        )
+                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)
@@ -342,14 +360,14 @@ class _AtomArrayBase(Copyable, metaclass=abc.ABCMeta):
                 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")
@@ -358,7 +376,7 @@ class _AtomArrayBase(Copyable, metaclass=abc.ABCMeta):
         for name in self._annot.keys():
             attr.append(name)
         return attr
-    
+
     def __eq__(self, item):
         """
         See Also
@@ -376,30 +394,31 @@ class _AtomArrayBase(Copyable, metaclass=abc.ABCMeta):
             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):
-        if type(self) != type(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)
+            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 = 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())
@@ -407,29 +426,29 @@ class _AtomArrayBase(Copyable, metaclass=abc.ABCMeta):
             if category in arr_categories:
                 annot = self._annot[category]
                 arr_annot = array._annot[category]
-                concat._annot[category] = np.concatenate((annot,arr_annot))
-        
+                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)
+                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
-    
+
     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])
@@ -437,23 +456,23 @@ class _AtomArrayBase(Copyable, metaclass=abc.ABCMeta):
             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
@@ -463,19 +482,19 @@ class Atom(Copyable):
     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"] = ""
@@ -500,17 +519,17 @@ class Atom(Copyable):
         """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}'
+            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})'    
+        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]
@@ -518,7 +537,7 @@ class Atom(Copyable):
             raise AttributeError(
                 f"'{type(self).__name__}' object has no attribute '{attr}'"
             )
-        
+
     def __setattr__(self, attr, value):
         if attr == "_annot":
             super().__setattr__(attr, value)
@@ -526,16 +545,18 @@ class Atom(Copyable):
             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}"
-    
+        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
@@ -547,18 +568,18 @@ class Atom(Copyable):
             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*
@@ -573,14 +594,14 @@ class AtomArray(_AtomArrayBase):
     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,
@@ -611,7 +632,7 @@ class AtomArray(_AtomArrayBase):
     ----------
     length : int
         The fixed amount of atoms in the array.
-    
+
     Attributes
     ----------
     {annot} : ndarray
@@ -629,44 +650,44 @@ class AtomArray(_AtomArrayBase):
         Shape of the atom array.
         The single value in the tuple is
         the length of the atom 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.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:
@@ -676,13 +697,13 @@ class AtomArray(_AtomArrayBase):
 
     def __repr__(self):
         """Represent AtomArray as a string for debugging."""
-        atoms = ''
+        atoms = ""
         for i in range(0, self.array_length()):
             if len(atoms) == 0:
-                atoms = '\n\t' + self.get_atom(i).__repr__()
+                atoms = "\n\t" + self.get_atom(i).__repr__()
             else:
-                atoms = atoms + ',\n\t' + self.get_atom(i).__repr__()
-        return f'array([{atoms}\n])'
+                atoms = atoms + ",\n\t" + self.get_atom(i).__repr__()
+        return f"array([{atoms}\n])"
 
     @property
     def shape(self):
@@ -703,33 +724,33 @@ class AtomArray(_AtomArrayBase):
         --------
         array_length
         """
-        return self.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`. 
+            Atom at position `index`.
         """
         kwargs = {}
         for name, annotation in self._annot.items():
             kwargs[name] = annotation[index]
-        return Atom(coord = self._coord[index], kwargs=kwargs)
-    
+        return Atom(coord=self._coord[index], kwargs=kwargs)
+
     def __iter__(self):
         """
         Iterate through the array.
-        
+
         Yields
         ------
         atom : Atom
@@ -738,16 +759,16 @@ class AtomArray(_AtomArrayBase):
         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
@@ -763,16 +784,14 @@ class AtomArray(_AtomArrayBase):
                 # If first index is "...", just ignore the first index
                 return self.__getitem__(index[1])
             else:
-                raise IndexError(
-                    "'AtomArray' does not accept multidimensional indices"
-                )
+                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
@@ -781,38 +800,38 @@ class AtomArray(_AtomArrayBase):
             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
@@ -824,15 +843,15 @@ class AtomArray(_AtomArrayBase):
         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.
         """
         return "\n".join([str(atom) for atom in self])
-    
+
     def __copy_create__(self):
         return AtomArray(self.array_length())
 
@@ -841,7 +860,7 @@ 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 occuring in every array in
     the stack at differing coordinates. This situation arises e.g. in
     NMR-elucidated or simulated structures. Since the annotations are
@@ -849,7 +868,7 @@ class AtomArrayStack(_AtomArrayBase):
     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
@@ -857,24 +876,24 @@ class AtomArrayStack(_AtomArrayBase):
     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 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,)
@@ -892,15 +911,15 @@ class AtomArrayStack(_AtomArrayBase):
         Shape of the stack.
         The numbers correspond to the stack depth
         and array length, respectively.
-    
+
     See also
     --------
     AtomArray
-    
+
     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")
@@ -925,40 +944,40 @@ class AtomArrayStack(_AtomArrayBase):
       [5. 6. 7.]
       [6. 7. 8.]]]
     """
-    
+
     def __init__(self, depth, length):
         super().__init__(length)
-        if depth == None or length == None:
+        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 = ''
+        arrays = ""
         for i in range(0, self.stack_depth()):
             if len(arrays) == 0:
-                arrays = '\n\t' + self.get_array(i).__repr__()
+                arrays = "\n\t" + self.get_array(i).__repr__()
             else:
-                arrays = arrays + ',\n\t' + self.get_array(i).__repr__()
-        return f'stack([{arrays}\n])'
+                arrays = arrays + ",\n\t" + self.get_array(i).__repr__()
+        return f"stack([{arrays}\n])"
 
     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`. 
+            AtomArray at position `index`.
         """
         array = AtomArray(self.array_length())
         for name in self._annot:
@@ -970,14 +989,14 @@ class AtomArrayStack(_AtomArrayBase):
             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
@@ -1005,7 +1024,7 @@ class AtomArrayStack(_AtomArrayBase):
     def __iter__(self):
         """
         Iterate through the array.
-        
+
         Yields
         ------
         array : AtomArray
@@ -1014,17 +1033,17 @@ class AtomArrayStack(_AtomArrayBase):
         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
@@ -1033,7 +1052,7 @@ class AtomArrayStack(_AtomArrayBase):
             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.  
+            instance is returned.
         """
         if isinstance(index, numbers.Integral):
             return self.get_array(index)
@@ -1050,7 +1069,7 @@ class AtomArrayStack(_AtomArrayBase):
                 if isinstance(index[1], numbers.Integral):
                     # Prevent reduction in dimensionality
                     # in second dimension
-                    new_stack = self._subarray(slice(index[1], index[1]+1))
+                    new_stack = self._subarray(slice(index[1], index[1] + 1))
                 else:
                     new_stack = self._subarray(index[1])
                 if index[0] is not Ellipsis:
@@ -1065,14 +1084,13 @@ class AtomArrayStack(_AtomArrayBase):
             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
@@ -1081,26 +1099,20 @@ class AtomArrayStack(_AtomArrayBase):
             The atom array to be set.
         """
         if not self.equal_annotations(array):
-            raise ValueError(
-                "The stack and the array have unequal annotations"
-            )
+            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"
-            )
+            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__}'"
-            )
-        
+            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
@@ -1109,14 +1121,12 @@ class AtomArrayStack(_AtomArrayBase):
         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__}'"
-            )
-    
+            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
@@ -1124,16 +1134,16 @@ class AtomArrayStack(_AtomArrayBase):
         """
         # 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
@@ -1145,20 +1155,20 @@ class AtomArrayStack(_AtomArrayBase):
         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):
-            string += "Model " + str(i+1) + "\n"
+            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())
 
@@ -1166,23 +1176,23 @@ class AtomArrayStack(_AtomArrayBase):
 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")
@@ -1204,7 +1214,7 @@ def array(atoms):
     array = AtomArray(len(atoms))
     # Add all (also optional) annotation categories
     for name in names:
-        array.add_annotation(name, dtype=type(atoms[0]._annot[name])) 
+        array.add_annotation(name, dtype=type(atoms[0]._annot[name]))
     # Add all atoms to AtomArray
     for i in range(len(atoms)):
         for name in names:
@@ -1216,23 +1226,23 @@ def array(atoms):
 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")
@@ -1272,7 +1282,7 @@ def stack(arrays):
     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] 
+    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
@@ -1296,14 +1306,14 @@ def repeat(atoms, coord):
         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
     --------
 
@@ -1336,7 +1346,7 @@ def repeat(atoms, coord):
         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
@@ -1358,24 +1368,24 @@ def repeat(atoms, coord):
             )
         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', "
             f"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):
+        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
 
 
@@ -1383,7 +1393,7 @@ 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)
@@ -1393,7 +1403,7 @@ def from_template(template, coord, box=None):
         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
@@ -1409,7 +1419,7 @@ def from_template(template, coord, box=None):
 
     # 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)
@@ -1417,30 +1427,30 @@ def from_template(template, coord, box=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