Python io.zopen函数代码示例

 def test_zopen(self):
     with zopen(os.path.join(test_dir, "myfile_gz.gz"), mode="rt") as f:
         self.assertEqual(f.read(), "HelloWorld.\n\n")
     with zopen(os.path.join(test_dir, "myfile_bz2.bz2"), mode="rt") as f:
         self.assertEqual(f.read(), "HelloWorld.\n\n")
     with zopen(os.path.join(test_dir, "myfile_bz2.bz2"), "rt") as f:
         self.assertEqual(f.read(), "HelloWorld.\n\n")
     with zopen(os.path.join(test_dir, "myfile"), mode="rt") as f:
         self.assertEqual(f.read(), "HelloWorld.\n\n")

    def read_excitation_energies(self):
        """
        Read a excitation energies after a TD-DFT calculation.

        Returns:

            A list: A list of tuple for each transition such as 
                    [(energie (eV), lambda (nm), oscillatory strength), ... ]
        """
        float_patt = re.compile("\s*([+-]?\d+\.\d+)")

        transitions = list()

        # read in file
        with zopen(self.filename, "r") as f:
            line = f.readline()
            td = False
            while line != "":
                if re.search("^\sExcitation energies and oscillator strengths:", line):
                    td = True
   
                if td:
                    if re.search("^\sExcited State\s*\d", line):
                        val = [float(v) for v in float_patt.findall(line)]
                        transitions.append(tuple(val[0:3]))
                line = f.readline()
        return transitions

    def atoms_string_from_file(filename):
        """
        Reads atomic shells from file such as feff.inp or ATOMS file
        The lines are arranged as follows:

        x y z   ipot    Atom Symbol   Distance   Number

        with distance being the shell radius and ipot an integer identifying
        the potential used.

        Args:
            filename: File name containing atomic coord data.

        Returns:
            Atoms string.
        """
        with zopen(filename, "r") as fobject:
            f = fobject.readlines()
            coords = 0
            atoms_str = []

            for line in f:
                if coords == 0:
                    find_atoms = line.find("ATOMS")
                    if find_atoms >= 0:
                        coords = 1
                if coords == 1:
                    atoms_str.append(line.replace("\r", ""))

        return ''.join(atoms_str)

def dumpfn(obj, fn, *args, **kwargs):
    """
    Dump to a json/yaml directly by filename instead of a File-like object.
    For YAML, PyYAML must be installed. The file type is automatically
    detected. YAML is assumed if the filename contains "yaml" (lower or upper
    case). Otherwise, json is always assumed. Furthermore, if pyyaml is
    compiled with the LibYAML library, the CDumper is automatically chosen
    for ~10x faster parsing.

    Args:
        obj (object): Object to dump.
        fn (str): filename.
        \*args: Any of the args supported by json/yaml.dump.
        \*\*kwargs: Any of the kwargs supported by json/yaml.dump.

    Returns:
        (object) Result of json.load.
    """
    with zopen(fn, "wt") as fp:
        if "yaml" in fn.lower():
            if yaml is None:
                raise RuntimeError("Loading of YAML files is not "
                                   "possible as PyYAML is not installed.")
            if "Dumper" not in kwargs:
                kwargs["Dumper"] = Dumper
            yaml.dump(obj, fp, *args, **kwargs)
        else:
            fp.write("%s" % json.dumps(obj, *args, **kwargs))

def parse_lammps_dumps(file_pattern):
    """
    Generator that parses dump file(s).

    Args:
        file_pattern (str): Filename to parse. The timestep wildcard
            (e.g., dump.atom.'*') is supported and the files are parsed
            in the sequence of timestep.

    Yields:
        LammpsDump for each available snapshot.

    """
    files = glob.glob(file_pattern)
    if len(files) > 1:
        pattern = r"%s" % file_pattern.replace("*", "([0-9]+)")
        pattern = pattern.replace("\\", "\\\\")
        files = sorted(files,
                       key=lambda f: int(re.match(pattern, f).group(1)))

    for fname in files:
        with zopen(fname, "rt") as f:
            dump_cache = []
            for line in f:
                if line.startswith("ITEM: TIMESTEP"):
                    if len(dump_cache) > 0:
                        yield LammpsDump.from_string("".join(dump_cache))
                    dump_cache = [line]
                else:
                    dump_cache.append(line)
            yield LammpsDump.from_string("".join(dump_cache))

