def test_read_after_close():
f = NetCDFTrajectoryFile(get_fn('mdcrd.nc'))
yield lambda: eq(f.n_atoms, 223)
yield lambda: eq(f.n_frames, 101)
f.close()
# should be an ioerror if you read a file that's closed
assert_raises(IOError, lambda: f.read())
def test_read_after_close():
f = NetCDFTrajectoryFile(get_fn('mdcrd.nc'))
yield lambda: eq(f.n_atoms, 223)
yield lambda: eq(f.n_frames, 101)
f.close()
# should be an ioerror if you read a file that's closed
assert_raises(IOError, lambda: f.read())
def test_read_write_1():
xyz = np.random.randn(100, 3, 3)
time = np.random.randn(100)
boxlengths = np.random.randn(100, 3)
boxangles = np.random.randn(100, 3)
with NetCDFTrajectoryFile(temp, 'w', force_overwrite=True) as f:
f.write(xyz, time, boxlengths, boxangles)
with NetCDFTrajectoryFile(temp) as f:
a, b, c, d = f.read()
yield lambda: eq(a, xyz)
yield lambda: eq(b, time)
yield lambda: eq(c, boxlengths)
yield lambda: eq(d, boxangles)
def test_read_chunk_1(get_fn):
with NetCDFTrajectoryFile(get_fn('mdcrd.nc')) as f:
a, b, c, d = f.read(10)
e, f, g, h = f.read()
assert eq(len(a), 10)
assert eq(len(b), 10)
assert eq(len(e), 101 - 10)
assert eq(len(f), 101 - 10)
xyz = NetCDFTrajectoryFile(get_fn('mdcrd.nc')).read()[0]
assert eq(a, xyz[0:10])
assert eq(e, xyz[10:])
def test_read_chunk_2():
with NetCDFTrajectoryFile(get_fn('mdcrd.nc')) as f:
a, b, c, d = f.read(10)
e, f, g, h = f.read(100000000000)
yield lambda: eq(len(a), 10)
yield lambda: eq(len(b), 10)
yield lambda: eq(len(e), 101 - 10)
yield lambda: eq(len(f), 101 - 10)
xyz = NetCDFTrajectoryFile(get_fn('mdcrd.nc')).read()[0]
yield lambda: eq(a, xyz[0:10])
yield lambda: eq(e, xyz[10:])
def test_shape(get_fn):
xyz, time, boxlength, boxangles = NetCDFTrajectoryFile(get_fn('mdcrd.nc')).read()
assert eq(xyz.shape, (101, 223, 3))
assert eq(time.shape, (101,))
assert eq(boxlength, None)
assert eq(boxangles, None)
def test_do_overwrite():
with open(temp, 'w') as f:
f.write('a')
with pytest.raises(IOError):
with NetCDFTrajectoryFile(temp, 'w', force_overwrite=False) as f:
f.write(np.random.randn(10, 5, 3))
def test_read_write_2():
xyz = np.random.randn(5, 22, 3)
time = np.random.randn(5)
with NetCDFTrajectoryFile(temp, 'w', force_overwrite=True) as f:
f.write(xyz, time)
with NetCDFTrajectoryFile(temp) as f:
rcoord, rtime, rlengths, rangles = f.read()
yield lambda: eq(rcoord, xyz)
yield lambda: eq(rtime, time)
yield lambda: eq(rlengths, None)
yield lambda: eq(rangles, None)
t = md.load(temp, top=get_fn('native.pdb'))
eq(t.unitcell_angles, None)
eq(t.unitcell_lengths, None)
def _nc_a_traj(info, gen, gen_fn, has_overlapping_frames):
out_fn = "{outdir}/{gen}.{outext}".format(gen=gen, **info['cnv2'])
log.debug("Converting to netcdf {} {}".format(gen_fn, out_fn))
with XTCTrajectoryFile(gen_fn, 'r') as xtc:
with NetCDFTrajectoryFile(out_fn, 'w') as nc:
_nc_a_chunk(xtc, nc, has_overlapping_frames)
return out_fn
def test_write_3():
xyz = np.random.randn(100, 3, 3)
time = np.random.randn(100)
with NetCDFTrajectoryFile(temp, 'w', force_overwrite=True) as f:
# you can't supply cell_lengths without cell_angles
assert_raises(ValueError, lambda: f.write(np.random.randn(100, 3, 3), cell_lengths=np.random.randn(100, 3)))
# or the other way aroun
assert_raises(ValueError, lambda: f.write(np.random.randn(100, 3, 3), cell_angles=np.random.randn(100, 3)))
def test_read_write_25():
xyz = np.random.randn(100, 3, 3)
time = np.random.randn(100)
with NetCDFTrajectoryFile(temp, 'w', force_overwrite=True) as f:
f.write(xyz, time)
f.write(xyz, time)
with NetCDFTrajectoryFile(temp) as f:
a, b, c, d = f.read()
yield lambda: eq(a[0:100], xyz)
yield lambda: eq(b[0:100], time)
yield lambda: eq(c, None)
yield lambda: eq(d, None)
yield lambda: eq(a[100:], xyz)
yield lambda: eq(b[100:], time)
yield lambda: eq(c, None)
yield lambda: eq(d, None)
def test_ragged_1():
# try first writing no cell angles/lengths, and then adding some
xyz = np.random.randn(100, 3, 3)
time = np.random.randn(100)
cell_lengths = np.random.randn(100, 3)
cell_angles = np.random.randn(100, 3)
with NetCDFTrajectoryFile(temp, 'w', force_overwrite=True) as f:
f.write(xyz, time)
assert_raises(ValueError, lambda: f.write(xyz, time, cell_lengths, cell_angles))
def test_ragged_2():
# try first writing no cell angles/lengths, and then adding some
xyz = np.random.randn(100, 3, 3)
time = np.random.randn(100)
cell_lengths = np.random.randn(100, 3)
cell_angles = np.random.randn(100, 3)
# from mdtraj.formats import HDF5TrajectoryFile
with NetCDFTrajectoryFile(temp, 'w', force_overwrite=True) as f:
f.write(xyz, time, cell_lengths, cell_angles)
with pytest.raises(ValueError):
f.write(xyz, time)
def test_read_after_close(get_fn):
f = NetCDFTrajectoryFile(get_fn('mdcrd.nc'))
assert eq(f.n_atoms, 223)
assert eq(f.n_frames, 101)
f.close()
# should be an IOError if you read a file that's closed
with pytest.raises(IOError):
f.read()
def test_reporter():
tempdir = os.path.join(dir, 'test1')
os.makedirs(tempdir)
pdb = PDBFile(get_fn('native.pdb'))
forcefield = ForceField('amber99sbildn.xml', 'amber99_obc.xml')
# NO PERIODIC BOUNARY CONDITIONS
system = forcefield.createSystem(pdb.topology,
nonbondedMethod=CutoffNonPeriodic,
nonbondedCutoff=1.0 * nanometers,
constraints=HBonds,
rigidWater=True)
integrator = LangevinIntegrator(300 * kelvin, 1.0 / picoseconds,
2.0 * femtoseconds)
integrator.setConstraintTolerance(0.00001)
platform = Platform.getPlatformByName('Reference')
simulation = Simulation(pdb.topology, system, integrator, platform)
simulation.context.setPositions(pdb.positions)
simulation.context.setVelocitiesToTemperature(300 * kelvin)
hdf5file = os.path.join(tempdir, 'traj.h5')
ncfile = os.path.join(tempdir, 'traj.nc')
dcdfile = os.path.join(tempdir, 'traj.dcd')
reporter = HDF5Reporter(hdf5file,
2,
coordinates=True,
time=True,
cell=True,
potentialEnergy=True,
kineticEnergy=True,
temperature=True,
velocities=True)
reporter2 = NetCDFReporter(ncfile,
2,
coordinates=True,
time=True,
cell=True)
reporter3 = DCDReporter(dcdfile, 2)
simulation.reporters.append(reporter)
simulation.reporters.append(reporter2)
simulation.reporters.append(reporter3)
simulation.step(100)
reporter.close()
reporter2.close()
with HDF5TrajectoryFile(hdf5file) as f:
got = f.read()
yield lambda: eq(got.temperature.shape, (50, ))
yield lambda: eq(got.potentialEnergy.shape, (50, ))
yield lambda: eq(got.kineticEnergy.shape, (50, ))
yield lambda: eq(got.coordinates.shape, (50, 22, 3))
yield lambda: eq(got.velocities.shape, (50, 22, 3))
yield lambda: eq(got.cell_lengths, None)
yield lambda: eq(got.cell_angles, None)
yield lambda: eq(got.time, 0.002 * 2 * (1 + np.arange(50)))
yield lambda: f.topology == md.load(get_fn('native.pdb')).top
with NetCDFTrajectoryFile(ncfile) as f:
xyz, time, cell_lengths, cell_angles = f.read()
yield lambda: eq(cell_lengths, None)
yield lambda: eq(cell_angles, None)
yield lambda: eq(time, 0.002 * 2 * (1 + np.arange(50)))
hdf5_traj = md.load(hdf5file)
dcd_traj = md.load(dcdfile, top=get_fn('native.pdb'))
netcdf_traj = md.load(ncfile, top=get_fn('native.pdb'))
# we don't have to convert units here, because md.load already
# handles that
assert hdf5_traj.unitcell_vectors is None
yield lambda: eq(hdf5_traj.xyz, netcdf_traj.xyz)
yield lambda: eq(hdf5_traj.unitcell_vectors, netcdf_traj.unitcell_vectors)
yield lambda: eq(hdf5_traj.time, netcdf_traj.time)
yield lambda: eq(dcd_traj.xyz, hdf5_traj.xyz)
def test_reporter_subset():
tempdir = os.path.join(dir, 'test2')
os.makedirs(tempdir)
pdb = PDBFile(get_fn('native2.pdb'))
pdb.topology.setUnitCellDimensions([2, 2, 2])
forcefield = ForceField('amber99sbildn.xml', 'amber99_obc.xml')
system = forcefield.createSystem(pdb.topology,
nonbondedMethod=CutoffPeriodic,
nonbondedCutoff=1 * nanometers,
constraints=HBonds,
rigidWater=True)
integrator = LangevinIntegrator(300 * kelvin, 1.0 / picoseconds,
2.0 * femtoseconds)
integrator.setConstraintTolerance(0.00001)
platform = Platform.getPlatformByName('Reference')
simulation = Simulation(pdb.topology, system, integrator, platform)
simulation.context.setPositions(pdb.positions)
simulation.context.setVelocitiesToTemperature(300 * kelvin)
hdf5file = os.path.join(tempdir, 'traj.h5')
ncfile = os.path.join(tempdir, 'traj.nc')
dcdfile = os.path.join(tempdir, 'traj.dcd')
atomSubset = [0, 1, 2, 4, 5]
reporter = HDF5Reporter(hdf5file,
2,
coordinates=True,
time=True,
cell=True,
potentialEnergy=True,
kineticEnergy=True,
temperature=True,
velocities=True,
atomSubset=atomSubset)
reporter2 = NetCDFReporter(ncfile,
2,
coordinates=True,
time=True,
cell=True,
atomSubset=atomSubset)
reporter3 = DCDReporter(dcdfile, 2, atomSubset=atomSubset)
simulation.reporters.append(reporter)
simulation.reporters.append(reporter2)
simulation.reporters.append(reporter3)
simulation.step(100)
reporter.close()
reporter2.close()
reporter3.close()
t = md.load(get_fn('native.pdb'))
t.restrict_atoms(atomSubset)
with HDF5TrajectoryFile(hdf5file) as f:
got = f.read()
eq(got.temperature.shape, (50, ))
eq(got.potentialEnergy.shape, (50, ))
eq(got.kineticEnergy.shape, (50, ))
eq(got.coordinates.shape, (50, len(atomSubset), 3))
eq(got.velocities.shape, (50, len(atomSubset), 3))
eq(got.cell_lengths, 2 * np.ones((50, 3)))
eq(got.cell_angles, 90 * np.ones((50, 3)))
eq(got.time, 0.002 * 2 * (1 + np.arange(50)))
assert f.topology == md.load(get_fn('native.pdb'),
atom_indices=atomSubset).topology
with NetCDFTrajectoryFile(ncfile) as f:
xyz, time, cell_lengths, cell_angles = f.read()
eq(cell_lengths, 20 * np.ones((50, 3)))
eq(cell_angles, 90 * np.ones((50, 3)))
eq(time, 0.002 * 2 * (1 + np.arange(50)))
eq(xyz.shape, (50, len(atomSubset), 3))
hdf5_traj = md.load(hdf5file)
dcd_traj = md.load(dcdfile, top=hdf5_traj)
netcdf_traj = md.load(ncfile, top=hdf5_traj)
# we don't have to convert units here, because md.load already handles
# that
eq(hdf5_traj.xyz, netcdf_traj.xyz)
eq(hdf5_traj.unitcell_vectors, netcdf_traj.unitcell_vectors)
eq(hdf5_traj.time, netcdf_traj.time)
eq(dcd_traj.xyz, hdf5_traj.xyz)
eq(dcd_traj.unitcell_vectors, hdf5_traj.unitcell_vectors)
def test_read_chunk_3():
# too big of a chunk should not be an issue
a = NetCDFTrajectoryFile(get_fn('mdcrd.nc')).read(1000000000)
b = NetCDFTrajectoryFile(get_fn('mdcrd.nc')).read()
eq(a[0], b[0])
def test_n_atoms():
with NetCDFTrajectoryFile(temp, 'w', force_overwrite=True) as f:
f.write(np.random.randn(1, 11, 3))
with NetCDFTrajectoryFile(temp) as f:
eq(f.n_atoms, 11)
def test_do_overwrite():
with open(temp, 'w') as f:
f.write('a')
with NetCDFTrajectoryFile(temp, 'w', force_overwrite=True) as f:
f.write(np.random.randn(10, 5, 3))