Ejemplo n.º 1
0
    def test_constructor(self):
        pname = fn("tz2.ortho.parm7")
        tname = fn("tz2.ortho.nc")
        traj = pt.Trajectory(tname, top=pname)
        assert traj.n_frames == 10

        traj = pt.Trajectory([tname, tname], top=pname)
        on_dis_traj = pt.iterload([tname, tname], pname)
        assert traj.n_frames == on_dis_traj.n_frames == 20
        aa_eq(traj.unitcells, on_dis_traj.unitcells)
Ejemplo n.º 2
0
    def test_raise_construtor(self):
        self.assertRaises(ValueError, lambda: pt.Trajectory(pt.trajectory))
        # raise if filename is not string or list of string
        empty_traj = pt.Trajectory()
        empty_traj.top = pt.datafiles.load_tz2()[:].top
        self.assertRaises(ValueError, lambda: empty_traj.load(pt.Trajectory))

        # raise if empty Topology
        xyz = np.arange(90).astype('f8').reshape(3, 10, 3)
        self.assertRaises(ValueError, lambda: pt.Trajectory(xyz=xyz, top=pt.Topology()))
Ejemplo n.º 3
0
    def test_closest_compared_to_cpptraj(self):
        trajin = fn('tz2.ortho.nc')
        parm = fn('tz2.ortho.parm7')
        traj = pt.iterload(trajin, parm, frame_slice=[
            (0, 2),
        ])
        state = pt.load_cpptraj_state('''
        parm {}
        trajin {} 1 2
        autoimage
        closest 100 :1-13
        createcrd mycrd'''.format(parm, trajin))
        state.run()

        fi, top = pt.closest(traj(autoimage=True),
                             mask=':1-13',
                             n_solvents=100)
        xyz = pt.get_coordinates(fi)
        t0 = pt.Trajectory(xyz=xyz, top=top)
        aa_eq(state.data['mycrd'].xyz, t0.xyz)

        # dtype = 'trajectory'
        t1 = pt.closest(traj(autoimage=True),
                        mask=':1-13',
                        n_solvents=100,
                        dtype='trajectory')
        aa_eq(state.data['mycrd'].xyz, t1.xyz)
Ejemplo n.º 4
0
def test_get_coordinates_trajectory():
    '''mutable pytraj.Trajectory
    '''
    traj = pt.Trajectory(xyz=traj_tz2_ortho.xyz, top=traj_tz2_ortho.top)
    # make a different copy since ``traj`` is mutable
    traj2 = traj.copy()

    # all coordinates
    xyz = pt.get_coordinates(traj)
    aa_eq(traj.xyz, xyz)

    # given frames
    xyz = pt.get_coordinates(traj, frame_indices=[0, 5])
    aa_eq(traj[[0, 5]].xyz, xyz)

    # given frames, autoimage=True
    xyz = pt.get_coordinates(traj, frame_indices=[0, 5], autoimage=True)
    aa_eq(traj2[[0, 5]].autoimage().xyz, xyz)

    # given frames, autoimage=True, rmsfit=ref
    ref = traj[-3]
    pt.autoimage(ref, top=traj.top)
    xyz = pt.get_coordinates(traj,
                             frame_indices=[0, 5],
                             autoimage=True,
                             rmsfit=ref)
    aa_eq(traj2[[0, 5]].autoimage().superpose(ref=ref).xyz, xyz)
Ejemplo n.º 5
0
def shrink(xtc,gro,simname):
	if os.path.exists(simname + '.xtc'):
		print simname + '_ai.xtc', "exists"
	else:
		print "writing", simname + '.xtc'

		traj = md.load_xtc(xtc,top=gro)
		traj = pt.Trajectory(xyz=m_traj.xyz.astype('f8'), top=topology_name)
		traj.autoimage()
		traj.save(simname + '_ai.xtc')

		traj_pt = []

		traj = md.load_xtc(simname + '_ai.xtc',top=gro)
		exit()

		with md.formats.XTCTrajectoryFile(simname + '.xtc', 'w') as f:
			for fr in range(50000)[::100]:
				print "frame", fr
				structure = traj[fr]
				peptide = structure.atom_slice(structure.topology.select('protein'))[0]
				peptide.center_coordinates()
				f.write(peptide.xyz*10)
			peptide.save(simname + '.pdb')
	return None
