Python mdtraj.load_frame函数代码示例

def sample_clusters(clusterer_dir, features_dir, traj_dir, save_dir, n_samples):
	clusters_map = dist_to_means(clusterer_dir, features_dir)
	if not os.path.exists(save_dir): os.makedirs(save_dir)
	
	#non_palm = get_traj_no_palm(traj_dir)

	trajectories = get_trajectory_files(traj_dir)

	for cluster in range(0, len(clusters_map.keys())):
		for s in range(0, n_samples):
			sample = clusters_map[cluster][s]
			traj_id = sample[0]
			frame = sample[1]
			traj = trajectories[traj_id]

			top = md.load_frame(traj, index=frame).topology
			indices = [a.index for a in top.atoms if str(a.residue)[0:3] != "SOD" and str(a.residue)[0:3] != "CLA" and a.residue.resSeq < 341]

			conformation = md.load_frame(traj, index=frame, atom_indices=indices)
			conformation.save_pdb("%s/cluster%d_sample%d.pdb" %(save_dir, cluster, s))
	
	remove_ter(save_dir)
	reorder(save_dir)
	#remove_palm(save_dir)
	new_dir = reimage(save_dir)

def map_drawn_samples(selected_pairs_by_state, trajectories, top=None):
    """Lookup trajectory frames using pairs of (trajectory, frame) indices.

    Parameters
    ----------
    selected_pairs_by_state : np.ndarray, dtype=int, shape=(n_states, n_samples, 2)
        selected_pairs_by_state[state, sample] gives the (trajectory, frame)
        index associated with a particular sample from that state.
    trajectories : list(md.Trajectory) or list(np.ndarray) or list(filenames)
        The trajectories assocated with sequences,
        which will be used to extract coordinates of the state centers
        from the raw trajectory data.  This can also be a list of np.ndarray
        objects or filenames.  If they are filenames, mdtraj will be used to load
    top : md.Topology, optional, default=None
        Use this topology object to help mdtraj load filenames

    Returns
    -------
    frames_by_state : mdtraj.Trajectory
        Output will be a list of trajectories such that frames_by_state[state]
        is a trajectory drawn from `state` of length `n_samples`.  If trajectories
        are numpy arrays, the output will be numpy arrays instead of md.Trajectories
    
    Examples
    --------
    >>> selected_pairs_by_state = hmm.draw_samples(sequences, 3)
    >>> samples = map_drawn_samples(selected_pairs_by_state, trajectories)
    
    Notes
    -----
    YOU are responsible for ensuring that selected_pairs_by_state and 
    trajectories correspond to the same dataset!
    
    See Also
    --------
    utils.map_drawn_samples : Extract conformations from MD trajectories by index.
    ghmm.GaussianFusionHMM.draw_samples : Draw samples from GHMM    
    ghmm.GaussianFusionHMM.draw_centroids : Draw centroids from GHMM    
    """

    frames_by_state = []

    for state, pairs in enumerate(selected_pairs_by_state):
        if isinstance(trajectories[0], str):
            import mdtraj as md
            if top:
                process = lambda x, frame: md.load_frame(x, frame, top=top)
            else:
                process = lambda x, frame: md.load_frame(x, frame)
        else:
            process = lambda x, frame: x[frame]

        frames = [process(trajectories[trj], frame) for trj, frame in pairs]
        try:  # If frames are mdtraj Trajectories
            state_trj = frames[0][0:0].join(frames)  # Get an empty trajectory with correct shape and call the join method on it to merge trajectories
        except AttributeError:
            state_trj = np.array(frames)  # Just a bunch of np arrays
        frames_by_state.append(state_trj)
    
    return frames_by_state

def test_load_frame():
    files = [
        "frame0.nc",
        "frame0.h5",
        "frame0.xtc",
        "frame0.trr",
        "frame0.dcd",
        "frame0.mdcrd",
        "frame0.binpos",
        "frame0.xyz",
        "frame0.lammpstrj",
    ]
    if not (on_win and on_py3):
        files.append("legacy_msmbuilder_trj0.lh5")

    trajectories = [md.load(get_fn(f), top=get_fn("native.pdb")) for f in files]
    rand = [np.random.randint(len(t)) for t in trajectories]
    frames = [md.load_frame(get_fn(f), index=r, top=get_fn("native.pdb")) for f, r in zip(files, rand)]

    for traj, frame, r, f in zip(trajectories, frames, rand, files):

        def test():
            eq(traj[r].xyz, frame.xyz)
            eq(traj[r].unitcell_vectors, frame.unitcell_vectors)
            eq(traj[r].time, frame.time, err_msg="%d, %d: %s" % (traj[r].time[0], frame.time[0], f))

        test.description = "test_load_frame: %s" % f
        yield test

    t1 = md.load(get_fn("2EQQ.pdb"))
    r = np.random.randint(len(t1))
    t2 = md.load_frame(get_fn("2EQQ.pdb"), r)
    eq(t1[r].xyz, t2.xyz)

