def test_read_5():
with DCDTrajectoryFile(fn_dcd) as f:
xyz_ref, box_lengths_ref, box_angles_ref = f.read()
with DCDTrajectoryFile(fn_dcd) as f:
xyz, box_lengths, box_angles = f.read(atom_indices=[1, 2, 5])
yield lambda: eq(xyz_ref[:, [1, 2, 5], :], xyz)
def get_frames_num(filepath, file_type, ext=None):
if file_type == 'coor':
if ext is None:
traj = mdtraj_load(filepath)
else:
ext2 = ext.lower()
if ext2 == 'pdb':
traj = mdtraj_load_pdb(filepath)
else:
raise ValueError('Extension "' + ext2 + '" not implemented.')
numframes = traj.n_frames
elif file_type == 'traj':
if ext is None:
trajfile = mdtraj_open(filepath)
else:
ext2 = ext.lower()
if ext2 == 'dcd':
trajfile = DCDTrajectoryFile(filepath)
elif ext2 == 'xtc':
trajfile = XTCTrajectoryFile(filepath)
else:
raise ValueError('Extension "' + ext + '" not implemented.')
numframes = mdtraj_get_frames_num(trajfile)
trajfile.close()
return numframes
def test_seek():
reference = DCDTrajectoryFile(get_fn('frame0.dcd')).read()[0]
with DCDTrajectoryFile(get_fn('frame0.dcd')) as f:
eq(f.tell(), 0)
eq(f.read(1)[0][0], reference[0])
eq(f.tell(), 1)
xyz = f.read(1)[0][0]
eq(xyz, reference[1])
eq(f.tell(), 2)
f.seek(0)
eq(f.tell(), 0)
xyz = f.read(1)[0][0]
eq(f.tell(), 1)
eq(xyz, reference[0])
f.seek(5)
eq(f.read(1)[0][0], reference[5])
eq(f.tell(), 6)
f.seek(-5, 1)
eq(f.tell(), 1)
eq(f.read(1)[0][0], reference[1])
def test_read_5(get_fn):
fn_dcd = get_fn('frame0.dcd')
with DCDTrajectoryFile(fn_dcd) as f:
xyz_ref, box_lengths_ref, box_angles_ref = f.read()
with DCDTrajectoryFile(fn_dcd) as f:
xyz, box_lengths, box_angles = f.read(atom_indices=[1, 2, 5])
assert eq(xyz_ref[:, [1, 2, 5], :], xyz)
def test_read_6():
with DCDTrajectoryFile(fn_dcd) as f:
xyz_ref, box_lengths_ref, box_angles_ref = f.read()
with DCDTrajectoryFile(fn_dcd) as f:
xyz, box_lengths, box_angles = f.read(
atom_indices=slice(None, None, 2))
yield lambda: eq(xyz_ref[:, ::2, :], xyz)
def test_read_2():
"DCDReader: check nframes"
xyz1, box_lengths1, box_angles1 = DCDTrajectoryFile(fn_dcd).read()
xyz2, box_lengths2, box_angles2 = DCDTrajectoryFile(fn_dcd).read(10000)
yield lambda: eq(xyz1, xyz2)
yield lambda: eq(box_lengths1, box_lengths2)
yield lambda: eq(box_angles1, box_angles2)
def test_write_1():
xyz = np.array(np.random.randn(500, 10, 3), dtype=np.float32)
with DCDTrajectoryFile(temp, 'w') as f:
f.write(xyz)
with DCDTrajectoryFile(temp) as f:
xyz2 = f.read()[0]
eq(xyz, xyz2)
def test_read_2(get_fn):
# check nframes
fn_dcd = get_fn('frame0.dcd')
xyz1, box_lengths1, box_angles1 = DCDTrajectoryFile(fn_dcd).read()
xyz2, box_lengths2, box_angles2 = DCDTrajectoryFile(fn_dcd).read(10000)
assert eq(xyz1, xyz2)
assert eq(box_lengths1, box_lengths2)
assert eq(box_angles1, box_angles2)
def test_read_6(get_fn):
fn_dcd = get_fn('frame0.dcd')
with DCDTrajectoryFile(fn_dcd) as f:
xyz_ref, box_lengths_ref, box_angles_ref = f.read()
with DCDTrajectoryFile(fn_dcd) as f:
xyz, box_lengths, box_angles = f.read(
atom_indices=slice(None, None, 2))
assert eq(xyz_ref[:, ::2, :], xyz)
def test_write_0():
with DCDTrajectoryFile(fn_dcd) as f:
xyz = f.read()[0]
with DCDTrajectoryFile(temp, 'w') as f:
f.write(xyz)
with DCDTrajectoryFile(temp) as f:
xyz2 = f.read()[0]
eq(xyz, xyz2)
def test_write_1(tmpdir):
fn = '{}/x.dcd'.format(tmpdir)
xyz = np.array(np.random.randn(500, 10, 3), dtype=np.float32)
with DCDTrajectoryFile(fn, 'w') as f:
f.write(xyz)
with DCDTrajectoryFile(fn) as f:
xyz2 = f.read()[0]
eq(xyz, xyz2)
def test_read_stride_2():
"Read dcd with stride when n_frames is supplied (different code path)"
with DCDTrajectoryFile(fn_dcd) as f:
xyz1, box_lengths1, box_angles1 = f.read()
with DCDTrajectoryFile(fn_dcd) as f:
xyz2, box_lengths2, box_angles2 = f.read(n_frames=1000, stride=2)
yield lambda: eq(xyz1[::2], xyz2)
yield lambda: eq(box_lengths1[::2], box_lengths2)
yield lambda: eq(box_angles1[::2], box_angles2)
def test_read_3():
"DCDReader: check streaming read of frames 1 at a time"
xyz_ref, box_lengths_ref, box_angles_ref = DCDTrajectoryFile(fn_dcd).read()
reader = DCDTrajectoryFile(fn_dcd)
for i in range(len(xyz_ref)):
xyz, box_lenths, box_angles = reader.read(1)
eq(xyz_ref[np.newaxis, i], xyz)
eq(box_lengths_ref[np.newaxis, i], box_lenths)
eq(box_angles_ref[np.newaxis, i], box_angles)
def test_read_stride():
"Read dcd with stride"
with DCDTrajectoryFile(fn_dcd) as f:
xyz1, box_lengths1, box_angles1 = f.read()
with DCDTrajectoryFile(fn_dcd) as f:
xyz2, box_lengths2, box_angles2 = f.read(stride=2)
yield lambda: eq(xyz1[::2], xyz2)
yield lambda: eq(box_lengths1[::2], box_lengths2)
yield lambda: eq(box_angles1[::2], box_angles2)
def test_read_3():
"DCDReader: check streaming read of frames 1 at a time"
xyz_ref, box_lengths_ref, box_angles_ref = DCDTrajectoryFile(fn_dcd).read()
reader = DCDTrajectoryFile(fn_dcd)
for i in range(len(xyz_ref)):
xyz, box_lenths, box_angles = reader.read(1)
eq(xyz_ref[np.newaxis, i], xyz)
eq(box_lengths_ref[np.newaxis, i], box_lenths)
eq(box_angles_ref[np.newaxis, i], box_angles)
def test_read_stride(get_fn):
# Read dcd with stride
fn_dcd = get_fn('frame0.dcd')
with DCDTrajectoryFile(fn_dcd) as f:
xyz1, box_lengths1, box_angles1 = f.read()
with DCDTrajectoryFile(fn_dcd) as f:
xyz2, box_lengths2, box_angles2 = f.read(stride=2)
assert eq(xyz1[::2], xyz2)
assert eq(box_lengths1[::2], box_lengths2)
assert eq(box_angles1[::2], box_angles2)
def test_read_3(get_fn):
# Check streaming read of frames 1 at a time
fn_dcd = get_fn('frame0.dcd')
xyz_ref, box_lengths_ref, box_angles_ref = DCDTrajectoryFile(fn_dcd).read()
reader = DCDTrajectoryFile(fn_dcd)
for i in range(len(xyz_ref)):
xyz, box_lenths, box_angles = reader.read(1)
eq(xyz_ref[np.newaxis, i], xyz)
eq(box_lengths_ref[np.newaxis, i], box_lenths)
eq(box_angles_ref[np.newaxis, i], box_angles)
def test_read_stride_2(get_fn):
# Read dcd with stride when n_frames is supplied (different code path)
fn_dcd = get_fn('frame0.dcd')
with DCDTrajectoryFile(fn_dcd) as f:
xyz1, box_lengths1, box_angles1 = f.read()
with DCDTrajectoryFile(fn_dcd) as f:
xyz2, box_lengths2, box_angles2 = f.read(n_frames=1000, stride=2)
assert eq(xyz1[::2], xyz2)
assert eq(box_lengths1[::2], box_lengths2)
assert eq(box_angles1[::2], box_angles2)
def test_read_3(get_fn):
# Check streaming read of frames 1 at a time
fn_dcd = get_fn('frame0.dcd')
xyz_ref, box_lengths_ref, box_angles_ref = DCDTrajectoryFile(fn_dcd).read()
reader = DCDTrajectoryFile(fn_dcd)
for i in range(len(xyz_ref)):
xyz, box_lenths, box_angles = reader.read(1)
eq(xyz_ref[np.newaxis, i], xyz)
eq(box_lengths_ref[np.newaxis, i], box_lenths)
eq(box_angles_ref[np.newaxis, i], box_angles)
def test_write_0(tmpdir, get_fn):
fn_dcd = get_fn('frame0.dcd')
fn = '{}/x.dcd'.format(tmpdir)
with DCDTrajectoryFile(fn_dcd) as f:
xyz = f.read()[0]
with DCDTrajectoryFile(fn, 'w') as f:
f.write(xyz)
with DCDTrajectoryFile(fn) as f:
xyz2 = f.read()[0]
eq(xyz, xyz2)
def __init__(self, file, mode='r'):
super().__init__()
try:
from mdtraj.formats import DCDTrajectoryFile
except:
raise ImportError('Currently mstools use mdtraj to parse DCD format. Cannot import mdtraj')
if mode == 'r':
self._dcd = DCDTrajectoryFile(file, mode='r')
elif mode == 'w':
self._dcd = DCDTrajectoryFile(file, mode='w')
else:
raise Exception('Appending not supported for DCD')
def test_do_overwrite(tmpdir):
fn = '{}/x.dcd'.format(tmpdir)
with open(fn, 'w') as f:
f.write('a')
with DCDTrajectoryFile(fn, 'w', force_overwrite=True) as f:
f.write(np.random.randn(10, 5, 3))
def test_write_3():
"checking"
xyz = np.array(np.random.randn(500, 10, 3), dtype=np.float32)
box_lengths = 25 * np.ones((600, 3), dtype=np.float32)
with DCDTrajectoryFile(temp, 'w') as f:
f.write(xyz, box_lengths)
def write_trajectory_dcd(netcdf_filename, topology, pdb_trajectory_filename, dcd_trajectory_filename):
"""
Write trajectory.
Parameters
----------
netcdf_filename : str
NetCDF filename.
topology : Topology
Topology object
pdb_trajectory_filename : str
PDB trajectory output filename
dcd_trajectory_filename : str
Output trajectory filename.
"""
ncfile = netCDF4.Dataset(netcdf_filename, 'r')
[nsamples, nstates] = ncfile.variables['logZ'].shape
# Write reference.pdb file
from simtk.openmm.app import PDBFile
outfile = open(pdb_trajectory_filename, 'w')
positions = unit.Quantity(ncfile.variables['positions'][0,:,:], unit.nanometers)
PDBFile.writeFile(topology, positions, file=outfile)
outfile.close()
# TODO: Export as DCD trajectory with MDTraj
from mdtraj.formats import DCDTrajectoryFile
with DCDTrajectoryFile(dcd_trajectory_filename, 'w') as f:
f.write(ncfile.variables['positions'][:,:,:] * 10.0) # angstroms
def test_read_4():
"DCDReader: check streaming read followed by reading the 'rest'"
xyz_ref, box_lengths_ref, box_angles_ref = DCDTrajectoryFile(fn_dcd).read()
reader = DCDTrajectoryFile(fn_dcd)
for i in range(int(len(xyz_ref) / 2)):
xyz, box_lenths, box_angles = reader.read(1)
eq(xyz_ref[np.newaxis, i], xyz)
eq(box_lengths_ref[np.newaxis, i], box_lenths)
eq(box_angles_ref[np.newaxis, i], box_angles)
xyz_rest, box_rest, angles_rest = reader.read()
yield lambda: eq(xyz_ref[i + 1:], xyz_rest)
yield lambda: eq(box_lengths_ref[i + 1:], box_rest)
yield lambda: eq(box_angles_ref[i + 1:], angles_rest)
yield lambda: len(xyz_ref) == i + len(xyz_rest)
def test_dont_overwrite(tmpdir):
fn = '{}/x.dcd'.format(tmpdir)
with open(fn, 'w') as f:
f.write('a')
with pytest.raises(IOError):
with DCDTrajectoryFile(fn, 'w', force_overwrite=False) as f:
f.write(np.random.randn(10, 5, 3))
def test_read_4():
"DCDReader: check streaming read followed by reading the 'rest'"
xyz_ref, box_lengths_ref, box_angles_ref = DCDTrajectoryFile(fn_dcd).read()
reader = DCDTrajectoryFile(fn_dcd)
for i in range(int(len(xyz_ref) / 2)):
xyz, box_lenths, box_angles = reader.read(1)
eq(xyz_ref[np.newaxis, i], xyz)
eq(box_lengths_ref[np.newaxis, i], box_lenths)
eq(box_angles_ref[np.newaxis, i], box_angles)
xyz_rest, box_rest, angles_rest = reader.read()
yield lambda: eq(xyz_ref[i + 1 :], xyz_rest)
yield lambda: eq(box_lengths_ref[i + 1 :], box_rest)
yield lambda: eq(box_angles_ref[i + 1 :], angles_rest)
yield lambda: len(xyz_ref) == i + len(xyz_rest)
def test_write_3(tmpdir):
fn = '{}/x.dcd'.format(tmpdir)
xyz = np.array(np.random.randn(500, 10, 3), dtype=np.float32)
box_lengths = 25 * np.ones((600, 3), dtype=np.float32)
with DCDTrajectoryFile(fn, 'w') as f:
with pytest.raises(ValueError):
f.write(xyz, box_lengths)
def test_read_1():
with BINPOSTrajectoryFile(fn_binpos, chunk_size_multiplier=0.5) as f:
xyz = f.read()
with DCDTrajectoryFile(fn_dcd) as f:
xyz2 = f.read()[0]
xyz3 = io.loadh(get_fn('frame0.binpos.h5'), 'xyz')
yield lambda: eq(xyz[1:], xyz2)
yield lambda: eq(xyz, xyz3)
def test_tell():
with DCDTrajectoryFile(get_fn('frame0.dcd')) as f:
eq(f.tell(), 0)
f.read(101)
eq(f.tell(), 101)
f.read(3)
eq(f.tell(), 104)
def test_read_4(get_fn):
# Check streaming read followed by reading the 'rest'
fn_dcd = get_fn('frame0.dcd')
xyz_ref, box_lengths_ref, box_angles_ref = DCDTrajectoryFile(fn_dcd).read()
reader = DCDTrajectoryFile(fn_dcd)
for i in range(len(xyz_ref) // 2):
xyz, box_lenths, box_angles = reader.read(1)
eq(xyz_ref[np.newaxis, i], xyz)
eq(box_lengths_ref[np.newaxis, i], box_lenths)
eq(box_angles_ref[np.newaxis, i], box_angles)
xyz_rest, box_rest, angles_rest = reader.read()
i = len(xyz_ref) // 2
assert eq(xyz_ref[i:], xyz_rest)
assert eq(box_lengths_ref[i:], box_rest)
assert eq(box_angles_ref[i:], angles_rest)
assert len(xyz_ref) == i + len(xyz_rest)
def test_read_4(get_fn):
# Check streaming read followed by reading the 'rest'
fn_dcd = get_fn('frame0.dcd')
xyz_ref, box_lengths_ref, box_angles_ref = DCDTrajectoryFile(fn_dcd).read()
reader = DCDTrajectoryFile(fn_dcd)
for i in range(len(xyz_ref) // 2):
xyz, box_lenths, box_angles = reader.read(1)
eq(xyz_ref[np.newaxis, i], xyz)
eq(box_lengths_ref[np.newaxis, i], box_lenths)
eq(box_angles_ref[np.newaxis, i], box_angles)
xyz_rest, box_rest, angles_rest = reader.read()
i = len(xyz_ref) // 2
assert eq(xyz_ref[i:], xyz_rest)
assert eq(box_lengths_ref[i:], box_rest)
assert eq(box_angles_ref[i:], angles_rest)
assert len(xyz_ref) == i + len(xyz_rest)
def test_ragged_1():
# try first writing no cell angles/lengths, and then adding some
xyz = np.random.randn(100, 5, 3)
cell_lengths = np.random.randn(100, 3)
cell_angles = np.random.randn(100, 3)
with DCDTrajectoryFile(temp, 'w', force_overwrite=True) as f:
f.write(xyz)
assert_raises(ValueError,
lambda: f.write(xyz, cell_lengths, cell_angles))
def test_read_2(get_fn):
fn_binpos = get_fn('frame0.binpos')
fn_dcd = get_fn('frame0.dcd')
with BINPOSTrajectoryFile(fn_binpos, chunk_size_multiplier=10) as f:
xyz = f.read()
with DCDTrajectoryFile(fn_dcd) as f:
xyz2 = f.read()[0]
xyz3 = io.loadh(get_fn('frame0.binpos.h5'), 'xyz')
assert eq(xyz[1:], xyz2)
assert eq(xyz, xyz3)
def test_write_2():
xyz = np.array(np.random.randn(500, 10, 3), dtype=np.float32)
box_lengths = 25 * np.ones((500, 3), dtype=np.float32)
box_angles = 90 * np.ones((500, 3), dtype=np.float32)
box_lengths[0, 0] = 10.0
f = DCDTrajectoryFile(temp, "w")
f.write(xyz, box_lengths, box_angles)
f.close()
f = DCDTrajectoryFile(temp)
xyz2, box_lengths2, box_angles2 = f.read()
f.close()
yield lambda: eq(xyz, xyz2)
yield lambda: eq(box_lengths, box_lengths2)
yield lambda: eq(box_angles, box_angles2)
def test_write_2(tmpdir):
fn = '{}/x.dcd'.format(tmpdir)
xyz = np.array(np.random.randn(500, 10, 3), dtype=np.float32)
box_lengths = 25 * np.ones((500, 3), dtype=np.float32)
box_angles = 90 * np.ones((500, 3), dtype=np.float32)
box_lengths[0, 0] = 10.0
f = DCDTrajectoryFile(fn, 'w')
f.write(xyz, box_lengths, box_angles)
f.close()
f = DCDTrajectoryFile(fn)
xyz2, box_lengths2, box_angles2 = f.read()
f.close()
assert eq(xyz, xyz2)
assert eq(box_lengths, box_lengths2)
assert eq(box_angles, box_angles2)
def test_write_closed(get_fn):
fn_dcd = get_fn('frame0.dcd')
with pytest.raises(IOError):
f = DCDTrajectoryFile(fn_dcd, 'w')
f.close()
f.write(np.random.randn(10, 3, 3))
def test_read_closed(get_fn):
fn_dcd = get_fn('frame0.dcd')
with pytest.raises(IOError):
f = DCDTrajectoryFile(fn_dcd)
f.close()
f.read()
def test_write_closed():
f = DCDTrajectoryFile(fn_dcd, "w")
f.close()
f.write(np.random.randn(10, 3, 3))
def test_read_closed():
f = DCDTrajectoryFile(fn_dcd)
f.close()
f.read()
