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!')

 def test_g_GetRandomConfs(self):
     P1 = Project.load_from(ProjectFn)
     Assignments = io.loadh("Data/Assignments.Fixed.h5", 'arr_0')
     
     # make a predictable stream of random numbers by seeding the RNG with 42
     random_source = np.random.RandomState(42)
     randomconfs = GetRandomConfs.run(P1, Assignments, NumRandomConformations, random_source)
     
     reference = Trajectory.load_trajectory_file(os.path.join(ReferenceDir, "2RandomConfs.lh5"))
     self.assert_trajectories_equal(reference, randomconfs)

    def _load_traj(self, file_list):
        """
        Load a set of xtc or dcd files as a single trajectory

        Note that the ordering of `file_list` is relevant, as the trajectories
        are catted together.

        Returns
        -------
        traj : msmbuilder.Trajectory
        """

        if self.input_traj_ext == '.xtc':
            traj = Trajectory.load_from_xtc(file_list, PDBFilename=self.conf_filename,
                        discard_overlapping_frames=True)
        elif self.input_traj_ext == '.dcd':
            traj = Trajectory.load_from_xtc(file_list, PDBFilename=self.conf_filename)
        else:
            raise ValueError()
        return traj

def main(modeldir, start, type):
    start=int(start)
    data=dict()
    project=Project.load_from('%s/ProjectInfo.yaml' % modeldir.split('Data')[0])
    files=glob.glob('%s/fkbp*xtal.pdb' % modeldir.split('Data')[0])
    pdb=files[0]
    unbound=numpy.loadtxt('%s/tpt-%s/unbound_%s_states.txt' % (modeldir, type, type), dtype=int)
    T=mmread('%s/tProb.mtx' % modeldir)
    startstate=unbound[start]
    ass=io.loadh('%s/Assignments.Fixed.h5' % modeldir)

    steps=100000
    print "on start state %s" % startstate
    if os.path.exists('%s/tpt-%s/movie_state%s_1millisec.states.dat' % (modeldir, type, startstate)):
        print "loading from states"
        traj=numpy.loadtxt('%s/tpt-%s/movie_state%s_1millisec.states.dat' % (modeldir, type, startstate))
    else:
        traj=msm_analysis.sample(T, int(startstate),int(steps))
        numpy.savetxt('%s/tpt-%s/movie_state%s_1millisec.states.dat' % (modeldir, type, startstate), traj)
    print "checking for chkpt file"
    checkfile=glob.glob('%s/tpt-%s/movie_state%s_*chkpt' % (modeldir, type, startstate))
    if len(checkfile) > 0:
        movie=Trajectory.load_from_xtc(checkfile[0], PDBFilename=pdb)
        n=int(checkfile[0].split('xtc.state')[1].split('chkpt')[0])
        os.system('mv %s %s.chkpt.cp' % (checkfile[0], checkfile[0].split('.xtc')[0]))
        print "checkpointing at state index %s out of %s" % (n, len(traj))
        checkfile=checkfile[0]
        restart=True
    else:
        restart=False
        n=0
        movie=project.empty_traj()
    while n < len(traj):
        print "on state %s" % n
        state=int(traj[n])
        t=project.get_random_confs_from_states(ass['arr_0'], [int(state),], 10)
        if n==0:
            movie['XYZList']=t[0]['XYZList']
            n+=1
            continue
        elif n % 100==0:
            movie['XYZList']=numpy.vstack((movie['XYZList'], t[0]['XYZList']))
            if restart==True:
                os.system('mv %s %s.chkpt.cp' % (checkfile, checkfile.split('.xtc')[0]))
            movie.save_to_xtc('%s/tpt-%s/movie_state%s_1millisec.xtc.state%schkpt' % (modeldir, type, startstate, n))
            checkfile='%s/tpt-%s/movie_state%s_1millisec.xtc.state%schkpt' % (modeldir, type, startstate, n)
            n+=1
            continue
        elif n!=0:
            movie['XYZList']=numpy.vstack((movie['XYZList'], t[0]['XYZList']))
            n+=1
            continue
    movie.save_to_xtc('%s/tpt-%s/movie_state%s_1millisec.xtc' % (modeldir, type, startstate))

def get_project_object( traj_directory, conf_filename, out_filename=None ):
    """
    This function constructs a msmbuilder.Project object 
    given a directory of trajectories saved as .lh5's. 

    Note that this is only really necessary when a script
    like ConvertDataToLHDF.py converts the data but fails
    to write out the ProjectInfo.yaml file.

    This function can also be used to combine two projects
    by copying and renaming the trajectories in a new 
    folder. Though, it's probably more efficient to just
    do some bash stuff to cat the ProjectInfo.yaml's 
    together and rename the trajectories.
    
    Inputs:
    -------
    1) traj_directory : directory to find the trajectories
    2) conf_filename : file to find the conformation
    3) out_filename [ None ] : if None, then this function 
        does not save the project file, but if given, the
        function will save the project file and also
        return the object

    Outputs:
    -------
    project : msmbuilder.Project object corresponding to 
        your project.
    """

    traj_paths = sorted( os.listdir( traj_directory ), key=keynat ) # relative to the traj_directory
    traj_paths = [ os.path.join( traj_directory, filename ) for filename in traj_paths ] # relative to current directory

    traj_lengths = []

    for traj_filename in traj_paths: # Get the length of each trajectory
        logger.info( traj_filename )
        traj_lengths.append( Trajectory.load_from_lhdf( traj_filename, JustInspect=True )[0] ) 
        # With JustInspect=True this just returns the shape of the XYZList

    project = Project({'conf_filename': conf_filename,
                       'traj_lengths': traj_lengths,
                       'traj_paths': traj_paths,
                       'traj_errors': [None] * len(traj_paths),
                       'traj_converted_from': [ [None] ] * len(traj_paths) })

    if out_filename is None:
        return project
    else:
        project.save( out_filename )
        logger.info('Saved project file to %s', out_filename)
        return project

def test_gpurmsd():
    traj = Trajectory.load_trajectory_file(trj_path)    

    gpurmsd = GPURMSD()
    ptraj = gpurmsd.prepare_trajectory(traj)
    gpurmsd._gpurmsd.print_params()
    gpu_distances = gpurmsd.one_to_all(ptraj, ptraj, 0)

    cpurmsd = RMSD()
    ptraj = cpurmsd.prepare_trajectory(traj)
    cpu_distances = cpurmsd.one_to_all(ptraj, ptraj, 0)
    
    npt.assert_array_almost_equal(cpu_distances, gpu_distances, decimal=4)