<commit_before>"""Helper module for structmap package. Contains a collection of tools
for analysing a pdb file.
"""
from __future__ import absolute_import, division, print_function

from Bio.SeqIO import PdbIO
from Bio.SeqUtils import seq1, ProtParamData
from Bio.Data.SCOPData import protein_letters_3to1
from scipy.spatial import distance
import numpy as np
from .seqtools import _construct_sub_align, blast_sequences
from . import gentests

ss_lookup_dict = {
    'H': 0,
    'B': 1,
    'E': 2,
    'G': 3,
    'I': 4,
    'T': 5,
    'S': 6,
    '-': 7,
    0: 'H',
    1: 'B',
    2: 'E',
    3: 'G',
    4: 'I',
    5: 'T',
    6: 'S',
    7: '-'
    }

def _euclidean_distance_matrix(chain, selector='all'):
    """Compute the Euclidean distance matrix for all atoms in a pdb chain.
    Return the euclidean distance matrix and a reference list of all atoms
    in the model (positionally matched to the euclidean matrix).
    Optional parameter is the atom in each residue with which to compute a
    distance matrix. The default setting is 'all', which gets all
    non-heterologous atoms. Other potential options include 'CA', 'CB' etc.
    If an atom is not found within a residue object, then method reverts to
    using 'CA'.
    """
    reference = []
    coords = []
    #Filter on non-HET atoms
    for residue in (x for x in chain if x.get_id()[0] == ' '):
        #If selecting based on all atoms within residue
        if selector == 'all':
            for atom in residue:
                coords.append(atom.get_coord())
                reference.append(atom.get_full_id()[3][1])
        #If measuring distance on particular atoms
        else:
            if selector in residue:
                select_atom = selector
            #Revert to carbon alpha if atom is not found
            elif 'CA' in residue:
                select_atom = 'CA'
            #if CA is not found, do not include residue in distance matrix
            else:
                continue
            coords.append(residue[select_atom].get_coord())
            reference.append(residue[select_atom].get_full_id()[3][1])
    #Convert to a np array, and compute Euclidean distance.
    coord_array = np.array(coords)
    euclid_mat = distance.pdist(coord_array, 'euclidean')
    #Convert to squareform matrix
    euclid_mat = distance.squareform(euclid_mat)
    ref_array = np.array(reference)
    return euclid_mat, ref_array

def nearby(chain, radius=15, selector='all'):
    """
    Take a Bio.PDB chain object, and find all residues within a radius of a
    given residue. Return a dictionary containing nearby residues for each
    residue in the chain.
    Optional parameter is the atom with which to compute distance. By default
    this is 'all', which gets all non-heterologous atoms. Other potential
    options include 'CA', 'CB' etc. If an atom is not found within a residue
    object, then method reverts to using 'CA'.
    """
    #Small hack to allow method to return central atom alone if radius is set
    #to zero.
    if radius == 0:
        radius = 0.01
    #Setup variables
    ref_dict = {}
    euclidean_distance, reference = _euclidean_distance_matrix(chain, selector)
    within_radius = euclidean_distance < radius
    near_map = within_radius * reference
    #Iterate over all atoms in Euclidean distance matrix.
    for i, atom in enumerate(near_map):
        if atom[i] not in ref_dict:
            ref_dict[atom[i]] = atom[np.nonzero(atom)]
        else:
            ref_dict[atom[i]] = np.append(ref_dict[atom[i]],
                                          atom[np.nonzero(atom)])
    #Remove non-unique values.
    for key in ref_dict:<commit_msg>Change calculation of nearby residues to greater than *or equal to* radius<commit_after>"""Helper module for structmap package. Contains a collection of tools
for analysing a pdb file.
"""
from __future__ import absolute_import, division, print_function

from Bio.SeqIO import PdbIO
from Bio.SeqUtils import seq1, ProtParamData
from Bio.Data.SCOPData import protein_letters_3to1
from scipy.spatial import distance
import numpy as np
from .seqtools import _construct_sub_align, blast_sequences
from . import gentests

ss_lookup_dict = {
    'H': 0,
    'B': 1,
    'E': 2,
    'G': 3,
    'I': 4,
    'T': 5,
    'S': 6,
    '-': 7,
    0: 'H',
    1: 'B',
    2: 'E',
    3: 'G',
    4: 'I',
    5: 'T',
    6: 'S',
    7: '-'
    }

def _euclidean_distance_matrix(chain, selector='all'):
    """Compute the Euclidean distance matrix for all atoms in a pdb chain.
    Return the euclidean distance matrix and a reference list of all atoms
    in the model (positionally matched to the euclidean matrix).
    Optional parameter is the atom in each residue with which to compute a
    distance matrix. The default setting is 'all', which gets all
    non-heterologous atoms. Other potential options include 'CA', 'CB' etc.
    If an atom is not found within a residue object, then method reverts to
    using 'CA'.
    """
    reference = []
    coords = []
    #Filter on non-HET atoms
    for residue in (x for x in chain if x.get_id()[0] == ' '):
        #If selecting based on all atoms within residue
        if selector == 'all':
            for atom in residue:
                coords.append(atom.get_coord())
                reference.append(atom.get_full_id()[3][1])
        #If measuring distance on particular atoms
        else:
            if selector in residue:
                select_atom = selector
            #Revert to carbon alpha if atom is not found
            elif 'CA' in residue:
                select_atom = 'CA'
            #if CA is not found, do not include residue in distance matrix
            else:
                continue
            coords.append(residue[select_atom].get_coord())
            reference.append(residue[select_atom].get_full_id()[3][1])
    #Convert to a np array, and compute Euclidean distance.
    coord_array = np.array(coords)
    euclid_mat = distance.pdist(coord_array, 'euclidean')
    #Convert to squareform matrix
    euclid_mat = distance.squareform(euclid_mat)
    ref_array = np.array(reference)
    return euclid_mat, ref_array

def nearby(chain, radius=15, selector='all'):
    """
    Take a Bio.PDB chain object, and find all residues within a radius of a
    given residue. Return a dictionary containing nearby residues for each
    residue in the chain.
    Optional parameter is the atom with which to compute distance. By default
    this is 'all', which gets all non-heterologous atoms. Other potential
    options include 'CA', 'CB' etc. If an atom is not found within a residue
    object, then method reverts to using 'CA'.
    """
    #Setup variables
    ref_dict = {}
    euclidean_distance, reference = _euclidean_distance_matrix(chain, selector)
    within_radius = euclidean_distance <= radius
    near_map = within_radius * reference
    #Iterate over all atoms in Euclidean distance matrix.
    for i, atom in enumerate(near_map):
        if atom[i] not in ref_dict:
            ref_dict[atom[i]] = atom[np.nonzero(atom)]
        else:
            ref_dict[atom[i]] = np.append(ref_dict[atom[i]],
                                          atom[np.nonzero(atom)])
    #Remove non-unique values.
    for key in ref_dict:
{'fraction_non_alphanumeric': 0.08777472527472527, 'fraction_numerical': 0.007692307692307693, 'mean_word_length': 2.6133995037220843, 'pattern_counts': {'":': 0, '<': 5, '<?xml version=': 0, '>': 3, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 4, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'dir': 'git-commits-cleaned', 'id': '1963813', 'n_tokens_mistral': 2205, 'n_tokens_neox': 2078, 'n_words': 803}