madman.helpers.ase.workflows.ground

  1r"""Electronic ground state."""
  2
  3import os
  4from argparse import ArgumentParser
  5from collections.abc import Mapping
  6from functools import cache
  7from numbers import Real
  8from typing import Any, Literal, Optional, Union
  9
 10import numpy as np
 11import seaborn as sns
 12import yaml
 13from ase import Atoms
 14from ase.filters import FrechetCellFilter
 15from ase.io import Trajectory, read
 16from ase.optimize import BFGS
 17from ase.parallel import paropen, parprint, world
 18from gpaw.calculator import GPAW
 19from gpaw.occupations import FermiDirac
 20from gpaw.wavefunctions.pw import PW
 21from matplotlib import pyplot as plt
 22from pydantic import BaseModel, Field, confloat, conint, validator
 23
 24from madman.helpers.ase.workflows.converge import MultivariateConvergence
 25from madman.utilities import 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 ConvergenceParameterSettings(BaseModel):
 36    r"""Total energy convergence parameter settings."""
 37
 38    value_min: confloat(ge=0.0, allow_inf_nan=False) = Field(frozen=True)
 39    r"""Parameter minimum value."""
 40
 41    value_max: Optional[confloat(ge=0.0, allow_inf_nan=False)] = Field(
 42        None, frozen=True
 43    )
 44    r"""Parameter maximum value."""
 45
 46    value_spc: Optional[confloat(gt=0.0, allow_inf_nan=False)] = Field(
 47        None, frozen=True
 48    )
 49    r"""Pamater value spacing."""
 50
 51    threshold: Optional[confloat(gt=0.0, allow_inf_nan=False)] = Field(
 52        None, frozen=True
 53    )
 54    r"""Convergence threshold."""
 55
 56    stability: Optional[conint(gt=0)] = Field(None, frozen=True)
 57    r"""Convergence stability."""
 58
 59
 60class ConvergenceSettings(BaseModel):
 61    r"""Total energy convergence settings."""
 62
 63    basis: Literal["dzp"] = Field("dzp", frozen=True)
 64    r"""LCAO basis set."""
 65
 66    ecut: Optional[
 67        Union[
 68            ConvergenceParameterSettings, confloat(gt=0.0, allow_inf_nan=False)
 69        ]
 70    ] = Field(None, frozen=True)
 71    r"""Energy cutoff [$\mathrm{eV}$]."""
 72
 73    kptden: Union[
 74        ConvergenceParameterSettings, confloat(gt=0.0, allow_inf_nan=False)
 75    ] = Field(1.0, frozen=True)
 76    r"""$\mathbf{k}$-point density [$\mathrm{Å}$]."""
 77
 78    smear: Union[
 79        ConvergenceParameterSettings,
 80        confloat(ge=0.0, allow_inf_nan=False),
 81    ] = Field(0.0, frozen=True)
 82    r"""Smearing [$\mathrm{eV}$]."""
 83
 84    niter_max: conint(gt=0) = Field(1, frozen=True)
 85    r"""Maximum number of iterations."""
 86
 87    fig_root: Optional[str] = Field(None, frozen=True)
 88    r"""Root to save convergence plots."""
 89
 90    cache_root: Optional[str] = Field(None, frozen=True)
 91    r"""Root to cache electronic ground states."""
 92
 93    @validator("ecut", "kptden", "smear", pre=True)
 94    def parse_param_map(cls, value: Any) -> Any:
 95        r"""Parse convergence parameter map."""
 96        if isinstance(value, Mapping):
 97            value = ConvergenceParameterSettings(**value)
 98        return value
 99
