Add preliminary version of plot examples and tests for vqp and orgmag

abipy/abilab.py

@@ -78,7 +78,6 @@
 from abipy.eph.gpath import GpathFile
 from abipy.wannier90 import WoutFile, AbiwanFile, AbiwanRobot
 from abipy.electrons.lobster import CoxpFile, ICoxpFile, LobsterDoscarFile, LobsterInput, LobsterAnalyzer
-
 from abipy.dynamics.cpx import EvpFile
 #try:
 #    from abipy.ml.aseml import AseMdLog

abipy/electrons/orbmag.py

@@ -1,5 +1,5 @@
 # coding: utf-8
-"""Classes for the analysis of electronic fatbands and projected DOSes."""
+"""Post-processing tools to analyze orbital magnetism."""
 from __future__ import annotations
 
 import numpy as np
@@ -12,12 +12,10 @@
 from abipy.core.structure import Structure
 from abipy.tools.numtools import BzRegularGridInterpolator
 from abipy.electrons.ebands import ElectronBands, ElectronsReader
-from abipy.core.kpoints import kpoints_indices, kmesh_from_mpdivs, map_grid2ibz #, IrredZone
+from abipy.core.kpoints import kpoints_indices, kmesh_from_mpdivs, map_grid2ibz
 #from abipy.tools.numtools import gaussian
 from abipy.tools.typing import Figure
 from abipy.tools.plotting import set_axlims, get_ax_fig_plt, get_axarray_fig_plt, add_fig_kwargs, Marker
-#from abipy.tools.plotting import (set_axlims, add_fig_kwargs, get_ax_fig_plt
-#    rotate_ticklabels, set_visible, plot_unit_cell, set_ax_xylabels)
 
 
 def print_options_decorator(**kwargs):
@@ -39,10 +37,9 @@ def wrapper(*args, **kwargs_inner):
     return decorator
 
 
-
 class OrbmagAnalyzer:
     """
-    This object gather three ORBMAG.nc files, post-process the data and
+    This object gathers three ORBMAG.nc files, post-processes the data and
     provides tools to analyze/plot the results.
 
     Usage example:
@@ -51,7 +48,6 @@ class OrbmagAnalyzer:
 
         from abipy.electrons.orbmag import OrbmagAnalyzer
         orban = OrbmagAnalyzer(["gso_DS1_ORBMAG.nc", "gso_DS2_ORBMAG.nc", "gso_DS3_ORBMAG.nc"])
-
         print(orban)
         orban.report_eigvals(report_type="S")
         orban.plot_fatbands("bands_GSR.nc")
@@ -64,12 +60,14 @@ def __init__(self, filepaths: list):
         """
         if not isinstance(filepaths, (list, tuple)):
             raise TypeError(f"Expecting list or tuple with paths but got {type(filepaths)=}")
+
         if len(filepaths) != 3:
             raise ValueError(f"{len(filepaths)=} != 3")
 
-        # TODO: One should store the direction in the netcdf file
+        # @JOE TODO: One should store the direction in the netcdf file
         # so that we can check that the files are given in the right order.
         self.orb_files = [OrbmagFile(path) for path in filepaths]
+        self.verbose = 0
 
         # This piece of code is taken from merge_orbmag_mesh. The main difference
         # is that here ncroots[0] is replaced by the reader instance of the first OrbmagFile.
@@ -133,6 +131,18 @@ def __init__(self, filepaths: list):
     #    lines = []; app = lines.append
     #    return "\n".join(lines)
 
+    def __enter__(self):
+        return self
+
+    def __exit__(self, exc_type, exc_val, exc_tb) -> None:
+        """Activated at the end of the with statement. It automatically closes all the files."""
+        self.close()
+
+    def close(self):
+        """Close all the files."""
+        for orb in self.orb_files:
+            orb.close()
+
     @lazy_property
     def structure(self):
         """Structure object."""
@@ -142,11 +152,18 @@ def structure(self):
             raise RuntimeError("ORBMAG.nc files have different structures")
         return structure
 
