Python msmbuilder.Trajectory类代码示例

def run(project, atom_indices=None, traj_fn = 'all'):

    n_atoms = project.load_conf()['XYZList'].shape[1]

    if traj_fn.lower() == 'all':

        SASA = np.ones((project.n_trajs, np.max(project.traj_lengths), n_atoms)) * -1

        for traj_ind in xrange(project.n_trajs):
            traj_asa = []
            logger.info("Working on Trajectory %d", traj_ind)
            traj_fn = project.traj_filename(traj_ind)
            chunk_ind = 0
            for traj_chunk in Trajectory.enum_chunks_from_lhdf( traj_fn, AtomIndices=atom_indices ):
                #print chunk_ind
                traj_asa.extend(asa.calculate_asa(traj_chunk, n_sphere_points = 24))
                chunk_ind += 1
            SASA[traj_ind, 0:project.traj_lengths[traj_ind]] = traj_asa

    else:
        traj_asa = []
        for traj_chunk in Trajectory.enum_chunks_from_lhdf( traj_fn, AtomIndices=atom_indices ):
            traj_asa.extend( asa.calculate_asa( traj_chunk ) )

        SASA = np.array(traj_asa)

    return SASA

def main(modeldir, genfile,  type, write=False):
    data=dict()
    pops=numpy.loadtxt('%s/Populations.dat' % modeldir)
    map=numpy.loadtxt('%s/Mapping.dat' % modeldir)
    frames=numpy.where(map!=-1)[0]

    unbound=numpy.loadtxt('%s/tpt-rmsd-%s/unbound_%s_states.txt' % (modeldir, type, type), dtype=int)
    bound=numpy.loadtxt('%s/tpt-rmsd-%s/bound_%s_states.txt' % (modeldir, type, type), dtype=int)

    dir=modeldir.split('Data')[0]
    name=glob.glob('%s/fkbp*xtal*pdb' % dir)
    pdb=Trajectory.load_from_pdb(name[0])
    paths=io.loadh('%s/tpt-rmsd-%s/Paths.h5' % (modeldir, type))

    committors=numpy.loadtxt('%s/commitor_states.txt' % modeldir, dtype=int)
    colors=['red', 'orange', 'green', 'cyan', 'blue', 'purple']
    colors=colors*40
    if type=='strict':
        ref=5
    elif type=='super-strict':
        ref=3
    elif type=='medium':
        ref=10
    elif type=='loose':
        ref=15
    #for p in range(0, 3):
    for p in range(0, 1):
        path=paths['Paths'][p]
        print "Bottleneck", paths['Bottlenecks'][p]
        flux=paths['fluxes'][p]/paths['fluxes'][0]
        if flux < 0.2:
            break
        print "flux %s" % flux
        frames=numpy.where(path!=-1)[0]
        path=numpy.array(path[frames], dtype=int)
        print path
        if write==True:
            size=(paths['fluxes'][p]/paths['fluxes'][0])*1000
            traj=Trajectory.load_from_xtc('%s/tpt-rmsd-%s/path%s_sample20.xtc' % (modeldir, type, p), Conf=pdb)
            data=build_metric(dir, pdb, traj)
            dir=modeldir.split('Data')[0]
            for op in sorted(data.keys()):
            #for op in residues:
                pylab.figure()
                pylab.scatter(data['rmsd'], data[op], c=colors[p], alpha=0.7) #, s=size)
                for j in paths['Bottlenecks'][p]:
                    frame=numpy.where(paths['Paths'][p]==j)[0]
                    pylab.scatter(data['rmsd'][frame*20], data[op][frame*20], marker='x', c='k', alpha=0.7, s=50)
                    location=numpy.where(committors==paths['Paths'][p][frame])[0]
                    if location.size:
                        print "path %s state %s bottleneck in committors" % (p, j)
                        print data['rmsd'][frame*20], data[op][frame*20]
                pylab.title('path %s' % p)
                pylab.xlabel('P-L RMSD')
                #pylab.xlabel('P-L COM')
                pylab.ylabel(op)
                pylab.xlim(0,max(data['rmsd'])+5)
                #pylab.ylim(0,max(data[op])+5)
                pylab.show()

