Python mdtraj.compute_contacts函数代码示例

def shukla_coords(trajectories,KER,Aloop,SRC2):

    difference = []
    rmsd = []

    for traj in trajectories:

        # append difference
        k295e310 = md.compute_contacts(traj, [KER[0]])
        e310r409 = md.compute_contacts(traj, [KER[1]])
        difference.append(10*(e310r409[0] - k295e310[0])) # 10x because mdtraj is naturally in nm

        # append rmsd
        Activation_Loop_SRC2 = SRC2.top.select("backbone and (resid %s to %s)" %(140,160))
        Activation_Loop_kinase = traj.top.select("backbone and (resid %s to %s)" %(Aloop[0],Aloop[1]))

        SRC2_cut = SRC2.atom_slice(Activation_Loop_SRC2)
        traj_cut = traj.atom_slice(Activation_Loop_kinase)

        rmsd.append(10*(md.rmsd(traj_cut,SRC2_cut,frame=0))) # 10x because mdtraj is naturaly in nm

    # flatten list of arrays
    flattened_difference = np.asarray([val for sublist in difference for val in sublist])
    flattened_rmsd = np.asarray([val for sublist in rmsd for val in sublist])

    return [flattened_rmsd, flattened_difference]

def test_trek():
    # setup
    with open("processed/p9761/24/7/info.json") as f:
        info = json.load(f)

    info = trajprocess.postprocess.stp(info, 'trek')

    # check stp cleanup
    assert not os.path.exists('{workdir}/stp/0/'.format(**info['path']))

    # check stp results
    traj = mdtraj.load(info['stp']['gens'][0], top=info['stp']['outtop'])
    assert traj.n_atoms == 30962
    assert len(traj) == 7

    # do ctr
    info = trajprocess.postprocess.ctr(info, "trek")

    # check ctr info
    assert not os.path.exists("{workdir}/cpptraj.tmp".format(**info['path']))
    assert not os.path.exists(
        "{workdir}/ctr/cpptraj.tmp".format(**info['path']))
    traj2 = mdtraj.load(info['ctr']['gens'][0], top=info['stp']['outtop'])

    # check ctr results
    # Trek has 518 protein residues
    pairs = np.random.randint(0, 518, (20, 2))
    cont1, _ = mdtraj.compute_contacts(traj, pairs)
    cont2, _ = mdtraj.compute_contacts(traj2, pairs)

    np.testing.assert_array_almost_equal(cont1, cont2, decimal=4)

def catkhrd(trajectories):

     # define empty lists

     D218 = []
     D222 = []

     for traj in trajectories:

          #append h188s218 difference

          h188s218 = md.compute_contacts(traj, [[120,151]],scheme='ca')
          D218.append(h188s218[0])

          #append k97s222 difference

          k97s222 = md.compute_contacts(traj, [[29,155]],scheme='ca')
          D222.append(k97s222[0])

     #flatten these lists of arrays

     flattened_h188s218 = np.asarray([val for sublist in D218 for val in sublist])
     flattened_k97s222 = np.asarray([val for sublist in D222 for val in sublist])

     return [flattened_h188s218, flattened_k97s222]

def test_contact_0():

    pdb = md.load(get_fn('bpti.pdb'))
    contacts = np.loadtxt(get_fn('contacts.dat')).astype(int)

    ca, ca_pairs = md.compute_contacts(pdb, contacts, scheme='ca')
    closest, closest_pairs = md.compute_contacts(pdb, contacts, scheme='closest')
    closest_heavy, closest_heavy_pairs = md.compute_contacts(pdb, contacts, scheme='closest-heavy')
    sidechain, sidechain_pairs = md.compute_contacts(pdb, contacts, scheme='sidechain')
    sidechain_heavy, sidechain_heavy_pairs = md.compute_contacts(pdb, contacts, scheme='sidechain-heavy')

    ref_ca = np.loadtxt(get_fn('cc_ca.dat'))
    ref_closest = np.loadtxt(get_fn('cc_closest.dat'))
    ref_closest_heavy = np.loadtxt(get_fn('cc_closest-heavy.dat'))
    ref_sidechain = np.loadtxt(get_fn('cc_sidechain.dat'))
    ref_sidechain_heavy = np.loadtxt(get_fn('cc_sidechain-heavy.dat'))

    eq(ref_ca, ca.flatten())
    eq(ref_closest, closest.flatten())
    eq(ref_closest_heavy, closest_heavy.flatten())
    eq(ref_sidechain, sidechain.flatten())
    eq(ref_sidechain_heavy, sidechain_heavy.flatten())
    eq(contacts, ca_pairs)
    eq(contacts, closest_pairs)
    eq(contacts, closest_heavy_pairs)
    eq(contacts, sidechain_pairs)
    eq(contacts, sidechain_heavy_pairs)