def test_residues_map_num_atoms(traj_file_1, traj_file_2, residues, residues_map):
	traj_1 = md.load_frame(traj_file_1, index = 0)
	traj_2 = md.load_frame(traj_file_2, index = 0)
	top1 = traj_1.topology
	top2 = traj_2.topology
	for residue in residues:
		new_residue = residues_map[residue]
		atoms = [a.index for a in top1.atoms if a.residue.resSeq == residue and a.residue.is_protein]
		len1 = len(atoms)
		atoms = [a.index for a in top2.atoms if a.residue.resSeq == new_residue and a.residue.is_protein]
		len2 = len(atoms)
		if (len1 != len2) or (len1 == len2):
			print("Atom number %d %d doesn't match for residue %d" %(len1, len2, residue))
	return

def test_residues_map(traj_file_1, traj_file_2, residues, residues_map):
	traj_1 = md.load_frame(traj_file_1, index = 0)
	traj_2 = md.load_frame(traj_file_2, index = 0)
	top1 = traj_1.topology
	top2 = traj_2.topology
	for residue in residues:
		new_residue = residues_map[residue]
		print("Original residues:")
		residues = [r for r in top1.residues if r.resSeq == residue and r.is_protein]
		print(residues[0])
		print("New residues:")
		residues = [r for r in top2.residues if r.resSeq == new_residue and r.is_protein]
		print(residues[0])
	return

def read_and_featurize(filename, dihedrals=['chi2'], stride=10):
	#print("reading and featurizing %s" %(filename))
	top = md.load_frame(filename, 0).topology
	#print("got top")
	atom_indices = [a.index for a in top.atoms if a.residue.resSeq == 93 and a.residue != "POPC" and str(a.residue)[0] == "H"]
	print(len(atom_indices))
	#atom_indices = [a.index for a in top.atoms if a.residue.chain.index == 0 and a.residue.resSeq != 93 and a.residue != "POPC" and a.residue.resSeq != 130 and a.residue.resSeq != 172 and a.residue.resSeq != 79 and a.residue.resSeq != 341]
	#print("got indices")
	traj = md.load(filename, stride=1000, atom_indices=atom_indices)
	#print("got traj")
	featurizer = DihedralFeaturizer(types = dihedrals)
	features = featurizer.transform(traj_list = traj)
	#print(np.shape(features))
	#print("finished featurizing")

	directory = filename.split("/")
	condition = directory[len(directory)-2]
	dcd_file = directory[len(directory)-1]
	new_file = "%s_features_stride%d.h5" %(dcd_file.rsplit( ".", 1 )[ 0 ] , stride)
	new_root_dir = "/scratch/users/enf/b2ar_analysis/subsampled_features"
	new_condition_dir = "%s/%s" %(new_root_dir, condition)

	new_file_full = "%s/%s/%s" %(new_root_dir, condition, new_file)
	#print("saving features as %s" %new_file_full)

	verbosedump(features, new_file_full)
	return features

def loadFrames(confs_by_state):
    """
    input is array of arrays
    """
    frames = []
    for elem in confs_by_state:
        trajFrames = []
        for trajFrame in elem:
            file = os.path.basename(trajFrame[0])
            frame = trajFrame[1]

            regex = "(.*)_traj.*_(\d*).xtc"
            m = re.match(regex,file)
            projectName = m.group(1)
            trajNum = m.group(2)

            #now find the actual trajectory
            #TODO also get the regular traj
            originalTraj = "../%s/analysis/full/traj_full_%s.xtc"%(projectName,trajNum)

            #load the ref
            ref = "../%s/analysis/full/ref.pdb"%projectName
            print ("loading %s frame %s"%(originalTraj,frame))
            loadedFrame = md.load_frame(originalTraj,frame,top=ref)

            trajFrames.append(loadedFrame)

        frames.append(trajFrames)

    return frames

