Python MDAnalysis.Universe类代码示例

 def test_bonds(self):
     u = Universe(self.filename, guess_bonds=True)
     # need to force topology to load before querying individual atom bonds
     u.build_topology()
     bonds0 = u.select_atoms("segid B and (altloc A)")[0].bonds
     bonds1 = u.select_atoms("segid B and (altloc B)")[0].bonds
     assert_equal(len(bonds0), len(bonds1))

def calc_rama(grof, xtcf, btime, etime):
    u = Universe(grof, xtcf)

    resname_query = 'resname GLY or resname VAL or resname PRO'
    atoms = u.selectAtoms(resname_query)
    resname = atoms.resnames()[0] # [0] because .resnames() returns a list of one element
    resid = atoms.resids()[0] # [0] because .resnames() returns a list of one element

    phi_query = ('(resname ACE and name C) or '
                 '(resname GLY or resname VAL or resname PRO and '
                 '(name N or name CA or name C))')

    psi_query = ('(resname GLY or resname VAL or resname PRO and (name N or name CA or name C or name NT)) or '
                 '(resname NH2 and name N)')

    # MDAnalysis will convert the unit of length to angstrom, though in Gromacs the unit is nm
    phi = u.selectAtoms(phi_query)
    psi = u.selectAtoms(psi_query)

    for _ in phi.atoms:
        print _

    for _ in psi.atoms:
        print _


    for ts in u.trajectory:
        if btime > ts.time:
            continue
        if etime > 0 and etime < ts.time:
            break

        yield '{0:.3f}  {1:.3f}  {2}-{3}\n'.format(
            phi.dihedral(), psi.dihedral(), resname, resid)
        U.print_progress(ts)

def main():
    arg_parser = argparse.ArgumentParser(description='通过给定残基名称，残基内原子数目，两个原子在残基内的索引(从0开始)，计算所有残基内这两个原子之间的直线距离。')
    arg_parser.add_argument('resname', action='store', help='残基名称')
    arg_parser.add_argument('atoms_num', type=int, action='store', help='残基内原子数目')
    arg_parser.add_argument('index1', type=int, action='store', help='第一个原子的索引，索引从0开始')
    arg_parser.add_argument('index2', type=int, action='store', help='第二个原子的索引，索引从0开始')
    arg_parser.add_argument('topology_file', action='store', help='拓扑文件，例如gro, pdb')
    args = arg_parser.parse_args()

    resname, atoms_num, index1, index2 = args.resname, args.atoms_num, args.index1, args.index2

    universe = Universe(args.topology_file)
    atom_groups = universe.selectAtoms("resname " + resname)
    if len(atom_groups) % atoms_num != 0:
        print("拓扑文件内对应残基原子总数不是所给原子数目的整数倍，请给予正确的原子数目。")
        exit(1)

    atoms1 = []
    atoms2 = []
    for i in range(0, len(atom_groups), atoms_num):
        atoms1.append(atom_groups[i:i + atoms_num][index1])
        atoms2.append(atom_groups[i:i + atoms_num][index2])

    dists = dist(AtomGroup(atoms1), AtomGroup(atoms2))
    print("The distance between atoms %s and %s is:" % (index1, index2))
    for i in dists[2]:
        print(i)
    print("The average distance between atoms %s and %s is:" % (index1, index2))
    print(np.average(dists[2]))