def shukla_coords_byrun(files,KER,Aloop,SRC2):

    difference = []
    rmsd = []

    difference_combinetrajs = []
    rmsd_combinetrajs = []

    path_base = files.split('*')[0]
    clone0_files = "%s/*clone0.h5" % path_base
    globfiles = glob(clone0_files)

    runs_list = []

    for filename in globfiles:
        run_string = re.search('run([^-]+)',filename).group(1)
        run = int(run_string)
        if run not in runs_list:
            runs_list.append(run)
        runs_list.sort()


    for run in runs_list:

        trajectories = dataset.MDTrajDataset("%s/run%d-clone*1.h5" % (path_base,run))
        print "Run %s has %s trajectories." % (run,len(trajectories))

        for traj in trajectories:

            # append difference
            k295e310 = md.compute_contacts(traj, [KER[0]])
            e310r409 = md.compute_contacts(traj, [KER[1]])
            difference_combinetrajs.append(10*(e310r409[0] - k295e310[0])) # 10x because mdtraj is naturally in nm

            # append rmsd
            Activation_Loop_SRC2 = SRC2.top.select("backbone and (resid %s to %s)" %(Aloop[0],Aloop[1]))
            Activation_Loop_kinase = traj.top.select("backbone and (resid %s to %s)" %(Aloop[0],Aloop[1]))

            SRC2_cut = SRC2.atom_slice(Activation_Loop_SRC2)
            traj_cut = traj.atom_slice(Activation_Loop_kinase)

            rmsd_combinetrajs.append(10*(md.rmsd(traj_cut,SRC2_cut,frame=0))) # 10x because mdtraj is naturaly in nm

        # flatten list of arrays
        difference_combinetrajs = np.asarray([val for sublist in difference_combinetrajs for val in sublist])
        rmsd_combinetrajs = np.asarray([val for sublist in rmsd_combinetrajs for val in sublist])

        difference.append(difference_combinetrajs)
        difference_combinetrajs = []

        rmsd.append(rmsd_combinetrajs)
        rmsd_combinetrajs = []

    return [rmsd, difference]

