Python termcolor.cprint函数代码示例

def abinb_mkstemp(force_abinb_workdir=False, use_relpath=False, **kwargs):
    """
    Invoke mkstep with kwargs, return the (fd, name) of the temporary file.
    kwargs are passed to ``mkstemp`` except for ``dir`` if we are inside a jupyter notebook.

    Args:
        use_abipy_nbworkdir:
        use_relpath: Return relative path (os.path.relpath) if True else absolute (default)
            Relative paths are required if we are gonna use the temporary file in
            notebooks or in web browers.
            In this case, the caller is responsbile for calling the function with the correct flag.

    .. example:

        _, filename = abinb_mkstep(suffix="." + ext, text=True)
    """
    if in_notebook() or force_abinb_workdir:
        d = kwargs.pop("dir", None)
        if d is not None:
            cprint("Files should be created inside abipy_nbworkdir if we are inside jupyter or force_abinb_workdir",
                    "yellow")
        fd, path = tempfile.mkstemp(dir=get_abinb_workdir(), **kwargs)
    else:
        fd, path = tempfile.mkstemp(**kwargs)

    if use_relpath:
        path = os.path.relpath(path)

    return fd, path

    def plot_line_neighbors(self, site_index, radius, num=200, max_nn=10, fontsize=12, **kwargs):
        """
        Plot (interpolated) density/potential in real space along the lines connecting
        an atom specified by ``site_index`` and all neighbors within a sphere of given ``radius``.

        .. warning::

            This routine can produce lots of plots!
            Be careful with the value of ``radius``. See also ``max_nn``.

        Args:
            site_index: Index of the atom in the structure.
            radius: Radius of the sphere in Angstrom
            num: Number of points sampled along the line.
            max_nn: By default, only the first `max_nn` neighbors are showed.
            fontsize: legend and title fontsize

        Return: |matplotlib-Figure|
        """
        site = self.structure[site_index]
        nn_list = self.structure.get_neighbors(site, radius, include_index=True)
        if not nn_list:
            cprint("Zero neighbors found for radius %s Ang. Returning None." % radius, "yellow")
            return None
        # Sorte sites by distance.
        nn_list = list(sorted(nn_list, key=lambda t: t[1]))

        if max_nn is not None and len(nn_list) > max_nn:
            cprint("For radius %s, found %s neighbors but only max_nn %s sites are show." %
                    (radius, len(nn_list), max_nn), "yellow")
            nn_list = nn_list[:max_nn]

        # Get grid of axes.
        nrows, ncols = len(nn_list), 1
        ax_list, fig, plt = get_axarray_fig_plt(None, nrows=nrows, ncols=ncols,
                                                sharex=True, sharey=True, squeeze=True)
        ax_list = ax_list.ravel()

        interpolator = self.get_interpolator()

        for i, (nn, ax) in enumerate(zip(nn_list, ax_list)):
            nn_site, nn_dist, nn_sc_index  = nn
            title = "%s, %s, dist=%.3f A" % (nn_site.species_string, str(nn_site.frac_coords), nn_dist)

            r = interpolator.eval_line(site.frac_coords, nn_site.frac_coords, num=num, kpoint=None)

            for ispden in range(self.nspden):
                ax.plot(r.dist, r.values[ispden],
                        label=latexlabel_ispden(ispden, self.nspden) if i == 0 else None)

            ax.set_title(title, fontsize=fontsize)
            ax.grid(True)

            if i == nrows - 1:
                ax.set_xlabel("Distance from site_index %s [Angstrom]" % site_index)
                ax.set_ylabel(self.latex_label)
                if self.nspden > 1:
                    ax.legend(loc="best", fontsize=fontsize, shadow=True)

        return fig

    def get_dims_dataframe(self, with_time=True, index=None):
        """
        Build and return |pandas-DataFrame| with the dimensions of the calculation.

        Args:
            with_time: True if walltime and cputime should be added
            index: Index of the dataframe. Use relative paths of files if None.
        """
        rows, my_index = [], []
        for i, abo in enumerate(self.abifiles):
            try:
                dims_dataset, spg_dataset = abo.get_dims_spginfo_dataset()
            except Exception as exc:
                cprint("Exception while trying to get dimensions from %s\n%s" % (abo.relpath, str(exc)), "yellow")
                continue

            for dtindex, dims in dims_dataset.items():
                dims = dims.copy()
                dims.update({"dtset": dtindex})
                # Add walltime and cputime in seconds
                if with_time:
                    dims.update(OrderedDict([(k, getattr(abo, k)) for k in
                        ("overall_cputime", "proc0_cputime", "overall_walltime", "proc0_walltime")]))
                rows.append(dims)
                my_index.append(abo.relpath if index is None else index[i])

        return pd.DataFrame(rows, index=my_index, columns=list(rows[0].keys()))

    def get_dataframe(self, with_geo=True, with_dims=True, abspath=False, funcs=None):
        """
        Return a |pandas-DataFrame| with the most important results and the filenames as index.

        Args:
            with_geo: True if structure info should be added to the dataframe
            with_dims: True if dimensions should be added
            abspath: True if paths in index should be absolute. Default: Relative to getcwd().
            funcs: Function or list of functions to execute to add more data to the DataFrame.
                Each function receives a |GsrFile| object and returns a tuple (key, value)
                where key is a string with the name of column and value is the value to be inserted.
        """
        rows, row_names = [], []
        for label, abo in self.items():
            row_names.append(label)
            d = OrderedDict()

            if with_dims:
                dims_dataset, spg_dataset = abo.get_dims_spginfo_dataset()
                if len(dims_dataset) > 1:
                    cprint("Multiple datasets are not supported. ARGH!", "yellow")
                d.update(dims_dataset[1])

            # Add info on structure.
            if with_geo and abo.run_completed:
                d.update(abo.final_structure.get_dict4pandas(with_spglib=True))

            # Execute functions
            if funcs is not None: d.update(self._exec_funcs(funcs, abo))
            rows.append(d)

        row_names = row_names if not abspath else self._to_relpaths(row_names)
        return pd.DataFrame(rows, index=row_names, columns=list(rows[0].keys()))

    def make_open_notebook(self, nbpath=None, foreground=False):
        """
        Generate an ipython notebook and open it in the browser.

        Args:
            nbpath: If nbpath is None, a temporay file is created.
            foreground: By default, jupyter is executed in background and stdout, stderr are redirected
            to devnull. Use foreground to run the process in foreground

        Return:
            system exit code.

        Raise:
            RuntimeError if jupyter is not in $PATH
        """
        nbpath = self.write_notebook(nbpath=nbpath)

        if foreground:
            cmd = "jupyter notebook %s" % nbpath
            return os.system(cmd)

        else:
            cmd = "jupyter notebook %s &> /dev/null &" % nbpath
            print("Executing:", cmd)

            import subprocess
            try:
                from subprocess import DEVNULL # py3k
            except ImportError:
                DEVNULL = open(os.devnull, "wb")

            process = subprocess.Popen(cmd.split(), shell=False, stdout=DEVNULL) #, stderr=DEVNULL)
            cprint("pid: %s" % str(process.pid), "yellow")

    def read_doses(self):
        """
        Return a |AttrDict| with the DOSes available in the file. Empty dict if
        DOSes are not available.
        """
        if "gruns_nomega" not in self.rootgrp.dimensions:
            cprint("File `%s` does not contain ph-DOSes, returning empty dict" % self.path, "yellow")
            return {}

        # Read q-point sampling used to compute DOSes.
        qptrlatt = self.read_value("gruns_qptrlatt")
        shifts = self.read_value("gruns_shiftq")
        qsampling = KSamplingInfo.from_kptrlatt(qptrlatt, shifts, kptopt=1)

        frac_coords_ibz = self.read_value("gruns_qibz")
        weights = self.read_value("gruns_wtq")
        qpoints = IrredZone(self.structure.reciprocal_lattice, frac_coords_ibz,
                            weights=weights, names=None, ksampling=qsampling)

        # DOSes are in 1/Hartree.
        d = AttrDict(wmesh=self.read_value("gruns_omega_mesh") * abu.Ha_eV, qpoints=qpoints)

        for dos_name in _ALL_DOS_NAMES:
            dos_idos = self.read_value(dos_name)
            dos_idos[0] *= abu.eV_Ha  # Here we convert to eV. IDOS are not changed.
            d[dos_name] = dos_idos

        return d

    def make_and_open_notebook(self, nbpath=None, foreground=False):  # pragma: no cover
        """
        Generate an jupyter_ notebook and open it in the browser.

        Args:
            nbpath: If nbpath is None, a temporay file is created.
            foreground: By default, jupyter is executed in background and stdout, stderr are redirected
            to devnull. Use foreground to run the process in foreground

        Return:
            system exit code.

        Raise:
            `RuntimeError` if jupyter_ is not in $PATH
        """
        nbpath = self.write_notebook(nbpath=nbpath)

        if which("jupyter") is None:
            raise RuntimeError("Cannot find jupyter in $PATH. Install it with `conda install jupyter or `pip install jupyter`")

        if foreground:
            return os.system("jupyter notebook %s" % nbpath)
        else:
            fd, tmpname = tempfile.mkstemp(text=True)
            print(tmpname)
            cmd = "jupyter notebook %s" % nbpath
            print("Executing:", cmd)
            print("stdout and stderr redirected to %s" % tmpname)
            import subprocess
            process = subprocess.Popen(cmd.split(), shell=False, stdout=fd, stderr=fd)
            cprint("pid: %s" % str(process.pid), "yellow")
            return 0

    def write_cube(self, filename=None, spin="total"):
        """Write density in CUBE format."""
        if filename is None:
            filename = self.basename.replace(".nc", ".cube")
            cprint("Writing density in CUBE format to file: %s" % filename, "yellow")

        return self.density.export_to_cube(filename, spin=spin)

    def from_files(cls, filenames, labels=None, abspath=False):
        """
        Build a Robot from a list of `filenames`.
        if labels is None, labels are automatically generated from absolute paths.

        Args:
            abspath: True if paths in index should be absolute. Default: Relative to `top`.
        """
        filenames = list_strings(filenames)
        from abipy.abilab import abiopen
        filenames = [f for f in filenames if cls.class_handles_filename(f)]
        items = []
        for i, f in enumerate(filenames):
            try:
                abifile = abiopen(f)
            except Exception as exc:
                cprint("Exception while opening file: `%s`" % str(f), "red")
                cprint(exc, "red")
                abifile = None

            if abifile is not None:
                label = abifile.filepath if labels is None else labels[i]
                items.append((label, abifile))

        new = cls(*items)
        if labels is None and not abspath: new.trim_paths(start=None)
        return new

    def plot_unfolded(self, kbounds, klabels, ylims=None, dist_tol=1e-12, verbose=0,
                      colormap="afmhot", facecolor="black", ax=None, fontsize=12, **kwargs):
        r"""
        Plot unfolded band structure with spectral weights.

        Args:
            klabels: dictionary whose keys are tuple with the reduced coordinates of the k-points.
                The values are the labels. e.g. ``klabels = {(0.0,0.0,0.0): "$\Gamma$", (0.5,0,0): "L"}``.
            ylims: Set the data limits for the y-axis. Accept tuple e.g. ``(left, right)``
                or scalar e.g. ``left``. If left (right) is None, default values are used
            dist_tol: A point is considered to be on the path if its distance from the line
                is less than dist_tol.
            verbose: Verbosity level.
            colormap: Have a look at the colormaps here and decide which one you like:
                http://matplotlib.sourceforge.net/examples/pylab_examples/show_colormaps.html
            facecolor:
            ax: |matplotlib-Axes| or None if a new figure should be created.
            fontsize: Legend and title fontsize.

        Returns: |matplotlib-Figure|
	"""
        cart_bounds = [self.pc_lattice.reciprocal_lattice.get_cartesian_coords(c)
                       for c in np.reshape(kbounds, (-1, 3))]
        uf_cart = self.uf_kpoints.get_cart_coords()

        p = find_points_along_path(cart_bounds, uf_cart, dist_tol)
        if len(p.ikfound) == 0:
            cprint("Warning: find_points_along_path returned zero points along the path. Try to increase dist_tol.", "yellow")
            return None
        if verbose:
            uf_frac_coords = np.reshape([k.frac_coords for k in self.uf_kpoints], (-1, 3))
            fcoords = uf_frac_coords[p.ikfound]
            print("Found %d points along input k-path" % len(fcoords))
            print("k-points of path in reduced coordinates:")
            print(fcoords)

        fact = 8.0
        e0 = self.ebands.fermie
        ax, fig, plt = get_ax_fig_plt(ax=ax)
        ax.set_facecolor(facecolor)

        xs = np.tile(p.dist_list, self.nband)
        marker_spin = {0: "^", 1: "v"} if self.nss == 2 else {0: "o"}
        for spin in range(self.nss):
            ys = self.uf_eigens[spin, p.ikfound, :] - e0
            ws = self.uf_weights[spin, p.ikfound, :]
            s = ax.scatter(xs, ys.T, s=fact * ws.T, c=ws.T,
                           marker=marker_spin[spin], label=None if self.nss == 1 else "spin %s" % spin,
                           linewidth=1, edgecolors='none', cmap=plt.get_cmap(colormap))
            plt.colorbar(s, ax=ax, orientation='vertical')

        ax.set_xticks(p.path_ticks, minor=False)
        ax.set_xticklabels(klabels, fontdict=None, minor=False, size=kwargs.pop("klabel_size", "large"))
        ax.grid(True)
        ax.set_ylabel('Energy (eV)')
        set_axlims(ax, ylims, "y")
        if self.nss == 2: ax.legend(loc="best", fontsize=fontsize, shadow=True)

        return fig

