madman.helpers.ase.workflows.ebands

  1r"""Electronic band structure."""
  2
  3import os
  4from argparse import ArgumentParser
  5from collections.abc import Mapping
  6from math import ceil
  7from typing import Optional
  8
  9import seaborn as sns
 10import yaml
 11from ase.parallel import paropen, parprint
 12from gpaw.calculator import GPAW
 13from matplotlib import pyplot as plt
 14from pydantic import BaseModel, Field, confloat, conint, constr, field_validator
 15
 16from madman.utilities import redirect_to
 17
 18
 19sns.set_theme(
 20    context="talk",
 21    style="white",
 22    rc={"figure.titlesize": "medium", "axes.formatter.useoffset": False},
 23)
 24
 25
 26class Config(BaseModel):
 27    r"""Configuration."""
 28
 29    ground_path: constr(pattern=r".*\.gpw$") = Field(frozen=True)
 30    r"""Path to electronic ground state."""
 31
 32    kptden: confloat(gt=0.0, allow_inf_nan=False) = Field(10.0, frozen=True)
 33    r"""$\mathbf{k}$-point density [$\mathrm{Å}$]."""
 34
 35    nbands_empty: conint(ge=0) = Field(0, frozen=True)
 36    r"""Number of empty bands."""
 37
 38    energy_min: confloat(lt=0.0, allow_inf_nan=False) = Field(-5.0, frozen=True)
 39    r"""Electron energy minimum w.r.t. Fermi level [$\mathrm{eV}$]."""
 40
 41    energy_max: confloat(gt=0.0, allow_inf_nan=False) = Field(5.0, frozen=True)
 42    r"""Electron energy maximum w.r.t. Fermi level [$\mathrm{eV}$]."""
 43
 44    ebands_save_prefix: Optional[str] = Field(None, frozen=True)
 45    r"""Prefix to save electronic band structure."""
 46
 47    ebands_plot_prefix: Optional[str] = Field(None, frozen=True)
 48    r"""Prefix to save electronic band structure plot."""
 49
 50    log_prefix: Optional[str] = Field(None, frozen=True)
 51    r"""Prefix to save GPAW log."""
 52
 53    parallel: Optional[Mapping] = Field(None, frozen=True)
 54    r"""GPAW parallel options."""
 55
 56    @field_validator("ground_path")
 57    def validate_ground_path(cls, value: str) -> str:
 58        r"""Validate `ground_path`.
 59
 60        Args:
 61            value: Path to load electronic ground state.
 62
 63        Returns:
 64            Path to load electronic ground state.
 65
 66        Raises:
 67            FileNotFoundError: If `value` is not a path to an existing file.
 68        """
 69        if not os.path.isfile(value):
 70            raise FileNotFoundError
 71        return value
 72
 73    @field_validator("ebands_save_prefix", "ebands_plot_prefix", "log_prefix")
 74    def create_tree(cls, value: Optional[str]) -> Optional[str]:
 75        r"""Create tree for prefix or path.
 76
 77        Args:
 78            value: Path or prefix of file to save.
 79
 80        Returns:
 81            Path or prefix of file to save.
 82        """
 83        if value:
 84            tree = os.path.dirname(value)
 85            if tree:
 86                os.makedirs(tree, exist_ok=True)
 87        return value
 88
 89
 90def calc_ebands(config: Config) -> None:
 91    r"""Calculate electronic band structure.
 92
 93    Args:
 94        config: Configuration.
 95    """
 96    ground = GPAW(config.ground_path)
 97    parprint("Electronic ground state imported...")
 98
 99    kpath = ground.atoms.cell.bandpath(density=config.kptden)
100    nelectrons = ground.get_number_of_electrons()
101    nbands_occupied = ceil(0.5 * nelectrons)
102    nbands = nbands_occupied + config.nbands_empty
103
104    with redirect_to(config.log_prefix, mode="w"):
105        ebands_calc = ground.fixed_density(
106            kpts=kpath.kpts,
107            symmetry="off",
108            nbands=nbands_occupied + 2 * config.nbands_empty,
109            convergence={"bands": nbands},
110            parallel=config.parallel,
111        )
112    ebands = ebands_calc.band_structure().subtract_reference()
113    parprint("Electronic band structure calculated...")
114
115    if config.ebands_save_prefix:
116        ebands.write(f"{config.ebands_save_prefix}.json")
117        parprint("Electronic band structure saved...")
118
119    fig, ax = plt.subplots(tight_layout=True)
120    ebands.plot(
121        ax=ax,
122        ylabel=r"$\epsilon_{\text{KS}}$ / $\mathrm{eV}$",
123        emin=config.energy_min,
124        emax=config.energy_max,
125        colors="k",
126    )
127    ax.set_xticklabels(
128        [
129            xlabel.get_text().replace(",", "|")
130            for xlabel in ax.get_xmajorticklabels()
131        ]
132    )
133
134    if config.ebands_plot_prefix:
135        fig.savefig(f"{config.ebands_plot_prefix}.svg")
136        parprint("Electronic band structure plot saved...")
137
138
139def calc_ebands_cli() -> None:
140    r"""Calculate electronic band structure - CLI interface."""
141    parser = ArgumentParser(description="Calculate electronic band structure")
142    parser.add_argument(
143        "config",
144        nargs="?",
145        default="./config.yml",
146        help="configuration file",
147    )
148    args = parser.parse_args()
149    with paropen(args.config, "r") as stream:
150        config = yaml.safe_load(stream)
151    config = Config(**config)
152    calc_ebands(config)