biotite 1.5.0__cp314-cp314-manylinux_2_24_x86_64.manylinux_2_28_x86_64.whl

biotite/structure/rdf.py

@@ -0,0 +1,245 @@
+# 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 provides functions to calculate the radial distribution function.
+"""
+
+__name__ = "biotite.structure"
+__author__ = "Daniel Bauer, Patrick Kunzmann"
+__all__ = ["rdf"]
+
+from numbers import Integral
+import numpy as np
+from biotite.structure.atoms import AtomArray, coord, stack
+from biotite.structure.box import box_volume
+from biotite.structure.celllist import CellList
+from biotite.structure.geometry import displacement
+from biotite.structure.util import vector_dot
+
+
+def rdf(
+    center, atoms, selection=None, interval=(0, 10), bins=100, box=None, periodic=False
+):
+    r"""
+    Compute the radial distribution function *g(r)* (RDF) for one or
+    multiple given central positions based on a given system of
+    particles.
+
+    Parameters
+    ----------
+    center : Atom or AtomArray or AtomArrayStack or ndarray, dtype=float
+        Coordinates or atoms(s) to use as origin(s) for RDF calculation.
+
+        - If a single :class:`Atom` or a :class:`ndarray` with shape
+          *(3,)* is given, the RDF is only calculated for this position.
+        - If an :class:`AtomArray` or a :class:`ndarray` with shape
+          *(n,3)* is given, the calculated RDF histogram is an average
+          over *n* postions.
+        - If an :class:`AtomArrayStack` or a :class:`ndarray` with shape
+          *(m,n,3)* is given, different centers are used for each model
+          *m*.
+          The calculated RDF histogram is an average over *m*
+          models and *n* positions.
+          This requires `atoms` to be an :class:`AtomArrayStack`.
+
+    atoms : AtomArray or AtomArrayStack
+        The distribution is calculated based on these atoms.
+        When an an :class:`AtomArrayStack` is provided, the RDF
+        histogram is averaged over all models.
+        Please not that `atoms` must have an associated box,
+        unless `box` is set.
+    selection : ndarray, dtype=bool, shape=(n,), optional
+        Boolean mask for `atoms` to limit the RDF calculation to
+        specific atoms.
+    interval : tuple, optional
+        The range in which the RDF is calculated.
+    bins : int or sequence of scalars or str, optional
+        Bins for the RDF.
+
+        - If `bins` is an `int`, it defines the number of bins for the
+          given `interval`.
+        - If `bins` is a sequence, it defines the bin edges, ignoring
+          the `interval` parameter. The output `bins` has the length
+          of this input parameter reduced by one.
+        - If `bins` is a string, it defines the function used to
+          calculate the bins.
+
+        See `numpy.histogram()` for further details.
+    box : ndarray, shape=(3,3) or shape=(m,3,3), optional
+        If this parameter is set, the given box is used instead of the
+        `box` attribute of `atoms`.
+        Must have shape *(3,3)* if atoms is an :class:`AtomArray` or
+        *(m,3,3)* if atoms is an :class:`AtomArrayStack`, respectively.
+    periodic : bool, optional
+        Defines if periodic boundary conditions are taken into account.
+
+    Returns
+    -------
+    bins : ndarray, dtype=float, shape=n
+        The centers of the histogram bins.
+        The length of the array is given by the `bins` input parameter.
+    rdf : ndarry, dtype=float, shape=n
+        RDF values for every bin.
+
+    Notes
+    -----
+    Since the RDF depends on the average particle density of the system,
+    this function strictly requires an box.
+
+    Examples
+    --------
+    Calculate the oxygen-oxygen radial distribution function of water.
+    The range of the histogram starts at 0.2 Å, in order to
+    ignore the counts for the density for each oxygen to itself.
+
+    >>> from os.path import join
+    >>> waterbox = load_structure(join(path_to_structures, "waterbox.gro"))
+    >>> oxygens = waterbox[:, waterbox.atom_name == 'OW']
+    >>> bins, g_r = rdf(oxygens, oxygens, bins=49, interval=(0.2, 10), periodic=True)
+
+    Print the RDF depending on the radius. Bins are in Å.
+
+    >>> for x, y in zip(bins, g_r):
+    ...     print(f"{x:.2f} {y:.2f}")
+    0.30 0.00
+    0.50 0.00
+    0.70 0.04
+    0.90 0.02
+    1.10 0.03
+    1.30 0.06
+    1.50 0.03
+    1.70 0.04
+    1.90 0.04
+    2.10 0.04
+    2.30 0.04
+    2.50 0.16
+    2.70 1.99
+    2.90 2.22
+    3.10 1.34
+    3.30 1.04
+    3.50 0.97
+    3.70 0.94
+    3.90 0.97
+    4.10 0.94
+    4.30 0.98
+    4.50 0.97
+    4.70 0.96
+    4.90 0.99
+    5.10 0.99
+    5.30 1.02
+    5.50 1.02
+    5.70 0.99
+    5.90 0.98
+    6.10 0.98
+    6.30 0.99
+    6.50 1.02
+    6.70 1.02
+    6.90 1.00
+    7.10 1.01
+    7.30 1.01
+    7.50 1.00
+    7.70 1.01
+    7.90 0.99
+    8.10 0.99
+    8.30 0.99
+    8.50 0.99
+    8.70 0.99
+    8.90 1.00
+    9.10 1.01
+    9.30 1.01
+    9.50 1.00
+    9.70 1.00
+    9.90 0.99
+
+    Find the radius for the first solvation shell.
+    In this simple case, the density peak is identified by finding
+    the maximum of the function.
+
+    >>> peak_position = np.argmax(g_r)
+    >>> print(f"{bins[peak_position]/10:.2f} nm")
+    0.29 nm
+    """
+    if isinstance(atoms, AtomArray):
+        # Reshape always to a stack for easier calculation
+        atoms = stack([atoms])
+    if selection is not None:
+        atoms = atoms[..., selection]
+
+    atom_coord = atoms.coord
+
+    if box is None:
+        if atoms.box is None:
+            raise ValueError("A box must be supplied")
+        else:
+            box = atoms.box
+    elif box.ndim == 2 and atoms.stack_depth() == 1:
+        box = box[np.newaxis, :, :]
+
+    center = coord(center)
+    if center.ndim == 1:
+        center = center.reshape((1, 1) + center.shape)
+    elif center.ndim == 2:
+        center = center.reshape((1,) + center.shape)
+
+    if box.shape[0] != center.shape[0] or box.shape[0] != atom_coord.shape[0]:
+        raise ValueError("Center, box, and atoms must have the same model count")
+
+    # Calculate distance histogram
+    edges = _calculate_edges(interval, bins)
+    # Make histogram of quared distances to save computation time
+    # of sqrt calculation
+    sq_edges = edges**2
+    threshold_dist = edges[-1]
+    cell_size = threshold_dist
+    disp = []
+    for i in range(atoms.stack_depth()):
+        # Use cell list to efficiently preselect atoms that are in range
+        # of the desired bin range
+        cell_list = CellList(atom_coord[i], cell_size, periodic, box[i])
+        # 'cell_radius=1' is used in 'get_atoms_in_cells()'
+        # This is enough to find all atoms that are in the given
+        # interval (and more), since the size of each cell is as large
+        # as the last edge of the bins
+        near_atom_mask = cell_list.get_atoms_in_cells(center[i], as_mask=True)
+        # Calculate distances of each center to preselected atoms
+        # for each center
+        for j in range(center.shape[1]):
+            dist_box = box[i] if periodic else None
+            # Calculate squared distances
+            disp.append(
+                displacement(
+                    center[i, j], atom_coord[i, near_atom_mask[j]], box=dist_box
+                )
+            )
+    # Make one array from multiple arrays with different length
+    disp = np.concatenate(disp)
+    sq_distances = vector_dot(disp, disp)
+    hist, _ = np.histogram(sq_distances, bins=sq_edges)
+
+    # Normalize with average particle density (N/V) in each bin
+    bin_volume = (4 / 3 * np.pi * np.power(edges[1:], 3)) - (
+        4 / 3 * np.pi * np.power(edges[:-1], 3)
+    )
+    n_frames = len(atoms)
+    volume = box_volume(box).mean()
+    density = atoms.array_length() / volume
+    g_r = hist / (bin_volume * density * n_frames)
+
+    # Normalize with number of centers
+    g_r /= center.shape[1]
+
+    bin_centers = (edges[:-1] + edges[1:]) * 0.5
+
+    return bin_centers, g_r
+
+
+def _calculate_edges(interval, bins):
+    if isinstance(bins, Integral):
+        if bins < 1:
+            raise ValueError("At least one bin is required")
+        return np.linspace(*interval, bins + 1)
+    else:
+        # 'bins' contains edges
+        return np.array(bins, dtype=float)

biotite/structure/repair.py

@@ -0,0 +1,302 @@
+# 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 functionalities for repairing malformed structures.
+"""
+
+__name__ = "biotite.structure"
+__author__ = "Patrick Kunzmann, Daniel Bauer"
+__all__ = ["create_continuous_res_ids", "infer_elements", "create_atom_names"]
+
+import warnings
+from collections import Counter
+import numpy as np
+from biotite.structure.atoms import AtomArray, AtomArrayStack
+from biotite.structure.chains import get_chain_starts
+from biotite.structure.residues import get_residue_starts
+
+
+def create_continuous_res_ids(atoms, restart_each_chain=True):
+    """
+    Create an array of continuous residue IDs for a given structure.
+
+    This means that residue IDs are incremented by 1 for each residue.
+
+    Parameters
+    ----------
+    atoms : AtomArray or AtomArrayStack
+        The atoms for which the continuous residue IDs should be created.
+    restart_each_chain : bool, optional
+        If true, the residue IDs are reset to 1 for each chain.
+
+    Returns
+    -------
+    res_ids : ndarray, dtype=int
+        The continuous residue IDs.
+
+    Examples
+    --------
+
+    >>> # Remove a residue to make the residue IDs discontinuous
+    >>> atom_array = atom_array[atom_array.res_id != 5]
+    >>> res_ids, _ = get_residues(atom_array)
+    >>> print(res_ids)
+    [ 1  2  3  4  6  7  8  9 10 11 12 13 14 15 16 17 18 19 20]
+    >>> atom_array.res_id = create_continuous_res_ids(atom_array)
+    >>> res_ids, _ = get_residues(atom_array)
+    >>> print(res_ids)
+    [ 1  2  3  4  5  6  7  8  9 10 11 12 13 14 15 16 17 18 19]
+    """
+    res_ids_diff = np.zeros(atoms.array_length(), dtype=int)
+    res_starts = get_residue_starts(atoms)
+    res_ids_diff[res_starts] = 1
+    res_ids = np.cumsum(res_ids_diff)
+
+    if restart_each_chain:
+        chain_starts = get_chain_starts(atoms)
+        for start in chain_starts:
+            res_ids[start:] -= res_ids[start] - 1
+
+    return res_ids
+
+
+def infer_elements(atoms):
+    """
+    Infer the elements of atoms based on their atom name.
+
+    Parameters
+    ----------
+    atoms : AtomArray or AtomArrayStack or array-like of str
+        The atoms for which the elements should be inferred.
+        Alternatively the atom names can be passed directly.
+
+    Returns
+    -------
+    elements : ndarray, dtype=str
+        The inferred elements.
+
+    See Also
+    --------
+    create_atoms_names : The opposite of this function.
+
+    Examples
+    --------
+
+    >>> print(infer_elements(atom_array)[:10])
+    ['N' 'C' 'C' 'O' 'C' 'C' 'O' 'N' 'H' 'H']
+    >>> print(infer_elements(["CA", "C", "C1", "OD1", "HD21", "1H", "FE"]))
+    ['C' 'C' 'C' 'O' 'H' 'H' 'FE']
+    """
+    if isinstance(atoms, (AtomArray, AtomArrayStack)):
+        atom_names = atoms.atom_name
+    else:
+        atom_names = atoms
+    return np.array([_guess_element(name) for name in atom_names])
+
+
+def create_atom_names(atoms):
+    """
+    Create atom names for a single residue based on elements.
+
+    The atom names are simply enumerated separately for each element.
+
+    Parameters
+    ----------
+    atoms : AtomArray or AtomArrayStack or array-like of str
+        The atoms for which the atom names should be created.
+        Alternatively the elements can be passed directly.
+
+    Returns
+    -------
+    atom_names : ndarray, dtype=str
+        The atom names.
+
+    See Also
+    --------
+    infer_elements : The opposite of this function.
+
+    Notes
+    -----
+    The atom names created this way may differ from the ones in the
+    original source, as different schemes for atom naming exist.
+    This function only ensures that the created atom names are unique.
+    This is e.g. necessary for writing bonds to PDBx files.
+
+    Note that this function should be used only on single residues,
+    otherwise enumeration would continue in the next residue.
+
+    Examples
+    --------
+
+    >>> atoms = residue("URA")  # Uracil
+    >>> print(atoms.element)
+    ['N' 'C' 'O' 'N' 'C' 'O' 'C' 'C' 'H' 'H' 'H' 'H']
+    >>> print(create_atom_names(atoms))
+    ['N1' 'C1' 'O1' 'N2' 'C2' 'O2' 'C3' 'C4' 'H1' 'H2' 'H3' 'H4']
+    """
+    if isinstance(atoms, (AtomArray, AtomArrayStack)):
+        elements = atoms.element
+    else:
+        elements = atoms
+
+    atom_names = np.zeros(len(elements), dtype="U6")
+    element_counter = Counter()
+    for i, elem in enumerate(elements):
+        element_counter[elem] += 1
+        atom_names[i] = f"{elem}{element_counter[elem]}"
+    return atom_names
+
+
+_elements = [
+    elem.upper()
+    for elem in [
+        "H",
+        "He",
+        "Li",
+        "Be",
+        "B",
+        "C",
+        "N",
+        "O",
+        "F",
+        "Ne",
+        "Na",
+        "Mg",
+        "Al",
+        "Si",
+        "P",
+        "S",
+        "Cl",
+        "Ar",
+        "K",
+        "Ca",
+        "Sc",
+        "Ti",
+        "V",
+        "Cr",
+        "Mn",
+        "Fe",
+        "Co",
+        "Ni",
+        "Cu",
+        "Zn",
+        "Ga",
+        "Ge",
+        "As",
+        "Se",
+        "Br",
+        "Kr",
+        "Rb",
+        "Sr",
+        "Y",
+        "Zr",
+        "Nb",
+        "Mo",
+        "Tc",
+        "Ru",
+        "Rh",
+        "Pd",
+        "Ag",
+        "Cd",
+        "In",
+        "Sn",
+        "Sb",
+        "Te",
+        "I",
+        "Xe",
+        "Cs",
+        "Ba",
+        "La",
+        "Ce",
+        "Pr",
+        "Nd",
+        "Pm",
+        "Sm",
+        "Eu",
+        "Gd",
+        "Tb",
+        "Dy",
+        "Ho",
+        "Er",
+        "Tm",
+        "Yb",
+        "Lu",
+        "Hf",
+        "Ta",
+        "W",
+        "Re",
+        "Os",
+        "Ir",
+        "Pt",
+        "Au",
+        "Hg",
+        "Tl",
+        "Pb",
+        "Bi",
+        "Po",
+        "At",
+        "Rn",
+        "Fr",
+        "Ra",
+        "Ac",
+        "Th",
+        "Pa",
+        "U",
+        "Np",
+        "Pu",
+        "Am",
+        "Cm",
+        "Bk",
+        "Cf",
+        "Es",
+        "Fm",
+        "Md",
+        "No",
+        "Lr",
+        "Rf",
+        "Db",
+        "Sg",
+        "Bh",
+        "Hs",
+        "Mt",
+        "Ds",
+        "Rg",
+        "Cn",
+        "Nh",
+        "Fl",
+        "Mc",
+        "Lv",
+        "Ts",
+        "Og",
+    ]
+]
+
+
+def _guess_element(atom_name):
+    # remove digits (1H -> H)
+    elem = "".join([i for i in atom_name if not i.isdigit()])
+    elem = elem.upper()
+    if len(elem) == 0:
+        return ""
+
+    # Some often used elements for biomolecules
+    if (
+        elem.startswith("C")
+        or elem.startswith("N")
+        or elem.startswith("O")
+        or elem.startswith("S")
+        or elem.startswith("H")
+    ):
+        return elem[0]
+
+    # Exactly match element abbreviations
+    try:
+        return _elements[_elements.index(elem[:2])]
+    except ValueError:
+        try:
+            return _elements[_elements.index(elem[0])]
+        except ValueError:
+            warnings.warn(f"Could not infer element for '{atom_name}'")
+            return ""