Ejemplo n.º 6
0
    def test_indexing(self):
        traj = pt.iterload(fn('Tc5b.x'), fn('Tc5b.top'))
        traj2 = pt.TrajectoryIterator()
        traj2.top = pt.load_topology(fn('Tc5b.top'))
        traj2._load(fn('Tc5b.x'))
        farray = traj2[[0, 9, 1]]
        assert farray.n_frames == 3
        assert traj2[0].atom(0) == farray[0].atom(0)
        assert traj2[9].atom(0) == farray[1].atom(0)
        assert traj2[1].atom(0) == farray[2].atom(0)

        arr = np.asarray(traj2[0]._buffer1d[:])
        frame0 = traj2[0]
        arr0 = np.asarray(frame0._buffer1d[:])

        mat0 = np.asmatrix(arr0).reshape(304, 3)
        mat0[:, 0] = np.asmatrix(list(range(304))).reshape(304, 1)
        assert frame0[0, 0] == 0.
        assert frame0[1, 0] == 1.
        assert frame0[2, 0] == 2.

        # raise if size = 0
        traj3 = pt.Trajectory()
        assert traj3.n_frames == 0, 'empty Trajectory, n_frames must be 0'
        self.assertRaises(IndexError, lambda: traj3[0])
        self.assertRaises(IndexError, lambda: traj3.__setitem__(0, traj[3]))
Ejemplo n.º 7
0
    def test_slice_and_iterate_trajectory_with_force_and_velocity(self):
        traj_on_disk = pt.iterload(fn('tz2.nc'), fn('tz2.parm7'))
        n_atoms = traj_on_disk.n_atoms
        n_frames = traj_on_disk.n_frames
        xyz = traj_on_disk.xyz
        force = np.random.rand(n_frames * n_atoms * 3).reshape(
            (n_frames, n_atoms, 3))
        velocity = np.random.rand(n_frames * n_atoms * 3).reshape(
            (n_frames, n_atoms, 3))
        top = traj_on_disk.top
        traj = pt.Trajectory(xyz=xyz, top=top, force=force, velocity=velocity)
        aa_eq(traj.forces, force)
        aa_eq(traj.velocities, velocity)

        # iterating
        for index, frame in enumerate(traj):
            aa_eq(frame.force, traj.forces[index])
            aa_eq(frame.velocity, traj.velocities[index])

        # slicing
        aa_eq(traj[0].force, traj.forces[0])
        aa_eq(traj[0].velocity, traj.velocities[0])
        aa_eq(traj[::2].forces, traj.forces[::2])
        aa_eq(traj[::2].velocities, traj.velocities[::2])
        ca_indices = traj_on_disk.top.select('@CA')
        aa_eq(traj['@CA'].velocities, traj.velocities[:, ca_indices])
        aa_eq(traj['@CA'].forces, traj.forces[:, ca_indices])
Ejemplo n.º 8
0
 def test_slice_basic(self):
     traj2 = pt.Trajectory()
     traj2.top = pt.load_topology(fn('Tc5b.top'))
     traj2.load(fn('Tc5b.x'))
     traj2.load(fn('Tc5b.x'))
     traj2.load(fn('Tc5b.x'))
     traj2.load(fn('Tc5b.x'))
     fsub = traj2[2:10]
     fsub[0][0] = 100.
Ejemplo n.º 9
0
    def test_allocate_frames(self):
        traj = pt.iterload(fn('Tc5b.x'), fn('Tc5b.top'))
        traj2 = pt.Trajectory()
        traj2._allocate(traj.n_frames, traj.n_atoms)
        assert (traj2.shape == traj.shape)

        traj2.top = traj.top.copy()
        traj2.xyz = traj.xyz[:]
        aa_eq(traj2.xyz, traj.xyz)
