madman.helpers.ase.workflows.edos

  1r"""Electron density of states."""
  2
  3import os
  4from argparse import ArgumentParser
  5from collections.abc import Mapping
  6from functools import cache
  7from math import ceil
  8from typing import Any, Optional, Union
  9
 10import numpy as np
 11import seaborn as sns
 12import xarray as xr
 13import yaml
 14from ase.parallel import paropen, parprint, world
 15from gpaw.calculator import GPAW
 16from matplotlib import pyplot as plt
 17from pydantic import BaseModel, Field, confloat, conint, constr, field_validator
 18
 19from madman.helpers.ase.workflows.converge import MultivariateConvergence
 20from madman.utilities import gen_regular_grid, redirect_to
 21
 22
 23sns.set_theme(
 24    context="talk",
 25    style="white",
 26    rc={"figure.titlesize": "medium", "axes.formatter.useoffset": False},
 27)
 28
 29
 30atm_ssh_seq: list[str, str, str, str] = ["s", "p", "d", "f"]
 31r"""Atomic subshells (index = l quantum number)."""
 32
 33
 34atm_ssh_to_harm_poly_seq: dict[str, list[str, ...]] = {
 35    "s": [""],
 36    "p": ["y", "z", "x"],
 37    "d": ["xy", "yz", "z2", "xz", "x2-y2"],
 38    "f": ["y3", "xyz", "yz2", "z3", "xz2", "x2z-y2z", "x3"],
 39}
 40r"""Atomic subshells to harmonic polynomials (index = GPAW m quantum number)."""
 41
 42
 43class kptdenConvergenceSettings(BaseModel):
 44    r"""$\mathbf{k}$-point density convergence settings."""
 45
 46    value_min: confloat(gt=0.0, allow_inf_nan=False) = Field(frozen=True)
 47    r"""Minimum value [$\mathrm{Å}$]."""
 48
 49    value_max: Optional[confloat(gt=0.0, allow_inf_nan=False)] = Field(
 50        None, frozen=True
 51    )
 52    r"""Maximum value [$\mathrm{Å}$]."""
 53
 54    value_spc: Optional[confloat(gt=0.0, allow_inf_nan=False)] = Field(
 55        None, frozen=True
 56    )
 57    r"""Value spacing [$\mathrm{Å}$]."""
 58
 59    threshold: Optional[confloat(gt=0.0, allow_inf_nan=False)] = Field(
 60        None, frozen=True
 61    )
 62    r"""Convergence threshold [$\mathrm{eV^{-1} \, Å^{-1}}$]."""
 63
 64    stability: Optional[conint(gt=0)] = Field(None, frozen=True)
 65    r"""Convergence stability."""
 66
 67
 68class nbandsConvergenceSettings(BaseModel):
 69    r"""Number of empty bands convergence settings."""
 70
 71    value_min: conint(ge=0) = Field(frozen=True)
 72    r"""Minimum value."""
 73
 74    value_max: Optional[conint(ge=0)] = Field(None, frozen=True)
 75    r"""Maximum value."""
 76
 77    value_spc: Optional[conint(gt=0)] = Field(None, frozen=True)
 78    r"""Value spacing."""
 79
 80    threshold: Optional[confloat(gt=0.0, allow_inf_nan=False)] = Field(
 81        None, frozen=True
 82    )
 83    r"""Convergence threshold [$\mathrm{eV^{-1}}$]."""
 84
 85    stability: Optional[conint(gt=0)] = Field(None, frozen=True)
 86    r"""Convergence stability."""
 87
 88
 89class Config(BaseModel):
 90    r"""Configuration."""
 91
 92    ground_path: constr(pattern=r".*\.gpw$") = Field(frozen=True)
 93    r"""Path to electronic ground state."""
 94
 95    kptden: Union[
 96        kptdenConvergenceSettings,
 97        Mapping,
 98        confloat(gt=0.0, allow_inf_nan=False),
 99    ] = Field(10.0, frozen=True)
