hydrogen bond code

hb.py

import numpy as np
from msmbuilder.metrics import AbstractDistanceMetric, Vectorized
# multithreaded distance and angle engines
from msmbuilder.geometry import angle, contact
import itertools

class HydrogenBond(Vectorized, AbstractDistanceMetric):
    """ 
    Distance metric class for calculating distance between frames based
    on their hydrogen bond contact maps, which are NON-SYMMETRIC matrices
    corresponding to whether donor i is hydrogen bonding with acceptor j.
    """

    def __init__( self, metric='matching', radius_cutoff=0.32, angle_cutoff=2*np.pi/3):
        """ Create a distance metric that will measure the distance between frames
        based on differences in their hydrogen bonding patterns"""
        # RTM: I think using the Jaccard distance is more justified here that the matching distance

        super(HydrogenBond, self).__init__(metric)
        self.radius_cutoff = float(radius_cutoff)
        self.angle_cutoff = float(angle_cutoff)

    def prepare_trajectory( self, traj ):
        """ For each (donor, acceptor) pair, calculate a boolean value indicating that the distance
        from donor to acceptor is small enough, and that the donor-H-angle distance is large enough
        
        Returns
        -------
        prepared_traj : np.ndarray, dtype=np.bool, shape=(n_frames, n_donors*n_acceptors)
            prepared_traj[i] gives a list of booleans indicating the presence or absense of each of the possible hbonds
        """
        # Construct the list of donors: i.e. the N-H (amide H's)
        # _ainds correspond to atom indices in the trajectory
    
        # this logic for atom selection is extremely fragile, because it relies on the atom names and residue names, which
        # are not always correct and vary from MD package to MD package. Be careful.
        donorH_ainds = np.where( ( (traj['AtomNames'] == 'H')|(traj['AtomNames'] =='HN' ) )&(traj['ResidueNames']!='PRO') )[0]
        # Remove the n-terminus for now, since the naming is screwy
        donor_ainds = np.where( (traj['AtomNames'] == 'N')&(traj['ResidueNames']!='PRO') )[0][1:]
        # This excludes the C-Terminus which has naming issues...
        acceptor_ainds = np.where( traj['AtomNames'] == 'O' )[0]
        
        n_donors = len(donorH_ainds)
        n_acceptors = len(acceptor_ainds)
        
        #print donorH_ainds.shape, donor_ainds.shape
        #print donorH_ainds, donor_ainds
        if len( donor_ainds ) != n_donors:
            raise Exception('Number of donor atoms and donor hydrogens are not the same...')
        
        donor_acceptor_pairs = np.array(list(itertools.product(donorH_ainds, acceptor_ainds)), dtype=np.int32)
        product = itertools.product(range(n_donors), range(n_acceptors))
        donor_H_acceptor_tripples = np.array([(donor_ainds[i], donorH_ainds[i], acceptor_ainds[j]) for (i,j) in product], dtype=np.int32)
        #print "Calculating H-Bonds for %d combinations" % len( donor_acceptor_pairs )

        self.last_donor_ainds = donor_ainds
        self.last_donorH_ainds = donorH_ainds
        self.last_acceptor_ainds = acceptor_ainds

        distances_satisfy = contact.atom_distances(traj['XYZList'], donor_acceptor_pairs) <= self.radius_cutoff
        angle_satisfy = angle.bond_angles(traj['XYZList'], donor_H_acceptor_tripples) >= self.angle_cutoff

        return np.logical_and(distances_satisfy, angle_satisfy)