Ejemplo n.º 10
0
 def test_slice_basic(self):
     traj2 = pt.Trajectory()
     traj2.top = pt.load_topology("./data/Tc5b.top")
     traj2.load("./data/Tc5b.x")
     traj2.load("./data/Tc5b.x")
     traj2.load("./data/Tc5b.x")
     traj2.load("./data/Tc5b.x")
     fsub = traj2[2:10]
     fsub[0][0] = 100.
Ejemplo n.º 11
0
    def test_closest(self):
        # raise if not has solvent
        traj0 = pt.iterload(fn('tz2.nc'),
                            fn('tz2.parm7'),
                            frame_slice=[(0, 2)])
        with pytest.raises(RuntimeError):
            pt.closest(traj0)

        traj = pt.iterload(fn('tz2.ortho.nc'),
                           fn('tz2.ortho.parm7'),
                           frame_slice=[(0, 2)])
        fi, top = pt.closest(traj)

        coords = []
        for frame in fi:
            coords.append(frame.xyz.copy())
            assert isinstance(frame, pt.Frame), 'must be Frame'

        # make a Trajectory
        fi, top = pt.closest(traj)
        xyz = pt.get_coordinates(fi)
        t0 = pt.Trajectory(xyz=xyz, top=top)
        aa_eq(np.array(coords), t0.xyz)

        # test write to disk
        fi, top = pt.closest(traj)
        new_traj = pt.Trajectory(top=top)
        for frame in fi:
            new_traj.append(frame)

        pt.write_traj(outputname('fi.nc'), new_traj, overwrite=True)
        # load back
        t1 = pt.load(outputname('fi.nc'), top=top)
        aa_eq(t0.xyz, t1.xyz)

        # make sure n_sovent=10 (default)
        n_solvents = 0
        for mol in top.mols:
            if mol.is_solvent():
                n_solvents += 1
        assert n_solvents == 10, 'must be 10 solvents'
Ejemplo n.º 12
0
def pack(traj, mol, n_copies, unitcells=None, ig=8888, grid_spacing=0.2):
    '''

    Parameters
    ----------
    traj : pytraj.Trajectory
    mol : pytraj.Trajectory
    unitcells : None or np.ndarray, dim=2
        if None, use box info from traj else use it
    n_copies : number of `mol`
    ig : int
        randome seed
    grid_spacing : float
    '''
    add_to_box_exe = amberbin('AddToBox') or 'AddToBox'
    input_pdb = 'input.pdb'
    mol_pdb = 'mol.pdb'
    out_pdb = 'out.pdb'

    assert mol.n_frames == 1
    
    total_n_atoms = traj.n_atoms + mol.n_atoms * n_copies
    new_traj_xyz = np.empty((traj.n_frames, total_n_atoms, 3), dtype='f8')

    with tempfolder():
        mol.save(mol_pdb, overwrite=True)
        for index, frame in enumerate(traj):
            pytraj.write_traj(input_pdb, traj=frame, top=traj.top, overwrite=True)
            parm = parmed.load_file(input_pdb)
            if unitcells is not None:
                parm.box = unitcells[index]
            else:
                assert frame.has_box(), 'must have box info'
                parm.box = frame.box
            # add remark 290
            parm.save(input_pdb, overwrite=True)
            command = [
                add_to_box_exe,
                '-c', input_pdb,
                '-a', mol_pdb,
                '-na', str(n_copies),
                '-IG', str(ig),
                '-G', str(grid_spacing),
                '-o', out_pdb
            ]
            try:
                subprocess.check_output(command, stderr=subprocess.STDOUT)
                new_traj_xyz[index] = pytraj.load(out_pdb).xyz
            except subprocess.CalledProcessError as e:
                return e.output.decode()
        top = pytraj.load_topology(out_pdb)
    return pytraj.Trajectory(xyz=new_traj_xyz, top=top)
