Source code for menten_gcn.wrappers

import numpy as np
import math

    from pyrosetta import rosetta
except BaseException:
    rosetta = None

    import mdtraj as md
except BaseException:
    md = None

    import Bio as bp
    bp = None

def estimate_CB_xyz(C_xyz, N_xyz, CA_xyz):
    # ICOOR_INTERNAL   CB  -122.800000   69.625412    1.521736   CA    N   C
    That is, CB is found 1.52 Ang from CA, at an angle of 70 degrees from the CA-N line, and a dihedral of -123 degrees for CB-CA-N-C.

    # Credit to Rohit Bhattacharya from
    Nerf method of finding 4th coord (d)
    in cartesian space
    a, b, c : coords of 3 points
    l : bond length between c and d
    theta : bond angle between b, c, d (in degrees)
    chi : dihedral using a, b, c, d (in degrees)
    d : tuple of (x, y, z) in cartesian space

    length = 1.521736
    theta = 69.625412
    chi = -122.800000

    a = C_xyz
    b = N_xyz
    c = CA_xyz

    # calculate unit vectors AB and BC
    ab_unit = (b - a) / np.linalg.norm(b - a)
    bc_unit = (c - b) / np.linalg.norm(c - b)

    # calculate unit normals n = AB x BC
    # and p = n x BC
    n_unit = np.cross(ab_unit, bc_unit)
    n_unit = n_unit / np.linalg.norm(n_unit)
    p_unit = np.cross(n_unit, bc_unit)

    # create rotation matrix [BC; p; n] (3x3)
    M = np.array([bc_unit, p_unit, n_unit]).T

    # convert degrees to radians
    theta = np.pi / 180 * theta
    chi = np.pi / 180 * chi

    # calculate coord pre rotation matrix
    d2 = [-length * np.cos(theta), length * np.sin(theta) * np.cos(chi), length * np.sin(theta) * np.sin(chi)]

    # calculate with rotation as our final output
    return c +, d2)

