Python mdtraj.load函数代码示例

def test_load():
    filenames = [
        "frame0.xtc",
        "frame0.trr",
        "frame0.dcd",
        "frame0.binpos",
        "traj.h5",
        "frame0.nc",
        "traj.h5",
        "frame0.lammpstrj",
        "frame0.xyz",
    ]
    num_block = 3
    for filename in filenames:
        t0 = md.load(get_fn(filename), top=nat, discard_overlapping_frames=True)
        t1 = md.load(get_fn(filename), top=nat, discard_overlapping_frames=False)
        t2 = md.load([get_fn(filename) for i in xrange(num_block)], top=nat, discard_overlapping_frames=False)
        t3 = md.load([get_fn(filename) for i in xrange(num_block)], top=nat, discard_overlapping_frames=True)

        # these don't actually overlap, so discard_overlapping_frames should
        # have no effect. the overlap is between the last frame of one and the
        # first frame of the next.
        yield lambda: eq(t0.n_frames, t1.n_frames)
        yield lambda: eq(t0.n_frames * num_block, t2.n_frames)
        yield lambda: eq(t3.n_frames, t2.n_frames)

            def check():
                out1 = md.load(os.path.join(output_dir, fn2), **load_kwargs_check1)
                out2 = md.load(os.path.join(output_dir, 'subset.' + fn2), **load_kwargs_check2)
                out3 = md.load(os.path.join(output_dir, 'stride.' + fn2), **load_kwargs_check1)

                if ext1 in ['.lh5'] or ext2 in ['.lh5']:
                    decimal = 3
                else:
                    decimal = 6
                eq(out1.xyz, TRAJ.xyz, decimal=decimal)
                eq(out2.xyz, TRAJ.xyz[:, atom_indices], decimal=decimal)
                eq(out3.xyz, TRAJ.xyz[::3], decimal=decimal)

                if ext1 not in ['.binpos', '.lh5'] and ext2 not in ['.binpos', '.lh5']:
                    # binpos doesn't save unitcell information
                    eq(out1.unitcell_vectors, TRAJ.unitcell_vectors, decimal=2)
                    eq(out2.unitcell_vectors, TRAJ.unitcell_vectors, decimal=2)
                    eq(out3.unitcell_vectors, TRAJ.unitcell_vectors[::3], decimal=2)

                if all(e in ['.xtc', '.trr', '.nc', '.h5'] for e in [ext1, ext2]):
                    # these formats contain time information
                    eq(out1.time, TRAJ.time)
                    eq(out2.time, TRAJ.time)
                    eq(out3.time, TRAJ.time[::3])

                if ext2 in ['.pdb', '.h5', '.lh5']:
                    # these formats contain a topology in the file that was
                    # read from disk
                    eq(out1.topology, TRAJ.topology)
                    eq(out2.topology, TRAJ.topology.subset(atom_indices))
                    eq(out3.topology, TRAJ.topology)

def calc_obs(traj):
    arg_cz_id = 2442
    glu_cd_id = 862
    lys_nz_id = 634
    tyr_oh_id = 2019
    inactive = mdt.load("./topologies/inactive.pdb")
    active = mdt.load("./topologies/active.pdb")

    aloop_atoms_list = [i.index for residue in np.arange(147, 168) for i in inactive.topology.residue(residue).atoms]
    all_heavy = [i.index for i in inactive.topology.atoms if i.residue.is_protein and i.element.name != "hydrogen"]
    print("Processing %s" % traj)
    # load the trajectory
    trj = mdt.load(traj, atom_indices=np.arange(inactive.n_atoms))

    inactive_rms = mdt.rmsd(trj, inactive, atom_indices=all_heavy)
    active_rms = mdt.rmsd(trj, active, atom_indices=all_heavy)
    aloop_rms = mdt.rmsd(trj, inactive, frame=0, atom_indices=aloop_atoms_list)
    distances = mdt.compute_distances(trj, np.vstack(([arg_cz_id, glu_cd_id], [lys_nz_id, glu_cd_id])))
    return dict(
        fname=os.path.basename(traj),
        inactive_rmsd=inactive_rms,
        active_rmsd=active_rms,
        aloop_inactive_rmsd=aloop_rms,
        glu_arg=distances[:, 0],
        gly_lys=distances[:, 1],
    )

def calculate_rmsd(trajectory, topology, reference):
    import mdtraj
    traj = mdtraj.load(trajectory, top=topology)
    ref = mdtraj.load(reference)
    rmsd = mdtraj.rmsd(traj, ref)
    data = {"step": str(traj.n_frames), "rmsd": str(rmsd[-1])}
    return data

