def dosplot(filename=None, code='vasp', prefix=None, directory=None,
elements=None, lm_orbitals=None, atoms=None, subplot=False,
shift=True, total_only=False, plot_total=True, legend_on=True,
legend_frame_on=False, legend_cutoff=3., gaussian=None, height=6.,
width=8., xmin=-6., xmax=6., num_columns=2, colours=None, yscale=1,
xlabel='Energy (eV)', ylabel='Arb. units',
style=None, no_base_style=False,
image_format='pdf', dpi=400, plt=None, fonts=None):
"""A script to plot the density of states from a vasprun.xml file.
Args:
filename (:obj:`str`, optional): Path to a vasprun.xml file (can be
gzipped).
code (:obj:`str`, optional): Electronic structure code used ('vasp' or
'questaal')
prefix (:obj:`str`, optional): Prefix for file names.
directory (:obj:`str`, optional): The directory in which to save files.
elements (:obj:`dict`, optional): The elements and orbitals to extract
from the projected density of states. Should be provided as a
:obj:`dict` with the keys as the element names and corresponding
values as a :obj:`tuple` of orbitals. For example, the following
would extract the Bi s, px, py and d orbitals::
{'Bi': ('s', 'px', 'py', 'd')}
If an element is included with an empty :obj:`tuple`, all orbitals
for that species will be extracted. If ``elements`` is not set or
set to ``None``, all elements for all species will be extracted.
lm_orbitals (:obj:`dict`, optional): The orbitals to decompose into
their lm contributions (e.g. p -> px, py, pz). Should be provided
as a :obj:`dict`, with the elements names as keys and a
:obj:`tuple` of orbitals as the corresponding values. For example,
the following would be used to decompose the oxygen p and d
orbitals::
{'O': ('p', 'd')}
atoms (:obj:`dict`, optional): Which atomic sites to use when
calculating the projected density of states. Should be provided as
a :obj:`dict`, with the element names as keys and a :obj:`tuple` of
:obj:`int` specifying the atomic indices as the corresponding
values. The elemental projected density of states will be summed
only over the atom indices specified. If an element is included
with an empty :obj:`tuple`, then all sites for that element will
be included. The indices are 0 based for each element specified in
the POSCAR. For example, the following will calculate the density
of states for the first 4 Sn atoms and all O atoms in the
structure::
{'Sn': (1, 2, 3, 4), 'O': (, )}
If ``atoms`` is not set or set to ``None`` then all atomic sites
for all elements will be considered.
subplot (:obj:`bool`, optional): Plot the density of states for each
element on separate subplots. Defaults to ``False``.
shift (:obj:`bool`, optional): Shift the energies such that the valence
band maximum (or Fermi level for metals) is at 0 eV. Defaults to
``True``.
total_only (:obj:`bool`, optional): Only extract the total density of
states. Defaults to ``False``.
plot_total (:obj:`bool`, optional): Plot the total density of states.
Defaults to ``True``.
legend_on (:obj:`bool`, optional): Plot the graph legend. Defaults
to ``True``.
legend_frame_on (:obj:`bool`, optional): Plot a frame around the
graph legend. Defaults to ``False``.
legend_cutoff (:obj:`float`, optional): The cut-off (in % of the
maximum density of states within the plotting range) for an
elemental orbital to be labelled in the legend. This prevents
the legend from containing labels for orbitals that have very
little contribution in the plotting range.
gaussian (:obj:`float`, optional): Broaden the density of states using
convolution with a gaussian function. This parameter controls the
sigma or standard deviation of the gaussian distribution.
height (:obj:`float`, optional): The height of the plot.
width (:obj:`float`, optional): The width of the plot.
xmin (:obj:`float`, optional): The minimum energy on the x-axis.
xmax (:obj:`float`, optional): The maximum energy on the x-axis.
num_columns (:obj:`int`, optional): The number of columns in the
legend.
colours (:obj:`dict`, optional): Use custom colours for specific
element and orbital combinations. Specified as a :obj:`dict` of
:obj:`dict` of the colours. For example::
{
'Sn': {'s': 'r', 'p': 'b'},
'O': {'s': '#000000'}
}
The colour can be a hex code, series of rgb value, or any other
format supported by matplotlib.
xlabel (:obj:`str`, optional): Label/units for x-axis (i.e. energy)
ylabel (:obj:`str`, optional): Label/units for y-axis (i.e. DOS)
yscale (:obj:`float`, optional): Scaling factor for the y-axis.
style (:obj:`list` or :obj:`str`, optional): (List of) matplotlib style
specifications, to be composed on top of Sumo base style.
no_base_style (:obj:`bool`, optional): Prevent use of sumo base style.
This can make alternative styles behave more predictably.
image_format (:obj:`str`, optional): The image file format. Can be any
format supported by matplotlib, including: png, jpg, pdf, and svg.
Defaults to pdf.
dpi (:obj:`int`, optional): The dots-per-inch (pixel density) for
the image.
plt (:obj:`matplotlib.pyplot`, optional): A
:obj:`matplotlib.pyplot` object to use for plotting.
fonts (:obj:`list`, optional): Fonts to use in the plot. Can be a
a single font, specified as a :obj:`str`, or several fonts,
specified as a :obj:`list` of :obj:`str`.
Returns:
A matplotlib pyplot object.
"""
if code.lower() == 'vasp':
if not filename:
if os.path.exists('vasprun.xml'):
filename = 'vasprun.xml'
elif os.path.exists('vasprun.xml.gz'):
filename = 'vasprun.xml.gz'
else:
logging.error('ERROR: No vasprun.xml found!')
sys.exit()
dos, pdos = load_dos(filename, elements, lm_orbitals, atoms, gaussian,
total_only)
elif code.lower() == 'questaal':
if filename:
pdos_file = filename
ext = pdos_file.split('.')[-1]
else:
pdos_candidates = glob('dos.*')
for candidate in pdos_candidates:
if candidate.split('.')[-1] in ('pdf', 'png', 'svg',
'jpg', 'jpeg'):
continue
elif candidate.split('.')[-1].lower() in ('gz', 'z', 'bz2'):
pdos_file = candidate
ext = candidate.split('.')[-2]
break
else:
pdos_file = candidate
ext = candidate.split('.')[-1]
break
else:
raise ValueError("No questaal dos file found")
if os.path.exists('tdos.{}'.format(ext)):
tdos_file = 'tdos.{}'.format(ext)
else:
tdos_file = None
if os.path.exists('site.{}'.format(ext)):
site_file = 'site.{}'.format(ext)
else:
site_file = None
if shift:
logging.warning("Fermi level shift requested, but not implemented "
"for Questaal DOS.")
dos, pdos = sumo.io.questaal.read_dos(
pdos_file=pdos_file, tdos_file=tdos_file, site_file=site_file,
ry=True, gaussian=gaussian, total_only=total_only,
elements=elements, lm_orbitals=lm_orbitals, atoms=atoms)
save_files = False if plt else True # don't save if pyplot object provided
plotter = SDOSPlotter(dos, pdos)
plt = plotter.get_plot(subplot=subplot, width=width, height=height,
xmin=xmin, xmax=xmax, yscale=yscale,
colours=colours, plot_total=plot_total,
legend_on=legend_on, num_columns=num_columns,
legend_frame_on=legend_frame_on,
xlabel=xlabel, ylabel=ylabel,
legend_cutoff=legend_cutoff, dpi=dpi, plt=plt,
fonts=fonts, style=style,
no_base_style=no_base_style)
if save_files:
basename = 'dos.{}'.format(image_format)
filename = '{}_{}'.format(prefix, basename) if prefix else basename
if directory:
filename = os.path.join(directory, filename)
plt.savefig(filename, format=image_format, dpi=dpi,
bbox_inches='tight')
write_files(dos, pdos, prefix=prefix, directory=directory)
else:
return plt
def dosplot(
filename=None,
code="vasp",
prefix=None,
directory=None,
elements=None,
lm_orbitals=None,
atoms=None,
spin=None,
subplot=False,
shift=True,
total_only=False,
plot_total=True,
legend_on=True,
legend_frame_on=False,
legend_cutoff=3.0,
gaussian=None,
colours=None,
height=6.0,
width=8.0,
xmin=-6.0,
xmax=6.0,
num_columns=2,
yscale=1,
xlabel="Energy (eV)",
ylabel="Arb. units",
style=None,
no_base_style=False,
image_format="pdf",
dpi=400,
plt=None,
fonts=None,
):
"""A script to plot the density of states from a vasprun.xml file.
Args:
filename (:obj:`str`, optional): Path to a DOS data file (can be
gzipped). The preferred file type depends on the electronic
structure code: vasprun.xml (VASP); *.bands (CASTEP); dos.*
(Questaal).
code (:obj:`str`, optional): Electronic structure code used ('vasp',
'castep' or 'questaal'). Note that for Castep only a rough TDOS
is available, assembled by sampling the eigenvalues.
prefix (:obj:`str`, optional): Prefix for file names.
directory (:obj:`str`, optional): The directory in which to save files.
elements (:obj:`dict`, optional): The elements and orbitals to extract
from the projected density of states. Should be provided as a
:obj:`dict` with the keys as the element names and corresponding
values as a :obj:`tuple` of orbitals. For example, the following
would extract the Bi s, px, py and d orbitals::
{'Bi': ('s', 'px', 'py', 'd')}
If an element is included with an empty :obj:`tuple`, all orbitals
for that species will be extracted. If ``elements`` is not set or
set to ``None``, all elements for all species will be extracted.
lm_orbitals (:obj:`dict`, optional): The orbitals to decompose into
their lm contributions (e.g. p -> px, py, pz). Should be provided
as a :obj:`dict`, with the elements names as keys and a
:obj:`tuple` of orbitals as the corresponding values. For example,
the following would be used to decompose the oxygen p and d
orbitals::
{'O': ('p', 'd')}
atoms (:obj:`dict`, optional): Which atomic sites to use when
calculating the projected density of states. Should be provided as
a :obj:`dict`, with the element names as keys and a :obj:`tuple` of
:obj:`int` specifying the atomic indices as the corresponding
values. The elemental projected density of states will be summed
only over the atom indices specified. If an element is included
with an empty :obj:`tuple`, then all sites for that element will
be included. The indices are 0 based for each element specified in
the POSCAR. For example, the following will calculate the density
of states for the first 4 Sn atoms and all O atoms in the
structure::
{'Sn': (1, 2, 3, 4), 'O': (, )}
If ``atoms`` is not set or set to ``None`` then all atomic sites
for all elements will be considered.
spin (:obj:`Spin`, optional): Plot only one spin channel from a
spin-polarised calculation; "up" or "1" for spin up only, "down" or
"-1" for spin down only. Defaults to ``None``.
subplot (:obj:`bool`, optional): Plot the density of states for each
element on separate subplots. Defaults to ``False``.
shift (:obj:`bool`, optional): Shift the energies such that the valence
band maximum (or Fermi level for metals) is at 0 eV. Defaults to
``True``.
total_only (:obj:`bool`, optional): Only extract the total density of
states. Defaults to ``False``.
plot_total (:obj:`bool`, optional): Plot the total density of states.
Defaults to ``True``.
legend_on (:obj:`bool`, optional): Plot the graph legend. Defaults
to ``True``.
legend_frame_on (:obj:`bool`, optional): Plot a frame around the
graph legend. Defaults to ``False``.
legend_cutoff (:obj:`float`, optional): The cut-off (in % of the
maximum density of states within the plotting range) for an
elemental orbital to be labelled in the legend. This prevents
the legend from containing labels for orbitals that have very
little contribution in the plotting range.
gaussian (:obj:`float`, optional): Broaden the density of states using
convolution with a gaussian function. This parameter controls the
sigma or standard deviation of the gaussian distribution.
height (:obj:`float`, optional): The height of the plot.
width (:obj:`float`, optional): The width of the plot.
xmin (:obj:`float`, optional): The minimum energy on the x-axis.
xmax (:obj:`float`, optional): The maximum energy on the x-axis.
num_columns (:obj:`int`, optional): The number of columns in the
legend.
colours (:obj:`dict`, optional): Use custom colours for specific
element and orbital combinations. Specified as a :obj:`dict` of
:obj:`dict` of the colours. For example::
{
'Sn': {'s': 'r', 'p': 'b'},
'O': {'s': '#000000'}
}
The colour can be a hex code, series of rgb value, or any other
format supported by matplotlib.
xlabel (:obj:`str`, optional): Label/units for x-axis (i.e. energy)
ylabel (:obj:`str`, optional): Label/units for y-axis (i.e. DOS)
yscale (:obj:`float`, optional): Scaling factor for the y-axis.
style (:obj:`list` or :obj:`str`, optional): (List of) matplotlib style
specifications, to be composed on top of Sumo base style.
no_base_style (:obj:`bool`, optional): Prevent use of sumo base style.
This can make alternative styles behave more predictably.
image_format (:obj:`str`, optional): The image file format. Can be any
format supported by matplotlib, including: png, jpg, pdf, and svg.
Defaults to pdf.
dpi (:obj:`int`, optional): The dots-per-inch (pixel density) for
the image.
plt (:obj:`matplotlib.pyplot`, optional): A
:obj:`matplotlib.pyplot` object to use for plotting.
fonts (:obj:`list`, optional): Fonts to use in the plot. Can be a
a single font, specified as a :obj:`str`, or several fonts,
specified as a :obj:`list` of :obj:`str`.
Returns:
A matplotlib pyplot object.
"""
if code.lower() == "vasp":
if not filename:
if os.path.exists("vasprun.xml"):
filename = "vasprun.xml"
elif os.path.exists("vasprun.xml.gz"):
filename = "vasprun.xml.gz"
else:
logging.error("ERROR: No vasprun.xml found!")
sys.exit()
dos, pdos = load_dos(
filename, elements, lm_orbitals, atoms, gaussian, total_only
)
elif code.lower() == "castep":
if filename:
bands_file = filename
else:
band_candidates = glob("*.bands")
if len(band_candidates) == 0:
logging.error("ERROR: No *.bands file found!")
sys.exit()
elif len(band_candidates) == 1:
bands_file = band_candidates[0]
else:
logging.error("ERROR: Too many *.bands files found!")
sys.exit()
pdos_file = _replace_ext(bands_file, "pdos_bin")
cell_file = _replace_ext(bands_file, "cell")
pdos_file = pdos_file if os.path.isfile(pdos_file) else None
cell_file = cell_file if os.path.isfile(cell_file) else None
if not total_only:
# Check if both pdos_bin and cell files are present.
# If not, we cannot plot the PDOS
if pdos_file is not None:
if cell_file is None:
logging.info(
"Plotting PDOS requires the .cell file to be "
"present; falling back to TDOS."
)
pdos_file = None
else:
logging.info(
f"Found PDOS binary file {pdos_file}; "
"including PDOS in the plot."
)
else:
logging.info("PDOS not available, falling back to TDOS.")
dos, pdos = sumo.io.castep.read_dos(
bands_file,
pdos_file=pdos_file,
cell_file=cell_file,
gaussian=gaussian,
emin=xmin,
emax=xmax,
lm_orbitals=lm_orbitals,
elements=elements,
total_only=total_only,
atoms=atoms,
)
elif code.lower() == "questaal":
if filename:
pdos_file = filename
ext = pdos_file.split(".")[-1]
else:
pdos_candidates = glob("dos.*")
for candidate in pdos_candidates:
if candidate.split(".")[-1] in ("pdf", "png", "svg", "jpg", "jpeg"):
continue
elif candidate.split(".")[-1].lower() in ("gz", "z", "bz2"):
pdos_file = candidate
ext = candidate.split(".")[-2]
break
else:
pdos_file = candidate
ext = candidate.split(".")[-1]
break
else:
raise ValueError("No questaal dos file found")
if os.path.exists(f"tdos.{ext}"):
tdos_file = f"tdos.{ext}"
else:
tdos_file = None
if os.path.exists(f"site.{ext}"):
site_file = f"site.{ext}"
else:
site_file = None
if shift:
logging.warning(
"Fermi level shift requested, but not implemented " "for Questaal DOS."
)
dos, pdos = sumo.io.questaal.read_dos(
pdos_file=pdos_file,
tdos_file=tdos_file,
site_file=site_file,
ry=True,
gaussian=gaussian,
total_only=total_only,
elements=elements,
lm_orbitals=lm_orbitals,
atoms=atoms,
)
else:
logging.error(f"ERROR: Unrecognised code: {code}")
return
save_files = False if plt else True # don't save if pyplot object provided
spin = string_to_spin(spin) # Convert spin name to pymatgen Spin object
plotter = SDOSPlotter(dos, pdos)
plt = plotter.get_plot(
subplot=subplot,
width=width,
height=height,
xmin=xmin,
xmax=xmax,
yscale=yscale,
colours=colours,
plot_total=plot_total,
legend_on=legend_on,
num_columns=num_columns,
legend_frame_on=legend_frame_on,
xlabel=xlabel,
ylabel=ylabel,
legend_cutoff=legend_cutoff,
dpi=dpi,
plt=plt,
fonts=fonts,
style=style,
no_base_style=no_base_style,
spin=spin,
)
if save_files:
basename = f"dos.{image_format}"
filename = f"{prefix}_{basename}" if prefix else basename
if directory:
filename = os.path.join(directory, filename)
plt.savefig(filename, format=image_format, dpi=dpi, bbox_inches="tight")
write_files(dos, pdos, prefix=prefix, directory=directory)
else:
return plt