-    #@print_options_decorator(precision=2, suppress=True)
+    @lazy_property
+    def has_timrev(self):
+        """True if time-reversal symmetry is used in the BZ sampling."""
+        has_timrev = self.orb_files[0].ebands.has_timrev
+        if any(orb_file.ebands.has_timrev != has_timrev for orb_file in self.orb_files[1:]):
+            raise RuntimeError("ORBMAG.nc files have different values of timrev")
+
+    @print_options_decorator(precision=2, suppress=True)
     def report_eigvals(self, report_type):
         """
         """
-        np.set_printoptions(precision=2)
+        #np.set_printoptions(precision=2)
 
         terms = ['CC   ','VV1  ','VV2  ','NL   ','LR   ','A0An ']
         # Shape: (orbmag_nterms, nsppol, nkpt, mband, ndir, ndir)
@@ -199,6 +216,7 @@ def ngkpt_and_shifts(self) -> tuple:
 
             if ngkpt is None:
                 raise ValueError("Non diagonal k-meshes are not supported!")
+
             if len(shifts) > 1:
                 raise ValueError("Multiple shifts are not supported!")
 
@@ -211,6 +229,21 @@ def ngkpt_and_shifts(self) -> tuple:
 
         return ngkpt, shifts
 
+    def get_value(self, what: str, spin: int, ikpt: int, band: int) -> float:
+        """
+        Postprocess orbmag_merge_sigij_mesh to compute the quantity associated
+        to `what` for the given (spin, ikpt, band).
+        """
+        if what == "isotropic":
+            vals = self.orbmag_merge_sigij_mesh[:, spin, ikpt, band].sum(axis=0)
+            eigenvalues = -1.0E6 * vals
+            return eigenvalues.sum() / 3.0
+            #span = eigenvalues.max() - eigenvalues.min()
+            #skew = 3.0 * (eigenvalues.sum() - eigenvalues.max() - eigenvalues.min() - isotropic) / span
+        #elif what == "foobar":
+
+        raise ValueError(f"Invalid {what=}")
+
     def insert_inbox(self, what: str, spin: int) -> tuple:
         """
         Return data, ngkpt, shifts where data is a (mband, nkx, nky, nkz)) array
@@ -233,35 +266,79 @@ def insert_inbox(self, what: str, spin: int) -> tuple:
         data = np.empty(shape, dtype=float)
 
         for iband in range(self.mband):
-            for ikpt, k_inds in zip(range(self.nkpt),  k_indices, strict=True):
+            for ikpt, k_inds in zip(range(self.nkpt), k_indices, strict=True):
                 ix, iy, iz = k_inds
-                vals = self.orbmag_merge_sigij_mesh[:, spin, ikpt, iband].sum(axis=0)
-
-                if what == "isotropic":
-                    eigenvalues = -1.0E6 * vals
-                    value = eigenvalues.sum() / 3.0
-                    #span = eigenvalues.max() - eigenvalues.min()
-                    #skew = 3.0 * (eigenvalues.sum() - eigenvalues.max() - eigenvalues.min() - isotropic) / span
-                else:
-                    raise ValueError(f"Invalid {what=}")
-
+                value = self.get_value(what, spin, ikpt, iband)
                 data[iband, ix, iy, iz] = value
 
         return data, ngkpt, shifts
 
+    def get_skw_interpolator(self, what: str, lpratio: int, filter_params: None):
+        """
+        """
+        orb = self.orb_files[0]
+        ebands = orb.ebands.kpoints
+        kpoints = orb.ebands.kpoints
+
+        # Get symmetries from abinit spacegroup (read from file).
+        if (abispg := self.structure.abi_spacegroup) is None:
+            abispg = self.structure.spgset_abi_spacegroup(has_timerev=self.has_timrev)
+        fm_symrel = [s for (s, afm) in zip(abispg.symrel, abispg.symafm, strict=True) if afm == 1]
+
+        from abipy.core.skw import SkwInterpolator
+        cell = (self.structure.lattice.matrix, self.structure.frac_coords, self.structure.atomic_numbers)
+
+        interp_spin = [None for _ in range(self.nsppol)]
+
+        for spin in range(self.nsppol):
+            values_kb = np.empty((self.nkpt, self.mband))
+            for ikpt in range(self.nkpt):
+                for iband in range(self.mband):
+                    values_kb[ikpt, iband] = self.get_value(what, spin, ikpt, band)
+
+            skw = SkwInterpolator(lpratio, kpoints.frac_coords, self.eigens[:,:,bstart:bstop], ebands.fermie, ebands.nelect,
+                                  cell, fm_symrel, self.has_timrev,
+                                  filter_params=filter_params, verbose=self.verbose)
+            interp_spin[spin] = skw
+            #skw.eval_sk(spin, kpt, der1=None, der2=None) -> np.ndarray:
+
+        return interp_spin
+
+    @lazy_property
+    def has_full_bz(self) -> bool:
+        """True if k-points cover the full BZ."""
+        ngkpt, shifts = self.ngkpt_and_shifts
+        return np.product(ngkpt) * len(shifts) == self.nkpt
+
     def get_bz_interpolator_spin(self, what: str, interp_method: str) -> list[BzRegularGridInterpolator]:
         """
