madman.helpers.ase.workflows.eps

  1r"""Dielectric response."""
  2
  3from __future__ import annotations
  4import os
  5from argparse import ArgumentParser
  6from collections.abc import Mapping, Sequence
  7from functools import cache
  8from itertools import product
  9from math import ceil
 10from typing import Any, Literal, Optional, Union
 11
 12import numpy as np
 13import seaborn as sns
 14import xarray as xr
 15import yaml
 16from ase.parallel import paropen, parprint
 17from gpaw.bztools import find_high_symmetry_monkhorst_pack
 18from gpaw.calculator import GPAW
 19from gpaw.response.df import DielectricFunction
 20from matplotlib import pyplot as plt
 21from pydantic import BaseModel, Field, confloat, conint, constr, field_validator
 22from scipy.interpolate import Akima1DInterpolator
 23
 24from madman.helpers.ase.workflows.converge import MultivariateConvergence
 25from madman.utilities import gen_regular_grid, redirect_to
 26
 27
 28sns.set_theme(
 29    context="talk",
 30    style="white",
 31    rc={"figure.titlesize": "medium", "axes.formatter.useoffset": False},
 32)
 33
 34
 35class Config(BaseModel):
 36    r"""Configuration."""
 37
 38    ground_path: constr(pattern=r".*\.gpw$") = Field(frozen=True)
 39    r"""Path to electronic ground state."""
 40
 41    kptden: Union[
 42        ConvergenceParameterSettingsReal,
 43        Mapping,
 44        confloat(gt=0.0, allow_inf_nan=False),
 45    ] = Field(10.0, frozen=True)
 46    r"""Convergence settings for $\mathbf{k}$-point density [$\mathrm{Å}$]."""
 47
 48    nbands_empty: Union[
 49        ConvergenceParameterSettingsIntegral,
 50        Mapping,
 51        conint(ge=0),
 52    ] = Field(10, frozen=True)
 53    r"""Convergence settings for number of empty bands."""
 54
 55    domega0: Union[
 56        ConvergenceParameterSettingsReal,
 57        Mapping,
 58        confloat(gt=0.0, allow_inf_nan=False),
 59    ] = Field(0.05, frozen=True)
 60    r"""Convergence settings for $\Delta \omega_0$ [$\mathrm{eV}$]."""
 61
 62    omega2: Union[
 63        ConvergenceParameterSettingsReal,
 64        Mapping,
 65        confloat(gt=0.0, allow_inf_nan=False),
 66    ] = Field(10.0, frozen=True)
 67    r"""Convergence settings for $\omega_2$ [$\mathrm{eV}$]."""
 68
 69    ecut: Union[
 70        ConvergenceParameterSettingsReal,
 71        Mapping,
 72        confloat(gt=0.0, allow_inf_nan=False),
 73    ] = Field(50.0, frozen=True)
 74    r"""Convergence settings for energy cutoff [$\mathrm{eV}$]."""
 75
 76    energy_min: confloat(ge=0.0, allow_inf_nan=False) = Field(0.0, frozen=True)
 77    r"""Photon energy minimum [$\mathrm{eV}$]."""
 78
 79    energy_max: confloat(gt=0.0, allow_inf_nan=False) = Field(5.0, frozen=True)
 80    r"""Photon energy maximum [$\mathrm{eV}$]."""
 81
 82    energy_spc: confloat(gt=0.0, allow_inf_nan=False) = Field(0.01, frozen=True)
 83    r"""Photon energy spacing [$\mathrm{eV}$]."""
 84
 85    directions: Sequence[Literal["x", "y", "z"], ...] = Field(
 86        {"x", "y", "z"}, frozen=True
 87    )
 88    r"""Dielectric function directions."""
 89
 90    strict: bool = Field(False, frozen=True)
 91    r"""If True, max. slope is used for convergence; otherwise, avg. is used."""
 92
 93    niter_max: conint(gt=0) = Field(1, frozen=True)
 94    r"""Maximum number of iterations for convergence."""
 95
 96    ks_calc_prefix: Optional[str] = Field(None, frozen=True)
 97    r"""Prefix to save Kohn-Sham orbitals calculator."""
 98
 99    eps_save_prefix: Optional[str] = Field(None, frozen=True)