Ejemplo n.º 13
0
    def test_ired_need_lapack_cpptraj(self):
        state = pt.load_cpptraj_state(txt)
        state.run()
        xyz = state.data['CRD1'].xyz
        top = state.data['CRD1'].top
        traj = pt.Trajectory(xyz=xyz, top=top)
        state_vecs = state.data[1:-3].values

        h_indices = pt.select_atoms('@H', traj.top)
        n_indices = pt.select_atoms('@H', traj.top) - 1
        nh_indices = list(zip(n_indices, h_indices))
        mat_ired = pt.ired_vector_and_matrix(traj,
                                             mask=nh_indices,
                                             order=2)[-1]
        mat_ired /= mat_ired[0, 0]

        # matired: make sure to reproduce cpptraj output
        aa_eq(mat_ired, state.data['matired'].values)

        # get modes
        modes = state.data[-2]
        cpp_eigenvalues = modes.eigenvalues
        cpp_eigenvectors = modes.eigenvectors
        evals, evecs = np.linalg.eigh(mat_ired)

        # need to sort a bit
        evals = evals[::-1]
        # cpptraj's eigvenvalues
        aa_eq(evals, cpp_eigenvalues)

        # cpptraj's eigvenvectors
        # use absolute values to avoid flipped sign
        # from Dan Roe
        # In practice, the "sign" of an eigenvector depends on the math library used to calculate it.
        # This is in fact why the modes command displacement test is disabled for cpptraj.
        # I bet if you use a different math library (e.g. use your system BLAS/LAPACK instead of the one bundled with Amber
        # or vice versa) you will get different signs.
        # Bottom line is that eigenvector sign doesn't matter.

        aa_eq(np.abs(evecs[:, ::-1].T), np.abs(cpp_eigenvectors), decimal=4)
        data = _ired(state_vecs, modes=(cpp_eigenvalues, cpp_eigenvectors))
        order_s2 = data['IRED_00127[S2]']

        # load cpptraj's output and compare to pytraj' values for S2 order paramters
        cpp_order_s2 = np.loadtxt(os.path.join(cpptraj_test_dir, 'Test_IRED',
                                               'orderparam.save')).T[-1]
        aa_eq(order_s2, cpp_order_s2, decimal=5)
Ejemplo n.º 14
0
    def test_append_trajectory(self):
        # test append
        traj = pt.Trajectory()
        t = pt.iterload(fn('Tc5b.x'), fn('Tc5b.top'))
        traj.top = t.top

        # append single Frame
        traj.append(t[0])
        assert traj.n_frames == 1

        # append xyz
        traj.append(t.xyz[:])
        assert traj.n_frames == t.n_frames + 1

        # append TrajectoryIterator
        traj.append(t)
        assert traj.n_frames == t.n_frames * 2 + 1

        # append frame_iter
        traj.append(t.iterframe())
        assert traj.n_frames == t.n_frames * 3 + 1

        # append pt.iterframe_master
        traj.append(pt.iterframe_master(t))
        assert traj.n_frames == t.n_frames * 4 + 1

        # append itself
        NFrames = traj.n_frames
        traj.append(traj)
        assert traj.n_frames == NFrames * 2

        # append itself frame_iter
        traj.append(traj.iterframe(stop=2))
        assert traj.n_frames == NFrames * 2 + 2

        # append pt.iterframe_master for itself
        NFrames = traj.n_frames
        traj.append(pt.iterframe_master(traj))
        assert traj.n_frames == NFrames * 2

        # append pt.iterframe_master for itself + other
        n0 = traj.n_frames
        n1 = t.n_frames
        traj.append(pt.iterframe_master([traj, t]))
        assert traj.n_frames == 2 * n0 + n1
Ejemplo n.º 15
0
    def test_autoimage_rotatedihedral(self):
        traj = pt.iterload(fn('tz2.ortho.nc'), fn('tz2.ortho.parm7'))
        farray = traj[:]

        t0trajectory = pt.Trajectory(traj)
        aa_eq(farray.unitcells, t0trajectory.unitcells)

        # autoimage
        farray.autoimage()
        t0trajectory.autoimage()
        aa_eq(farray.xyz, t0trajectory.xyz)

        # rotate_dihedral
        pt.rotate_dihedral(t0trajectory, '3:phi:120')
        pt.rotate_dihedral(farray, '3:phi:120')
        aa_eq(farray.xyz, t0trajectory.xyz)

        aa_eq(
            pt.calc_phi(t0trajectory, '3').values,
            [120 for _ in range(t0trajectory.n_frames)])