def plot_gpu_cmd_correlation():
    traj1 = Trajectory.load_trajectory_file(ww_1, Conf=ww_conf)
    traj1_copy = Trajectory.load_trajectory_file(ww_1, Conf=ww_conf)
    #traj2 = Trajectory.load_trajectory_file(ww_2, Conf=ww_conf)
    #traj2_copy = Trajectory.load_trajectory_file(ww_2, Conf=ww_conf)

    def gpudist(t):
        gpurmsd = GPURMSD()
        pt = gpurmsd.prepare_trajectory(t)
        gpurmsd._gpurmsd.print_params()
        return gpurmsd.one_to_all(pt, pt, 0)
    def cpudist(t):
        rmsd = RMSD()
        pt = rmsd.prepare_trajectory(t)
        return rmsd.one_to_all(pt, pt, 0)
    g1 = gpudist(traj1) #, gpudist(traj2)
    c1 = cpudist(traj1_copy) #, cpudist(traj2_copy)

    pp.subplot(231)
    pp.plot(c1)
    pp.title('cpu rmsd drift along traj')
    pp.xlabel('frame index')
    pp.xlabel('cpurmsd($X_{0}$, $X_{frame_index}$)')

    pp.subplot(232)
    pp.scatter(g1, c1)
    pp.xlabel('gpu rmsd')
    pp.ylabel('cpu rmsd')

    pp.subplot(233)
    pp.plot(g1)
    pp.title('gpu rmsd drift along traj')
    pp.xlabel('frame index')
    pp.xlabel('gpurmsd($X_{0}$, $X_{frame_index}$)')


    #PLOT c2 and g2 in the lower portion of the graph

    #pp.subplot(234)
    #pp.plot(c2)
    #pp.title('cpu rmsd drift along pre-aligned traj')
    #pp.xlabel('frame index')
    #pp.xlabel('cpurmsd($X_{0}$, $X_{frame_index}$)')

    #pp.subplot(235)
    #pp.scatter(g2, c2)
    #pp.xlabel('gpu rmsd')
    #pp.ylabel('cpu rmsd')

    #pp.subplot(236)
    #pp.plot(g2)
    #pp.title('gpu rmsd drift along pre-aligned traj')
    #pp.xlabel('frame index')
    #pp.xlabel('gpurmsd($X_{0}$, $X_{frame_index}$)')

    #pp.subplots_adjust(hspace=0.4)
    #pp.savefig('gpucpu_correlation.png')
    pp.show()

def test_xtc_dcd():
    pdb_filename = get("native.pdb", just_filename=True)
    xtc_filename = get('RUN00_frame0.xtc', just_filename=True)
    dcd_filename = get('RUN00_frame0.dcd', just_filename=True)
    r_xtc = Trajectory.load_from_xtc(xtc_filename, pdb_filename)
    r_dcd = Trajectory.load_from_dcd(dcd_filename, pdb_filename)

    x_xtc = r_xtc["XYZList"]
    x_dcd = r_dcd["XYZList"]

    eq(x_xtc, x_dcd, decimal=4)