-        Build and return an interpolator for
+        Build and return a list with nsppol interpolators.
 
         Args:
             what: Strings defining the quantity to insert in the box.
             interp_method: The method of interpolation. Supported are “linear”, “nearest”,
                 “slinear”, “cubic”, “quintic” and “pchip”.
         """
+        # This part is tricky as we have to handle different cases:
+        #
+        # 1) k-mesh defined in terms of ngkpt and one shift.
+        #    In this case we can interpolate easily although we have to take into account
+        #    if the k-points have been reduced by symmetry or not.
+        #    If we have the full BZ, we can use BzRegularGridInterpolator
+        #    else we have to resort to StarFunction interpolation
+        #
+        # 2) k-mesh defined in terms of ngkpt and multiple shifts.
+        #    If we have the IBZ, we have to resort to StarFunction interpolation
+        #    Handling the full BZ is not yet possible due to an internal limitation of BzRegularGridInterpolator
+
         interp_spin = [None for _ in range(self.nsppol)]
-        for spin in range(self.nsppol):
-            data, ngkpt, shifts = self.insert_inbox(what, spin)
-            interp_spin[spin] = BzRegularGridInterpolator(self.structure, shifts, data, method=interp_method)
+
+        if self.has_full_bz:
+            for spin in range(self.nsppol):
+                data, ngkpt, shifts = self.insert_inbox(what, spin)
+                interp_spin[spin] = BzRegularGridInterpolator(self.structure, shifts, data, method=interp_method)
+        else:
+            raise NotImplementedError("k-points must cover the full BZ.")
 
         return interp_spin
 
@@ -272,7 +349,7 @@ def plot_fatbands(self, ebands_kpath,
                       marker_alpha=0.5, fontsize=12, interp_method="linear",
                       ax_mat=None, **kwargs) -> Figure:
         """
-        Plot fatbands.
+        Plot fatbands ...
 
         Args:
             ebands_kpath: ElectronBands instance with energies along a k-path
@@ -359,9 +436,6 @@ def __init__(self, filepath: str):
         super().__init__(filepath)
         self.r = ElectronsReader(filepath)
 
-        # Initialize the electron bands from file
-        self.natom = len(self.structure)
-
     @lazy_property
     def ebands(self) -> ElectronBands:
         """|ElectronBands| object."""

abipy/eph/tests/test_vpq.py

@@ -1,15 +1,20 @@
 """Tests for varpeq module."""
+import pytest
+import os
 import abipy.data as abidata
 
 from abipy.core.testing import AbipyTest
 from abipy.eph.vpq import VpqFile
 
+filepath = "/Users/giantomassi/git_repos/abinit/_build/tests/POLARON/varpeq6/out_VPQ.nc"
+
 
 class VarpeqTest(AbipyTest):
 
+    @pytest.mark.xfail(condition=not os.path.exists(filepath), reason=f"{filepath=} does not exist")
     def test_varpeq_file(self):
         """Testing VpqFile."""
-        filepath = "/Users/giantomassi/git_repos/abinit/_build/tests/POLARON/varpeq6/out_VPQ.nc"
+
         with VpqFile(filepath) as vpq:
             repr(vpq)
             str(vpq)

abipy/eph/vpq.py

@@ -117,10 +117,10 @@ class VpqFile(AbinitNcFile, Has_Structure, Has_ElectronBands, NotebookWriter):
 
     .. code-block:: python
 