def reproject_oldata():
  r1 = redis.StrictRedis(port=6390, decode_responses=True)
  cache = redis.StrictRedis(host='bigmem0006', port=6380, decode_responses=True)
  execlist = r1.hgetall('anl_sequence')
  keyorder = ['jc_'+i[0] for i in sorted(execlist.items(), key=lambda x:x[1])]
  # skip first 100 (non-sampled)
  pts = []
  bad_ref = 0
  miss = 0
  for key in keyorder:
    conf = r1.hgetall(key)
    src = int(conf['src_index'])
    ref = r1.lindex('xid:reference', src)
    if ref is not None:
      fileno, frame = eval(ref)
      ckey = 'sim:%s' % conf['name']
      xyz = cache.lindex(ckey, frame)
      if xyz is not None:
        pts.append(pickle.loads(xyz))
      else:
        tr = md.load_frame(conf['dcd'], frame, top=conf['pdb'])
        if len(tr.xyz) == 0:
          miss += 1
        else:
          pts.append(tr.xyz[0])
    else:
      bad_ref += 1
  traj = md.Trajectory(pts, deshaw.topo_prot.top)
  alpha = datareduce.filter_alpha(traj)
  return alpha

def save_pdb(traj_dir, clusterer, i):
	location = clusterer.cluster_ids_[i,:]
	traj = get_trajectory_files(traj_dir)[location[0]]
	print("traj = %s, frame = %d" %(traj, location[1]))
	conformation = md.load_frame(traj, location[1])
	conformation.save_pdb("/scratch/users/enf/b2ar_analysis/clusters_1000_allprot/%d.pdb" %i)
	return None

def save_features_to_residues_map(traj_file, contact_residues, feature_residues_csv, cutoff, residues_map = None, exacycle = False):
	if residues_map is not None:
		contact_residues = [r for r in contact_residues if r in residues_map.keys()]
		if exacycle: contact_residues = [residues_map[key] for key in contact_residues]

	traj = md.load_frame(traj_file, 0)
	#traj = fix_traj(traj)
	top = traj.topology 
	residue_pairs, residue_infos = compute_contacts_below_cutoff([traj_file, [0]], cutoff = cutoff, contact_residues = contact_residues, anton = False)
	if exacycle:
		reverse_residues_map = {v: k for k, v in residues_map.items()}
		new_residue_pairs = []
		for residue_pair in residue_pairs:
			new_residue_pair = [reverse_residues_map[residue_pair[0]], reverse_residues_map[residue_pair[1]]]
			new_residue_pairs.append(new_residue_pair)
		residue_pairs = new_residue_pairs

		new_reisdue_infos = []
		for residue_info in residue_infos:
			new_residue_info = [(reverse_residues_map[residue_info[0][0]], residue_info[0][1], residue_info[0][2]), (reverse_residues_map[residue_info[1][0]], residue_info[1][1], residue_info[1][2])]
			new_residue_infos.append(new_residue_info)
		residue_infos = new_reisdue_infos

	print("There are: %d residue pairs" %len(residue_pairs))
	f = open(feature_residues_csv, "wb")
	f.write("feature, residue.1.resSeq, residue.1.res, residue.1.chain, residue.2.resSeq, residue.2.res, residue.2.chain,\n")
	for i in range(0, len(residue_infos)):
		f.write("%d, %d, %d, %d, %d, %d, %d,\n" %(i, residue_infos[i][0][0], residue_infos[i][0][1], residue_infos[i][0][2], residue_infos[i][1][0], residue_infos[i][1][1], residue_infos[i][1][2]))
	f.close()
	return 