100    r"""$\mathbf{k}$-point density [$\mathrm{Å}$] convergence settings."""
101
102    nbands_empty: Union[
103        nbandsConvergenceSettings,
104        Mapping,
105        conint(ge=0),
106    ] = Field(0, frozen=True)
107    r"""Number of empty bands convergence settings."""
108
109    strict: bool = Field(False, frozen=True)
110    r"""If True, max. slope is used for convergence; otherwise, avg. is used."""
111
112    niter_max: conint(gt=0) = Field(1, frozen=True)
113    r"""Maximum number of iterations for convergence."""
114
115    energy_min: confloat(lt=0.0, allow_inf_nan=False) = Field(-5.0, frozen=True)
116    r"""Electron energy minimum w.r.t. Fermi level [$\mathrm{eV}$]."""
117
118    energy_max: confloat(gt=0.0, allow_inf_nan=False) = Field(5.0, frozen=True)
119    r"""Electron energy maximum w.r.t. Fermi level [$\mathrm{eV}$]."""
120
121    energy_spc: confloat(gt=0.0, allow_inf_nan=False) = Field(0.01, frozen=True)
122    r"""Electron energy spacing [$\mathrm{eV}$]."""
123
124    energy_spc_prj: Optional[confloat(gt=0.0, allow_inf_nan=False)] = Field(None, frozen=True)
125    r"""Electron energy spacing for projected density of states."""
126
127    projected: bool = Field(False, frozen=True)
128    r"""If True, also projected electron density of states is calculated."""
129
130    gpaw_calc_prefix: Optional[str] = Field(None, frozen=True)
131    r"""Prefix to save GPAW calculator."""
132
133    edos_save_prefix: Optional[str] = Field(None, frozen=True)
134    r"""Prefix to save electron density of states."""
135
136    edos_plot_prefix: Optional[str] = Field(None, frozen=True)
137    r"""Prefix to save electron density of states plots."""
138
139    conv_plot_root: Optional[str] = Field(None, frozen=True)
140    r"""Root to save convergence plots."""
141
142    log_prefix: Optional[str] = Field(None, frozen=True)
143    r"""Prefix to save GPAW log."""
144
145    cache_root: Optional[str] = Field(None, frozen=True)
146    r"""Root to cache electron densitysity of states."""
147
148    parallel: Optional[Mapping] = Field(None, frozen=True)
149    r"""GPAW parallel options."""
150
151    @field_validator("ground_path")
152    def validate_ground_path(cls, value: str) -> str:
153        r"""Validate `ground_path`.
154
155        Args:
156            value: Path to load electronic ground state.
157
158        Returns:
159            Path to load electronic ground state.
160
161        Raises:
162            FileNotFoundError: If `value` is not a path to an existing file.
163        """
164        if not os.path.isfile(value):
165            raise FileNotFoundError
166        return value
167
168    @field_validator("kptden")
169    def parse_kptden_map(cls, value: Any) -> Any:
170        r"""Parse $\mathrm{k}$-point density convergence settings map."""
171        if isinstance(value, Mapping):
172            value = kptdenConvergenceSettings(**value)
173        return value
174
175    @field_validator("nbands_empty")
176    def parse_nbands_empty_map(cls, value: Any) -> Any:
177        r"""Parse number of empty bands convergence settings map."""
178        if isinstance(value, Mapping):
179            value = nbandsConvergenceSettings(**value)
180        return value
181
182    @field_validator(
183        "gpaw_calc_prefix",
184        "edos_save_prefix",
185        "edos_plot_prefix",
186        "conv_plot_root",
187        "log_prefix",
188        "cache_root",
189    )
190    def create_tree(cls, value: Optional[str]) -> Optional[str]:
191        r"""Create tree for prefix or path.
192
193        Args:
194            value: Path or prefix of file to save.
195
196        Returns:
197            Path or prefix of file to save.
198        """
199        if value:
200            tree = os.path.dirname(value)
201            if tree:
202                os.makedirs(tree, exist_ok=True)
203        return value
204
205
206def calc_edos(config: Config) -> None:
207    r"""Calculate electron densitysity of states.
208
209    Args:
210        config: Configuration.
211    """
212    with redirect_to(config.log_prefix, mode="w"):
213        pass
214    parprint("GPAW log initialized...")
215
216    calc_cache = {}
217    calc_iter_i = 0
218    parprint("GPAW calculator cache initialized...")
219
220    ground = GPAW(config.ground_path)
221    parprint("Electronic ground state imported...")
222
223    nelectrons = ground.get_number_of_electrons()
224    nbands_occupied = ceil(0.5 * nelectrons)
225    parprint(f"Number of occupied bands: {nbands_occupied}...")
226
227    kptden_d = config.kptden
228    if isinstance(kptden_d, float):
229        kptden_d = {"value_min": kptden_d}
230    else:
231        kptden_d = kptden_d.dict()
232    kptden_d.update(
233        {
234            "direction": 1,
235            "criterion": f"slope-abs-{"max" if config.strict else "avg"}",
236            "symb": r"\lambda_{\mathbf{k}}",
237            "unit": r"Å",
238        }
239    )
240    nbands_empty_d = config.nbands_empty
241    if isinstance(nbands_empty_d, int):
242        nbands_empty_d = {"value_min": nbands_empty_d}
243    else:
244        nbands_empty_d = nbands_empty_d.dict()
245    nbands_empty_d.update(
246        {
247            "direction": 1,
248            "criterion": f"slope-abs-{"max" if config.strict else "avg"}",
249            "symb": r"N^{*}",
250            "unit": r"1",
251        }
252    )
253    params = {"kptden": kptden_d, "nbands_empty": nbands_empty_d}
254    parprint("Electron density of states convergence parameters set up...")
255
256    energy_grid = gen_regular_grid(
257        config.energy_min, config.energy_max, config.energy_spc
258    )
259    parprint("Electron energy grid set up...")
260
261    @cache
262    def objective(kptden, nbands_empty):
263        nonlocal calc_iter_i
264        nbands_empty = int(nbands_empty)
265        nbands = nbands_occupied + nbands_empty
266        with redirect_to(config.log_prefix, mode="a"):
267            calc = ground.fixed_density(
268                kpts={"density": kptden, "gamma": True},
269                nbands=nbands + nbands_empty,
270                convergence={"bands": nbands},
271                parallel=config.parallel,
272            )
273        calc_iter_i += 1
274        key = [("kptden", kptden), ("nbands_empty", nbands_empty)]
275        if config.cache_root is None:
276            calc_cache.update({tuple(sorted(key)): calc})
277        else:
278            calc_path = os.path.join(config.cache_root, f"{calc_iter_i}.gpw")
279            calc.write(calc_path)
280            calc_cache.update({tuple(sorted(key)): calc_path})
281        edos_calc = calc.dos()
282        edos = edos_calc.raw_dos(energy_grid, width=0.0)
283        return edos
284
285    parprint("Electron density of states objective function set up...")
286
287    convergence = MultivariateConvergence(
288        objective,
289        params,
290        crop=True,
291        req_sc=True,
292        niter_max=config.niter_max,
293        obj_symb=r"\text{Electron DOS}",
294        obj_unit=r"1/eV",
295    )
296    parprint("Electron density of states convergence set up...")
297
298    parprint("Electron density of states convergence started...")
299    convergence.run()
300
301    if config.conv_plot_root is not None:
302        value_plots_map = convergence.plot("obj-value")
303        slope_plots_map = convergence.plot("obj-slope")
304
305        ndigits = len(str(config.niter_max))
306        for qnty, plots_map in zip(
307            ["edos", "medos"], [value_plots_map, slope_plots_map]
308        ):
309            for param, plots_seq in plots_map.items():
310                figdir = os.path.join(
311                    config.conv_plot_root, f"{qnty}-vs-{param}"
312                )
313                os.makedirs(figdir, exist_ok=True)
314                for i, plots in enumerate(plots_seq):
315                    (plot,) = tuple(plots)
316                    figpath = os.path.join(figdir, f"{i:0{ndigits}.0f}.svg")
317                    fig = plot.get_figure()
318                    fig.savefig(figpath)
319                    plt.close(fig)
320        parprint("Converge plots saved...")
321
322    if convergence.converged is False:
323        parprint("Electron density of states convergence failed!")
324    else:
325        parprint("Electron density of states converged...")
326
327    params_opt = convergence.values_opt
328    key = tuple(sorted(params_opt.items()))
329    if key not in calc_cache:
330        objective(**params_opt)
331
332    if config.cache_root is None:
333        calc = calc_cache[key]
334    else:
335        calc = GPAW(calc_cache[key])
336    if config.gpaw_calc_prefix is not None:
337        calc.write(f"{config.gpaw_calc_prefix}.gpw")
338        parprint("GPAW calculator saved...")
339
340    edos_calc = calc.dos()
341    edos = edos_calc.raw_dos(energy_grid, width=0.0)
342    parprint("Electron density of states calculated...")
343
344    edos_da = xr.DataArray(edos, coords=[energy_grid], dims=["energy"])
345    edos_da.to_netcdf(f"{config.edos_save_prefix}_total.nc")
346    parprint("Electron density of states saved...")
347
348    fig, ax = plt.subplots(tight_layout=True)
349    ax.set_xlabel(r"$\epsilon_{\text{KS}}$ / $\mathrm{eV}$")
350    ax.set_ylabel(r"$\text{Electron DOS}$ / $\mathrm{eV^{-1}}$")
351    ax.plot(edos_da.energy.values, edos_da.values, color="k")
352    if config.edos_save_prefix is not None:
353        fig.savefig(f"{config.edos_plot_prefix}_total.svg")
354        parprint("Electron density of states plot saved...")
355
356    if config.projected is True:
357        if config.energy_spc_prj is not None:
358            energy_grid = gen_regular_grid(
359                config.energy_min, config.energy_max, config.energy_spc_prj
360            )
361        parprint("Calculating projected electron density of states...")
362        edos_proj_da_seq = []
363        atm_spc_seq = calc.atoms.get_chemical_symbols()
364        for atm_i, atm_spc in enumerate(atm_spc_seq):
365            for atm_ssh in atm_ssh_seq:
366                for harm_poly in sum(atm_ssh_to_harm_poly_seq.values(), []):
367                    lqntnum = atm_ssh_seq.index(atm_ssh)
368                    try:
369                        mqntnum = atm_ssh_to_harm_poly_seq[atm_ssh].index(
370                            harm_poly
371                        )
372                        edos_proj = edos_calc.raw_pdos(
373                            energy_grid,
374                            a=atm_i,
375                            l=lqntnum,
376                            m=mqntnum,
377                            width=0.0,
378                        )
379                    except ValueError:
380                        edos_proj = np.zeros(energy_grid.shape)
381                    edos_proj_da = xr.DataArray(
382                        edos_proj, coords=[energy_grid], dims=["energy"]
383                    )
384                    edos_proj_da = (
385                        edos_proj_da.assign_coords(atm_spc=atm_spc)
386                        .expand_dims("atm_spc")
387                        .assign_coords(atm_i=atm_i)
388                        .expand_dims("atm_i")
389                        .assign_coords(atm_ssh=atm_ssh)
390                        .expand_dims("atm_ssh")
391                        .assign_coords(harm_poly=harm_poly)
392                        .expand_dims("harm_poly")
393                    )
394                    edos_proj_da_seq.append(edos_proj_da)
395        parprint("Projected electron density of states calculated...")
396
397        edos_proj_da = xr.combine_by_coords(edos_proj_da_seq)
398        if config.edos_save_prefix is not None:
399            edos_proj_da.to_netcdf(f"{config.edos_save_prefix}_projected.nc")
400            parprint("Projected electron density of states saved...")
401
402        edos_proj_vs_atm_spc_da = edos_proj_da.groupby("atm_spc").sum(
403            dim=["atm_i", "atm_ssh", "harm_poly"]
404        )
405        fig, ax = plt.subplots(tight_layout=True)
406        edos_proj_vs_atm_spc_da.plot(x="energy", hue="atm_spc", ax=ax)
407        ax.set_title("")
408        ax.set_xlabel(r"$\epsilon_{\text{KS}}$ / $\mathrm{eV}$")
409        ax.set_ylabel(r"$\text{Electron DOS}$ / $\mathrm{eV^{-1}}$")
410        legend = ax.get_legend()
411        legend.set_loc("upper left")
412        legend.set_bbox_to_anchor((1.0, 1.0))
413        legend.set_frame_on(False)
414        legend.set_title(None)
415        if config.edos_plot_prefix is not None:
416            fig.savefig(f"{config.edos_plot_prefix}_atm-spc.svg")
417            parprint(
418                "Projected (atomic specie) electron density of states plot"
419                + " saved..."
420            )
421
422        edos_proj_vs_atm_ssh_da = edos_proj_da.groupby("atm_ssh").sum(
423            dim=["atm_spc", "atm_i", "harm_poly"]
424        )
425        fig, ax = plt.subplots(tight_layout=True)
426        edos_proj_vs_atm_ssh_da.plot(x="energy", hue="atm_ssh", ax=ax)
427        ax.set_title("")
428        ax.set_xlabel(r"$\epsilon_{\text{KS}}$ / $\mathrm{eV}$")
429        ax.set_ylabel(r"$\text{Electron DOS}$ / $\mathrm{eV^{-1}}$")
430        legend = ax.get_legend()
431        legend.set_loc("upper left")
432        legend.set_bbox_to_anchor((1.0, 1.0))
433        legend.set_frame_on(False)
434        legend.set_title(None)
435        if config.edos_plot_prefix is not None:
436            fig.savefig(f"{config.edos_plot_prefix}_atm-ssh.svg")
437            parprint(
438                "Projected (atomic subshell) electron density of states plot"
439                + " saved..."
440            )
441
442
443def calc_edos_cli() -> None:
444    r"""Calculate electron density of states - CLI interface."""
445    parser = ArgumentParser(description="Calculate electron density of states")
446    parser.add_argument(
447        "config",
448        nargs="?",
449        default="./config.yml",
450        help="configuration file",
451    )
452    args = parser.parse_args()
453    with paropen(args.config, "r") as stream:
454        config = yaml.safe_load(stream)
455    config = Config(**config)
456    calc_edos(config)