How to use mdtraj - 10 common examples

def test_atom_indices_1(get_fn):
    atom_indices = np.arange(10)
    top = md.load(get_fn('native.pdb'))
    t0 = md.load(get_fn('frame0.mdcrd'), top=top)
    t1 = md.load(get_fn('frame0.mdcrd'), top=top, atom_indices=atom_indices)

    eq(t0.xyz[:, atom_indices], t1.xyz)

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

def test_molecules(get_fn):
    top = md.load(get_fn('4OH9.pdb')).topology
    molecules = top.find_molecules()
    assert sum(len(mol) for mol in molecules) == top.n_atoms
    assert sum(1 for mol in molecules if len(mol) > 1) == 2  # All but two molecules are water

def test_trajectory_save_load(get_fn):
    t = md.load(get_fn('native.pdb'))
    t.unitcell_lengths = 1 * np.ones((1, 3))
    t.unitcell_angles = 90 * np.ones((1, 3))

    t.save(temp)
    t2 = md.load(temp, top=t.topology)

    eq(t.xyz, t2.xyz)
    eq(t.unitcell_lengths, t2.unitcell_lengths)

def test_4():
    # using a really big box, we should get the same results with and without
    # pbcs
    box = np.array([[100, 0, 0], [0, 200, 0], [0, 0, 300]])
    box = np.zeros((N_FRAMES, 3, 3)) + box  # broadcast it out
    a = _displacement_mic(xyz, pairs, box, False)
    b = _displacement(xyz, pairs)
    eq(a, b, decimal=3)

def test_multiread(get_fn):
    reference = md.load(get_fn('frame0.mdcrd'), top=get_fn('native.pdb'))
    with MDCRDTrajectoryFile(get_fn('frame0.mdcrd'), n_atoms=22) as f:
        xyz0, box0 = f.read(n_frames=1)
        xyz1, box1 = f.read(n_frames=1)

    eq(reference.xyz[0], xyz0[0] / 10)
    eq(reference.xyz[1], xyz1[0] / 10)

def test_read_0(get_fn):
    with XYZTrajectoryFile(get_fn('frame0.xyz')) as f:
        xyz = f.read()
    with XYZTrajectoryFile(get_fn('frame0.xyz')) as f:
        xyz3 = f.read(stride=3)
    eq(xyz[::3], xyz3)

def test_box(get_fn):
    t = md.load(get_fn('native.pdb'))
    assert eq(t.unitcell_vectors, None)
    assert eq(t.unitcell_lengths, None)
    assert eq(t.unitcell_angles, None)
    assert eq(t.unitcell_volumes, None)

    t.unitcell_vectors = np.array([[1, 0, 0], [0, 1, 0], [0, 0, 1]]).reshape(1, 3, 3)
    assert eq(np.array([1.0, 1.0, 1.0]), t.unitcell_lengths[0])
    assert eq(np.array([90.0, 90.0, 90.0]), t.unitcell_angles[0])
    assert eq(np.array([1.0]), t.unitcell_volumes)

def test_seek(get_fn):
    reference = md.load(get_fn('frame0.xyz'), top=get_fn('native.pdb'))

    with XYZTrajectoryFile(get_fn('frame0.xyz')) as f:
        f.seek(1)
        eq(1, f.tell())
        xyz1 = f.read(n_frames=1)
        eq(reference.xyz[1], xyz1[0] / 10)

        f.seek(10)
        eq(10, f.tell())
        xyz10 = f.read(n_frames=1)
        eq(reference.xyz[10], xyz10[0] / 10)
        eq(11, f.tell())

        f.seek(-8, 1)
        xyz3 = f.read(n_frames=1)
        eq(reference.xyz[3], xyz3[0] / 10)

        f.seek(4, 1)
        xyz8 = f.read(n_frames=1)
        eq(reference.xyz[8], xyz8[0] / 10)

def test_read_chunk1(get_fn, fn_xtc):
    with XTCTrajectoryFile(fn_xtc, 'r', chunk_size_multiplier=0.5) as f:
        xyz, time, step, box = f.read()

    iofile = io.loadh(get_fn('frame0.xtc.h5'), deferred=False)
    assert eq(xyz, iofile['xyz'])
    assert eq(step, iofile['step'])
    assert eq(box, iofile['box'])
    assert eq(time, iofile['time'])