Skip to main content
`Ctrl`+`K`

#

# Source code for ase.spectrum.band_structure

    # fmt: off

    import numpy as np

    import ase  # Annotations
    from ase.calculators.calculator import PropertyNotImplementedError
    from ase.utils import jsonable

    def calculate_band_structure(atoms, path=None, scf_kwargs=None,
                                 bs_kwargs=None, kpts_tol=1e-6, cell_tol=1e-6):
        """Calculate band structure.

        The purpose of this function is to abstract a band structure calculation
        so the workflow does not depend on the calculator.

        First trigger SCF calculation if necessary, then set arguments
        on the calculator for band structure calculation, then return
        calculated band structure.

        The difference from get_band_structure() is that the latter
        expects the calculation to already have been done."""
        if path is None:
            path = atoms.cell.bandpath()

        from ase.lattice import celldiff  # Should this be a method on cell?
        if any(path.cell.any(1) != atoms.pbc):
            raise ValueError('The band path\'s cell, {}, does not match the '
                             'periodicity {} of the atoms'
                             .format(path.cell, atoms.pbc))
        cell_err = celldiff(path.cell, atoms.cell.uncomplete(atoms.pbc))
        if cell_err > cell_tol:
            raise ValueError('Atoms and band path have different unit cells.  '
                             'Please reduce atoms to standard form.  '
                             'Cell lengths and angles are {} vs {}'
                             .format(atoms.cell.cellpar(), path.cell.cellpar()))

        calc = atoms.calc
        if calc is None:
            raise ValueError('Atoms have no calculator')

        if scf_kwargs is not None:
            calc.set(**scf_kwargs)

        # Proposed standard mechanism for calculators to advertise that they
        # use the bandpath keyword to handle band structures rather than
        # a double (SCF + BS) run.
        use_bandpath_kw = getattr(calc, 'accepts_bandpath_keyword', False)
        if use_bandpath_kw:
            calc.set(bandpath=path)
            atoms.get_potential_energy()
            return calc.band_structure()

        atoms.get_potential_energy()

        if hasattr(calc, 'get_fermi_level'):
            # What is the protocol for a calculator to tell whether
            # it has fermi_energy?
            eref = calc.get_fermi_level()
        else:
            eref = 0.0

        if bs_kwargs is None:
            bs_kwargs = {}

        calc.set(kpts=path, **bs_kwargs)
        calc.results.clear()  # XXX get rid of me

        # Calculators are too inconsistent here:
        # * atoms.get_potential_energy() will fail when total energy is
        #   not in results after BS calculation (Espresso)
        # * calc.calculate(atoms) doesn't ask for any quantity, so some
        #   calculators may not calculate anything at all
        # * 'bandstructure' is not a recognized property we can ask for
        try:
            atoms.get_potential_energy()
        except PropertyNotImplementedError:
            pass

        ibzkpts = calc.get_ibz_k_points()
        kpts_err = np.abs(path.kpts - ibzkpts).max()
        if kpts_err > kpts_tol:
            raise RuntimeError('Kpoints of calculator differ from those '
                               'of the band path we just used; '
                               'err={} > tol={}'.format(kpts_err, kpts_tol))

        bs = get_band_structure(atoms, path=path, reference=eref)
        return bs

    def get_band_structure(atoms=None, calc=None, path=None, reference=None):
        """Create band structure object from Atoms or calculator."""
        # path and reference are used internally at the moment, but
        # the exact implementation will probably change.  WIP.
        #
        # XXX We throw away info about the bandpath when we create the calculator.
        # If we have kept the bandpath, we can provide it as an argument here.
        # It would be wise to check that the bandpath kpoints are the same as
        # those stored in the calculator.
        atoms = atoms if atoms is not None else calc.atoms
        calc = calc if calc is not None else atoms.calc

        kpts = calc.get_ibz_k_points()

        energies = []
        for s in range(calc.get_number_of_spins()):
            energies.append([calc.get_eigenvalues(kpt=k, spin=s)
                             for k in range(len(kpts))])
        energies = np.array(energies)

        if path is None:
            from ase.dft.kpoints import (
                BandPath,
                find_bandpath_kinks,
                resolve_custom_points,
            )
            standard_path = atoms.cell.bandpath(npoints=0)
            # Kpoints are already evaluated, we just need to put them into
            # the path (whether they fit our idea of what the path is, or not).
            #
            # Depending on how the path was established, the kpoints might
            # be valid high-symmetry points, but since there are multiple
            # high-symmetry points of each type, they may not coincide
            # with ours if the bandpath was generated by another code.
            #
            # Here we hack it so the BandPath has proper points even if they
            # come from some weird source.
            #
            # This operation (manually hacking the bandpath) is liable to break.
            # TODO: Make it available as a proper (documented) bandpath method.
            kinks = find_bandpath_kinks(atoms.cell, kpts, eps=1e-5)
            pathspec, special_points = resolve_custom_points(
                kpts[kinks], standard_path.special_points, eps=1e-5)
            path = BandPath(standard_path.cell,
                            kpts=kpts,
                            path=pathspec,
                            special_points=special_points)

        # XXX If we *did* get the path, now would be a good time to check
        # that it matches the cell!  Although the path can only be passed
        # because we internally want to not re-evaluate the Bravais
        # lattice type.  (We actually need an eps parameter, too.)

        if reference is None:
            # Fermi level should come from the GS calculation, not the BS one!
            reference = calc.get_fermi_level()

        if reference is None:
            # Fermi level may not be available, e.g., with non-Fermi smearing.
            # XXX Actually get_fermi_level() should raise an error when Fermi
            # level wasn't available, so we should fix that.
            reference = 0.0

        return BandStructure(path=path,
                             energies=energies,
                             reference=reference)

    class BandStructurePlot:
        def __init__(self, bs):
            self.bs = bs
            self.ax = None
            self.xcoords = None

        def plot(self, ax=None, *, spin=None, emin=-10, emax=5, filename=None,
                 show=False, ylabel=None, colors=None, point_colors=None,
                 label=None, loc=None,
                 cmap=None, cmin=-1.0, cmax=1.0, sortcolors=False,
                 colorbar=True, clabel='$s_z$', cax=None,
                 **plotkwargs):
            """Plot band-structure.

            ax: Axes
                MatPlotLib Axes object.  Will be created if not supplied.
            spin: int or None
                If given, only plot the specified spin channel.
                If None, plot all spins.
                Default: None, i.e., plot all spins.
            emin, emax: float
                Minimum and maximum energy above reference.
            filename: str
                If given, write image to a file.
            show: bool
                Show the image (not needed in notebooks).
            ylabel: str
                The label along the y-axis.  Defaults to 'energies [eV]'
            colors: sequence of str
                A sequence of one or two color specifications, depending on
                whether there is spin.
                Default: green if no spin, yellow and blue if spin is present.
            point_colors: ndarray
                An array of numbers of the shape (nspins, n_kpts, nbands) which
                are then mapped onto colors by the colormap (see ``cmap``).
                ``colors`` and ``point_colors`` are mutually exclusive
            label: str or list of str
                Label for the curves on the legend.  A string if one spin is
                present, a list of two strings if two spins are present.
                Default: If no spin is given, no legend is made; if spin is
                present default labels 'spin up' and 'spin down' are used, but
                can be suppressed by setting ``label=False``.
            loc: str
                Location of the legend.

            If ``point_colors`` is given, the following arguments can be specified.

            cmap:
                Only used if colors is an array of numbers.  A matplotlib
                colormap object, or a string naming a standard colormap.
                Default: The matplotlib default, typically 'viridis'.
            cmin, cmax: float
                Minimal and maximal values used for colormap translation.
                Default: -1.0 and 1.0
            colorbar: bool
                Whether to make a colorbar.
            clabel: str
                Label for the colorbar (default 's_z', set to None to suppress.
            cax: Axes
                Axes object used for plotting colorbar.  Default: split off a
                new one.
            sortcolors (bool or callable):
                Sort points so highest color values are in front.  If a callable is
                given, then it is called on the color values to determine the sort
                order.

            Any additional keyword arguments are passed directly to matplotlib's
            plot() or scatter() methods, depending on whether point_colors is
            given.
            """
            import matplotlib.pyplot as plt

            if colors is not None and point_colors is not None:
                raise ValueError("Don't give both 'color' and 'point_color'")

            if self.ax is None:
                ax = self.prepare_plot(ax, emin, emax, ylabel)

            if spin is None:
                e_skn = self.bs.energies
            elif spin not in [0, 1]:
                raise ValueError(f"spin should be 0 or 1, not {spin}")
            else:
                # Select only one spin channel.
                e_skn = self.bs.energies[spin, np.newaxis]

            nspins = len(e_skn)

            if point_colors is None:
                # Normal band structure plot
                if colors is None:
                    if len(e_skn) == 1:
                        colors = 'g'
                    else:
                        colors = 'yb'
                elif (len(colors) != nspins):
                    raise ValueError(
                        "colors should be a sequence of {nspin} colors"
                    )

                # Default values for label
                if label is None and nspins == 2:
                    label = ['spin up', 'spin down']

                if label:
                    if nspins == 1 and isinstance(label, str):
                        label = [label]
                    elif len(label) != nspins:
                        raise ValueError(
                            f'label should be a list of {nspins} strings'
                        )

                for spin, e_kn in enumerate(e_skn):
                    kwargs = dict(color=colors[spin])
                    kwargs.update(plotkwargs)
                    lbl = None   # Retain lbl=None if label=False
                    if label:
                        lbl = label[spin]
                    ax.plot(self.xcoords, e_kn[:, 0], label=lbl, **kwargs)

                    for e_k in e_kn.T[1:]:
                        ax.plot(self.xcoords, e_k, **kwargs)
                show_legend = label is not None or nspins == 2

            else:
                # A color per datapoint.
                kwargs = dict(vmin=cmin, vmax=cmax, cmap=cmap, s=1)
                kwargs.update(plotkwargs)
                shape = e_skn.shape
                xcoords = np.zeros(shape)
                xcoords += self.xcoords[np.newaxis, :, np.newaxis]
                if sortcolors:
                    if callable(sortcolors):
                        perm = sortcolors(point_colors).argsort(axis=None)
                    else:
                        perm = point_colors.argsort(axis=None)
                    e_skn = e_skn.ravel()[perm].reshape(shape)
                    point_colors = point_colors.ravel()[perm].reshape(shape)
                    xcoords = xcoords.ravel()[perm].reshape(shape)

                things = ax.scatter(xcoords, e_skn, c=point_colors, **kwargs)
                if colorbar:
                    cbar = plt.colorbar(things, cax=cax)
                    if clabel:
                        cbar.set_label(clabel)
                show_legend = False

            self.finish_plot(filename, show, loc, show_legend)

            return ax

        def prepare_plot(self, ax=None, emin=-10, emax=5, ylabel=None):
            import matplotlib.pyplot as plt
            if ax is None:
                ax = plt.figure().add_subplot(111)

            def pretty(kpt):
                if kpt == 'G':
                    kpt = r'$\Gamma$'
                elif len(kpt) == 2:
                    kpt = kpt[0] + '$_' + kpt[1] + '$'
                return kpt

            self.xcoords, label_xcoords, orig_labels = self.bs.get_labels()
            label_xcoords = list(label_xcoords)
            labels = [pretty(name) for name in orig_labels]

            i = 1
            while i < len(labels):
                if label_xcoords[i - 1] == label_xcoords[i]:
                    labels[i - 1] = labels[i - 1] + ',' + labels[i]
                    labels.pop(i)
                    label_xcoords.pop(i)
                else:
                    i += 1

            for x in label_xcoords[1:-1]:
                ax.axvline(x, color='0.5')

            ylabel = ylabel if ylabel is not None else 'energies [eV]'

            ax.set_xticks(label_xcoords)
            ax.set_xticklabels(labels)
            ax.set_ylabel(ylabel)
            ax.axhline(self.bs.reference, color='k', ls=':')
            ax.axis(xmin=0, xmax=self.xcoords[-1], ymin=emin, ymax=emax)
            self.ax = ax
            return ax

        def finish_plot(self, filename, show, loc, show_legend=False):
            import matplotlib.pyplot as plt

            if show_legend:
                leg = plt.legend(loc=loc)
                leg.get_frame().set_alpha(1)

            if filename:
                plt.savefig(filename)

            if show:
                plt.show()

    @jsonable('bandstructure')
    class BandStructure:
        """A band structure consists of an array of eigenvalues and a bandpath.

        BandStructure objects support JSON I/O.
        """

        def __init__(self, path, energies, reference=0.0):
            self._path = path
            self._energies = np.asarray(energies)
            assert self.energies.shape[0] in [1, 2]  # spins x kpts x bands
            assert self.energies.shape[1] == len(path.kpts)
            assert np.isscalar(reference)
            self._reference = reference

        @property
        def energies(self) -> np.ndarray:
            """The energies of this band structure.

            This is a numpy array of shape (nspins, nkpoints, nbands)."""
            return self._energies

        @property
        def path(self) -> 'ase.dft.kpoints.BandPath':
            """The :class:`~ase.dft.kpoints.BandPath` of this band structure."""
            return self._path

        @property
        def reference(self) -> float:
            """The reference energy.

            Semantics may vary; typically a Fermi energy or zero,
            depending on how the band structure was created."""
            return self._reference

        def subtract_reference(self) -> 'BandStructure':
            """Return new band structure with reference energy subtracted."""
            return BandStructure(self.path, self.energies - self.reference,
                                 reference=0.0)

        def todict(self):
            return dict(path=self.path,
                        energies=self.energies,
                        reference=self.reference)

        def get_labels(self, eps=1e-5):
            """"See :func:`ase.dft.kpoints.labels_from_kpts`."""
            return self.path.get_linear_kpoint_axis(eps=eps)

        def plot(self, *args, **kwargs):
            """Plot this band structure."""
            bsp = BandStructurePlot(self)
            return bsp.plot(*args, **kwargs)

        def __repr__(self):
            return ('{}(path={!r}, energies=[{} values], reference={})'
                    .format(self.__class__.__name__, self.path,
                            '{}x{}x{}'.format(*self.energies.shape),
                            self.reference))