1# coding: utf-8
2"""
3Objects related to the computation of Debye-Waller tensors from the generalized phonon DOS.
4"""
5import numpy as np
6import abipy.core.abinit_units as abu
7
8from collections import OrderedDict
9from monty.string import list_strings, marquee
10from monty.collections import dict2namedtuple
11from monty.termcolor import cprint
12from abipy.core.mixins import Has_Structure
13from abipy.tools.plotting import add_fig_kwargs, set_axlims, get_axarray_fig_plt, set_visible
14from abipy.tools.printing import print_dataframe
15
16
17class _Component(object):
18    """
19    Object used to select/plot the components of the DW tensor.
20    """
21
22    def __init__(self, name, ij, **plot_kwargs):
23        self.name = name
24        self.ij = ij
25        self.i, self.j = None, None
26        if self.ij is not None:
27            self.i, self.j = self.ij[0], self.ij[1]
28        self.plot_kwargs = plot_kwargs
29
30    def get_tavg_label(self, what, with_units=False):
31        n = dict(displ="U", vel="V")[what]
32        unit = ""
33        if with_units:
34            unit = r"\;%s" % (dict(displ=r"\AA^2", vel="v")[what])
35
36        if self.name == "trace":
37            return r"$\langle {%s}^2 \rangle%s$" % (n, unit)
38        else:
39            return r"$\langle {%s}_{%s} \rangle%s$" % (n, self.name, unit)
40
41    def eval33w(self, mat33w):
42        #assert mat33w.shape[:2] == (3, 3)
43        if self.ij is not None:
44            return mat33w[self.i, self.j]
45        if self.name == "trace":
46            return mat33w.trace()
47
48        raise TypeError("Don't know how to extract data for `%s`" % self.name)
49
50
51class MsqDos(Has_Structure):
52    """
53    This object stores the generalized phonon DOS in CARTESIAN coords.
54    This DOS-like quantity allows one to calculate Debye Waller factors as a function of Temperature
55    by integration with 1/omega and the Bose-Einstein factor.
56    This object is usually instanciated via the read_msq_dos method of the PhdosReader
57    and stored as attribute of the PhdosFile instance.
58
59    See :cite:`Lee1995` for the further details about the internal implementation and
60    :cite:`Trueblood1996` for the different conventions used by crystallographers.
61
62    See also http://atztogo.github.io/phonopy/formulation.html#thermal-displacement
63    """
64    C = _Component
65    ALL_COMPS = OrderedDict([
66        ("trace", C(name="trace", ij=None, color="k")),
67        ("xx", C(name="xx", ij=(0, 0), color="r", ls="-")),
68        ("yy", C(name="yy", ij=(1, 1), color="g", ls="-")),
69        ("zz", C(name="zz", ij=(2, 2), color="b", ls="-")),
70        ("xy", C(name="xy", ij=(0, 1), color="c", ls="--")),
71        ("xz", C(name="xz", ij=(0, 2), color="m", ls="--")),
72        ("yz", C(name="yz", ij=(1, 2), color="y", ls="--")),
73        # Symmetric components.
74        #("yx", (1, 0)),
75        #("zx", (2, 0)),
76        #("zy", (2, 1)),
77    ])
78    del C
79
80    def __init__(self, structure, wmesh, values, amu_symbol):
81        """
82        Arg:
83            structure: |Structure| object.
84            wmesh: Frequency mesh in eV
85            values: (natom, 3, 3, nomega) array with generalized DOS in Cartesian coordinates.
86            amu_symbol: Dictionary element.symbol -> mass in atomic units.
87        """
88        self._structure = structure
89        self.wmesh = wmesh * abu.eV_Ha ####
90        self.nw = len(self.wmesh)
91        self.values = values * abu.Ha_eV ###
92        self.amu_symbol = amu_symbol
93        assert len(self.values) == len(self.structure)
94
95    @property
96    def structure(self):
97        """|Structure| object."""
98        return self._structure
99
100    def __str__(self):
101        """Invoked by str"""
102        return self.to_string()
103
104    def to_string(self, verbose=0):
105        """
106        Human-readable string with useful information on the object.
107
108        Args:
109            verbose: Verbosity level.
110        """
111        lines = []; app = lines.append
112
113        app(self.structure.to_string(verbose=verbose, title="Structure"))
114        app("")
115        app(marquee(r"Fullfilment of \int dw g_ij(w) = \delta_ij", mark="="))
116        app("")
117        from scipy.integrate import simps
118        for iatom, site in enumerate(self.structure):
119            d = simps(self.values[iatom], x=self.wmesh)
120            app("For site: %s" % site)
121            app(str(d))
122            app("Trace: %.4f, determinant: %.4f" % (d.trace(), np.linalg.det(d)))
123            app("")
124
125        for fmt in ("cartesian", "cif"):
126            df = self.get_dataframe(temp=300, view="inequivalent", fmt=fmt, verbose=verbose)
127            s = print_dataframe(df, title="Format: %s" % fmt, file="string")
128            lines.extend(s.splitlines())
129
130        title = marquee("Constraints on tensor components in reduced coords induced by site symmetries", mark="=")
131        s = print_dataframe(self.structure.site_symmetries.get_tensor_rank2_dataframe(), file="string", title=title)
132        lines.extend(s.splitlines())
133
134        #if verbose > 1:
135            #max_err = self.check_site_symmetries(verbose=verbose)
136
137        return "\n".join(lines)
138
139    def get_json_doc(self, tstart=0, tstop=600, num=11):
140        """
141        Return dictionary with results. Used by emmet builder.
142
143        Args:
144            tstart: The starting value (in Kelvin) of the temperature mesh.
145            tstop: The end value (in Kelvin) of the mesh.
146            num: int, optional Number of samples to generate.
147        """
148        tmesh = np.linspace(tstart, tstop, num=num)
149        # (natom, 3, 3, nt)
150        ucart_t = self.get_msq_tmesh(tmesh, what_list=("displ")).displ
151
152        # Convert tensor to U_CIF format
153        ucif_t = np.empty_like(ucart_t)
154        for it in range(num):
155            ucif_t[:,:,:,it] = self.convert_ucart(ucart_t[:,:,:,it], fmt="cif")
156
157        jdoc = {
158            "natom": len(self.structure),
159            "nomega": self.nw,             # Number of frequencies
160            "ntemp": len(tmesh),            # Number of temperatures
161            "tmesh": tmesh,                 # Temperature mesh in K
162            "wmesh": self.wmesh,            # Frequency mesh in ??
163            "gdos_aijw": self.values,       # Generalized DOS in Cartesian coords. (natom, 3, 3, nomega) array.
164            "amu_symbol": self.amu_symbol,  # Dict symbol --> Atomic mass in a.u.
165            "structure": self.structure,    # Structure object
166            "ucif_t": ucif_t,               # U tensors (natom, 3, 3, ntemp)  as a function of T for T in tmesh in CIF format.
167            "ucif_string_t300k": self.get_cif_string(temp=300, symprec=None),  # String with U tensor at T=300 in Cif format.
168        }
169
170        return jdoc
171
172    def get_msq_tmesh(self, tmesh, iatom_list=None, what_list=("displ", "vel")):
173        """
174        Compute mean square displacement for each atom in `iatom_list` as a function of T.
175        in Cartesian coordinates and atomic-units.
176
177        Args:
178            tmesh: array-like with temperatures in Kelvin.
179            iatom_list: List of atom sites to compute. None if all aomts are wanted.
180            what_list: "displ" for displacement, "vel" for velocity tensor.
181
182        Return:
183            namedtuple with (tmesh=tmesh, displ=msq_d, vel=msq_v)
184
185            msq_d = np.empty((natom, 3, 3, nt))
186        """
187        tmesh = np.array(tmesh)
188        nt = len(tmesh)
189
190        # Frequency mesh starts at iomin to avoid 1/0 and ignore eventual negative frequencies.
191        for iomin, w in enumerate(self.wmesh):
192            if w > 1e-12: break
193        else:
194            raise ValueError("Cannot find index such that wmesh[i] > 1e-12 !!!")
195        wvals = self.wmesh[iomin:]
196        nw = len(wvals)
197
198        # We will compute: Ucart(T, k, ij) = 1/M_k \int dw (n(w) + 1/2) g_ij(w) / w for the k-atom in a.u.
199        # Calculate Bose-Einstein occupation factors only once for each T (instead of for each atom).
200        npht = np.zeros((nt, nw))
201        for it, temp in enumerate(tmesh):
202            npht[it] = abu.occ_be(wvals, temp * abu.kb_HaK) + 0.5
203
204        natom = len(self.structure)
205        msq_d = np.empty((natom, 3, 3, nt))
206        msq_v = np.empty((natom, 3, 3, nt))
207        what_list = list_strings(what_list)
208
209        # Perform frequency integration to get tensor(T)
210        from scipy.integrate import simps
211        if iatom_list is not None: iatom_list = set(iatom_list)
212        for iatom in range(natom):
213            if iatom_list is not None and iatom not in iatom_list: continue
214            symbol = self.structure[iatom].specie.symbol
215            for it in range(nt):
216                fn = self.values[iatom, :, :, iomin:] * npht[it]
217                if "displ" in what_list:
218                    # Mean square displacement for each atom as a function of T (bohr^2).
219                    ys = fn / wvals
220                    fact = 1.0 / (self.amu_symbol[symbol] * abu.amu_emass)
221                    msq_d[iatom, :, :, it] = simps(ys, x=wvals) * fact * abu.Bohr_Ang ** 2
222                if "vel" in what_list:
223                    # Mean square velocity for each atom as a function of T (bohr^2/atomic time unit^2)"
224                    ys = fn * wvals
225                    fact = 1.0 / (self.amu_symbol[symbol] * abu.amu_emass)
226                    msq_v[iatom, :, :, it] = simps(ys, x=wvals) * fact # * abu.velocity_at_to_si ** 2
227
228        return dict2namedtuple(tmesh=tmesh, displ=msq_d, vel=msq_v)
229
230    def convert_ucart(self, ucart_mat, fmt):
231        """
232        Convert the U tensor from Cartesian coordinates to format `fmt`
233        Return new tensor. See also :cite:`Grosse-Kunstleve2002`.
234
235        Args:
236            ucart_mat: (natom,3,3) array with tensor in Cartesian coords.
237            fmt: Output format. Available options: "cif", "ustar", "beta", "cartesian"
238
239        Return: (natom, 3, 3) tensor.
240        """
241        natom = len(self.structure)
242        if fmt == "cartesian":
243            return ucart_mat.copy()
244
245        #elif fmt == "B":
246        #    # Eq 7
247        #    return ucart_mat * 8 * np.pi**2
248
249        elif fmt in ("cif", "ustar", "beta"):
250            # Build A matrix
251            amat = self.structure.lattice.matrix.T
252            ainv = np.linalg.inv(amat)
253            new_mat = np.zeros_like(ucart_mat)
254            # Eq 3b: A^-1 U_Cart A^-T
255            for iatom in range(natom):
256                new_mat[iatom] = np.matmul(ainv, np.matmul(ucart_mat[iatom], ainv.T))
257
258            # Now we have Ustar
259            if fmt == "ustar": return new_mat
260            if fmt == "beta": return new_mat * 8 * np.pi**2  # Eq 6
261
262            # CIF Format Eq 4a
263            # Build N matrix (no 2 pi factor)
264            lengths = self.structure.lattice.reciprocal_lattice_crystallographic.lengths
265            ninv = np.diag(1.0 / np.array(lengths, dtype=float))
266            # N^-1 U_star N^-T
267            for iatom in range(natom):
268                new_mat[iatom] = np.matmul(ninv, np.matmul(new_mat[iatom], ninv.T))
269
270            return new_mat
271
272        raise ValueError("Invalid format: `%s`" % str(fmt))
273
274    def get_dataframe(self, temp=300, fmt="cartesian", view="inequivalent", what="displ", decimals=4,
275                      select_symbols=None, verbose=0):
276        """
277        Return |pandas-DataFrame| with cartesian tensor components as columns and (inequivalent) sites along the rows.
278
279        Args:
280            temp: Temperature in Kelvin.
281            fmt: "cartesian" for elements in Cartesian coordinates, "cif" for results in reduced coordinates
282            view: "inequivalent" to show only inequivalent atoms. "all" for all sites.
283            what: "displ" for displament, "vel" for velocity.
284            decimals: Number of decimal places to round to.
285                If decimals is negative, it specifies the number of positions to the left of the decimal point.
286            select_symbols: String or list of strings with chemical symbols. Used to select only atoms of this type.
287            verbose: Verbosity level.
288
289        Return: |pandas-DataFrame|
290        """
291        # Select atoms.
292        aview = self._get_atomview(view, select_symbols=select_symbols, verbose=verbose)
293
294        # [natom, 3, 3, nt=1]
295        msq = self.get_msq_tmesh([float(temp)], iatom_list=aview.iatom_list, what_list=what)
296        ucart = getattr(msq, what)
297        natom = len(self.structure)
298        ucart = np.reshape(ucart, (natom, 3, 3))
299        values = ucart
300        if what == "displ":
301            values = self.convert_ucart(ucart, fmt)
302
303        columns = ["xx", "yy", "zz", "yz", "xz", "xy"]
304        inds = [(0, 0), (1, 1), (2, 2), (1, 2), (0, 2), (0, 1)]
305        rows = []
306        for (iatom, wlabel) in zip(aview.iatom_list, aview.wyck_labels):
307            site = self.structure[iatom]
308            d = OrderedDict()
309            d["element"] = site.specie.symbol
310            d["site_index"] = iatom
311            d["frac_coords"] = np.round(site.frac_coords, decimals=decimals)
312            d["cart_coords"] = np.round(site.coords, decimals=decimals)
313            d["wyckoff"] = wlabel
314            if fmt == "cartesian":
315                d["isotropic"] = ucart[iatom].trace() / 3.0
316                d["determinant"] = np.linalg.det(ucart[iatom])
317            for col, ind in zip(columns, inds):
318                d[col] = values[iatom, ind[0], ind[1]]
319            rows.append(d)
320
321        import pandas as pd
322        return pd.DataFrame(rows, columns=list(rows[0].keys()) if rows else None)
323
324    def write_cif_file(self, filepath, temp=300, symprec=None):
325        """
326        Write CIF file with structure and anisotropic U tensor in CIF format.
327
328        Args:
329            filepath: Name of CIF file. If None, a temporary filepath is created.
330            temp: Temperature in Kelvin (used to compute U).
331            symprec (float): If not none, finds the symmetry of the structure
332                and writes the cif with symmetry information. Passes symprec to the SpacegroupAnalyzer
333
334        Return: Filepath
335        """
336        if filepath is None:
337            import tempfile
338            _, filepath = tempfile.mkstemp(suffix=".cif", text=True)
339
340        with open(filepath, "wt") as fh:
341            fh.write(self.get_cif_string(temp=temp, symprec=symprec))
342
343        return filepath
344
345    #def jsmol(self, temp=300, symprec=None, verbose=0): # pragma: no cover
346    #    """
347    #    Args:
348    #        symprec (float): If not none, finds the symmetry of the structure
349    #            and writes the cif with symmetry information. Passes symprec to the SpacegroupAnalyzer
350    #        verbose: Verbosity level.
351    #    """
352    #    cif_string = self.get_cif_string(temp=temp, symprec=symprec)
353
354    #    try:
355    #      from jupyter_jsmol import JsmolView
356    #    except ImportError:
357    #        raise ImportError("jupyter_jsmol is not installed. See https://github.com/fekad/jupyter-jsmol")
358
359    #    jsmol = JsmolView(color='white')
360    #    from IPython.display import display, HTML
361    #    FIXME TEMPORARY HACK
362    #    display(HTML('<script type="text/javascript" src="/nbextensions/jupyter-jsmol/jsmol/JSmol.min.js"></script>'))
363    #    display(jsmol)
364
365    #    cmd = 'load inline "%s" {1 1 1}; ellipsoid;' % cif_string
366    #    if verbose: print("executing cmd:", cmd)
367    #    jsmol.script(cmd)
368
369    #    return jsmol
370
371    def vesta_open(self, temp=300): # pragma: no cover
372        """
373        Visualize termal displacement ellipsoids at temperature `temp` (Kelvin) with Vesta_ application.
374        """
375        filepath = self.write_cif_file(filepath=None, temp=temp)
376        cprint("Writing structure + Debye-Waller tensor in CIF format for T = %s (K) to file: %s" % (temp, filepath), "green")
377        cprint("In the Vesta GUI, select: Properties -> Atoms -> Show as displament ellipsoids.", "green")
378        from abipy.iotools import Visualizer
379        visu = Visualizer.from_name("vesta")
380
381        return visu(filepath)()
382
383    def get_cif_string(self, temp=300, symprec=None):
384        """
385        Return string with structure and anisotropic U tensor in CIF format at temperature `temp` in Kelvin
386
387        Args:
388            symprec (float): If not none, finds the symmetry of the structure
389                and writes the cif with symmetry information. Passes symprec
390                to the SpacegroupAnalyzer
391        """
392        # Get string with structure in CIF format.
393        # Don't use symprec because it changes the order of the sites
394        # and we must be consistent with site_labels when writing aniso_U terms!
395        from pymatgen.io.cif import CifWriter
396        cif = CifWriter(self.structure, symprec=symprec)
397
398        aniso_u = """loop_
399_atom_site_aniso_label
400_atom_site_aniso_U_11
401_atom_site_aniso_U_22
402_atom_site_aniso_U_33
403_atom_site_aniso_U_23
404_atom_site_aniso_U_13
405_atom_site_aniso_U_12""".splitlines()
406
407        # Compute U tensor in CIF format (reduced coords)
408        natom = len(self.structure)
409        msq = self.get_msq_tmesh([float(temp)], what_list="displ")
410        ucart = getattr(msq, "displ")
411        ucart = np.reshape(ucart, (natom, 3, 3))
412        ucif = self.convert_ucart(ucart, fmt="cif")
413
414        # Add matrix elements. Use 0 based index
415        for iatom, site in enumerate(self.structure):
416            site_label = "%s%d" % (site.specie.symbol, iatom)
417            m = ucif[iatom]
418            aniso_u.append("%s %10.5f %10.5f %10.5f %10.5f %10.5f %10.5f" %
419                    (site_label, m[0, 0], m[1, 1], m[2, 2], m[1, 2], m[0, 2], m[0, 1]))
420
421        return str(cif) + "\n".join(aniso_u)
422
423    def check_site_symmetries(self, temp=300, verbose=0):
424        """
425        Check site symmetries of the displacement tensor at temperature `temp` in Kelvin.
426        Return maximum error.
427        """
428        msq = self.get_msq_tmesh([float(temp)], what_list="displ")
429        ucart = getattr(msq, "displ")
430        natom = len(self.structure)
431        ucart = np.reshape(ucart, (natom, 3, 3))
432
433        return self.structure.site_symmetries.check_site_symmetries(ucart, verbose=verbose)
434
435    def _get_components(self, components):
436        """
437        Return list of components to analyze from user input.
438        """
439        if components == "all":
440            return list(self.ALL_COMPS.values())
441        elif components == "upper":
442            return [self.ALL_COMPS[c] for c in ("xx", "yy", "zz", "yz", "xz", "xy")]
443        elif components == "diag":
444            return [self.ALL_COMPS[c] for c in ("xx", "yy", "zz")]
445        elif components == "offdiag":
446            return [self.ALL_COMPS[c] for c in ("xy", "xz", "yz")]
447        else:
448            return [self.ALL_COMPS[c] for c in list_strings(components)]
449
450    @add_fig_kwargs
451    def plot(self, components="upper", view="inequivalent", units="eV", select_symbols=None,
452             xlims=None, ylims=None, sharey=False, fontsize=8, verbose=0, **kwargs):
453        """
454        Plot the generalized phonon DOS g_ij(omega, atom) for each atom in the unit cell.
455        One subplot per atom. Each subplot shows the 9 independent components of the symmetric tensor.
456        as a function of frequency. By default, only "inequivalent" atoms are shown.
457
458        Args:
459            view: "inequivalent" to show only inequivalent atoms. "all" for all sites.
460            components: List of cartesian tensor components to plot e.g. ["xx", "xy"].
461                "all" for all components. "upper" for the upper triangle, "diag" for diagonal elements.
462            units: Units energy axis. Possible values in ("eV", "meV", "Ha", "cm-1", "Thz").
463                Case-insensitive.
464            select_symbols: String or list of strings with chemical symbols. Used to select only atoms of this type.
465            xlims: Set the data limits for the x-axis. Accept tuple e.g. ``(left, right)``
466                   or scalar e.g. ``left``. If left (right) is None, default values are used.
467            ylims: Set the data limits for the y-axis.
468            sharey: True if y-axis should be shared.
469            fontsize: Legend and title fontsize.
470            verbose: Verbosity level.
471
472        Returns: |matplotlib-Figure|
473        """
474        # TODO Decide units for internal arrays.
475        factor = abu.phfactor_ev2units(units)
476
477        # Select atoms.
478        aview = self._get_atomview(view, select_symbols, verbose=verbose)
479
480        num_plots = len(aview.iatom_list)
481        nrows, ncols = 1, 1
482        if num_plots > 1:
483            ncols = 2
484            nrows = num_plots // ncols + num_plots % ncols
485
486        ax_list, fig, plt = get_axarray_fig_plt(None, nrows=nrows, ncols=ncols,
487                                                sharex=True, sharey=sharey, squeeze=True)
488        ax_list = np.reshape(ax_list, (nrows, ncols)).ravel()
489        # don't show the last ax if num_plots is odd.
490        if num_plots % ncols != 0: ax_list[-1].axis("off")
491
492        xx = self.wmesh * factor
493        components = self._get_components(components)
494
495        # For each atom in the view.
496        for ix, (ax, iatom, site_label) in enumerate(zip(ax_list, aview.iatom_list, aview.site_labels)):
497            irow, icol = divmod(ix, ncols)
498            ax.grid(True)
499            set_axlims(ax, xlims, "x")
500            set_axlims(ax, ylims, "y")
501            ax.set_title(site_label, fontsize=fontsize)
502            #site = self.structure[iatom]
503            #color = cmap(float(iatom) / max((len(iatom_list) - 1), 1))
504
505            # Plot components for this atom on the same ax.
506            for c in components:
507                yw = c.eval33w(self.values[iatom])
508                label = r"$G_{%s}$" % c.name if ix == 0 else None
509                ax.plot(xx, yw / factor, label=label, **c.plot_kwargs)
510
511            # Handle labels.
512            if irow == nrows - 1:
513                ax.set_xlabel('Frequency %s' % abu.phunit_tag(units))
514            else:
515                set_visible(ax, False, "xlabel", "xticklabels")
516
517            if ix == 0:
518                ax.set_ylabel(r"$g_{ij}(\omega)$ 1/%s (Cart coords)" % abu.phunit_tag(units))
519                ax.legend(loc="best", fontsize=fontsize, shadow=True)
520
521        return fig
522
523    @add_fig_kwargs
524    def plot_tensor(self, tstart=0, tstop=600, num=50, components="all", what="displ", view="inequivalent",
525                    select_symbols=None, colormap="jet", xlims=None, ylims=None, fontsize=10, verbose=0, **kwargs):
526        """
527        Plot tensor(T) for each atom in the unit cell.
528        One subplot for each component, each subplot show all inequivalent sites.
529        By default, only "inequivalent" atoms are shown.
530
531        Args:
532            tstart: The starting value (in Kelvin) of the temperature mesh.
533            tstop: The end value (in Kelvin) of the mesh.
534            num: int, optional Number of samples to generate.
535            components: "all" for all components. "diag" for diagonal elements, "offdiag" for off-diagonal terms only.
536            what: "displ" for displament, "vel" for velocity.
537            view: "inequivalent" to show only inequivalent atoms. "all" for all sites.
538            select_symbols: String or list of strings with chemical symbols. Used to select only atoms of this type.
539            colormap: matplotlib colormap.
540            xlims: Set the data limits for the x-axis. Accept tuple e.g. ``(left, right)``
541                   or scalar e.g. ``left``. If left (right) is None, default values are used.
542            ylims: Set the data limits for the y-axis. Accept tuple e.g. ``(left, right)``
543                   or scalar e.g. ``left``. If left (right) is None, default values are used
544            fontsize: Legend and title fontsize.
545            verbose: Verbosity level.
546
547        Returns: |matplotlib-Figure|
548        """
549        # Select atoms.
550        aview = self._get_atomview(view, select_symbols=select_symbols, verbose=verbose)
551
552        # One subplot for each component
553        diag = ["xx", "yy", "zz"]
554        offdiag = ["xy", "xz", "yz"]
555        components = {
556            "all": diag + offdiag, "diag": diag, "offdiag": offdiag,
557        }[components]
558
559        components = self._get_components(components)
560        shape = np.reshape(components, (-1, 3)).shape
561        nrows, ncols = shape[0], shape[1]
562
563        ax_list, fig, plt = get_axarray_fig_plt(None, nrows=nrows, ncols=ncols,
564                                                sharex=True, sharey=True, squeeze=True)
565        ax_list = np.reshape(ax_list, (nrows, ncols)).ravel()
566        cmap = plt.get_cmap(colormap)
567
568        # Compute U(T)
569        tmesh = np.linspace(tstart, tstop, num=num)
570        msq = self.get_msq_tmesh(tmesh, iatom_list=aview.iatom_list, what_list=what)
571        # [natom,3,3,nt] array
572        values = getattr(msq, what)
573
574        for ix, (ax, comp) in enumerate(zip(ax_list, components)):
575            irow, icol = divmod(ix, ncols)
576            ax.grid(True)
577            set_axlims(ax, xlims, "x")
578            set_axlims(ax, ylims, "y")
579            ylabel = comp.get_tavg_label(what, with_units=True)
580            ax.set_ylabel(ylabel, fontsize=fontsize)
581
582            # Plot this component for all inequivalent atoms on the same subplot.
583            for ii, (iatom, site_label) in enumerate(zip(aview.iatom_list, aview.site_labels)):
584                color = cmap(float(ii) / max((len(aview.iatom_list) - 1), 1))
585                ys = comp.eval33w(values[iatom])
586                ax.plot(msq.tmesh, ys, label=site_label if ix == 0 else None,
587                        color=color) #, marker="o")
588                if ix == 0:
589                    ax.legend(loc="best", fontsize=fontsize, shadow=True)
590
591            if irow == 1:
592                ax.set_xlabel('Temperature (K)')
593            else:
594                set_visible(ax, False, "xlabel", "xticklabels")
595
596        return fig
597
598    @add_fig_kwargs
599    def plot_uiso(self, tstart=0, tstop=600, num=50, what="displ", view="inequivalent",
600                  select_symbols=None, colormap="jet", xlims=None, ylims=None, sharey=False,
601                  fontsize=10, verbose=0, **kwargs):
602        """
603        Plot phonon PJDOS for each atom in the unit cell.
604        One subplot for each component, each subplot show all inequivalent sites.
605        By default, only "inequivalent" atoms are shown.
606
607        comparison of Ueq values, which
608        are calculated as the mean of the diagonal elements of the harmonic ADP tensor, (d)
609        comparison of the ADP anisotropy factor, which is defined as the ratio between maximum Uii
610        and minimum Uii values. A ratio of 1 would correspond to an isotropic displacement.
611
612        Args:
613            tstart: The starting value (in Kelvin) of the temperature mesh.
614            tstop: The end value (in Kelvin) of the mesh.
615            num: int, optional Number of samples to generate.
616            components: "all" for all components. "diag" for diagonal elements, "offdiag" for off-diagonal terms only.
617            what: "displ" for displament, "vel" for velocity.
618            view: "inequivalent" to show only inequivalent atoms. "all" for all sites.
619            select_symbols: String or list of strings with chemical symbols. Used to select only atoms of this type.
620            colormap: matplotlib colormap.
621            xlims: Set the data limits for the x-axis. Accept tuple e.g. ``(left, right)``
622                   or scalar e.g. ``left``. If left (right) is None, default values are used.
623            ylims: Set the data limits for the y-axis. Accept tuple e.g. ``(left, right)``
624                   or scalar e.g. ``left``. If left (right) is None, default values are used
625            sharey: True if y-axis should be shared.
626            fontsize: Legend and title fontsize.
627            verbose: Verbosity level.
628
629        Returns: |matplotlib-Figure|
630        """
631        # Select atoms.
632        aview = self._get_atomview(view, select_symbols=select_symbols, verbose=verbose)
633
634        ax_list, fig, plt = get_axarray_fig_plt(None, nrows=2, ncols=1,
635                                                sharex=True, sharey=sharey, squeeze=True)
636        cmap = plt.get_cmap(colormap)
637
638        # Compute U(T)
639        tmesh = np.linspace(tstart, tstop, num=num)
640        msq = self.get_msq_tmesh(tmesh, iatom_list=aview.iatom_list, what_list=what)
641        # [natom, 3, 3, nt]
642        values = getattr(msq, what)
643        ntemp = len(msq.tmesh)
644
645        for ix, ax in enumerate(ax_list):
646            ax.grid(True)
647            set_axlims(ax, xlims, "x")
648            set_axlims(ax, ylims, "y")
649            if what == "displ":
650                ylabel = r"$U_{iso}\;(\AA^2)$" if ix == 0 else \
651                         r"Anisotropy factor ($\dfrac{\epsilon_{max}}{\epsilon_{min}}}$)"
652            elif what == "vel":
653                ylabel = r"$V_{iso}\;(m/s)^2$" if ix == 0 else \
654                         r"Anisotropy factor ($\dfrac{\epsilon_{max}}{\epsilon_{min}}}$)"
655            else:
656                raise ValueError("Unknown value for what: `%s`" % str(what))
657            ax.set_ylabel(ylabel, fontsize=fontsize)
658
659            # Plot this component for all inequivalent atoms on the same subplot.
660            for ii, (iatom, site_label) in enumerate(zip(aview.iatom_list, aview.site_labels)):
661                color = cmap(float(ii) / max((len(aview.iatom_list) - 1), 1))
662                #msq.displ[iatom, 3, 3, nt]
663                if ix == 0:
664                    # ISO calculated as the mean of the diagonal elements of the harmonic ADP tensor
665                    ys = np.trace(values[iatom]) / 3.0
666                elif ix == 1:
667                    # Ratio between maximum Uii and minimum Uii values.
668                    # A ratio of 1 would correspond to an isotropic displacement.
669                    ys = np.empty(ntemp)
670                    for itemp in range(ntemp):
671                        eigs = np.linalg.eigvalsh(values[iatom, :, :, itemp], UPLO='U')
672                        ys[itemp] = eigs.max() / eigs.min()
673                else:
674                    raise ValueError("Invalid ix index: `%s" % ix)
675
676                ax.plot(msq.tmesh, ys, label=site_label if ix == 0 else None,
677                        color=color) #, marker="o")
678                if ix == 0:
679                    ax.legend(loc="best", fontsize=fontsize, shadow=True)
680
681            if ix == len(ax_list) - 1:
682                ax.set_xlabel("Temperature (K)")
683            else:
684                set_visible(ax, False, "xlabel", "xticklabels")
685
686        return fig
687