How to use the mdtraj.testing.get_fn function in mdtraj

HMM.transmat_ = np.array([[0.9, 0.1, 0.0, 0.0],
                              [0.1, 0.7, 0.2, 0.0],
                              [0.0, 0.1, 0.8, 0.1],
                              [0.0, 0.1, 0.1, 0.8]])
    HMM.means_ = np.array([[-10, -10, -10],
                           [-5,  -5,   -5],
                           [5,    5,    5],
                           [10,  10,   10]])
    HMM.covars_ = np.array([[0.1, 0.1, 0.1],
                            [0.5, 0.5, 0.5],
                            [1, 1, 1],
                            [4, 4, 4]])
    HMM.startprob_ = np.array([1, 1, 1, 1]) / 4.0

    # get a 1 atom topology
    topology = md.load(get_mdtraj_fn('native.pdb')).restrict_atoms([1]).topology

    # generate the trajectories and save them to disk
    for i in range(10):
        d, s = HMM.sample(100)
        t = md.Trajectory(xyz=d.reshape(len(d), 1, 3), topology=topology)
        t.save(os.path.join(DATADIR, 'Trajectory%d.h5' % i))

"""
Test the cython dcd module

Note, this file cannot be located in the dcd subdirectory, because that
directory is not a python package (it has no __init__.py) and is thus tests
there are not discovered by nose
"""

import tempfile, os
import numpy as np
from mdtraj import dcd, io
from mdtraj.testing import get_fn, eq
import warnings

fn_dcd = get_fn('frame0.dcd')
pdb = get_fn('native.pdb')

temp = tempfile.mkstemp(suffix='.dcd')[1]
def teardown_module(module):
    """remove the temporary file created by tests in this file 
    this gets automatically called by nose"""
    os.unlink(temp)

def test_read():
    xyz = dcd.read_xyz(fn_dcd)
    xyz2 = io.loadh(get_fn('frame0.dcd.h5'), 'xyz')

    eq(xyz, xyz2)
    
def test_write_0():
    xyz = dcd.read_xyz(fn_dcd)

def test_find_cluster_centers():

    N_TRJS = 20
    many_trjs = [md.load(get_fn('frame0.xtc'), top=get_fn('native.pdb'))
                 for i in range(N_TRJS)]

    distances = np.ones((N_TRJS, len(many_trjs[0])))

    center_inds = [(0, 0), (5, 2), (15, 300)]

    for ind in center_inds:
        distances[center_inds[0], center_inds[1]] = 0

    centers = find_cluster_centers(many_trjs, distances)

    # we should get back the same number of centers as there are zeroes
    # in the distances matrix
    assert_equal(len(centers), np.count_nonzero(distances == 0))

    for indx in center_inds:

def test_reassign_script():

    topologies = [get_fn('native.pdb'), get_fn('native.pdb')]
    trajectories = [get_fn('frame0.xtc'), get_fn('frame0.xtc')]

    centers = [c for c in md.load(trajectories[0], top=topologies[0])[::50]]

    with tempfile.NamedTemporaryFile(suffix='.pkl') as ctrs_f:
        pickle.dump(centers, ctrs_f)
        ctrs_f.flush()

        runhelper(
            ['--centers', ctrs_f.name,
             '--trajectories', trajectories,
             '--topology', topologies])

def test_reassign_script_multitop():

    xtc2 = os.path.join(cards.__path__[0], 'test_data', 'trj0.xtc')
    top2 = os.path.join(cards.__path__[0], 'test_data', 'PROT_only.pdb')

    topologies = [get_fn('native.pdb'), top2]
    trajectories = [
        [get_fn('frame0.xtc'), get_fn('frame0.xtc')],
        [xtc2, xtc2]]

    atoms = '(name N or name C or name CA or name H or name O) and (residue 2)'
    centers = [c for c in md.load(trajectories[0], top=topologies[0])[::50]]

    with tempfile.NamedTemporaryFile(suffix='.pkl') as ctrs_f:
        pickle.dump(centers, ctrs_f)
        ctrs_f.flush()

        runhelper(
            ['--centers', ctrs_f.name,
             '--trajectories', trajectories[0],
             '--topology', topologies[0],
             '--trajectories', trajectories[1],
             '--topology', topologies[1],
             '--atoms', atoms])

