"""
== ==========
R
r
G
g
h
x r or h
== ==========
"""
from __future__ import annotations
from collections import defaultdict
from typing import DefaultDict
import numpy as np
from gpaw.core.arrays import DistributedArrays
from gpaw.core.atom_arrays import AtomArrays
from gpaw.core.uniform_grid import UGArray
from gpaw.new import trace, zips
from gpaw.new.potential import Potential
from gpaw.new.xc import Functional
from gpaw.setup import Setup
from gpaw.spinorbit import soc as soc_terms
from gpaw.typing import Array1D, Array2D, Array3D
from gpaw.utilities import pack_hermitian, pack_density, unpack_hermitian
from gpaw.yml import indent
from gpaw.mpi import MPIComm, serial_comm
from gpaw.new.external_potential import ExternalPotential
[docs]class PotentialCalculator:
def __init__(self,
xc: Functional,
poisson_solver,
setups: list[Setup],
*,
fracpos_ac: Array2D,
external_potential: ExternalPotential | None = None,
soc: bool = False):
self.poisson_solver = poisson_solver
self.xc = xc
self.setups = setups
self.external_potential = external_potential or ExternalPotential()
self.fracpos_ac = fracpos_ac
self.soc = soc
def __str__(self):
return (f'{self.poisson_solver}\n'
f'xc functional:\n{indent(self.xc)}\n')
[docs] def calculate_pseudo_potential(self,
density,
ibzwfs,
vHt_x: DistributedArrays | None
) -> tuple[dict[str, float],
UGArray,
UGArray,
DistributedArrays]:
raise NotImplementedError
[docs] def calculate_charges(self, vHt_x):
raise NotImplementedError
[docs] def restrict(self, a_r, a_R=None):
raise NotImplementedError
[docs] def calculate_without_orbitals(self,
density,
ibzwfs=None,
vHt_x: DistributedArrays | None = None,
kpt_band_comm: MPIComm | None = None
) -> tuple[Potential, AtomArrays]:
xc = self.xc
if xc.exx_fraction != 0.0:
from gpaw.new.xc import create_functional
self.xc = create_functional('PBE', xc.grid)
potential, Q_al = self.calculate(density, ibzwfs, vHt_x, kpt_band_comm)
if xc.exx_fraction != 0.0:
self.xc = xc
return potential, Q_al
[docs] def calculate(self,
density,
ibzwfs=None,
vHt_x: DistributedArrays | None = None,
kpt_band_comm: MPIComm | None = None
) -> tuple[Potential, AtomArrays]:
energies, vt_sR, dedtaut_sr, vHt_x = self.calculate_pseudo_potential(
density, ibzwfs, vHt_x)
e_kinetic = 0.0
for spin, (vt_R, nt_R) in enumerate(zips(vt_sR, density.nt_sR)):
e_kinetic -= vt_R.integrate(nt_R)
if spin < density.ndensities:
e_kinetic += vt_R.integrate(density.nct_R)
if dedtaut_sr is not None:
dedtaut_sR = self.restrict(dedtaut_sr)
for dedtaut_R, taut_R in zips(dedtaut_sR,
density.taut_sR):
e_kinetic -= dedtaut_R.integrate(taut_R)
e_kinetic += dedtaut_R.integrate(density.tauct_R)
else:
dedtaut_sR = None
energies['kinetic'] = e_kinetic
if kpt_band_comm is None:
if ibzwfs is None:
kpt_band_comm = serial_comm
else:
kpt_band_comm = ibzwfs.kpt_band_comm
Q_aL = self.calculate_charges(vHt_x)
dH_asii, corrections = calculate_non_local_potential(
self.setups,
density,
self.xc,
self.external_potential,
Q_aL,
self.soc,
kpt_band_comm)
energies['spinorbit'] = 0.0
for key, e in corrections.items():
if 0:
print(f'{key:10} {energies[key]:15.9f} {e:15.9f}')
energies[key] += e
return Potential(vt_sR, dH_asii, dedtaut_sR, energies, vHt_x), Q_aL
@trace
def calculate_non_local_potential(setups,
density,
xc,
ext_pot,
Q_aL,
soc: bool,
kpt_band_comm: MPIComm
) -> tuple[AtomArrays,
dict[str, float]]:
dtype = float if density.ncomponents < 4 else complex
D_asii = density.D_asii.to_xp(np)
dH_asii = D_asii.layout.new(dtype=dtype).empty(density.ncomponents)
Q_aL = Q_aL.to_xp(np)
energy_corrections: DefaultDict[str, float] = defaultdict(float)
rank = 0
for a, D_sii in D_asii.items():
if rank % kpt_band_comm.size == kpt_band_comm.rank:
Q_L = Q_aL[a]
setup = setups[a]
dH_sii, corrections = calculate_non_local_potential1(
setup, xc, ext_pot, D_sii, Q_L, soc)
dH_asii[a][:] = dH_sii
for key, e in corrections.items():
energy_corrections[key] += e
else:
dH_asii[a][:] = 0.0
rank += 1
kpt_band_comm.sum(dH_asii.data)
# Sum over domain:
names = ['kinetic', 'coulomb', 'zero', 'xc', 'external', 'spinorbit']
energies = np.array([energy_corrections[name] for name in names])
density.D_asii.layout.atomdist.comm.sum(energies)
kpt_band_comm.sum(energies)
return (dH_asii.to_xp(density.D_asii.layout.xp),
dict(zips(names, energies)))
def calculate_non_local_potential1(setup: Setup,
xc: Functional,
ext_pot,
D_sii: Array3D,
Q_L: Array1D,
soc: bool) -> tuple[Array3D,
dict[str, float]]:
ncomponents = len(D_sii)
ndensities = 2 if ncomponents == 2 else 1
D_sp = np.array([pack_density(D_ii.real) for D_ii in D_sii])
D_p = D_sp[:ndensities].sum(0)
dH_p = (setup.K_p + setup.M_p +
setup.MB_p + 2.0 * setup.M_pp @ D_p +
setup.Delta_pL @ Q_L)
e_kinetic = setup.K_p @ D_p + setup.Kc
e_zero = setup.MB + setup.MB_p @ D_p
e_coulomb = setup.M + D_p @ (setup.M_p + setup.M_pp @ D_p)
dH_sp = np.zeros_like(D_sp, dtype=float if ncomponents < 4 else complex)
e_soc = 0.
if soc:
dHsoc_sii = soc_terms(setup, xc.xc, D_sp)
e_soc += (D_sii[1:4] * dHsoc_sii).sum().real
dH_sp[1:4] = pack_hermitian(dHsoc_sii)
dH_sp[:ndensities] = dH_p
e_xc = xc.calculate_paw_correction(setup, D_sp, dH_sp)
e_external = ext_pot.add_paw_correction(setup.Delta_pL[:, 0], dH_sp)
dH_sii = unpack_hermitian(dH_sp)
if setup.hubbard_u is not None:
eU, dHU_sii = setup.hubbard_u.calculate(setup, D_sii)
e_xc += eU
dH_sii += dHU_sii
e_kinetic -= (D_sii * dH_sii).sum().real
return dH_sii, {'kinetic': e_kinetic,
'coulomb': e_coulomb,
'zero': e_zero,
'xc': e_xc,
'external': e_external,
'spinorbit': e_soc}