Python utils.enter_temp_directory函数代码示例

def _test_against_vmd(pdb):
    # this is probably not cross-platform compatible. I assume that the exact
    # path to this CHARMM topology that is included with VMD depends on
    # the install mechanism, especially for bundled mac or windows installers
    VMD_ROOT = os.path.join(os.path.dirname(os.path.realpath(VMD)), '..')
    top_paths = [os.path.join(r, f) for (r, _, fs) in os.walk(VMD_ROOT) for f in fs
         if 'top_all27_prot_lipid_na.inp' in f]
    assert len(top_paths) >= 0
    top = os.path.abspath(top_paths[0]).replace(" ", "\\ ")

    TEMPLATE = '''
package require psfgen
topology %(top)s
pdbalias residue HIS HSE
pdbalias atom ILE CD1 CD
segment U {pdb %(pdb)s}
coordpdb %(pdb)s U
guesscoord
writepdb out.pdb
writepsf out.psf
exit
    ''' % {'top': top, 'pdb' : pdb}

    with enter_temp_directory():
        with open('script.tcl', 'w') as f:
            f.write(TEMPLATE)
        os.system(' '.join([VMD, '-e', 'script.tcl', '-dispdev', 'none']))
        out_pdb = md.load('out.pdb')
        out_psf = md.load_psf('out.psf')

        # make sure the two topologies are equal
        eq(out_pdb.top, out_psf)

def test_iterload():
    t_ref = md.load(get_fn('frame0.h5'))[:20]
    with enter_temp_directory():
        for ext in t_ref._savers().keys():

            # only a 1 frame per file format
            if ext in ('.ncrst', '.rst7'):
                continue
            if ext in ('.lh5') and (on_win and on_py3):
                continue

            fn = 'temp%s' % ext
            t_ref.save(fn)

            def test():
                for stride in [1, 2, 3]:
                    loaded = md.load(fn, top=t_ref, stride=stride)
                    iterloaded = functools.reduce(lambda a, b: a.join(b), md.iterload(fn, top=t_ref, stride=stride, chunk=6))
                    eq(loaded.xyz, iterloaded.xyz)
                    eq(loaded.time, iterloaded.time)
                    eq(loaded.unitcell_angles, iterloaded.unitcell_angles)
                    eq(loaded.unitcell_lengths, iterloaded.unitcell_lengths)

            test.description = 'test_iterload: %s' % ext
            yield test

def test_dihedral_1():
    pymol = find_executable('pymol')
    if pymol is None:
        raise SkipTest("pymol executable not found")
    
    xyz = '''MODEL        0
ATOM      1    A ACE     1       4.300  13.100   8.600  1.00  0.00
ATOM      2    B ACE     1       5.200  13.600   8.800  1.00  0.00
ATOM      3    C ACE     1       4.900  14.300   9.600  1.00  0.00
ATOM      4    D ACE     1       5.600  14.200   7.900  1.00  0.00
    '''
    script = '''
with open('output.txt', 'w') as f:
    f.write('%f' % cmd.get_dihedral('1/A', '1/B', '1/C', '1/D'))
'''
    
    with enter_temp_directory():
        with open('xyz.pdb', 'w') as f:
            f.write(xyz)
        with open('pymolscript.py', 'w') as f:
            f.write(script)

        os.system('%s %s -cr %s' % (pymol, 'xyz.pdb', 'pymolscript.py'))
        with open('output.txt') as f:
            pymol_value = np.deg2rad(float(f.read()))
        t = md.load('xyz.pdb')

    mdtraj_value = md.compute_dihedrals(t, [[0,1,2,3]])[0,0]
    
    np.testing.assert_array_almost_equal(pymol_value, mdtraj_value)

    def run(self):
        # Transform paths in absolute paths since we'll change the working directory
        input_files = {local + os.path.splitext(path)[1]: os.path.abspath(path)
                       for local, path in self._file_paths.items() if 'moli' in local}
        output_files = {local + os.path.splitext(path)[1]: os.path.abspath(path)
                        for local, path in self._file_paths.items() if 'molo' in local}

        # Resolve all the names in the script
        local_files = {local: local + os.path.splitext(path)[1]
                       for local, path in self._file_paths.items()}
        script = self._script.format(**local_files) + 'quit\n'

        with enter_temp_directory():
            # Copy input files
            for local_file, file_path in input_files.items():
                shutil.copy(file_path, local_file)

            # Save script and run tleap
            with open('leap.in', 'w') as f:
                f.write(script)
            subprocess.check_output(['tleap', '-f', 'leap.in'])

            #Copy back output files
            for local_file, file_path in output_files.items():
                shutil.copy(local_file, file_path)