def test_featurizer():
    fn = get_mdtraj_fn('1bpi.pdb')
    with tempdir():
        shell('hmsm atomindices -o alpha.dat --alpha -a -p %s' % fn)
        shell('hmsm featurizer --top %s -o alpha.pickl -a alpha.dat' % fn)
        f = np.load('alpha.pickl')
        eq(f.atom_indices, np.loadtxt('alpha.dat', int))

    with tempdir():
        shell('hmsm atomindices -o alphapairs.dat --alpha -d -p %s' % fn)
        shell('hmsm featurizer --top %s -o alpha.pickl -d alphapairs.dat' % fn)
        f = np.load('alpha.pickl')
        eq(f.pair_indices, np.loadtxt('alphapairs.dat', int))

def test_rmsd_cluster_multitop():

    expected_size = (3, (501, 501, 5001))

    # trj is length 5001
    xtc2 = os.path.join(cards.__path__[0], 'test_data', 'trj0.xtc')
    top2 = os.path.join(cards.__path__[0], 'test_data', 'PROT_only.pdb')

    runhelper([
        '--trajectories', get_fn('frame0.xtc'), get_fn('frame0.xtc'),
        '--trajectories', xtc2,
        '--topology', get_fn('native.pdb'),
        '--topology', top2,
        '--atoms', '(name N or name C or name CA or name H or name O) '
                   'and (residue 2)',
        '--rmsd-cutoff', '0.1',
        '--algorithm', 'khybrid'],
        expected_size=expected_size)

def test_reassignment_function_heterogenous():

    xtc2 = os.path.join(cards.__path__[0], 'test_data', 'trj0.xtc')
    top2 = os.path.join(cards.__path__[0], 'test_data', 'PROT_only.pdb')

    topologies = [get_fn('native.pdb'), top2]
    trajectories = [
        [get_fn('frame0.xtc'), get_fn('frame0.xtc')],
        [xtc2, xtc2]]

    atoms = '(name N or name C or name CA or name H or name O) and (residue 2)'
    top = md.load(topologies[0]).top
    centers = [c.atom_slice(top.select(atoms)) for c
               in md.load(trajectories[0][0], top=topologies[0])[::50]]

    assigns, dists = reassign.reassign(
        topologies, trajectories, atoms, centers)

    assert_is(type(assigns), ra.RaggedArray)
    assert_array_equal(assigns.lengths, [501, 501, 5001, 5001])
    assert_equal(len(assigns), 4)

def test_reassign_script():

    topologies = [get_fn('native.pdb'), get_fn('native.pdb')]
    trajectories = [get_fn('frame0.xtc'), get_fn('frame0.xtc')]

    centers = [c for c in md.load(trajectories[0], top=topologies[0])[::50]]

    with tempfile.NamedTemporaryFile(suffix='.pkl') as ctrs_f:
        pickle.dump(centers, ctrs_f)
        ctrs_f.flush()

        runhelper(
            ['--centers', ctrs_f.name,
             '--trajectories', trajectories,
             '--topology', topologies])

"""Rapid calculation of the RMSD drift of a simulation.
"""

import mdtraj as md

# Find two files that are distributed with MDTraj for testing purposes -- 
# we can us them to make our plot

import mdtraj.testing
crystal_fn = mdtraj.testing.get_fn('native.pdb')
trajectory_fn = mdtraj.testing.get_fn('frame0.xtc')

# Load up the trajectories from disk

crystal = md.load(crystal_fn)
trajectory = md.load(trajectory_fn, top=crystal)  # load the xtc. the crystal structure defines the topology
print trajectory


# RMSD with exchangeable hydrogen atoms is generally not a good idea
# so let's take a look at just the heavy atoms

rmsds_to_crystal = md.rmsd(trajectory, crystal, 0)
heavy_atoms = [atom.index for atom in crystal.topology.atoms if atom.element.symbol != 'H']
heavy_rmds_to_crystal = md.rmsd(trajectory, crystal, 0, atom_indices=heavy_atoms)

# Plot the results