Source code for menten_gcn.decorators.geometry

import numpy as np
import math

from menten_gcn.decorators.base import *


def angle_rad(a, b, c):
    # https://stackoverflow.com/questions/35176451/python-code-to-calculate-angle-between-three-point-using-their-3d-coordinates
    ba = a - b
    bc = c - b
    cosine_angle = np.dot(ba, bc) / (np.linalg.norm(ba) * np.linalg.norm(bc))
    angle = np.arccos(cosine_angle)
    return angle


def dihedral_rad(p0, p1, p2, p3):
    # https://stackoverflow.com/questions/20305272/dihedral-torsion-angle-from-four-points-in-cartesian-coordinates-in-python
    # Praxeolitic formula 1 sqrt, 1 cross product
    b0 = -1.0 * (p1 - p0)
    b1 = p2 - p1
    b2 = p3 - p2

    # normalize b1 so that it does not influence magnitude of vector
    # rejections that come next
    b1 /= np.linalg.norm(b1)

    # vector rejections
    # v = projection of b0 onto plane perpendicular to b1
    #   = b0 minus component that aligns with b1
    # w = projection of b2 onto plane perpendicular to b1
    #   = b2 minus component that aligns with b1
    v = b0 - np.dot(b0, b1) * b1
    w = b2 - np.dot(b2, b1) * b1

    # angle between v and w in a plane is the torsion angle
    # v and w may not be normalized but that's fine since tan is y/x
    x = np.dot(v, w)
    y = np.dot(np.cross(b1, v), w)
    # return np.degrees(np.arctan2(y, x))
    return np.arctan2(y, x)


