index | modules | gitlab

# Source code for ase.spectrum.doscollection

    import collections
    from functools import reduce, singledispatch
    from typing import (Any, Dict, Iterable, List, Optional,
                        overload, Sequence, TypeVar, Union)

    import numpy as np
    from ase.spectrum.dosdata import DOSData, RawDOSData, GridDOSData, Info
    from ase.utils.plotting import SimplePlottingAxes

    # This import is for the benefit of type-checking / mypy
    if False:
        import matplotlib.axes

    class DOSCollection(collections.abc.Sequence):
        """Base class for a collection of DOSData objects"""
        def __init__(self, dos_series: Iterable[DOSData]) -> None:
            self._data = list(dos_series)

        def _sample(self,
                    energies: Sequence[float],
                    width: float = 0.1,
                    smearing: str = 'Gauss') -> np.ndarray:
            """Sample the DOS data at chosen points, with broadening

            This samples the underlying DOS data in the same way as the .sample()
            method of those DOSData items, returning a 2-D array with columns
            corresponding to x and rows corresponding to the collected data series.

            Args:
                energies: energy values for sampling
                width: Width of broadening kernel
                smearing: selection of broadening kernel (only "Gauss" is currently
                    supported)

            Returns:
                Weights sampled from a broadened DOS at values corresponding to x,
                in rows corresponding to DOSData entries contained in this object
            """

            if len(self) == 0:
                raise IndexError("No data to sample")

            return np.asarray(
                [data._sample(energies, width=width, smearing=smearing)
                 for data in self])

        def plot(self,
                 npts: int = 1000,
                 xmin: float = None,
                 xmax: float = None,
                 width: float = 0.1,
                 smearing: str = 'Gauss',
                 ax: 'matplotlib.axes.Axes' = None,
                 show: bool = False,
                 filename: str = None,
                 mplargs: dict = None) -> 'matplotlib.axes.Axes':
            """Simple plot of collected DOS data, resampled onto a grid

            If the special key 'label' is present in self.info, this will be set
            as the label for the plotted line (unless overruled in mplargs). The
            label is only seen if a legend is added to the plot (i.e. by calling
            `ax.legend()`).

            Args:
                npts, xmin, xmax: output data range, as passed to self.sample_grid
                width: Width of broadening kernel, passed to self.sample_grid()
                smearing: selection of broadening kernel for self.sample_grid()
                ax: existing Matplotlib axes object. If not provided, a new figure
                    with one set of axes will be created using Pyplot
                show: show the figure on-screen
                filename: if a path is given, save the figure to this file
                mplargs: additional arguments to pass to matplotlib plot command
                    (e.g. {'linewidth': 2} for a thicker line).

            Returns:
                Plotting axes. If "ax" was set, this is the same object.
            """
            return self.sample_grid(npts,
                                    xmin=xmin, xmax=xmax,
                                    width=width, smearing=smearing
                                    ).plot(npts=npts,
                                           xmin=xmin, xmax=xmax,
                                           width=width, smearing=smearing,
                                           ax=ax, show=show, filename=filename,
                                           mplargs=mplargs)

        def sample_grid(self,
                        npts: int,
                        xmin: float = None,
                        xmax: float = None,
                        padding: float = 3,
                        width: float = 0.1,
                        smearing: str = 'Gauss',
                        ) -> 'GridDOSCollection':
            """Sample the DOS data on an evenly-spaced energy grid

            Args:
                npts: Number of sampled points
                xmin: Minimum sampled energy value; if unspecified, a default is
                    chosen
                xmax: Maximum sampled energy value; if unspecified, a default is
                    chosen
                padding: If xmin/xmax is unspecified, default value will be padded
                    by padding * width to avoid cutting off peaks.
                width: Width of broadening kernel, passed to self.sample_grid()
                smearing: selection of broadening kernel, for self.sample_grid()

            Returns:
                (energy values, sampled DOS)
            """
            if len(self) == 0:
                raise IndexError("No data to sample")

            if xmin is None:
                xmin = (min(min(data.get_energies()) for data in self)
                        - (padding * width))
            if xmax is None:
                xmax = (max(max(data.get_energies()) for data in self)
                        + (padding * width))

            return GridDOSCollection(
                [data.sample_grid(npts, xmin=xmin, xmax=xmax, width=width,
                                  smearing=smearing)
                 for data in self])

        @classmethod
        def from_data(cls,
                      energies: Sequence[float],
                      weights: Sequence[Sequence[float]],
                      info: Sequence[Info] = None) -> 'DOSCollection':
            """Create a DOSCollection from data sharing a common set of energies

            This is a convenience method to be used when all the DOS data in the
            collection has a common energy axis. There is no performance advantage
            in using this method for the generic DOSCollection, but for
            GridDOSCollection it is more efficient.

            Args:
                energy: common set of energy values for input data
                weights: array of DOS weights with rows corresponding to different
                    datasets
                info: sequence of info dicts corresponding to weights rows.

            Returns:
                Collection of DOS data (in RawDOSData format)
            """

            info = cls._check_weights_and_info(weights, info)

            return cls(RawDOSData(energies, row_weights, row_info)
                       for row_weights, row_info in zip(weights, info))

        @staticmethod
        def _check_weights_and_info(weights: Sequence[Sequence[float]],
                                    info: Union[Sequence[Info], None],
                                    ) -> Sequence[Info]:
            if info is None:
                info = [{} for _ in range(len(weights))]
            else:
                if len(info) != len(weights):
                    raise ValueError("Length of info must match number of rows in "
                                     "weights")
            return info

        @overload
        def __getitem__(self, item: int) -> DOSData:
            ...

        @overload  # noqa F811
        def __getitem__(self, item: slice) -> 'DOSCollection':  # noqa F811
            ...

        def __getitem__(self, item):  # noqa F811
            if isinstance(item, int):
                return self._data[item]
            elif isinstance(item, slice):
                return type(self)(self._data[item])
            else:
                raise TypeError("index in DOSCollection must be an integer or "
                                "slice")

        def __len__(self) -> int:
            return len(self._data)

        def _almost_equals(self, other: Any) -> bool:
            """Compare with another DOSCollection for testing purposes"""
            if not isinstance(other, type(self)):
                return False
            elif not len(self) == len(other):
                return False
            else:
                return all([a._almost_equals(b) for a, b in zip(self, other)])

        def total(self) -> DOSData:
            """Sum all the DOSData in this Collection and label it as 'Total'"""
            data = self.sum_all()
            data.info.update({'label': 'Total'})
            return data

        def sum_all(self) -> DOSData:
            """Sum all the DOSData contained in this Collection"""
            if len(self) == 0:
                raise IndexError("No data to sum")
            elif len(self) == 1:
                data = self[0].copy()
            else:
                data = reduce(lambda x, y: x + y, self)
            return data

        D = TypeVar('D', bound=DOSData)

        @staticmethod
        def _select_to_list(dos_collection: Sequence[D],         # Bug in flakes
                            info_selection: Dict[str, str],      # misses 'D' def
                            negative: bool = False) -> List[D]:  # noqa: F821
            query = set(info_selection.items())

            if negative:
                return [data for data in dos_collection
                        if not query.issubset(set(data.info.items()))]
            else:
                return [data for data in dos_collection
                        if query.issubset(set(data.info.items()))]

        def select(self, **info_selection: str) -> 'DOSCollection':
            """Narrow DOSCollection to items with specified info

            For example, if ::

              dc = DOSCollection([DOSData(x1, y1, info={'a': '1', 'b': '1'}),
                                  DOSData(x2, y2, info={'a': '2', 'b': '1'})])

            then ::

              dc.select(b='1')

            will return an identical object to dc, while ::

              dc.select(a='1')

            will return a DOSCollection with only the first item and ::

              dc.select(a='2', b='1')

            will return a DOSCollection with only the second item.

            """

            matches = self._select_to_list(self, info_selection)
            return type(self)(matches)

        def select_not(self, **info_selection: str) -> 'DOSCollection':
            """Narrow DOSCollection to items without specified info

            For example, if ::

              dc = DOSCollection([DOSData(x1, y1, info={'a': '1', 'b': '1'}),
                                  DOSData(x2, y2, info={'a': '2', 'b': '1'})])

            then ::

              dc.select_not(b='2')

            will return an identical object to dc, while ::

              dc.select_not(a='2')

            will return a DOSCollection with only the first item and ::

              dc.select_not(a='1', b='1')

            will return a DOSCollection with only the second item.

            """
            matches = self._select_to_list(self, info_selection, negative=True)
            return type(self)(matches)

        def sum_by(self, *info_keys: str) -> 'DOSCollection':
            """Return a DOSCollection with some data summed by common attributes

            For example, if ::

              dc = DOSCollection([DOSData(x1, y1, info={'a': '1', 'b': '1'}),
                                  DOSData(x2, y2, info={'a': '2', 'b': '1'}),
                                  DOSData(x3, y3, info={'a': '2', 'b': '2'})])

            then ::

              dc.sum_by('b')

            will return a collection equivalent to ::

              DOSCollection([DOSData(x1, y1, info={'a': '1', 'b': '1'})
                             + DOSData(x2, y2, info={'a': '2', 'b': '1'}),
                             DOSData(x3, y3, info={'a': '2', 'b': '2'})])

            where the resulting contained DOSData have info attributes of
            {'b': '1'} and {'b': '2'} respectively.

            dc.sum_by('a', 'b') on the other hand would return the full three-entry
            collection, as none of the entries have common 'a' *and* 'b' info.

            """

            def _matching_info_tuples(data: DOSData):
                """Get relevent dict entries in tuple form

                e.g. if data.info = {'a': 1, 'b': 2, 'c': 3}
                     and info_keys = ('a', 'c')

                     then return (('a', 1), ('c': 3))
                """
                matched_keys = set(info_keys) & set(data.info)
                return tuple(sorted([(key, data.info[key])
                                     for key in matched_keys]))

            # Sorting inside info matching helps set() to remove redundant matches;
            # combos are then sorted() to ensure consistent output across sessions.
            all_combos = map(_matching_info_tuples, self)
            unique_combos = sorted(set(all_combos))

            # For each key/value combination, perform a select() to obtain all
            # the matching entries and sum them together.
            collection_data = [self.select(**dict(combo)).sum_all()
                               for combo in unique_combos]
            return type(self)(collection_data)

        def __add__(self, other: Union['DOSCollection', DOSData]
                    ) -> 'DOSCollection':
            """Join entries between two DOSCollection objects of the same type

            It is also possible to add a single DOSData object without wrapping it
            in a new collection: i.e. ::

              DOSCollection([dosdata1]) + DOSCollection([dosdata2])

            or ::

              DOSCollection([dosdata1]) + dosdata2

            will return ::

              DOSCollection([dosdata1, dosdata2])

            """
            return _add_to_collection(other, self)

    @singledispatch
    def _add_to_collection(other: DOSCollection,
                           collection: DOSCollection) -> DOSCollection:
        if isinstance(other, type(collection)):
            return type(collection)(list(collection) + list(other))
        elif isinstance(other, DOSCollection):
            raise TypeError("Only DOSCollection objects of the same type may "
                            "be joined with '+'.")
        else:
            raise TypeError("DOSCollection may only be joined to DOSData or "
                            "DOSCollection objects with '+'.")

    @_add_to_collection.register(DOSData)
    def _add_data(other: DOSData, collection: DOSCollection) -> DOSCollection:
        """Return a new DOSCollection with an additional DOSData item"""
        return type(collection)(list(collection) + [other])

    class RawDOSCollection(DOSCollection):
        def __init__(self, dos_series: Iterable[RawDOSData]) -> None:
            super().__init__(dos_series)
            for dos_data in self:
                if not isinstance(dos_data, RawDOSData):
                    raise TypeError("RawDOSCollection can only store "
                                    "RawDOSData objects.")

    class GridDOSCollection(DOSCollection):
        def __init__(self, dos_series: Iterable[GridDOSData],
                     energies: Optional[Sequence[float]] = None) -> None:
            dos_list = list(dos_series)
            if energies is None:
                if len(dos_list) == 0:
                    raise ValueError("Must provide energies to create a "
                                     "GridDOSCollection without any DOS data.")
                self._energies = dos_list[0].get_energies()
            else:
                self._energies = np.asarray(energies)

            self._weights = np.empty((len(dos_list), len(self._energies)), float)
            self._info = []

            for i, dos_data in enumerate(dos_list):
                if not isinstance(dos_data, GridDOSData):
                    raise TypeError("GridDOSCollection can only store "
                                    "GridDOSData objects.")
                if (dos_data.get_energies().shape != self._energies.shape
                    or not np.allclose(dos_data.get_energies(), self._energies)):
                    raise ValueError("All GridDOSData objects in GridDOSCollection"
                                     " must have the same energy axis.")
                self._weights[i, :] = dos_data.get_weights()
                self._info.append(dos_data.info)

        def get_energies(self) -> Sequence[float]:
            return self._energies.copy()

        def get_all_weights(self) -> Sequence[Sequence[float]]:
            return self._weights.copy()

        def __len__(self) -> int:
            return self._weights.shape[0]

        @overload  # noqa F811
        def __getitem__(self, item: int) -> DOSData:
            ...

        @overload  # noqa F811
        def __getitem__(self, item: slice) -> 'GridDOSCollection':  # noqa F811
            ...

        def __getitem__(self, item):  # noqa F811
            if isinstance(item, int):
                return GridDOSData(self._energies, self._weights[item, :],
                                   info=self._info[item])
            elif isinstance(item, slice):
                return type(self)([self[i] for i in range(len(self))[item]])
            else:
                raise TypeError("index in DOSCollection must be an integer or "
                                "slice")

        @classmethod
        def from_data(cls,
                      energies: Sequence[float],
                      weights: Sequence[Sequence[float]],
                      info: Sequence[Info] = None) -> 'GridDOSCollection':
            """Create a GridDOSCollection from data with a common set of energies

            This convenience method may also be more efficient as it limits
            redundant copying/checking of the data.

            Args:
                energies: common set of energy values for input data
                weights: array of DOS weights with rows corresponding to different
                    datasets
                info: sequence of info dicts corresponding to weights rows.

            Returns:
                Collection of DOS data (in RawDOSData format)
            """

            weights_array = np.asarray(weights, dtype=float)
            if len(weights_array.shape) != 2:
                raise IndexError("Weights must be a 2-D array or nested sequence")
            if weights_array.shape[0] < 1:
                raise IndexError("Weights cannot be empty")
            if weights_array.shape[1] != len(energies):
                raise IndexError("Length of weights rows must equal size of x")

            info = cls._check_weights_and_info(weights, info)

            dos_collection = cls([GridDOSData(energies, weights_array[0])])
            dos_collection._weights = weights_array
            dos_collection._info = list(info)

            return dos_collection

        def select(self, **info_selection: str) -> 'DOSCollection':
            """Narrow GridDOSCollection to items with specified info

            For example, if ::

              dc = GridDOSCollection([GridDOSData(x, y1,
                                                  info={'a': '1', 'b': '1'}),
                                      GridDOSData(x, y2,
                                                  info={'a': '2', 'b': '1'})])

            then ::

              dc.select(b='1')

            will return an identical object to dc, while ::

              dc.select(a='1')

            will return a DOSCollection with only the first item and ::

              dc.select(a='2', b='1')

            will return a DOSCollection with only the second item.

            """

            matches = self._select_to_list(self, info_selection)
            if len(matches) == 0:
                return type(self)([], energies=self._energies)
            else:
                return type(self)(matches)

        def select_not(self, **info_selection: str) -> 'DOSCollection':
            """Narrow GridDOSCollection to items without specified info

            For example, if ::

              dc = GridDOSCollection([GridDOSData(x, y1,
                                                  info={'a': '1', 'b': '1'}),
                                      GridDOSData(x, y2,
                                                  info={'a': '2', 'b': '1'})])

            then ::

              dc.select_not(b='2')

            will return an identical object to dc, while ::

              dc.select_not(a='2')

            will return a DOSCollection with only the first item and ::

              dc.select_not(a='1', b='1')

            will return a DOSCollection with only the second item.

            """
            matches = self._select_to_list(self, info_selection, negative=True)
            if len(matches) == 0:
                return type(self)([], energies=self._energies)
            else:
                return type(self)(matches)

        def plot(self,
                 npts: int = 0,
                 xmin: float = None,
                 xmax: float = None,
                 width: float = None,
                 smearing: str = 'Gauss',
                 ax: 'matplotlib.axes.Axes' = None,
                 show: bool = False,
                 filename: str = None,
                 mplargs: dict = None) -> 'matplotlib.axes.Axes':
            """Simple plot of collected DOS data, resampled onto a grid

            If the special key 'label' is present in self.info, this will be set
            as the label for the plotted line (unless overruled in mplargs). The
            label is only seen if a legend is added to the plot (i.e. by calling
            `ax.legend()`).

            Args:
                npts:
                    Number of points in resampled x-axis. If set to zero (default),
                    no resampling is performed and the stored data is plotted
                    directly.
                xmin, xmax:
                    output data range; this limits the resampling range as well as
                    the plotting output
                width: Width of broadening kernel, passed to self.sample()
                smearing: selection of broadening kernel, passed to self.sample()
                ax: existing Matplotlib axes object. If not provided, a new figure
                    with one set of axes will be created using Pyplot
                show: show the figure on-screen
                filename: if a path is given, save the figure to this file
                mplargs: additional arguments to pass to matplotlib plot command
                    (e.g. {'linewidth': 2} for a thicker line).

            Returns:
                Plotting axes. If "ax" was set, this is the same object.
            """

            # Apply defaults if necessary
            npts, width = GridDOSData._interpret_smearing_args(npts, width)

            if npts:
                assert isinstance(width, float)
                dos = self.sample_grid(npts,
                                       xmin=xmin, xmax=xmax,
                                       width=width, smearing=smearing)
            else:
                dos = self

            energies, all_y = dos._energies, dos._weights

            all_labels = [DOSData.label_from_info(data.info) for data in self]

            with SimplePlottingAxes(ax=ax, show=show, filename=filename) as ax:
                self._plot_broadened(ax, energies, all_y, all_labels, mplargs)

            return ax

        @staticmethod
        def _plot_broadened(ax: 'matplotlib.axes.Axes',
                            energies: Sequence[float],
                            all_y: np.ndarray,
                            all_labels: Sequence[str],
                            mplargs: Union[Dict, None]):
            """Plot DOS data with labels to axes

            This is separated into another function so that subclasses can
            manipulate broadening, labels etc in their plot() method."""
            if mplargs is None:
                mplargs = {}

            all_lines = ax.plot(energies, all_y.T, **mplargs)
            for line, label in zip(all_lines, all_labels):
                line.set_label(label)
            ax.legend()

            ax.set_xlim(left=min(energies), right=max(energies))
            ax.set_ylim(bottom=0)