Ejemplo n.º 16
0
    def test_closest_compared_to_cpptraj(self):
        traj = pt.iterload(fn('tz2.ortho.nc'), fn('tz2.ortho.parm7'))
        state = pt.load_cpptraj_state(
            '''
        autoimage
        closest 100 :1-13
        createcrd mycrd''', traj)
        state.run()

        fi, top = pt.closest(traj(autoimage=True),
                             mask=':1-13',
                             n_solvents=100)
        xyz = pt.get_coordinates(fi)
        t0 = pt.Trajectory(xyz=xyz, top=top)
        aa_eq(state.data['mycrd'].xyz, t0.xyz)

        # dtype = 'trajectory'
        t1 = pt.closest(traj(autoimage=True),
                        mask=':1-13',
                        n_solvents=100,
                        dtype='trajectory')
        aa_eq(state.data['mycrd'].xyz, t1.xyz)
Ejemplo n.º 17
0
def cal_dih(top, trajectory, dih_mask, dihnum):
    '''
    top is an amber topology file
    trajectory is an amber trajectorty, it can be both netcdf or ascii
    dih_mask is the dihedral atom number mask  e.g @2 @3 @5 @8
    dihnum is the number of dihedral being fitted
    cal_dih() returns dihedral values for a given mask
    '''
    # first load trajectory

    traj = pt.Trajectory(trajectory, top)
    # if only one dihedral then
    if dihnum == 1:
        dih_data = pt.dihedral(traj, '%s' % (dih_mask))
        return dih_data
    # we have multiple dihedrals
    else:
        dih_data = []
        # calculate all dihedrals
        for mask in dih_mask:
            data = pt.dihedral(traj, '%s' % (mask))
            dih_data.append(data)
        return dih_data
    def test_buffer_not_c_contiguous(self):
        # source code was lightly adapted from jmborr
        # https://github.com/Amber-MD/pytraj/issues/991
        traj = pt.load(fn('issue991/short.dcd'),
                       fn('issue991/pdb.gz'),
                       mask='(!:1-256)&(@H1,@H2,@H3,@H4,@H5)')

        # Trajectory of the center of mass of the first two residues
        xyz = np.empty((2, traj.n_frames, 3))
        xyz[0] = pt.center_of_mass(traj, mask=':1')
        xyz[1] = pt.center_of_mass(traj, mask=':2')

        # "topology" using the first two atoms
        minitop = traj.top['@1,2']
        minitraj = pt.Trajectory(xyz=np.ascontiguousarray(
            xyz.transpose((1, 0, 2))),
                                 top=minitop)

        # Save trajectory
        # make sure there is no ValueError
        # something is wrong with pdb, crd extension when loading with
        # minitop (poor topology?)
        # make issue in cpptraj too?

        exts = ['nc', 'dcd', 'mdcrd', 'crd', 'pdb', 'trr']
        exts.remove('mdcrd')
        exts.remove('crd')

        for ext in exts:
            junk_fn = 'junk.' + ext
            pt.write_traj(filename=junk_fn, traj=minitraj, overwrite=True)

            # load coord back to make sure we correctly write it
            new_traj = pt.iterload(junk_fn, minitop)
            # mdcrd, crd, pdb has only 3 digits after decimal
            decimal = 5 if ext in ['nc', 'dcd', 'trr'] else 3
            aa_eq(minitraj.xyz, new_traj.xyz, decimal=decimal)
Ejemplo n.º 19
0
    def test_closest(self):
        # raise if not has solvent
        traj0 = pt.iterload(fn('tz2.nc'), fn('tz2.parm7'))
        self.assertRaises(RuntimeError, lambda: pt.closest(traj0))

        traj = pt.iterload(fn('tz2.ortho.nc'), fn('tz2.ortho.parm7'))
        fi, top = pt.closest(traj)

        coords = []
        for frame in fi:
            coords.append(frame.xyz.copy())
            assert isinstance(frame, pt.Frame), 'must be Frame'

        # make a Trajectory
        fi, top = pt.closest(traj)
        xyz = pt.get_coordinates(fi)
        t0 = pt.Trajectory(xyz=xyz, top=top)
        aa_eq(np.array(coords), t0.xyz)

        # test write to disk
        fi, top = pt.closest(traj)

        pt.write_traj('output/fi.nc', fi, top=top, overwrite=True)
        # load back
        t1 = pt.load('output/fi.nc', top=top)
        aa_eq(t0.xyz, t1.xyz)

        # make sure n_sovent=10 (default)
        n_solvents = 0
        for mol in top.mols:
            if mol.is_solvent():
                n_solvents += 1
        assert n_solvents == 10, 'must be 10 solvents'

        fi, top = pt.closest(traj)
        pt.write_traj('output/test.pdb', next(fi), top=top, overwrite=True)