def gen_hbond_map(xpm, ndx, grof):
    xpm = objs.XPM(xpm)
    hbndx = objs.HBNdx(ndx)

    univ = Universe(grof)
    pro_atoms = univ.selectAtoms('protein and not resname ACE and not resname NH2')
    hbonds_by_resid = hbndx.map_id2resid(pro_atoms)

    # pl: peptide length
    pl = pro_atoms.residues.numberOfResidues()

    hblist = []
    for i, j in zip(hbonds_by_resid, xpm.color_count):
        # j[1] is the probability of hbonds, while j[0] = 1 - j[1]
        # format: [resid of donor, resid of acceptor]
        # -1 is because resid in MDAnalysis starts from 1, minus so as to fit
        # -into hb_map initialized by hb_map
        hblist.append([i[0]-1, i[1]-1, j[1]])

    # +1: for missing resname ACE, such that it's easier to proceed in the next
    # step
    pl1 = pl + 1
    hb_map = np.zeros((pl1, pl1))
    for _ in hblist:
        hb_map[_[0]][_[1]] = _[2]

    return hb_map

 def test_write_selection(self):
     ref = Universe(mol2_molecule)
     gr0 = ref.select_atoms("name C*")
     gr0.write(self.outfile)
     u = Universe(self.outfile)
     gr1 = u.select_atoms("name C*")
     assert_equal(len(gr0), len(gr1))

def count_interactions(grof, xtcf, btime, cutoff, debug):
    u = Universe(grof, xtcf)
    un_query = ('(resname PRO and (name CB or name CG or name CD)) or'
                '(resname VAL and (name CG1 or name CG2)) or'
                '(resname GLY and name CA) or'
                '(resname ALA and name CB)')
    vp_query = ('name OW')
    # MDAnalysis will convert the unit of length to angstrom, though in Gromacs the unit is nm
    un_atoms = u.selectAtoms(un_query)
    for ts in u.trajectory:
        if ts.time >= btime:
            numcount = 0
            tropo_vp_atoms = u.selectAtoms(
                '({0}) and around 8 ({1})'.format(vp_query, un_query))
            # different from when calculating unun, there is no overlap atom
            # between un_atoms & tropo_vp_atoms
            for ai in un_atoms:
                for aj in tropo_vp_atoms:
                    d = np.linalg.norm(ai.pos - aj.pos)
                    if d <= cutoff:
                        numcount += 1
            yield '{0:10.0f}{1:8d}\n'.format(ts.time,  numcount)
        # per 100 frames, num of frames changes with the size of xtc file, for debugging
        if debug and ts.frame % 2 == 0: 
            print "time: {0:10.0f}; step: {1:10d}; frame: {2:10d}".format(ts.time, ts.step, ts.frame)

def count_interactions(grof, xtcf, btime, etime, cutoff):
    cutoff = cutoff * 10 # * 10: convert from nm to angstrom to work with MDAnalysis
    u = Universe(grof, xtcf)
    query = ('(resname PRO and (name CB or name CG or name CD)) or'
             '(resname VAL and (name CG1 or name CG2)) or'
             '(resname GLY and name CA) or'
             '(resname ALA and name CB)')
    # MDAnalysis will convert the unit of length to angstrom, though in Gromacs the unit is nm
    atoms = u.selectAtoms(query)
    for ts in u.trajectory:
        if btime > ts.time:
            continue
        if etime > 0 and etime < ts.time:
            break

        numcount = 0
        for i, ai in enumerate(atoms):
            for j, aj in enumerate(atoms):
                # to avoid counting the same pair twices,
                # the 2 resid cannot be neigbors
                if i < j and abs(ai.resid - aj.resid) >= 2: 
                    d = np.linalg.norm(ai.pos - aj.pos)
                    if d <= cutoff:
                        numcount += 1
        yield '{0:10.0f}{1:8d}\n'.format(ts.time,  numcount)
        utils.print_progress(ts)

