Source code for lineshape_tools.plot

"""Various plotting utilities."""

from __future__ import annotations

import logging
from typing import TYPE_CHECKING

from lineshape_tools.constants import omega2eV
from lineshape_tools.lineshape import convert_A_to_L, get_phonon_spec_func
from lineshape_tools.phonon import get_ipr, get_phonons

logger = logging.getLogger(__name__)

if TYPE_CHECKING:
    from pathlib import Path

    import numpy as np


def _make_subplots(
    spec_funcs: list[tuple],
    dE: float,
    emission: bool,
    omega_max: float,
    omega_mult: float,
    figsize: tuple[float, float],
    skip_lim_adjust: bool,
    ax=None,
):
    """Helper function to make subplot-type plot."""
    if ax is not None:
        fig = ax[0].get_figure()
    else:
        import matplotlib.pyplot as plt

        fig, ax = plt.subplots(1, 3, figsize=figsize, constrained_layout=True)

    for spec_func in spec_funcs:
        w, dos, S, tw, L = spec_func

        if dos is not None:
            ax[0].plot(w, dos)

        if S is not None:
            ax[1].plot(w, S)

        if L is not None:
            ax[2].plot(tw, L)

    if not skip_lim_adjust:
        for a in ax:
            a.set_ylim((0, a.get_ylim()[1]))

        for a in ax[:2]:
            a.set_xlim((0.0, omega_max + 0.005))

        if emission:
            ax[2].set_xlim((dE - omega_mult * omega_max - 0.02, dE + 0.02))
        else:
            ax[2].set_xlim((dE - 0.02, dE + omega_mult * omega_max + 0.02))

    ax[0].set_xlabel(r"$\hbar\omega$ [eV]")
    ax[0].set_ylabel(r"DOS $\rho(\hbar\omega)$ [eV$^{-1}$]")
    ax[1].set_xlabel(r"$\hbar\omega$ [eV]")
    ax[1].set_ylabel(r"Spectral Density $S(\hbar\omega)$ [eV$^{-1}$]")
    ax[2].set_xlabel(r"Energy [eV]")
    if emission:
        ax[2].set_ylabel(r"Luminescence $L(\hbar\omega)$ [eV$^{-1}$]")
    else:
        ax[2].set_ylabel(r"Absorption $L(\hbar\omega)$ [eV$^{-1}$]")
    return fig, ax


def _make_inset(
    spec_funcs: list[tuple],
    dE: float,
    emission: bool,
    omega_max: float,
    omega_mult: float,
    figsize: tuple[float, float],
    skip_lim_adjust: bool,
    ax=None,
):
    """Helper function to make inset-type plot."""
    if ax is not None:
        fig = ax[0].get_figure()
    else:
        import matplotlib.pyplot as plt

        fig, ax1 = plt.subplots(figsize=figsize, constrained_layout=True)
        ax = [ax1, fig.add_axes((0.17, 0.55, 0.4, 0.35))]

    for spec_func in spec_funcs:
        w, _, S, tw, L = spec_func

        if L is not None:
            ax[0].plot(tw, L)

        if S is not None:
            ax[1].plot(w, S)

    if not skip_lim_adjust:
        for a in ax:
            a.set_ylim((0, a.get_ylim()[1]))

        ax[1].set_xlim((0.0, omega_max + 0.005))

        if emission:
            ax[0].set_xlim((dE - omega_mult * omega_max - 0.02, dE + 0.02))
        else:
            ax[0].set_xlim((dE - 0.02, dE + omega_mult * omega_max + 0.02))

    ax[0].set_xlabel(r"Energy [eV]")
    if emission:
        ax[0].set_ylabel(r"Luminescence $L(\hbar\omega)$ [eV$^{-1}$]")
    else:
        ax[0].set_ylabel(r"Absorption $L(\hbar\omega)$ [eV$^{-1}$]")
    ax[1].set_xlabel(r"$\hbar\omega$ [eV]")
    ax[1].set_ylabel(r"$S(\hbar\omega)$ [eV$^{-1}$]")
    return fig, ax


def _make_single(
    func_type: str,
    spec_funcs: list[tuple],
    dE: float,
    emission: bool,
    omega_max: float,
    omega_mult: float,
    figsize: tuple[float, float],
    skip_lim_adjust: bool,
    ax=None,
):
    """Helper function to make a single-function plot."""
    if ax is not None:
        fig = ax.get_figure()
    else:
        import matplotlib.pyplot as plt

        fig, ax = plt.subplots(figsize=figsize, constrained_layout=True)

    for spec_func in spec_funcs:
        w, dos, S, tw, L = spec_func

        if func_type[0].lower() == "d":
            ax.plot(w, dos)
        elif func_type[0].lower() == "s":
            ax.plot(w, S)
        else:
            ax.plot(tw, L)

    if not skip_lim_adjust:
        ax.set_ylim((0, ax.get_ylim()[1]))

        if func_type[0].lower() in ("d", "s"):
            ax.set_xlim((0.0, omega_max + 0.005))
        else:
            if emission:
                ax.set_xlim((dE - omega_mult * omega_max - 0.02, dE + 0.02))
            else:
                ax.set_xlim((dE - 0.02, dE + omega_mult * omega_max + 0.02))

    if func_type[0].lower() in ("d", "s"):
        ax.set_xlabel(r"$\hbar\omega$ [eV]")
    else:
        ax.set_xlabel(r"Energy [eV]")

    if func_type[0].lower() == "d":
        ax.set_ylabel(r"DOS $\rho(\hbar\omega)$ [eV$^{-1}$]")
    elif func_type[0].lower() == "s":
        ax.set_ylabel(r"Spectral Density $S(\hbar\omega)$ [eV$^{-1}$]")
    else:
        if emission:
            ax.set_ylabel(r"Luminescence $L(\hbar\omega)$ [eV$^{-1}$]")
        else:
            ax.set_ylabel(r"Absorption $L(\hbar\omega)$ [eV$^{-1}$]")
    return fig, ax