-        from abipy.eph.varpeq import VpqFile
-        with VpqFile("out_VPQ.nc") as varpeq:
-            print(varpeq)
-            for polaron in varpeq.polaron_spin:
+        from abipy.eph.vpq import VpqFile
+        with VpqFile("out_VPQ.nc") as vpq:
+            print(vpq)
+            for polaron in vpq.polaron_spin:
                 print(polaron)
                 polaron.plot_scf_cycle()
 
@@ -154,7 +154,7 @@ def close(self) -> None:
     @lazy_property
     def polaron_spin(self) -> list[Polaron]:
         """List of Polaron objects, one for each spin (if any)."""
-        return [Polaron.from_varpeq(self, spin) for spin in range(self.r.nsppol)]
+        return [Polaron.from_vpq(self, spin) for spin in range(self.r.nsppol)]
 
     @lazy_property
     def params(self) -> dict:
@@ -236,14 +236,14 @@ class Polaron:
     spin: int          # Spin index.
     nstates: int       # Number of polaronic states for this spin.
     nb: int            # Number of bands in A_kn.
-    nk: int            # Number of k-points in A_kn, (including filtering if any).
+    nk: int            # Number of k-points in A_kn (including filtering if any).
     nq: int            # Number of q-points in B_qnu (including filtering if any).
     bstart: int        # First band starts at bstart.
     bstop: int         # Last band (python convention)
     varpeq: VpqFile
 
     @classmethod
-    def from_varpeq(cls, varpeq: VpqFile, spin: int) -> Polaron:
+    def from_vpq(cls, varpeq: VpqFile, spin: int) -> Polaron:
         """
         Build an istance from a VpqFile and the spin index.
         """

abipy/examples/plot/plot_orbmag.py

@@ -0,0 +1,29 @@
+#!/usr/bin/env python
+r"""
+Orbital magnetism
+=================
+
+This example shows how to plot ...
+See ...
+"""
+from abipy.electrons.orbmag import OrbmagAnalyzer
+#import abipy.data as abidata
+
+#%%
+# We start by defining a list with the paths to the PHDOS.nc files
+# In this case, for simplicity, we use the same file but we must
+# use different labels when adding them to the plotter with the add_phdos method.
+
+import os
+root = "/Users/giantomassi/git_repos/ABIPY_WORKS/JOE_ORBMAG"
+filepaths = [os.path.join(root, s) for s in ["gso_DS12_ORBMAG.nc", "gso_DS22_ORBMAG.nc", "gso_DS32_ORBMAG.nc"]]
+
+orban = OrbmagAnalyzer(filepaths)
+print(orban)
+
+choices = ['S','T','B','TB']
+orban.report_eigvals(report_type="S")
+
+#%%
+# Plot the phonon frequencies. Note that the labels for the q-points
+orban.plot_fatbands(os.path.join(root, "bandso_DS1_GSR.nc"))

abipy/examples/plot/plot_vpq.py

@@ -0,0 +1,36 @@
+#!/usr/bin/env python
+r"""
+e-ph matrix along a path in the BZ
+==================================
+
+This example shows how to use the VPQ.nc file produced by ABINIT
+to plot e-ph matrix elements along a path in the BZ.
+"""
+import os
+import abipy.data as abidata
+
+from abipy.eph.vpq import VpqFile
+
+filepath = "/Users/giantomassi/git_repos/abinit/_build/tests/POLARON/varpeq6/out_VPQ.nc"
+
+#%%
+# Here we use one of the GPATH files shipped with abipy
+# with the e-ph matrix g(k,q) with q along a path and k fixed.
+# Replace the call to abidata.ref_file("") with the path to your GPATH.nc
+
+varpeq = VpqFile(abidata.ref_file(filepath))
+print(varpeq)
+
+#%%
+# To plot the e-ph matrix elements as a function of q.
+for polaron in varpeq.polaron_spin:
+    df = polaron.get_final_results_df(with_params=True)
+    polaron.plot_scf_cycle()
+    #polaron.plot_ank_with_ebands(ebands_kpath, ebands_kmesh=None)
+    #polaron.plot_bqnu_with_ddb("in_DDB", with_phdos=True)
+    #polaron.plot_bqnu_with_phbands(phbands_qpath)
+
+
+#%%
+# For the meaning of the different arguments, see the docstring
+#print(varpeq.plot_g_qpath.__doc__)