def test_iterframe_from_array(self): traj = pt.iterload(fn('Tc5b.x'), fn('Tc5b.top')) # no mass fi = pt.iterframe_from_array(traj.xyz, traj.n_atoms, range(traj.n_frames)) for f_traj, (_, f_fi) in zip(traj, fi): aa_eq(f_traj.xyz, f_fi.xyz) # f_fi mass must be [1.0, ...] aa_eq(f_fi.mass, [1.0 for _ in range(traj.n_atoms)]) # mass from Topology fi = pt.iterframe_from_array(traj.xyz, traj.n_atoms, range(traj.n_frames), traj.top) for f_traj, (_, f_fi) in zip(traj, fi): aa_eq(f_traj.xyz, f_fi.xyz) aa_eq(f_fi.mass, f_traj.mass) aa_eq(f_fi.mass, traj.top.mass) # mass from array fi = pt.iterframe_from_array(traj.xyz, traj.n_atoms, range(traj.n_frames), mass=traj.top.mass) for f_traj, (_, f_fi) in zip(traj, fi): aa_eq(f_traj.xyz, f_fi.xyz) aa_eq(f_fi.mass, f_traj.mass) aa_eq(f_fi.mass, traj.top.mass)
def test_ired_simple_for_coverage(): ''' ''' traj = pt.iterload(traj_dir, parm_dir) 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)) vecs_and_mat = pt.ired_vector_and_matrix(traj, nh_indices, dtype='tuple') vecs_and_mat = pt.ired_vector_and_matrix(traj, nh_indices, dtype='tuple') state_vecs = vecs_and_mat[0] mat_ired = vecs_and_mat[1] # get eigenvalues and eigenvectors modes = pt.matrix.diagonalize(mat_ired, n_vecs=len(state_vecs)) evals, evecs = modes data_0 = _ired( state_vecs, modes=(evals, evecs), NHbond=True, tcorr=10000, tstep=1.) data_1 = _ired(state_vecs, modes=modes, NHbond=True, tcorr=10000, tstep=1) for d0, d1 in zip(data_0, data_1): if d0.dtype not in [ 'modes', ]: aa_eq(d0.values, d1.values) else: # modes # values: tuple aa_eq(d0.values[0], d1.values[0]) aa_eq(d0.values[1], d1.values[1]) # try different dtype out_try_new_dtype = pt.ired_vector_and_matrix( traj, nh_indices, dtype='cpptraj_dataset')
def test_hbonds_from_pdb(self): traj = pt.load(fn('1L2Y.pdb')) hb = pt.search_hbonds(traj) state = pt.load_cpptraj_state(''' parm {pdb} trajin {pdb} hbond series '''.format(pdb=fn('1L2Y.pdb'))) state.run() for data_p, data_cpp in zip(hb.data, state.data[1:]): assert len( data_p) == traj.n_frames == 38, 'size of dataset must be 38' aa_eq(data_p, data_cpp) # make sure distances are smaller than cutoff distance_cutoff = 2.5 angle_cutoff = 135. hb = pt.search_hbonds(traj) distances = pt.distance(traj, hb.get_amber_mask()[0]) angles = pt.angles(traj, hb.get_amber_mask()[1]) dist_indices = np.where(distances > distance_cutoff) angle_indices = np.where(angles < angle_cutoff) saved_donor_acceptors = [ 'ASP9_OD2-ARG16_NH1-HH12', 'ASP9_OD2-ARG16_NH2-HH22', 'ASP9_OD2-ARG16_NE-HE', 'ASP9_OD2-ARG16_NH2-HH21', 'ASP9_OD2-ARG16_NH1-HH11' ] donor_acceptors = pt.search_hbonds(traj, ':9,16').donor_acceptor assert saved_donor_acceptors == donor_acceptors, 'saved_donor_acceptors' aa_eq(hb.total_solute_hbonds(), hb.data['total_solute_hbonds'])
def test_real_box(self): traj = pt.load(fn('tz2.ortho.nc'), fn('tz2.ortho.parm7')) trajiter = pt.iterload(fn('tz2.ortho.nc'), fn('tz2.ortho.parm7')) saved_box = Box( [3.94559740E+01, 4.68215170E+01, 4.04695410E+01, 90., 90., 90.]) aa_eq(traj.top.box.values, saved_box.values) for frame in traj: assert frame.box.type == 'ortho' aa_eq( frame.box.values, [35.2627796623, 41.8455476799, 36.168629529, 90.0, 90.0, 90.0], decimal=1) arr0 = traj.unitcells arr1 = trajiter.unitcells for b0, b1, frame in zip(arr0, arr1, trajiter): box = frame.box # FIXME: # python3) b2 = box.values aa_eq(b0, b1) aa_eq(b0, b2) # test assign Box with list/tuple b = Box(saved_box.values) b2 = Box((t for t in saved_box.values)) aa_eq(b.values, saved_box.values, decimal=7) aa_eq(b2.values, saved_box.values, decimal=7) # assign frame.box with list/tuple frame.box = b.values b3 = frame.box aa_eq(b3.values, saved_box.values, decimal=7)
def test_basic(self): '''slicing, select''' top = pt.load_topology(fn('Tc5b.top')) # assert isinstance(top[0], Atom) assert isinstance(top[:2], pt.Topology) assert isinstance(top[:1], pt.Topology) assert isinstance(top[range(10)], pt.Topology) assert isinstance(top[list(range(10))], pt.Topology) assert isinstance(top[np.array(range(10))], pt.Topology) assert top[0].name == top['@1'][0].name # mask, AtomMask, python array, list atm = top("@CA") indices = atm.indices for a1, a2, a3, a4 in zip(top['@CA'], top[atm], top[indices], top[list(indices)]): assert a1.name == a2.name == a3.name == a4.name == 'CA' # check len assert len(top[:]) == top.n_atoms assert len(top[:10]) == 10 # API top.bond_indices top.angle_indices top.dihedral_indices
def test_convert_to_dict_and_rebuild(self): '''test_convert_to_dict_and_rebuild ''' top = self.traj.top d = top.to_dict() new_top = pt.Topology() MOLNUM = 0 for idx, (aname, atype, charge, mass, resid, resname, mol_number) in enumerate( zip(d['atom_name'], d['atom_type'], d['atom_charge'], d[ 'atom_mass'], d['resid'], d['resname'], d[ 'mol_number'])): atom = pt.core.Atom( name=aname, type=atype, charge=charge, mass=mass, resid=resid) atom.set_mol(mol_number) residue = pt.core.Residue(resname, resid) new_top.add_atom(atom, residue) if idx == 0: new_top.start_new_mol() if mol_number > MOLNUM: new_top.start_new_mol() MOLNUM += 1 new_top.add_bonds(d['bond_index']) new_top.add_dihedrals(d['dihedral_index']) new_top.box = Box(top.box.values) assert_equal_topology(top, new_top, self.traj)
def test_0_trajiter(self): traj = io.load_sample_data("tz2") from pytraj.externals.six import zip for idx, (f0, f1) in enumerate(zip(traj, traj)): pass assert idx > 0 # tee ilist = it.tee(traj, 2) for idx, (f0, f1) in enumerate(zip(ilist[0], ilist[1])): assert f0.rmsd(f1) < 1E-5 assert idx == traj.n_frames - 1 # product it.product(traj, repeat=2)
def test_simplifed_topology(self): '''simplify''' top = pt.load_topology(fn('Tc5b.top')) sim_top = top.simplify() for atom, sim_atom in zip(top.atoms, sim_top.atoms): assert atom.resname == sim_atom.resname, 'equal resname' assert atom.name == sim_atom.name, 'equal resname' assert atom.type == sim_atom.type, 'equal resname' assert atom.charge == sim_atom.charge, 'equal resname' assert atom.mass == sim_atom.mass, 'equal resname'
def test_autoimage_and_slicing(self): traj_on_disk = pt.datafiles.load_tz2_ortho() traj_on_mem = traj_on_disk[:] aa_eq(traj_on_disk.xyz, traj_on_mem.xyz) aa_eq(traj_on_disk.autoimage().xyz, traj_on_mem.autoimage().xyz) aa_eq(traj_on_mem[:1].xyz, traj_on_disk[:1].xyz) aa_eq(traj_on_mem[:].xyz, traj_on_disk[:].xyz) from pytraj.externals.six import zip for f0, f1 in zip(traj_on_disk(0, 8, 2), traj_on_mem(0, 8, 2)): aa_eq(f0.xyz, f1.xyz)
def test_0(self): traj = pt.iterload(fn('tz2.ortho.nc'), fn('tz2.ortho.parm7')) fa1 = traj[:] # do inplace-autoimage for Trajectory fa1.autoimage() # fa2 will store frame when iterating `traj` (TrajectoryIterator) fa2 = Trajectory() fa2.top = traj.top.copy() # frame_iter for frame in traj(autoimage=True): fa2.append(frame) aa_eq(fa1.xyz, fa2.xyz) # frame_iter with mask and autoimage fa3 = fa2.copy() # fa3 is a copy of autoimaged fa2. then we strip all but CA atoms # just want to make sure we can use `mask` fa3.strip("!@CA") fa4 = Trajectory() fa4.top = fa3.top.copy() for frame in traj(mask='@CA', autoimage=True): fa4.append(frame) aa_eq(fa3.xyz, fa4.xyz) # frame_iter with mask and autoimage, and rmsfit # fa3 is a copy of autoimaged fa2. then we strip all but CA atoms # just want to make sure we can use `mask` fa3 = traj[:] fa3.autoimage() fa3.rmsfit(ref=5, mask='@CB') fa3.strip("!@CA") fa4 = Trajectory() fa4.top = fa3.top.copy() ref1 = traj[5].copy() for frame in traj(mask='@CA', autoimage=True, rmsfit=(ref1, '@CB')): fa4.append(frame) aa_eq(fa3.xyz, fa4.xyz) for f0, f1 in zip(fa3, fa4): assert f0.rmsd_nofit(f1) < 1E-7
def test_0(self): # TrajectoryIterrator # status: failed from pytraj.externals.six import zip traj = pt.iterload(fn('tz2.ortho.nc'), fn('tz2.ortho.parm7')) pt.write_traj("./output/tz2.autoimage_with_rmsfit.nc", traj(autoimage=True, rmsfit=(0, '@CA,C,N')), overwrite=True) saved_traj = pt.load(fn('tz2.autoimage_with_rmsfit.nc'), traj.top) p_traj = pt.load('./output/tz2.autoimage_with_rmsfit.nc', traj.top) aa_eq(saved_traj.xyz, p_traj.xyz) for f1, f2 in zip(p_traj, saved_traj): pass
def test_on_disk_trajectory(self): # TrajectoryIterrator output = "ok_to_delete.nc" traj = pt.iterload(fn('tz2.ortho.nc'), fn('tz2.ortho.parm7')) pt.write_traj( output, traj(autoimage=True, rmsfit=(0, '@CA,C,N')), overwrite=True) saved_traj = pt.load(fn('tz2.autoimage_with_rmsfit.nc'), traj.top) p_traj = pt.load(output, traj.top) aa_eq(saved_traj.xyz, p_traj.xyz) # ensure iterable without memory freeing for f1, f2 in zip(p_traj, saved_traj): aa_eq(f1.xyz, f2.xyz)
def test_ired_need_lapack_cpptraj(): 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_eigenvectors, cpp_eigenvalues)) 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)
def merge_frame_from_trajs(trajlist): """ Examples -------- >>> import numpy as np >>> import pytraj as pt >>> traj0 = pt.load_sample_data('tz2')[:3] >>> traj1 = pt.load_sample_data('tz2')[3:6] >>> traj2 = pt.load_sample_data('tz2')[6:9] >>> print(traj0.n_atoms, traj1.n_atoms, traj2.n_atoms) 5293 5293 5293 >>> for frame in pt.tools.merge_frame_from_trajs((traj0, traj1, traj2)): print(frame) <Frame with 15879 atoms> <Frame with 15879 atoms> <Frame with 15879 atoms> """ for iterables in zip(*trajlist): yield merge_frames(iterables)
def test_0(self): from pytraj.externals.six import zip from pytraj import Trajectory, TrajectoryIterator # TrajectoryIterator object traj0 = pt.iterload(fn('tz2.ortho.nc'), fn('tz2.ortho.parm7')) assert isinstance(traj0, TrajectoryIterator) # Trajectory object traj1 = traj0[:] assert isinstance(traj1, Trajectory) # check same coords ref = traj0[0] for f0, f1 in zip(traj0(autoimage=True, rmsfit=(ref, '@CA,C,N')), traj1(autoimage=True, rmsfit=(ref, '@CA,C,N'))): aa_eq(f0.xyz, f1.xyz) assert f0.rmsd_nofit(f1) == 0.
def test_simplifed_topology(self): '''simplify''' top = pt.load_topology(fn('Tc5b.top')) sim_top = top.simplify() assert sim_top.select('@CA').tolist() == top.select('@CA').tolist() for atom, sim_atom in zip(top.atoms, sim_top.atoms): assert atom.resname == sim_atom.resname assert atom.name == sim_atom.name assert atom.type == sim_atom.type assert atom.charge == sim_atom.charge assert atom.mass == sim_atom.mass # API atom = sim_top.atoms[0] atom.residue atom.residue.name atom.residue.index atom.bond_partners
def n_grams(a, n): """n_grams Parameters ---------- a : sequence n : number of elements asarray : bool, default False if False: return an iterator if True: return a numpy array Examples -------- >>> list(n_grams([2, 3, 4 ,5], 2)) [(2, 3), (3, 4), (4, 5)] Notes ----- adapted from: http://sahandsaba.com/thirty-python-language-features-and-tricks-you-may-not-know.html """ z = (islice(a, i, None) for i in range(n)) return zip(*z)
def test_assert(self): traj = pt.iterload(fn('Tc5b.x'), fn('Tc5b.top')) fa = Trajectory.from_iterable(iterframe_master(traj), top=traj.top) for f0, f1 in zip(fa, traj): aa_eq(f0.xyz, f1.xyz)
def merge_trajs(traj1, traj2, start_new_mol=True, n_frames=None): """ Examples -------- >>> import pytraj as pt >>> import numpy as np >>> traj1 = pt.load_sample_data('ala3')[:1] >>> traj2 = pt.load_sample_data('tz2')[:1] >>> traj3 = merge_trajs(traj1, traj2) >>> # from frame_iter for saving memory >>> traj3 = merge_trajs((traj1(0, 10, 2), traj1.top), (traj2(100, 110, 2), traj2.top), n_frames=6) >>> # raise error if not having the same n_frames >>> traj4 = pt.load_sample_data('tz2')[:] >>> traj4.n_frames 10 >>> traj1.n_frames 1 >>> merge_trajs(traj1, traj4) Traceback (most recent call last): ... ValueError: must have the same n_frames Notes ----- Code might be changed """ from pytraj.externals.six import zip from pytraj import Trajectory import numpy as np if isinstance(traj1, (list, tuple)): n_frames_1 = n_frames top1 = traj1[1] _traj1 = traj1[0] else: n_frames_1 = traj1.n_frames top1 = traj1.top _traj1 = traj1 if isinstance(traj2, (list, tuple)): n_frames_2 = n_frames top2 = traj2[1] # example: (traj(0, 5), traj.top) _traj2 = traj2[0] else: n_frames_2 = traj2.n_frames top2 = traj2.top _traj2 = traj2 if n_frames_1 != n_frames_2: raise ValueError("must have the same n_frames") traj = Trajectory() traj._allocate(n_frames_1, top1.n_atoms + top2.n_atoms) # merge Topology top = top1.copy() if start_new_mol: top.start_new_mol() top.join(top2) traj.top = top # update coords for f1, f2, frame in zip(_traj1, _traj2, traj): frame.xyz = np.vstack((f1.xyz, f2.xyz)) return traj
def test_4_trajiter(self): traj = pt.load_sample_data("tz2") for idx, (f0, f1) in enumerate(zip(traj, traj)): f0.rmsd(f1)
def test_0(self): traj = pt.iterload(fn('Tc5b.x'), fn('Tc5b.top')) farray = traj[:] for i, f0 in enumerate(traj): for j, x in enumerate(f0.xyz): if np.all(np.abs(x - 5.707) < 1E-3): pass for frame in traj.iterframe(): pass for frame0 in farray.iterframe(): pass i = 0 for frame0 in farray.iterframe(start=0, step=1): i += 1 aa_eq(traj[-1].xyz, frame0.xyz) start, stop, step = 2, 8, 4 indices = list(range(start, stop, step)) for idx, frame0, f in zip(indices, farray.iterframe(start, stop, step), traj[indices]): aa_eq(frame0.xyz, f.xyz) aa_eq(traj[6].xyz, frame0.xyz) for frame0 in farray.iterframe(start=2, step=2): pass aa_eq(traj[8].xyz, frame0.xyz) traj[6][0] for frame0 in traj.iterframe(start=2, step=4, stop=8): pass for frame0 in traj.iterframe(start=2, step=4, stop=8): pass aa_eq(traj[6].xyz, frame0.xyz) for frame0 in traj.iterframe(start=2, step=2): pass aa_eq(traj[8].xyz, frame0.xyz) count = 0 for frame0 in traj.iterframe(start=2): count += 1 aa_eq(traj[-1].xyz, frame0.xyz) count = 0 for frame0 in traj.iterframe(start=2, stop=7): count += 1 aa_eq(traj[6].xyz, frame0.xyz) for frame0 in traj.iterframe(): pass aa_eq(traj[-1].xyz, frame0.xyz) for frame0 in farray.iterframe(): pass aa_eq(traj[-1].xyz, frame0.xyz) for frame0 in traj(): pass aa_eq(traj[-1].xyz, frame0.xyz) for frame0 in farray(): pass aa_eq(farray[-1].xyz, frame0.xyz) count = 0 for frame0 in traj(start=2, stop=7): count += 1 aa_eq(traj[6].xyz, frame0.xyz) count = 0 for frame0 in farray(start=2, stop=7): count += 1 aa_eq(traj[6].xyz, frame0.xyz) count = 0 for frame0 in farray(2, 7, 1): count += 1 aa_eq(traj[6].xyz, frame0.xyz) count = 0 for frame0 in farray(2, 7, 2): count += 1 aa_eq(traj[6].xyz, frame0.xyz)