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))