Spaces:
Runtime error
Runtime error
| #@title get secondary structure (SSE) from given PDB file | |
| #@markdown So far it seems the best solution is to steal code from biotite | |
| #@markdown which calculates the SSE of a peptide chain based on the P-SEA algorithm (Labesse 1997) | |
| # CODE FROM BIOKITE | |
| # From Krypton | |
| import numpy as np | |
| import random | |
| import torch | |
| def vector_dot(v1,v2): | |
| return (v1*v2).sum(axis=-1) | |
| def norm_vector(v): | |
| factor = np.linalg.norm(v, axis=-1) | |
| if isinstance(factor, np.ndarray): | |
| v /= factor[..., np.newaxis] | |
| else: | |
| v /= factor | |
| return v | |
| def coord(x): | |
| return np.asarray(x) | |
| def displacement(atoms1, atoms2): | |
| v1 = coord(atoms1) | |
| v2 = coord(atoms2) | |
| if len(v1.shape) <= len(v2.shape): | |
| diff = v2 - v1 | |
| else: | |
| diff = -(v1 - v2) | |
| return diff | |
| def distance(atoms1, atoms2): | |
| diff = displacement(atoms1, atoms2) | |
| return np.sqrt(vector_dot(diff, diff)) | |
| def angle(atoms1, atoms2, atoms3): | |
| v1 = displacement(atoms1, atoms2) | |
| v2 = displacement(atoms3, atoms2) | |
| norm_vector(v1) | |
| norm_vector(v2) | |
| return np.arccos(vector_dot(v1,v2)) | |
| def dihedral(atoms1, atoms2, atoms3, atoms4): | |
| v1 = displacement(atoms1, atoms2) | |
| v2 = displacement(atoms2, atoms3) | |
| v3 = displacement(atoms3, atoms4) | |
| norm_vector(v1) | |
| norm_vector(v2) | |
| norm_vector(v3) | |
| n1 = np.cross(v1, v2) | |
| n2 = np.cross(v2, v3) | |
| # Calculation using atan2, to ensure the correct sign of the angle | |
| x = vector_dot(n1,n2) | |
| y = vector_dot(np.cross(n1,n2), v2) | |
| return np.arctan2(y,x) | |
| def replace_letters(arr): | |
| # Create a dictionary that maps the letters 'a', 'b', and 'c' to the corresponding numbers | |
| letter_to_number = {'a': 0, 'b': 1, 'c': 2} | |
| # Create a new array that will hold the numbers | |
| nums = [] | |
| # Loop through the input array and replace the letters with the corresponding numbers | |
| for letter in arr: | |
| if letter in letter_to_number: | |
| nums.append(letter_to_number[letter]) | |
| else: | |
| nums.append(letter) | |
| return np.array(nums) | |
| def replace_with_mask(arr, percentage, replace_loops=False): | |
| # Make sure the percentage is between 0 and 100 | |
| percentage = min(max(percentage, 0), 100) | |
| # Calculate the number of values to replace | |
| num_to_replace = int(len(arr) * percentage / 100) | |
| # Choose a random subset of the array to replace | |
| replace_indices = random.sample(range(len(arr)), num_to_replace) | |
| # Replace the values at the chosen indices with the number 3 | |
| for i in replace_indices: | |
| arr[i] = 3 | |
| if replace_loops: | |
| for i in arr: | |
| if arr[i] == 2: | |
| arr[i] = 3 | |
| return arr | |
| def annotate_sse(ca_coord, percentage_mask=0, replace_loops=False): | |
| _radians_to_angle = 2*np.pi/360 | |
| _r_helix = ((89-12)*_radians_to_angle, (89+12)*_radians_to_angle) | |
| _a_helix = ((50-20)*_radians_to_angle, (50+20)*_radians_to_angle) | |
| _d2_helix = ((5.5-0.5), (5.5+0.5)) | |
| _d3_helix = ((5.3-0.5), (5.3+0.5)) | |
| _d4_helix = ((6.4-0.6), (6.4+0.6)) | |
| _r_strand = ((124-14)*_radians_to_angle, (124+14)*_radians_to_angle) | |
| _a_strand = ((-180)*_radians_to_angle, (-125)*_radians_to_angle, | |
| (145)*_radians_to_angle, (180)*_radians_to_angle) | |
| _d2_strand = ((6.7-0.6), (6.7+0.6)) | |
| _d3_strand = ((9.9-0.9), (9.9+0.9)) | |
| _d4_strand = ((12.4-1.1), (12.4+1.1)) | |
| # Filter all CA atoms in the relevant chain. | |
| d2i_coord = np.full(( len(ca_coord), 2, 3 ), np.nan) | |
| d3i_coord = np.full(( len(ca_coord), 2, 3 ), np.nan) | |
| d4i_coord = np.full(( len(ca_coord), 2, 3 ), np.nan) | |
| ri_coord = np.full(( len(ca_coord), 3, 3 ), np.nan) | |
| ai_coord = np.full(( len(ca_coord), 4, 3 ), np.nan) | |
| # The distances and angles are not defined for the entire interval, | |
| # therefore the indices do not have the full range | |
| # Values that are not defined are NaN | |
| for i in range(1, len(ca_coord)-1): | |
| d2i_coord[i] = (ca_coord[i-1], ca_coord[i+1]) | |
| for i in range(1, len(ca_coord)-2): | |
| d3i_coord[i] = (ca_coord[i-1], ca_coord[i+2]) | |
| for i in range(1, len(ca_coord)-3): | |
| d4i_coord[i] = (ca_coord[i-1], ca_coord[i+3]) | |
| for i in range(1, len(ca_coord)-1): | |
| ri_coord[i] = (ca_coord[i-1], ca_coord[i], ca_coord[i+1]) | |
| for i in range(1, len(ca_coord)-2): | |
| ai_coord[i] = (ca_coord[i-1], ca_coord[i], | |
| ca_coord[i+1], ca_coord[i+2]) | |
| d2i = distance(d2i_coord[:,0], d2i_coord[:,1]) | |
| d3i = distance(d3i_coord[:,0], d3i_coord[:,1]) | |
| d4i = distance(d4i_coord[:,0], d4i_coord[:,1]) | |
| ri = angle(ri_coord[:,0], ri_coord[:,1], ri_coord[:,2]) | |
| ai = dihedral(ai_coord[:,0], ai_coord[:,1], | |
| ai_coord[:,2], ai_coord[:,3]) | |
| sse = np.full(len(ca_coord), "c", dtype="U1") | |
| # Annotate helices | |
| # Find CA that meet criteria for potential helices | |
| is_pot_helix = np.zeros(len(sse), dtype=bool) | |
| for i in range(len(sse)): | |
| if ( | |
| d3i[i] >= _d3_helix[0] and d3i[i] <= _d3_helix[1] | |
| and d4i[i] >= _d4_helix[0] and d4i[i] <= _d4_helix[1] | |
| ) or ( | |
| ri[i] >= _r_helix[0] and ri[i] <= _r_helix[1] | |
| and ai[i] >= _a_helix[0] and ai[i] <= _a_helix[1] | |
| ): | |
| is_pot_helix[i] = True | |
| # Real helices are 5 consecutive helix elements | |
| is_helix = np.zeros(len(sse), dtype=bool) | |
| counter = 0 | |
| for i in range(len(sse)): | |
| if is_pot_helix[i]: | |
| counter += 1 | |
| else: | |
| if counter >= 5: | |
| is_helix[i-counter : i] = True | |
| counter = 0 | |
| # Extend the helices by one at each end if CA meets extension criteria | |
| i = 0 | |
| while i < len(sse): | |
| if is_helix[i]: | |
| sse[i] = "a" | |
| if ( | |
| d3i[i-1] >= _d3_helix[0] and d3i[i-1] <= _d3_helix[1] | |
| ) or ( | |
| ri[i-1] >= _r_helix[0] and ri[i-1] <= _r_helix[1] | |
| ): | |
| sse[i-1] = "a" | |
| sse[i] = "a" | |
| if ( | |
| d3i[i+1] >= _d3_helix[0] and d3i[i+1] <= _d3_helix[1] | |
| ) or ( | |
| ri[i+1] >= _r_helix[0] and ri[i+1] <= _r_helix[1] | |
| ): | |
| sse[i+1] = "a" | |
| i += 1 | |
| # Annotate sheets | |
| # Find CA that meet criteria for potential strands | |
| is_pot_strand = np.zeros(len(sse), dtype=bool) | |
| for i in range(len(sse)): | |
| if ( d2i[i] >= _d2_strand[0] and d2i[i] <= _d2_strand[1] | |
| and d3i[i] >= _d3_strand[0] and d3i[i] <= _d3_strand[1] | |
| and d4i[i] >= _d4_strand[0] and d4i[i] <= _d4_strand[1] | |
| ) or ( | |
| ri[i] >= _r_strand[0] and ri[i] <= _r_strand[1] | |
| and ( (ai[i] >= _a_strand[0] and ai[i] <= _a_strand[1]) | |
| or (ai[i] >= _a_strand[2] and ai[i] <= _a_strand[3])) | |
| ): | |
| is_pot_strand[i] = True | |
| # Real strands are 5 consecutive strand elements, | |
| # or shorter fragments of at least 3 consecutive strand residues, | |
| # if they are in hydrogen bond proximity to 5 other residues | |
| pot_strand_coord = ca_coord[is_pot_strand] | |
| is_strand = np.zeros(len(sse), dtype=bool) | |
| counter = 0 | |
| contacts = 0 | |
| for i in range(len(sse)): | |
| if is_pot_strand[i]: | |
| counter += 1 | |
| coord = ca_coord[i] | |
| for strand_coord in ca_coord: | |
| dist = distance(coord, strand_coord) | |
| if dist >= 4.2 and dist <= 5.2: | |
| contacts += 1 | |
| else: | |
| if counter >= 4: | |
| is_strand[i-counter : i] = True | |
| elif counter == 3 and contacts >= 5: | |
| is_strand[i-counter : i] = True | |
| counter = 0 | |
| contacts = 0 | |
| # Extend the strands by one at each end if CA meets extension criteria | |
| i = 0 | |
| while i < len(sse): | |
| if is_strand[i]: | |
| sse[i] = "b" | |
| if d3i[i-1] >= _d3_strand[0] and d3i[i-1] <= _d3_strand[1]: | |
| sse[i-1] = "b" | |
| sse[i] = "b" | |
| if d3i[i+1] >= _d3_strand[0] and d3i[i+1] <= _d3_strand[1]: | |
| sse[i+1] = "b" | |
| i += 1 | |
| sse=replace_letters(sse) | |
| sse=replace_with_mask(sse, percentage_mask, replace_loops=replace_loops) | |
| sse=torch.nn.functional.one_hot(torch.tensor(sse), num_classes=4) | |
| return sse | |