def test_phase_creation():
    """Phases are initialized correctly by Yank.create()."""
    phase_name = 'my-solvent-phase'
    toluene = testsystems.TolueneImplicit()
    protocol = AbsoluteAlchemicalFactory.defaultSolventProtocolImplicit()
    atom_indices = find_components(toluene.system, toluene.topology, 'resname TOL')

    phase = AlchemicalPhase(phase_name, toluene.system, toluene.topology,
                            toluene.positions, atom_indices, protocol)
    thermodynamic_state = ThermodynamicState(temperature=300.0*unit.kelvin)

    # Create new simulation.
    with enter_temp_directory():
        output_dir = 'output'
        utils.config_root_logger(verbose=False)
        yank = Yank(store_directory=output_dir)
        yank.create(thermodynamic_state, phase)

        # Netcdf dataset has been created
        nc_path = os.path.join(output_dir, phase_name + '.nc')
        assert os.path.isfile(nc_path)

        # Read data
        try:
            nc_file = netcdf.Dataset(nc_path, mode='r')
            metadata_group = nc_file.groups['metadata']
            serialized_topology = metadata_group.variables['topology'][0]
        finally:
            nc_file.close()

        # Topology has been stored correctly
        deserialized_topology = utils.deserialize_topology(serialized_topology)
        assert deserialized_topology == mdtraj.Topology.from_openmm(toluene.topology)

def chemical_shifts_shiftx2(trj, pH=5.0, temperature=298.00):
    """Predict chemical shifts of a trajectory using ShiftX2.

    Parameters
    ----------
    trj : Trajectory
        Trajectory to predict shifts for.
    pH : float, optional, default=5.0
        pH value which gets passed to the ShiftX2 predictor.
    temperature : float, optional, default=298.00
        Temperature which gets passed to the ShiftX2 predictor.

    Returns
    -------
    results : pandas DataFrame
        Dataframe containing results, with index consisting of
        (resSeq, atom_name) pairs and columns for each frame in trj.

    Notes
    -----
    You must have ShiftX2 available on your path; see (http://www.shiftx2.ca/).

    Chemical shift prediction is for PROTEIN atoms; trajectory objects
    with ligands, solvent, ions, or other non-protein components may give
    UNKNOWN RESULTS.

    Please cite the appropriate reference below.

    References
    ----------
    .. [1] Beomsoo Han, Yifeng Liu, Simon Ginzinger, and David Wishart.
       "SHIFTX2: significantly improved protein chemical shift
       prediction." J. Biomol. NMR, 50, 1 43-57 (2011)
    """
    pd = import_('pandas')
    binary = find_executable(SHIFTX2)
    if binary is None:
        raise OSError('External command not found. Looked for %s in PATH. `chemical_shifts_shiftx2` requires the external program SHIFTX2, available at http://www.shiftx2.ca/' % ', '.join(SHIFTX2))

    results = []
    with enter_temp_directory():
        for i in range(trj.n_frames):
            trj[i].save("./trj%d.pdb" % i)
        cmd = "%s -b 'trj*.pdb' -p %.1f -t %.2f" % (binary, pH, temperature)

        return_flag = os.system(cmd)

        if return_flag != 0:
            raise(IOError("Could not successfully execute command '%s', check your ShiftX2 installation or your input trajectory." % cmd))

        for i in range(trj.n_frames):
            d = pd.read_csv("./trj%d.pdb.cs" % i)
            d.rename(columns={"NUM": "resSeq", "RES": "resName", "ATOMNAME": "name"}, inplace=True)
            d["frame"] = i
            results.append(d)

    results = pd.concat(results)
    results = results.pivot_table(rows=["resSeq", "name"], cols="frame", values="SHIFT")
    return results