def test_asa_3():

    traj_ref = np.loadtxt( os.path.join(reference_dir(),'g_sas_ref.dat'))
    Conf = Trajectory.load_from_pdb(os.path.join( fixtures_dir(), 'native.pdb'))

    traj = Trajectory.load_trajectory_file( os.path.join(fixtures_dir(), 'trj0.xtc') , Conf=Conf)
    traj_asa = calculate_asa(traj, probe_radius=0.14, n_sphere_points = 960)
    
    # the algorithm used by gromacs' g_sas is slightly different than the one
    # used here, so the results are not exactly the same -- see the comments
    # in src/python/geomtry/asa.py or the readme file src/ext/asa/README.txt
    # for details
    npt.assert_array_almost_equal(traj_asa, traj_ref, decimal=2)    

    def load_frame(self, traj_index, frame_index):
        """Load one or more specified frames.

        Example
        -------
        >>> project = Project.load_from('ProjectInfo.yaml')
        >>> foo = project.load_frame(1,10)
        >>> bar = Trajectory.read_frame(TrajFilename=project.traj_filename(1),
            WhichFrame=10)
        >>> np.all(foo['XYZList'] == bar)
        True

        Parameters
        ----------
        traj_index : int, [int]
            Index or indices of the trajectories to pull from
        frame_index : int, [int]
            Index or indices of the frames to pull from

        Returns
        -------
        traj : msmbuilder.Trajectory
            A trajectory object containing the requested frame(s).
        """

        if np.isscalar(traj_index) and np.isscalar(frame_index):
            xyz = Trajectory.read_frame(TrajFilename=self.traj_filename(traj_index),
                WhichFrame=frame_index)
            xyzlist = np.array([xyz])
        else:
            traj_index = np.array(traj_index)
            frame_index = np.array(frame_index)
            if not (traj_index.ndim == 1 and np.all(traj_index.shape == frame_index.shape)):
                raise ValueError('traj_index and frame_index must be 1D and have the same length')

            xyzlist = []
            for i,j in zip(traj_index, frame_index):
                if j >= self.traj_lengths[i]:
                    raise ValueError('traj %d too short (%d) to contain a frame %d' % (i, self.traj_lengths[i], j))
                xyz = Trajectory.read_frame(TrajFilename=self.traj_filename(i),
                    WhichFrame=j)
                xyzlist.append(xyz)
            xyzlist = np.array(xyzlist)

        conf = self.load_conf()
        conf['XYZList'] = xyzlist

        return conf

    def test(self):
        from msmbuilder.scripts.SaveStructures import save

        project = get('ProjectInfo.yaml')
        assignments = get('Assignments.h5')['arr_0']
        which_states = [0, 1, 2]
        list_of_trajs = project.get_random_confs_from_states(assignments, 
            which_states, num_confs=2, replacement=True,
            random=np.random.RandomState(42))

        assert isinstance(list_of_trajs, list)
        assert isinstance(list_of_trajs[0], Trajectory)
        eq(len(list_of_trajs), len(which_states))
        for t in list_of_trajs:
            eq(len(t), 2)

        print list_of_trajs[0].keys()
        # sep, tps, one
        save(list_of_trajs, which_states, style='sep', format='lh5', outdir=self.td)
        save(list_of_trajs, which_states, style='tps', format='lh5', outdir=self.td)
        save(list_of_trajs, which_states, style='one', format='lh5', outdir=self.td)

        names = ['State0-0.lh5', 'State0-1.lh5', 'State0.lh5', 'State1-0.lh5',
                'State1-1.lh5', 'State1.lh5', 'State2-0.lh5', 'State2-1.lh5',
                'State2.lh5']

        for name in names:
            t = Trajectory.load_trajectory_file(pjoin(self.td, name))
            eq(t, get('save_structures/' + name))

    def save(self):
        "Save the trajs as a n MSMBuilder project"
        
        traj_dir = pjoin(self.project_dir, 'Trajectories')
        if not os.path.exists(traj_dir):
            os.makedirs(traj_dir)

        t = Trajectory.load_trajectory_file(self.conf_filename)

        traj_paths = []
        for i, xyz in enumerate(self.trajectories):
            t['IndexList'] = None # bug in msmbuilder
            t['XYZList'] = xyz

            traj_paths.append(pjoin(traj_dir, 'trj%d.lh5' % i))
            t.save(traj_paths[-1])

        p = Project({'conf_filename': os.path.abspath(self.conf_filename),
            'traj_lengths': self.n_frames*np.ones(self.n_trajs),
            'traj_paths': [os.path.abspath(e) for e in traj_paths],
            'traj_converted_from': [[] for i in range(self.n_trajs)],
            'traj_errors': [None for i in range(self.n_trajs)],
            }, project_dir=self.project_dir, validate=True)
        p.save(pjoin(self.project_dir,'Project.yaml'))

        # just check again
        p = Project.load_from(pjoin(self.project_dir,'Project.yaml'))
        p._validate()
        assert np.all((p.load_traj(0)['XYZList'] - self.trajectories[0])**2 < 1e-6)