100    r"""Prefix to save dielectric function."""
101
102    eps_plot_prefix: Optional[str] = Field(None, frozen=True)
103    r"""Prefix to save dielectric function plot."""
104
105    conv_plot_root: Optional[str] = Field(None, frozen=True)
106    r"""Root to save convergence plots."""
107
108    log_prefix: Optional[str] = Field(None, frozen=True)
109    r"""Prefix to save GPAW log."""
110
111    cache_root: str = Field("./tmpdata/", frozen=True)
112    r"""Root to cache electron densitysity of states."""
113
114    @field_validator("ground_path")
115    def validate_ground_path(cls, value):
116        r"""Validate path to electronic ground state."""
117        if not os.path.isfile(value):
118            raise FileNotFoundError
119        return value
120
121    @field_validator("kptden", "domega0", "omega2", "ecut")
122    def parse_conv_param_cfg_real(cls, value):
123        r"""Parse convergence parameter settings (real-valued)."""
124        if isinstance(value, Mapping):
125            value = ConvergenceParameterSettingsReal(**value)
126        return value
127
128    @field_validator("nbands_empty")
129    def parse_conv_param_cfg_int(cls, value):
130        r"""Parse convergence parameter settings (integral-valued)."""
131        if isinstance(value, Mapping):
132            value = ConvergenceParameterSettingsIntegral(**value)
133        return value
134
135    @field_validator("directions", mode="before")
136    def parse_directions(cls, value):
137        if isinstance(value, str):
138            value = [c for c in value]
139        return sorted(set(value))
140
141    @field_validator(
142        "ks_calc_prefix",
143        "eps_save_prefix",
144        "eps_plot_prefix",
145        "conv_plot_root",
146        "log_prefix",
147        "cache_root",
148    )
149    def create_tree(cls, value):
150        r"""Create tree for prefix or path."""
151        if value:
152            tree = os.path.dirname(value)
153            if tree:
154                os.makedirs(tree, exist_ok=True)
155        return value
156
157
158class ConvergenceParameterSettings(BaseModel):
159    r"""Convergence parameter settings."""
160
161    threshold: Optional[confloat(gt=0.0, allow_inf_nan=False)] = Field(
162        None, frozen=True
163    )
164    r"""Convergence threshold."""
165
166    stability: Optional[conint(gt=0)] = Field(None, frozen=True)
167    r"""Convergence stability."""
168
169
170class ConvergenceParameterSettingsReal(ConvergenceParameterSettings):
171    r"""Convergence parameter settings (real-valued)."""
172
173    value_min: confloat(gt=0.0, allow_inf_nan=False) = Field(frozen=True)
174    r"""Parameter minimum value."""
175
176    value_max: Optional[confloat(gt=0.0, allow_inf_nan=False)] = Field(
177        None, frozen=True
178    )
179    r"""Parameter maximum value."""
180
181    value_spc: Optional[confloat(gt=0.0, allow_inf_nan=False)] = Field(
182        None, frozen=True
183    )
184    r"""Paramenter value spacing."""
185
186
187class ConvergenceParameterSettingsIntegral(ConvergenceParameterSettings):
188    r"""Convergence parameter settings (integral-valued)."""
189
190    value_min: conint(gt=0) = Field(frozen=True)
191    r"""Parameter minimum value."""
192
193    value_max: Optional[conint(gt=0)] = Field(None, frozen=True)
194    r"""Parameter maximum value."""
195
196    value_spc: Optional[conint(gt=0)] = Field(None, frozen=True)
197    r"""Paramenter value spacing."""
198
199
200def ret_params(config: Config) -> dict:
201    r"""Return convergence parameters.
202
203    Args:
204        config: Configuration.
205
206    Returns:
207        Convergence parameters.
208
209    Raises:
210        TypeError: If configuration is not a Config instance.
211    """
212    if not isinstance(config, Config):
213        raise TypeError("Not a Config instance!")
214
215    kptden_d = config.kptden
216    if isinstance(kptden_d, float):
217        kptden_d = {"value_min": kptden_d}
218    else:
219        kptden_d = kptden_d.dict()
220    kptden_d.update(
221        {
222            "direction": 1,
223            "criterion": f"slope-abs-{"max" if config.strict else "avg"}",
224            "symb": r"\lambda_{\mathbf{k}}",
225            "unit": r"Å",
226        }
227    )
228    nbands_empty_d = config.nbands_empty
229    if isinstance(nbands_empty_d, int):
230        nbands_empty_d = {"value_min": nbands_empty_d}
231    else:
232        nbands_empty_d = nbands_empty_d.dict()
233    nbands_empty_d.update(
234        {
235            "direction": 1,
236            "criterion": f"slope-abs-{"max" if config.strict else "avg"}",
237            "symb": r"N^{*}",
238            "unit": r"1",
239        }
240    )
241    domega0_d = config.domega0
242    if isinstance(domega0_d, float):
243        domega0_d = {"value_min": domega0_d}
244    else:
245        domega0_d = domega0_d.dict()
246    domega0_d.update(
247        {
248            "direction": -1,
249            "criterion": f"slope-abs-{"max" if config.strict else "avg"}",
250            "symb": r"\Delta{\omega_{0}}",
251            "unit": r"eV",
252        }
253    )
254    omega2_d = config.omega2
255    if isinstance(omega2_d, float):
256        omega2_d = {"value_min": omega2_d}
257    else:
258        omega2_d = omega2_d.dict()
259    omega2_d.update(
260        {
261            "direction": 1,
262            "criterion": f"slope-abs-{"max" if config.strict else "avg"}",
263            "symb": r"\omega_{2}",
264            "unit": r"eV",
265        }
266    )
267    ecut_d = config.ecut
268    if isinstance(ecut_d, float):
269        ecut_d = {"value_min": ecut_d}
270    else:
271        ecut_d = ecut_d.dict()
272    ecut_d.update(
273        {
274            "direction": 1,
275            "criterion": f"slope-abs-{"max" if config.strict else "avg"}",
276            "symb": r"\epsilon_{\text{cut}}",
277            "unit": r"eV",
278        }
279    )
280    params = {
281        "kptden": kptden_d,
282        "nbands_empty": nbands_empty_d,
283        "domega0": domega0_d,
284        "omega2": omega2_d,
285        "ecut": ecut_d,
286    }
287    return params
288
289
290def calc_eps(config: Config) -> None:
291    r"""Calculate dielectric function.
292
293    Args:
294        config: Configuration.
295
296    Raises:
297        TypeError: If configuration is not a Config instance.
298    """
299    if not isinstance(config, Config):
300        raise TypeError("Not a Config instance!")
301
302    with redirect_to(config.log_prefix, mode="w"):
303        pass
304    parprint("GPAW log initialized...")
305
306    cache_ks = {}
307    ks_iter_i = 0
308    parprint("Cache initialized...")
309
310    ground = GPAW(config.ground_path)
311    parprint("Electronic ground state imported...")
312
313    nelectrons = ground.get_number_of_electrons()
314    nbands_occupied = ceil(0.5 * nelectrons)
315    parprint(f"Number of occupied bands: {nbands_occupied}...")
316
317    params = ret_params(config)
318    parprint("Dielectric function convergence parameters set up...")
319
320    energy_grid = gen_regular_grid(
321        config.energy_min, config.energy_max, config.energy_spc
322    )
323    parprint("Photon energy grid set up...")
324
325    @cache
326    def objective(kptden, nbands_empty, domega0, omega2, ecut):
327        nonlocal ks_iter_i
328        nbands_empty = int(nbands_empty) # FIXME: Should be fixed in `convergence.py` module
329
330        key_ks = tuple(sorted([("kptden", kptden), ("nbands_empty", nbands_empty)]))
331        if key_ks not in cache_ks:
332            path_ks = os.path.join(config.cache_root, f"ks{ks_iter_i}.gpw")
333            nbands = nbands_occupied + nbands_empty
334            with redirect_to(config.log_prefix, mode="a"):
335                calc_ks = ground.fixed_density(kpts={"density": kptden})
336                calc_ks.write(path_ks, "all")
337            cache_ks.update({key_ks: path_ks})
338            ks_iter_i += 1
339        path_ks = cache_ks[key_ks]
340        freqs = {"type": "nonlinear", "domega0": domega0, "omega2": omega2}
341        with redirect_to(config.log_prefix, mode="a"):
342            calc_eps = DielectricFunction(
343                calc=path_ks,
344                frequencies=freqs,
345                ecut=ecut,
346                rate="eta",
347            )
348        # The dielectric function rows are sorted according to:
349        # - direction (x, y, z)
350        # - local field correction (0: off; 1: on)
351        # - real / imaginary part
352        eps_arr = []
353        for direction in config.directions:
354            with redirect_to(config.log_prefix, mode="a"):
355                eps = calc_eps.get_dielectric_function(
356                    direction=direction,
357                    filename=os.path.join(config.cache_root, "eps.csv"),
358                )
359            for has_loc_field_corr in [0, 1]:
360                eps_re = Akima1DInterpolator(
361                    x=calc_eps.get_frequencies(),
362                    y=eps[has_loc_field_corr].real,
363                    method="akima",
364                    extrapolate=False,
365                )(energy_grid)
366                eps_arr.append(eps_re)
367                eps_im = Akima1DInterpolator(
368                    x=calc_eps.get_frequencies(),
369                    y=eps[has_loc_field_corr].imag,
370                    method="akima",
371                    extrapolate=False,
372                )(energy_grid)
373                eps_arr.append(eps_im)
374        return np.stack(eps_arr, dtype=float)
375
376    parprint("Dielectric function objective function set up...")
377
378    obj_symb = [
379        f"{part}\\left[\\epsilon_{{{direction}}}^{{\\text{{{has_loc_field_corr}}}}}\\right]"
380        for direction in config.directions
381        for has_loc_field_corr in ["", "LFC"]
382        for part in ["\\Re", "\\Im"]
383    ]
384    convergence = MultivariateConvergence(
385        objective,
386        params,
387        crop=True,
388        req_sc=True,
389        niter_max=config.niter_max,
390        obj_symb=obj_symb,
391        obj_unit=r"1",
392    )
393    parprint("Dielectric function convergence set up...")
394
395    parprint("Dielectric function convergence started...")
396    convergence.run()
397    if convergence.converged is False:
398        parprint("Dielectric function convergence failed!")
399    else:
400        parprint("Dielectric function converged...")
401
402    params_opt = convergence.values_opt
403    eps_arr = objective(**params_opt)
404    parprint("Dielectric function calculated...")
405
406    if config.eps_save_prefix is not None:
407        eps_da = xr.DataArray(
408            eps_arr.reshape(len(config.directions), 2, 2, -1),
409            coords=[config.directions, [0, 1], ["re", "im"], energy_grid],
410            dims=["direction", "has_loc_field_corr", "part", "energy"],
411        )
412        eps_da.to_netcdf(f"{config.eps_save_prefix}.nc")
413        parprint("Dielectric function saved...")
414
415    if config.eps_plot_prefix is not None:
416        fig, ax = plt.subplots(tight_layout=True)
417        ax.set_xlabel(r"$\hbar \, \omega$ / $\mathrm{eV}$")
418        ax.set_ylabel(r"$\epsilon$")
419        for eps, label in zip(eps_arr, obj_symb):
420            ax.plot(energy_grid, eps, label=f"${label}$")
421        ax.legend(
422            loc="upper center",
423            bbox_to_anchor=[1.0, 1.0],
424            frameon=True,
425            ncols=len(config.directions),
426            edgecolor="k",
427            fontsize="xx-small",
428        )
429        fig.savefig(f"{config.eps_plot_prefix}.svg")
430        parprint("Dielectric function plot saved...")
431
432    if config.conv_plot_root is not None:
433        value_plots_map = convergence.plot("obj-value")
434        slope_plots_map = convergence.plot("obj-slope")
435
436        ndigits = len(str(config.niter_max))
437
438        for qnty, plots_map in zip(
439            ["eps", "meps"], [value_plots_map, slope_plots_map]
440        ):
441            for param, plots_seq in plots_map.items():
442                for i, plots in enumerate(plots_seq):
443                    for (direction, has_loc_field_corr, part), plot in zip(
444                        product(config.directions, [0, 1], ["re", "im"]), plots
445                    ):
446                        figdir = os.path.join(
447                            config.conv_plot_root,
448                            f"{qnty}-vs-{param}",
449                            f"{direction}-{has_loc_field_corr}-{part}",
450                        )
451                        os.makedirs(figdir, exist_ok=True)
452                        figpath = os.path.join(figdir, f"{i:0{ndigits}.0f}.svg")
453                        fig = plot.get_figure()
454                        fig.savefig(figpath)
455                        plt.close(fig)
456        parprint("Converge plots saved...")
457
458    return None
459
460def calc_eps_cli() -> None:
461    r"""Calculate electron density of states - CLI interface."""
462    parser = ArgumentParser(description="Calculate electron density of states")
463    parser.add_argument(
464        "config",
465        nargs="?",
466        default="./config.yml",
467        help="configuration file",
468    )
469    args = parser.parse_args()
470    with paropen(args.config, "r") as stream:
471        config = yaml.safe_load(stream)
472    config = Config(**config)
473    calc_eps(config)