Python testing.get_fn函数代码示例

def test_iterload_skip():
    files = [
        "frame0.nc",
        "frame0.h5",
        "frame0.xtc",
        "frame0.trr",
        "frame0.dcd",
        "frame0.binpos",
        "frame0.xyz",
        "frame0.lammpstrj",
    ]
    if not (on_win and on_py3):
        files.append("legacy_msmbuilder_trj0.lh5")

    err_msg = "failed for file %s with chunksize %i and skip %i"

    for file in files:
        for cs in [0, 1, 11, 100]:
            for skip in [0, 1, 20, 101]:
                print("testing file %s with skip=%i" % (file, skip))
                t_ref = md.load(get_fn(file), top=get_fn("native.pdb"))
                t = functools.reduce(
                    lambda a, b: a.join(b), md.iterload(get_fn(file), skip=skip, top=get_fn("native.pdb"), chunk=cs)
                )
                eq(t_ref.xyz[skip:], t.xyz, err_msg=err_msg % (file, cs, skip))
                eq(t_ref.time[skip:], t.time, err_msg=err_msg % (file, cs, skip))
                eq(t_ref.topology, t.topology, err_msg=err_msg % (file, cs, skip))

def test_load():
    filenames = [
        "frame0.xtc",
        "frame0.trr",
        "frame0.dcd",
        "frame0.binpos",
        "traj.h5",
        "frame0.nc",
        "traj.h5",
        "frame0.lammpstrj",
        "frame0.xyz",
    ]
    num_block = 3
    for filename in filenames:
        t0 = md.load(get_fn(filename), top=nat, discard_overlapping_frames=True)
        t1 = md.load(get_fn(filename), top=nat, discard_overlapping_frames=False)
        t2 = md.load([get_fn(filename) for i in xrange(num_block)], top=nat, discard_overlapping_frames=False)
        t3 = md.load([get_fn(filename) for i in xrange(num_block)], top=nat, discard_overlapping_frames=True)

        # these don't actually overlap, so discard_overlapping_frames should
        # have no effect. the overlap is between the last frame of one and the
        # first frame of the next.
        yield lambda: eq(t0.n_frames, t1.n_frames)
        yield lambda: eq(t0.n_frames * num_block, t2.n_frames)
        yield lambda: eq(t3.n_frames, t2.n_frames)

def test_seek():
    reference = BINPOSTrajectoryFile(get_fn('frame0.binpos')).read()
    with BINPOSTrajectoryFile(get_fn('frame0.binpos')) as f:
        xyz = f.read(1)[0]
        eq(xyz, reference[0])
        eq(f.tell(), 1)
        
        xyz = f.read(1)[0]
        eq(xyz, reference[1])
        eq(f.tell(), 2)
        
        f.seek(0)
        eq(f.tell(), 0)
        xyz = f.read(1)[0]
        eq(f.tell(), 1)
        eq(xyz, reference[0])
        
        f.seek(5)
        eq(f.read(1)[0], reference[5])
        eq(f.tell(), 6)

        f.seek(-5, 1)
        eq(f.tell(), 1)
        xyz = f.read(1)[0]
        eq(xyz, reference[1])

        f.seek(0, 2)
        eq(f.tell(), len(reference))

def test_dihedral_pbc():
    traj_uncorrected = md.load(get_fn('1am7_uncorrected.xtc'), top=get_fn('1am7_protein.pdb'))
    traj_corrected = md.load(get_fn('1am7_corrected.xtc'), top=get_fn('1am7_protein.pdb'))

    epsilon = 6E-3

    ang1 = md.geometry.compute_phi(traj_uncorrected, opt=False, periodic=True)[1]
    ang2 = md.geometry.compute_phi(traj_corrected, opt=False, periodic=True)[1]
    assert np.max(np.abs(ang1 - ang2)) < epsilon
    
    ang1 = md.geometry.compute_phi(traj_uncorrected, opt=True, periodic=True)[1]
    ang2 = md.geometry.compute_phi(traj_corrected, opt=True, periodic=True)[1]
    assert np.max(np.abs(ang1 - ang2)) < epsilon

    ang1 = md.geometry.compute_phi(traj_uncorrected, opt=True, periodic=True)[1]
    ang2 = md.geometry.compute_phi(traj_corrected, opt=True, periodic=False)[1]
    assert np.max(np.abs(ang1 - ang2)) < epsilon

    ang1 = md.geometry.compute_phi(traj_uncorrected, opt=False, periodic=True)[1]
    ang2 = md.geometry.compute_phi(traj_corrected, opt=False, periodic=False)[1]
    assert np.max(np.abs(ang1 - ang2)) < epsilon

    ang1 = md.geometry.compute_phi(traj_uncorrected, opt=True, periodic=True)[1]
    ang2 = md.geometry.compute_phi(traj_corrected, opt=False, periodic=False)[1]
    assert np.max(np.abs(ang1 - ang2)) < epsilon

    ang1 = md.geometry.compute_phi(traj_uncorrected, opt=False, periodic=True)[1]
    ang2 = md.geometry.compute_phi(traj_corrected, opt=True, periodic=True)[1]
    assert np.max(np.abs(ang1 - ang2)) < epsilon
    
    ang1 = md.geometry.compute_phi(traj_uncorrected, opt=False, periodic=True)[1]
    ang2 = md.geometry.compute_phi(traj_corrected, opt=True, periodic=False)[1]
    assert np.max(np.abs(ang1 - ang2)) < epsilon