def read_and_featurize(traj_file, features_dir = None, condition=None, dihedral_types = ["phi", "psi", "chi1", "chi2"], dihedral_residues = None, resSeq_pairs = None, iterative = True):

	a = time.time()
	dihedral_indices = []
	residue_order = []
	if len(dihedral_residues) > 0:
		for dihedral_type in dihedral_types:
			if dihedral_type == "phi": dihedral_indices.append(phi_indices(fix_topology(top), dihedral_residues))
			if dihedral_type == "psi": dihedral_indices.append(psi_indices(fix_topology(top), dihedral_residues))
			if dihedral_type == "chi1": dihedral_indices.append(chi1_indices(fix_topology(top), dihedral_residues))
			if dihedral_type == "chi2": dihedral_indices.append(chi2_indices(fix_topology(top), dihedral_residues))

		#print("new features has dim %d" %(2*len(phi_tuples) + 2*len(psi_tuples) + 2*len(chi2_tuples)))

		#print("feauturizing manually:")
		dihedral_angles = []

		for dihedral_type in dihedral_indices:
			angles = np.transpose(ManualDihedral.compute_dihedrals(traj=traj,indices=dihedral_type))
			dihedral_angles.append(np.sin(angles))
			dihedral_angles.append(np.cos(angles))

		manual_features = np.transpose(np.concatenate(dihedral_angles))

	if len(resSeq_pairs) > 0:
		top = md.load_frame(traj_file, index=0).topology
		resIndex_pairs = convert_resSeq_to_resIndex(top, resSeq_pairs)
		contact_features = []
		if iterative:
			try:
				for chunk in md.iterload(traj_file, chunk = 1000):
				#	chunk = fix_traj(chunk)
				#chunk = md.load(traj_file,stride=1000)
				#print(resIndex_pairs[0:10])
					chunk_features = md.compute_contacts(chunk, contacts = resIndex_pairs, scheme = 'closest-heavy', ignore_nonprotein=False)[0]
					print(np.shape(chunk_features))
					contact_features.append(chunk_features)
				contact_features = np.concatenate(contact_features)
			except Exception,e:
				print str(e)
				print("Failed")
				return
				#traj = md.load(traj_file)
				#contact_features = md.compute_contacts(chunk, contacts = contact_residue_pairs, scheme = 'closest-heavy', ignore_nonprotein=False)[0]
		else:
			try:
				traj = md.load(traj_file)
				contact_features =  md.compute_contacts(traj, contacts = resIndex_pairs, scheme = 'closest-heavy', ignore_nonprotein=False)[0]
			except Exception,e:
				print str(e)
				print("Failed for traj")
				return

    def partial_transform(self, traj):
        """Featurize an MD trajectory into a vector space derived from
        residue-residue distances

        Parameters
        ----------
        traj : mdtraj.Trajectory
            A molecular dynamics trajectory to featurize.

        Returns
        -------
        features : np.ndarray, dtype=float, shape=(n_samples, n_features)
            A featurized trajectory is a 2D array of shape
            `(length_of_trajectory x n_features)` where each `features[i]`
            vector is computed by applying the featurization function
            to the `i`th snapshot of the input trajectory.

        See Also
        --------
        transform : simultaneously featurize a collection of MD trajectories
        """

        distances, _ = md.compute_contacts(traj, self.contacts,
                                           self.scheme, self.ignore_nonprotein)
        return self._transform(distances)