def chemical_shifts_ppm(trj):
    """Predict chemical shifts of a trajectory using ppm.

    Parameters
    ----------
    trj : Trajectory
        Trajectory to predict shifts for.

    Returns
    -------
    results : pandas.DataFrame
        Dataframe containing results, with index consisting of
        (resSeq, atom_name) pairs and columns for each frame in trj.

    Notes
    -----
    You must have ppm available on your path; see
    (http://spin.ccic.ohio-state.edu/index.php/download/index).

    Chemical shift prediction is for PROTEIN atoms; trajectory objects
    with ligands, solvent, ions, or other non-protein components may give
    UNKNOWN RESULTS.

    Please cite the appropriate reference below.

    References
    ----------
    .. [1] Li, DW, and Bruschweiler, R. "PPM: a side-chain and backbone chemical
       shift predictor for the assessment of protein conformational ensembles."
       J Biomol NMR. 2012 Nov;54(3):257-65.
    """
    pd = import_('pandas')
    binary = find_executable(PPM)

    first_resSeq = trj.top.residue(0).resSeq

    if binary is None:
        raise OSError('External command not found. Looked for %s in PATH. `chemical_shifts_ppm` requires the external program PPM, available at http://spin.ccic.ohio-state.edu/index.php/download/index' % ', '.join(PPM))

    with enter_temp_directory():
        trj.save("./trj.pdb")
        cmd = "%s -pdb trj.pdb -mode detail" % binary

        return_flag = os.system(cmd)

        if return_flag != 0:
            raise(IOError("Could not successfully execute command '%s', check your PPM installation or your input trajectory." % cmd))

        d = pd.read_table("./bb_details.dat", index_col=False, header=None, sep="\s*").drop([3], axis=1)

        d = d.rename(columns={0: "resSeq", 1: "resName", 2: "name"})
        d["resSeq"] += first_resSeq - 1  # Fix bug in PPM that reindexes to 1
        d = d.drop("resName", axis=1)
        d = d.set_index(["resSeq", "name"])
        d.columns = np.arange(trj.n_frames)
        d.columns.name = "frame"

    return d

def test_load_freesolv_gaffmol2_vs_sybylmol2_vs_obabelpdb():
    with enter_temp_directory():

        tar_filename = "freesolve_v0.3.tar.bz2"
        tar = tarfile.open(get_fn(tar_filename), mode="r:bz2")
        tar.extractall()
        tar.close()

        with open("./v0.3/database.pickle", 'rb') as f:
            database = pickle.load(f)

        for key in database:
            gaff_filename = "./v0.3/mol2files_gaff/%s.mol2" % key
            pdb_filename = "./v0.3/mol2files_gaff/%s.pdb" % key
            sybyl_filename = "./v0.3/mol2files_sybyl/%s.mol2" % key

            cmd = "obabel -imol2 %s -opdb > %s 2>/dev/null" % (sybyl_filename, pdb_filename)
            assert os.system(cmd) == 0

            t_pdb = md.load(pdb_filename)
            t_gaff = md.load(gaff_filename)
            t_sybyl = md.load(sybyl_filename)

            eq(t_pdb.n_atoms, t_gaff.n_atoms)
            eq(t_pdb.n_atoms, t_sybyl.n_atoms)

            eq(t_pdb.n_frames, t_gaff.n_frames)
            eq(t_pdb.n_frames, t_gaff.n_frames)

            eq(t_pdb.xyz, t_gaff.xyz, decimal=4)
            eq(t_pdb.xyz, t_sybyl.xyz, decimal=4)

            top_pdb, bonds_pdb = t_pdb.top.to_dataframe()
            top_gaff, bonds_gaff = t_gaff.top.to_dataframe()
            top_sybyl, bonds_sybyl = t_sybyl.top.to_dataframe()

            eq(top_sybyl.name.values, top_pdb.name.values)
            # eq(top_gaff.name.values, top_sybyl.name.values)  # THEY CAN HAVE DIFFERENT NAMES, so this isn't TRUE!

            def make_bonds_comparable(bond_array):
                """Create a bond connectivity matrix from a numpy array of atom pairs.  Avoids having to compare the order in which bonds are listed."""
                n_bonds = len(bond_array)
                data = np.ones(n_bonds)
                i = bond_array[:, 0]
                j = bond_array[:, 1]
                matrix = scipy.sparse.coo_matrix((data, (i, j)), shape=(t_pdb.n_atoms, t_pdb.n_atoms)).toarray()
                return matrix + matrix.T  # Symmetrize to account for (a ~ b) versus (b ~ a)

            bond_matrix_pdb = make_bonds_comparable(bonds_pdb)
            bond_matrix_gaff = make_bonds_comparable(bonds_gaff)
            bond_matrix_sybyl = make_bonds_comparable(bonds_sybyl)

            eq(bond_matrix_pdb, bond_matrix_gaff)
            eq(bond_matrix_pdb, bond_matrix_sybyl)