def test_load_psf():
    top = psf.load_psf(get_fn('ala_ala_ala.psf'))
    ref_top = md.load(get_fn('ala_ala_ala.pdb')).topology
    eq(top, ref_top)
    # Test the XPLOR psf format parsing
    top2 = psf.load_psf(get_fn('ala_ala_ala.xpsf'))
    eq(top2, ref_top)

def test_seek_read_mode():
    # Test the seek/tell capacity of the different TrajectoryFile objects in
    # read mode. Basically, we just seek around the files and read different
    # segments, keeping track of our location manually and checking with both
    # tell() and by checking that the right coordinates are actually returned
    files = [(md.formats.NetCDFTrajectoryFile, 'frame0.nc'),
             (md.formats.HDF5TrajectoryFile, 'frame0.h5'),
             (md.formats.XTCTrajectoryFile, 'frame0.xtc'),
             (md.formats.TRRTrajectoryFile, 'frame0.trr'),
             (md.formats.DCDTrajectoryFile, 'frame0.dcd'),
             (md.formats.MDCRDTrajectoryFile, 'frame0.mdcrd'),
             (md.formats.BINPOSTrajectoryFile, 'frame0.binpos'),
             (md.formats.LH5TrajectoryFile, 'legacy_msmbuilder_trj0.lh5'),
             (md.formats.DTRTrajectoryFile,'frame0.dtr/clickme.dtr'),]

    for a, b in files:
        point = 0
        xyz = md.load(get_fn(b), top=get_fn('native.pdb')).xyz
        length = len(xyz)
        kwargs = {}
        if a is md.formats.MDCRDTrajectoryFile:
            kwargs = {'n_atoms': 22}

        with a(get_fn(b), **kwargs) as f:
            for i in range(100):
                r = np.random.rand()
                if r < 0.25:
                    offset = np.random.randint(-5, 5)
                    if 0 < point + offset < length:
                        point += offset
                        f.seek(offset, 1)
                    else:
                        f.seek(0)
                        point = 0
                if r < 0.5:
                    offset = np.random.randint(1, 10)
                    if point + offset < length:
                        read = f.read(offset)
                        if a not in [md.formats.BINPOSTrajectoryFile, md.formats.LH5TrajectoryFile]:
                            read = read[0]
                        readlength = len(read)
                        read = mdtraj.utils.in_units_of(read, f.distance_unit, 'nanometers')
                        eq(xyz[point:point+offset], read)
                        point += readlength
                elif r < 0.75:
                    offset = np.random.randint(low=-100, high=0)
                    try:
                        f.seek(offset, 2)
                        point = length + offset
                    except NotImplementedError:
                        # not all of the *TrajectoryFiles currently support
                        # seeking from the end, so we'll let this pass if they
                        # say that they dont implement this.
                        pass
                else:
                    offset = np.random.randint(100)
                    f.seek(offset, 0)
                    point = offset

                eq(f.tell(), point)

def test_read_1():
    with MDCRDTrajectoryFile(get_fn("frame0.mdcrd"), n_atoms=22) as f:
        xyz, _ = f.read()
    with MDCRDTrajectoryFile(get_fn("frame0.mdcrd"), n_atoms=22) as f:
        xyz3, _ = f.read(stride=3)

    eq(xyz[::3], xyz3)

def test_atom_indices_1():
    atom_indices = np.arange(10)
    top = md.load(get_fn("native.pdb"))
    t0 = md.load(get_fn("frame0.mdcrd"), top=top)
    t1 = md.load(get_fn("frame0.mdcrd"), top=top, atom_indices=atom_indices)

    eq(t0.xyz[:, atom_indices], t1.xyz)