def get_J3_HN_HA(traj, top, frame=None,  model="Habeck", outname = None):
    '''Compute J3_HN_HA for frames in a trajectories.
    Parameters
    ----------
    traj: trajectory file
    top: topology file
    frame: specific frame for computing
    model: Karplus coefficient models ["Ruterjans1999","Bax2007","Bax1997","Habeck" ,"Vuister","Pardi"]
    outname: if not None, the output will be saved and a file name (in the format of string) is required.
    '''
    J=[]
    if frame is None:
            t = md.load(traj,top=top)
            J = compute_J3_HN_HA(t, model = model)
    if frame is not None:
            for i in range(len(frame)):
                t = md.load(traj,top=top)[frame[i]]
                d = compute_J3_HN_HA(t, model = model)
                if i == 0:
                        J.append(d[0])
                        J.append(d[1])
                else:
                        J.append(d[1])
    if outname is not None:
            print('saving output file...')
            np.save(outname, J)
            print('Done!')

    return J

def fetch_fs_peptide(data_home=None, download_if_missing=True):
    """Loader for the Fs peptide dataset

    Parameters
    ----------
    data_home : optional, default: None
        Specify another download and cache folder for the datasets. By default
        all mixtape data is stored in '~/mixtape_data' subfolders.

    download_if_missing: optional, True by default
        If False, raise a IOError if the data is not locally available
        instead of trying to download the data from the source site.
    """
    data_home = get_data_home(data_home=data_home)
    if not exists(data_home):
        makedirs(data_home)

    data_dir = join(data_home, TARGET_DIRECTORY)
    if not exists(data_dir):
        print('downloading fs peptide from %s to %s' % (DATA_URL, data_home))
        fhandle = urlopen(DATA_URL)
        buf = BytesIO(fhandle.read())
        zip_file = ZipFile(buf)
        makedirs(data_dir)
        for name in zip_file.namelist():
            zip_file.extract(name, path=data_dir)

    top = md.load(join(data_dir, 'fs_peptide.pdb'))
    trajectories = []
    for fn in sorted(glob(join(data_dir, 'trajectory*.xtc'))):
        print('loading %s...' % basename(fn))
        trajectories.append(md.load(fn, top=top))

    return Bunch(trajectories=trajectories, DESCR=__doc__)

def test_vsite_elements(get_fn):
    #  Test case for issue #265
    pdb_filename = get_fn('GG-tip4pew.pdb')
    trj = md.load(pdb_filename)
    trj.save_hdf5(temp)
    
    trj2 = md.load(temp, top=pdb_filename)

def MDTRAJwrite(mol, filename):
    try:
        import mdtraj as md
        from htmd.util import tempname
        ext = os.path.splitext(filename)[1][1:]
        if ext == 'gz':
            pieces = filename.split('.')
            ext = '{}.{}'.format(pieces[-2], pieces[-1])

        if ext in _MDTRAJ_TOPOLOGY_SAVERS:
            tmppdb = tempname(suffix='.pdb')
            mol.write(tmppdb)
            traj = md.load(tmppdb)
            os.remove(tmppdb)
        elif ext in _MDTRAJ_TRAJECTORY_SAVERS:
            tmppdb = tempname(suffix='.pdb')
            tmpxtc = tempname(suffix='.xtc')
            mol.write(tmppdb)
            mol.write(tmpxtc)
            traj = md.load(tmpxtc, top=tmppdb)
            os.remove(tmppdb)
            os.remove(tmpxtc)
        else:
            raise ValueError('Unknown file type for file {}'.format(filename))
        # traj.xyz = np.swapaxes(np.swapaxes(self.coords, 1, 2), 0, 1) / 10
        # traj.time = self.time
        # traj.unitcell_lengths = self.box.T / 10
        traj.save(filename)
    except Exception as e:
        raise ValueError('MDtraj reader failed for file {} with error "{}"'.format(filename, e))

def main():
    p = args.pro
    ppref = p.split('/')[-1].split('.')[0]
    l = args.lig
    lpref = l.split('.')[0]

    if args.clean_lig:
        mdtraj.load(l).save_pdb('tidy_' + l)

    comb = sr + 'clc/dock/analysis/comb_pl.py ' + p + ' ' + l
    call_chimera(comb)

    # keep remarks w binding info
    fn = ppref + '_' + lpref + '.pdb'
    get_remarks = 'grep REMARK ' + l + '>> ' + fn
    call_cl(get_remarks)

    ref = args.ref
    if ref != None:
        import numpy as np
        # read in combined pdb
        lp = mdtraj.load(fn)
        # find distance between reference residue and ligand
        pi = lp.topology.select('resid ' + ref + ' and name C')
        li = lp.topology.select('not protein and name S')
        lq = lp.atom_slice(li).xyz[0][0]
        pq = lp.atom_slice(pi).xyz[0][0]
        dist = np.linalg.norm(pq-lq)
        # log results
        with open('dists.dat', 'a') as f:
            f.write(lpref + ': ' + str(dist) + '\n')
        f.close()
        # remove comb pdb if beyond cutoff distance
        if dist > args.cut:
            os.remove(fn)