def smiles_to_antechamber(
    smiles_string,
    gaff_mol2_filename,
    frcmod_filename,
    residue_name="MOL",
    strictStereo=False,
):
    """Build a molecule from a smiles string and run antechamber,
    generating GAFF mol2 and frcmod files from a smiles string.  Charges
    will be generated using the OpenEye QuacPac AM1-BCC implementation.

    Parameters
    ----------
    smiles_string : str
        Smiles string of molecule to construct and charge
    gaff_mol2_filename : str
        Filename of mol2 file output of antechamber, with charges
        created from openeye
    frcmod_filename : str
        Filename of frcmod file output of antechamber.  Most likely
        this file will be almost empty, at least for typical molecules.
    residue_name : str, optional, default="MOL"
        OpenEye writes mol2 files with <0> as the residue / ligand name.
        This chokes many mol2 parsers, so we replace it with a string of
        your choosing.  This might be useful for downstream applications
        if the residue names are required to be unique.
    strictStereo : bool, optional, default=False
        If False, permits smiles strings with unspecified stereochemistry.
        See https://docs.eyesopen.com/omega/usage.html
    """
    oechem = import_("openeye.oechem")
    if not oechem.OEChemIsLicensed():
        raise (ImportError("Need License for oechem!"))

    # Get the absolute path so we can find these filenames from inside a temporary directory.
    gaff_mol2_filename = os.path.abspath(gaff_mol2_filename)
    frcmod_filename = os.path.abspath(frcmod_filename)

    m = smiles_to_oemol(smiles_string)
    m = get_charges(m, strictStereo=strictStereo, keep_confs=1)

    with enter_temp_directory():  # Avoid dumping 50 antechamber files in local directory.
        _unused = molecule_to_mol2(m, "./tmp.mol2", residue_name=residue_name)
        net_charge = oechem.OENetCharge(m)
        tmp_gaff_mol2_filename, tmp_frcmod_filename = run_antechamber(
            "tmp", "./tmp.mol2", charge_method=None, net_charge=net_charge
        )  # USE OE AM1BCC charges!
        shutil.copy(tmp_gaff_mol2_filename, gaff_mol2_filename)
        shutil.copy(tmp_frcmod_filename, frcmod_filename)

def test_against_gmx():
    t1 = md.load(get_fn('frame0.pdb'))

    # generated by converting frame0.pdb to gro with g_trajconv
    t2 = md.load(get_fn('frame0.gro'))

    with enter_temp_directory():
        t1.save('temp.gro')
        t3 = md.load('temp.gro')

    eq(t1.xyz, t2.xyz)
    eq(t1.xyz, t3.xyz)
    eq(t1.time, t2.time)
    eq(t1.time, t3.time)
    eq(t1.unitcell_vectors, t2.unitcell_vectors)
    eq(t1.unitcell_vectors, t3.unitcell_vectors)

def test_no_alchemical_atoms():
    """Test whether Yank raises exception when no alchemical atoms are specified."""
    toluene = testsystems.TolueneImplicit()

    # Create parameters. With the exception of atom_indices, all other
    # parameters must be legal, we don't want to catch an exception
    # different than the one we are testing.
    phase = AlchemicalPhase(name='solvent-implicit', reference_system=toluene.system,
                            reference_topology=toluene.topology,
                            positions=toluene.positions, atom_indices={'ligand': []},
                            protocol=AbsoluteAlchemicalFactory.defaultSolventProtocolImplicit())
    thermodynamic_state = ThermodynamicState(temperature=300.0*unit.kelvin)

    # Create new simulation.
    with enter_temp_directory():
        yank = Yank(store_directory='output')
        yank.create(thermodynamic_state, phase)

def test_force_overwrite():
    t_ref = md.load(get_fn('native2.pdb'), no_boxchk=True)
    for ext in t_ref._savers().keys():
        with enter_temp_directory():
            def test_1():
                fn = 'temp-1%s' % ext
                open(fn, 'w').close()
                t_ref.save(fn, force_overwrite=True)
            def test_2():
                fn = 'temp-2%s' % ext
                open(fn, 'w').close()
                assert_raises(IOError,
                              lambda: t_ref.save(fn, force_overwrite=False))
            test_1.description = 'test_force_overwrite (1): %s' % ext
            test_2.description = 'test_force_overwrite (2): %s' % ext
            yield test_1
            yield test_2

