def make_fake_topology(n_atoms): '''make fake Topology, just for writing xyz array to supported formats (netcdf, dcd, trr, ...) >>> import pytraj as pt >>> top = pt.tools.make_fake_topology(100) >>> top.n_atoms 100 >>> isinstance(top, pt.Topology) True >>> import numpy as np >>> xyz = np.random.rand(10*100*3).reshape(10, 100, 3) >>> traj0 = pt.Trajectory(xyz=xyz, top=top) >>> pt.write_traj('output/test.nc', traj0, overwrite=True) >>> traj = pt.iterload('output/test.nc', top=top) >>> traj.n_atoms 100 ''' from pytraj import Atom, Residue, Topology top = Topology() for _ in range(n_atoms): atom = Atom(name='X', type='X', charge=0., mass=0., resid=0) residue = Residue('Y', 0) top.add_atom(atom, residue) return top
def split_traj_by_residues(traj, start=0, stop=-1, step=1): '''return a generator Examples -------- >>> import pytraj as pt >>> traj = pt.datafiles.load_rna() >>> g = pt.tools.split_traj_by_residues(traj) >>> t0 = next(g) >>> print(t0.top.n_residues) 1 ''' from pytraj.externals.six.moves import range from pytraj.utils.cyutils import get_positive_idx _stop = get_positive_idx(stop, traj.top.n_residues) for i in range(start, _stop, step): j = ':' + str(i + 1) # example: traj[':3'] yield traj[j]
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)