def test_seek():
    reference = TRRTrajectoryFile(get_fn("frame0.trr")).read()[0]
    with TRRTrajectoryFile(get_fn("frame0.trr")) as f:
        eq(len(f), len(reference))
        eq(len(f.offsets), len(reference))

        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_seek():
    reference = XTCTrajectoryFile(get_fn('frame0.xtc')).read()[0]
    with XTCTrajectoryFile(get_fn('frame0.xtc')) 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_seek():
    reference = XTCTrajectoryFile(get_fn('frame0.xtc')).read()[0]
    with XTCTrajectoryFile(get_fn('frame0.xtc')) 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) # offset array is going to be built
        assert len(f.offsets) == len(reference)
        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_kappa():
    traj = md.load(get_fn("tip3p_300K_1ATM.xtc"), top=get_fn("tip3p_300K_1ATM.pdb"))
    kappa = md.geometry.isothermal_compressability_kappa_T(traj, temperature)
    reference = 2.05427E-10 * 1E5  # m^3 / J to 1 / bar.  Data from gromacs.  See above comment

    # 20% tolerance.
    assert abs((kappa - reference) / reference) < 2E-1, "Compressability tolerance not met!"

def deprecated_test_fah_core17_1():
    from mdtraj.utils import six
    from mdtraj.testing import get_fn, eq
    filename = get_fn('frame0.xtc')
    tempdir = tempfile.mkdtemp()
    tar_filename = os.path.join(tempdir, "results-000.tar.bz2")
    archive = tarfile.open(tar_filename, mode='w:bz2')

    tar = tarfile.open(tar_filename, "w:bz2")
    tar.add(filename, arcname="positions.xtc")
    tar.close()

    shutil.copy(tar_filename, os.path.join(tempdir, "results-001.tar.bz2"))

    trj0 = md.load(get_fn("frame0.xtc"), top=get_fn("frame0.h5"))
    output_filename = os.path.join(tempdir, "traj.h5")
    fah.concatenate_core17(tempdir, trj0, output_filename)

    trj = md.load(output_filename)
    eq(trj.n_atoms, trj0.n_atoms)
    eq(trj.n_frames, trj0.n_frames * 2)

    shutil.copy(tar_filename, os.path.join(tempdir, "results-002.tar.bz2"))

    fah.concatenate_core17(tempdir, trj0, output_filename)
    # Should notice the new file and append it to the HDF file.

    trj = md.load(output_filename)
    eq(trj.n_atoms, trj0.n_atoms)
    eq(trj.n_frames, trj0.n_frames * 3)

def test_lprmsd_5():
    t = md.load(get_fn('frame0.h5'))
    t1 = md.load(get_fn('frame0.h5'))

    r = md.rmsd(t, t1, 0)
    a = md.lprmsd(t, t1, 0, permute_groups=[[]], superpose=True)
    eq(a, r, decimal=3)

def test_lprmsd_3():
    # resolve rotation and permutation togetehr
    t1 = md.load(get_fn('alanine-dipeptide-explicit.pdb'))[0]
    t2 = md.load(get_fn('alanine-dipeptide-explicit.pdb'))[0]

    h2o_o_indices = [a.index for a in t1.topology.atoms if a.residue.name == 'HOH' and a.element.symbol == 'O'][0:20]
    h2o_h_indices = [a.index for a in t1.topology.atoms if a.residue.name == 'HOH' and a.element.symbol == 'H'][0:20]
    backbone_indices = [a.index for a in t1.topology.atoms if a.element.symbol in ['C', 'N']][:5]

    # scramble two groups of indices
    t2.xyz[:, random.permutation(h2o_o_indices)] = t2.xyz[:, h2o_o_indices]
    t2.xyz[:, random.permutation(h2o_h_indices)] = t2.xyz[:, h2o_h_indices]

    # offset the backbone indices slightly
    t2.xyz[:, backbone_indices] += 0.001 * random.randn(len(backbone_indices), 3)

    # rotate everything
    rot = rotation_matrix_from_quaternion(uniform_quaternion())
    t2.xyz[0].dot(rot)

    print('raw distinguishable indices', backbone_indices)
    atom_indices = np.concatenate((h2o_o_indices, backbone_indices))
    value = md.lprmsd(t2, t1, atom_indices=atom_indices, permute_groups=[h2o_o_indices])
    t1.xyz[:, h2o_o_indices] += random.randn(len(h2o_o_indices), 3)

    print('final value', value)
    assert value[0] < 1e-2