def find_most_important_residues_in_tIC(traj_file, tica_object, tic_features_csv, contact_residues,tic_residue_csv, feature_coefs_csv, duplicated_feature_coefs_csv, cutoff):
	try:
		tica = verboseload(tica_object)
	except:
		tica = load_dataset(tica_object)
	print traj_file
	traj = md.load_frame(traj_file, 0)
	#traj = fix_traj(traj)
	top = traj.topology 
	#residue_pairs = compute_contacts_below_cutoff([traj_file, [0]], cutoff = cutoff, contact_residues = contact_residues, anton = True)
	residue_pairs = generate_features(tic_features_csv)
	new_residue_pairs = []
	for pair in residue_pairs:
		new_residue_pairs.append(("%s%d.%d" %(pair[0][2], pair[0][1], pair[0][0])), ("%s%d.%d" %(pair[1][2], pair[1][1], pair[1][0])))
	residue_pairs = new_residue_pairs
	#print traj_file

	
	top_indices_per_tIC = {}
	feature_coefs_per_tIC = {}
	duplicated_feature_coefs_per_tIC = {}


	#for each tIC:
		#for each feature, get the absolute component value
		#add to feature_coefs_per_tIC dictionary the absolute coefficient for that tIC
		#duplicate them for the analysis where we look at residues individually
		#sort by absolute coefficient value

	#for each tIC:
		#

	for i in range(0, np.shape(tica.components_)[0]):
		print i
		index_components = [(j,abs(tica.components_[i][j])) for j in range(0,np.shape(tica.components_)[1])]
		feature_coefs_per_tIC[i] = [component[1] for component in index_components]
		duplicated_feature_coefs_per_tIC[i] = [j for k in feature_coefs_per_tIC[i] for j in (k, k)] 
		index_components = sorted(index_components, key= lambda x: x[1],reverse=True)
		print(index_components[0:10])
		list_i = [index_components[j][0] for j in range(0,len(index_components))]
		top_indices_per_tIC[i] = list_i
	
	top_residues_per_tIC = {}
	for i in range(0, np.shape(tica.components_)[0]):
		top_residues_per_tIC[i] = []
		for index in top_indices_per_tIC[i]:
			residues = residue_pairs[index]
			top_residues_per_tIC[i].append(residues)
		top_residues_per_tIC[i] = [item for sublist in top_residues_per_tIC[i] for item in sublist]

	residue_list = residue_pairs

	feature_coefs_per_tIC["residues_0"] = [pair[0] for pair in residue_list]
	feature_coefs_per_tIC["residues_1"] = [pair[1] for pair in residue_list]
	duplicated_feature_coefs_per_tIC["residues"] = [residue for residue_pair in residue_list for residue in residue_pair]

	write_map_to_csv(tic_residue_csv, top_residues_per_tIC, [])
	write_map_to_csv(feature_coefs_csv, feature_coefs_per_tIC, [])
	write_map_to_csv(duplicated_feature_coefs_csv, duplicated_feature_coefs_per_tIC, [])
	return

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 rmsd_to_structure(clusters_dir, ref_dir, text):
	pdbs = get_trajectory_files(clusters_dir)

	ref = md.load_frame(ref_dir, index=0)
	rmsds = np.zeros(shape=(len(pdbs),2))

	for i in range(0,len(pdbs)):
		print i 
		pdb_file = pdbs[i]
		pdb = md.load_frame(pdb_file, index=0)
		rmsd = md.rmsd(pdb, ref, 0)
		rmsds[i,0] = i
		rmsds[i,1] = rmsd[0]

	rmsd_file = "%s/%s_rmsds.csv" %(clusters_dir, text)
	np.savetxt(rmsd_file, rmsds, delimiter=",")

def read_and_featurize_divided(filename, dihedrals=['phi', 'psi', 'chi2'], stride=10):
	#print("reading and featurizing %s" %(filename))

	traj_top = md.load_frame(filename,0).topology
	atom_indices = [a.index for a in traj_top.atoms if a.residue.name[0:2] != "HI"]

	traj = md.load(filename,atom_indices=atom_indices)
	#print("got traj")
	featurizer = DihedralFeaturizer(types = dihedrals)
	features = featurizer.transform(traj_list = traj)
	#print(np.shape(features))
	#print("finished featurizing")

	directory = filename.split("/")
	condition = directory[len(directory)-2]
	dcd_file = directory[len(directory)-1]
	new_file = "%s_features_stride%d.h5" %(dcd_file.rsplit( ".", 1 )[ 0 ] , stride)
	new_root_dir = "/scratch/users/enf/b2ar_analysis/subsampled_features"
	new_condition_dir = "%s/%s" %(new_root_dir, condition)

	new_file_full = "%s/%s/%s" %(new_root_dir, condition, new_file)
	#print("saving features as %s" %new_file_full)

	verbosedump(features, new_file_full)
	return features

def timefld(n):
    start = dt.datetime.now()
    tr = md.load_frame("bpti-all-1%03d.dcd" % n, 23, top=pdb)
    tr.atom_slice(tr.top.select("protein"), inplace=True)
    end = dt.datetime.now()
    print("Time: ", (end - start).total_seconds())
    return tr