from __future__ import print_function
from parmed.utils.netcdf import netcdf_file
import pytraj as pt
import numpy as np

fh = netcdf_file('data/mdfrc',
                 mmap=False)  # change mdfrc to your force filename
forces = fh.variables['forces']

# do similar thing for velocities

top = pt.load_topology(
    'data/mdfrc.prmtop')  # change prmtop to your parm7 filename
traj = pt.Trajectory(xyz=np.ascontiguousarray(forces[:], dtype='f8'), top=top)
print('original trajectory', traj)

# strip atom atoms
# change ':1-100' to your mask (e.g ':WAT')
new_traj = pt.strip(':1-100', traj)
print('stripped trajectory', new_traj)

# write trajectory
new_traj.save('test.nc', overwrite=True)
Ejemplo n.º 21
0
# load pytraj
import pytraj as pt
import numpy as np

# load mdtraj
import mdtraj as md
from mdtraj.testing import get_fn
# load h5 file 
#filename, topology_name = get_fn('tfit_10.xtc'), get_fn('prot0.pdb')
m_traj = md.load('tfit_10.xtc', top='prot0.pdb')
m_traj
traj = pt.Trajectory(xyz=m_traj.xyz.astype('f8'), top='prot0.pdb')
traj

# creat N-H vector pairs

# select H (backbone) indices with given residue number (from previous step)
#H_mask = ':' + ','.join(str(i) for i in resnums) + '@H'
#print('H_mask: ', H_mask)
h_indices = pt.select_atoms(traj.top, '@H')

# select N (backbone) indices
n_indices = h_indices - 1

# create pairs
nh_pairs = list(zip(n_indices, h_indices))
nh_pairs[:3]

# do calculation
s2 = pt.nh_order_parameters(traj, nh_pairs, tcorr=52801.0, tstep=1.0)
st20 = pt.nh_order_parameters(traj, nh_pairs, tcorr=20000.0, tstep=1.0)
Ejemplo n.º 22
0
frame = traj[1]
print(frame)

# now you want to rotate omega dihedral angle for residue 3?

# make a list to store new frames
traj_xyz = []

for deg in range(-180, 180, 5):
    # since Frame does not hold Topology information, you need to
    # pass it
    pt.rotate_dihedral(frame, '3:omega:' + str(deg), top=traj.top)
    # always make a copy since
    # the `rotate_dihedral` inplace-change coords of frame
    traj_xyz.append(frame.xyz.copy())

# write multiple pdbs to a single file
# just like multiple pdbs when you load from rcsb
# you can use this pdb file to view in VMD without loading
# prmtop/gro/psf/...

new_traj = pt.Trajectory(xyz=np.asarray(traj_xyz), top=traj.top)
pt.write_traj(
    "./output/test0.pdb",
    new_traj,
    options='model',
    overwrite=True)

# you can use VMD to open the new file
# vmd ./output/test0.pdb
Ejemplo n.º 23
0
import pytraj as pt

try:
    import mdtraj as md

    # load mdtraj object
    m_traj = md.load_netcdf('../tests/data/tz2.ortho.nc',
                            '../tests/data/tz2.ortho.parm7')
    print(m_traj)

    # convert to pytraj object
    # you can use a pdb file, a mol2 file, ... as Topology too
    # as long as pytraj/cpptraj supports
    traj = pt.Trajectory(xyz=m_traj.xyz.astype('f8'),
                         top='../tests/data/tz2.ortho.parm7')
    print(traj)

    # perform 'action' on traj
    traj.autoimage()

    # copy data back to mdtraj object
    m_traj.xyz = traj.xyz[:]

    # mdtraj has very fast rmsd calculation, you can pass pytraj'traj object
    # to mdtraj to 'borrow' its action too.
    # note that pytraj/cpptraj use Angstrom for unit while mdtraj use nm
    print(md.rmsd(traj, traj, 0))