def test_1():
    fn = get_fn("frame0.h5")
    with enter_temp_directory():
        assert os.system("mixtape DRIDFeaturizer --trjs {} --out a.pkl".format(fn)) == 0
        assert os.system("mixtape DihedralFeaturizer --types phi psi --trjs {} --out b.pkl".format(fn)) == 0

        assert os.system("mixtape tICA --inp a.pkl --out ticamodel.pkl --transformed tics.pkl") == 0
        assert (
            os.system(
                "mixtape KMeans --random_state 0 --n_init 1 --inp b.pkl --out kmeans.pkl --transformed labels.pkl"
            )
            == 0
        )

        kmeans0 = verboseload("labels.pkl")
        kmeans1 = KMeans(random_state=0, n_init=1).fit_predict(verboseload("b.pkl"))
        tica0 = verboseload("tics.pkl")
        tica1 = tICA().fit_transform(verboseload("a.pkl"))

    eq(kmeans0[0], kmeans1[0])
    eq(tica0[0], tica1[0])

def test_save_load():
    # this cycles all the known formats you can save to, and then tries
    # to reload, using just a single-frame file.

    t_ref = md.load(get_fn('native.pdb'))
    t_ref.unitcell_vectors = np.array([[[1,0,0], [0,1,0], [0,0,1]]])
    with enter_temp_directory():
        for ext in t_ref._savers().keys():
            def test():
                fn = 'temp%s' % ext
                t_ref.save(fn)
                t = md.load(fn, top=t_ref.topology)

                eq(t.xyz, t_ref.xyz)
                eq(t.time, t_ref.time)
                if t._have_unitcell:
                    eq(t.unitcell_angles, t_ref.unitcell_angles)
                    eq(t.unitcell_lengths, t_ref.unitcell_lengths)

            test.description = 'test_save_load: %s' % ext
            yield test

def chemical_shifts_spartaplus(trj, rename_HN=True):
    """Predict chemical shifts of a trajectory using SPARTA+.

    Parameters
    ----------
    trj : Trajectory
        Trajectory to predict shifts for.
    rename_HN : bool, optional, default=True
        SPARTA+ calls the amide proton "HN" instead of the standard "H".
        When True, this option renames the output as "H" to match the PDB
        and BMRB nomenclature.

    Returns
    -------
    results : pandas.DataFrame
        Dataframe containing results, with index consisting of
        (resSeq, atom_name) pairs and columns for each frame in trj.

    Notes
    -----
    You must have SPARTA+ available on your path; see
    (http://spin.niddk.nih.gov/bax/software/SPARTA+/). Also, the SPARTAP_DIR
    environment variable must be set so that SPARTA+ knows where to find
    its database files.

    Chemical shift prediction is for PROTEIN atoms; trajectory objects
    with ligands, solvent, ions, or other non-protein components may give
    UNKNOWN RESULTS.

    Please cite the appropriate reference below.

    References
    ----------
    .. [1] Shen, Y., and Bax, Ad. "SPARTA+: a modest improvement in empirical
       NMR chemical shift prediction by means of an artificial neural network."
       J. Biomol. NMR, 48, 13-22 (2010)
    """
    pd = import_('pandas')
    binary = find_executable(SPARTA_PLUS)
    if binary is None:
        raise OSError('External command not found. Looked for %s in PATH. `chemical_shifts_spartaplus` requires the external program SPARTA+, available at http://spin.niddk.nih.gov/bax/software/SPARTA+/' % ', '.join(SPARTA_PLUS))

    names = ["resSeq", "resName", "name", "SS_SHIFT", "SHIFT", "RC_SHIFT", "HM_SHIFT", "EF_SHIFT", "SIGMA"]

    with enter_temp_directory():
        for i in range(trj.n_frames):
            trj[i].save("./trj%d.pdb" % i)

        cmd = "%s -in %s" % (binary, ' '.join("trj%d.pdb" % i for i in range(trj.n_frames)))

        return_flag = os.system(cmd)

        if return_flag != 0:
            raise(IOError("Could not successfully execute command '%s', check your SPARTA+ installation or your input trajectory." % cmd))

        lines_to_skip = _get_lines_to_skip("trj0_pred.tab")

        results = []
        for i in range(trj.n_frames):
            d = pd.read_table("./trj%d_pred.tab" % i, names=names, header=None, sep="\s*", skiprows=lines_to_skip)
            d["frame"] = i
            results.append(d)

    results = pd.concat(results)

    if rename_HN:
        results.name[results.name == "HN"] = "H"

    results = results.pivot_table(rows=["resSeq", "name"], cols="frame", values="SHIFT")

    return results