Source code for gpaw.hybrids.eigenvalues

"""Calculate non self-consistent eigenvalues for hybrid functionals."""
from __future__ import annotations
import functools
import json
from pathlib import Path
from typing import Generator, List, Optional, Tuple, Union

import numpy as np
from ase.units import Ha
from gpaw.calculator import GPAW as GPAWOld
from gpaw import GPAW
from import ASECalculator
from gpaw.kpt_descriptor import KPointDescriptor
from gpaw.mpi import serial_comm
from import PWDescriptor
from import PWLFC
from gpaw.typing import Array3D
from gpaw.xc import XC
from gpaw.xc.kernel import XCNull
from import vxc

from . import parse_name
from .coulomb import coulomb_interaction
from .kpts import RSKPoint, get_kpt, to_real_space
from .paw import calculate_paw_stuff
from .symmetry import Symmetry

[docs]def non_self_consistent_eigenvalues( calc: Union[GPAWOld, ASECalculator, str, Path], xcname: str, n1: int = 0, n2: int = 0, kpt_indices: List[int] = None, snapshot: Union[str, Path] = None, ftol: float = 1e-9) -> tuple[Array3D, Array3D, Array3D]: """Calculate non self-consistent eigenvalues for a hybrid functional. Based on a self-consistent DFT calculation (calc). Only eigenvalues n1 to n2 - 1 for the IBZ indices in kpt_indices are calculated (default is all bands and all k-points). EXX integrals involving states with occupation numbers less than ftol are skipped. Use snapshot='name.json' to get snapshots for each k-point finished. Returns three (nspins, nkpts, n2 - n1)-shaped ndarrays with contributions to the eigenvalues in eV: >>> nsceigs = non_self_consistent_eigenvalues >>> eig_dft, vxc_dft, vxc_hyb = nsceigs('<gpw-file>', xcname='PBE0') >>> eig_hyb = eig_dft - vxc_dft + vxc_hyb """ if not isinstance(calc, (GPAWOld, ASECalculator)): if calc == '<gpw-file>': # for doctest return (np.zeros((1, 1, 1)), np.zeros((1, 1, 1)), np.zeros((1, 1, 1))) calc = GPAW(Path(calc), txt=None, parallel={'band': 1, 'kpt': 1}) assert isinstance(calc, (GPAWOld, ASECalculator)) wfs = calc.wfs if n2 <= 0: n2 += if kpt_indices is None: kpt_indices = np.arange(wfs.kd.nibzkpts).tolist() path = Path(snapshot) if snapshot is not None else None e_dft_sin = np.zeros(0) v_dft_sin = np.zeros(0) # sl=semilocal, nl=nonlocal v_hyb_sl_sin = np.zeros(0) v_hyb_nl_sin: Optional[List[List[np.ndarray]]] = None if path: e_dft_sin, v_dft_sin, v_hyb_sl_sin, v_hyb_nl_sin = read_snapshot(path) xcname, exx_fraction, omega = parse_name(xcname) if v_dft_sin.size == 0: xc = XC(xcname) e_dft_sin, v_dft_sin, v_hyb_sl_sin = _semi_local( calc, xc, n1, n2, kpt_indices) write_snapshot(e_dft_sin, v_dft_sin, v_hyb_sl_sin, v_hyb_nl_sin, path, # Non-local hybrid contribution if v_hyb_nl_sin is None: v_hyb_nl_sin = [[] for s in range(wfs.nspins)] # Find missing indices: kpt_indices_s = [kpt_indices[len(v_hyb_nl_in):] for v_hyb_nl_in in v_hyb_nl_sin] if any(len(kpt_indices) > 0 for kpt_indices in kpt_indices_s): for s, v_hyb_nl_n in _non_local(calc, n1, n2, kpt_indices_s, ftol, omega): v_hyb_nl_sin[s].append(v_hyb_nl_n * exx_fraction) write_snapshot(e_dft_sin, v_dft_sin, v_hyb_sl_sin, v_hyb_nl_sin, path, return (e_dft_sin * Ha, v_dft_sin * Ha, (v_hyb_sl_sin + v_hyb_nl_sin) * Ha)
def _semi_local(calc: GPAWOld | ASECalculator, xc, n1: int, n2: int, kpt_indices: List[int] ) -> Tuple[np.ndarray, np.ndarray, np.ndarray]: wfs = calc.wfs nspins = wfs.nspins e_dft_sin = np.array([[calc.get_eigenvalues(k, spin)[n1:n2] for k in kpt_indices] for spin in range(nspins)]) v_dft_sin = vxc(calc.gs_adapter(), n1=n1, n2=n2)[:, kpt_indices] if isinstance(xc.kernel, XCNull): v_hyb_sl_sin = np.zeros_like(v_dft_sin) else: v_hyb_sl_sin = vxc(calc.gs_adapter(), xc, n1=n1, n2=n2)[:, kpt_indices] return e_dft_sin / Ha, v_dft_sin / Ha, v_hyb_sl_sin / Ha def _non_local(calc: GPAWOld | ASECalculator, n1: int, n2: int, kpt_indices_s: List[List[int]], ftol: float, omega: float) -> Generator[Tuple[int, np.ndarray], None, None]: wfs = calc.wfs kd = wfs.kd dens = calc.density nocc = max(((kpt.f_n / kpt.weight) > ftol).sum() for kpt in wfs.kpt_u) nocc = kd.comm.max_scalar( coulomb = coulomb_interaction(omega,, kd) sym = Symmetry(kd) paw_s = calculate_paw_stuff(wfs, dens) for spin, kpt_indices in enumerate(kpt_indices_s): if len(kpt_indices) == 0: continue kpts2 = [get_kpt(wfs, k, spin, 0, nocc) for k in range(kd.nibzkpts)] for i in kpt_indices: kpt1 = get_kpt(wfs, i, spin, n1, n2) v_n = _calculate_eigenvalues( kpt1, kpts2, paw_s[spin], kd, coulomb, sym, wfs, calc.spos_ac) yield spin, v_n def _calculate_eigenvalues(kpt1, kpts2, paw, kd, coulomb, sym, wfs, spos_ac): pd = kpt1.psit.pd gd = comm = size = comm.size rank = comm.rank nsym = len(kd.symmetry.op_scc) assert len(kpts2) == kd.nibzkpts N1 = len(kpt1.psit.array) N2 = len(kpts2[0].psit.array) size1, size2 = layout(N1, N2, size) assert size1 * size2 == size B1 = (N1 + size1 - 1) // size1 B2 = (N2 + size2 - 1) // size2 rank1, rank2 = divmod(rank, size2) n1a = min(B1 * rank1, N1) n1b = min(n1a + B1, N1) n2a = min(B2 * rank2, N2) n2b = min(n2a + B2, N2) u1_nR = to_real_space(kpt1.psit, n1a, n1b) proj1all = kpt1.proj.broadcast() proj1 = proj1all.view(n1a, n1b) E_n = np.zeros(N1) e_n = E_n[n1a:n1b] e_nn = np.empty((n1b - n1a, n2b - n2a)) for i2, kpt2 in enumerate(kpts2): u2_nR = to_real_space(kpt2.psit, n2a, n2b) rskpt0 = RSKPoint(u2_nR, kpt2.proj.broadcast().view(n2a, n2b), kpt2.f_n[n2a:n2b], kpt2.k_c, kpt2.weight) for k, i in enumerate(kd.bz2ibz_k): if i != i2: continue s = kd.sym_k[k] + kd.time_reversal_k[k] * nsym rskpt2 = sym.apply_symmetry(s, rskpt0, wfs, spos_ac) q_c = rskpt2.k_c - kpt1.k_c qd = KPointDescriptor([-q_c]) pd12 = PWDescriptor(pd.ecut, gd, pd.dtype, kd=qd) ghat = PWLFC([data.ghat_l for data in wfs.setups], pd12) ghat.set_positions(spos_ac) v_G = coulomb.get_potential(pd12) Q_annL = [np.einsum('mi, ijL, nj -> mnL', proj1[a], Delta_iiL, rskpt2.proj[a].conj()) for a, Delta_iiL in enumerate(paw.Delta_aiiL)] rho_nG = ghat.pd.empty(n2b - n2a, u1_nR.dtype) for n1, u1_R in enumerate(u1_nR): for u2_R, rho_G in zip(rskpt2.u_nR, rho_nG): rho_G[:] = ghat.pd.fft(u1_R * u2_R.conj()) ghat.add(rho_nG, {a: Q_nnL[n1] for a, Q_nnL in enumerate(Q_annL)}) for n2, rho_G in enumerate(rho_nG): vrho_G = v_G * rho_G e = ghat.pd.integrate(rho_G, vrho_G).real e_nn[n1, n2] = e / kd.nbzkpts e_n -= for a, VV_ii in paw.VV_aii.items(): P_ni = proj1all[a] vv_n = np.einsum('ni, ij, nj -> n', P_ni.conj(), VV_ii, P_ni).real if paw.VC_aii[a] is not None: vc_n = np.einsum('ni, ij, nj -> n', P_ni.conj(), paw.VC_aii[a], P_ni).real else: vc_n = 0.0 E_n -= (2 * vv_n + vc_n) return E_n def write_snapshot(e_dft_sin: np.ndarray, v_dft_sin: np.ndarray, v_hyb_sl_sin: np.ndarray, v_hyb_nl_sin: Optional[List[List[np.ndarray]]], path: Optional[Path], comm) -> None: """Write to json-file what has been calculated so far.""" if comm.rank == 0 and path: dct = {'e_dft_sin': e_dft_sin.tolist(), 'v_dft_sin': v_dft_sin.tolist(), 'v_hyb_sl_sin': v_hyb_sl_sin.tolist()} if v_hyb_nl_sin is not None: dct['v_hyb_nl_sin'] = [[v_n.tolist() for v_n in v_in] for v_in in v_hyb_nl_sin] path.write_text(json.dumps(dct, indent=0)) def read_snapshot(snapshot: Path ) -> Tuple[np.ndarray, np.ndarray, np.ndarray, Optional[List[List[np.ndarray]]]]: """Read from json-file what has already been calculated.""" if snapshot.is_file(): dct = json.loads(snapshot.read_text()) v_hyb_nl_sin = dct.get('v_hyb_nl_sin') if v_hyb_nl_sin is not None: v_hyb_nl_sin = [[np.array(v_n) for v_n in v_in] for v_in in v_hyb_nl_sin] return (np.array(dct['e_dft_sin']), np.array(dct['v_dft_sin']), np.array(dct['v_hyb_sl_sin']), v_hyb_nl_sin) return np.array([[[]]]), np.array([[[]]]), np.array([[[]]]), None @functools.lru_cache() def layout(n1: int, n2: int, size: int) -> Tuple[int, int]: """Distribute n1*n2 matrix over s1*s2=size blocks. Returns s1, s2. >>> layout(10, 10, 8) (4, 2) """ candidates: List[Tuple[float, int, int]] = [] for s1 in range(1, size + 1): s2, r = divmod(size, s1) if r > 0: continue fitness = (1 - idle(n1, s1)) * (1 - idle(n2, s2)) candidates.append((fitness, s1, s2)) return max(candidates)[1:] def idle(n: int, s: int) -> float: """Idle fraction (helper function for layout() function).""" b = (n + s - 1) // s return 1 - n / (b * s)