def test_traj_0():
    
    aind = np.unique( np.random.randint( 22, size=4) )
    stride = np.random.randint(1, 100 )
    
    r_traj = get('Trajectories/trj0.lh5')

    r_traj.restrict_atom_indices( aind )

    r_traj['XYZList'] = r_traj['XYZList'][ ::stride ]

    traj = Trajectory.load_from_lhdf(get('Trajectories/trj0.lh5', just_filename=True),
        Stride=stride, AtomIndices=aind)

    # make sure we loaded the right number of atoms
    assert traj['XYZList'].shape[1] == len(aind)

    for key in traj.keys():
        if key in ['SerializerFilename'] :
            continue
        
        if key in ['IndexList']:
            for row, r_row in zip( traj[key], r_traj[key] ):
                eq(row, r_row)
        elif key == 'XYZList':
            eq(traj[key], r_traj[key])
        else:
            eq(traj[key], r_traj[key])

    def __init__(self, structure_or_filename, metric, max_distance):
        """Create an explosion validator
    
        Checks the distance from every frame to a structure and
        watches for things that are too far away
        
        Parameters
        ----------
        structure_or_filename : {msmbuilder.Trajectory, str}
            The structure to measure distances to, either as a trajectory (the first
            frame is the only one that counts) or a path to a trajectory
            on disk that can be loaded
        metric : msmbuilder distance metric
            Metric by which you want to measure distance
        max_distance : float
            The threshold distance, above which a ValidationError
            will be thrown
        """

        if isinstance(structure_or_filename, Trajectory):
            conf = structure_or_filename
        elif isinstance(structure_or_filename, basestring):
            conf = Trajectory.load_trajectory_file(structure_or_filename)

        self.max_distance = max_distance
        self.metric = metric
        self._pconf = self.metric.prepare_trajectory(conf)

def load_gens(gens_fn, conf_fn, metric):
    """Setup a worker by adding pgens to its global namespace
    
    This is necessary because pgens are not necessarily picklable, so we can't
    just prepare them on the master and then push them to the remote workers --
    instead we want to actually load the pgens from disk and prepare them on
    the remote node
    """
    from msmbuilder import Trajectory
    
    global PGENS, CONF, METRIC, PREPARED
    
    METRIC = metric
    CONF = Trajectory.load_trajectory_file(conf_fn)
    gens = Trajectory.load_trajectory_file(gens_fn)
    PGENS = metric.prepare_trajectory(gens)
    PREPARED = True

    def test_c_Cluster(self):
        # We need to be sure to skip the stochastic k-mediods
        cmd = "Cluster.py -p {project} -s {stride} rmsd -a {atomindices} kcenters -d {rmsdcutoff}".format(project=ProjectFn, stride=Stride, atomindices="AtomIndices.dat", rmsdcutoff=RMSDCutoff)
        print cmd

        os.system(cmd)
        
        try:
            os.remove(os.path.join(WorkingDir, 'Data', 'Assignments.h5'))
            os.remove(os.path.join(WorkingDir, 'Data', 'Assignments.h5.distances'))
        except:
            pass

        
        G   = Trajectory.load_trajectory_file(GensPath)
        r_G = Trajectory.load_trajectory_file(ReferenceDir +'/'+ GensPath)
        self.assert_trajectories_equal(G, r_G)

def test_asa_2():
    t = Trajectory.load_trajectory_file(os.path.join(fixtures_dir(), 'trj0.lh5'))
    val1 = np.sum(calculate_asa(t[0])) # calculate only frame 0
    val2 = np.sum(calculate_asa(t)[0]) # calculate on all frames
    true_frame_0_asa = 2.859646797180176
    
    npt.assert_approx_equal(true_frame_0_asa, val1)
    npt.assert_approx_equal(true_frame_0_asa, val2)

 def _eval_traj_shapes(self):
     lengths = np.zeros(self.n_trajs)
     n_atoms = np.zeros(self.n_trajs)
     conf = self.load_conf()
     for i in xrange(self.n_trajs):
         shape = Trajectory.load_trajectory_file(self.traj_filename(i), JustInspect=True, Conf=conf)
         lengths[i] = shape[0]
         n_atoms[i] = shape[1]
     return lengths, n_atoms

def main(args, metric):
    assignments_path = os.path.join(args.output_dir, "Assignments.h5")
    distances_path = os.path.join(args.output_dir, "Assignments.h5.distances")
    project = Project.load_from(args.project)
    gens = Trajectory.load_trajectory_file(args.generators)
    
    # this runs assignment and prints them to disk
    assign_with_checkpoint(metric, project, gens, assignments_path, distances_path)

    logger.info('All Done!')