def compute_contacts_below_cutoff(traj_file_frame, cutoff = 100000.0, contact_residues = [], anton = False):
	traj_file = traj_file_frame[0]
	frame = md.load_frame(traj_file, index = 0)
	#frame = fix_traj(frame)
	top = frame.topology
	
	distance_residues = []
	res_indices = []
	resSeq_to_resIndex = {}
	residue_full_infos = []

	for i in range(0, len(contact_residues)):
		residue = contact_residues[i]
		indices = [r.index for r in top.residues if r.resSeq == residue[1] and r.chainid == residue[0] and not r.is_water]
		if len(indices) == 0:
			print("No residues in trajectory for residue %d" %residue)
			continue
		else:
			ind = indices[0]
			for j in indices:
				if j != ind: 
					#print("Warning: multiple res objects for residue %d " %residue)
					if "CB" in [str(a) for a in r.atoms for r in top.residues if r.index == ind]:
						ind = j
			res_indices.append(ind)
			distance_residues.append(residue)
			resSeq_to_resIndex[residue] = ind
	
	resSeq_combinations = itertools.combinations(distance_residues, 2)
	res_index_combinations = []
	resSeq_pairs = [c for c in resSeq_combinations]
	for combination in resSeq_pairs:
		res0 = combination[0]
		res1 = combination[1]
		res_index0 = resSeq_to_resIndex[res0]
		res_index1 = resSeq_to_resIndex[res1]
		res_index_combinations.append((res_index0, res_index1))


	final_resSeq_pairs = []
	final_resIndex_pairs = []

	distances = md.compute_contacts(frame, contacts = res_index_combinations, scheme = 'closest-heavy', ignore_nonprotein=False)[0]
	#print(distances)
	print(np.shape(distances))
	for i in range(0, len(distances[0])):
		distance = distances[0][i]
		#print(distance)
		if distance < cutoff:
			final_resIndex_pairs.append(res_index_combinations[i])
			final_resSeq_pairs.append(resSeq_pairs[i])

	for pair in final_resIndex_pairs:
		info0 = [(r.resSeq, r.name, r.chain.index) for r in top.residues if r.index == pair[0]]
		info1 = [(r.resSeq, r.name, r.chain.index) for r in top.residues if r.index == pair[1]]
		residue_full_infos.append((info0, info1))

	print(len(final_resSeq_pairs))
	print(len(final_resIndex_pairs))
	return((final_resSeq_pairs, residue_full_infos))

    def partial_transform(self, traj):
        """Featurize an MD trajectory into a vector space derived from
        residue-residue distances

        Parameters
        ----------
        traj : mdtraj.Trajectory
            A molecular dynamics trajectory to featurize.

        Returns
        -------
        features : np.ndarray, dtype=float, shape=(n_samples, n_features)
            A featurized trajectory is a 2D array of shape
            `(length_of_trajectory x n_features)` where each `features[i]`
            vector is computed by applying the featurization function
            to the `i`th snapshot of the input trajectory.

        See Also
        --------
        transform : simultaneously featurize a collection of MD trajectories
        """

        # check to make sure topologies are consistent with the reference frame
        try:
            assert traj.top == self.reference_frame.top
        except:
            warnings.warn("The topology of the trajectory is not" +
                          "the same as that of the reference frame," +
                          "which might give meaningless results.")
        distances, _ = md.compute_contacts(traj, self.contacts,
                                        self.scheme, ignore_nonprotein=False,
                                        periodic = self.periodic)
        return self._transform(distances)

    def describe_features(self, traj):
        """Return a list of dictionaries describing the features in Contacts."""
        x = []
        # fill in the atom indices using just the first frame
        distances, residue_indices = md.compute_contacts(traj, self.contacts, self.scheme, self.ignore_nonprotein)
        n = residue_indices.shape[0]
        aind = ["N/A"] * n
        resSeq = [np.array([traj.top.residue(j).resSeq for j in i]) for i in residue_indices]
        resid = [np.array([traj.top.residue(j).index for j in i]) for i in residue_indices]
        resnames = [[traj.topology.residue(j).name for j in i] for i in resid]
        bigclass = [self.contacts] * n
        smallclass = [self.scheme] * n
        otherInfo = [self.ignore_nonprotein] * n

        for i in range(n):
            d_i = dict(
                resname=resnames[i],
                atomind=aind[i],
                resSeq=resSeq[i],
                resid=resid[i],
                otherInfo=otherInfo[i],
                bigclass=bigclass[i],
                smallclass=smallclass[i],
            )
            x.append(d_i)

        return x

def plot_native_state_contact_map(title):

    colors = [('white')] + [(cm.jet(i)) for i in xrange(1,256)]
    new_map = matplotlib.colors.LinearSegmentedColormap.from_list('new_map', colors, N=256)

    if os.path.exists("contact_pairs.dat") and os.path.exists("contact_probabilities.dat"):
        pairs = np.loadtxt("contact_pairs.dat")
        probability = np.loadtxt("contact_probabilities.dat")
    else:
        print "  Loading BeadBead.dat"
        beadbead = np.loadtxt("BeadBead.dat",dtype=str)
        sigij = beadbead[:,5].astype(float)
        epsij = beadbead[:,6].astype(float)
        deltaij = beadbead[:,7].astype(float)
        interaction_numbers = beadbead[:,4].astype(str)
        pairs = beadbead[:,:2].astype(int)
        pairs -= np.ones(pairs.shape,int)
        np.savetxt("contact_pairs.dat",pairs)

        print "  Computing distances with mdtraj..."
        traj = md.load("traj.xtc",top="Native.pdb")
        distances = md.compute_contacts(traj,pairs)
        contacts = (distances[0][:] <= 1.2*sigij).astype(int)
        print "  Computing contact probability..."
        probability = sum(contacts.astype(float))/contacts.shape[0]
        np.savetxt("contact_probabilities.dat",probability)

    Qref = np.loadtxt("Qref_cryst.dat")
    C = np.zeros(Qref.shape,float)

    for k in range(len(pairs)):
        C[pairs[k][0],pairs[k][1]] = probability[k]

    print "  Plotting..."
    plt.figure()
    plt.subplot(1,1,1,aspect=1)
    ax = plt.subplot(1,1,1,aspect=1)
    plt.pcolor(C,cmap=new_map)
    for k in range(len(pairs)):
        if probability[k] > 0.01:
            plt.plot(pairs[k][1],pairs[k][0],marker='s',ms=3.0,markeredgecolor=new_map(probability[k]),color=new_map(probability[k]))
        else:
            continue
    plt.xlim(0,len(Qref))
    plt.ylim(0,len(Qref))
    #plt.text(10,70,name.upper(),fontsize=70,color="r")
    ax = plt.gca()
    cbar = plt.colorbar()
    cbar.set_clim(0,1)
    cbar.set_label("Contact probability",fontsize=20)
    cbar.ax.tick_params(labelsize=20)
    plt.xlabel("Residue i",fontsize=20)
    plt.ylabel("Residue j",fontsize=20)
    #plt.title("Native State Contact Map "+title,fontsize=20)
    plt.title(title)
    for label in ax.get_xticklabels() + ax.get_yticklabels():
        label.set_fontsize(15)
    print "  Saving..."
    plt.savefig("native_state_contact_map.pdf")