def check_ordered_structure(structure):
    """Print a warning and sys.exit 1 if structure is disordered."""
    if not structure.is_ordered:
        cprint("""
Cannot handle disordered structure with fractional site occupancies.
Use OrderDisorderedStructureTransformation or EnumerateStructureTransformation
to build an appropriate supercell from partial occupancies.""", color="magenta")
        sys.exit(1)

def handle_overwrite(path, options):
    """Exit 1 if file ``path`` exists and not options.force else return path."""
    name_parts = os.path.splitext(path)
    print("Writing %s file:" % name_parts[-1].replace("." , "").upper())
    if os.path.exists(path) and not options.force:
        cprint("Cannot overwrite pre-existent file. Use `-f` options.", "red")
        sys.exit(1)
    return path

    def interpolate_ebands(self, vertices_names=None, line_density=20, ngkpt=None, shiftk=(0, 0, 0), kpoints=None):
        """
        Build new |ElectronBands| object by interpolating the KS Hamiltonian with Wannier functions.
        Supports k-path via (vertices_names, line_density), IBZ mesh defined by ngkpt and shiftk
        or input list of kpoints.

        Args:
            vertices_names: Used to specify the k-path for the interpolated QP band structure
                List of tuple, each tuple is of the form (kfrac_coords, kname) where
                kfrac_coords are the reduced coordinates of the k-point and kname is a string with the name of
                the k-point. Each point represents a vertex of the k-path. ``line_density`` defines
                the density of the sampling. If None, the k-path is automatically generated according
                to the point group of the system.
            line_density: Number of points in the smallest segment of the k-path. Used with ``vertices_names``.
            ngkpt: Mesh divisions. Used if bands should be interpolated in the IBZ.
            shiftk: Shifts for k-meshs. Used with ngkpt.
            kpoints: |KpointList| object taken e.g from a previous ElectronBands.
                Has precedence over vertices_names and line_density.

        Returns: |ElectronBands| object with Wannier-interpolated energies.
        """
        # Need KpointList object.
        if kpoints is None:
            if ngkpt is not None:
                # IBZ sampling
                kpoints = IrredZone.from_ngkpt(self.structure, ngkpt, shiftk, kptopt=1, verbose=0)
            else:
                # K-Path
                if vertices_names is None:
                    vertices_names = [(k.frac_coords, k.name) for k in self.structure.hsym_kpoints]
                kpoints = Kpath.from_vertices_and_names(self.structure, vertices_names, line_density=line_density)

        nk = len(kpoints)
        eigens = np.zeros((self.nsppol, nk, self.mwan))

        # Interpolate Hamiltonian for each kpoint and spin.
        start = time.time()
        write_warning = True
        for spin in range(self.nsppol):
            num_wan = self.nwan_spin[spin]
            for ik, kpt in enumerate(kpoints):
                oeigs = self.hwan.eval_sk(spin, kpt.frac_coords)
                eigens[spin, ik, :num_wan] = oeigs
                if num_wan < self.mwan:
                    # May have different number of wannier functions if nsppol == 2.
                    # Here I use the last value to fill eigens matrix (not very clean but oh well).
                    eigens[spin, ik, num_wan:self.mwan] = oeigs[-1]
                    if write_warning:
                        cprint("Different number of wannier functions for spin. Filling last bands with oeigs[-1]",
                               "yellow")
                        write_warning = False

        print("Interpolation completed in %.3f [s]" % (time.time() - start))
        occfacts = np.zeros_like(eigens)

        return ElectronBands(self.structure, kpoints, eigens, self.ebands.fermie,
                             occfacts, self.ebands.nelect, self.nspinor, self.nspden,
                             smearing=self.ebands.smearing)

    def write_chgcar(self, filename=None):
        """
        Write density in CHGCAR format. Return :class:`ChgCar` instance.
        """
        if filename is None:
            filename = self.basename.replace(".nc", "_CHGCAR")
            cprint("Writing density in CHGCAR format to file: %s" % filename, "yellow")

        return self.density.to_chgcar(filename=filename)

def remove_disordered(structures, paths):
    """Remove disordered structures and print warning message."""
    slist = []
    for s, p in zip(structures, paths):
        if not s.is_ordered:
            cprint("Removing disordered structure: %s found in %s" % (s.formula, p), "magenta")
        else:
            slist.append(s)
    return slist