# coding: utf-8
"""
Tools to analyze the V1QAVG file produced by the E-PH code (eph_task +15 or -15)
"""
from __future__ import annotations
import numpy as np
from collections import OrderedDict
from monty.string import list_strings, marquee
from monty.functools import lazy_property
from abipy.core.structure import Structure
from abipy.tools.plotting import add_fig_kwargs, get_ax_fig_plt, get_axarray_fig_plt
from abipy.tools.typing import Figure
from abipy.core.mixins import AbinitNcFile, Has_Structure, NotebookWriter
from abipy.core.kpoints import Kpath
from abipy.abio.robots import Robot
from abipy.iotools import ETSF_Reader
def _get_style(reim, what, marker=None, markersize=None, alpha=1.0) -> str:
lw = 1
symbol, linestyle, linewidth = {
"v1scf_avg": (r"v1_{\bf q}", "-", lw),
"v1scf_abs_avg": (r"|v1_{\bf q}|", "-", lw),
"v1lr_avg": (r"v1_{\mathrm{lr}}", "--", lw),
"v1lr_abs_avg": (r"|v1_{\mathrm{lr}}|", "--", lw),
"v1scfmlr_avg": (r"(v1_{\bf q} - v1_{\mathrm{lr}})", "-.", lw),
"v1scfmlr_abs_avg": (r"(|v1_{\bf q} - v1_{\mathrm{lr}|})", "-.", lw),
"v1scf_gsmall": (r"v1_{\bf q}(G)", "-", lw),
"v1lr_gsmall": (r"v1lr_{\bf q}(G)", "-.", lw),
"v1scfmlr_gsmall": (r"(v1_{\bf q}(G) - v1_{\mathrm{lr}}(G))", "-.", lw),
}[what]
return dict(
marker=marker,
markersize=markersize,
label=r"$\langle \%s %s \rangle$" % ({0: "Re", 1: "Im"}[reim], symbol),
color={0: "blue", 1: "red"}[reim],
linestyle=linestyle,
linewidth=linewidth,
alpha=alpha,
)
[docs]
class V1qAvgFile(AbinitNcFile, Has_Structure, NotebookWriter):
"""
The V1QAVG.nc file contains the average over the unit cell of the periodic part the DFPT scattering potential.
This file is produced by the E-PH code by setting eph_task to +15 or -15.
If eph_task is +15, the input DVDB contains a q-mesh and the potentials are interpolated on a list of q-points
(usually a q-path) specified by the user. In this case the V1QAVG.nc file also contains an extra array
with Max_r |W(R, r)|, useful to study the decay of the scattering potentials in R-space.
If eph_task is -15, the netcdf file contains the average for the q-points found in the DVDB file.
This option is usually used to visualize the ab-initio potentials and compare then with the model for the LR part.
"""
def __init__(self, filepath: str):
super().__init__(filepath)
self.reader = r = ETSF_Reader(filepath)
# Read medadata
self.has_zeff = bool(r.read_value("has_zeff"))
self.has_dielt = bool(r.read_value("has_dielt"))
self.has_quadrupoles = bool(r.read_value("has_quadrupoles"))
self.has_efield = bool(r.read_value("has_efield", default=False))
self.dvdb_add_lr = r.read_value("dvdb_add_lr")
self.symv1scf = r.read_value("symv1scf")
self.interpolated = bool(r.read_value("interpolated"))
self.qdamp = r.read_value("qdamp")
[docs]
@lazy_property
def structure(self) -> Structure:
"""|Structure| object."""
return self.reader.read_structure()
[docs]
@lazy_property
def qpoints(self) -> Kpath:
"""List of q-points."""
frac_coords = self.reader.read_value('qpoints')
return Kpath(self.structure.reciprocal_lattice, frac_coords, ksampling=None)
[docs]
@lazy_property
def has_maxw(self) -> bool:
"""True if ncfile contains Max_r |W(R, r)|"""
return "maxw" in self.reader.rootgrp.variables
[docs]
def close(self) -> None:
self.reader.close()
[docs]
@lazy_property
def params(self) -> dict:
"""Dict with parameters that might be subject to convergence studies."""
return {}
def __str__(self) -> str:
return self.to_string()
[docs]
def to_string(self, verbose=0) -> str:
"""String representation."""
lines = []; app = lines.append
app(marquee("File Info", mark="="))
app(self.filestat(as_string=True))
app("")
app(self.structure.to_string(verbose=verbose, title="Structure"))
app("")
app(self.qpoints.to_string(verbose=verbose, title="Q-path"))
app("")
app("has_dielt: %s, has_zeff: %s, has_quadrupoles: %s, has_efield: %s" % (
self.has_dielt, self.has_zeff, self.has_quadrupoles, self.has_efield))
app("dvdb_add_lr: %s, symv1scf: %s, interpolated: %s, qdamp: %s" % (
self.dvdb_add_lr, self.symv1scf, self.interpolated, self.qdamp))
return "\n".join(lines)
[docs]
def make_ticks_and_labels(self) -> tuple:
"""
Find the k-point names in the pymatgen database.
"""
od = OrderedDict()
# If the first or the last k-point are not recognized in findname_in_hsym_stars
# matplotlib won't show the full band structure along the k-path
# because the labels are not defined. So we have to make sure that
# the labels for the extrema of the path are always defined.
od[0] = " "
for idx, qpoint in enumerate(self.qpoints):
name = qpoint.name if qpoint.name is not None else self.structure.findname_in_hsym_stars(qpoint)
if name:
od[idx] = name
if qpoint.name is None: qpoint.set_name(name)
last = len(self.qpoints) - 1
if last not in od: od[last] = " "
return list(od.keys()), list(od.values())
#def xsf_write(self, idir=0, ipert=0, ispden=0):
# natom = len(self.structure)
# ip = idir + 3 * ipert
# r = self.reader
# ngfft = r.read_value("ngfft")
# from abipy.tools.numtools import transpose_last3dims
# from abipy.iotools import xsf
# def write(datar, filename):
# with open(filename, mode="wt") as fh:
# xsf.xsf_write_structure(fh, self.structure)
# xsf.xsf_write_data(fh, self.structure, datar, add_replicas=True, cplx_mode="abs")
# # nctkarr_t("v1r_interpolated", "dp", "two, nfft, nspden, natom3")
# var = r.read_variable("v1r_interpolated")
# datar = var[ip, ispden, :, 0] + 1j * var[ip, ispden, :, 1]
# # Because we are reading Fortran data (z,y,x) and xsf_write_data expects (..., x, y, z)
# #print(datar.shape)
# datar = np.reshape(datar, np.flip(ngfft))
# datar = transpose_last3dims(datar)
# write(datar, "v1r_interp_idir%d_ipert%d.xsf" % (idir, ipert))
# var = r.read_variable("v1r_lrmodel")
# datar = var[ip, ispden, :, 0] + 1j * var[ip, ispden, :, 1]
# datar = np.reshape(datar, np.flip(ngfft))
# datar = transpose_last3dims(datar)
# write(datar, "v1r_lr_idir%d_ipert%d.xsf" % (idir, ipert))
[docs]
@add_fig_kwargs
def plot(self, what_list="all", ispden=0, fontsize=6, sharey=False, **kwargs) -> Figure:
"""
Plot
Args:
what_list:
ispden: Spin density component to plot.
fontsize: fontsize for legends and titles
sharey: True to share y-axes.
Return: |matplotlib-Figure|
"""
#all_varnames = ["v1scf_avg", "v1lr_avg", "v1scfmlr_avg", "v1scfmlr_abs_avg", "v1scf_abs_avg"]
what_list = list_strings(what_list) if what_list != "all" else ["v1scf_avg", "v1lr_avg"]
data = {}
for vname in what_list:
# Fortran array: nctkarr_t("v1scf_avg", "dp", "two, nspden, three, natom, nqpt")
data[vname] = self.reader.read_value(vname)
# Build [natom, 3] grid plot.
natom = len(self.structure)
nrows, ncols = natom, 3
ax_mat, fig, plt = get_axarray_fig_plt(None, nrows=nrows, ncols=ncols,
sharex=True, sharey=sharey, squeeze=False)
xs = np.arange(len(self.qpoints))
ticks, labels = self.make_ticks_and_labels()
for ip, ax in enumerate(ax_mat.ravel()):
idir = ip % 3
iat = (ip - idir) // 3
for reim in (0, 1):
for vname in what_list:
ys = data[vname][:, iat, idir, ispden, reim]
ax.plot(xs, ys, **_get_style(reim, vname))
ax.grid(True)
if iat == natom - 1: ax.set_xlabel("Q Wave Vector")
if idir == 0: ax.set_ylabel(r"(Hartree/Bohr)")
if ticks:
ax.set_xticks(ticks, minor=False)
ax.set_xticklabels(labels, fontdict=None, minor=False, size=kwargs.get("qlabel_size", "large"))
ax.set_xlim(ticks[0], ticks[-1])
if ip == 0:
ax.legend(loc="best", fontsize=fontsize, shadow=True)
site = self.structure[iat]
s = "%s [%.3f, %.3f, %.3f]" % (site.specie.symbol, site.frac_coords[0], site.frac_coords[1], site.frac_coords[2])
ax.set_title("idir: %d, iat: %d, %s" % (idir, iat, s), fontsize=fontsize)
return fig
[docs]
@add_fig_kwargs
def plot_gvec(self, gvec, ispden=0, fontsize=6, sharey=False, **kwargs) -> Figure:
"""
Plot
Args:
gvec: G vector in reduced coordinates.
ispden: Spin density component to plot.
fontsize: fontsize for legends and titles
sharey: True to share y-axes.
Return: |matplotlib-Figure|
"""
gsmall = self.reader.read_value("gsmall")
for ig, g in enumerate(gsmall):
if np.all(g == gvec): break
else:
raise RuntimeError("Cannot find gvec %s in gsmall array" % str(gvec))
# nctkarr_t("v1scf_gsmall", "dp", "two, ngsmall, nspden, three, natom, nqpt"), &
# nctkarr_t("v1lr_gsmall", "dp", "two, ngsmall, nspden, three, natom, nqpt"), &
what_list = ["v1scf_gsmall", "v1lr_gsmall"]
data = {}
for vname in what_list:
ncvar = self.reader.read_variable(vname)
data[vname] = ncvar[..., ig, :]
# Build [natom, 3] grid plot.
natom = len(self.structure)
nrows, ncols = natom, 3
ax_mat, fig, plt = get_axarray_fig_plt(None, nrows=nrows, ncols=ncols,
sharex=True, sharey=sharey, squeeze=False)
xs = np.arange(len(self.qpoints))
ticks, labels = self.make_ticks_and_labels()
for ip, ax in enumerate(ax_mat.ravel()):
idir = ip % 3
iat = (ip - idir) // 3
for reim in (0, 1):
for vname in what_list:
ys = data[vname][:, iat, idir, ispden, reim]
ax.plot(xs, ys, **_get_style(reim, vname))
# plot difference.
ys = (data["v1scf_gsmall"][:, iat, idir, ispden, reim]
- data["v1lr_gsmall"][:, iat, idir, ispden, reim]) # * 10
ax.plot(xs, ys, **_get_style(reim, "v1scfmlr_gsmall"))
ax.grid(True)
if iat == natom - 1: ax.set_xlabel("Q Wave Vector")
if idir == 0: ax.set_ylabel(r"(Hartree/Bohr)")
if ticks:
ax.set_xticks(ticks, minor=False)
ax.set_xticklabels(labels, fontdict=None, minor=False, size=kwargs.get("qlabel_size", "large"))
ax.set_xlim(ticks[0], ticks[-1])
if ip == 0:
ax.legend(loc="best", fontsize=fontsize, shadow=True)
site = self.structure[iat]
s = "%s [%.3f, %.3f, %.3f]" % (site.specie.symbol, site.frac_coords[0], site.frac_coords[1], site.frac_coords[2])
ax.set_title("idir: %d, iat: %d, %s" % (idir, iat, s), fontsize=fontsize)
fig.suptitle("G = %s" % str(gvec))
return fig
[docs]
@add_fig_kwargs
def plot_maxw(self, scale="semilogy", ax=None, fontsize=8, **kwargs) -> Figure:
"""
Plot the decay of max_{r,idir,ipert} |W(R,r,idir,ipert)|
Args:
scale: "semilogy", "loglog" or "plot".
ax: |matplotlib-Axes| or None if a new figure should be created.
fontsize: fontsize for legends and titles
Return: |matplotlib-Figure|
"""
if not self.has_maxw: return None
ax, fig, plt = get_ax_fig_plt(ax=ax)
# Fortran array: nctkarr_t("maxw", "dp", "nrpt, natom3")
rmod = self.reader.read_value("rmod")
maxw = self.reader.read_value("maxw")
data = np.max(maxw, axis=0)
f = {"plot": ax.plot, "semilogy": ax.semilogy, "loglog": ax.loglog}[scale]
f(rmod, data, marker="o", ls=":", lw=0, **kwargs)
ax.grid(True)
ax.set_ylabel(r"$Max_{({\bf{r}}, idir, ipert)} \| W({\bf{r}}, {\bf{R}}, idir, ipert) \|$")
ax.set_xlabel(r"$\|{\bf{R}}\|$ (Bohr)")
#if kwargs.pop("with_title", True):
# ax.set_title("dvdb_add_lr %d, qdamp: %s, symv1scf: %d" % (self.dvdb_add_lr, self.qdamp, self.symv1scf),
# fontsize=fontsize)
return fig
[docs]
@add_fig_kwargs
def plot_maxw_perts(self, scale="semilogy", sharey=False, fontsize=8, **kwargs) -> Figure:
"""
Plot the decay of max_r |W(R,r,idir,ipert)| for the individual atomic perturbations.
Args:
scale: "semilogy", "loglog" or "plot".
sharey: True is y-axes should be shared.
fontsize: fontsize for legends and titles.
Return: |matplotlib-Figure|
"""
if not self.has_maxw: return None
# Build grid of plots.
natom = len(self.structure)
ncols, nrows = (2, natom // 2) if natom % 2 == 0 else (1, natom)
ax_list, fig, plt = get_axarray_fig_plt(None, nrows=nrows, ncols=ncols,
sharex=True, sharey=sharey, squeeze=False)
ax_list = ax_list.ravel()
# Fortran array: nctkarr_t("maxw", "dp", "nrpt, natom3")
nrpt = self.reader.read_dimvalue("nrpt")
rmod = self.reader.read_value("rmod")
maxw = np.reshape(self.reader.read_value("maxw"), (natom, 3, nrpt))
for iatom, ax in enumerate(ax_list.ravel()):
site = self.structure[iatom]
title = "{} [{:.4f} {:.4f} {:.4f}]".format(site.specie.symbol, *site.frac_coords)
ax.set_title(title, fontsize=fontsize)
f = {"plot": ax.plot, "semilogy": ax.semilogy, "loglog": ax.loglog}[scale]
f(rmod, maxw[iatom, 0], marker="o", ls=":", lw=0, label="$L_x$" if iatom == 0 else None)
f(rmod, maxw[iatom, 1], marker="o", ls=":", lw=0, label="$L_y$" if iatom == 0 else None)
f(rmod, maxw[iatom, 2], marker="o", ls=":", lw=0, label="$L_z$" if iatom == 0 else None)
ax.grid(True)
if iatom == 0:
ax.set_ylabel(r"$Max_{{\bf{r}}} \| W({\bf{r}}, {\bf{R}}) \|$")
ax.legend(loc="best", fontsize=fontsize, shadow=True)
if iatom == len(ax_list) - 1: ax.set_xlabel(r"$\|{\bf{R}}\|$ (Bohr)")
#fig.suptitle("dvdb_add_lr %d, qdamp: %s, symv1scf: %d" % (self.dvdb_add_lr, self.qdamp, self.symv1scf),
# fontsize=fontsize)
return fig
[docs]
def yield_figs(self, **kwargs): # pragma: no cover
"""
This function generates a predefined list of matplotlib figures with minimal input from the user.
"""
title = r"$\langle v1_{\bf q} \rangle \,vs\, \langle v1_{\bfq}^{\mathrm{LR}} \rangle$"
yield self.plot(title=title, show=False)
yield self.plot(what_list="v1scfmlr_avg", title=r"$v1_{\bf q} - v1_{\bf q}^{\mathrm{LR}}$", show=False)
#yield self.plot(what_list=["v1scf_abs_avg", "v1lr_abs_avg"], title=r"ABS", show=False)
#yield self.plot(what_list="v1scfmlr_abs_avg", title=r"$|v1_{\bf q} - v1_{\bf q}^{\mathrm{LR}|}$", show=False)
for gvec in [[0, 0, 0], [1, 0, 0], [0, 1, 1], [1, 1, 1], [2, 2, 2]]:
yield self.plot_gvec(gvec, show=False)
if self.has_maxw:
if kwargs.get("verbose", 0) > 0:
yield self.plot_maxw(show=False)
yield self.plot_maxw_perts(show=False)
else:
print("Use verbose to print the decay of W(r,R)")
[docs]
def write_notebook(self, nbpath=None) -> str:
"""
Write a jupyter notebook to ``nbpath``. If nbpath is None, a temporary file in the current
working directory is created. Return path to the notebook.
"""
nbformat, nbv, nb = self.get_nbformat_nbv_nb(title=None)
nb.cells.extend([
nbv.new_code_cell("ncfile = abilab.abiopen('%s')" % self.filepath),
nbv.new_code_cell("print(ncfile)"),
nbv.new_code_cell("ncfile.plot();"),
])
if self.has_maxw:
nbv.new_code_cell("ncfile.plot_maxw();"),
nbv.new_code_cell("ncfile.plot_maxw_perts();"),
return self._write_nb_nbpath(nb, nbpath)
[docs]
class V1qAvgRobot(Robot):
"""
This robot analyzes the results contained in multiple V1qAvgFile.nc files.
.. rubric:: Inheritance Diagram
.. inheritance-diagram:: V1qAvgRobot
"""
EXT = "V1QAVG"
# Absolute tolerance used to compare q-points.
atol = 1e-2
[docs]
@lazy_property
def qpoints(self):
"""List of q-points."""
if len(self) == 1: return self.abifiles[0].qpoints
if (any(len(ncfile.qpoints) != len(self.abifiles[0].qpoints) for ncfile in self.abifiles)):
raise RuntimeError("Assuming ncfiles with same number of q-points.\nFound %s" % (
str([len(ncfile.qpoints) for ncfile in self.abifiles])))
for abifile in self.abifiles[1:]:
if np.any(np.abs(abifile.qpoints.frac_coords - self.abifiles[0].qpoints.frac_coords) > self.atol):
for q1, q2 in zip(self.abifiles[0].qpoints, abifile.qpoints):
print("q1:", q1, ", q2:", q2)
raise RuntimeError("Found different q-points with tolerance: %s!" % self.atol)
return self.abifiles[0].qpoints
[docs]
@add_fig_kwargs
def plot(self, ispden=0, vname="v1scf_avg", sharey=False, fontsize=8, **kwargs) -> Figure:
"""
Plot
Args:
ispden: Spin density component to plot.
sharey: True to share y-axes.
fontsize: fontsize for legends and titles
Return: |matplotlib-Figure|
"""
# Caveat: No check is done on the consistency among structures.
ref_file = self.abifiles[0]
structure = ref_file.structure
natom = len(structure)
nrows, ncols = natom, 3
ax_mat, fig, plt = get_axarray_fig_plt(None, nrows=nrows, ncols=ncols,
sharex=True, sharey=sharey, squeeze=False)
xs = np.arange(len(self.qpoints))
ticks, labels = ref_file.make_ticks_and_labels()
data_file = {abilabel: abifile.reader.read_value(vname) for abilabel, abifile in self.items()}
style = dict(marker=".", markersize=2)
for ip, ax in enumerate(ax_mat.ravel()):
idir = ip % 3
iat = (ip - idir) // 3
for ifile, (abilabel, abifile) in enumerate(self.items()):
v1scf_avg = data_file[abilabel]
for reim in (0, 1):
style = _get_style(reim, vname)
label = "%s %s" % (style.pop("label"), abilabel)
style["linestyle"] = {0: "-", 1: "--", 2: "-.", 3: ":"}[ifile]
ax.plot(xs, v1scf_avg[:, iat, idir, ispden, reim], label=label if ip == 0 else None, **style)
if ip == 0:
ax.legend(loc="best", fontsize=fontsize, shadow=True)
ax.grid(True)
if iat == natom - 1: ax.set_xlabel("Q Wave Vector")
if idir == 0: ax.set_ylabel(r"(Hartree/Bohr)")
site = ref_file.structure[iat]
s = "%s [%.3f, %.3f, %.3f]" % (site.specie.symbol, site.frac_coords[0], site.frac_coords[1], site.frac_coords[2])
ax.set_title("idir: %d, iat: %d, %s" % (idir, iat, s), fontsize=fontsize)
if ticks:
ax.set_xticks(ticks, minor=False)
ax.set_xticklabels(labels, fontdict=None, minor=False, size=kwargs.get("qlabel_size", "large"))
ax.set_xlim(ticks[0], ticks[-1])
return fig
[docs]
@add_fig_kwargs
def plot_maxw(self, ax=None, **kwargs) -> Figure:
"""
Plot
Args:
ax: |matplotlib-Axes| or None if a new figure should be created.
fontsize: fontsize for legends and titles
Return: |matplotlib-Figure|
"""
if any(not abifile.has_maxw for abifile in self.abifiles): return None
ax, fig, plt = get_ax_fig_plt(ax=ax)
for label, abifile in self.items():
abifile.plot_maxw(ax=ax, label=label, with_title=False, show=False, **kwargs)
return fig
[docs]
def yield_figs(self, **kwargs): # pragma: no cover
"""
This function *generates* a predefined list of matplotlib figures with minimal input from the user.
"""
for vname in ("v1scf_avg", "v1scf_abs_avg"):
yield self.plot(vname=vname, title=vname, show=False)
if all(abifile.has_maxw for abifile in self.abifiles):
yield self.plot_maxw(self, show=False, **kwargs)
[docs]
def write_notebook(self, nbpath=None) -> str:
"""
Write a jupyter notebook to `nbpath`. If nbpath is None, a temporary file in the current
working directory is created. Return path to the notebook.
"""
nbformat, nbv, nb = self.get_nbformat_nbv_nb(title=None)
args = [(l, f.filepath) for l, f in self.items()]
nb.cells.extend([
#nbv.new_markdown_cell("# This is a markdown cell"),
nbv.new_code_cell("robot = abilab.V1qAvgRobot(*%s)\nrobot.trim_paths()\nrobot" % str(args)),
nbv.new_code_cell("robot.plot();"),
])
return self._write_nb_nbpath(nb, nbpath)
#class V1qAvgPlotter(object):
#
# def __init__(self, v1q_dfpt_path, v1q_frohl_path, v1q_frohl_quad_path, v1q_froh_quad_efield_path):
# self.v1q_dfpt = V1qAvgFile.from_file(v1q_dfpt_path)
# #v1q_frohl = V1qAvgFile.from_file(v1q_frohl_path) if v1q_frohl_path is not None else None
# #v1q_frohl_quad_path =
# #v1q_frohl_quad_efield_path =
#
# #def close(self):
# # for k in ("v1q_frohl", "v1q_frohl_quad_path", "v1q_frohl_quad_efield_path"):
# # ncfile = getattr(self, k)
# # if ncfile is not None: ncfile.close()
#
# @add_fig_kwargs
# def plot(self, iatom, idir, ax=None, ispden=0, fontsize=6, **kwargs):
#
# ax, fig, plt = get_ax_fig_plt(ax=ax)
#
# return fig