biotite 0.41.1__cp312-cp312-macosx_10_16_arm64.whl

biotite/structure/celllist.pyx

@@ -0,0 +1,849 @@
+# 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 allows efficient search of atoms in a defined radius around
+a location.
+"""
+
+__name__ = "biotite.structure"
+__author__ = "Patrick Kunzmann"
+__all__ = ["CellList"]
+
+cimport cython
+cimport numpy as np
+from libc.stdlib cimport realloc, malloc, free
+
+import numpy as np
+from .atoms import coord as to_coord
+from .atoms import AtomArrayStack
+from .box import repeat_box_coord, move_inside_box
+
+ctypedef np.uint64_t ptr
+ctypedef np.float32_t float32
+ctypedef np.uint8_t uint8
+
+
+cdef class CellList:
+    """
+    __init__(atom_array, cell_size, periodic=False, box=None, selection=None)
+    
+    This class enables the efficient search of atoms in vicinity of a
+    defined location.
+    
+    This class stores the indices of an atom array in virtual "cells",
+    each corresponding to a specific coordinate interval.
+    If the atoms in vicinity of a specific location are searched, only
+    the atoms in the relevant cells are checked.
+    Effectively this decreases the operation time for finding atoms
+    with a maximum distance to given coordinates from *O(n)* to *O(1)*,
+    after the :class:`CellList` has been created.
+    Therefore a :class:`CellList` saves calculation time in those
+    cases, where vicinity is checked for multiple locations.
+    
+    Parameters
+    ----------
+    atom_array : AtomArray or ndarray, dtype=float, shape=(n,3)
+        The :class:`AtomArray` to create the :class:`CellList` for.
+        Alternatively the atom coordinates are accepted directly.
+        In this case `box` must be set, if `periodic` is true.
+    cell_size : float
+        The coordinate interval each cell has for x, y and z axis.
+        The amount of cells depends on the range of coordinates in the
+        `atom_array` and the `cell_size`.
+    periodic : bool, optional
+        If true, the cell list considers periodic copies of atoms.
+        The periodicity is based on the `box` attribute of `atom_array`.
+        (Default: False)
+    box : ndarray, dtype=float, shape=(3,3), optional
+        If provided, the periodicity is based on this parameter instead
+        of the :attr:`box` attribute of `atom_array`.
+        Only has an effect, if `periodic` is ``True``.
+    selection : ndarray, dtype=bool, shape=(n,), optional
+        If provided, only the atoms masked by this array are stored in
+        the cell list. However, the indices stored in the cell list
+        will still refer to the original unfiltered `atom_array`.
+            
+    Examples
+    --------
+    
+    >>> cell_list = CellList(atom_array, cell_size=5)
+    >>> near_atoms = atom_array[cell_list.get_atoms(np.array([1,2,3]), radius=7.0)]
+    """
+    
+    # The atom coordinates
+    cdef float32[:,:] _coord
+    # A boolean mask that covers the selected atoms
+    cdef uint8[:] _selection
+    cdef bint _has_selection
+    # The cells to store the coordinates in; an ndarray of pointers
+    cdef ptr[:,:,:] _cells
+    # The amount elements in each C-array in '_cells'
+    cdef int[:,:,:] _cell_length
+    # The maximum value of '_cell_length' over all cells,
+    # required for worst case assumption on size of output arrays
+    cdef int _max_cell_length
+    # The length of the cell in each direction (x,y,z)
+    cdef float _cellsize
+    # The minimum and maximum coordinates for all atoms
+    # Used as origin ('_min_coord' is at _cells[0,0,0])
+    # and for bound checks
+    cdef float32[:] _min_coord
+    cdef float32[:] _max_coord
+    # Indicates whether the cell list takes periodicity into account
+    cdef bint _periodic
+    cdef np.ndarray _box
+    # The length of the array before appending periodic copies
+    # if 'periodic' is true
+    cdef int _orig_length
+    cdef float32[:] _orig_min_coord
+    cdef float32[:] _orig_max_coord
+    
+    
+    @cython.initializedcheck(False)
+    @cython.boundscheck(False)
+    @cython.wraparound(False)
+    def __cinit__(self, atom_array not None, float cell_size,
+                  bint periodic=False, box=None, np.ndarray selection=None):
+        cdef float32 x, y, z
+        cdef int i, j, k
+        cdef int atom_array_i
+        cdef int* cell_ptr = NULL
+        cdef int length
+
+        if isinstance(atom_array, AtomArrayStack):
+            raise TypeError("Expected 'AtomArray' but got 'AtomArrayStack'")
+        coord = to_coord(atom_array)
+        # the length of the array before appending periodic copies
+        # if 'periodic' is true
+        self._orig_length = coord.shape[0]
+        self._box = None
+        if coord.ndim != 2:
+            raise ValueError("Coordinates must have shape (n,3)")
+        if coord.shape[0] == 0:
+            raise ValueError("Coordinates must not be empty")
+        if coord.shape[1] != 3:
+            raise ValueError("Coordinates must have form (x,y,z)")
+        if np.isnan(coord).any():
+            raise ValueError("Coordinates contain NaN values")
+
+        if periodic:
+            if box is not None:
+                self._box = box
+            elif atom_array.box is not None:
+                if atom_array.box.shape != (3,3):
+                    raise ValueError(
+                        "Box has invalid shape"
+                    )
+                self._box = atom_array.box
+            else:
+                raise ValueError(
+                    "AtomArray must have a box to enable periodicity"
+                )
+            if np.isnan(self._box).any():
+                raise ValueError("Box contains NaN values")
+            coord = move_inside_box(coord, self._box)
+            coord, indices = repeat_box_coord(coord, self._box)
+        
+        if self._has_initialized_cells():
+            raise Exception("Duplicate call of constructor")
+        self._cells = None
+        if cell_size <= 0:
+            raise ValueError("Cell size must be greater than 0")
+        self._periodic = periodic
+        self._coord = coord.astype(np.float32, copy=False)
+        self._cellsize = cell_size
+        # calculate how many cells are required for each dimension
+        min_coord = np.min(coord, axis=0).astype(np.float32)
+        max_coord = np.max(coord, axis=0).astype(np.float32)
+        self._min_coord = min_coord
+        self._max_coord = max_coord
+        cell_count = (((max_coord - min_coord) / cell_size) +1).astype(int)
+        if self._periodic:
+            self._orig_min_coord = np.min(coord[:self._orig_length], axis=0) \
+                                   .astype(np.float32)
+            self._orig_max_coord = np.max(coord[:self._orig_length], axis=0) \
+                                   .astype(np.float32)
+        
+        # ndarray of pointers to C-arrays
+        # containing indices to atom array
+        self._cells = np.zeros(cell_count, dtype=np.uint64)
+        # Stores the length of the C-arrays
+        self._cell_length = np.zeros(cell_count, dtype=np.int32)
+        
+        # Prepare selection
+        if selection is not None:
+            self._has_selection = True
+            self._selection = np.frombuffer(selection, dtype=np.uint8)
+            if self._selection.shape[0] != self._orig_length:
+                raise IndexError(
+                    f"Atom array has length {self._orig_length}, "
+                    f"but selection has length {self._selection.shape[0]}"
+                )
+        else:
+            self._has_selection = False
+        
+        # Fill cells
+        for atom_array_i in range(self._coord.shape[0]):
+            # Only put selected atoms into cell list
+            if not self._has_selection \
+               or self._selection[atom_array_i % self._orig_length]:
+                    x = self._coord[atom_array_i, 0]
+                    y = self._coord[atom_array_i, 1]
+                    z = self._coord[atom_array_i, 2]
+                    # Get cell indices for coordinates
+                    self._get_cell_index(x, y, z, &i, &j, &k)
+                    # Increment cell length and reallocate
+                    length = self._cell_length[i,j,k] + 1
+                    cell_ptr = <int*>self._cells[i,j,k]
+                    cell_ptr = <int*>realloc(cell_ptr, length * sizeof(int))
+                    if not cell_ptr:
+                        raise MemoryError()
+                    # Potentially increase max cell length
+                    if length > self._max_cell_length:
+                        self._max_cell_length = length
+                    # Store atom array index in respective cell
+                    cell_ptr[length-1] = atom_array_i
+                    # Store new cell pointer and length
+                    self._cell_length[i,j,k] = length
+                    self._cells[i,j,k] = <ptr> cell_ptr
+            
+    
+    def __dealloc__(self):
+        if self._has_initialized_cells():
+            deallocate_ptrs(self._cells)
+    
+    
+    @cython.initializedcheck(False)
+    @cython.boundscheck(False)
+    @cython.wraparound(False)
+    def create_adjacency_matrix(self, float32 threshold_distance):
+        """
+        create_adjacency_matrix(threshold_distance)
+        
+        Create an adjacency matrix for the atoms in this cell list.
+
+        An adjacency matrix depicts which atoms *i* and *j* have
+        a distance lower than a given threshold distance.
+        The values in the adjacency matrix ``m`` are
+        ``m[i,j] = 1 if distance(i,j) <= threshold else 0``
+        
+        Parameters
+        ----------
+        threshold_distance : float
+            The threshold distance. All atom pairs that have a distance
+            lower than this value are indicated by ``True`` values in
+            the resulting matrix.
+        
+        Returns
+        -------
+        matrix : ndarray, dtype=bool, shape=(n,n)
+            An *n x n* adjacency matrix.
+            If a `selection` was given to the constructor of the
+            :class:`CellList`, the rows and columns corresponding to
+            atoms, that are not masked by the selection, have all
+            elements set to ``False``.
+        
+        Notes
+        -----
+        The highest performance is achieved when the the cell size is
+        equal to the threshold distance. However, this is purely
+        optinal: The resulting adjacency matrix is the same for every
+        cell size.
+
+        Although the adjacency matrix should be symmetric in most cases,
+        it may occur that ``m[i,j] != m[j,i]``, when ``distance(i,j)``
+        is very close to the `threshold_distance` due to numerical
+        errors.
+        The matrix can be symmetrized with ``numpy.maximum(a, a.T)``.
+
+        Examples
+        --------
+        Create adjacency matrix for CA atoms in a structure:
+
+        >>> atom_array = atom_array[atom_array.atom_name == "CA"]
+        >>> cell_list = CellList(atom_array, 5)
+        >>> matrix = cell_list.create_adjacency_matrix(5)
+        """
+        if threshold_distance < 0:
+            raise ValueError("Threshold must be a positive value")
+        cdef int i=0
+        
+        # Get atom position for all original positions
+        # (no periodic copies)
+        coord = np.asarray(self._coord[:self._orig_length])
+
+        if self._has_selection:
+            selection = np.asarray(self._selection, dtype=bool)
+            # Create matrix with all elements set to False
+            matrix = np.zeros(
+                (self._orig_length, self._orig_length), dtype=bool
+            )
+            # Set only those rows that belong to masked atoms
+            matrix[selection, :] = self.get_atoms(
+                coord[selection], threshold_distance, as_mask=True
+            )
+            return matrix
+        else:
+            return self.get_atoms(coord, threshold_distance, as_mask=True)
+        
+    
+    @cython.initializedcheck(False)
+    @cython.boundscheck(False)
+    @cython.wraparound(False)
+    def get_atoms(self, np.ndarray coord, radius, bint as_mask=False):
+        """
+        get_atoms(coord, radius, as_mask=False)
+        
+        Find atoms with a maximum distance from given coordinates.
+        
+        Parameters
+        ----------
+        coord : ndarray, dtype=float, shape=(3,) or shape=(m,3)
+            The central coordinates, around which the atoms are
+            searched.
+            If a single position is given, the indices of atoms in its
+            radius are returned.
+            Multiple positions (2-D :class:`ndarray`) have a vectorized
+            behavior:
+            Each row in the resulting :class:`ndarray` contains the
+            indices for the corresponding position.
+            Since the positions may have different amounts of adjacent
+            atoms, trailing `-1` values are used to indicate nonexisting
+            indices.
+        radius : float or ndarray, shape=(n,), dtype=float, optional
+            The radius around `coord`, in which the atoms are searched,
+            i.e. all atoms in `radius` distance to `coord` are returned.
+            Either a single radius can be given as scalar, or individual
+            radii for each position in `coord` can be provided as
+            :class:`ndarray`.
+        as_mask : bool, optional
+            If true, the result is returned as boolean mask, instead
+            of an index array.
+        
+        Returns
+        -------
+        indices : ndarray, dtype=int32, shape=(p,) or shape=(m,p)
+            The indices of the atom array, where the atoms are in the
+            defined `radius` around `coord`.
+            If `coord` contains multiple positions, this return value is
+            two-dimensional with trailing `-1` values for empty values.
+            Only returned with `as_mask` set to false.
+        mask : ndarray, dtype=bool, shape=(m,n) or shape=(n,)
+            Same as `indices`, but as boolean mask.
+            The values are true for atoms in the atom array,
+            that are in the defined vicinity.
+            Only returned with `as_mask` set to true.
+            
+        See Also
+        --------
+        get_atoms_in_cells
+
+        Notes
+        -----
+        In case of a :class:`CellList` with `periodic` set to `True`:
+        If more than one periodic copy of an atom is within the
+        threshold radius, the returned `indices` array contains the
+        corresponding index multiple times.
+        Please use ``numpy.unique()``, if this is undesireable.
+
+        Examples
+        --------
+        Get adjacent atoms for a single position:
+
+        >>> cell_list = CellList(atom_array, 3)
+        >>> pos = np.array([1.0, 2.0, 3.0])
+        >>> indices = cell_list.get_atoms(pos, radius=2.0)
+        >>> print(indices)
+        [102 104 112]
+        >>> print(atom_array[indices])
+            A       6  TRP CE3    C         0.779    0.524    2.812
+            A       6  TRP CZ3    C         1.439    0.433    4.053
+            A       6  TRP HE3    H        -0.299    0.571    2.773
+        >>> indices = cell_list.get_atoms(pos, radius=3.0)
+        >>> print(atom_array[indices])
+            A       6  TRP CD2    C         1.508    0.564    1.606
+            A       6  TRP CE3    C         0.779    0.524    2.812
+            A       6  TRP CZ3    C         1.439    0.433    4.053
+            A       6  TRP HE3    H        -0.299    0.571    2.773
+            A       6  TRP HZ3    H         0.862    0.400    4.966
+            A       3  TYR CZ     C        -0.639    3.053    5.043
+            A       3  TYR HH     H         1.187    3.395    5.567
+            A      19  PRO HD2    H         0.470    3.937    1.260
+            A       6  TRP CE2    C         2.928    0.515    1.710
+            A       6  TRP CH2    C         2.842    0.407    4.120
+            A      18  PRO HA     H         2.719    3.181    1.316
+            A      18  PRO HB3    H         2.781    3.223    3.618
+            A      18  PRO CB     C         3.035    4.190    3.187
+        
+        Get adjacent atoms for mutliple positions:
+
+        >>> cell_list = CellList(atom_array, 3)
+        >>> pos = np.array([[1.0,2.0,3.0], [2.0,3.0,4.0], [3.0,4.0,5.0]])
+        >>> indices = cell_list.get_atoms(pos, radius=3.0)
+        >>> print(indices)
+        [[ 99 102 104 112 114  45  55 290 101 105 271 273 268  -1  -1]
+         [104 114  45  46  55  44  54 105 271 273 265 268 269 272 275]
+         [ 46  55 273 268 269 272 274 275  -1  -1  -1  -1  -1  -1  -1]]
+        >>> # Convert to list of arrays and remove trailing -1
+        >>> indices = [row[row != -1] for row in indices]
+        >>> for row in indices:
+        ...     print(row)
+        [ 99 102 104 112 114  45  55 290 101 105 271 273 268]
+        [104 114  45  46  55  44  54 105 271 273 265 268 269 272 275]
+        [ 46  55 273 268 269 272 274 275]
+        """
+        cdef int i=0, j=0
+        cdef int array_i = 0
+        cdef int max_array_length = 0
+        cdef int coord_index
+        cdef float32 x1, y1, z1, x2, y2, z2
+        cdef float32 sq_dist
+        cdef float32 sq_radius
+        cdef float32[:] sq_radii
+        cdef np.ndarray cell_radii
+        
+        cdef int[:,:] all_indices
+        cdef int[:,:] indices
+        cdef float32[:,:] coord_v
+
+        if len(coord) == 0:
+            return _empty_result(as_mask)
+        
+        # Handle periodicity for the input coordinates
+        if self._periodic:
+            coord = move_inside_box(coord, self._box)
+        # Convert input parameters into a uniform format
+        coord, radius, is_multi_coord, is_multi_radius \
+            = _prepare_vectorization(coord, radius, np.float32)
+        if is_multi_radius:
+            sq_radii = radius * radius
+            cell_radii = np.ceil(radius / self._cellsize).astype(np.int32)
+        else:
+            # All radii are equal
+            sq_radii = np.full(
+                len(coord), radius[0]*radius[0], dtype=np.float32
+            )
+            cell_radii = np.full(
+                len(coord),
+                int(np.ceil(radius[0] / self._cellsize)),
+                dtype=np.int32
+            )
+
+        # Get indices for adjacent atoms, based on a cell radius
+        all_indices = self._get_atoms_in_cells(
+            coord, cell_radii, is_multi_radius
+        )
+        # These have to be narrowed down in the next step
+        # using the Euclidian distance
+        
+        # Filter all indices from all_indices
+        # where squared distance is smaller than squared radius
+        # Using the squared distance is computationally cheaper than
+        # calculating the sqaure root for every distance
+        indices = np.full(
+            (all_indices.shape[0], all_indices.shape[1]), -1, dtype=np.int32
+        )
+        coord_v = coord
+        for i in range(all_indices.shape[0]):
+            sq_radius = sq_radii[i]
+            x1 = coord_v[i,0]
+            y1 = coord_v[i,1]
+            z1 = coord_v[i,2]
+            array_i = 0
+            for j in range(all_indices.shape[1]):
+                coord_index = all_indices[i,j]
+                if coord_index != -1:
+                    x2 = self._coord[coord_index, 0]
+                    y2 = self._coord[coord_index, 1]
+                    z2 = self._coord[coord_index, 2]
+                    sq_dist = squared_distance(x1, y1, z1, x2, y2, z2)
+                    if sq_dist <= sq_radius:
+                        indices[i, array_i] = coord_index
+                        array_i += 1
+            if array_i > max_array_length:
+                max_array_length = array_i
+        
+        return self._post_process(
+            np.asarray(indices)[:, :max_array_length],
+            as_mask, is_multi_coord
+        )
+    
+    
+    @cython.boundscheck(False)
+    @cython.wraparound(False)
+    def get_atoms_in_cells(self, np.ndarray coord,
+                           cell_radius=1, bint as_mask=False):
+        """
+        get_atoms_in_cells(coord, cell_radius=1, as_mask=False)
+        
+        Find atoms with a maximum cell distance from given
+        coordinates.
+        
+        Instead of using the radius as maximum euclidian distance to the
+        given coordinates,
+        the radius is measured as the amount of cells:
+        A radius of 0 means, that only the atoms in the same cell
+        as the given coordinates are considered. A radius of 1 means,
+        that the atoms indices from this cell and the 8 surrounding
+        cells are returned and so forth.
+        This is more efficient than `get_atoms()`.
+        
+        Parameters
+        ----------
+        coord : ndarray, dtype=float, shape=(3,) or shape=(m,3)
+            The central coordinates, around which the atoms are
+            searched.
+            If a single position is given, the indices of atoms in its
+            cell radius are returned.
+            Multiple positions (2-D :class:`ndarray`) have a vectorized
+            behavior:
+            Each row in the resulting :class:`ndarray` contains the
+            indices for the corresponding position.
+            Since the positions may have different amounts of adjacent
+            atoms, trailing `-1` values are used to indicate nonexisting
+            indices.
+        cell_radius : int or ndarray, shape=(n,), dtype=int, optional
+            The radius around `coord` (in amount of cells), in which
+            the atoms are searched. This does not correspond to the
+            Euclidian distance used in `get_atoms()`. In this case, all
+            atoms in the cell corresponding to `coord` and in adjacent
+            cells are returned.
+            Either a single radius can be given as scalar, or individual
+            radii for each position in `coord` can be provided as
+            :class:`ndarray`.
+            By default atoms are searched in the cell of `coord`
+            and directly adjacent cells (cell_radius = 1).
+        as_mask : bool, optional
+            If true, the result is returned as boolean mask, instead
+            of an index array.
+        
+        Returns
+        -------
+        indices : ndarray, dtype=int32, shape=(p,) or shape=(m,p)
+            The indices of the atom array, where the atoms are in the
+            defined `radius` around `coord`.
+            If `coord` contains multiple positions, this return value is
+            two-dimensional with trailing `-1` values for empty values.
+            Only returned with `as_mask` set to false.
+        mask : ndarray, dtype=bool, shape=(m,n) or shape=(n,)
+            Same as `indices`, but as boolean mask.
+            The values are true for atoms in the atom array,
+            that are in the defined vicinity.
+            Only returned with `as_mask` set to true.
+
+        See Also
+        --------
+        get_atoms
+
+        Notes
+        -----
+        In case of a :class:`CellList` with `periodic` set to `True`:
+        If more than one periodic copy of an atom is within the
+        threshold radius, the returned `indices` array contains the
+        corresponding index multiple times.
+        Please use ``numpy.unique()``, if this is undesireable.
+        """
+        # This function is a thin wrapper around the private method
+        # with the same name, with addition of handling periodicty
+        # and the ability to return a mask instead of indices
+
+        if len(coord) == 0:
+            return _empty_result(as_mask)
+
+        # Handle periodicity for the input coordinates
+        if self._periodic:
+            coord = move_inside_box(coord, self._box)
+        # Convert input parameters into a uniform format
+        coord, cell_radius, is_multi_coord, is_multi_radius \
+            = _prepare_vectorization(coord, cell_radius, np.int32)
+        # Get adjacent atom indices
+        array_indices = self._get_atoms_in_cells(
+            coord, cell_radius, is_multi_radius
+        )
+        return self._post_process(array_indices, as_mask, is_multi_coord)
+    
+    
+    @cython.boundscheck(False)
+    @cython.wraparound(False)
+    def _get_atoms_in_cells(self,
+                            np.ndarray coord,
+                            np.ndarray cell_radii,
+                            bint is_multi_radius):
+        """
+        Get the indices of atoms in `cell_radii` adjacency of `coord`.
+        
+        Parameters
+        ----------
+        coord : ndarray, dtype=float32, shape=(n,3)
+            The position to find adjacent atoms for.
+        cell_radii : ndarray, dtype=int32, shape=(n)
+            The radius for each position.
+        is_multi_radius : bool
+            True indicates, that all values in `cell_radii` are the
+            same.
+        
+        Returns
+        -------
+        array_indices : ndarray, dtype=int32, shape=(m,p)
+            Indices of adjancent atoms.
+        """
+
+        cdef int max_cell_radius
+        if is_multi_radius:
+            max_cell_radius = np.max(cell_radii)
+        else:
+            # All radii are equal
+            max_cell_radius = cell_radii[0]
+        # Worst case assumption on index array length requirement:
+        # At maximum, the amount of adjacent atoms can only be the
+        # maximum amount of atoms per cell times the amount of cells
+        # Since the cells extend in 3 dimensions the amount of cells is
+        # (2*r + 1)**3
+        cdef int length = (2*max_cell_radius + 1)**3 * self._max_cell_length
+        array_indices = np.full((len(coord), length), -1, dtype=np.int32)
+        # Fill index array
+        cdef int max_array_length \
+            = self._find_adjacent_atoms(coord, array_indices, cell_radii)
+        return array_indices[:, :max_array_length]
+    
+    
+    @cython.boundscheck(False)
+    @cython.wraparound(False)
+    cdef int _find_adjacent_atoms(self,
+                                  float32[:,:] coord,
+                                  int[:,:] indices,
+                                  int[:] cell_radius):
+        """
+        This method fills the given empty index array
+        with actual indices of adjacent atoms.
+    
+        Since the length of 'indices' (second dimension) is
+        the worst case assumption, this method returns the actual
+        required length, i.e. the highest length of all arrays
+        in this 'array of arrays'.
+        """
+        cdef int length
+        cdef int* list_ptr
+        cdef float32 x, y,z
+        cdef int i=0, j=0, k=0
+        cdef int adj_i, adj_j, adj_k
+        cdef int pos_i, array_i, cell_i
+        cdef int max_array_length = 0
+        cdef int cell_r
+        
+        cdef ptr[:,:,:] cells = self._cells
+        cdef int[:,:,:] cell_length = self._cell_length
+
+        for pos_i in range(coord.shape[0]):
+            array_i = 0
+            cell_r = cell_radius[pos_i]
+            x = coord[pos_i, 0]
+            y = coord[pos_i, 1]
+            z = coord[pos_i, 2]
+            self._get_cell_index(x, y, z, &i, &j, &k)
+            # Look into cells of the indices and adjacent cells
+            # in all 3 dimensions
+
+            for adj_i in range(i-cell_r, i+cell_r+1):
+                if (adj_i >= 0 and adj_i < cells.shape[0]):
+                    for adj_j in range(j-cell_r, j+cell_r+1):
+                        if (adj_j >= 0 and adj_j < cells.shape[1]):
+                            for adj_k in range(k-cell_r, k+cell_r+1):
+                                if (adj_k >= 0 and adj_k < cells.shape[2]):
+                                    # Fill index array
+                                    # with indices in cell
+                                    list_ptr = <int*>cells[adj_i, adj_j, adj_k]
+                                    length = cell_length[adj_i, adj_j, adj_k]
+                                    for cell_i in range(length):
+                                        indices[pos_i, array_i] = \
+                                            list_ptr[cell_i]
+                                        array_i += 1
+            if array_i > max_array_length:
+                max_array_length = array_i
+        return max_array_length
+    
+
+    @cython.boundscheck(False)
+    @cython.wraparound(False)
+    def _post_process(self,
+                     np.ndarray indices,
+                     bint as_mask,
+                     bint is_multi_coord):
+        """
+        Post process the resulting indices of adjacent atoms,
+        including periodicity handling and optional conversion into a
+        boolean matrix.
+        """
+        # Handle periodicity for the output indices
+        if self._periodic:
+            # Map indices of repeated coordinates to original
+            # coordinates, i.e. the coordinates in the central box
+            # -> Remainder of dividing index by original array length
+            # Furthermore this ensures, that the indices have valid
+            # values for '_as_mask()'
+            indices[indices != -1] %= self._orig_length
+        if as_mask:
+            matrix = self._as_mask(indices)
+            if is_multi_coord:
+                return matrix
+            else:
+                return matrix[0]
+        else:
+            if is_multi_coord:
+                return indices
+            else:
+                return indices[0]
+    
+    
+    @cython.initializedcheck(False)
+    @cython.boundscheck(False)
+    @cython.wraparound(False)
+    @cython.cdivision(True)
+    cdef inline void _get_cell_index(self, float32 x, float32 y, float32 z,
+                             int* i, int* j, int* k):
+        i[0] = <int>((x - self._min_coord[0]) / self._cellsize)
+        j[0] = <int>((y - self._min_coord[1]) / self._cellsize)
+        k[0] = <int>((z - self._min_coord[2]) / self._cellsize)
+    
+    @cython.initializedcheck(False)
+    @cython.boundscheck(False)
+    @cython.wraparound(False)
+    cdef inline bint _check_coord(self, float32 x, float32 y, float32 z):
+        if x < self._min_coord[0] or x > self._max_coord[0]:
+            return False
+        if y < self._min_coord[1] or y > self._max_coord[1]:
+            return False
+        if z < self._min_coord[2] or z > self._max_coord[2]:
+            return False
+        return True
+    
+    @cython.initializedcheck(False)
+    @cython.boundscheck(False)
+    @cython.wraparound(False)
+    cdef np.ndarray _as_mask(self, int[:,:] indices):
+        cdef int i,j
+        cdef int index
+        cdef uint8[:,:] matrix = np.zeros(
+            (indices.shape[0], self._orig_length), dtype=np.uint8
+        )
+        # Fill matrix
+        for i in range(indices.shape[0]):
+            for j in range(indices.shape[1]):
+                index = indices[i,j]
+                if index == -1:
+                    # End of list -> jump to next position
+                    break
+                matrix[i, index] = True
+        return np.asarray(matrix, dtype=bool)
+    
+    cdef inline bint _has_initialized_cells(self):
+        # Memoryviews are not initialized on class creation
+        # This method checks if the _cells memoryview was initialized
+        # and is not None
+        try:
+            if self._cells is not None:
+                return True
+            else:
+                return False
+        except AttributeError:
+            return False
+
+
+def _empty_result(as_mask):
+    """
+    Create return value for :func:`get_atoms()` and
+    :func:`get_atoms_in_cells()`, if no coordinates are given.
+    """
+    if as_mask:
+        return np.array([], dtype=bool)
+    else:
+        return np.array([], dtype=np.int32)
+
+
+def _prepare_vectorization(np.ndarray coord, radius, radius_dtype):
+    """
+    Since `get_atoms()` and `get_atoms_in_cells()`, may take different
+    amount of dimensions for the coordinates and the radius to enable
+    vectorized compuation, each of these functions would need to handle
+    the different cases.
+
+    This function converts the input radius and coordinates into a
+    uniform format and also return, whether single/multiple
+    radii/coordinates were given.
+
+    The shapes before and after conversion are:
+       
+       - coord: (3, ), radius: scalar -> coord: (1,3), radius: (1,)
+       - coord: (n,3), radius: scalar -> coord: (n,3), radius: (n,)
+       - coord: (n,3), radius: (n,  ) -> coord: (n,3), radius: (n,)
+    
+    Thes resulting values have the same dimensionality for all cases and
+    can be handeled uniformly by `get_atoms()` and
+    `get_atoms_in_cells()`.
+    """
+    cdef bint is_multi_coord
+    cdef bint is_multi_radius
+
+    if coord.ndim == 1 and coord.shape[0] == 3:
+        # Single position
+        coord = coord[np.newaxis, :].astype(np.float32, copy=False)
+        is_multi_coord = False
+    elif coord.ndim == 2 and coord.shape[1] == 3:
+        # Multiple positions
+        coord = coord.astype(np.float32, copy=False)
+        is_multi_coord = True
+    else:
+        raise ValueError(
+            f"Invalid shape for input coordinates"
+        )
+    
+    if isinstance(radius, np.ndarray):
+        # Multiple radii
+        # Check whether amount of coordinates match amount of radii
+        if not is_multi_coord:
+            raise ValueError(
+                "Cannot accept array of radii, if a single position is given"
+            )
+        if radius.ndim != 1:
+            raise ValueError("Array of radii must be one-dimensional")
+        if radius.shape[0] != coord.shape[0]:
+            raise ValueError(
+                f"Amount of radii ({radius.shape[0]}) "
+                f"and coordinates ({coord.shape[0]}) are not equal"
+            )
+        if (radius < 0).any():
+            raise ValueError("Radii must be a positive values")
+        radius = radius.astype(radius_dtype, copy=False)
+        is_multi_radius = True
+    else:
+        # Single radius
+        if radius < 0:
+            raise ValueError("Radius must be a positive value")
+        # If only a single integer is given,
+        # create numpy array filled with identical values
+        # with the same length as the coordiantes
+        radius = np.full(coord.shape[0], radius, dtype=radius_dtype)
+        is_multi_radius = False
+
+    return coord, radius, is_multi_coord, is_multi_radius
+
+
+cdef inline void deallocate_ptrs(ptr[:,:,:] ptrs):
+    cdef int i, j, k
+    cdef int* cell_ptr
+    # Free cell pointers
+    for i in range(ptrs.shape[0]):
+        for j in range(ptrs.shape[1]):
+            for k in range(ptrs.shape[2]):
+                cell_ptr = <int*>ptrs[i,j,k]
+                free(cell_ptr)
+
+
+cdef inline float32 squared_distance(float32 x1, float32 y1, float32 z1,
+                    float32 x2, float32 y2, float32 z2):
+    cdef float32 diff_x = x2 - x1
+    cdef float32 diff_y = y2 - y1
+    cdef float32 diff_z = z2 - z1
+    return diff_x*diff_x + diff_y*diff_y + diff_z*diff_z