import collections
from functools import reduce, singledispatch
from typing import (Any, Dict, Iterable, List, Optional, Sequence, TypeVar,
Union, overload)
import numpy as np
from ase.spectrum.dosdata import DOSData, Floats, GridDOSData, Info, RawDOSData
from ase.utils.plotting import SimplePlottingAxes
# This import is for the benefit of type-checking / mypy
if False:
import matplotlib.axes
[docs]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: Floats,
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])
[docs] 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)
[docs] 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])
[docs] @classmethod
def from_data(cls,
energies: Floats,
weights: Sequence[Floats],
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[Floats],
info: Optional[Sequence[Info]],
) -> 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 len(self) != len(other):
return False
else:
return all(a._almost_equals(b) for a, b in zip(self, other))
[docs] 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
[docs] 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()))]
[docs] 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)
[docs] 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)
[docs] 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: Union[DOSData, 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])
[docs]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.")
[docs]class GridDOSCollection(DOSCollection):
def __init__(self, dos_series: Iterable[GridDOSData],
energies: Optional[Floats] = 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) -> Floats:
return self._energies.copy()
def get_all_weights(self) -> Union[Sequence[Floats], np.ndarray]:
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")
[docs] @classmethod
def from_data(cls,
energies: Floats,
weights: Sequence[Floats],
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
[docs] 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)
[docs] 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)
[docs] 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: Floats,
all_y: np.ndarray,
all_labels: Sequence[str],
mplargs: Optional[Dict]):
"""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)