def main(opts):
    print 'Loading atom indices file for trajectories', opts.ndx
    ndx = np.loadtxt(opts.ndx, dtype=np.int)

    print 'Loading cells from', opts.cells
    cells = mdtraj.load(opts.topol, atom_indices=ndx)
    cells.xyz = load_cells_gps(opts.cells)

    print 'Loading trajectories', ' '.join(opts.trajs)
    traj = mdtraj.load(opts.trajs, top=opts.topol, atom_indices=ndx)

    print 'Assigning to {} cells'.format(len(cells))
    rmsds = -np.ones((len(cells), len(traj)))
    for i in xrange(len(cells)):
        rmsds[i] = mdtraj.rmsd(traj, cells, frame=i)
    rmsds = rmsds.T
    A = -np.ones((len(traj),), dtype=np.int)
    for f in xrange(len(traj)):
        A[f] = rmsds[f].argmin()

    np.savetxt(opts.assignments, A, fmt='%d')

    print 'Computing populations'
    P = np.bincount(A)
    np.savetxt(opts.populations, P, fmt='%d')

def test_load_combination(ref_traj, get_fn):
    # Test that the load function's stride and atom_indices work across
    # all trajectory formats

    topology = md.load(get_fn('native.pdb')).topology
    ainds = np.array([a.index for a in topology.atoms if a.element.symbol == 'C'])

    no_kwargs = md.load(get_fn(ref_traj.fn), top=topology)
    strided3 = md.load(get_fn(ref_traj.fn), top=topology, stride=3)
    subset = md.load(get_fn(ref_traj.fn), top=topology, atom_indices=ainds)

    # test 1
    t1 = no_kwargs
    t2 = strided3
    assert eq(t1.xyz[::3], t2.xyz)
    assert eq(t1.time[::3], t2.time)
    if t1.unitcell_vectors is not None:
        assert eq(t1.unitcell_vectors[::3], t2.unitcell_vectors)
    assert eq(t1.topology, t2.topology)

    # test 2
    t1 = no_kwargs
    t2 = subset
    assert eq(t1.xyz[:, ainds, :], t2.xyz)
    assert eq(t1.time, t2.time)
    if t1.unitcell_vectors is not None:
        assert eq(t1.unitcell_vectors, t2.unitcell_vectors)
    assert eq(t1.topology.subset(ainds), t2.topology)

def test_pdbwrite(get_fn):
    pdb = get_fn('native.pdb')
    p = load(pdb)
    p.save(temp)

    r = load(temp)
    eq(p.xyz, r.xyz)

def deprecated_test_fah_core17_1():
    from mdtraj.utils import six
    from mdtraj.testing import get_fn, eq
    filename = get_fn('frame0.xtc')
    tempdir = tempfile.mkdtemp()
    tar_filename = os.path.join(tempdir, "results-000.tar.bz2")
    archive = tarfile.open(tar_filename, mode='w:bz2')

    tar = tarfile.open(tar_filename, "w:bz2")
    tar.add(filename, arcname="positions.xtc")
    tar.close()

    shutil.copy(tar_filename, os.path.join(tempdir, "results-001.tar.bz2"))

    trj0 = md.load(get_fn("frame0.xtc"), top=get_fn("frame0.h5"))
    output_filename = os.path.join(tempdir, "traj.h5")
    fah.concatenate_core17(tempdir, trj0, output_filename)

    trj = md.load(output_filename)
    eq(trj.n_atoms, trj0.n_atoms)
    eq(trj.n_frames, trj0.n_frames * 2)

    shutil.copy(tar_filename, os.path.join(tempdir, "results-002.tar.bz2"))

    fah.concatenate_core17(tempdir, trj0, output_filename)
    # Should notice the new file and append it to the HDF file.

    trj = md.load(output_filename)
    eq(trj.n_atoms, trj0.n_atoms)
    eq(trj.n_frames, trj0.n_frames * 3)

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 _trj_load(file, top):
    if os.path.isdir(file):
        return mdt.load(os.path.join(file,"positions.xtc"),top=top)
    else:
        os.system("tar -xvjf %s"%file)
        return mdt.load("positions.xtc", top=top)
    return