def count_interactions(A):
    logger.debug('loading {0}'.format(A.grof))
    univ = Universe(A.grof)
    logger.debug('loaded {0}'.format(A.grof))

    pro_atoms = univ.selectAtoms('protein and not resname ACE and not resname NH2')
    pl = pro_atoms.residues.numberOfResidues()
    # +1: for missing resname ACE, such that it's easier to proceed in the next
    # step

    logger.debug('loading {0}, {1}'.format(A.grof, A.xtcf))
    u = Universe(A.grof, A.xtcf)
    logger.debug('loaded {0}, {1}'.format(A.grof, A.xtcf))

    # Just for reference to the content of query when then code was first
    # written and used
    # query = ('(resname PRO and (name CB or name CG or name CD)) or'
    #          '(resname VAL and (name CG1 or name CG2)) or'
    #          '(resname GLY and name CA) or'
    #          '(resname ALA and name CB)')

    query = A.query
    atoms = u.selectAtoms(query)
    logger.info('Number of atoms selected: {0}'.format(atoms.numberOfAtoms()))

    # MDAnalysis will convert the unit of length to angstrom, though in Gromacs
    # the unit is nm
    cutoff = A.cutoff * 10
    nres_away = A.nres_away
    btime = A.btime
    etime = A.etime
    nframe = 0
    unun_map = None
    for ts in u.trajectory:
        if btime > ts.time:
            continue
        if etime > 0 and etime < ts.time:
            break

        nframe += 1
        map_ = np.zeros((pl+1, pl+1))                   # map for a single frame
        for i, ai in enumerate(atoms):
            ai_resid = ai.resid
            for j, aj in enumerate(atoms):
                aj_resid = aj.resid
                # to avoid counting the same pair twices,
                # the 2 resid cannot be neigbors
                if i < j and aj_resid - ai_resid >= nres_away:
                    d = np.linalg.norm(ai.pos - aj.pos)
                    if d <= cutoff:
                        # -1: resid in MDAnalysis starts from 1
                        map_[ai_resid-1][aj_resid-1] += 1
        if unun_map is None:
            unun_map = map_
        else:
            unun_map = unun_map + map_
        utils.print_progress(ts)
    sys.stdout.write("\n")
    return unun_map / float(nframe)

def main(struct):
    u = Universe(struct)

    phi = u.selectAtoms(PHI_SEL)
    psi = u.selectAtoms(PSI_SEL)
    
    print u.filename
    print 'phi: {0:8.2f}'.format(phi.dihedral())
    print 'psi: {0:8.2f}'.format(psi.dihedral())
    print 

 def test_atomgroups(self):
     u = Universe(self.filename)
     segidB0 = len(u.select_atoms("segid B and (not altloc B)"))
     segidB1 = len(u.select_atoms("segid B and (not altloc A)"))
     assert_equal(segidB0, segidB1)
     altlocB0 = len(u.select_atoms("segid B and (altloc A)"))
     altlocB1 = len(u.select_atoms("segid B and (altloc B)"))
     assert_equal(altlocB0, altlocB1)
     sum = len(u.select_atoms("segid B"))
     assert_equal(sum, segidB0 + altlocB0)

def calc_rg(grof, xtcf, btime, debug):
    u = Universe(grof, xtcf)
    query = 'name CA'
    # MDAnalysis will convert the unit of length to angstrom, though in Gromacs the unit is nm
    atoms = u.selectAtoms(query)
    natoms = atoms.numberOfAtoms()
    for ts in u.trajectory:
        if ts.time >= btime:
            com = atoms.centerOfMass()                                # center of mass
            _sum = sum((sum(i**2 for i in (a.pos - com)) for a in atoms))
            rg = np.sqrt(_sum / natoms)
            yield '{0:10.0f}{1:15.6f}\n'.format(ts.time, rg)
        # per 100 frames, num of frames changes with the size of xtc file, for debugging
        if debug and ts.frame % 2 == 0: 
            print "time: {0:10.0f}; step: {1:10d}; frame: {2:10d}".format(ts.time, ts.step, ts.frame)

 def __setstate__(self, dict):
     self.__dict__.update(dict)
     #reconstruct the universe
     self.u = Universe(dict['structure_filename'], dict['trajectory_filename'])
     apply_mass_map(self.u, dict['mass_map'])
     self.u.trajectory.periodic = dict['trajectory_is_periodic']
     for f in self.tar_forces + self.ref_forces:
         cat = f.get_category()
         if(not (cat is None)):
             self.ref_cats.append(cat)
         f.setup_hook(self.u)