[docs]class CACA_dist(Decorator): """ Measures distance between the two C-Alpha atoms of each residue - 0 Node Features - 1 Edge Feature Parameters --------- use_nm: bool If true (default), measure distance in Angstroms. Otherwise use nanometers. """ def __init__(self, use_nm: bool = False): self.use_nm = use_nm
[docs] def get_version_name(self): return "CACA_dist"
[docs] def n_edge_features(self): return 1
[docs] def calc_edge_features(self, wrapped_pose, resid1: int, resid2: int, dict_cache=None): xyz1 = wrapped_pose.get_atom_xyz(resid1, "CA") xyz2 = wrapped_pose.get_atom_xyz(resid2, "CA") distance = [np.linalg.norm(xyz1 - xyz2)] if self.use_nm: distance[0] = distance[0] / 10.0 # distance is symmetric, so return it twice return distance, distance
[docs] def describe_edge_features(self): if self.use_nm: d_units = "nanometers" else: d_units = "Angstroms" return [ "Euclidean distance between the CA atoms of each residue, measured in " + d_units, ]
[docs]class CBCB_dist(Decorator): """ Measures distance between the two C-Beta atoms of each residue. Note: We will calculate the "ideal ALA" CB location even if this residue has a CB atom. This may sound silly but it is intended to prevents noise from different native amino acid types. - 0 Node Features - 1 Edge Feature Parameters --------- use_nm: bool If true (default), measure distance in Angstroms. Otherwise use nanometers. """ def __init__(self, use_nm: bool = False): self.use_nm = use_nm
[docs] def get_version_name(self): return "CBCB_dist"
[docs] def n_edge_features(self): return 1
[docs] def calc_edge_features(self, wrapped_pose, resid1: int, resid2: int, dict_cache=None): CB1_xyz = wrapped_pose.approximate_ALA_CB(resid1) CB2_xyz = wrapped_pose.approximate_ALA_CB(resid2) distance = [np.linalg.norm(CB1_xyz - CB2_xyz)] if self.use_nm: distance[0] = distance[0] / 10.0 return distance, distance
[docs] def describe_edge_features(self): if self.use_nm: d_units = "nanometers" else: d_units = "Angstroms" return [ "Euclidean distance between the CB atoms of each residue, measured in " + d_units + ". In the case of GLY, use an estimate of ALA's CB position", ]
[docs]class PhiPsiRadians(Decorator): """ Returns the phi and psi values of each residue position. - 2-4 Node Features - 0 Edge Features Parameters --------- sincos: bool Return the sine and cosine of phi and psi instead of just the raw values. """ def __init__(self, sincos: bool = False): self.sincos = sincos
[docs] def get_version_name(self): return "PhiPsiRadians"
[docs] def n_node_features(self): if self.sincos: return 4 else: return 2
[docs] def calc_node_features(self, wrapped_pose, resid, dict_cache=None): if self.sincos: phi_rad, psi_rad = wrapped_pose.get_phi_psi(resid) vec = np.zeros(shape=(4)) if not wrapped_pose.resid_is_N_term(resid): vec[0] = math.sin(phi_rad) vec[1] = math.cos(phi_rad) if not wrapped_pose.resid_is_C_term(resid): vec[2] = math.sin(psi_rad) vec[3] = math.cos(psi_rad) return vec else: return wrapped_pose.get_phi_psi(resid)
[docs] def describe_node_features(self): if self.sincos: return [ "Sine of Phi of the given residue, 0 if N-term. Ranges from -1 to 1", "Cosine of Phi of the given residue, 0 if N-term. Ranges from -1 to 1", "Sine of Psi of the given residue, 0 if C-term. Ranges from -1 to 1", "Cosine of Psi of the given residue, 0 if C-term. Ranges from -1 to 1", ] else: return [ "Phi of the given residue, measured in radians. Spans from -pi to pi", "Psi of the given residue, measured in radians. Spans from -pi to pi", ]
[docs]class trRosettaEdges(Decorator): """ Use the residue pair geometries used in this paper: https://www.pnas.org/content/117/3/1496/tab-figures-data - 0 Node Features - 4-7 Edge Features Parameters --------- sincos: bool Return the sine and cosine of phi and psi instead of just the raw values. use_nm: bool If true, measure distance in Angstroms. Otherwise use nanometers. Note: This default value does not match the default of other decorators. This is for the sake of matching the trRosetta paper. """ def __init__(self, sincos: bool = False, use_nm: bool = False): self.sincos = sincos self.use_nm = use_nm def get_version_name(self): return "trRosettaEdges" def n_node_features(self): return 0 def n_edge_features(self): if self.sincos: return 7 else: return 4 def calc_edge_features(self, wrapped_pose, resid1, resid2, dict_cache=None): # See Fig1 of https://www.biorxiv.org/content/10.1101/846279v1.full.pdf CA1_xyz = wrapped_pose.get_atom_xyz(resid1, "CA") N1_xyz = wrapped_pose.get_atom_xyz(resid1, "N") CB1_xyz = wrapped_pose.approximate_ALA_CB(resid1) CA2_xyz = wrapped_pose.get_atom_xyz(resid2, "CA") N2_xyz = wrapped_pose.get_atom_xyz(resid2, "N") CB2_xyz = wrapped_pose.approximate_ALA_CB(resid2) CB_distance = np.linalg.norm(CB1_xyz - CB2_xyz) if self.use_nm: CB_distance = CB_distance / 10.0 omega_rad = dihedral_rad(CA1_xyz, CB1_xyz, CB2_xyz, CA2_xyz) theta_12 = dihedral_rad(N1_xyz, CA1_xyz, CB1_xyz, CB2_xyz) theta_21 = dihedral_rad(N2_xyz, CA2_xyz, CB2_xyz, CB1_xyz) phi_12 = angle_rad(CA1_xyz, CB1_xyz, CB2_xyz) phi_21 = angle_rad(CA2_xyz, CB2_xyz, CB1_xyz) f_12 = [CB_distance, omega_rad, theta_12, phi_12] f_21 = [CB_distance, omega_rad, theta_21, phi_21] if self.sincos: return self.convert_edges_to_sincos(f_12, f_21) else: return f_12, f_21 def convert_edges_to_sincos(self, f_12, f_21): f_12_prime = np.zeros(shape=(7)) f_12_prime[0] = f_12[0] f_12_prime[1] = math.sin(f_12[1]) f_12_prime[2] = math.cos(f_12[1]) f_12_prime[3] = math.sin(f_12[2]) f_12_prime[4] = math.cos(f_12[2]) f_12_prime[5] = math.sin(f_12[3]) f_12_prime[6] = math.cos(f_12[3]) f_21_prime = np.zeros(shape=(7)) f_21_prime[0] = f_21[0] f_21_prime[1] = math.sin(f_21[1]) f_21_prime[2] = math.cos(f_21[1]) f_21_prime[3] = math.sin(f_21[2]) f_21_prime[4] = math.cos(f_21[2]) f_21_prime[5] = math.sin(f_21[3]) f_21_prime[6] = math.cos(f_21[3]) np.testing.assert_almost_equal(f_12_prime[0:2], f_21_prime[0:2], decimal=4) return f_12_prime, f_21_prime def describe_edge_features(self): if self.use_nm: d_units = "nanometers" else: d_units = "Angstroms" if self.sincos: return [ "Euclidean distance between the two CB atoms of each residue, measured in " + d_units + " (symmetric)", "Sine of CA-CB-CB-CA torsion angle, spans from -1 to 1 (symmetric)", "Cosine of CA-CB-CB-CA torsion angle, spans from -1 to 1 (symmetric)", "Sine of N1-CA1-CB1-CB2 torsion angle, spans from -1 to 1 (asymmetric)", "Cosine of N1-CA1-CB1-CB2 torsion angle, spans from -1 to 1 (asymmetric)", "Sine of CA1-CB1-CB2 bond angle, spans from 0 to 1 (asymmetric)", "Cosine of CA1-CB1-CB2 bond angle, spans from -1 to 1 (asymmetric)", ] else: return [ "Euclidean distance between the two CB atoms of each residue, measured in " + d_units + " (symmetric)", "CA-CB-CB-CA torsion angle in radians, spans from -pi to pi (symmetric)", "N1-CA1-CB1-CB2 torsion angle in radians, spans from -pi to pi (asymmetric)", "CA1-CB1-CB2 bond angle in radians, spans from 0 to pi (asymmetric)", ]
class _SimpleBBGeometry_v0(CombinedDecorator): def __init__(self, use_nm=False): decorators = [PhiPsiRadians(sincos=False), CBCB_dist(use_nm=use_nm)] CombinedDecorator.__init__(self, decorators) def get_version_name(self): return "_SimpleBBGeometry_v0" class _StandardBBGeometry_v0(CombinedDecorator): def __init__(self, use_nm=False): decorators = [PhiPsiRadians(sincos=True), trRosettaEdges(sincos=False, use_nm=use_nm)] CombinedDecorator.__init__(self, decorators) def get_version_name(self): return "_StandardBBGeometry_v0" class _AdvancedBBGeometry_v0(CombinedDecorator): def __init__(self, use_nm=False): decorators = [PhiPsiRadians(sincos=True), CACA_dist(use_nm=use_nm), trRosettaEdges(sincos=True, use_nm=use_nm)] CombinedDecorator.__init__(self, decorators) def get_version_name(self): return "_AdvancedBBGeometry_v0"
[docs]class ChiAngleDecorator(Decorator): """ Returns the chi values of each residue position. Ranges from -pi to pi or -1 to 1 if sincos=True. WARNING: This can behave inconsistantly for proton chis accross modeling frameworks. Rosetta adds hydrogens when they are absent from the input file but MDtraj does not. This results in Rosetta calculating a chi value in some cases that MDtraj skips! - 0-8 Node Features - 0 Edge Features Parameters --------- chi1: bool Include chi1's value chi2: bool Include chi2's value chi3: bool Include chi3's value chi4: bool Include chi4's value sincos: bool Return the sine and cosine of chi instead of just the raw values """ def __init__(self, chi1: bool = True, chi2: bool = True, chi3: bool = True, chi4: bool = True, sincos: bool = True): self.chis = [] if chi1: self.chis.append(1) if chi2: self.chis.append(2) if chi3: self.chis.append(3) if chi4: self.chis.append(4) self.sincos = sincos print(("ChiAngleDecorator: Warning, different protein representations (i.e., Rosetta vs MDTraj)" " represent proton chis differently if the hydrogen atoms are missing from the input file."))
[docs] def get_version_name(self): return "ChiAngleDecorator"
[docs] def n_edge_features(self): return 0
[docs] def n_node_features(self): if self.sincos: return len(self.chis) * 2 else: return len(self.chis)
[docs] def calc_node_features(self, wrapped_pose, resid, dict_cache=None): f = [] for chi in self.chis: chi_rad, is_valid = wrapped_pose.get_chi(resid, chi) #print( chi_rad ) if self.sincos: if is_valid: f.append(math.sin(chi_rad)) f.append(math.cos(chi_rad)) else: f.append(0) f.append(0) else: if is_valid: f.append(chi_rad) else: f.append(-5) return f
[docs] def describe_node_features(self): desc = [] for chi in self.chis: if self.sincos: desc.append( "Sine of chi angle number " + str(chi) + " of each residue. Spans from -1 to 1, 0 if chi angle is not valid for this residue") desc.append( "Cosine of chi angle number " + str(chi) + " of each residue. Spans from -1 to 1, 0 if chi angle is not valid for this residue") else: desc.append( "Chi angle number " + str(chi) + " of each residue measured in radians. Spans from -pi to pi, -5 if chi angle is not valid for this residue") return desc
[docs]class SimpleBBGeometry(_SimpleBBGeometry_v0): """ Meta-decorator that combines PhiPsiRadians(sincos=False) and CBCB_dist - 2 Node Features - 1 Edge Feature Parameters --------- use_nm: bool If true, measure distance in Angstroms. Otherwise use nanometers. """ pass
[docs]class StandardBBGeometry(_StandardBBGeometry_v0): """ Meta-decorator that combines PhiPsiRadians(sincos=True) and trRosettaEdges(sincos=False) - 4 Node Features - 4 Edge Features Parameters --------- use_nm: bool If true, measure distance in Angstroms. Otherwise use nanometers. """ pass
[docs]class AdvancedBBGeometry(_AdvancedBBGeometry_v0): """ Meta-decorator that combines PhiPsiRadians(sincos=True), CACA_dist, and trRosettaEdges(sincos=True) - 4 Node Features - 8 Edge Features Parameters --------- use_nm: bool If true, measure all distances in Angstroms. Otherwise use nanometers. """ pass