"""
This module contains objects for analyzing
the GWAN.nc file with the e-ph vertex in the Wannier representation.
For a theoretical introduction see :cite:`Giustino2017`
"""
from __future__ import annotations
import dataclasses
import numpy as np
import pandas as pd
#import abipy.core.abinit_units as abu
from monty.string import marquee #, list_strings
from monty.functools import lazy_property
from monty.termcolor import cprint
from abipy.core.structure import Structure
from abipy.core.kpoints import kpoints_indices
from abipy.core.mixins import AbinitNcFile, Has_Structure, Has_ElectronBands, Has_Header #, NotebookWriter
from abipy.tools.typing import PathLike
from abipy.tools.numtools import BzRegularGridInterpolator, nparr_to_df
#from abipy.tools.plotting import (add_fig_kwargs, get_ax_fig_plt, get_axarray_fig_plt, set_axlims, set_visible,
# rotate_ticklabels, ax_append_title, set_ax_xylabels, linestyles)
#from abipy.tools import duck
from abipy.electrons.ebands import ElectronBands, RobotWithEbands
#from abipy.tools.typing import Figure
from abipy.abio.robots import Robot
from abipy.eph.common import BaseEphReader
[docs]
class GwanFile(AbinitNcFile, Has_Header, Has_Structure, Has_ElectronBands): # , NotebookWriter):
"""
This file stores the e-ph matrix elements in the wannier representation
and provides methods to analyze and plot results.
Usage example:
.. code-block:: python
with GwanFile("out_GWAN.nc") as gwan:
print(gwan)
for spin in range(gwan.nsppol):
# Extract the object storing the g for this spin.
#gqk = gstore.gqk_spin[spin]
#print(gqk)
# Get a Dataframe with g(k, q) for all modes and bands.
#df = gqk.get_gdf_at_qpt_kpt([1/2, 0, 0], [0, 0, 0])
#print(df)
.. rubric:: Inheritance Diagram
.. inheritance-diagram:: GwanFile
"""
[docs]
@classmethod
def from_file(cls, filepath: PathLike) -> GwanFile:
"""Initialize the object from a netcdf file."""
return cls(filepath)
def __init__(self, filepath: PathLike):
super().__init__(filepath)
self.r = GwanReader(filepath)
[docs]
@lazy_property
def ebands(self) -> ElectronBands:
"""|ElectronBands| object."""
return self.r.read_ebands()
@property
def structure(self) -> Structure:
"""|Structure| object."""
return self.ebands.structure
[docs]
def close(self) -> None:
"""Close the file."""
self.r.close()
#@lazy_property
#def gqk_spin(self) -> list:
# return [Gqk.from_gstore(self, spin) for spin in range(self.nsppol)]
[docs]
@lazy_property
def params(self) -> dict:
"""dict with the convergence parameters, e.g. ``nbsum``."""
#od = OrderedDict([
# ("nbsum", self.nbsum),
# ("nqibz", self.r.nqibz),
#])
## Add EPH parameters.
#od.update(self.r.common_eph_params)
od = {}
return od
def __str__(self) -> str:
return self.to_string()
[docs]
def to_string(self, verbose=0) -> str:
"""String representation with verbosiy level ``verbose``."""
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.ebands.to_string(with_structure=False, verbose=verbose, title="Electronic Bands"))
if verbose > 1:
app("")
app(self.hdr.to_string(verbose=verbose, title="Abinit Header"))
app(f"nsppol: {self.r.nsppol}")
app(f"gstore_completed: {bool(self.r.completed)}")
app(f"gstore_cplex: {self.r.cplex}")
app(f"gstore_kptopt: {self.r.kptopt}")
app(f"gstore_qptopt: {self.r.qptopt}")
#for spin in range(self.r.nsppol):
# app(f"gstore_brange_spin[{spin}]: {self.r.brange_spin[spin]}")
# app(f"gstore_erange_spin[{spin}]: {self.r.erange_spin[spin]}")
# app(f"gstore_glob_spin_nq[{spin}]: {self.r.glob_spin_nq[spin]}")
return "\n".join(lines)
#def write_epw_hdf5(self, filepath: PathLike) -> None:
[docs]
@dataclasses.dataclass(kw_only=True)
class Gqk:
"""
This object stores the e-ph matrix elements (g or g^2) and the matrix elements
of the velocity operator for a given spin.
"""
cplex: int # 1 if |g|^2 is stored
# 2 if complex valued g (mind the gauge)
spin: int # Spin index.
nb: int # Number of bands
bstart: int
#bstop: int
glob_nk: int # Total number of k/q points in global matrix.
glob_nq: int # Note that k-points/q-points can be filtered.
# Use kzone, qzone and kfilter to interpret these dimensions.
gstore: GstoreFile
gvals: np.ndarray | None
g2: np.ndarray | None
vk_cart_ibz: np.ndarray | None
vkmat_cart_ibz: np.ndarray | None
[docs]
@classmethod
def from_gstore(cls, gstore: GstoreFile, spin: int):
"""
Build an istance from a GstoreFile and the spin index.
"""
ncr = gstore.r
path = f"gqk_spin{spin+1}"
cplex = ncr.read_dimvalue("gstore_cplex")
nb = ncr.read_dimvalue("nb", path=path)
glob_nk = ncr.read_dimvalue("glob_nk", path=path)
glob_nq = ncr.read_dimvalue("glob_nq", path=path)
# Read e-ph matrix elements
# nctkarr_t("gvals", "dp", "gstore_cplex, nb_kq, nb_k, natom3, glob_nk, glob_nq)
# Have to transpose the (nb_kq, nb_k) submatrix written by Fortran.
g2, gvals = None, None
if cplex == 1:
g2 = ncr.read_value("gvals", path=path).transpose(0, 1, 2, 4, 3, 5).copy()
elif cplex == 2:
gvals = ncr.read_value("gvals", path=path).transpose(0, 1, 2, 4, 3, 5).copy()
gvals = gvals[...,0] + 1j*gvals[...,1]
vk_cart_ibz, vkmat_cart_ibz = None, None
if ncr.with_vk == 1:
# nctk_def_arrays(spin_ncid, nctkarr_t("vk_cart_ibz", "dp", "three, nb, gstore_nkibz"))
vk_cart_ibz = ncr.read_value("vk_cart_ibz", path=path)
if ncr.with_vk == 2:
# Full (nb x nb) matrix.
# Have to transpose (nb_kq, nb_k) submatrix written by Fortran.
# nctk_def_arrays(spin_ncid, nctkarr_t("vkmat_cart_ibz", "dp", "two, three, nb, nb, gstore_nkibz"))
vkmat_cart_ibz = ncr.read_value("vkmat_cart_ibz", path=path).transpose(0, 1, 3, 2, 4).copy()
vkmat_cart_ibz = vkmat_cart_ibz[...,0] + 1j*vkmat_cart_ibz[...,1]
# Note conversion between Fortran and python indexing.
bstart = ncr.read_value("bstart", path=path) - 1
#bstop = ncr.read_value("stop", path=path)
data = locals()
return cls(**{k: data[k] for k in [field.name for field in dataclasses.fields(Gqk)]})
def __str__(self) -> str:
return self.to_string()
[docs]
def to_string(self, verbose=0) -> str:
"""String representation with verbosiy level ``verbose``."""
lines = []; app = lines.append
app(marquee(f"Gqk for spin: {self.spin}", mark="="))
app(f"cplex: {self.cplex}")
app(f"nb: {self.nb}")
app(f"bstart: {self.bstart}")
app(f"glob_nk: {self.glob_nk}")
app(f"glob_nq: {self.glob_nq}")
return "\n".join(lines)
@property
def structure(self):
return self.gstore.structure
[docs]
def get_dataframe(self, what: str = "g2") -> pd.DataFrame:
"""
Build and return a dataframe with all the |g(k,q)|^2 if what == "g2" or
all |v_nk|^2 if what == "v2".
"""
if what == "g2":
g2 = self.g2 if self.g2 is not None else np.abs(self.gvals) ** 2
df = nparr_to_df("g2", g2, ["iq", "ik", "imode", "m_kq", "n_k"])
elif what == "v2":
if self.vk_cart_ibz is None:
raise ValueError("vk_cart_ibz is not available in GSTORE!")
# Compute the squared norm of each vector
v2 = np.sum(self.vk_cart_ibz ** 2, axis=2)
df = nparr_to_df("v2", v2, ["ik", "n_k"])
else:
raise ValueError(f"Invalid {what=}")
#df["m_kq"] += bstart_mkq
#df["n_k"] += bstart_nk
return df
[docs]
def get_g2q_interpolator_kpoint(self, kpoint, method="linear", check_mesh=1):
"""
"""
r = self.gstore.r
# Find the index of the kpoint.
ik_g, kpoint = r.find_ik_glob_kpoint(kpoint, self.spin)
# Compute indices of qpoints in the ngqpt mesh.
ngqpt, shifts = r.ngqpt, [0, 0, 0]
q_indices = kpoints_indices(r.qbz, ngqpt, check_mesh=check_mesh)
natom3 = 3 * len(self.structure)
nb = self.nb
nx, ny, nz = ngqpt
# (glob_nq, glob_nk, natom3, m_kq, n_k)
g2 = self.g2 if self.g2 is not None else np.abs(self.gvals) ** 2
g2_qph_mn = g2[:,ik_g]
# Insert g2 in g2_grid
g2_grid = np.empty((nb, nb, natom3, nx, ny, nz))
for nu in range(natom3):
for g2_mn, q_inds in zip(g2_qph_mn[:,nu], q_indices):
ix, iy, iz = q_inds
g2_grid[:, :, nu, ix, iy, iz] = g2_mn
return BzRegularGridInterpolator(self.structure, shifts, g2_grid, method=method)
[docs]
def get_g_qpt_kpt(self, qpoint, kpoint, what) -> np.ndarray:
"""
Return numpy array with e-ph matrix elements the for the given (qpoint, kpoint) pair.
Args:
what="g2" for |g(k,q)|^2, "g" for g(k,q)
"""
# Find the internal indices of (qpoint, kpoint)
iq_g, qpoint = self.gstore.r.find_iq_glob_qpoint(qpoint, self.spin)
ik_g, kpoint = self.gstore.r.find_ik_glob_kpoint(kpoint, self.spin)
if what == "g2":
g2 = self.g2 if self.g2 is not None else np.abs(self.gvals) ** 2
return g2[iq_g, ik_g]
if what == "g":
if self.cplex != 2:
raise ValueError("Gstore file stores g2 instead of complex g")
return self.gvals[iq_g, ik_g]
raise ValueError(f"Invalid {what=}")
[docs]
def get_gdf_at_qpt_kpt(self, qpoint, kpoint, what="g2") -> pd.DataFrame:
"""
Build and return a dataframe with the |g(k,q)|^2 for the given (qpoint, kpoint) pair.
Args:
what="g2" for |g(k,q)|^2, "g" for g(k,q)
"""
g2_slice = self.get_g_qpt_kpt(qpoint, kpoint, what)
df = nparr_to_df(what, g2_slice, ["imode", "m_kq", "n_k"])
#df["m_kq"] += bstart_mkq
#df["n_k"] += bstart_nk
return df
[docs]
def neq(self, other: Gqk, verbose: int) -> int:
"""
Helper function to compare two GQK objects.
"""
# This dimensions must agree in order to have a meaningfull comparison.
# so raise immediately if not equal.
aname_list = ["cplex", "spin", "nb", "glob_nk", "glob_nq"]
for aname in aname_list:
val1, val2 = getattr(self, aname), getattr(other, aname)
if isinstance(val1, (str, int, float)):
eq = val1 == val2
elif isinstance(val1, np.ndarray):
eq = np.allclose(val1, val2)
else:
raise TypeError(f"Don't know how to handle comparison for type: {type(val1)}")
if not eq:
raise RuntimeError(f"Different values of {aname=}, {val1=}, {val2=}")
ierr = 0
kws = dict(verbose=verbose) # , atol= rtol)
# Compare v_nk or v_mn_k.
if self.vk_cart_ibz is not None:
if not _allclose("vk_cart_ibz", self.vk_cart_ibz, other.vk_cart_ibz, **kws): ierr += 1
if self.vkmat_cart_ibz is not None:
if not _allclose("vkmat_cart_ibz", self.vkmat_cart_ibz, other.vkmat_cart_ibz, **kws): ierr += 1
# Compare g or g^2.
if self.g2 is not None:
if not _allclose("g2", self.g2, other.g2, **kws): ierr += 1
if self.gvals is not None:
if not _allclose("gvals", self.gvals, other.gvals, **kws): ierr += 1
return ierr
[docs]
class GwanReader(BaseEphReader):
"""
Reads data from file and constructs objects.
.. rubric:: Inheritance Diagram
.. inheritance-diagram:: GstoreReader
"""
def __init__(self, filepath: PathLike):
super().__init__(filepath)
# Read important dimensions.
self.nsppol = self.read_dimvalue("number_of_spins")
self.cplex = self.read_dimvalue("gstore_cplex")
self.nkbz = self.read_dimvalue("gstore_nkbz")
self.nkibz = self.read_dimvalue("gstore_nkibz")
self.nqbz = self.read_dimvalue("gstore_nqbz")
self.nqibz = self.read_dimvalue("gstore_nqibz")
# Read important variables.
self.completed = self.read_value("gstore_completed")
self.done_spin_qbz = self.read_value("gstore_done_qbz_spin")
self.with_vk = self.read_value("gstore_with_vk")
self.qptopt = self.read_value("gstore_qptopt")
self.kptopt = self.read_value("kptopt")
self.kzone = self.read_string("gstore_kzone")
self.qzone = self.read_string("gstore_qzone")
self.kfilter = self.read_string("gstore_kfilter")
self.gmode = self.read_string("gstore_gmode")
# Note conversion Fortran --> C for the isym index.
self.brange_spin = self.read_value("gstore_brange_spin")
self.brange_spin[:,0] -= 1
self.erange_spin = self.read_value("gstore_erange_spin")
# Total number of k/q points for each spin after filtering (if any)
self.glob_spin_nq = self.read_value("gstore_glob_nq_spin")
self.glob_nk_spin = self.read_value("gstore_glob_nk_spin")
# K-points and q-points in the IBZ
self.kibz = self.read_value("reduced_coordinates_of_kpoints")
self.qibz = self.read_value("gstore_qibz")
# K-points and q-points in the BZ
self.kbz = self.read_value("gstore_kbz")
self.qbz = self.read_value("gstore_qbz")
self.ngqpt = self.read_value("gstore_ngqpt")
# Mapping BZ --> IBZ. Note conversion Fortran --> C for the isym index.
# nctkarr_t("gstore_kbz2ibz", "i", "six, gstore_nkbz"), &
# nctkarr_t("gstore_qbz2ibz", "i", "six, gstore_nqbz"), &
self.kbz2ibz = self.read_value("gstore_kbz2ibz")
self.kbz2ibz[:,0] -= 1
self.qbz2ibz = self.read_value("gstore_qbz2ibz")
self.qbz2ibz[:,0] -= 1
# Mapping q/k points in gqk --> BZ. Note conversion Fortran --> C for indexing.
# nctkarr_t("gstore_qglob2bz", "i", "gstore_max_nq, number_of_spins"), &
# nctkarr_t("gstore_kglob2bz", "i", "gstore_max_nk, number_of_spins") &
self.qglob2bz = self.read_value("gstore_qglob2bz")
self.qglob2bz -= 1
self.kglob2bz = self.read_value("gstore_kglob2bz")
self.kglob2bz -= 1
[docs]
def find_iq_glob_qpoint(self, qpoint, spin: int):
"""
Find the internal index of the qpoint needed to access the gvals array.
"""
qpoint = np.asarray(qpoint)
for iq_g, iq_bz in enumerate(self.qglob2bz[spin]):
if np.allclose(qpoint, self.qbz[iq_bz]):
#print(f"Found {qpoint = } with index {iq_g = }")
return iq_g, qpoint
raise ValueError(f"Cannot find {qpoint = } in GSTORE.nc")
[docs]
def find_ik_glob_kpoint(self, kpoint, spin: int):
"""Find the internal indices of the kpoint needed to access the gvals array."""
kpoint = np.asarray(kpoint)
for ik_g, ik_bz in enumerate(self.kglob2bz[spin]):
if np.allclose(kpoint, self.kbz[ik_bz]):
#print(f"Found {kpoint = } with index {ik_g = }")
return ik_g, kpoint
raise ValueError(f"Cannot find {kpoint = } in GSTORE.nc")
# TODO: This fix to read groups should be imported in pymatgen.
[docs]
@lazy_property
def path2group(self) -> dict:
return self.rootgrp.groups
[docs]
class GstoreRobot(Robot, RobotWithEbands):
"""
This robot analyzes the results contained in multiple GSTORE.nc files.
Usage example:
.. code-block:: python
robot = GstoreRobot.from_files([
"t04o_GSTORE.nc",
"t05o_GSTORE.nc",
])
robot.neq(verbose=1)
.. rubric:: Inheritance Diagram
.. inheritance-diagram:: GstoreRobot
"""
EXT = "GSTORE"
[docs]
def neq(self, ref_basename: str | None = None, verbose: int = 0) -> int:
"""
Compare all GSTORE.nc files stored in the GstoreRobot
"""
# Find reference gstore. By default the first file in the robot is used.
ref_gstore = self._get_ref_abifile_from_basename(ref_basename)
exc_list = []
ierr = 0
for other_gstore in self.abifiles:
if ref_gstore.filepath == other_gstore.filepath:
continue
print("Comparing: ", ref_gstore.basename, " with: ", other_gstore.basename)
try:
ierr += self._neq_two_gstores(ref_gstore, other_gstore, verbose)
cprint("EQUAL", color="green")
except Exception as exc:
exc_list.append(str(exc))
for exc in exc_list:
cprint(exc, color="red")
return ierr
@staticmethod
def _neq_two_gstores(gstore1: GstoreFile, gstore2: GstoreFile, verbose: int) -> int:
"""
Helper function to compare two GSTORE files.
"""
# These quantities must be the same to have a meaningfull comparison.
aname_list = ["structure", "nsppol", "cplex", "nkbz", "nkibz",
"nqbz", "nqibz", "completed", "kzone", "qzone", "kfilter", "gmode",
"brange_spin", "erange_spin", "glob_spin_nq", "glob_nk_spin",
]
for aname in aname_list:
self._compare_attr_name(aname, gstore1, gstore2)
# Now compare the gkq objects for each spin.
ierr = 0
for spin in range(gstore1.nsppol):
gqk1, gqk2 = gstore1.gqk_spin[spin], gstore2.gqk_spin[spin]
ierr += gqk1.neq(gqk2, verbose)
return ierr
[docs]
def yield_figs(self, **kwargs): # pragma: no cover
"""
This function *generates* a predefined list of matplotlib figures with minimal input from the user.
Used in abiview.py to get a quick look at the results.
"""
#for fig in self.get_ebands_plotter().yield_figs(): yield fig
[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.GstoreRobot(*%s)\nrobot.trim_paths()\nrobot" % str(args)),
#nbv.new_code_cell("ebands_plotter = robot.get_ebands_plotter()"),
])
# Mixins
#nb.cells.extend(self.get_baserobot_code_cells())
#nb.cells.extend(self.get_ebands_code_cells())
return self._write_nb_nbpath(nb, nbpath)