def test_read_3(get_fn): # Test read with n_frames fn_dtr = get_fn('frame0.dtr') dtr_traj = DTRTrajectoryFile(fn_dtr) dtr_traj.seek(1) xyz, times, cell_lens, cell_angles = dtr_traj.read(n_frames=900) eq(len(xyz), 500)
def test_read_7(): 'test two full read' reader = DTRTrajectoryFile(fn_dtr) xyz, times, cell_lens, cell_angles = reader.read() xyz, times, cell_lens, cell_angles = reader.read() eq(len(xyz), 0) eq(len(times), 0) eq(len(cell_lens), 0) eq(len(cell_angles), 0)
def test_read_1(get_fn): # test read with n_frame fn_dtr = get_fn('frame0.dtr') xyz, times, cell_lens, cell_angles = DTRTrajectoryFile(fn_dtr).read() xyz2, times2, cell_lens2, cell_angles2 = DTRTrajectoryFile(fn_dtr).read(n_frames=501) eq(xyz, xyz2) eq(times, times2) eq(cell_lens, cell_lens2) eq(cell_angles, cell_angles2)
def test_read_7(get_fn): # Test two full reads fn_dtr = get_fn('frame0.dtr') reader = DTRTrajectoryFile(fn_dtr) xyz, times, cell_lens, cell_angles = reader.read() xyz, times, cell_lens, cell_angles = reader.read() eq(len(xyz), 0) eq(len(times), 0) eq(len(cell_lens), 0) eq(len(cell_angles), 0)
def test_read_2(): """ test read with atom indices""" indices = np.array([0, 3, 12, 4]) xyz, times, cell_lens, cell_angles = DTRTrajectoryFile(fn_dtr).read() xyz2, times2, cell_lens2, cell_angles2 = DTRTrajectoryFile(fn_dtr).read( atom_indices=indices) eq(xyz[:, indices, :], xyz2) eq(times, times2) eq(cell_lens, cell_lens2) eq(cell_angles, cell_angles2)
def test_read_2(get_fn): # test read with atom indices fn_dtr = get_fn('frame0.dtr') indices = np.array([0, 3, 12, 4]) xyz, times, cell_lens, cell_angles = DTRTrajectoryFile(fn_dtr).read() xyz2, times2, cell_lens2, cell_angles2 = DTRTrajectoryFile(fn_dtr).read(atom_indices=indices) eq(xyz[:, indices, :], xyz2) eq(times, times2) eq(cell_lens, cell_lens2) eq(cell_angles, cell_angles2)
def test_read(get_fn): # Test the default read and compare against reference trajectory in dcd format fn_dtr = get_fn('frame0.dtr') fn_dcd = get_fn('frame0.dcd') dtr_traj = DTRTrajectoryFile(fn_dtr) eq(len(dtr_traj), 501) xyz, times, cell_lens, cell_angles = dtr_traj.read() xyz2, cell_lens2, cell_angles2 = DCDTrajectoryFile(fn_dcd).read() eq(xyz, xyz2) eq(cell_lens, cell_lens2) eq(cell_angles, cell_angles2)
def test_read(): """ test the default read and compare against reference trajectory in dcd format """ dtr_traj = DTRTrajectoryFile(fn_dtr) eq(len(dtr_traj), 501) xyz, times, cell_lens, cell_angles = dtr_traj.read() xyz2, cell_lens2, cell_angles2 = DCDTrajectoryFile(fn_dcd).read() eq(xyz, xyz2) eq(cell_lens, cell_lens2) eq(cell_angles, cell_angles2)
def test_read_stride(): "Read dtr with stride" with DTRTrajectoryFile(fn_dtr) as f: xyz1, times1, box_lengths1, box_angles1 = f.read() with DTRTrajectoryFile(fn_dtr) as f: xyz2, times2, box_lengths2, box_angles2 = f.read(stride=2) yield lambda: eq(xyz1[::2], xyz2) yield lambda: eq(times1[::2], times2) yield lambda: eq(box_lengths1[::2], box_lengths2) yield lambda: eq(box_angles1[::2], box_angles2)
def test_read_5(): "check streaming read of frames 1 at a time" xyz_ref, times_ref, box_lengths_ref, box_angles_ref = DTRTrajectoryFile(fn_dtr).read() reader = DTRTrajectoryFile(fn_dtr) for i in range(len(xyz_ref)): xyz, times, box_lenths, box_angles = reader.read(1) eq(xyz_ref[np.newaxis, i], xyz) eq(times_ref[np.newaxis, i], times) eq(box_lengths_ref[np.newaxis, i], box_lenths) eq(box_angles_ref[np.newaxis, i], box_angles)
def test_read_4(get_fn): # Read dtr with stride and n_frames fn_dtr = get_fn('frame0.dtr') with DTRTrajectoryFile(fn_dtr) as f: xyz1, times1, box_lengths1, box_angles1 = f.read() with DTRTrajectoryFile(fn_dtr) as f: xyz2, times2, box_lengths2, box_angles2 = f.read(n_frames=300, stride=2) assert eq(xyz1[::2], xyz2) assert eq(times1[::2], times2) assert eq(box_lengths1[::2], box_lengths2) assert eq(box_angles1[::2], box_angles2)
def test_write_1(): "test write" xyz, times, cell_lens, cell_angles = DTRTrajectoryFile(fn_dtr).read() xyz += 1 DTRTrajectoryFile(temp, 'w').write(xyz,cell_lengths=cell_lens, cell_angles=cell_angles, times=times) xyz2, times2, cell_lens2, cell_angles2 = DTRTrajectoryFile(temp).read() eq(xyz, xyz2) eq(times, times2) eq(cell_lens, cell_lens2) eq(cell_angles, cell_angles2)
def test_write_1(tmpdir, get_fn): # Test write fn_dtr = get_fn('frame0.dtr') fn = '{}/x.dtr'.format(tmpdir) xyz, times, cell_lens, cell_angles = DTRTrajectoryFile(fn_dtr).read() xyz += 1 DTRTrajectoryFile(fn, 'w').write(xyz, cell_lengths=cell_lens, cell_angles=cell_angles, times=times) xyz2, times2, cell_lens2, cell_angles2 = DTRTrajectoryFile(fn).read() eq(xyz, xyz2) eq(times, times2) eq(cell_lens, cell_lens2) eq(cell_angles, cell_angles2)
def test_read_4(): """Read dtr with stride and n_frames""" # dtr_traj = DTRTrajectoryFile(fn_dtr) # dtr_traj.seek(1) # xyz, times, cell_lens, cell_angles = dtr_traj.read(n_frames=300, stride=2) # eq(len(xyz), 251) with DTRTrajectoryFile(fn_dtr) as f: xyz1, times1, box_lengths1, box_angles1 = f.read() with DTRTrajectoryFile(fn_dtr) as f: xyz2, times2, box_lengths2, box_angles2 = f.read(n_frames=300, stride=2) yield lambda: eq(xyz1[::2], xyz2) yield lambda: eq(times1[::2], times2) yield lambda: eq(box_lengths1[::2], box_lengths2) yield lambda: eq(box_angles1[::2], box_angles2)
def test_write_3(): "test a random write operation" xyz = np.array(np.random.uniform(low=-50, high=-50, size=(3, 17, 3)), dtype=np.float32) times = np.array([1, 23.0, 48.0], dtype=np.float64) cell_lengths=np.array(np.random.uniform(low=100, high=200, size=(3, 3)), dtype=np.float32) cell_angles=np.array([[90, 90, 90], [80, 100, 120], [120, 90, 80]], dtype=np.float32) with DTRTrajectoryFile(temp, 'w') as f: f.write(xyz, cell_lengths=cell_lengths, cell_angles=cell_angles, times=times) with DTRTrajectoryFile(temp) as f: xyz2, times2, cell_lengths2, cell_angles2 = f.read() eq(xyz, xyz2)
def test_write_2(): """ test two separate write call """ xyz, times, cell_lens, cell_angles = DTRTrajectoryFile(fn_dtr).read() writer = DTRTrajectoryFile(temp, 'w') writer.write(xyz,cell_lengths=cell_lens, cell_angles=cell_angles, times=times) n_frames = len(xyz) times += 50.0 writer.write(xyz,cell_lengths=cell_lens, cell_angles=cell_angles, times=times) # # try to write frames with different number of atoms # assert_raises(ValueError, writer.write, xyz[:,10:,:], # cell_lengths=cell_lens, # cell_angles=cell_angles, # times=times) writer.close() xyz2, times2, cell_lens2, cell_angles2 = DTRTrajectoryFile(temp).read() eq(len(xyz2), n_frames*2) eq(xyz, xyz2[n_frames:]) eq(times, times2[n_frames:]) eq(cell_lens, cell_lens2[n_frames:]) eq(cell_angles, cell_angles2[n_frames:])
def test_read_6(): "DTRReader: check streaming read followed by reading the 'rest'" xyz_ref, times_ref, box_lengths_ref, box_angles_ref = DTRTrajectoryFile(fn_dtr).read() reader = DTRTrajectoryFile(fn_dtr) for i in range(int(len(xyz_ref)/2)): xyz, times, box_lenths, box_angles = reader.read(1) eq(xyz_ref[np.newaxis, i], xyz) eq(times_ref[np.newaxis, i], times) eq(box_lengths_ref[np.newaxis, i], box_lenths) eq(box_angles_ref[np.newaxis, i], box_angles) xyz_rest, times_rest, box_rest, angles_rest = reader.read() yield lambda: eq(xyz_ref[i+1:], xyz_rest) yield lambda: eq(times_ref[i+1:], times_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_4(): "test write error" xyz = np.array(np.random.uniform(low=-50, high=-50, size=(3, 17, 3)), dtype=np.float32) times = np.array([1, 23.0, 48.0], dtype=np.float64) cell_lengths=np.array(np.random.uniform(low=100, high=200, size=(3, 3)), dtype=np.float32) cell_angles=np.array([[90, 90, 90], [80, 100, 120], [120, 90, 80]], dtype=np.float32) bad_times = np.array([21, 3.0, 48.0], dtype=np.float64) f = DTRTrajectoryFile(temp, 'w') assert_raises(ValueError, f.write, xyz, cell_lengths=cell_lengths) assert_raises(ValueError, f.write, xyz, cell_angles=cell_angles) assert_raises(ValueError, f.write, xyz, times=times) assert_raises(ValueError, f.write, xyz, cell_lengths=cell_lengths, cell_angles=cell_angles, times=bad_times) f.close()
def test_read_6(get_fn): # Check streaming read followed by reading the 'rest' fn_dtr = get_fn('frame0.dtr') xyz_ref, times_ref, box_lengths_ref, box_angles_ref = DTRTrajectoryFile(fn_dtr).read() reader = DTRTrajectoryFile(fn_dtr) for i in range(len(xyz_ref) // 2): xyz, times, box_lenths, box_angles = reader.read(1) eq(xyz_ref[np.newaxis, i], xyz) eq(times_ref[np.newaxis, i], times) eq(box_lengths_ref[np.newaxis, i], box_lenths) eq(box_angles_ref[np.newaxis, i], box_angles) xyz_rest, times_rest, box_rest, angles_rest = reader.read() i = len(xyz_ref) // 2 assert eq(xyz_ref[i:], xyz_rest) assert eq(times_ref[i:], times_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_tell(): with DTRTrajectoryFile(fn_dtr) as f: last = len(f) eq(f.tell(), 0) f.read(2) eq(f.tell(), 2) f.read(100) eq(f.tell(), 102) f.seek(600) eq(f.tell(), last)
def test_write_2(tmpdir, get_fn): # Test two separate write call fn_dtr = get_fn('frame0.dtr') fn = '{}/x.dtr'.format(tmpdir) xyz, times, cell_lens, cell_angles = DTRTrajectoryFile(fn_dtr).read() writer = DTRTrajectoryFile(fn, 'w') writer.write(xyz, cell_lengths=cell_lens, cell_angles=cell_angles, times=times) n_frames = len(xyz) times += 50.0 writer.write(xyz, cell_lengths=cell_lens, cell_angles=cell_angles, times=times) writer.close() xyz2, times2, cell_lens2, cell_angles2 = DTRTrajectoryFile(fn).read() eq(len(xyz2), n_frames * 2) eq(xyz, xyz2[n_frames:]) eq(times, times2[n_frames:]) eq(cell_lens, cell_lens2[n_frames:]) eq(cell_angles, cell_angles2[n_frames:])
def test_tell(get_fn): fn_dtr = get_fn('frame0.dtr') with DTRTrajectoryFile(fn_dtr) as f: last = len(f) eq(f.tell(), 0) f.read(2) eq(f.tell(), 2) f.read(100) eq(f.tell(), 102) f.seek(600) eq(f.tell(), last)
def test_read_3(): """test read with n_frames""" dtr_traj = DTRTrajectoryFile(fn_dtr) dtr_traj.seek(1) xyz, times, cell_lens, cell_angles = dtr_traj.read(n_frames=900) eq(len(xyz), 500)
def test_write_4(tmpdir): # Test write error xyz = np.array(np.random.uniform(low=-50, high=-50, size=(3, 17, 3)), dtype=np.float32) times = np.array([1, 23.0, 48.0], dtype=np.float64) cell_lengths = np.array(np.random.uniform(low=100, high=200, size=(3, 3)), dtype=np.float32) cell_angles = np.array([[90, 90, 90], [80, 100, 120], [120, 90, 80]], dtype=np.float32) bad_times = np.array([21, 3.0, 48.0], dtype=np.float64) fn = '{}/x.dtr'.format(tmpdir) f = DTRTrajectoryFile(fn, 'w') with pytest.raises(ValueError): f.write(xyz, cell_lengths=cell_lengths) with pytest.raises(ValueError): f.write(xyz, cell_angles=cell_angles) with pytest.raises(ValueError): f.write(xyz, times=times) with pytest.raises(ValueError): f.write(xyz, cell_lengths=cell_lengths, cell_angles=cell_angles, times=bad_times) f.close()
def test_read_closed(): f = DTRTrajectoryFile(fn_dtr) f.close() f.read()
def test_read_closed(get_fn): fn_dtr = get_fn('frame0.dtr') with pytest.raises(IOError): f = DTRTrajectoryFile(fn_dtr) f.close() f.read()