1# coding: utf-8
2"""Classes for the analysis of GW calculations."""
3import sys
4import copy
5import warnings
6import numpy as np
7import pandas as pd
8
9from collections import namedtuple, OrderedDict
10from io import StringIO
11from tabulate import tabulate
12from monty.string import list_strings, is_string, marquee
13from monty.collections import dict2namedtuple
14from monty.functools import lazy_property
15from monty.termcolor import cprint
16from monty.bisect import find_le, find_ge
17from abipy.core.func1d import Function1D
18from abipy.core.kpoints import Kpoint, KpointList, Kpath, IrredZone, has_timrev_from_kptopt
19from abipy.core.mixins import AbinitNcFile, Has_Structure, Has_ElectronBands, NotebookWriter
20from abipy.iotools import ETSF_Reader
21from abipy.tools.plotting import (ArrayPlotter, add_fig_kwargs, get_ax_fig_plt, get_axarray_fig_plt, Marker,
22    set_axlims, set_visible, rotate_ticklabels)
23from abipy.tools import duck
24from abipy.abio.robots import Robot
25from abipy.electrons.ebands import ElectronBands, RobotWithEbands
26from abipy.electrons.scissors import Scissors
27
28import logging
29logger = logging.getLogger(__name__)
30
31__all__ = [
32    "QPState",
33    "SigresFile",
34    "SigresRobot",
35]
36
37
38class QPState(namedtuple("QPState", "spin kpoint band e0 qpe qpe_diago vxcme sigxme sigcmee0 vUme ze0")):
39    """
40    Quasi-particle result for given (spin, kpoint, band).
41
42    .. Attributes:
43
44        spin: spin index (C convention, i.e >= 0)
45        kpoint: |Kpoint| object.
46        band: band index. (C convention, i.e >= 0).
47        e0: Initial KS energy.
48        qpe: Quasiparticle energy (complex) computed with the perturbative approach.
49        qpe_diago: Quasiparticle energy (real) computed by diagonalizing the self-energy.
50        vxcme: Matrix element of vxc[n_val] with nval the valence charge density.
51        sigxme: Matrix element of Sigma_x.
52        sigcmee0: Matrix element of Sigma_c(e0) with e0 being the KS energy.
53        vUme: Matrix element of the vU term of the LDA+U Hamiltonian.
54        ze0: Renormalization factor computed at e=e0.
55
56    .. note::
57
58        Energies are in eV.
59    """
60    @property
61    def qpeme0(self):
62        """E_QP - E_0"""
63        return self.qpe - self.e0
64
65    @property
66    def re_qpe(self):
67        """Real part of the QP energy."""
68        return self.qpe.real
69
70    @property
71    def imag_qpe(self):
72        """Imaginay part of the QP energy."""
73        return self.qpe.imag
74
75    @property
76    def skb(self):
77        """Tuple with (spin, kpoint, band)"""
78        return self.spin, self.kpoint, self.band
79
80    #def __str__(self):
81    #    return self.to_string()
82
83    #def to_string(self, verbose=0, title=None):
84    #    """
85    #    String representation with verbosity level ``verbose`` and optional ``title``.
86    #    """
87    #    s = str(self.get_dataframe())
88    #    return "\n".join([marquee(title, mark="="), s]) if title is not None else s
89
90    def copy(self):
91        """Return Shallow copy."""
92        d = {f: copy.copy(getattr(self, f)) for f in self._fields}
93        return self.__class__(**d)
94
95    @classmethod
96    def get_fields(cls, exclude=()):
97        fields = list(cls._fields) + ["qpeme0"]
98        for e in exclude:
99            fields.remove(e)
100        return tuple(fields)
101
102    def as_dict(self, **kwargs):
103        """
104        Convert self into a dictionary.
105        """
106        od = OrderedDict(zip(self._fields, self))
107        od["qpeme0"] = self.qpeme0
108        return od
109
110    def to_strdict(self, fmt=None):
111        """
112        Ordered dictionary mapping fields --> strings.
113        """
114        d = self.as_dict()
115        for k, v in d.items():
116            if duck.is_intlike(v):
117                d[k] = "%d" % int(v)
118            elif isinstance(v, Kpoint):
119                d[k] = "%s" % v
120            elif np.iscomplexobj(v):
121                if abs(v.imag) < 1.e-3:
122                    d[k] = "%.2f" % v.real
123                else:
124                    d[k] = "%.2f%+.2fj" % (v.real, v.imag)
125            else:
126                try:
127                    d[k] = "%.2f" % v
128                except TypeError as exc:
129                    #print("k", k, str(exc))
130                    d[k] = str(v)
131        return d
132
133    @property
134    def tips(self):
135        """Bound method of self that returns a dictionary with the description of the fields."""
136        return self.__class__.TIPS()
137
138    @classmethod
139    def TIPS(cls):
140        """
141        Class method that returns a dictionary with the description of the fields.
142        The string are extracted from the class doc string.
143        """
144        try:
145            return cls._TIPS
146
147        except AttributeError:
148            # Parse the doc string.
149            cls._TIPS = _TIPS = {}
150            lines = cls.__doc__.splitlines()
151
152            for i, line in enumerate(lines):
153                if line.strip().startswith(".. Attributes"):
154                    lines = lines[i+1:]
155                    break
156
157            def num_leadblanks(string):
158                """Returns the number of the leading whitespaces."""
159                return len(string) - len(string.lstrip())
160
161            for field in cls._fields:
162                for i, line in enumerate(lines):
163
164                    if line.strip().startswith(field + ":"):
165                        nblanks = num_leadblanks(line)
166                        desc = []
167                        for s in lines[i+1:]:
168                            if nblanks == num_leadblanks(s) or not s.strip():
169                                break
170                            desc.append(s.lstrip())
171
172                        _TIPS[field] = "\n".join(desc)
173
174            diffset = set(cls._fields) - set(_TIPS.keys())
175            if diffset:
176                raise RuntimeError("The following fields are not documented: %s" % str(diffset))
177
178            return _TIPS
179
180    @classmethod
181    def get_fields_for_plot(cls, with_fields, exclude_fields):
182        """
183        Return list of QPState fields to plot from input arguments.
184        """
185        all_fields = list(cls.get_fields(exclude=["spin", "kpoint"]))[:]
186
187        # Initialize fields.
188        if is_string(with_fields) and with_fields == "all":
189            fields = all_fields
190        else:
191            fields = list_strings(with_fields)
192            for f in fields:
193                if f not in all_fields:
194                    raise ValueError("Field %s not in allowed values %s" % (f, all_fields))
195
196        # Remove entries
197        if exclude_fields:
198            if is_string(exclude_fields):
199                exclude_fields = exclude_fields.split()
200            for e in exclude_fields:
201                try:
202                    fields.remove(e)
203                except ValueError:
204                    pass
205
206        return fields
207
208
209class QPList(list):
210    """
211    A list of quasiparticle corrections for a given spin.
212    """
213    def __init__(self, *args, **kwargs):
214        super().__init__(*args)
215        self.is_e0sorted = kwargs.get("is_e0sorted", False)
216
217    def __repr__(self):
218        return "<%s at %s, len=%d>" % (self.__class__.__name__, id(self), len(self))
219
220    def __str__(self):
221        """String representation."""
222        return self.to_string()
223
224    def to_table(self):
225        """Return a table (list of list of strings)."""
226        header = QPState.get_fields(exclude=["spin", "kpoint"])
227        table = [header]
228        for qp in self:
229            d = qp.to_strdict(fmt=None)
230            table.append([d[k] for k in header])
231
232        return tabulate(table, tablefmt="plain")
233
234    def to_string(self, verbose=0):
235        """String representation."""
236        return self.to_table()
237
238    def copy(self):
239        """Copy of self."""
240        return self.__class__([qp.copy() for qp in self], is_e0sorted=self.is_e0sorted)
241
242    def sort_by_e0(self):
243        """Return a new object with the E0 energies sorted in ascending order."""
244        return self.__class__(sorted(self, key=lambda qp: qp.e0), is_e0sorted=True)
245
246    def get_e0mesh(self):
247        """Return the E0 energies."""
248        if not self.is_e0sorted:
249            raise ValueError("QPState corrections are not sorted. Use sort_by_e0")
250
251        return np.array([qp.e0 for qp in self])
252
253    def get_field(self, field):
254        """|numpy-array| containing the values of field."""
255        return np.array([getattr(qp, field) for qp in self])
256
257    def get_skb_field(self, skb, field):
258        """Return the value of field for the given spin kp band tuple, None if not found"""
259        for qp in self:
260            if qp.skb == skb:
261                return getattr(qp, field)
262        return None
263
264    def get_qpenes(self):
265        """Return an array with the :class:`QPState` energies."""
266        return self.get_field("qpe")
267
268    def get_qpeme0(self):
269        """Return an arrays with the :class:`QPState` corrections."""
270        return self.get_field("qpeme0")
271
272    def merge(self, other, copy=False):
273        """
274        Merge self with other. Return new :class:`QPList` object
275
276        Raise:
277            `ValueError` if merge cannot be done.
278        """
279        skb0_list = [qp.skb for qp in self]
280        for qp in other:
281            if qp.skb in skb0_list:
282                raise ValueError("Found duplicated (s,b,k) indexes: %s" % str(qp.skb))
283
284        qps = self.copy() + other.copy() if copy else self + other
285        return self.__class__(qps)
286
287    @add_fig_kwargs
288    def plot_qps_vs_e0(self, with_fields="all", exclude_fields=None, fermie=None,
289                       ax_list=None, sharey=False, xlims=None, fontsize=12, **kwargs):
290        """
291        Plot the QP results as function of the initial KS energy.
292
293        Args:
294            with_fields: The names of the qp attributes to plot as function of e0.
295                Accepts: List of strings or string with tokens separated by blanks.
296                See :class:`QPState` for the list of available fields.
297            exclude_fields: Similar to ``with_field`` but excludes fields.
298            fermie: Value of the Fermi level used in plot. None for absolute e0s.
299            ax_list: List of |matplotlib-Axes| for plot. If None, new figure is produced.
300            sharey: True if y-axis should be shared.
301            kwargs: linestyle, color, label, marker
302
303        Returns: |matplotlib-Figure|
304        """
305        # TODO: This is to maintain the old API.
306        fermi = kwargs.pop("fermi", None)
307        if fermi is not None:
308            fermie = fermi
309            warnings.warn("fermi keyword argument have been renamed fermie. Old arg will be removed in version 0.4")
310
311        fields = QPState.get_fields_for_plot(with_fields, exclude_fields)
312        if not fields: return None
313
314        num_plots, ncols, nrows = len(fields), 1, 1
315        if num_plots > 1:
316            ncols = 2
317            nrows = (num_plots // ncols) + (num_plots % ncols)
318
319        # Build grid of plots.
320        ax_list, fig, plt = get_axarray_fig_plt(ax_list, nrows=nrows, ncols=ncols,
321                                                sharex=True, sharey=sharey, squeeze=False)
322        ax_list = np.array(ax_list).ravel()
323
324        # Get qplist and sort it.
325        qps = self if self.is_e0sorted else self.sort_by_e0()
326        e0mesh = qps.get_e0mesh()
327        xlabel = r"$\epsilon_{KS}\;(eV)$"
328        #print("fermie", fermie)
329        if fermie is not None:
330            xlabel = r"$\epsilon_{KS}-\epsilon_F\;(eV)$"
331            e0mesh -= fermie
332
333        kw_linestyle = kwargs.pop("linestyle", "o")
334        kw_color = kwargs.pop("color", None)
335        kw_label = kwargs.pop("label", None)
336
337        for ii, (field, ax) in enumerate(zip(fields, ax_list)):
338            irow, icol = divmod(ii, ncols)
339            ax.grid(True)
340            if irow == nrows - 1:
341                ax.set_xlabel(xlabel)
342            ax.set_ylabel(field, fontsize=fontsize)
343            yy = qps.get_field(field)
344
345            # TODO real and imag?
346            ax.plot(e0mesh, yy.real, kw_linestyle, color=kw_color, label=kw_label, **kwargs)
347            set_axlims(ax, xlims, "x")
348
349        if kw_label:
350            ax_list[0].legend(loc="best", fontsize=fontsize, shadow=True)
351
352        # Get around a bug in matplotlib
353        if num_plots % ncols != 0: ax_list[-1].axis('off')
354
355        return fig
356
357    def build_scissors(self, domains, bounds=None, k=3, plot=False, **kwargs):
358        """
359        Construct a scissors operator by interpolating the QPState corrections
360        as function of the initial energies E0.
361
362        Args:
363            domains: list in the form [ [start1, stop1], [start2, stop2]
364                Domains should not overlap, cover e0mesh, and given in increasing order.
365                Holes are permitted but the interpolation will raise an exception if
366                the point is not in domains.
367            bounds: Specify how to handle out-of-boundary conditions, i.e. how to treat
368                energies that do not fall inside one of the domains (not used at present)
369            plot: If true, use matplolib plot to compare input data  and fit.
370
371        Return:
372            instance of :class:`Scissors`operator
373
374        Usage example:
375
376        .. code-block:: python
377
378            # Build the scissors operator.
379            scissors = qplist_spin[0].build_scissors(domains)
380
381            # Compute list of interpolated QP energies.
382            qp_enes = [scissors.apply(e0) for e0 in ks_energies]
383        """
384        # Sort QP corrections according to the initial KS energy.
385        qps = self.sort_by_e0()
386        e0mesh, qpcorrs = qps.get_e0mesh(), qps.get_qpeme0()
387
388        # Check domains.
389        domains = np.atleast_2d(domains)
390        dsize, dflat = domains.size, domains.ravel()
391
392        for idx, v in enumerate(dflat):
393            if idx == 0 and v > e0mesh[0]:
394                raise ValueError("min(e0mesh) %s is not included in domains" % e0mesh[0])
395            if idx == dsize-1 and v < e0mesh[-1]:
396                raise ValueError("max(e0mesh) %s is not included in domains" % e0mesh[-1])
397            if idx != dsize-1 and dflat[idx] > dflat[idx+1]:
398                raise ValueError("domain boundaries should be given in increasing order.")
399            if idx == dsize-1 and dflat[idx] < dflat[idx-1]:
400                raise ValueError("domain boundaries should be given in increasing order.")
401
402        # Create the sub_domains and the spline functions in each subdomain.
403        func_list, residues = [], []
404
405        if len(domains) == 2:
406            #print('forcing extremal point on the scissor')
407            ndom = 0
408        else:
409            ndom = 99
410
411        for dom in domains[:]:
412            ndom += 1
413            low, high = dom[0], dom[1]
414            start, stop = find_ge(e0mesh, low), find_le(e0mesh, high)
415
416            dom_e0 = e0mesh[start:stop+1]
417            dom_corr = qpcorrs[start:stop+1]
418
419            # todo check if the number of non degenerate data points > k
420            from scipy.interpolate import UnivariateSpline
421            w = len(dom_e0)*[1]
422            if ndom == 1:
423                w[-1] = 1000
424            elif ndom == 2:
425                w[0] = 1000
426            else:
427                w = None
428            f = UnivariateSpline(dom_e0, dom_corr, w=w, bbox=[None, None], k=k, s=None)
429            func_list.append(f)
430            residues.append(f.get_residual())
431
432        # Build the scissors operator.
433        sciss = Scissors(func_list, domains, residues, bounds)
434
435        # Compare fit with input data.
436        if plot:
437            title = kwargs.pop("title", None)
438            import matplotlib.pyplot as plt
439            plt.plot(e0mesh, qpcorrs, 'o', label="input data")
440            if title: plt.suptitle(title)
441            for dom in domains[:]:
442                plt.plot(2*[dom[0]], [min(qpcorrs), max(qpcorrs)])
443                plt.plot(2*[dom[1]], [min(qpcorrs), max(qpcorrs)])
444            intp_qpc = [sciss.apply(e0) for e0 in e0mesh]
445            plt.plot(e0mesh, intp_qpc, label="scissor")
446            plt.legend(bbox_to_anchor=(0.9, 0.2))
447            plt.show()
448
449        # Return the object.
450        return sciss
451
452
453class SelfEnergy(object):
454    """
455    This object stores the values of the electron-electron self-energy
456    and the associated spectral function as function of frequency
457    """
458    # Symbols used in matplotlib plots.
459    latex_symbol = dict(
460        re=r"$\Re{\Sigma(\omega)}$",
461        im=r"$\Im{\Sigma(\omega)}$",
462        spfunc=r"$A(\omega)}$",
463    )
464
465    def __init__(self, spin, kpoint, band, wmesh, sigmaxc_values, spfunc_values):
466        self.spin, self.kpoint, self.band = spin, kpoint, band
467        self.wmesh = np.array(wmesh)
468        self.xc = Function1D(self.wmesh, sigmaxc_values)
469        self.spfunc = Function1D(self.wmesh, spfunc_values)
470        assert len(sigmaxc_values) == len(spfunc_values)
471
472    def __str__(self):
473        return self.to_string()
474
475    def to_string(self, verbose=0, title=None):
476        """String representation."""
477        lines = []; app = lines.append
478        if title is not None: app(marquee(title, mark="="))
479        app("K-point: %s, band: %d, spin: %d" % (repr(self.kpoint), self.band, self.spin))
480        app("Number of frequencies: %d, from %.1f to %.1f (eV)" % (len(self.wmesh), self.wmesh[0], self.wmesh[-1]))
481
482        return "\n".join(lines)
483
484    def plot_ax(self, ax, what="a", fontsize=12, **kwargs):
485        """
486        Helper function to plot data on the axis ax with fontsize
487
488        Args:
489            ax: |matplotlib-Axes| or None if a new figure should be created.
490            what:
491            fontsize: legend and title fontsize.
492
493        Return: List of lines.
494        """
495        #if not kwargs: kwargs = {"color": "black", "linewidth": 2.0}
496        lines = []
497        extend = lines.extend
498        ax.grid(True)
499
500        if what == "s":
501            f = self.xc
502            label = kwargs.get("label", r"$\Sigma(\omega)$")
503            extend(f.plot_ax(ax, cplx_mode="re", label="Re " + label))
504            extend(f.plot_ax(ax, cplx_mode="im", label="Im " + label))
505            #ax.set_ylabel('Energy (eV)')
506
507        elif what == "a":
508            f = self.spfunc
509            label = kwargs.get("label", r"$A(\omega)$")
510            extend(f.plot_ax(ax, label=label))
511        else:
512            raise ValueError("Don't know how to handle what option %s" % str(what))
513
514        #ax.set_xlabel(r"$\omega - \espilon_{KS} (eV)")
515        ax.legend(loc="best", fontsize=fontsize, shadow=True)
516
517        return lines
518
519    @add_fig_kwargs
520    def plot(self, ax_list=None, what_list=("re", "im", "spfunc"), fermie=None,
521             xlims=None, fontsize=12, **kwargs):
522        """
523        Plot the real/imaginary part of self-energy as well as the spectral function
524
525        Args:
526            ax_list: List of |matplotlib-Axes| for plot. If None, new figure is produced.
527            what_list: List of strings selecting the quantity to plot.
528            fermie: Value of the Fermi level used in plot. None for absolute e0s.
529            xlims: Set the data limits for the x-axis. Accept tuple e.g. ``(left, right)``
530                or scalar e.g. ``left``. If left (right) is None, default values are used.
531            fontsize: legend and label fontsize.
532            kwargs: Keyword arguments passed to ax.plot
533
534        Returns: |matplotlib-Figure|
535        """
536        what_list = list_strings(what_list)
537        ax_list, fig, plt = get_axarray_fig_plt(ax_list, nrows=len(what_list), ncols=1,
538                                                sharex=True, sharey=False, squeeze=False)
539        ax_list = np.array(ax_list).ravel()
540        xlabel = r"$\omega\;(eV)$"
541        wmesh = self.wmesh
542        if fermie is not None:
543            xlabel = r"$\omega - \epsilon_F}\;(eV)$"
544            wmesh = self.wmesh - fermie
545
546        kw_color = kwargs.pop("color", None)
547        kw_label = kwargs.pop("label", None)
548        for i, (what, ax) in enumerate(zip(what_list, ax_list)):
549            ax.grid(True)
550            ax.set_ylabel(self.latex_symbol[what])
551            if (i == len(ax_list) - 1): ax.set_xlabel(xlabel)
552            ax.plot(wmesh, self._get_ys(what),
553                    color=kw_color,
554                    label=kw_label if i == 0 else None,
555            )
556            set_axlims(ax, xlims, "x")
557            if i == 0 and kw_label:
558                ax.legend(loc="best", shadow=True, fontsize=fontsize)
559
560        if "title" not in kwargs:
561            title = "K-point: %s, band: %d, spin: %d" % (repr(self.kpoint), self.band, self.spin)
562            fig.suptitle(title, fontsize=fontsize)
563
564        return fig
565
566    def _get_ys(self, what):
567        """Return array to plot from what string."""
568        return dict(
569            re=self.xc.values.real,
570            im=self.xc.values.imag,
571            spfunc=self.spfunc.values,
572        )[what]
573
574
575class SigresFile(AbinitNcFile, Has_Structure, Has_ElectronBands, NotebookWriter):
576    """
577    Container storing the GW results reported in the SIGRES.nc file.
578
579    Usage example:
580
581    .. code-block:: python
582
583        sigres = SigresFile("foo_SIGRES.nc")
584        sigres.plot_qps_vs_e0()
585
586    .. rubric:: Inheritance Diagram
587    .. inheritance-diagram:: SigresFile
588    """
589    # Markers used for up/down bands.
590    marker_spin = {0: "^", 1: "v"}
591    color_spin = {0: "k", 1: "r"}
592
593    @classmethod
594    def from_file(cls, filepath):
595        """Initialize an instance from file."""
596        return cls(filepath)
597
598    def __init__(self, filepath):
599        """Read data from the netcdf file path."""
600        super().__init__(filepath)
601
602        # Keep a reference to the SigresReader.
603        self.reader = reader = SigresReader(self.filepath)
604
605        self._structure = reader.read_structure()
606        self.gwcalctyp = reader.gwcalctyp
607        self.ibz = reader.ibz
608        self.gwkpoints = reader.gwkpoints
609        self.nkcalc = len(self.gwkpoints)
610
611        self.gwbstart_sk = reader.gwbstart_sk
612        self.gwbstop_sk = reader.gwbstop_sk
613        self.min_gwbstart = reader.min_gwbstart
614        self.max_gwbstart = reader.max_gwbstart
615        self.min_gwbstop = reader.min_gwbstop
616        self.max_gwbstop = reader.max_gwbstop
617
618        self._ebands = ebands = reader.ks_bands
619
620        qplist_spin = self.qplist_spin
621
622        # TODO handle the case in which nkptgw < nkibz
623        self.qpgaps = reader.read_qpgaps()
624        self.qpenes = reader.read_qpenes()
625        self.ksgaps = reader.read_ksgaps()
626
627    @property
628    def sigma_kpoints(self):
629        """The K-points where QP corrections have been calculated."""
630        return self.gwkpoints
631
632    def get_marker(self, qpattr):
633        """
634        Return :class:`Marker` object associated to the QP attribute qpattr.
635        Used to prepare plots of KS bands with markers.
636        """
637        # Each marker is a list of tuple(x, y, value)
638        x, y, s = [], [], []
639
640        for spin in range(self.nsppol):
641            for qp in self.qplist_spin[spin]:
642                ik = self.ebands.kpoints.index(qp.kpoint)
643                x.append(ik)
644                y.append(qp.e0)
645                size = getattr(qp, qpattr)
646                # Handle complex quantities
647                if np.iscomplex(size): size = size.real
648                s.append(size)
649
650        return Marker(*(x, y, s))
651
652    @lazy_property
653    def params(self):
654        """:class:`OrderedDict` with parameters that might be subject to convergence studies e.g ecuteps"""
655        od = self.get_ebands_params()
656        od.update(self.reader.read_params())
657        return od
658
659    def close(self):
660        """Close the netcdf file."""
661        self.reader.close()
662
663    def __str__(self):
664        return self.to_string()
665
666    def to_string(self, verbose=0):
667        """String representation with verbosity level ``verbose``."""
668        lines = []; app = lines.append
669
670        app(marquee("File Info", mark="="))
671        app(self.filestat(as_string=True))
672        app("")
673        app(self.structure.to_string(verbose=verbose, title="Structure"))
674        app("")
675        app(self.ebands.to_string(title="Kohn-Sham bands", with_structure=False))
676        app("")
677
678        # TODO: Finalize the implementation: add GW metadata.
679        app(marquee("QP direct gaps", mark="="))
680        for kgw in self.gwkpoints:
681            for spin in range(self.nsppol):
682                qp_dirgap = self.get_qpgap(spin, kgw)
683                app("QP_dirgap: %.3f (eV) for K-point: %s, spin: %s" % (qp_dirgap, repr(kgw), spin))
684                #ks_dirgap =
685        app("")
686
687        # Show QP results
688        strio = StringIO()
689        self.print_qps(precision=3, ignore_imag=verbose == 0, file=strio)
690        strio.seek(0)
691        app("")
692        app(marquee("QP results for each k-point and spin (all in eV)", mark="="))
693        app("".join(strio))
694        app("")
695
696        # TODO: Fix header.
697        #if verbose > 1:
698        #    app("")
699        #    app(marquee("Abinit Header", mark="="))
700        #    app(self.hdr.to_string(verbose=verbose))
701
702        return "\n".join(lines)
703
704    @property
705    def structure(self):
706        """|Structure| object."""
707        return self._structure
708
709    @property
710    def ebands(self):
711        """|ElectronBands| with the KS energies."""
712        return self._ebands
713
714    @property
715    def has_spectral_function(self):
716        """True if file contains spectral function data."""
717        return self.reader.has_spfunc
718
719    @lazy_property
720    def qplist_spin(self):
721        """Tuple of :class:`QPList` objects indexed by spin."""
722        return self.reader.read_allqps()
723
724    def get_qplist(self, spin, kpoint, ignore_imag=False):
725        """Return :class`QPList` for the given (spin, kpoint)"""
726        return self.reader.read_qplist_sk(spin, kpoint, ignore_imag=ignore_imag)
727
728    def get_qpcorr(self, spin, kpoint, band, ignore_imag=False):
729        """Returns the :class:`QPState` object for the given (s, k, b)"""
730        return self.reader.read_qp(spin, kpoint, band, ignore_imag=ignore_imag)
731
732    @lazy_property
733    def qpgaps(self):
734        """|numpy-array| of shape [nsppol, nkibz] with the QP direct gaps in eV."""
735        return self.reader.read_qpgaps()
736
737    def get_qpgap(self, spin, kpoint, with_ksgap=False):
738        """Return the QP gap in eV at the given (spin, kpoint)"""
739        k = self.reader.kpt2fileindex(kpoint)
740        if not with_ksgap:
741            return self.qpgaps[spin, k]
742        else:
743            return self.qpgaps[spin, k], self.ksgaps[spin, k]
744
745    def read_sigee_skb(self, spin, kpoint, band):
746        """"
747        Read self-energy(w) for (spin, kpoint, band)
748        Return: :class:`SelfEnergy` object
749        """
750        ik = self.reader.kpt2fileindex(kpoint)
751        kpoint = self.ibz[ik]
752        wmesh, sigxc_values = self.reader.read_sigmaw(spin, ik, band)
753        wmesh, spf_values = self.reader.read_spfunc(spin, ik, band)
754
755        return SelfEnergy(spin, kpoint, band, wmesh, sigxc_values, spf_values)
756
757    def print_qps(self, precision=3, ignore_imag=True, file=sys.stdout):
758        """
759        Print QP results to stream ``file``.
760
761        Args:
762            precision: Number of significant digits.
763            ignore_imag: True if imaginary part should be ignored.
764            file: Output stream.
765        """
766        from abipy.tools.printing import print_dataframe
767        keys = "band e0 qpe qpe_diago vxcme sigxme sigcmee0 vUme ze0".split()
768        for gwkpoint in self.gwkpoints:
769            for spin in range(self.nsppol):
770                df_sk = self.get_dataframe_sk(spin, gwkpoint, ignore_imag=ignore_imag)[keys]
771                print_dataframe(df_sk, title="K-point: %s, spin: %s" % (repr(gwkpoint), spin),
772                                precision=precision, file=file)
773
774    def get_points_from_ebands(self, ebands_kpath, dist_tol=1e-12, size=24, verbose=0):
775        """
776        Generate object storing the QP energies lying on the k-path used by ebands_kpath.
777        Mainly used to plot the QP energies in ebands.plot when the QP energies are interpolated with the SKW method.
778
779        Args:
780            ebands_kpath: |ElectronBands| object with the QP energies along an arbitrary k-path.
781            dist_tol: A point is considered to be on the path if its distance from the line
782                is less than dist_tol.
783            size: The marker size in points**2
784            verbose: Verbosity level
785
786        Return:
787
788        Example::
789
790            r = sigres.interpolate(lpratio=5,
791                                   ks_ebands_kpath=ks_ebands_kpath,
792                                   ks_ebands_kmesh=ks_ebands_kmesh
793                                   )
794            points = sigres.get_points_from_ebands(r.qp_ebands_kpath, size=24)
795            r.qp_ebands_kpath.plot(points=points)
796        """
797        kpath = ebands_kpath.kpoints
798        if not ebands_kpath.kpoints.is_path:
799            raise TypeError("Expecting band structure with a Kpath, got %s" % type(kpath))
800        if verbose:
801            print("Input kpath\n", ebands_kpath.kpoints)
802            print("gwkpoints included in GW calculation\n", self.gwkpoints)
803            print("lines\n", kpath.lines)
804            print("kpath.frac_bounds:\n", kpath.frac_bounds)
805            print("kpath.cart_bounds:\n", kpath.frac_bounds)
806
807        # Construct the stars of the k-points for all k-points included in the GW calculation.
808        # In principle, the input k-path is arbitrary and not necessarily in the IBZ used for GW
809        # so we have to build the k-stars and find the k-points lying along the path and keep
810        # track of the mapping kpt --> star --> kgw
811        gw_stars = [kpoint.compute_star(self.structure.abi_spacegroup.fm_symmops) for kpoint in self.gwkpoints]
812        cart_coords, back2istar = [], []
813        for istar, gw_star in enumerate(gw_stars):
814            cart_coords.extend([k.cart_coords for k in gw_star])
815            back2istar.extend([istar] * len(gw_star))
816        cart_coords = np.reshape(cart_coords, (-1, 3))
817
818        # Find (star) k-points on the path.
819        p = kpath.find_points_along_path(cart_coords, dist_tol=dist_tol)
820        if len(p.ikfound) == 0:
821            cprint("Warning: Found zero points lying on the input k-path. Try to increase dist_tol.", "yellow")
822            return Marker()
823
824        # Read complex GW energies from file.
825        qp_arr = self.reader.read_value("egw", cmode="c")
826
827        # Each marker is a list of tuple(x, y, value)
828        x, y, s = [], [], []
829        kpath_lenght = kpath.ds.sum()
830
831        for ik, dalong_path in zip(p.ikfound, p.dist_list):
832            istar = back2istar[ik]
833            # The k-point used in the GW calculation.
834            gwk = gw_stars[istar].base_point
835            # Indices needed to access SIGRES arrays (we have to live with this format)
836            ik_ibz = self.reader.kpt2fileindex(gwk)
837            ik_b = self.reader.gwkpt2seqindex(gwk)
838            for spin in range(self.nsppol):
839                # Need to select bands included in the GW calculation.
840                for qpe in qp_arr[spin, ik_ibz, self.gwbstart_sk[spin, ik_b]:self.gwbstop_sk[spin, ik_b]]:
841                    # Assume the path is properly normalized.
842                    x.append((dalong_path / kpath_lenght) * (len(kpath) - 1))
843                    y.append(qpe.real)
844                    s.append(size)
845
846        return Marker(*(x, y, s))
847
848    @add_fig_kwargs
849    def plot_qpgaps(self, ax=None, plot_qpmks=True, fontsize=8, **kwargs):
850        """
851        Plot the KS and the QP direct gaps for all the k-points and spins available on file.
852
853        Args:
854            ax: |matplotlib-Axes| or None if a new figure should be created.
855            plot_qpmks: If False, plot QP_gap, KS_gap else (QP_gap - KS_gap)
856            fontsize: legend and title fontsize.
857            kwargs: Passed to ax.plot method except for marker.
858
859        Returns: |matplotlib-Figure|
860        """
861        ax, fig, plt = get_ax_fig_plt(ax=ax)
862        label = kwargs.pop("label", None)
863        xs = np.arange(self.nkcalc)
864
865        # Add xticklabels from k-points.
866        tick_labels = [repr(k) for k in self.gwkpoints]
867        ax.set_xticks(xs)
868        ax.set_xticklabels(tick_labels, fontdict=None, rotation=30, minor=False, size="x-small")
869
870        for spin in range(self.nsppol):
871            qp_gaps, ks_gaps = map(np.array, zip(*[self.get_qpgap(spin, kgw, with_ksgap=True) for kgw in self.gwkpoints]))
872            if not plot_qpmks:
873                # Plot QP gaps
874                ax.plot(xs, qp_gaps, marker=self.marker_spin[spin],
875                        label=label if spin == 0 else None, **kwargs)
876                # Add KS gaps
877                #ax.scatter(xx, ks_gaps) #, label="KS gap %s" % label)
878            else:
879                # Plot QP_gap - KS_gap
880                ax.plot(xs, qp_gaps - ks_gaps, marker=self.marker_spin[spin],
881                    label=label if spin == 0 else None, **kwargs)
882
883            ax.grid(True)
884            ax.set_xlabel("K-point")
885            if plot_qpmks:
886                ax.set_ylabel("QP-KS gap (eV)")
887            else:
888                ax.set_ylabel("QP direct gap (eV)")
889            #ax.set_title("k:%s" % (repr(kgw)), fontsize=fontsize)
890            if label:
891                ax.legend(loc="best", fontsize=fontsize, shadow=True)
892
893        return fig
894
895    @add_fig_kwargs
896    def plot_qps_vs_e0(self, with_fields="all", exclude_fields=None, e0="fermie",
897                       xlims=None, sharey=False, ax_list=None, fontsize=8, **kwargs):
898        """
899        Plot QP result in SIGRES file as function of the KS energy.
900
901        Args:
902            with_fields: The names of the qp attributes to plot as function of eKS.
903                Accepts: List of strings or string with tokens separated by blanks.
904                See :class:`QPState` for the list of available fields.
905            exclude_fields: Similar to ``with_fields`` but excludes fields
906            e0: Option used to define the zero of energy in the band structure plot. Possible values:
907                - `fermie`: shift all eigenvalues to have zero energy at the Fermi energy (`self.fermie`).
908                -  Number e.g e0=0.5: shift all eigenvalues to have zero energy at 0.5 eV
909                -  None: Don't shift energies, equivalent to e0=0
910            ax_list: List of |matplotlib-Axes| for plot. If None, new figure is produced.
911            xlims: Set the data limits for the x-axis. Accept tuple e.g. ``(left, right)``
912                   or scalar e.g. ``left``. If left (right) is None, default values are used.
913            sharey: True if y-axis should be shared.
914            fontsize: Legend and title fontsize.
915
916        Returns: |matplotlib-Figure|
917        """
918        with_fields = QPState.get_fields_for_plot(with_fields, exclude_fields)
919
920        # Because qplist does not have the fermi level.
921        fermie = self.ebands.get_e0(e0) if e0 is not None else None
922        for spin in range(self.nsppol):
923            fig = self.qplist_spin[spin].plot_qps_vs_e0(
924                with_fields=with_fields, exclude_fields=exclude_fields, fermie=fermie,
925                xlims=xlims, sharey=sharey, ax_list=ax_list, fontsize=fontsize,
926                marker=self.marker_spin[spin], show=False, **kwargs)
927            ax_list = fig.axes
928
929        return fig
930
931    @add_fig_kwargs
932    def plot_spectral_functions(self, include_bands=None, fontsize=8, **kwargs):
933        """
934        Plot the spectral function for all k-points, bands and spins available in the SIGRES file.
935
936        Args:
937            include_bands: List of bands to include. Nonee means all.
938            fontsize: Legend and title fontsize.
939
940        Returns: |matplotlib-Figure|
941        """
942        if include_bands is not None:
943            include_bands = set(include_bands)
944
945        # Build grid of plots.
946        nrows, ncols = len(self.sigma_kpoints), 1
947        ax_list = None
948        ax_list, fig, plt = get_axarray_fig_plt(ax_list, nrows=nrows, ncols=ncols,
949                                                sharex=True, sharey=False, squeeze=False)
950        ax_list = np.array(ax_list).ravel()
951
952        for ik_gw, (kgw, ax) in enumerate(zip(self.sigma_kpoints, ax_list)):
953            for spin in range(self.nsppol):
954                for band in range(self.gwbstart_sk[spin, ik_gw], self.gwbstop_sk[spin, ik_gw]):
955                    if include_bands and band not in include_bands: continue
956                    sigw = self.read_sigee_skb(spin, kgw, band)
957                    label = r"$A(\omega)$: band: %d, spin: %d" % (spin, band)
958                    sigw.plot_ax(ax, what="a", label=label, fontsize=fontsize, **kwargs)
959
960            ax.set_title("K-point: %s" % repr(sigw.kpoint), fontsize=fontsize)
961            #if ik_gw == len(self.sigma_kpoints) - 1:
962
963        return fig
964
965    @add_fig_kwargs
966    def plot_eigvec_qp(self, spin=0, kpoint=None, band=None, **kwargs):
967        """
968
969        Args:
970            spin: Spin index
971            kpoint: K-point in self-energy. Accepts |Kpoint|, vector or index.
972                If None, all k-points for the given ``spin`` are shown.
973            band: band index. If None all bands are displayed else
974                only <KS_b|QP_{b'}> for the given b.
975            kwargs: Passed to plot method of :class:`ArrayPlotter`.
976
977        Returns: |matplotlib-Figure|
978        """
979        plotter = ArrayPlotter()
980        if kpoint is None:
981            for kpoint in self.ibz:
982                ksqp_arr = self.reader.read_eigvec_qp(spin, kpoint, band=band)
983                plotter.add_array(repr(kpoint), ksqp_arr)
984            return plotter.plot(show=False, **kwargs)
985
986        else:
987            ksqp_arr = self.reader.read_eigvec_qp(spin, kpoint, band=band)
988            plotter.add_array(repr(kpoint), ksqp_arr)
989            return plotter.plot(show=False, **kwargs)
990
991    @add_fig_kwargs
992    def plot_qpbands_ibz(self, e0="fermie", colormap="jet", ylims=None, fontsize=8, **kwargs):
993        r"""
994        Plot the KS band structure in the IBZ with the QP energies.
995
996        Args:
997            e0: Option used to define the zero of energy in the band structure plot.
998            colormap: matplotlib color map.
999            ylims: Set the data limits for the y-axis. Accept tuple e.g. ``(left, right)``
1000                   or scalar e.g. ``left``. If left (right) is None, default values are used.
1001            fontsize: Legend and title fontsize.
1002
1003        Returns: |matplotlib-Figure|
1004        """
1005        # Map sigma_kpoints to ebands.kpoints
1006        kcalc2ibz = np.empty(self.nkcalc, dtype=int)
1007        for ikc, sigkpt in enumerate(self.sigma_kpoints):
1008            kcalc2ibz[ikc] = self.ebands.kpoints.index(sigkpt)
1009
1010        # TODO: It seems there's a minor issue with fermie if SCF band structure.
1011        e0 = self.ebands.get_e0(e0)
1012        #print("e0",e0, self.ebands.fermie)
1013
1014        # Build grid with (1, nsppol) plots.
1015        nrows, ncols = 1, self.nsppol
1016        ax_list, fig, plt = get_axarray_fig_plt(None, nrows=nrows, ncols=ncols,
1017                                                sharex=True, sharey=True, squeeze=False)
1018        ax_list = np.array(ax_list).ravel()
1019        cmap = plt.get_cmap(colormap)
1020
1021        # Read QP energies: Fortran egw(nbnds,nkibz,nsppol)
1022        qpes = self.reader.read_value("egw", cmode="c") # * units.Ha_to_eV
1023        band_range = (self.reader.max_gwbstart, self.reader.min_gwbstop)
1024
1025        nb = self.reader.min_gwbstop - self.reader.min_gwbstart
1026        for spin, ax in zip(range(self.nsppol), ax_list):
1027            # Plot KS bands in the band range included in self-energy calculation.
1028            self.ebands.plot(ax=ax, e0=e0, spin=spin, band_range=band_range, show=False)
1029            # Extract QP in IBZ
1030            yvals = qpes[spin, kcalc2ibz, :].real - e0
1031            # Add (scattered) QP energies for the calculated k-points.
1032            for band in range(self.reader.max_gwbstart, self.reader.min_gwbstop):
1033                ax.scatter(kcalc2ibz, yvals[:, band],
1034                    color=cmap(band / nb), alpha=0.6, marker="o", s=20,
1035                )
1036            set_axlims(ax, ylims, "y")
1037
1038        return fig
1039
1040    @add_fig_kwargs
1041    def plot_ksbands_with_qpmarkers(self, qpattr="qpeme0", e0="fermie", fact=1000, ax=None, **kwargs):
1042        """
1043        Plot the KS energies as function of k-points and add markers whose size
1044        is proportional to the QPState attribute ``qpattr``
1045
1046        Args:
1047            qpattr: Name of the QP attribute to plot. See :class:`QPState`.
1048            e0: Option used to define the zero of energy in the band structure plot. Possible values:
1049                - ``fermie``: shift all eigenvalues to have zero energy at the Fermi energy (``self.fermie``).
1050                -  Number e.g ``e0 = 0.5``: shift all eigenvalues to have zero energy at 0.5 eV
1051                -  None: Don't shift energies, equivalent to ``e0 = 0``
1052            fact: Markers are multiplied by this factor.
1053            ax: |matplotlib-Axes| or None if a new figure should be created.
1054
1055        Returns: |matplotlib-Figure|
1056        """
1057        ax, fig, plt = get_ax_fig_plt(ax=ax)
1058
1059        gwband_range = self.min_gwbstart, self.max_gwbstop
1060        self.ebands.plot(band_range=gwband_range, e0=e0, ax=ax, show=False, **kwargs)
1061
1062        e0 = self.ebands.get_e0(e0)
1063        marker = self.get_marker(qpattr)
1064        pos, neg = marker.posneg_marker()
1065
1066        # Use different symbols depending on the value of s. Cannot use negative s.
1067        if pos:
1068            ax.scatter(pos.x, pos.y - e0, s=np.abs(pos.s) * fact, marker="^", label=qpattr + " >0")
1069        if neg:
1070            ax.scatter(neg.x, neg.y - e0, s=np.abs(neg.s) * fact, marker="v", label=qpattr + " <0")
1071
1072        return fig
1073
1074    def get_dataframe(self, ignore_imag=False):
1075        """
1076        Returns |pandas-DataFrame| with QP results for all k-points included in the GW calculation
1077
1078        Args:
1079            ignore_imag: Only real part is returned if ``ignore_imag``.
1080        """
1081        df_list = []
1082        for spin in range(self.nsppol):
1083            for gwkpoint in self.gwkpoints:
1084                df_sk = self.get_dataframe_sk(spin, gwkpoint, ignore_imag=ignore_imag)
1085                df_list.append(df_sk)
1086
1087        return pd.concat(df_list)
1088
1089    # FIXME: To maintain previous interface.
1090    to_dataframe = get_dataframe
1091
1092    def get_dataframe_sk(self, spin, kpoint, index=None, ignore_imag=False, with_params=True):
1093        """
1094        Returns |pandas-DataFrame| with QP results for the given (spin, k-point).
1095
1096        Args:
1097            ignore_imag: Only real part is returned if ``ignore_imag``.
1098            with_params: True to include convergence paramenters.
1099        """
1100        rows, bands = [], []
1101        # bstart and bstop depends on kpoint.
1102        ik_gw = self.reader.gwkpt2seqindex(kpoint)
1103        for band in range(self.gwbstart_sk[spin, ik_gw], self.gwbstop_sk[spin, ik_gw]):
1104            bands.append(band)
1105            # Build dictionary with the QP results.
1106            qpstate = self.reader.read_qp(spin, kpoint, band, ignore_imag=ignore_imag)
1107            d = qpstate.as_dict()
1108            # Add other entries that may be useful when comparing different calculations.
1109            if with_params:
1110                d.update(self.params)
1111            rows.append(d)
1112
1113        import pandas as pd
1114        index = len(bands) * [index] if index is not None else bands
1115        return pd.DataFrame(rows, index=index, columns=list(rows[0].keys()))
1116
1117    #def plot_matrix_elements(self, mel_name, spin, kpoint, *args, **kwargs):
1118    #   matrix = self.reader.read_mel(mel_name, spin, kpoint):
1119    #   return plot_matrix(matrix, *args, **kwargs)
1120
1121    #def plot_mlda_to_qps(self, spin, kpoint, *args, **kwargs):
1122    #    matrix = self.reader.read_mlda_to_qps(spin, kpoint)
1123    #    return plot_matrix(matrix, *args, **kwargs)
1124
1125    def interpolate(self, lpratio=5, ks_ebands_kpath=None, ks_ebands_kmesh=None, ks_degatol=1e-4,
1126                    vertices_names=None, line_density=20, filter_params=None, only_corrections=False, verbose=0):
1127        """
1128        Interpolate the GW corrections in k-space on a k-path and, optionally, on a k-mesh.
1129
1130        Args:
1131            lpratio: Ratio between the number of star functions and the number of ab-initio k-points.
1132                The default should be OK in many systems, larger values may be required for accurate derivatives.
1133            ks_ebands_kpath: KS |ElectronBands| on a k-path. If present,
1134                the routine interpolates the QP corrections and apply them on top of the KS band structure
1135                This is the recommended option because QP corrections are usually smoother than the
1136                QP energies and therefore easier to interpolate. If None, the QP energies are interpolated
1137                along the path defined by ``vertices_names`` and ``line_density``.
1138            ks_ebands_kmesh: KS |ElectronBands| on a homogeneous k-mesh. If present, the routine
1139                interpolates the corrections on the k-mesh (used to compute QP the DOS)
1140            ks_degatol: Energy tolerance in eV. Used when either ``ks_ebands_kpath`` or ``ks_ebands_kmesh`` are given.
1141                KS energies are assumed to be degenerate if they differ by less than this value.
1142                The interpolator may break band degeneracies (the error is usually smaller if QP corrections
1143                are interpolated instead of QP energies). This problem can be partly solved by averaging
1144                the interpolated values over the set of KS degenerate states.
1145                A negative value disables this ad-hoc symmetrization.
1146            vertices_names: Used to specify the k-path for the interpolated QP band structure
1147                when ``ks_ebands_kpath`` is None.
1148                It's a list of tuple, each tuple is of the form (kfrac_coords, kname) where
1149                kfrac_coords are the reduced coordinates of the k-point and kname is a string with the name of
1150                the k-point. Each point represents a vertex of the k-path. ``line_density`` defines
1151                the density of the sampling. If None, the k-path is automatically generated according
1152                to the point group of the system.
1153            line_density: Number of points in the smallest segment of the k-path. Used with ``vertices_names``.
1154            filter_params: TO BE DESCRIBED
1155            only_corrections: If True, the output contains the interpolated QP corrections instead of the QP energies.
1156                Available only if ks_ebands_kpath and/or ks_ebands_kmesh are used.
1157            verbose: Verbosity level
1158
1159        Returns:
1160
1161            :class:`namedtuple` with the following attributes::
1162
1163                * qp_ebands_kpath: |ElectronBands| with the QP energies interpolated along the k-path.
1164                * qp_ebands_kmesh: |ElectronBands| with the QP energies interpolated on the k-mesh.
1165                    None if ``ks_ebands_kmesh`` is not passed.
1166                * ks_ebands_kpath: |ElectronBands| with the KS energies interpolated along the k-path.
1167                * ks_ebands_kmesh: |ElectronBands| with the KS energies on the k-mesh..
1168                    None if ``ks_ebands_kmesh`` is not passed.
1169                * interpolator: |SkwInterpolator| object.
1170        """
1171        # TODO: Consistency check.
1172        errlines = []
1173        eapp = errlines.append
1174        if len(self.gwkpoints) != len(self.ibz):
1175            eapp("QP energies should be computed for all k-points in the IBZ but nkibz != nkptgw")
1176        if len(self.gwkpoints) == 1:
1177            eapp("QP Interpolation requires nkptgw > 1.")
1178        #if (np.any(self.gwbstop_sk[0, 0] != self.gwbstop_sk):
1179        #    cprint("Highest bdgw band is not constant over k-points. QP Bands will be interpolated up to...")
1180        #if (np.any(self.gwbstart_sk[0, 0] != self.gwbstart_sk):
1181        #if (np.any(self.gwbstart_sk[0, 0] != 0):
1182        if errlines:
1183            raise ValueError("\n".join(errlines))
1184
1185        # Get symmetries from abinit spacegroup (read from file).
1186        abispg = self.structure.abi_spacegroup
1187        fm_symrel = [s for (s, afm) in zip(abispg.symrel, abispg.symafm) if afm == 1]
1188
1189        if ks_ebands_kpath is None:
1190            # Generate k-points for interpolation. Will interpolate all bands available in the sigres file.
1191            bstart, bstop = 0, -1
1192            if vertices_names is None:
1193                vertices_names = [(k.frac_coords, k.name) for k in self.structure.hsym_kpoints]
1194            kpath = Kpath.from_vertices_and_names(self.structure, vertices_names, line_density=line_density)
1195            kfrac_coords, knames = kpath.frac_coords, kpath.names
1196
1197        else:
1198            # Use list of k-points from ks_ebands_kpath.
1199            ks_ebands_kpath = ElectronBands.as_ebands(ks_ebands_kpath)
1200            kfrac_coords = [k.frac_coords for k in ks_ebands_kpath.kpoints]
1201            knames = [k.name for k in ks_ebands_kpath.kpoints]
1202
1203            # Find the band range for the interpolation.
1204            bstart, bstop = 0, ks_ebands_kpath.nband
1205            bstop = min(bstop, self.min_gwbstop)
1206            if ks_ebands_kpath.nband < self.min_gwbstop:
1207                cprint("Number of bands in KS band structure smaller than the number of bands in GW corrections", "red")
1208                cprint("Highest GW bands will be ignored", "red")
1209
1210        # Interpolate QP energies if ks_ebands_kpath is None else interpolate QP corrections
1211        # and re-apply them on top of the KS band structure.
1212        gw_kcoords = [k.frac_coords for k in self.gwkpoints]
1213
1214        # Read GW energies from file (real part) and compute corrections if ks_ebands_kpath.
1215        egw_rarr = self.reader.read_value("egw", cmode="c").real
1216        if ks_ebands_kpath is not None:
1217            if ks_ebands_kpath.structure != self.structure:
1218                cprint("sigres.structure and ks_ebands_kpath.structures differ. Check your files!", "red")
1219            egw_rarr -= self.reader.read_value("e0")
1220
1221        # Note there's no guarantee that the gwkpoints and the corrections have the same k-point index.
1222        # Be careful because the order of the k-points and the band range stored in the SIGRES file may differ ...
1223        qpdata = np.empty(egw_rarr.shape)
1224        for gwk in self.gwkpoints:
1225            ik_ibz = self.reader.kpt2fileindex(gwk)
1226            for spin in range(self.nsppol):
1227                qpdata[spin, ik_ibz, :] = egw_rarr[spin, ik_ibz, :]
1228
1229        # Build interpolator for QP corrections.
1230        from abipy.core.skw import SkwInterpolator
1231        cell = (self.structure.lattice.matrix, self.structure.frac_coords, self.structure.atomic_numbers)
1232        qpdata = qpdata[:, :, bstart:bstop]
1233        # Old sigres files do not have kptopt.
1234        has_timrev = has_timrev_from_kptopt(self.reader.read_value("kptopt", default=1))
1235
1236        skw = SkwInterpolator(lpratio, gw_kcoords, qpdata, self.ebands.fermie, self.ebands.nelect,
1237                              cell, fm_symrel, has_timrev,
1238                              filter_params=filter_params, verbose=verbose)
1239
1240        if ks_ebands_kpath is None:
1241            # Interpolate QP energies.
1242            eigens_kpath = skw.interp_kpts(kfrac_coords).eigens
1243        else:
1244            # Interpolate QP energies corrections and add them to KS.
1245            ref_eigens = ks_ebands_kpath.eigens[:, :, bstart:bstop]
1246            qp_corrs = skw.interp_kpts_and_enforce_degs(kfrac_coords, ref_eigens, atol=ks_degatol).eigens
1247            eigens_kpath = qp_corrs if only_corrections else ref_eigens + qp_corrs
1248
1249        # Build new ebands object with k-path.
1250        kpts_kpath = Kpath(self.structure.reciprocal_lattice, kfrac_coords, weights=None, names=knames)
1251        occfacts_kpath = np.zeros(eigens_kpath.shape)
1252
1253        # Finding the new Fermi level of the interpolated bands is not trivial, in particular if metallic.
1254        # because one should first interpolate the QP bands on a mesh. Here I align the QP bands
1255        # at the HOMO of the KS bands.
1256        homos = ks_ebands_kpath.homos if ks_ebands_kpath is not None else self.ebands.homos
1257        qp_fermie = max([eigens_kpath[e.spin, e.kidx, e.band] for e in homos])
1258        #qp_fermie = self.ebands.fermie
1259        #qp_fermie = 0.0
1260
1261        qp_ebands_kpath = ElectronBands(self.structure, kpts_kpath, eigens_kpath, qp_fermie, occfacts_kpath,
1262                                        self.ebands.nelect, self.ebands.nspinor, self.ebands.nspden,
1263                                        smearing=self.ebands.smearing)
1264
1265        qp_ebands_kmesh = None
1266        if ks_ebands_kmesh is not None:
1267            # Interpolate QP corrections on the same k-mesh as the one used in the KS run.
1268            ks_ebands_kmesh = ElectronBands.as_ebands(ks_ebands_kmesh)
1269            if bstop > ks_ebands_kmesh.nband:
1270                raise ValueError("Not enough bands in ks_ebands_kmesh, found %s, minimum expected %d\n" % (
1271                    ks_ebands_kmesh.nband, bstop))
1272            if ks_ebands_kpath.structure != self.structure:
1273                cprint("sigres.structure and ks_ebands_kpath.structures differ. Check your files!", "red")
1274                #raise ValueError("sigres.structure and ks_ebands_kmesh.structures differ. Check your files!")
1275
1276            # K-points and weight for DOS are taken from ks_ebands_kmesh
1277            dos_kcoords = [k.frac_coords for k in ks_ebands_kmesh.kpoints]
1278            dos_weights = [k.weight for k in ks_ebands_kmesh.kpoints]
1279
1280            # Interpolate QP corrections from bstart to bstop
1281            ref_eigens = ks_ebands_kmesh.eigens[:, :, bstart:bstop]
1282            qp_corrs = skw.interp_kpts_and_enforce_degs(dos_kcoords, ref_eigens, atol=ks_degatol).eigens
1283            eigens_kmesh = qp_corrs if only_corrections else ref_eigens + qp_corrs
1284
1285            # Build new ebands object with k-mesh
1286            kpts_kmesh = IrredZone(self.structure.reciprocal_lattice, dos_kcoords, weights=dos_weights,
1287                                   names=None, ksampling=ks_ebands_kmesh.kpoints.ksampling)
1288            occfacts_kmesh = np.zeros(eigens_kmesh.shape)
1289            qp_ebands_kmesh = ElectronBands(self.structure, kpts_kmesh, eigens_kmesh, qp_fermie, occfacts_kmesh,
1290                                            self.ebands.nelect, self.ebands.nspinor, self.ebands.nspden,
1291                                            smearing=self.ebands.smearing)
1292
1293        return dict2namedtuple(qp_ebands_kpath=qp_ebands_kpath,
1294                               qp_ebands_kmesh=qp_ebands_kmesh,
1295                               ks_ebands_kpath=ks_ebands_kpath,
1296                               ks_ebands_kmesh=ks_ebands_kmesh,
1297                               interpolator=skw,
1298                               )
1299
1300    def yield_figs(self, **kwargs):  # pragma: no cover
1301        """
1302        This function *generates* a predefined list of matplotlib figures with minimal input from the user.
1303        Used in abiview.py to get a quick look at the results.
1304        """
1305        yield self.plot_qpgaps(show=False)
1306        yield self.plot_qps_vs_e0(show=False)
1307        yield self.plot_qpbands_ibz(show=False)
1308        yield self.plot_ksbands_with_qpmarkers(show=False)
1309        if self.has_spectral_function:
1310            yield self.plot_spectral_functions(include_bands=None, show=False)
1311
1312    def write_notebook(self, nbpath=None):
1313        """
1314        Write a jupyter_ notebook to ``nbpath``. If nbpath is None, a temporay file in the current
1315        working directory is created. Return path to the notebook.
1316        """
1317        nbformat, nbv, nb = self.get_nbformat_nbv_nb(title=None)
1318
1319        nb.cells.extend([
1320            nbv.new_code_cell("sigres = abilab.abiopen('%s')" % self.filepath),
1321            nbv.new_code_cell("print(sigres)"),
1322            nbv.new_code_cell("sigres.plot_qps_vs_e0();"),
1323            nbv.new_code_cell("sigres.plot_spectral_functions(spin=0, kpoint=[0, 0, 0], bands=0);"),
1324            nbv.new_code_cell("#sigres.plot_ksbands_with_qpmarkers(qpattr='qpeme0', fact=100);"),
1325            nbv.new_code_cell("r = sigres.interpolate(ks_ebands_kpath=None, ks_ebands_kmesh=None); print(r.interpolator)"),
1326            nbv.new_code_cell("r.qp_ebands_kpath.plot();"),
1327            nbv.new_code_cell("""
1328if r.ks_ebands_kpath is not None:
1329    plotter = abilab.ElectronBandsPlotter()
1330    plotter.add_ebands("KS", r.ks_ebands_kpath) # dos=r.ks_ebands_kmesh.get_edos())
1331    plotter.add_ebands("GW (interpolated)", r.qp_ebands_kpath) # dos=r.qp_ebands_kmesh.get_edos()))
1332    plotter.ipw_select_plot()"""),
1333        ])
1334
1335        return self._write_nb_nbpath(nb, nbpath)
1336
1337
1338class SigresReader(ETSF_Reader):
1339    r"""
1340    This object provides method to read data from the SIGRES file produced ABINIT.
1341
1342    .. rubric:: Inheritance Diagram
1343    .. inheritance-diagram:: SigresReader
1344
1345    # See 70gw/m_sigma_results.F90
1346
1347    # Name of the diagonal matrix elements stored in the file.
1348    # b1gw:b2gw,nkibz,nsppol*nsig_ab))
1349    #_DIAGO_MELS = [
1350    #    "sigxme",
1351    #    "vxcme",
1352    #    "vUme",
1353    #    "dsigmee0",
1354    #    "sigcmee0",
1355    #    "sigxcme",
1356    #    "ze0",
1357    #]
1358
1359    integer :: b1gw,b2gw      ! min and Max gw band indeces over spin and k-points (used to dimension)
1360    integer :: gwcalctyp      ! Flag defining the calculation type.
1361    integer :: nkptgw         ! No. of points calculated
1362    integer :: nkibz          ! No. of irreducible k-points.
1363    integer :: nbnds          ! Total number of bands
1364    integer :: nomega_r       ! No. of real frequencies for the spectral function.
1365    integer :: nomega_i       ! No. of frequencies along the imaginary axis.
1366    integer :: nomega4sd      ! No. of real frequencies to evaluate the derivative of $\Sigma(E)$.
1367    integer :: nsig_ab        ! 1 if nspinor=1,4 for noncollinear case.
1368    integer :: nsppol         ! No. of spin polarizations.
1369    integer :: usepawu        ! 1 if we are using LDA+U as starting point (only for PAW)
1370
1371    real(dp) :: deltae       ! Frequency step for the calculation of d\Sigma/dE
1372    real(dp) :: maxomega4sd  ! Max frequency around E_ks for d\Sigma/dE.
1373    real(dp) :: maxomega_r   ! Max frequency for spectral function.
1374    real(dp) :: scissor_ene  ! Scissor energy value. zero for None.
1375
1376    integer,pointer :: maxbnd(:,:)
1377    ! maxbnd(nkptgw,nsppol)
1378    ! Max band index considered in GW for this k-point.
1379
1380    integer,pointer :: minbnd(:,:)
1381    ! minbnd(nkptgw,nsppol)
1382    ! Min band index considered in GW for this k-point.
1383
1384    real(dp),pointer :: degwgap(:,:)
1385    ! degwgap(nkibz,nsppol)
1386    ! Difference btw the QPState and the KS optical gap.
1387
1388    real(dp),pointer :: egwgap(:,:)
1389    ! egwgap(nkibz,nsppol))
1390    ! QPState optical gap at each k-point and spin.
1391
1392    real(dp),pointer :: en_qp_diago(:,:,:)
1393    ! en_qp_diago(nbnds,nkibz,nsppol))
1394    ! QPState energies obtained from the diagonalization of the Hermitian approximation to Sigma (QPSCGW)
1395
1396    real(dp),pointer :: e0(:,:,:)
1397    ! e0(nbnds,nkibz,nsppol)
1398    ! KS eigenvalues for each band, k-point and spin. In case of self-consistent?
1399
1400    real(dp),pointer :: e0gap(:,:)
1401    ! e0gap(nkibz,nsppol),
1402    ! KS gap at each k-point, for each spin.
1403
1404    real(dp),pointer :: omega_r(:)
1405    ! omega_r(nomega_r)
1406    ! real frequencies used for the self energy.
1407
1408    real(dp),pointer :: kptgw(:,:)
1409    ! kptgw(3,nkptgw)
1410    ! ! TODO there is a similar array in sigma_parameters
1411    ! List of calculated k-points.
1412
1413    real(dp),pointer :: sigxme(:,:,:)
1414    ! sigxme(b1gw:b2gw,nkibz,nsppol*nsig_ab))
1415    ! Diagonal matrix elements of $\Sigma_x$ i.e $\<nks|\Sigma_x|nks\>$
1416
1417    real(dp),pointer :: vxcme(:,:,:)
1418    ! vxcme(b1gw:b2gw,nkibz,nsppol*nsig_ab))
1419    ! $\<nks|v_{xc}[n_val]|nks\>$ matrix elements of vxc (valence-only contribution).
1420
1421    real(dp),pointer :: vUme(:,:,:)
1422    ! vUme(b1gw:b2gw,nkibz,nsppol*nsig_ab))
1423    ! $\<nks|v_{U}|nks\>$ for LDA+U.
1424
1425    complex(dpc),pointer :: degw(:,:,:)
1426    ! degw(b1gw:b2gw,nkibz,nsppol))
1427    ! Difference between the QPState and the KS energies.
1428
1429    complex(dpc),pointer :: dsigmee0(:,:,:)
1430    ! dsigmee0(b1gw:b2gw,nkibz,nsppol*nsig_ab))
1431    ! Derivative of $\Sigma_c(E)$ calculated at the KS eigenvalue.
1432
1433    complex(dpc),pointer :: egw(:,:,:)
1434    ! egw(nbnds,nkibz,nsppol))
1435    ! QPState energies, $\epsilon_{nks}^{QPState}$.
1436
1437    complex(dpc),pointer :: eigvec_qp(:,:,:,:)
1438    ! eigvec_qp(nbnds,nbnds,nkibz,nsppol))
1439    ! Expansion of the QPState amplitude in the KS basis set.
1440
1441    complex(dpc),pointer :: hhartree(:,:,:,:)
1442    ! hhartree(b1gw:b2gw,b1gw:b2gw,nkibz,nsppol*nsig_ab)
1443    ! $\<nks|T+v_H+v_{loc}+v_{nl}|mks\>$
1444
1445    complex(dpc),pointer :: sigcme(:,:,:,:)
1446    ! sigcme(b1gw:b2gw,nkibz,nomega_r,nsppol*nsig_ab))
1447    ! $\<nks|\Sigma_{c}(E)|nks\>$ at each nomega_r frequency
1448
1449    complex(dpc),pointer :: sigmee(:,:,:)
1450    ! sigmee(b1gw:b2gw,nkibz,nsppol*nsig_ab))
1451    ! $\Sigma_{xc}E_{KS} + (E_{QPState}- E_{KS})*dSigma/dE_KS
1452
1453    complex(dpc),pointer :: sigcmee0(:,:,:)
1454    ! sigcmee0(b1gw:b2gw,nkibz,nsppol*nsig_ab))
1455    ! Diagonal mat. elements of $\Sigma_c(E)$ calculated at the KS energy $E_{KS}$
1456
1457    complex(dpc),pointer :: sigcmesi(:,:,:,:)
1458    ! sigcmesi(b1gw:b2gw,nkibz,nomega_i,nsppol*nsig_ab))
1459    ! Matrix elements of $\Sigma_c$ along the imaginary axis.
1460    ! Only used in case of analytical continuation.
1461
1462    complex(dpc),pointer :: sigcme4sd(:,:,:,:)
1463    ! sigcme4sd(b1gw:b2gw,nkibz,nomega4sd,nsppol*nsig_ab))
1464    ! Diagonal matrix elements of \Sigma_c around the zeroth order eigenvalue (usually KS).
1465
1466    complex(dpc),pointer :: sigxcme(:,:,:,:)
1467    ! sigxme(b1gw:b2gw,nkibz,nomega_r,nsppol*nsig_ab))
1468    ! $\<nks|\Sigma_{xc}(E)|nks\>$ at each real frequency frequency.
1469
1470    complex(dpc),pointer :: sigxcmesi(:,:,:,:)
1471    ! sigxcmesi(b1gw:b2gw,nkibz,nomega_i,nsppol*nsig_ab))
1472    ! Matrix elements of $\Sigma_{xc}$ along the imaginary axis.
1473    ! Only used in case of analytical continuation.
1474
1475    complex(dpc),pointer :: sigxcme4sd(:,:,:,:)
1476    ! sigxcme4sd(b1gw:b2gw,nkibz,nomega4sd,nsppol*nsig_ab))
1477    ! Diagonal matrix elements of \Sigma_xc for frequencies around the zeroth order eigenvalues.
1478
1479    complex(dpc),pointer :: ze0(:,:,:)
1480    ! ze0(b1gw:b2gw,nkibz,nsppol))
1481    ! renormalization factor. $(1-\dfrac{\partial\Sigma_c} {\partial E_{KS}})^{-1}$
1482
1483    complex(dpc),pointer :: omega_i(:)
1484    ! omegasi(nomega_i)
1485    ! Frequencies along the imaginary axis used for the analytical continuation.
1486
1487    complex(dpc),pointer :: omega4sd(:,:,:,:)
1488    ! omega4sd(b1gw:b2gw,nkibz,nomega4sd,nsppol).
1489    ! Frequencies used to evaluate the Derivative of Sigma.
1490    """
1491    def __init__(self, path):
1492        self.ks_bands = ElectronBands.from_file(path)
1493        self.nsppol = self.ks_bands.nsppol
1494        super().__init__(path)
1495
1496        try:
1497            self.nomega_r = self.read_dimvalue("nomega_r")
1498        except self.Error:
1499            self.nomega_r = 0
1500
1501        #self.nomega_i = self.read_dim("nomega_i")
1502
1503        # Save important quantities needed to simplify the API.
1504        self.structure = self.read_structure()
1505
1506        self.gwcalctyp = self.read_value("gwcalctyp")
1507        self.usepawu = self.read_value("usepawu")
1508
1509        # 1) The K-points of the homogeneous mesh.
1510        self.ibz = self.ks_bands.kpoints
1511
1512        # 2) The K-points where QPState corrections have been calculated.
1513        gwred_coords = self.read_redc_gwkpoints()
1514        self.gwkpoints = KpointList(self.structure.reciprocal_lattice, gwred_coords)
1515        # Find k-point name
1516        for kpoint in self.gwkpoints:
1517            kpoint.set_name(self.structure.findname_in_hsym_stars(kpoint))
1518
1519        # minbnd[nkptgw,nsppol] gives the minimum band index computed
1520        # Note conversion between Fortran and python convention.
1521        self.gwbstart_sk = self.read_value("minbnd") - 1
1522        self.gwbstop_sk = self.read_value("maxbnd")
1523        # min and Max band index for GW corrections.
1524        self.min_gwbstart = np.min(self.gwbstart_sk)
1525        self.max_gwbstart = np.max(self.gwbstart_sk)
1526
1527        self.min_gwbstop = np.min(self.gwbstop_sk)
1528        self.max_gwbstop = np.max(self.gwbstop_sk)
1529
1530        self._egw = self.read_value("egw", cmode="c")
1531
1532        # Read and save important matrix elements.
1533        # All these arrays are dimensioned
1534        # vxcme(b1gw:b2gw,nkibz,nsppol*nsig_ab))
1535        self._vxcme = self.read_value("vxcme")
1536        self._sigxme = self.read_value("sigxme")
1537        self._hhartree = self.read_value("hhartree", cmode="c")
1538        self._vUme = self.read_value("vUme")
1539        #if self.usepawu == 0: self._vUme.fill(0.0)
1540
1541        # Complex arrays
1542        self._sigcmee0 = self.read_value("sigcmee0", cmode="c")
1543        self._ze0 = self.read_value("ze0", cmode="c")
1544
1545        # Frequencies for the spectral function.
1546        # Note that omega_r does not depend on (s, k, b).
1547        if self.has_spfunc:
1548            self._omega_r = self.read_value("omega_r")
1549            self._sigcme = self.read_value("sigcme", cmode="c")
1550            self._sigxcme = self.read_value("sigxcme", cmode="c")
1551
1552        # Self-consistent case
1553        self._en_qp_diago = self.read_value("en_qp_diago")
1554        # <KS|QPState>
1555        self._eigvec_qp = self.read_value("eigvec_qp", cmode="c")
1556
1557        #self._mlda_to_qp
1558
1559    #def is_selfconsistent(self, mode):
1560    #    return self.gwcalctyp
1561
1562    @property
1563    def has_spfunc(self):
1564        """True if self contains the spectral function."""
1565        return self.nomega_r
1566
1567    def kpt2fileindex(self, kpoint):
1568        """
1569        Helper function that returns the index of kpoint in the netcdf file.
1570        Accepts |Kpoint| instance or integer
1571
1572        Raise:
1573            `KpointsError` if kpoint cannot be found.
1574
1575        .. note::
1576
1577            This function is needed since arrays in the netcdf file are dimensioned
1578            with the total number of k-points in the IBZ.
1579        """
1580        if duck.is_intlike(kpoint): return int(kpoint)
1581        return self.ibz.index(kpoint)
1582
1583    def gwkpt2seqindex(self, gwkpoint):
1584        """
1585        This function returns the index of the GW k-point in (0:nkptgw)
1586        Used to access data in the arrays that are dimensioned [0:nkptgw] e.g. minbnd.
1587        """
1588        if duck.is_intlike(gwkpoint):
1589            return int(gwkpoint)
1590        else:
1591            return self.gwkpoints.index(gwkpoint)
1592
1593    def read_redc_gwkpoints(self):
1594        return self.read_value("kptgw")
1595
1596    def read_allqps(self, ignore_imag=False):
1597        """
1598        Return list with ``nsppol`` items. Each item is a :class:`QPList` with the QP results
1599
1600        Args:
1601            ignore_imag: Only real part is returned if ``ignore_imag``.
1602        """
1603        qps_spin = self.nsppol * [None]
1604
1605        for spin in range(self.nsppol):
1606            qps = []
1607            for gwkpoint in self.gwkpoints:
1608                ik = self.gwkpt2seqindex(gwkpoint)
1609                for band in range(self.gwbstart_sk[spin, ik], self.gwbstop_sk[spin, ik]):
1610                    qps.append(self.read_qp(spin, gwkpoint, band, ignore_imag=ignore_imag))
1611
1612            qps_spin[spin] = QPList(qps)
1613
1614        return tuple(qps_spin)
1615
1616    def read_qplist_sk(self, spin, kpoint, ignore_imag=False):
1617        """
1618        Read and return :class:`QPList` object for the given spin, kpoint.
1619
1620        Args:
1621            ignore_imag: Only real part is returned if ``ignore_imag``.
1622        """
1623        ik = self.gwkpt2seqindex(kpoint)
1624        bstart, bstop = self.gwbstart_sk[spin, ik], self.gwbstop_sk[spin, ik]
1625
1626        return QPList([self.read_qp(spin, kpoint, band, ignore_imag=ignore_imag)
1627                      for band in range(bstart, bstop)])
1628
1629    #def read_qpene(self, spin, kpoint, band)
1630
1631    def read_qpenes(self):
1632        return self._egw[:, :, :]
1633
1634    def read_qp(self, spin, kpoint, band, ignore_imag=False):
1635        """
1636        Return :class`QPState` for the given (spin, kpoint, band).
1637        Only real part is returned if ``ignore_imag``.
1638        """
1639        ik_file = self.kpt2fileindex(kpoint)
1640        # Must shift band index (see fortran code that allocates with mdbgw)
1641        ib_gw = band - self.min_gwbstart
1642
1643        def ri(a):
1644            return np.real(a) if ignore_imag else a
1645
1646        return QPState(
1647            spin=spin,
1648            kpoint=kpoint,
1649            band=band,
1650            e0=self.read_e0(spin, ik_file, band),
1651            qpe=ri(self._egw[spin, ik_file, band]),
1652            qpe_diago=ri(self._en_qp_diago[spin, ik_file, band]),
1653            # Note ib_gw index.
1654            vxcme=self._vxcme[spin, ik_file, ib_gw],
1655            sigxme=self._sigxme[spin, ik_file, ib_gw],
1656            sigcmee0=ri(self._sigcmee0[spin, ik_file, ib_gw]),
1657            vUme=self._vUme[spin, ik_file, ib_gw],
1658            ze0=ri(self._ze0[spin, ik_file, ib_gw]),
1659        )
1660
1661    def read_qpgaps(self):
1662        """Read the QP gaps. Returns [nsppol, nkibz] array with QP gaps in eV."""
1663        return self.read_value("egwgap")
1664
1665    def read_ksgaps(self):
1666        """Read the KS gaps. Returns [nsppol, nkibz] array with KS gaps in eV."""
1667        return self.read_value("e0gap")
1668
1669    def read_e0(self, spin, kfile, band):
1670        return self.ks_bands.eigens[spin, kfile, band]
1671
1672    def read_sigmaw(self, spin, kpoint, band):
1673        """Returns the real and the imaginary part of the self energy."""
1674        if not self.has_spfunc:
1675            raise ValueError("%s does not contain spectral function data." % self.path)
1676
1677        ik = self.kpt2fileindex(kpoint)
1678        # Must shift band index (see fortran code that allocates with mdbgw)
1679        ib_gw = band - self.min_gwbstart
1680        #ib_gw = band - self.gwbstart_sk[spin, self.gwkpt2seqindex(kpoint)]
1681
1682        return self._omega_r, self._sigxcme[spin,:, ik, ib_gw]
1683
1684    def read_spfunc(self, spin, kpoint, band):
1685        """
1686        Returns the spectral function.
1687
1688         one/pi * ABS(AIMAG(Sr%sigcme(ib,ikibz,io,is))) /
1689         ( (REAL(Sr%omega_r(io)-Sr%hhartree(ib,ib,ikibz,is)-Sr%sigxcme(ib,ikibz,io,is)))**2 &
1690        +(AIMAG(Sr%sigcme(ib,ikibz,io,is)))**2) / Ha_eV,&
1691        """
1692        if not self.has_spfunc:
1693            raise ValueError("%s does not contain spectral function data" % self.path)
1694
1695        ik = self.kpt2fileindex(kpoint)
1696        # Must shift band index (see fortran code that allocates with mdbgw)
1697        ib_gw = band - self.min_gwbstart
1698        #ib_gw = band - self.gwbstart_sk[spin, self.gwkpt2seqindex(kpoint)]
1699
1700        aim_sigc = np.abs(self._sigcme[spin,:,ik,ib_gw].imag)
1701        den = np.zeros(self.nomega_r)
1702
1703        for io, omega in enumerate(self._omega_r):
1704            den[io] = (omega - self._hhartree[spin,ik,ib_gw,ib_gw].real - self._sigxcme[spin,io,ik,ib_gw].real) ** 2 + \
1705                self._sigcme[spin,io,ik,ib_gw].imag ** 2
1706
1707        return self._omega_r, 1./np.pi * (aim_sigc/den)
1708
1709    def read_eigvec_qp(self, spin, kpoint, band=None):
1710        """
1711        Returns <KS|QPState> for the given spin, kpoint and band.
1712        If band is None, <KS_b|QP_{b'}> is returned.
1713        """
1714        ik = self.kpt2fileindex(kpoint)
1715        if band is not None:
1716            return self._eigvec_qp[spin, ik, :, band]
1717        else:
1718            return self._eigvec_qp[spin, ik, :, :]
1719
1720    def read_params(self):
1721        """
1722        Read the parameters of the calculation.
1723        Returns: OrderedDict with the value of the parameters.
1724        """
1725        param_names = [
1726            "ecutwfn", "ecuteps", "ecutsigx", "scr_nband", "sigma_nband",
1727            "gwcalctyp", "scissor_ene",
1728        ]
1729
1730        # Read data and convert to scalar to avoid problems with pandas dataframes.
1731        # Old sigres files may not have all the metadata.
1732        params = OrderedDict()
1733        for pname in param_names:
1734            v = self.read_value(pname, default=None)
1735            params[pname] = v if v is None else np.asarray(v).item()
1736
1737        # Other quantities that might be subject to convergence studies.
1738        #params["nkibz"] = len(self.ibz)
1739
1740        return params
1741
1742    #def read_mlda_to_qp(self, spin, kpoint, band=None):
1743    #    """Returns the unitary transformation KS --> QPS"""
1744    #    ik = self.kpt2fileindex(kpoint)
1745    #    if band is not None:
1746    #        return self._mlda_to_qp[spin,ik,:,band]
1747    #    else:
1748    #        return self._mlda_to_qp[spin,ik,:,:]
1749
1750    #def read_qprhor(self):
1751    #    """Returns the QPState density in real space."""
1752
1753
1754class SigresRobot(Robot, RobotWithEbands):
1755    """
1756    This robot analyzes the results contained in multiple SIGRES.nc files.
1757
1758    .. rubric:: Inheritance Diagram
1759    .. inheritance-diagram:: SigresRobot
1760    """
1761    # Try to have API similar to SigEPhRobot
1762    EXT = "SIGRES"
1763
1764    def __init__(self, *args):
1765        super().__init__(*args)
1766        if len(self.abifiles) in (0, 1): return
1767
1768        # TODO
1769        # Check dimensions and self-energy states and issue warning.
1770        warns = []; wapp = warns.append
1771        nc0 = self.abifiles[0]
1772        same_nsppol, same_nkcalc = True, True
1773        if any(nc.nsppol != nc0.nsppol for nc in self.abifiles):
1774            same_nsppol = False
1775            wapp("Comparing ncfiles with different values of nsppol.")
1776        if any(nc.nkcalc != nc0.nkcalc for nc in self.abifiles):
1777            same_nkcalc = False
1778            wapp("Comparing ncfiles with different number of k-points in self-energy. Doh!")
1779
1780        if same_nsppol and same_nkcalc:
1781            # FIXME
1782            # Different values of bstart_ks are difficult to handle
1783            # Because the high-level API assumes an absolute global index
1784            # Should decide how to treat this case: either raise or interpret band as an absolute band index.
1785            if any(np.any(nc.gwbstart_sk != nc0.gwbstart_sk) for nc in self.abifiles):
1786                wapp("Comparing ncfiles with different values of gwbstart_sk")
1787            if any(np.any(nc.gwbstop_sk != nc0.gwbstop_sk) for nc in self.abifiles):
1788                wapp("Comparing ncfiles with different values of gwbstop_sk")
1789
1790        if warns:
1791            for w in warns:
1792                cprint(w, color="yellow")
1793
1794    def _check_dims_and_params(self):
1795        """Test that nsppol, sigma_kpoints, tlist are consistent."""
1796        if not len(self.abifiles) > 1:
1797            return 0
1798
1799        nc0 = self.abifiles[0]
1800        errors = []
1801        eapp = errors.append
1802
1803        if any(nc.nsppol != nc0.nsppol for nc in self.abifiles[1:]):
1804            eapp("Files with different values of `nsppol`")
1805
1806        if any(nc.nkcalc != nc0.nkcalc for nc in self.abifiles[1:]):
1807            eapp("Files with different values of `nkcalc`")
1808        else:
1809            for nc in self.abifiles[1:]:
1810                for k0, k1 in zip(nc0.sigma_kpoints, nc.sigma_kpoints):
1811                    if k0 != k1:
1812                        eapp("Files with different values of `sigma_kpoints`")
1813
1814        if errors:
1815            raise ValueError("Cannot compare multiple SIGRES.nc files. Reason:\n %s" % "\n".join(errors))
1816
1817    def merge_dataframes_sk(self, spin, kpoint, **kwargs):
1818        for i, (label, sigr) in enumerate(self.items()):
1819            frame = sigr.get_dataframe_sk(spin, kpoint, index=label)
1820            if i == 0:
1821                table = frame
1822            else:
1823                table = table.append(frame)
1824
1825        return table
1826
1827    def get_qpgaps_dataframe(self, spin=None, kpoint=None, with_geo=False, abspath=False, funcs=None, **kwargs):
1828        """
1829        Return a |pandas-DataFrame| with the QP gaps for all files in the robot.
1830
1831        Args:
1832            spin: Spin index.
1833            kpoint
1834            with_geo: True if structure info should be added to the dataframe
1835            abspath: True if paths in index should be absolute. Default: Relative to getcwd().
1836            funcs: Function or list of functions to execute to add more data to the DataFrame.
1837                Each function receives a |SigresFile| object and returns a tuple (key, value)
1838                where key is a string with the name of column and value is the value to be inserted.
1839        """
1840        # TODO: Ideally one should select the k-point for which we have the fundamental gap for the given spin
1841        # TODO: In principle the SIGRES might have different k-points
1842        if spin is None: spin = 0
1843        if kpoint is None: kpoint = 0
1844
1845        attrs = [
1846            "nsppol",
1847            #"nspinor", "nspden", #"ecut", "pawecutdg",
1848            #"tsmear", "nkibz",
1849        ] + kwargs.pop("attrs", [])
1850
1851        rows, row_names = [], []
1852        for label, sigres in self.items():
1853            row_names.append(label)
1854            d = OrderedDict()
1855            for aname in attrs:
1856                d[aname] = getattr(sigres, aname, None)
1857
1858            qpgap = sigres.get_qpgap(spin, kpoint)
1859            d.update({"qpgap": qpgap})
1860
1861            # Add convergence parameters
1862            d.update(sigres.params)
1863
1864            # Add info on structure.
1865            if with_geo:
1866                d.update(sigres.structure.get_dict4pandas(with_spglib=True))
1867
1868            # Execute functions.
1869            if funcs is not None: d.update(self._exec_funcs(funcs, sigres))
1870            rows.append(d)
1871
1872        row_names = row_names if not abspath else self._to_relpaths(row_names)
1873        return pd.DataFrame(rows, index=row_names, columns=list(rows[0].keys()))
1874
1875    # An alias to have a common API for robots.
1876    get_dataframe = get_qpgaps_dataframe
1877
1878    @add_fig_kwargs
1879    def plot_qpgaps_convergence(self, plot_qpmks=True, sortby=None, hue=None, sharey=False, fontsize=8, **kwargs):
1880        """
1881        Plot the convergence of the direct QP gaps for all the k-points available in the robot.
1882
1883        Args:
1884            plot_qpmks: If False, plot QP_gap, KS_gap else (QP_gap - KS_gap)
1885            sortby: Define the convergence parameter, sort files and produce plot labels.
1886                Can be None, string or function. If None, no sorting is performed.
1887                If string and not empty it's assumed that the abifile has an attribute
1888                with the same name and `getattr` is invoked.
1889                If callable, the output of sortby(abifile) is used.
1890            hue: Variable that define subsets of the data, which will be drawn on separate lines.
1891                Accepts callable or string
1892                If string, it's assumed that the abifile has an attribute with the same name and getattr is invoked.
1893                If callable, the output of hue(abifile) is used.
1894            sharey: True if y-axis should be shared.
1895            fontsize: legend and label fontsize.
1896
1897        Returns: |matplotlib-Figure|
1898        """
1899        # Make sure that nsppol, sigma_kpoints are consistent.
1900        self._check_dims_and_params()
1901
1902        nc0 = self.abifiles[0]
1903        nsppol, sigma_kpoints = nc0.nsppol, nc0.sigma_kpoints
1904        # Build grid with (nkpt, 1) plots.
1905        ncols, nrows = 1, len(sigma_kpoints)
1906        ax_list, fig, plt = get_axarray_fig_plt(None, nrows=nrows, ncols=ncols,
1907                                                sharex=True, sharey=sharey, squeeze=False)
1908        ax_list = ax_list.ravel()
1909
1910        if hue is None:
1911            labels, ncfiles, params = self.sortby(sortby, unpack=True)
1912        else:
1913            groups = self.group_and_sortby(hue, sortby)
1914
1915        for ik, (kgw, ax) in enumerate(zip(sigma_kpoints, ax_list)):
1916            for spin in range(nsppol):
1917                ax.set_title("QP dirgap k:%s" % (repr(kgw)), fontsize=fontsize)
1918
1919                # Extract QP dirgap for [spin, kgw, itemp]
1920                if hue is None:
1921                    qp_gaps, ks_gaps = map(np.array, zip(*[ncfile.get_qpgap(spin, kgw, with_ksgap=True)
1922                        for ncfile in ncfiles]))
1923                    yvals = qp_gaps if not plot_qpmks else qp_gaps - ks_gaps
1924
1925                    if not is_string(params[0]):
1926                        ax.plot(params, yvals, marker=nc0.marker_spin[spin])
1927                    else:
1928                        # Must handle list of strings in a different way.
1929                        xn = range(len(params))
1930                        ax.plot(xn, yvals, marker=nc0.marker_spin[spin])
1931                        ax.set_xticks(xn)
1932                        ax.set_xticklabels(params, fontsize=fontsize)
1933                else:
1934                    for g in groups:
1935                        qp_gaps, ks_gaps = map(np.array, zip(*[ncfile.get_qpgap(spin, kgw, with_ksgap=True)
1936                            for ncfile in g.abifiles]))
1937                        yvals = qp_gaps if not plot_qpmks else qp_gaps - ks_gaps
1938                        label = "%s: %s" % (self._get_label(hue), g.hvalue)
1939                        ax.plot(g.xvalues, qp_gaps, marker=nc0.marker_spin[spin], label=label)
1940
1941            ax.grid(True)
1942            if ik == len(sigma_kpoints) - 1:
1943                ax.set_xlabel("%s" % self._get_label(sortby))
1944                if sortby is None: rotate_ticklabels(ax, 15)
1945            if ik == 0:
1946                if plot_qpmks:
1947                    ax.set_ylabel("QP-KS direct gap (eV)", fontsize=fontsize)
1948                else:
1949                    ax.set_ylabel("QP direct gap (eV)", fontsize=fontsize)
1950
1951            if hue is not None:
1952                ax.legend(loc="best", fontsize=fontsize, shadow=True)
1953
1954        return fig
1955
1956    @add_fig_kwargs
1957    def plot_qpdata_conv_skb(self, spin, kpoint, band, sortby=None, hue=None,
1958                            fontsize=8, **kwargs):
1959        """
1960        For each file in the robot, plot the convergence of the QP results for given (spin, kpoint, band)
1961
1962        Args:
1963            spin: Spin index.
1964            kpoint: K-point in self-energy. Accepts |Kpoint|, vector or index.
1965            band: Band index.
1966            sortby: Define the convergence parameter, sort files and produce plot labels.
1967                Can be None, string or function. If None, no sorting is performed.
1968                If string and not empty it's assumed that the abifile has an attribute
1969                with the same name and `getattr` is invoked.
1970                If callable, the output of sortby(abifile) is used.
1971            hue: Variable that define subsets of the data, which will be drawn on separate lines.
1972                Accepts callable or string
1973                If string, it's assumed that the abifile has an attribute with the same name and getattr is invoked.
1974                If callable, the output of hue(abifile) is used.
1975            what_list: List of strings selecting the quantity to plot.
1976            fontsize: legend and label fontsize.
1977
1978        Returns: |matplotlib-Figure|
1979        """
1980        # Make sure that nsppol and sigma_kpoints are consistent
1981        self._check_dims_and_params()
1982
1983        # TODO: Add more quantities DW, Fan(0)
1984        # TODO: Decide how to treat complex quantities, avoid annoying ComplexWarning
1985        # TODO: Format for g.hvalue
1986        # Treat fundamental gaps
1987        # Quantities to plot.
1988        what_list = ["re_qpe", "imag_qpe", "ze0"]
1989
1990        # Build grid plot.
1991        nrows, ncols = len(what_list), 1
1992        ax_list, fig, plt = get_axarray_fig_plt(None, nrows=nrows, ncols=ncols,
1993                                                sharex=True, sharey=False, squeeze=False)
1994        ax_list = ax_list.ravel()
1995
1996        nc0 = self.abifiles[0]
1997        ik = nc0.reader.gwkpt2seqindex(kpoint)
1998        kpoint = nc0.sigma_kpoints[ik]
1999
2000        # Sort and read QP data.
2001        if hue is None:
2002            labels, ncfiles, params = self.sortby(sortby, unpack=True)
2003            qplist = [ncfile.reader.read_qp(spin, kpoint, band) for ncfile in ncfiles]
2004        else:
2005            groups = self.group_and_sortby(hue, sortby)
2006            qplist_group = []
2007            for g in groups:
2008                lst = [ncfile.reader.read_qp(spin, kpoint, band) for ncfile in g.abifiles]
2009                qplist_group.append(lst)
2010
2011        for i, (ax, what) in enumerate(zip(ax_list, what_list)):
2012            if hue is None:
2013                # Extract QP data.
2014                yvals = [getattr(qp, what) for qp in qplist]
2015                if not is_string(params[0]):
2016                    ax.plot(params, yvals, marker=nc0.marker_spin[spin])
2017                else:
2018                    # Must handle list of strings in a different way.
2019                    xn = range(len(params))
2020                    ax.plot(xn, yvals, marker=nc0.marker_spin[spin])
2021                    ax.set_xticks(xn)
2022                    ax.set_xticklabels(params, fontsize=fontsize)
2023            else:
2024                for g, qplist in zip(groups, qplist_group):
2025                    # Extract QP data.
2026                    yvals = [getattr(qp, what) for qp in qplist]
2027                    label = "%s: %s" % (self._get_label(hue), g.hvalue)
2028                    ax.plot(g.xvalues, yvals, marker=nc0.marker_spin[spin], label=label)
2029
2030            ax.grid(True)
2031            ax.set_ylabel(what)
2032            if i == len(what_list) - 1:
2033                ax.set_xlabel("%s" % self._get_label(sortby))
2034                if sortby is None: rotate_ticklabels(ax, 15)
2035            if i == 0 and hue is not None:
2036                ax.legend(loc="best", fontsize=fontsize, shadow=True)
2037
2038        if "title" not in kwargs:
2039            title = "QP results spin: %s, k:%s, band: %s" % (spin, repr(kpoint), band)
2040            fig.suptitle(title, fontsize=fontsize)
2041
2042        return fig
2043
2044    @add_fig_kwargs
2045    def plot_qpfield_vs_e0(self, field, sortby=None, hue=None, fontsize=8,
2046                           sharey=False, colormap="jet", e0="fermie", **kwargs):
2047        """
2048        For each file in the robot, plot one of the attributes of :class:`QpState`
2049        as a function of the KS energy.
2050
2051        Args:
2052            field (str): String defining the attribute to plot.
2053            sharey: True if y-axis should be shared.
2054
2055        .. note::
2056
2057            For the meaning of the other arguments, see other robot methods.
2058
2059        Returns: |matplotlib-Figure|
2060        """
2061        import matplotlib.pyplot as plt
2062        cmap = plt.get_cmap(colormap)
2063
2064        if hue is None:
2065            ax_list = None
2066            lnp_list = self.sortby(sortby)
2067            for i, (label, ncfile, param) in enumerate(lnp_list):
2068                if sortby is not None:
2069                    label = "%s: %s" % (self._get_label(sortby), param)
2070                fig = ncfile.plot_qps_vs_e0(with_fields=list_strings(field),
2071                    e0=e0, ax_list=ax_list, color=cmap(i / len(lnp_list)), fontsize=fontsize,
2072                    sharey=sharey, label=label, show=False)
2073                ax_list = fig.axes
2074        else:
2075            # group_and_sortby and build (ngroups,) subplots
2076            groups = self.group_and_sortby(hue, sortby)
2077            nrows, ncols = 1, len(groups)
2078            ax_mat, fig, plt = get_axarray_fig_plt(None, nrows=nrows, ncols=ncols,
2079                                                   sharex=True, sharey=sharey, squeeze=False)
2080            for ig, g in enumerate(groups):
2081                subtitle = "%s: %s" % (self._get_label(hue), g.hvalue)
2082                ax_mat[0, ig].set_title(subtitle, fontsize=fontsize)
2083                for i, (nclabel, ncfile, param) in enumerate(g):
2084                    fig = ncfile.plot_qps_vs_e0(with_fields=list_strings(field),
2085                        e0=e0, ax_list=ax_mat[:, ig], color=cmap(i / len(g)), fontsize=fontsize,
2086                        sharey=sharey, label="%s: %s" % (self._get_label(sortby), param), show=False)
2087
2088                if ig != 0:
2089                    for ax in ax_mat[:, ig]:
2090                        set_visible(ax, False, "ylabel")
2091
2092        return fig
2093
2094    @add_fig_kwargs
2095    def plot_selfenergy_conv(self, spin, kpoint, band, sortby=None, hue=None,
2096                             colormap="jet", xlims=None, fontsize=8, **kwargs):
2097        """
2098        Plot the convergence of the e-e self-energy wrt to the ``sortby`` parameter.
2099        Values can be optionally grouped by `hue`.
2100
2101        Args:
2102            spin: Spin index.
2103            kpoint: K-point in self-energy (can be |Kpoint|, list/tuple or int)
2104            band: Band index.
2105            sortby: Define the convergence parameter, sort files and produce plot labels.
2106                Can be None, string or function. If None, no sorting is performed.
2107                If string and not empty it's assumed that the abifile has an attribute
2108                with the same name and `getattr` is invoked.
2109                If callable, the output of sortby(abifile) is used.
2110            hue: Variable that define subsets of the data, which will be drawn on separate lines.
2111                Accepts callable or string
2112                If string, it's assumed that the abifile has an attribute with the same name and getattr is invoked.
2113                If callable, the output of hue(abifile) is used.
2114            colormap: matplotlib color map.
2115            xlims: Set the data limits for the x-axis. Accept tuple e.g. ``(left, right)``
2116                   or scalar e.g. ``left``. If left (right) is None, default values are used.
2117            fontsize: Legend and title fontsize.
2118
2119        Returns: |matplotlib-Figure|
2120        """
2121        # Make sure that nsppol and sigma_kpoints are consistent
2122        self._check_dims_and_params()
2123        import matplotlib.pyplot as plt
2124        cmap = plt.get_cmap(colormap)
2125
2126        if hue is None:
2127            ax_list = None
2128            lnp_list = self.sortby(sortby)
2129            for i, (label, ncfile, param) in enumerate(lnp_list):
2130                sigma = ncfile.read_sigee_skb(spin, kpoint, band)
2131                fig = sigma.plot(ax_list=ax_list,
2132                    label=label, color=cmap(i/len(lnp_list)), show=False)
2133                ax_list = fig.axes
2134        else:
2135            # group_and_sortby and build (3, ngroups) subplots
2136            groups = self.group_and_sortby(hue, sortby)
2137            nrows, ncols = 3, len(groups)
2138            ax_mat, fig, plt = get_axarray_fig_plt(None, nrows=nrows, ncols=ncols,
2139                                                   sharex=True, sharey=True, squeeze=False)
2140            for ig, g in enumerate(groups):
2141                subtitle = "%s: %s" % (self._get_label(hue), g.hvalue)
2142                ax_mat[0, ig].set_title(subtitle, fontsize=fontsize)
2143                for i, (nclabel, ncfile, param) in enumerate(g):
2144                    sigma = ncfile.read_sigee_skb(spin, kpoint, band)
2145                    fig = sigma.plot(ax_list=ax_mat[:, ig],
2146                        label="%s: %s" % (self._get_label(sortby), param),
2147                        color=cmap(i / len(g)), show=False)
2148
2149            if ig != 0:
2150                for ax in ax_mat[:, ig]:
2151                    set_visible(ax, False, "ylabel")
2152
2153            for ax in ax_mat.ravel():
2154                set_axlims(ax, xlims, "x")
2155
2156        return fig
2157
2158    def yield_figs(self, **kwargs):  # pragma: no cover
2159        """
2160        This function *generates* a predefined list of matplotlib figures with minimal input from the user.
2161        """
2162        yield self.plot_qpgaps_convergence(plot_qpmks=True, show=False)
2163        #yield self.plot_qpdata_conv_skb(spin, kpoint, band, show=False)
2164        #yield self.plot_qpfield_vs_e0(field, show=False)
2165        #yield self.plot_selfenergy_conv(spin, kpoint, band, sortby=None, hue=None, show=False)
2166
2167    def write_notebook(self, nbpath=None):
2168        """
2169        Write a jupyter_ notebook to ``nbpath``. If nbpath is None, a temporay file in the current
2170        working directory is created. Return path to the notebook.
2171        """
2172        nbformat, nbv, nb = self.get_nbformat_nbv_nb(title=None)
2173
2174        args = [(l, f.filepath) for l, f in self.items()]
2175        nb.cells.extend([
2176            #nbv.new_markdown_cell("# This is a markdown cell"),
2177            #nbv.new_code_cell("plotter = robot.get_ebands_plotter()"),
2178            nbv.new_code_cell("robot = abilab.SigresRobot(*%s)\nrobot.trim_paths()\nrobot" % str(args)),
2179            nbv.new_code_cell("robot.get_qpgaps_dataframe(spin=None, kpoint=None, with_geo=False)"),
2180            nbv.new_code_cell("robot.plot_qpgaps_convergence(plot_qpmks=True, sortby=None, hue=None);"),
2181            nbv.new_code_cell("#robot.plot_qpdata_conv_skb(spin=0, kpoint=0, band=0, sortby=None, hue=None);"),
2182            nbv.new_code_cell("robot.plot_qpfield_vs_e0(field='qpeme0', sortby=None, hue=None);"),
2183            nbv.new_code_cell("#robot.plot_selfenergy_conv(spin=0, kpoint=0, band=0, sortby=None, hue=None);"),
2184        ])
2185
2186        # Mixins
2187        nb.cells.extend(self.get_baserobot_code_cells())
2188        nb.cells.extend(self.get_ebands_code_cells())
2189
2190        return self._write_nb_nbpath(nb, nbpath)
2191