def loadfn(fn, *args, **kwargs):
    """
    Loads json/yaml directly from a filename instead of a File-like object.
    For YAML, PyYAML must be installed. The file type is automatically
    detected. YAML is assumed if the filename contains "yaml" (lower or upper
    case). Otherwise, json is always assumed. Furthermore, if pyyaml is
    compiled with the LibYAML library, the CLoader is automatically chosen
    for ~10x faster parsing.

    Args:
        fn (str): filename
        \*args: Any of the args supported by json/yaml.load.
        \*\*kwargs: Any of the kwargs supported by json/yaml.load.

    Returns:
        (object) Result of json/yaml.load.
    """
    with zopen(fn) as fp:
        if "yaml" in fn.lower():
            if yaml is None:
                raise RuntimeError("Loading of YAML files is not "
                                   "possible as PyYAML is not installed.")
            if "Loader" not in kwargs:
                kwargs["Loader"] = Loader
            return yaml.load(fp, *args, **kwargs)
        else:
            return json.load(fp, *args, **kwargs)

    def gw_run(self):
        """
        Performs FIESTA (gw) run
        """

        if self.folder != os.getcwd():
            init_folder = os.getcwd()
            os.chdir(self.folder)

        with zopen(self.log_file, "w") as fout:
            subprocess.call(
                [
                    "mpirun",
                    "-n",
                    str(self.mpi_procs),
                    "fiesta",
                    str(self.grid[0]),
                    str(self.grid[1]),
                    str(self.grid[2]),
                ],
                stdout=fout,
            )

        if self.folder != os.getcwd():
            os.chdir(init_folder)

 def test_potcar_map(self):
     fe_potcar = zopen(os.path.join(test_dir, "POT_GGA_PAW_PBE", "POTCAR.Fe_pv.gz")).read().decode("utf-8")
     # specify V instead of Fe - this makes sure the test won't pass if the
     # code just grabs the POTCAR from the config file (the config file would
     # grab the V POTCAR)
     potcar = Potcar(["V"], sym_potcar_map={"V": fe_potcar})
     self.assertEqual(potcar.symbols, ["Fe_pv"], "Wrong symbols read in " "for POTCAR")

def write_mol(mol, filename):
    """
    Write a molecule to a file based on file extension. For example, anything
    ending in a "xyz" is assumed to be a XYZ file. Supported formats include
    xyz, Gaussian input (gjf|g03|g09|com|inp), and pymatgen's JSON serialized
    molecules.

    Args:
        mol (Molecule/IMolecule): Molecule to write
        filename (str): A filename to write to.
    """
    fname = os.path.basename(filename)
    if fnmatch(fname.lower(), "*.xyz*"):
        return XYZ(mol).write_file(filename)
    elif any([fnmatch(fname.lower(), "*.{}*".format(r))
              for r in ["gjf", "g03", "g09", "com", "inp"]]):
        return GaussianInput(mol).write_file(filename)
    elif fnmatch(fname, "*.json*") or fnmatch(fname, "*.mson*"):
        with zopen(filename, "wt") as f:
            return f.write(str2unicode(json.dumps(mol, cls=MontyEncoder)))
    else:
        m = re.search("\.(pdb|mol|mdl|sdf|sd|ml2|sy2|mol2|cml|mrv)",
                      filename.lower())
        if m:
            return BabelMolAdaptor(mol).write_file(filename, m.group(1))

    raise ValueError("Unrecognized file extension!")

def store_dataframe_as_json(dataframe, filename, compression=None,
                            orient='split'):
    """Store pandas dataframe as a json file.

    Automatically encodes pymatgen objects as dictionaries.

    Args:
        dataframe (Pandas.Dataframe): A pandas dataframe.
        filename (str): Path to json file.
        compression (str or None): A compression mode. Valid options are "gz",
            "bz2", and None. Defaults to None. If the filename does not end
            in with the correct suffix it will be added automatically.
        orient (str): Determines the format in which the dictionary data is
            stored. This takes the same set of arguments as the `orient` option
            in `pandas.DataFrame.to_dict()` function. 'split' is recommended
            as it is relatively space efficient and preserves the dtype
            of the index.
    """
    if compression not in ["gz", "bz2", None]:
        raise ValueError("Supported compression formats are 'gz' and 'bz2'.")

    if compression and not filename.lower().endswith(".{}".format(compression)):
        filename = "{}.{}".format(filename, compression)

    write_type = "wb" if compression else "w"

    with zopen(filename, write_type) as f:
        data = json.dumps(dataframe.to_dict(orient=orient), cls=MontyEncoder)
        if compression:
            data = data.encode()
        f.write(data)

    def read(self):
        """Read all snapshots from each dump file."""
        for dumpfile in self.dumpfiles:
            with zopen(dumpfile) as f:
                snapshot = self.read_snapshot(f)
                while snapshot is not None:
                    self.trajectory.append(snapshot)
                    print(snapshot.timestep, end=' ')
                    sys.stdout.flush()
                    snapshot = self.read_snapshot(f)
        print()

        self.trajectory.sort(key=attrgetter('timestep'))
        self.trajectory.cull()

        print("read {:d} snapshots".format(self.Nsnaps))

        self.trajectory.time_selection.all()
        self.trajectory.t0_snapshot = self.trajectory[0]

        if self.dumpattrs:
            print('Dumped Atom attributes: {}'.format(self.dumpattrs2str()))
        else:
            print('No dump column assignments')

        if 'x' not in self.dumpattrs or 'y' not in self.dumpattrs or \
                'z' not in self.dumpattrs:
            print('dump scaling status unknown')
        elif self.Nsnaps > 0:
            if self.scale_original:
                self.unscale()
            elif self.scale_original is None:
                print('dump scaling status unknown')
            else:
                print('dump is already unscaled')

