Source code for vibeqc.plot

"""Matplotlib plotters for periodic-system observables.

This module is a *thin* presentation layer on top of
:mod:`vibeqc.bands`. It imports matplotlib lazily so that vibe-qc itself
does not impose a hard matplotlib dependency — callers only pay the
import cost when they actually draw something.

All functions accept an optional ``ax`` (or ``axes``) so they compose
into existing figures, and return the matplotlib ``Figure`` so the
caller can save it.
"""

from __future__ import annotations

from typing import Optional, Tuple

import numpy as np

from .bands import BandStructure, DensityOfStates


__all__ = [
    "band_structure_figure",
    "dos_figure",
    "bands_dos_figure",
]


_HARTREE_TO_EV = 27.211386245988


def _require_matplotlib():
    try:
        import matplotlib.pyplot as plt   # noqa: F401
    except ImportError as e:                # pragma: no cover - import error
        raise ImportError(
            "vibeqc.plot requires matplotlib. Install with "
            "`pip install matplotlib`."
        ) from e
    import matplotlib.pyplot as plt
    return plt




def band_structure_figure(
    bs: BandStructure,
    *,
    ax=None,
    units: str = "eV",
    shift_to_fermi: bool = True,
    color: str = "tab:blue",
    linewidth: float = 1.0,
    title: Optional[str] = None,
):
    """Draw a band-structure plot.

    Parameters
    ----------
    bs
        :class:`vibeqc.BandStructure` from :func:`vibeqc.band_structure`.
    ax
        Existing matplotlib axes; a new figure is made if ``None``.
    units
        ``"eV"`` (default) or ``"Hartree"``.
    shift_to_fermi
        If True and ``bs.e_fermi`` is set, shift the y-axis so the
        Fermi level sits at zero.
    """
    plt = _require_matplotlib()
    if ax is None:
        fig, ax = plt.subplots(figsize=(5, 4))
    else:
        fig = ax.figure

    if units == "eV":
        scale = _HARTREE_TO_EV
        ylabel = "Energy (eV)"
    elif units in ("Hartree", "Ha", "hartree"):
        scale = 1.0
        ylabel = "Energy (Hartree)"
    else:
        raise ValueError(f"Unknown units: {units!r}")

    e0 = bs.e_fermi if (shift_to_fermi and bs.e_fermi is not None) else 0.0
    energies = (bs.energies - e0) * scale
    x = bs.kpath.distances

    for n in range(bs.n_bands):
        ax.plot(x, energies[:, n], color=color, lw=linewidth)

    if shift_to_fermi and bs.e_fermi is not None:
        ax.axhline(0.0, color="0.4", lw=0.6, ls="--")
        ylabel = ylabel.replace("Energy", "E − E_F")

    # High-symmetry tick marks.
    ax.set_xticks([d for d, _ in bs.kpath.labels])
    ax.set_xticklabels([lbl for _, lbl in bs.kpath.labels])
    for d, _ in bs.kpath.labels:
        ax.axvline(d, color="0.7", lw=0.5)

    ax.set_xlim(x.min(), x.max())
    ax.set_ylabel(ylabel)
    ax.set_xlabel("k-path")
    if title:
        ax.set_title(title)
    return fig





def dos_figure(
    dos: DensityOfStates,
    *,
    ax=None,
    units: str = "eV",
    shift_to_fermi: bool = True,
    orientation: str = "vertical",
    color: str = "tab:orange",
    fill: bool = True,
    title: Optional[str] = None,
):
    """Draw a density-of-states plot.

    ``orientation="horizontal"`` puts energy on the y-axis (so the
    figure can sit beside a band-structure panel).
    """
    plt = _require_matplotlib()
    if ax is None:
        fig, ax = plt.subplots(figsize=(4, 4))
    else:
        fig = ax.figure

    if units == "eV":
        scale = _HARTREE_TO_EV
        e_label = "Energy (eV)"
    elif units in ("Hartree", "Ha", "hartree"):
        scale = 1.0
        e_label = "Energy (Hartree)"
    else:
        raise ValueError(f"Unknown units: {units!r}")

    e0 = dos.e_fermi if (shift_to_fermi and dos.e_fermi is not None) else 0.0
    e = (dos.energies - e0) * scale
    d = dos.dos / scale  # density per unit energy → rescale with units

    if orientation == "vertical":
        ax.plot(e, d, color=color, lw=1.0)
        if fill:
            ax.fill_between(e, 0.0, d, color=color, alpha=0.25)
        ax.set_xlabel(("E − E_F " if (shift_to_fermi and dos.e_fermi is not None) else "") + e_label)
        ax.set_ylabel("DOS (states / energy / cell)")
        if shift_to_fermi and dos.e_fermi is not None:
            ax.axvline(0.0, color="0.4", lw=0.6, ls="--")
    elif orientation == "horizontal":
        ax.plot(d, e, color=color, lw=1.0)
        if fill:
            ax.fill_betweenx(e, 0.0, d, color=color, alpha=0.25)
        ax.set_ylabel(("E − E_F " if (shift_to_fermi and dos.e_fermi is not None) else "") + e_label)
        ax.set_xlabel("DOS")
        if shift_to_fermi and dos.e_fermi is not None:
            ax.axhline(0.0, color="0.4", lw=0.6, ls="--")
    else:
        raise ValueError(f"orientation must be 'vertical' or 'horizontal', got {orientation!r}")

    if title:
        ax.set_title(title)
    return fig





def bands_dos_figure(
    bs: BandStructure,
    dos: DensityOfStates,
    *,
    units: str = "eV",
    shift_to_fermi: bool = True,
    width_ratio: Tuple[float, float] = (3.0, 1.0),
    title: Optional[str] = None,
):
    """Combined band structure + DOS panel — the standard solid-state
    layout (bands on the left with shared y-axis to a horizontal DOS on
    the right).
    """
    plt = _require_matplotlib()
    fig, (ax_b, ax_d) = plt.subplots(
        1, 2, sharey=True,
        gridspec_kw={"width_ratios": list(width_ratio)},
        figsize=(7, 4),
    )
    band_structure_figure(
        bs, ax=ax_b, units=units, shift_to_fermi=shift_to_fermi,
    )
    dos_figure(
        dos, ax=ax_d, units=units, shift_to_fermi=shift_to_fermi,
        orientation="horizontal",
    )
    ax_d.set_ylabel("")
    if title:
        fig.suptitle(title)
    fig.tight_layout()
    return fig