except ImportError:
    print("does not have mdtraj")
Ejemplo n.º 24
0
    def __init__(self,
                 top,
                 trajs,
                 rotfit=False,
                 mask='@CA',
                 stride=1,
                 n_blocks=1,
                 bframes=50000,
                 skip=25000):
        #read sequence
        if trajs[0][-3:] == 'dcd':
            traj = pt.iterload(trajs[0], top)
        else:
            #filename, top_name = get_fn(trajs[0]), get_fn(top)
            m_traj = md.load(trajs[0], top=top)
            traj = pt.Trajectory(xyz=m_traj.xyz.astype('f8'), top=top)
        self.sequence = {res.index + 1: res.name for res in traj.top.residues}
        self.top = top
        self.trajs = trajs
        n_frames = len(traj)

        self.n_indices = pt.select_atoms('@N & !(:1) & !(:PRO)', traj.top)
        self.h_indices = self.n_indices + 1
        self.nh_pairs = np.array(list(zip(self.n_indices, self.h_indices)))

        res_list = sorted(
            [k for k, val in self.sequence.items() if val != 'PRO' and k != 1])

        self.acf = dict()
        self.S2 = dict()

        if n_blocks == 1:
            for tt, t in enumerate(self.trajs):
                if t[-3:] == 'dcd':
                    traj = pt.load(t, top, stride=stride)
                else:
                    #filename, top_name = get_fn(t), get_fn(top)
                    m_traj = md.load(t, top=top)
                    traj = pt.Trajectory(xyz=m_traj.xyz.astype('f8'), top=top)
                if rotfit == True:
                    #print('RMSD fitting on '+mask)
                    _ = traj.superpose(ref=-1, mask='@CA')
                data_vec = va.vector_mask(traj, self.nh_pairs, dtype='ndarray')
                self.acf[tt] = {
                    res_list[n]: pt.timecorr(vals,
                                             vals,
                                             order=2,
                                             tstep=1,
                                             tcorr=len(vals),
                                             norm=False,
                                             dtype='ndarray')
                    for n, vals in enumerate(data_vec)
                }
                self.S2[tt] = {
                    res_list[n]: S2(vals)
                    for n, vals in enumerate(data_vec)
                }
        else:
            index = 0
            bcount = 0
            for tt, t in enumerate(self.trajs):
                print(tt + 1, '/', len(self.trajs))
                while bcount < n_blocks:
                    first_frame = bcount * skip
                    last_frame = first_frame + bframes
                    if t[-3:] == 'dcd':
                        traj = pt.load(t,
                                       top,
                                       frame_indices=slice(
                                           first_frame, last_frame, stride))
                    else:
                        #filename, top_name = get_fn(t), get_fn(top)
                        m_traj = md.load(t, top=top)
                        m_traj = m_traj[first_frame:last_frame:stride]
                        traj = pt.Trajectory(xyz=m_traj.xyz.astype('f8'),
                                             top=top)
                    if rotfit == True:
                        _ = traj.superpose(ref=-1, mask='@CA')
                    data_vec = va.vector_mask(traj,
                                              self.nh_pairs,
                                              dtype='ndarray')
                    self.acf[index] = {
                        res_list[n]: pt.timecorr(vals,
                                                 vals,
                                                 order=2,
                                                 tstep=1,
                                                 tcorr=len(vals),
                                                 norm=False,
                                                 dtype='ndarray')
                        for n, vals in enumerate(data_vec)
                    }
                    self.S2[index] = {
                        res_list[n]: S2(vals)
                        for n, vals in enumerate(data_vec)
                    }
                    bcount += 1
                    index += 1
                bcount = 0
def load_mutable_traj(filename=filename):
    with Dataset(filename, 'r') as fh:
        c = fh['coordinates']
        traj = pt.Trajectory(xyz=c[:N], top=top.copy())
        traj.xyz[:] = traj.xyz.astype('f8')