def test_contact_3(get_fn):
    pdb = md.load(get_fn('bpti.pdb'))
    beta = 20
    dists, pairs = md.compute_contacts(pdb, soft_min=True, soft_min_beta=beta)

    maps = md.geometry.squareform(dists, pairs)
    for i, (r0, r1) in enumerate(pairs):
        for t in range(pdb.n_frames):
            assert np.allclose(beta / np.log(np.sum(np.exp(beta / maps[t, r0, r1]))), dists[t, i])

def shukla_coords(trajectories,KER,Aloop,SRC2):

    difference = []
    rmsd = []

    for traj in trajectories:

        # append difference
        k295e310 = md.compute_contacts(traj, [KER[0]])
        e310r409 = md.compute_contacts(traj, [KER[1]])
        difference.append(10*(e310r409[0] - k295e310[0])) # 10x because mdtraj is naturally in nm

        # append rmsd
        Activation_Loop_SRC2 = SRC2.top.select("backbone and (resid %s to %s)" %(Aloop[0],Aloop[1]))
        Activation_Loop_kinase = traj.top.select("backbone and (resid %s to %s)" %(Aloop[0],Aloop[1]))

        SRC2_cut = SRC2.atom_slice(Activation_Loop_SRC2)
        traj_cut = traj.atom_slice(Activation_Loop_kinase)

        rmsd.append(10*(md.rmsd(traj_cut,SRC2_cut,frame=0))) # 10x because mdtraj is naturaly in nm

    return [rmsd, difference]

def test_contact_1():
    pdb = md.load(get_fn('bpti.pdb'))
    dists, pairs = md.compute_contacts(pdb)
    for r0, r1 in pairs:
        # are these valid residue indices?
        pdb.topology.residue(r0)
        pdb.topology.residue(r1)

        assert not (abs(r0 - r1) < 3)

    maps = md.geometry.squareform(dists, pairs)
    for i, (r0, r1) in enumerate(pairs):
        for t in range(pdb.n_frames):
            eq(maps[t, r0, r1], dists[t, i])

def test_ContactFeaturizer_describe_features():
    scheme = np.random.choice(['ca','closest','closest-heavy'])
    feat = ContactFeaturizer(scheme=scheme, ignore_nonprotein=True)
    rnd_traj = np.random.randint(len(trajectories))
    features = feat.transform([trajectories[rnd_traj]])
    df = pd.DataFrame(feat.describe_features(trajectories[rnd_traj]))

    for f in range(25):
        f_index = np.random.choice(len(df))

        residue_ind = df.iloc[f_index].resids
        feature_value, _ = md.compute_contacts(trajectories[rnd_traj],
                                               contacts=[residue_ind],
                                               scheme=scheme)
        assert (features[0][:, f_index] == feature_value.flatten()).all()