[docs]class WrappedPose: """ This is the base class for all pose representations. The internal Menten GCN code will use API listed here """ def __init__(self, designable_resids=None): self.CB_estimates = None self.legal_nbrs = None self.designable_resids = designable_resids def get_legal_nbrs(self): return self.legal_nbrs def get_atom_xyz(self, resid, atomname): raise NotImplementedError def get_phi_psi(self, resid): # radians raise NotImplementedError def get_chi(self, resid, chi_number): # radians raise NotImplementedError def get_name1(self, resid): raise NotImplementedError def residues_are_polymer_bonded(self, resid1, resid2): raise NotImplementedError def n_residues(self): raise NotImplementedError def resid_is_N_term(self, resid): raise NotImplementedError def resid_is_C_term(self, resid): raise NotImplementedError def resids_are_same_chain(self, resid1, resid2): raise NotImplementedError def set_designable_resids(self, resids): self.designable_resids = resids def resid_is_designable(self, resid): assert self.designable_resids is not None return resid in self.designable_resids def approximate_ALA_CB(self, resid): assert hasattr(self, 'CB_estimates') # if not hasattr( self, 'CB_estimates' ): if self.CB_estimates is None: # Lazy initialization self.CB_estimates = [None for i in range(0, self.n_residues() + 1)] elif self.CB_estimates[resid] is not None: return self.CB_estimates[resid] get_xyz = self.get_atom_xyz self.CB_estimates[resid] = estimate_CB_xyz(get_xyz(resid, "C"), get_xyz(resid, "N"), get_xyz(resid, "CA")) return self.CB_estimates[resid]
[docs]class RosettaPoseWrapper(WrappedPose): """ This wrapper takes a rosetta pose and requires pyrosetta to be installed Parameters --------- pose: Pose Rosetta pose """ def __init__(self, pose): WrappedPose.__init__(self) if rosetta is None: print("RosettaPoseWrapper requires the pyrosetta library to be installed") raise ImportError assert isinstance(pose, rosetta.core.pose.Pose) self.pose = pose def get_atom_xyz(self, resid, atomname): xyz = self.pose.residue(resid).xyz(atomname) return np.asarray([xyz.x, xyz.y, xyz.z]) def get_phi_psi(self, resid): phipsi = np.asarray([self.pose.phi(resid), self.pose.psi(resid)]) phipsi[0] = math.radians(phipsi[0]) phipsi[1] = math.radians(phipsi[1]) return phipsi def get_chi(self, resid, chi_number): if self.pose.residue(resid).nchi() < chi_number: return 0, False chi_deg = self.pose.chi(chi_number, resid) chi_rad = math.radians(chi_deg) return chi_rad, True def get_name1(self, resid): return self.pose.residue(resid).name1() def residues_are_polymer_bonded(self, resid1, resid2): return self.pose.residue(resid1).is_polymer_bonded(resid2) def n_residues(self): return self.pose.size() def approximate_ALA_CB_via_mutation(self, resid): if not self.pose.residue(resid).name1() == 'G': print("RosettaPoseWrapper.approximate_ALA_CB is only setup for glycine right now") print(self.pose.residue(resid).name1()) assert False mutator = rosetta.protocols.simple_moves.MutateResidue(resid, 'ALA') mutator.apply(self.pose) xyz = self.get_atom_xyz(resid, "CB") mutator = rosetta.protocols.simple_moves.MutateResidue(resid, 'GLY') mutator.apply(self.pose) return xyz def resid_is_N_term(self, resid): return self.pose.residue(resid).is_lower_terminus() def resid_is_C_term(self, resid): return self.pose.residue(resid).is_upper_terminus() def resids_are_same_chain(self, resid1, resid2): return self.pose.chain(resid1) == self.pose.chain(resid2)
[docs]class MDTrajPoseWrapper(WrappedPose): """ This wrapper takes a MDTraj trajectory and requires MDTraj to be installed Parameters --------- mdtraj_trajectory: Trajectory Pose in MDTraj trajectory format """ def __init__(self, mdtraj_trajectory): WrappedPose.__init__(self) if md is None: print("MDTrajPoseWrapper requires the mdtraj library to be installed") raise ImportError assert isinstance(mdtraj_trajectory, md.Trajectory) assert mdtraj_trajectory.n_frames == 1 self.trajectory = mdtraj_trajectory # RADIANS: self.phi_atoms, self.phis = md.compute_phi(self.trajectory) self.psi_atoms, self.psis = md.compute_psi(self.trajectory) self.chis = [None, None, None, None, None] # Adding zero element just to make indexing easier self.chi_atoms = [None, None, None, None, None] self.chi_atoms[1], self.chis[1] = md.compute_chi1(self.trajectory) self.chi_atoms[2], self.chis[2] = md.compute_chi2(self.trajectory) self.chi_atoms[3], self.chis[3] = md.compute_chi3(self.trajectory) self.chi_atoms[4], self.chis[4] = md.compute_chi4(self.trajectory) def get_atom_xyz(self, resid, atomname): atom = self.trajectory.topology.residue(resid - 1).atom(atomname) all_xyzs = return all_xyzs[0][atom.index] * 10 # nm -> Å def _get_phi_or_psi_angle(self, resid, value_vec, atom_vec): assert len(value_vec[0]) == len(atom_vec) # value_vec.shape: (n_frames, n_phi) # value_vec.shape: (n_phi, 4) # Okay this is pretty inefficient top = self.trajectory.topology for i in range(0, len(value_vec[0])): atom_index = atom_vec[i][2] # Last-Middle atom atom = top.atom(atom_index) if atom.residue.index == resid - 1: return value_vec[0][i] elif atom.residue.index > resid - 1: return 0 return 0 def get_phi_psi(self, resid): phi_rad = self._get_phi_or_psi_angle(resid, self.phis, self.phi_atoms) psi_rad = self._get_phi_or_psi_angle(resid, self.psis, self.psi_atoms) return [phi_rad, psi_rad] def get_chi(self, resid, chi_number): # DOESN'T GIVE PROTON CHIs assert chi_number > 0 assert chi_number <= 4 # Okay this is pretty inefficient top = self.trajectory.topology for i in range(0, len(self.chi_atoms[chi_number])): atom_index = self.chi_atoms[chi_number][i][3] # Last atom atom = top.atom(atom_index) if atom.residue.index == resid - 1: return self.chis[chi_number][0][i], True elif atom.residue.index > resid - 1: return 0, False return 0, False def get_name1(self, resid): return self.trajectory.topology.residue(resid - 1).code def residues_are_polymer_bonded(self, resid1, resid2): if not self.resids_are_same_chain(resid1, resid2): return False if abs(resid1 - resid2) > 1: return False first_res = min(resid1, resid2) second_res = max(resid1, resid2) top = self.trajectory.topology C_atom_index = top.residue(first_res - 1).atom("C") N_atom_index = top.residue(second_res - 1).atom("N") for a1, a2 in top.bonds: if (a1 == C_atom_index and a2 == N_atom_index) or (a2 == C_atom_index and a1 == N_atom_index): return True return False def n_residues(self): return self.trajectory.topology.n_residues def resid_is_N_term(self, resid): top = self.trajectory.topology chain = top.residue(resid - 1).chain N_term_resid = chain.residue(0) return N_term_resid == resid - 1 def resid_is_C_term(self, resid): top = self.trajectory.topology chain = top.residue(resid - 1).chain C_term_resid = chain.residue(chain.n_residues - 1) return C_term_resid == resid - 1 def resids_are_same_chain(self, resid1, resid2): top = self.trajectory.topology return top.residue(resid1 - 1).chain.index == top.residue(resid2 - 1).chain.index