[docs] def plot_spec_funcs( dynmats: tuple | np.ndarray | Path | str | list[tuple | np.ndarray | Path | str], dq: np.ndarray | None, dE: float, gamma_zpl: float = 0.001, sigma_zpl: float = 0.0, sigma_psb: tuple[float, float] = (0.005, 0.001), gamma_psb: tuple[float, float] | None = None, emission: bool = True, omega_mult: float = 5.0, omega_max: float = 0.0, norm: str = "area", T: float = 0, plot_type: str = "subplot", figsize: tuple[float, float] = (8.0, 2.5), skip_lim_adjust: bool = False, ax=None, ): """Make a plot of the spectral functions and luminescence/absorption intensity. Args: dynmats (tuple | np.ndarray | Path | str): path to a dynamical matrix in .npz format or a np.ndarray corresponding to a dynamical matrix (shape 3N x 3N). If a tuple is given, assume that the spectral functions have already been obtained. The tuple should contain elements (w, dos, S, w_L, L) where w is the freq grid of dos/S and w_L is the freq grid of L. A list of dynmats can be provided instead. dq (np.ndarray): mass-weighted displacement vector in (amu^{1/2} Ang). Can be None if all dynmats that are provided are of type tuple (see above). dE (float): energy of the zero-phonon line. gamma_zpl (float): Lorentzian broadening in the ZPL to capture homogeneous broadening. sigma_zpl (float): Gaussian broadening in the ZPL to capture inhomogeneous broadening. sigma_psb (float, float): Gaussian broadening used to broaden the partial Huang-Rhys factors. The broadening factor is linearly interpolated from sigma_psb[0] at zero frequency to sigma_psb[1] at the highest (non-LVM) frequency. gamma_psb (float, float): Turns on Lorentzian broadening of local vibrational modes identified by their inverse participation ratio. gamma_psb[0] is ipr_cut and gamma_psb[1] is gamma_lvm. See :class:`Broadening`. emission (bool): determines if the intensity corresponds to emission or absorption. omega_mult (float): how many factors of the maximum phonon frequency from the ZPL will be plotted in the luminescence/absorption intensity. omega_max (float): maximum phonon frequency can be provided if known, used in determining the plot ranges. norm (str): normalization of luminescence (area or max). T (float): Temperature in kelvin. plot_type (str): type of plot to generate. "subplot" will generate a subplot for the dos, S and L, respectively. "inset" will plot L with S as an inset. To make a plot of a single type of function, specify "dos", "S", or "L". figsize (tuple): figsize passed to `matplotlib.pyplot.subplots`. skip_lim_adjust (bool): do not adjust xlim and ylim when making plots. ax (plt.Axes): matplotlib axes object (an array of them) if already made. """ if not isinstance(dynmats, list): dynmats = [dynmats] spec_funcs = [] for i, dynmat in enumerate(dynmats): if isinstance(dynmat, tuple): spec_funcs.append(dynmat) continue logger.info(f"dynmat {i} - diagonalizing") omega, U = get_phonons(dynmat) dq_k = U.T @ dq if gamma_psb is not None: logger.info(f"dynmat {i} - computing inverse participation ratios") ipr, ipr_cut, gamma_lvm = get_ipr(U), gamma_psb[0], gamma_psb[1] else: ipr, ipr_cut, gamma_lvm = None, None, None if omega2eV * omega.max() > omega_max: omega_max = omega2eV * omega.max() logger.info(f"dynmat{i} - evaluating spectral functions") w, dos, S, A = get_phonon_spec_func( dq_k, omega, sigma_psb=sigma_psb, gamma_zpl=gamma_zpl, sigma_zpl=sigma_zpl, ipr=ipr, ipr_cut=ipr_cut, gamma_lvm=gamma_lvm, T=T, ) tw, L = convert_A_to_L(w, A, dE, emission=emission, norm=norm) spec_funcs.append((w, dos, S, tw, L)) if omega_max <= 0.0 and not skip_lim_adjust: logger.warning("the maximum omega could not be determined, falling back to 0.1") omega_max = 0.1 if plot_type[0].lower() == "i": return _make_inset( spec_funcs, dE, emission, omega_max, omega_mult, figsize, skip_lim_adjust, ax=ax ) elif plot_type.lower()[:2] == "su": return _make_subplots( spec_funcs, dE, emission, omega_max, omega_mult, figsize, skip_lim_adjust, ax=ax ) elif plot_type[0].lower() in ("d", "s", "l"): return _make_single( plot_type, spec_funcs, dE, emission, omega_max, omega_mult, figsize, skip_lim_adjust, ax=ax, ) else: raise ValueError(f"unknown plot_type {plot_type}")