def sequence_spacing(grof, xtcf, btime, etime, peptide_length, atom_sel):
    u = Universe(grof, xtcf)
    # this selection part should be better customized
    # here, only have been backbone atoms are used, u.selectAtoms doesn't
    # include Hydrogen atoms
    # REMMEMBER: ARGS verification should be done in main ONLY!
    # range works like this:

    # in MDAnalysis, resid starts from 1, in sequence_spacing.py, we don't count
    # the C- and N- termini, so it's from 2 to peptide_len+2
    residues = [u.selectAtoms(atom_sel.format(i)) for i in range(2, peptide_length + 2)]
    ijdist_dict = {}
    for ts in u.trajectory:
        # btime, etime defaults to 0, if etime is 0, loop till the end of the
        # trajectory
        if btime > ts.time:
            continue
        if etime > 0 and etime < ts.time:
            break

        # the good stuff
        for i, resi in enumerate(residues):
            for j, resj in enumerate(residues):
                # to remove duplicate since resi & resj are within the same peptide 
                if i < j:
                    dij = abs(i - j)
                    d_atomi_atomj = []
                    # loop through every atom in both residues
                    for atomi in resi:
                        for atomj in resj:
                            d_atomi_atomj.append(
                                np.linalg.norm(atomi.pos - atomj.pos))
                # add the result to the dictionary
                    ij_dist = np.average(d_atomi_atomj)   # distance between i and j
                    if dij not in ijdist_dict.keys():
                        ijdist_dict[dij] = [ij_dist]
                    else:
                        ijdist_dict[dij].append(ij_dist)
        utils.print_progress(ts)

    return ijdist_dict

def getWaterCoorWithH(self,centre,psf,dcd,outputFile):
    rho=Universe(psf,dcd)
    H2OCoordinate=[]
    no=0
    title='resname'+'    '+'atomid'+'    '+'resnumber'+'    X    Y     Z   '+'   '+'segname'+'  '+'frameNo'+'   '+'centreNo'+'\n'
    outputFile.write(title)
    for oxygenInforSet in self:
        H2OCoordinateSet=[]
        
        
        print 'There were',len(oxygenInforSet),'waters in the'
        for oxygenInfor in oxygenInforSet:
##            no1+=1
##            print no1
            frameNo=oxygenInfor[-2]
            frameNo=int(frameNo)-1
            segName=oxygenInfor[-3]
            resNumber=oxygenInfor[2]
            frame=rho.trajectory[frameNo]
            infor='segid '+segName+' and resid '+resNumber
            selected=rho.selectAtoms(infor)
            atomID=[]
            for atoms in selected.atoms:
                ID=str(atoms).split()[2][:-1]
                atomID.append(ID)
            selectedResId=selected.resids()
            selectedResNa=selected.resnames()
            coordsOH1H2=selected.coordinates()
            for i in range(3):
                atomInfor=str(selectedResNa[0])+'    '+str(atomID[i])+'    '+str(resNumber)+'    '+str(coordsOH1H2[i])[1:-1]+'   '+segName+'    '+str(frameNo)+'    '+str(no)+'\n'
                outputFile.write(atomInfor)
            H2OCoordinateSet.append(coordsOH1H2)

        no+=1

        H2OCoordinate.append(H2OCoordinateSet)
        print no,'is finished'
    outputFile.close()
    return H2OCoordinate

def process_trajectory(args):
	start, end = args.region.split('-')
	plane = [int(x) for x in args.plane.split(',')]
	univ = Universe(args.topo,args.traj)
	results = []
	for ts in univ.trajectory:
		norm = get_normal(univ, atoms=plane)
		temp = get_region(univ, start, end)
		if norm is not None and temp is not None and len(temp) > 0:
			angle = get_vector(temp)
			name = univ.split('/')[-1]
			if angle is not None:
				results.append(dot(norm, angle))
	if args.graph and len(results) > 0:
		import matplotlib.pyplot as plt
		ysort = sorted(list(enumerate(results,start=1)), key=lambda kv: float(kv[1]))
		x, y = zip(*ysort)
		ind = arange(len(x))
		plt.plot(y, ind, color='r')
		plt.yticks(ind, x)
		print ysort
		plt.show()
	return