100    @validator("fig_root", "cache_root")
101    def create_root(cls, root: Any) -> Any:
102        r"""Create root, if needed."""
103        if isinstance(root, str) and root != "":
104            os.makedirs(root, exist_ok=True)
105        return root
106
107
108class RelaxationSettings(BaseModel):
109    r"""Atomic structure relaxation settings."""
110
111    algorithm: Optional[Literal["single-point", "iterative"]] = Field(
112        None, frozen=True
113    )
114    r"""Relaxation algorithm."""
115
116    force_max: confloat(gt=0.0, allow_inf_nan=False) = Field(0.01, frozen=True)
117    r"""Maximum force for equilibrium [$\mathrm{eV \, Å^{-1}}$]."""
118
119    niter_max: conint(gt=0) = Field(1, frozen=True)
120    r"""Maximum number of iterations."""
121
122    rlx_prefix: Optional[str] = Field(None, frozen=True)
123    r"""Prefix to save relaxed atomic structure."""
124
125    hst_prefix: Optional[str] = Field(None, frozen=True)
126    r"""Prefix to save relaxation history."""
127
128    fig_prefix: Optional[str] = Field(None, frozen=True)
129    r"""Prefix to save relaxation plot."""
130
131    @validator("rlx_prefix", "hst_prefix", "fig_prefix")
132    def create_tree(cls, prefix: Any) -> Any:
133        r"""Create prefix tree, if needed."""
134        if isinstance(prefix, str):
135            prefix_dir = os.path.dirname(prefix)
136            if prefix_dir != "":
137                os.makedirs(prefix_dir, exist_ok=True)
138        return prefix
139
140
141class ElectronicGroundState:
142    r"""Electronic ground state."""
143
144    def __init__(
145        self,
146        atm_str: Atoms,
147        *,
148        charge: Real = 0.0,
149        xc: Literal["LDA", "PBE"] = "PBE",
150        conv_cfg: Union[ConvergenceSettings, Mapping],
151        relax_cfg: Union[RelaxationSettings, Mapping],
152        parallel: Optional[Mapping] = None,
153        egs_prefix: Optional[str] = None,
154        log_prefix: Optional[str] = None,
155    ) -> None:
156        r"""Initialize ElectronicGroundState object.
157
158        Args:
159            atm_str: Atomic structure.
160            charge: Charge [$q$].
161            xc: Exchange-correlation functional.
162            conv_cfg: Total energy convergence settings.
163            relax_cfg: Atomic structure relaxation settings.
164            parallel: GPAW parallelization options.
165            egs_prefix: Prefix to save electronic ground state.
166            log_prefix: Prefix to save GPAW log.
167
168        Raises:
169            TypeError: If atomic structure is not a Atoms object.
170            TypeError: If charge is not real.
171            ValueError: If exchange-correlation functional is invalid.
172        """
173        if not isinstance(atm_str, Atoms):
174            raise TypeError("Invalid atomic structure!")
175        self._atm_str = atm_str
176
177        if not isinstance(charge, Real):
178            raise TypeError("Invalid charge!")
179        self._charge = float(charge)
180
181        if xc not in {"LDA", "PBE"}:
182            raise ValueError("Invalid exchange-correlation functional!")
183        self._xc = xc
184
185        if isinstance(conv_cfg, Mapping):
186            conv_cfg = ConvergenceSettings(**conv_cfg)
187        if not isinstance(conv_cfg, ConvergenceSettings):
188            raise TypeError("Invalid total energy convergence settings!")
189        self._conv_cfg = conv_cfg
190
191        if isinstance(relax_cfg, Mapping):
192            relax_cfg = RelaxationSettings(**relax_cfg)
193        if not isinstance(relax_cfg, RelaxationSettings):
194            raise TypeError("Invalid relaxation settings!")
195        self._relax_cfg = relax_cfg
196
197        if parallel is not None and not isinstance(parallel, Mapping):
198            raise TypeError("Invalid GPAW parallelization options!")
199        self._parallel = None if parallel is None else dict(parallel)
200
201        if egs_prefix is not None and not isinstance(egs_prefix, str):
202            raise TypeError("Invalid prefix to save electronic ground state!")
203        self._egs_prefix = egs_prefix
204
205        if log_prefix is not None and not isinstance(log_prefix, str):
206            raise TypeError("Invalid prefix to save GPAW log!")
207        self._log_prefix = log_prefix
208
209        for prefix in [self._egs_prefix, self._log_prefix]:
210            if isinstance(prefix, str):
211                prefix_dir = os.path.dirname(prefix)
212                if prefix_dir != "":
213                    os.makedirs(prefix_dir, exist_ok=True)
214        with redirect_to(self._log_prefix, mode="w"):
215            pass
216
217        self._cache = {}
218        self._history = {}
219        self._etot_iter_i = 0
220        self._relax_iter_i = 0
221
222        kptden_d = self._conv_cfg.kptden
223        if isinstance(kptden_d, float):
224            kptden_d = {"value_min": kptden_d}
225        else:
226            kptden_d = kptden_d.dict()
227        kptden_d.update(
228            {
229                "direction": 1,
230                "criterion": "slope-abs",
231                "symb": r"\lambda_{\mathbf{k}}",
232                "unit": r"Å",
233            }
234        )
235        smear_d = self._conv_cfg.smear
236        if isinstance(smear_d, float):
237            smear_d = {"value_min": smear_d}
238        else:
239            smear_d = smear_d.dict()
240        smear_d.update(
241            {
242                "direction": 1,
243                "criterion": "1st-min",
244                "symb": r"\epsilon_{\text{smear}}",
245                "unit": r"eV",
246            }
247        )
248        ecut_d = self._conv_cfg.ecut
249
250        if ecut_d is not None:
251            if isinstance(ecut_d, float):
252                ecut_d = {"value_min": ecut_d}
253            else:
254                ecut_d = ecut_d.dict()
255            ecut_d.update(
256                {
257                    "direction": 1,
258                    "criterion": "slope-neg",
259                    "symb": r"\epsilon_{\text{cut}}",
260                    "unit": r"eV",
261                }
262            )
263            params = {"ecut": ecut_d, "kptden": kptden_d, "smear": smear_d}
264
265            @cache
266            def objective(ecut, kptden, smear):
267                atm_str = self._atm_str.copy()
268                with redirect_to(self._log_prefix, mode="a"):
269                    etot = type(self).calc_etot(
270                        atm_str,
271                        charge=self._charge,
272                        xc=self._xc,
273                        ecut=ecut,
274                        kptden=kptden,
275                        smear=smear,
276                        parallel=self._parallel,
277                    )
278                self._etot_iter_i += 1
279                key = [("ecut", ecut), ("kptden", kptden), ("smear", smear)]
280                if self._conv_cfg.cache_root is None:
281                    self._cache.update({tuple(sorted(key)): atm_str.calc})
282                else:
283                    egs_path = os.path.join(
284                        self._conv_cfg.cache_root,
285                        f"{self._etot_iter_i}.gpw",
286                    )
287                    atm_str.calc.write(egs_path)
288                    self._cache.update({tuple(sorted(key)): egs_path})
289                return etot
290
291        else:
292            params = {"kptden": kptden_d, "smear": smear_d}
293
294            @cache
295            def objective(kptden, smear):
296                atm_str = self._atm_str.copy()
297                with redirect_to(self._log_prefix, mode="a"):
298                    etot = type(self).calc_etot(
299                        atm_str,
300                        charge=self._charge,
301                        xc=self._xc,
302                        basis=self._conv_cfg.basis,
303                        kptden=kptden,
304                        smear=smear,
305                        parallel=self._parallel,
306                    )
307                self._etot_iter_i += 1
308                key = [("kptden", kptden), ("smear", smear)]
309                if self._conv_cfg.cache_root is None:
310                    self._cache.update({tuple(sorted(key)): atm_str.calc})
311                else:
312                    egs_path = os.path.join(
313                        self._conv_cfg.cache_root,
314                        f"{self._etot_iter_i}.gpw",
315                    )
316                    atm_str.calc.write(egs_path)
317                    self._cache.update({tuple(sorted(key)): egs_path})
318                return etot
319
320        self._convergence = MultivariateConvergence(
321            objective,
322            params,
323            crop=False,
324            req_sc=True,
325            niter_max=self._conv_cfg.niter_max,
326            obj_symb=r"E_{\text{tot}}",
327            obj_unit=r"eV",
328        )
329        self._relaxation = None
330
331    @property
332    def atm_str(self) -> Atoms:
333        r"""Atomic structure."""
334        return self._atm_str
335
336    @property
337    def charge(self) -> float:
338        r"""Change [$q$]."""
339        return self._charge
340
341    @property
342    def xc(self) -> str:
343        r"""Exchange-correlation functional."""
344        return self._xc
345
346    @property
347    def conv_cfg(self) -> ConvergenceSettings:
348        r"""Total energy convergence settings."""
349        return self._conv_cfg
350
351    @property
352    def relax_cfg(self) -> RelaxationSettings:
353        r"""Atomic structure relaxation settings."""
354        return self._relax_cfg
355
356    @property
357    def parallel(self) -> dict:
358        r"""GPAW parallelization options."""
359        return self._parallel
360
361    @property
362    def egs_prefix(self) -> Optional[str]:
363        r"""Prefix to save electronic ground state."""
364        return self._egs_prefix
365
366    @property
367    def log_prefix(self) -> Optional[str]:
368        r"""Prefix to save GPAW logs."""
369        return self._log_prefix
370
371    @property
372    def cache(self) -> dict[tuple[tuple[str, float], ...], Union[GPAW, str]]:
373        r"""Electronic groud state cache."""
374        return self._cache
375
376    @property
377    def history(self) -> Optional[dict[int, list[FrechetCellFilter, ...]]]:
378        r"""Relaxation history."""
379        return self._history
380
381    @property
382    def etot_iter_i(self) -> int:
383        r"""Total energy iteration index."""
384        return self._etot_iter_i
385
386    @property
387    def relax_iter_i(self) -> int:
388        r"""Relaxation iteration index."""
389        return self._relax_iter_i
390
391    @property
392    def convergence(self) -> MultivariateConvergence:
393        r"""Total energy convergence."""
394        return self._convergence
395
396    @property
397    def relaxation(self) -> Optional[BFGS]:
398        r"""Atomic structure relaxation."""
399        return self._relaxation
400
401    def converge(self) -> None:
402        r"""Converge total energy.
403
404        Returns:
405            True if total energy convergence succeeded.
406        """
407        self._convergence.run()
408        params_opt = self._convergence.values_opt
409
410        key = tuple(sorted(params_opt.items()))
411        if key not in self._cache:
412            self._convergence.objective(**params_opt)
413        if self._conv_cfg.cache_root is None:
414            self._atm_str.calc = self._cache[key]
415        else:
416            self._atm_str.calc = GPAW(self._cache[key])
417        self.save_ground()
418
419        if self._convergence.converged is False:
420            parprint("Total energy convergence failed!")
421        else:
422            parprint("Total energy converged...")
423        self._relaxation = BFGS(
424            atoms=FrechetCellFilter(self._atm_str),
425            logfile="-",
426            trajectory=(
427                None
428                if self.relax_cfg.hst_prefix is None
429                else f"{self.relax_cfg.hst_prefix}.traj"
430            ),
431        )
432        return None
433
434
435    def relax(self) -> None:
436        r"""Relax atomic structure.
437
438        Returns:
439            True if atomic structure relaxation succeeded.
440        """
441        if self._relaxation is None:
442            parprint("Please, converge first! Skipping...")
443            return None
444
445        with redirect_to(self._log_prefix, mode="a"):
446            parprint("Relaxing atomic structure...")
447            relaxed = self._relaxation.run(fmax=self._relax_cfg.force_max)
448
449        self._atm_str = self._relaxation.atoms.atoms
450        self.save_relaxed()
451        self.save_ground()
452
453        if self._conv_cfg.cache_root is not None:
454            if world.rank == 0:
455                for egs_path in self._cache.values():
456                    os.remove(egs_path)
457
458        self._cache = {}
459        self._convergence.objective.cache_clear()
460
461        if isinstance(self._history, dict):
462            self._history.update({self._relax_iter_i: []})
463            for atm_str in Trajectory(f"{self.relax_cfg.hst_prefix}.traj"):
464                atm_str = FrechetCellFilter(atm_str)
465                self._history[self._relax_iter_i].append(atm_str)
466
467        if relaxed is False:
468            parprint("Atomic structure relaxation failed!")
469        else:
470            parprint("Atomic structure relaxed...")
471        return None
472
473    def calculate(self) -> None:
474        r"""Calculate electronic ground state.
475
476        Raises:
477            ValueError: If algorithm is invalid.
478        """
479        if self._relax_cfg.algorithm is None:
480            _ = self.converge()
481            self.plot_convergence()
482            return None
483
484        if self._relax_cfg.algorithm == "single-point":
485            self.converge()
486            self.plot_convergence()
487            self.relax()
488            self.plot_relaxation()
489            return None
490
491        if self._relax_cfg.algorithm == "iterative":
492            prev_values_opt = None
493            while self._relax_iter_i < self._relax_cfg.niter_max:
494                self.converge()
495                self.plot_convergence()
496                if self._convergence.values_opt == prev_values_opt:
497                    parprint("Self-consistency achieved...")
498                    self.plot_relaxation()
499                    return None
500                self.relax()
501                self.plot_relaxation()
502                prev_values_opt = self._convergence.values_opt.copy()
503                self._relax_iter_i += 1
504            parprint("Self-consistency not achived within iterations limit!")
505            return None
506
507        raise ValueError("Unknown algorithm!")
508
509    def save_ground(self) -> None:
510        r"""Save electronic ground state."""
511        if self._egs_prefix is None:
512            parprint(
513                "Prefix to save electronic ground state not specified!"
514                + " Not saving..."
515            )
516            return None
517
518        self._atm_str.calc.write(f"{self._egs_prefix}.gpw")
519        parprint("Electronic ground state saved...")
520        return None
521
522    def save_relaxed(self) -> None:
523        r"""Save relaxed atomic structure."""
524        if self._relax_cfg.rlx_prefix is None:
525            parprint(
526                "Prefix to save relaxed atomic structure not specified!"
527                + " Not saving..."
528            )
529            return None
530
531        for ext in [".traj", ".cif"]:
532            self._atm_str.write(f"{self._relax_cfg.rlx_prefix}{ext}")
533        parprint("Relaxed atomic structure saved...")
534        return None
535
536    def plot_convergence(self) -> None:
537        r"""Plot convergence at current relaxation step."""
538        if world.rank == 0:
539            value_plots_map = self._convergence.plot("obj-value")
540            slope_plots_map = self._convergence.plot("obj-slope")
541
542            if self._conv_cfg.fig_root is not None:
543                rnd = len(str(self._relax_cfg.niter_max))
544                cnd = len(str(self._conv_cfg.niter_max))
545
546                for qnty, plots_map in zip(
547                    ["etot", "metot"], [value_plots_map, slope_plots_map]
548                ):
549                    for param, plots_seq in plots_map.items():
550                        figdir = os.path.join(
551                            self._conv_cfg.fig_root, f"{qnty}-vs-{param}"
552                        )
553                        os.makedirs(figdir, exist_ok=True)
554                        for i, plots in enumerate(plots_seq):
555                            (plot,) = tuple(plots)
556                            fn = f"{self._relax_iter_i:0{rnd}.0f}_{i:0{cnd}.0f}"
557                            figpath = os.path.join(figdir, f"{fn}.svg")
558                            fig = plot.get_figure()
559                            fig.savefig(figpath)
560                            plt.close(fig)
561
562    def plot_relaxation(self) -> Optional[plt.Axes]:
563        r"""Plot relaxation.
564
565        Returns:
566            Plot figure.
567        """
568        if world.rank == 0:
569            if self._history is None:
570                parprint("Relaxation history not saved! Skipping...")
571                return None
572
573            fig, ax = plt.subplots(tight_layout=True)
574            ax.set_xlabel(r"#")
575            ax.set_ylabel(r"$F_{\text{max}}$ / $\mathrm{eV \, Å^{-1}}$")
576            ax.set_yscale("log")
577
578            f_max_vs_step = []
579            relax_iter_start_step = 0
580            for _, atm_str_vs_step in sorted(self._history.items()):
581                if relax_iter_start_step > 0:
582                    ax.axvline(relax_iter_start_step - 0.5, c="k", alpha=0.5)
583                for atm_str in atm_str_vs_step:
584                    forces = atm_str.get_forces()
585                    f_max = np.sqrt((forces**2).sum(axis=1).max())
586                    f_max_vs_step.append(f_max)
587                relax_iter_start_step += len(atm_str_vs_step)
588            ax.plot(f_max_vs_step, "o-", color="k")
589            ax.axhline(self._relax_cfg.force_max, ls="--", c="k")
590
591            if self._relax_cfg.fig_prefix is not None:
592                fig.savefig(f"{self._relax_cfg.fig_prefix}.svg")
593            return None
594
595    @staticmethod
596    def calc_etot(
597        atm_str: Atoms,
598        *,
599        charge: Real = 0.0,
600        xc: str = "PBE",
601        basis: Mapping | str = "dzp",
602        ecut: Optional[Real] = None,
603        kptden: Real = 1.0,
604        smear: Real = 0.0,
605        parallel: Optional[Mapping] = None,
606    ) -> float:
607        r"""Calculate total energy.
608
609        Args:
610            atm_str: Atomic structure.
611            charge: Charge [$q$].
612            xc: Exchange-correlation functional.
613            basis: LCAO basis set.
614            ecut: Energy cutoff [$\mathrm{eV}$].
615            kptden: $\mathbf{k}$-point density [$\mathrm{Å}$].
616            smear: Smearing [$\mathrm{eV}$].
617            parallel: GPAW parallel options.
618
619        Returns:
620            Total energy [eV].
621        """
622        gpaw_kwargs = {
623            "charge": charge,
624            "xc": xc,
625            "kpts": {"density": kptden, "even": True},
626            "occupations": FermiDirac(smear),
627            "parallel": parallel,
628        }
629        if ecut is not None:
630            gpaw_kwargs.update({"mode": PW(ecut)})
631        else:
632            gpaw_kwargs.update({"mode": "lcao", "basis": basis})
633        atm_str.calc = GPAW(**gpaw_kwargs)
634        etot = atm_str.get_potential_energy()
635        return etot
636
637
638def calc_ground_cli() -> None:
639    r"""Calculate electronic ground state - CLI interface."""
640    parser = ArgumentParser(description="Calculate electronic ground state")
641    parser.add_argument(
642        "config",
643        nargs="?",
644        default="./config.yml",
645        help="configuration file",
646    )
647    args = parser.parse_args()
648
649    with paropen(args.config, "r") as stream:
650        config = yaml.safe_load(stream)
651
652    atm_str_path = config.pop("atm_str_path")
653    atm_str = read(atm_str_path)
654
655    ground = ElectronicGroundState(atm_str, **config)
656    ground.calculate()