def read_mol(filename):
    """
    Reads a molecule based on file extension. For example, anything ending in
    a "xyz" is assumed to be a XYZ file. Supported formats include xyz,
    gaussian input (gjf|g03|g09|com|inp), Gaussian output (.out|and
    pymatgen's JSON serialized molecules. Using openbabel,
    many more extensions are supported but requires openbabel to be installed.

    Args:
        filename (str): A filename to read from.

    Returns:
        A Molecule object.
    """
    fname = os.path.basename(filename)
    if fnmatch(fname.lower(), "*.xyz*"):
        return XYZ.from_file(filename).molecule
    elif any([fnmatch(fname.lower(), "*.{}*".format(r)) for r in ["gjf", "g03", "g09", "com", "inp"]]):
        return GaussianInput.from_file(filename).molecule
    elif any([fnmatch(fname.lower(), "*.{}*".format(r)) for r in ["out", "lis", "log"]]):
        return GaussianOutput(filename).final_structure
    elif fnmatch(fname, "*.json*") or fnmatch(fname, "*.mson*"):
        with zopen(filename) as f:
            s = json.load(f, cls=MontyDecoder)
            if type(s) != Molecule:
                raise IOError("File does not contain a valid serialized " "molecule")
            return s
    else:
        m = re.search("\.(pdb|mol|mdl|sdf|sd|ml2|sy2|mol2|cml|mrv)", filename.lower())
        if m:
            return BabelMolAdaptor.from_file(filename, m.group(1)).pymatgen_mol

    raise ValueError("Unrecognized file extension!")

def write_structure(structure, filename):
    """
    Write a structure to a file based on file extension. For example, anything
    ending in a "cif" is assumed to be a Crystallographic Information Format
    file. Supported formats include CIF, POSCAR, CSSR and pymatgen's JSON
    serialized structures.

    Args:
        structure (Structure/IStructure): Structure to write
        filename (str): A filename to write to.
    """
    fname = os.path.basename(filename)
    if fnmatch(fname, "*.cif*"):
        writer = CifWriter(structure)
    elif fnmatch(fname, "POSCAR*") or fnmatch(fname, "CONTCAR*"):
        writer = Poscar(structure)
    elif fnmatch(fname.lower(), "*.cssr*"):
        writer = Cssr(structure)
    elif fnmatch(fname, "*.json*") or fnmatch(fname, "*.mson*"):
        with zopen(filename, "wt") as f:
            f.write(str2unicode(json.dumps(structure, cls=MontyEncoder)))
            return
    else:
        raise ValueError("Unrecognized file extension!")

    writer.write_file(filename)

 def run(self):
     """
     Performs actual nwchem run.
     """
     with zopen(self.output_file, 'w') as fout:
         return subprocess.Popen(self.nwchem_cmd + [self.input_file],
                                 stdout=fout)

def read_structure(filename):
    """
    Reads a structure based on file extension. For example, anything ending in
    a "cif" is assumed to be a Crystallographic Information Format file.
    Supported formats include CIF, POSCAR/CONTCAR, CHGCAR, LOCPOT,
    vasprun.xml, CSSR and pymatgen's JSON serialized structures.

    Args:
        filename (str): A filename to read from.

    Returns:
        A Structure object.
    """
    fname = os.path.basename(filename)
    if fnmatch(fname.lower(), "*.cif*"):
        parser = CifParser(filename)
        return parser.get_structures(True)[0]
    elif fnmatch(fname, "POSCAR*") or fnmatch(fname, "CONTCAR*"):
        return Poscar.from_file(filename, False).structure
    elif fnmatch(fname, "CHGCAR*") or fnmatch(fname, "LOCPOT*"):
        return Chgcar.from_file(filename).structure
    elif fnmatch(fname, "vasprun*.xml*"):
        return Vasprun(filename).final_structure
    elif fnmatch(fname.lower(), "*.cssr*"):
        cssr = Cssr.from_file(filename)
        return cssr.structure
    elif fnmatch(fname, "*.json*") or fnmatch(fname, "*.mson*"):
        with zopen(filename) as f:
            s = json.load(f, cls=PMGJSONDecoder)
            if type(s) != Structure:
                raise IOError("File does not contain a valid serialized "
                              "structure")
            return s
    raise ValueError("Unrecognized file extension!")