"""This module defines a Hamiltonian."""
import functools
import numpy as np
from ase.units import Ha
from gpaw.arraydict import ArrayDict
from gpaw.external import create_external_potential
from gpaw.lfc import LFC
from gpaw.poisson import PoissonSolver
from gpaw.spinorbit import soc
from gpaw.transformers import Transformer
from gpaw.utilities import (pack_atomic_matrices, pack_hermitian,
unpack_atomic_matrices, unpack_density,
unpack_hermitian)
from gpaw.utilities.partition import AtomPartition
ENERGY_NAMES = ['e_kinetic', 'e_coulomb', 'e_zero', 'e_external', 'e_xc',
'e_entropy', 'e_total_free', 'e_total_extrapolated']
def apply_non_local_hamilton(dH_asp, collinear, P, out=None):
if out is None:
out = P.new()
for a, I1, I2 in P.indices:
if collinear:
dH_ii = unpack_hermitian(dH_asp[a][P.spin])
out.array[:, I1:I2] = np.dot(P.array[:, I1:I2], dH_ii)
else:
dH_xp = dH_asp[a]
# We need the transpose because
# we are dotting from the left
dH_ii = unpack_hermitian(dH_xp[0]).T
dH_vii = [unpack_hermitian(dH_p).T for dH_p in dH_xp[1:]]
out.array[:, 0, I1:I2] = (np.dot(P.array[:, 0, I1:I2],
dH_ii + dH_vii[2]) +
np.dot(P.array[:, 1, I1:I2],
dH_vii[0] - 1j * dH_vii[1]))
out.array[:, 1, I1:I2] = (np.dot(P.array[:, 1, I1:I2],
dH_ii - dH_vii[2]) +
np.dot(P.array[:, 0, I1:I2],
dH_vii[0] + 1j * dH_vii[1]))
return out
# from gpaw.utilities.debug import frozen
# @frozen
[docs]class Hamiltonian:
def __init__(self, gd, finegd, nspins, collinear, setups, timer, xc, world,
redistributor, vext=None):
self.gd = gd
self.finegd = finegd
self.nspins = nspins
self.collinear = collinear
self.setups = setups
self.timer = timer
self.xc = xc
self.world = world
self.redistributor = redistributor
self.ncomponents = self.nspins if self.collinear else 1 + 3
self.atomdist = None
self.dH_asp = None
self.vt_xG = None
self.vt_sG = None
self.vt_vG = None
self.vHt_g = None
self.vt_xg = None
self.vt_sg = None
self.vt_vg = None
self.atom_partition = None
# Energy contributioons that sum up to e_total_free:
self.e_kinetic = None
self.e_coulomb = None
self.e_zero = None
self.e_external = None
self.e_xc = None
self.e_entropy = None
self.e_band = None
self.e_sic = None
self.e_total_free = None
self.e_total_extrapolated = None
self.e_kinetic0 = None
self.ref_vt_sG = None
self.ref_dH_asp = None
if isinstance(vext, dict):
vext = create_external_potential(**vext)
self.vext = vext # external potential
self.positions_set = False
self.spos_ac = None
self.soc = False
@property
def dH(self):
return functools.partial(apply_non_local_hamilton,
self.dH_asp, self.collinear)
def update_atomic_hamiltonians(self, value):
if isinstance(value, dict):
dtype = complex if self.soc else float
tmp = self.setups.empty_atomic_matrix(self.ncomponents,
self.atom_partition, dtype)
tmp.update(value)
value = tmp
assert isinstance(value, ArrayDict) or value is None, type(value)
if value is not None:
value.check_consistency()
self.dH_asp = value
def __str__(self):
s = 'Hamiltonian:\n'
s += (' XC and Coulomb potentials evaluated on a {}*{}*{} grid\n'
.format(*self.finegd.N_c))
s += ' Using the %s Exchange-Correlation functional\n' % self.xc.name
# We would get the description of the XC functional here,
# except the thing has probably not been fully initialized yet.
if self.vext is not None:
s += f' External potential:\n {self.vext}\n'
return s
def summary(self, wfs, log):
log('Energy contributions relative to reference atoms:',
f'(reference = {self.setups.Eref * Ha:.6f})\n')
energies = [('Kinetic: ', self.e_kinetic),
('Potential: ', self.e_coulomb),
('External: ', self.e_external),
('XC: ', self.e_xc),
('Entropy (-ST):', self.e_entropy),
('Local: ', self.e_zero),
('SIC: ', self.e_sic)]
for name, e in energies:
if name == 'SIC: ' and e is None:
continue
else:
log('%-14s %+11.6f' % (name, Ha * e))
log('--------------------------')
log('Free energy: %+11.6f' % (Ha * self.e_total_free))
log('Extrapolated: %+11.6f' % (Ha * self.e_total_extrapolated))
log()
self.xc.summary(log)
try:
workfunctions = self.get_workfunctions(wfs)
except ValueError:
pass
else:
log('Dipole-layer corrected work functions: {:.6f}, {:.6f} eV'
.format(*np.array(workfunctions) * Ha))
log()
[docs] def get_workfunctions(self, wfs):
"""
Returns the work functions, in Hartree, for a dipole-corrected
simulation. Returns None if no dipole correction is present.
(wfs can be obtained from calc.wfs)
"""
try:
dipole_correction = self.poisson.correction
except AttributeError:
raise ValueError(
'Work function not defined if no field-free region. Consider '
'using a dipole correction if you are looking for a '
'work function.')
c = self.poisson.c # index of axis perpendicular to dipole-layer
if not self.gd.pbc_c[c]:
# zero boundary conditions
vacuum = 0.0
else:
v_q = self.pd3.gather(self.vHt_q)
if self.pd3.comm.rank == 0:
axes = (c, (c + 1) % 3, (c + 2) % 3)
v_g = self.pd3.ifft(v_q, local=True).transpose(axes)
vacuum = v_g[0].mean()
else:
vacuum = np.nan
fermilevel = wfs.fermi_level
wf1 = vacuum - fermilevel + dipole_correction
wf2 = vacuum - fermilevel - dipole_correction
return np.array([wf1, wf2])
def set_positions_without_ruining_everything(self, spos_ac,
atom_partition):
self.spos_ac = spos_ac
rank_a = atom_partition.rank_a
# If both old and new atomic ranks are present, start a blank dict if
# it previously didn't exist but it will needed for the new atoms.
# XXX what purpose does this serve? In what case does it happen?
# How would one even go about figuring it out? Why does it all have
# to be so unreadable? -Ask
#
if (self.atom_partition is not None and
self.dH_asp is None and (rank_a == self.gd.comm.rank).any()):
self.update_atomic_hamiltonians({})
if self.atom_partition is not None and self.dH_asp is not None:
self.timer.start('Redistribute')
self.dH_asp.redistribute(atom_partition)
self.timer.stop('Redistribute')
self.atom_partition = atom_partition
self.atomdist = self.redistributor.get_atom_distributions(spos_ac)
def set_positions(self, spos_ac, atom_partition):
self.vbar.set_positions(spos_ac, atom_partition)
self.xc.set_positions(spos_ac)
self.set_positions_without_ruining_everything(spos_ac, atom_partition)
self.positions_set = True
def initialize(self):
self.vt_xg = self.finegd.empty(self.ncomponents)
self.vt_sg = self.vt_xg[:self.nspins]
self.vt_vg = self.vt_xg[self.nspins:]
self.vHt_g = self.finegd.zeros()
self.vt_xG = self.gd.empty(self.ncomponents)
self.vt_sG = self.vt_xG[:self.nspins]
self.vt_vG = self.vt_xG[self.nspins:]
[docs] def update(self, density, wfs=None, kin_en_using_band=True):
"""Calculate effective potential.
The XC-potential and the Hartree potential are evaluated on
the fine grid, and the sum is then restricted to the coarse
grid."""
self.timer.start('Hamiltonian')
if self.vt_sg is None:
with self.timer('Initialize Hamiltonian'):
self.initialize()
finegrid_energies = self.update_pseudo_potential(density)
coarsegrid_e_kinetic = self.calculate_kinetic_energy(density)
with self.timer('Calculate atomic Hamiltonians'):
W_aL = self.calculate_atomic_hamiltonians(density)
atomic_energies = self.update_corrections(density, W_aL)
# Make energy contributions summable over world:
finegrid_energies *= self.finegd.comm.size / self.world.size
coarsegrid_e_kinetic *= self.gd.comm.size / self.world.size
# (careful with array orderings/contents)
if 0:
print('kinetic', atomic_energies[0], coarsegrid_e_kinetic)
print('coulomb', atomic_energies[1], finegrid_energies[0])
print('zero', atomic_energies[2], finegrid_energies[1])
print('xc', atomic_energies[4], finegrid_energies[3])
print('external', atomic_energies[3], finegrid_energies[2])
energies = atomic_energies # kinetic, coulomb, zero, external, xc
energies[1:] += finegrid_energies # coulomb, zero, external, xc
energies[0] += coarsegrid_e_kinetic # kinetic
with self.timer('Communicate'): # time possible load imbalance
self.world.sum(energies)
if not kin_en_using_band:
assert wfs is not None
with self.timer('New Kinetic Energy'):
energies[0] = \
self.calculate_kinetic_energy_directly(density,
wfs)
(self.e_kinetic0, self.e_coulomb, self.e_zero,
self.e_external, self.e_xc) = energies
self.timer.stop('Hamiltonian')
def update_corrections(self, dens, W_aL):
self.timer.start('Atomic')
self.update_atomic_hamiltonians(None) # XXXX
e_kinetic = 0.0
e_coulomb = 0.0
e_zero = 0.0
e_external = 0.0
e_xc = 0.0
D_asp = self.atomdist.to_work(dens.D_asp)
dtype = complex if self.soc else float
dH_asp = self.setups.empty_atomic_matrix(self.ncomponents,
D_asp.partition, dtype)
for a, D_sp in D_asp.items():
W_L = W_aL[a]
setup = self.setups[a]
if self.nspins == 2:
D_p = D_sp.sum(0)
else:
D_p = D_sp[0]
dH_p = (setup.K_p + setup.M_p +
setup.MB_p + 2.0 * np.dot(setup.M_pp, D_p) +
np.dot(setup.Delta_pL, W_L))
e_kinetic += np.dot(setup.K_p, D_p) + setup.Kc
e_zero += setup.MB + np.dot(setup.MB_p, D_p)
e_coulomb += setup.M + np.dot(D_p, (setup.M_p +
np.dot(setup.M_pp, D_p)))
if self.soc:
dH_vii = soc(setup, self.xc, D_sp)
dH_sp = np.zeros_like(D_sp, dtype=complex)
dH_sp[1:] = pack_hermitian(dH_vii)
else:
dH_sp = np.zeros_like(D_sp)
if setup.hubbard_u is not None:
eU, dHU_sii = setup.hubbard_u.calculate(setup,
unpack_density(D_sp))
e_xc += eU
dH_sp += pack_hermitian(dHU_sii)
dH_sp[:self.nspins] += dH_p
if self.vext is not None:
self.vext.paw_correction(setup.Delta_pL[:, 0], dH_sp)
if self.vext and self.vext.get_name() == 'CDFTPotential':
# cDFT atomic hamiltonian, eq. 25
# energy correction added in cDFT main
h_cdft_a, h_cdft_b = self.vext.get_atomic_hamiltonians(
setups=setup.Delta_pL[:, 0], atom=a)
dH_sp[0] += h_cdft_a
dH_sp[1] += h_cdft_b
if self.ref_dH_asp:
assert self.collinear
dH_sp += self.ref_dH_asp[a]
dH_asp[a] = dH_sp
self.timer.start('XC Correction')
for a, D_sp in D_asp.items():
e_xc += self.xc.calculate_paw_correction(self.setups[a], D_sp,
dH_asp[a], a=a)
self.timer.stop('XC Correction')
for a, D_sp in D_asp.items():
e_kinetic -= (D_sp * dH_asp[a]).sum().real
self.update_atomic_hamiltonians(self.atomdist.from_work(dH_asp))
self.timer.stop('Atomic')
# Make corrections due to non-local xc:
# self.Enlxc = 0.0 # XXXxcfunc.get_non_local_energy()
e_kinetic += self.xc.get_kinetic_energy_correction() / self.world.size
return np.array([e_kinetic, e_coulomb, e_zero, e_external, e_xc])
[docs] def get_energy(self, e_entropy, wfs, kin_en_using_band=True, e_sic=None):
"""Sum up all eigenvalues weighted with occupation numbers"""
self.e_band = wfs.calculate_band_energy()
if kin_en_using_band:
self.e_kinetic = self.e_kinetic0 + self.e_band
else:
self.e_kinetic = self.e_kinetic0
self.e_entropy = e_entropy
if 0:
print(self.e_kinetic0,
self.e_band,
self.e_coulomb,
self.e_external,
self.e_zero,
self.e_xc,
self.e_entropy)
self.e_total_free = (self.e_kinetic + self.e_coulomb +
self.e_external + self.e_zero + self.e_xc +
self.e_entropy)
if e_sic is not None:
self.e_sic = e_sic
self.e_total_free += e_sic
self.e_total_extrapolated = (
self.e_total_free +
wfs.occupations.extrapolate_factor * e_entropy)
return self.e_total_free
def linearize_to_xc(self, new_xc, density):
# Store old hamiltonian
ref_vt_sG = self.vt_sG.copy()
ref_dH_asp = self.dH_asp.copy()
self.xc = new_xc
self.xc.set_positions(self.spos_ac)
self.update(density)
ref_vt_sG -= self.vt_sG
for a, dH_sp in self.dH_asp.items():
ref_dH_asp[a] -= dH_sp
self.ref_vt_sG = ref_vt_sG
self.ref_dH_asp = self.atomdist.to_work(ref_dH_asp)
def calculate_forces(self, dens, F_av):
ghat_aLv = dens.ghat.dict(derivative=True)
nct_av = dens.nct.dict(derivative=True)
vbar_av = self.vbar.dict(derivative=True)
self.calculate_forces2(dens, ghat_aLv, nct_av, vbar_av)
F_coarsegrid_av = np.zeros_like(F_av)
# Force from compensation charges:
_Q, Q_aL = dens.calculate_multipole_moments()
for a, dF_Lv in ghat_aLv.items():
F_av[a] += np.dot(Q_aL[a], dF_Lv)
# Force from smooth core charge:
for a, dF_v in nct_av.items():
F_coarsegrid_av[a] += dF_v[0]
# Force from zero potential:
for a, dF_v in vbar_av.items():
F_av[a] += dF_v[0]
self.xc.add_forces(F_av)
self.gd.comm.sum(F_coarsegrid_av, 0)
self.finegd.comm.sum(F_av, 0)
if self.vext:
if self.vext.get_name() == 'CDFTPotential':
F_av += self.vext.get_cdft_forces()
F_av += F_coarsegrid_av
def apply_local_potential(self, psit_nG, Htpsit_nG, s):
vt_G = self.vt_sG[s]
if psit_nG.ndim == 3:
Htpsit_nG += psit_nG * vt_G
else:
for psit_G, Htpsit_G in zip(psit_nG, Htpsit_nG):
Htpsit_G += psit_G * vt_G
[docs] def apply(self, a_xG, b_xG, wfs, kpt, calculate_P_ani=True):
"""Apply the Hamiltonian operator to a set of vectors.
Parameters:
a_nG: ndarray
Set of vectors to which the overlap operator is applied.
b_nG: ndarray, output
Resulting S times a_nG vectors.
wfs: WaveFunctions
Wave-function object defined in wavefunctions.py
kpt: KPoint object
k-point object defined in kpoint.py.
calculate_P_ani: bool
When True, the integrals of projector times vectors
P_ni = <p_i | a_nG> are calculated.
When False, existing P_ani are used
"""
wfs.kin.apply(a_xG, b_xG, kpt.phase_cd)
self.apply_local_potential(a_xG, b_xG, kpt.s)
shape = a_xG.shape[:-3]
P_axi = wfs.pt.dict(shape)
if calculate_P_ani: # TODO calculate_P_ani=False is experimental
wfs.pt.integrate(a_xG, P_axi, kpt.q)
else:
for a, P_ni in kpt.P_ani.items():
P_axi[a][:] = P_ni
for a, P_xi in P_axi.items():
dH_ii = unpack_hermitian(self.dH_asp[a][kpt.s])
P_axi[a] = np.dot(P_xi, dH_ii)
wfs.pt.add(b_xG, P_axi, kpt.q)
[docs] def get_xc_difference(self, xc, density):
"""Calculate non-selfconsistent XC-energy difference."""
if density.nt_sg is None:
density.interpolate_pseudo_density()
nt_sg = density.nt_sg
if hasattr(xc, 'hybrid'):
xc.calculate_exx()
finegd_e_xc = xc.calculate(density.finegd, nt_sg)
D_asp = self.atomdist.to_work(density.D_asp)
atomic_e_xc = 0.0
for a, D_sp in D_asp.items():
setup = self.setups[a]
atomic_e_xc += xc.calculate_paw_correction(setup, D_sp, a=a)
e_xc = finegd_e_xc + self.world.sum_scalar(atomic_e_xc)
return e_xc - self.e_xc
def estimate_memory(self, mem):
nbytes = self.gd.bytecount()
nfinebytes = self.finegd.bytecount()
arrays = mem.subnode('Arrays', 0)
arrays.subnode('vHt_g', nfinebytes)
arrays.subnode('vt_sG', self.nspins * nbytes)
arrays.subnode('vt_sg', self.nspins * nfinebytes)
self.xc.estimate_memory(mem.subnode('XC'))
self.poisson.estimate_memory(mem.subnode('Poisson'))
self.vbar.estimate_memory(mem.subnode('vbar'))
def write(self, writer):
# Write all eneriges:
for name in ENERGY_NAMES:
energy = getattr(self, name)
if energy is not None:
energy *= Ha
writer.write(name, energy)
writer.write(
potential=self.gd.collect(self.vt_xG) * Ha,
atomic_hamiltonian_matrices=pack_atomic_matrices(self.dH_asp) * Ha)
self.xc.write(writer.child('xc'))
if hasattr(self.poisson, 'write'):
self.poisson.write(writer.child('poisson'))
def read(self, reader):
h = reader.hamiltonian
# Read all energies:
for name in ENERGY_NAMES:
energy = h.get(name)
if energy is not None:
energy /= reader.ha
setattr(self, name, energy)
# Read pseudo potential on the coarse grid
# and broadcast on kpt/band comm:
self.initialize()
self.gd.distribute(h.potential / reader.ha, self.vt_xG)
self.atom_partition = AtomPartition(self.gd.comm,
np.zeros(len(self.setups), int),
name='hamiltonian-init-serial')
# Read non-local part of hamiltonian
self.update_atomic_hamiltonians({})
dH_sP = h.atomic_hamiltonian_matrices / reader.ha
if self.gd.comm.rank == 0:
self.update_atomic_hamiltonians(
unpack_atomic_matrices(dH_sP, self.setups))
if hasattr(self.poisson, 'read'):
self.poisson.read(reader)
self.poisson.set_grid_descriptor(self.finegd)
[docs] def calculate_kinetic_energy_directly(self, density, wfs):
"""
Calculate kinetic energy as 1/2 (nable psi)^2
it gives better estimate of kinetic energy during the SCF.
Important for direct min.
'calculate_kinetic_energy' method gives a correct
value of kinetic energy only at self-consistent solution.
:param density:
:param wfs:
:return: total kinetic energy
"""
# pseudo-part
if wfs.mode == 'lcao':
return self.calculate_kinetic_energy_using_kin_en_matrix(
density, wfs)
elif wfs.mode == 'pw':
e_kin = 0.0
for kpt in wfs.kpt_u:
for f, psit_G in zip(kpt.f_n, kpt.psit_nG):
if f > 1.0e-10:
G2_G = wfs.pd.G2_qG[kpt.q]
e_kin += f * wfs.pd.integrate(
0.5 * G2_G * psit_G, psit_G).real
else:
e_kin = 0.0
def Lapl(psit_G, kpt):
Lpsit_G = np.zeros_like(psit_G)
wfs.kin.apply(psit_G, Lpsit_G, kpt.phase_cd)
return Lpsit_G
for kpt in wfs.kpt_u:
for f, psit_G in zip(kpt.f_n, kpt.psit_nG):
if f > 1.0e-10:
e_kin += f * wfs.integrate(
Lapl(psit_G, kpt), psit_G, False)
e_kin = e_kin.real
e_kin = wfs.gd.comm.sum_scalar(e_kin)
e_kin = wfs.kd.comm.sum_scalar(e_kin) # ?
# paw corrections
e_kin_paw = 0.0
for a, D_sp in density.D_asp.items():
setup = wfs.setups[a]
D_p = D_sp.sum(0)
e_kin_paw += np.dot(setup.K_p, D_p) + setup.Kc
e_kin_paw = density.gd.comm.sum_scalar(e_kin_paw)
return e_kin + e_kin_paw
[docs] def calculate_kinetic_energy_using_kin_en_matrix(self, density,
wfs):
"""
E_k = sum_{M'M} rho_MM' T_M'M
better agreement between gradients of energy and
the total energy during the direct minimisation.
This is important when the line search is used.
Also avoids using the eigenvalues which are
not calculated during the direct minimisation.
'calculate_kinetic_energy' method gives a correct
value of kinetic energy only at self-consistent solution.
:param density:
:param wfs:
:return: total kinetic energy
"""
# pseudo-part
e_kinetic = 0.0
e_kin_paw = 0.0
for kpt in wfs.kpt_u:
# calculation of the density matrix directly
# can be expansive for a large scale
# as there are lot of empty states
# when the exponential transformation is used
# (n_bands=n_basis_functions.)
#
# rho_MM = \
# wfs.calculate_density_matrix(kpt.f_n, kpt.C_nM)
# e_kinetic += np.einsum('ij,ji->', kpt.T_MM, rho_MM)
#
# the code below is faster
self.timer.start('Pseudo part')
occ = kpt.f_n > 1e-10
x_nn = np.dot(kpt.C_nM[occ],
np.dot(kpt.T_MM,
kpt.C_nM[occ].T.conj())).real
e_kinetic += np.einsum('i,ii->', kpt.f_n[occ], x_nn)
self.timer.stop('Pseudo part')
# del rho_MM
e_kinetic = wfs.kd.comm.sum_scalar(e_kinetic)
# paw corrections
for a, D_sp in density.D_asp.items():
setup = wfs.setups[a]
D_p = D_sp.sum(0)
e_kin_paw += np.dot(setup.K_p, D_p) + setup.Kc
e_kin_paw = self.gd.comm.sum_scalar(e_kin_paw)
return e_kinetic.real + e_kin_paw
class RealSpaceHamiltonian(Hamiltonian):
def __init__(self, gd, finegd, nspins, collinear, setups, timer, xc, world,
vext=None,
psolver=None, stencil=3, redistributor=None,
charge: float = 0.0):
Hamiltonian.__init__(self, gd, finegd, nspins, collinear,
setups, timer, xc,
world, vext=vext,
redistributor=redistributor)
# Solver for the Poisson equation:
if psolver is None:
psolver = {}
if isinstance(psolver, dict):
psolver = PoissonSolver(**psolver)
self.poisson = psolver
self.poisson.set_grid_descriptor(self.finegd)
# Restrictor function for the potential:
self.restrictor = Transformer(self.finegd, self.redistributor.aux_gd,
stencil)
self.restrict = self.restrictor.apply
self.vbar = LFC(self.finegd, [[setup.vbar] for setup in setups],
forces=True)
self.vbar_g = None
self.npoisson = None
def restrict_and_collect(self, a_xg, b_xg=None, phases=None):
if self.redistributor.enabled:
tmp_xg = self.restrictor.apply(a_xg, output=None, phases=phases)
b_xg = self.redistributor.collect(tmp_xg, b_xg)
else:
b_xg = self.restrictor.apply(a_xg, output=b_xg, phases=phases)
return b_xg
def __str__(self):
s = Hamiltonian.__str__(self)
degree = self.restrictor.nn * 2 - 1
name = ['linear', 'cubic', 'quintic', 'heptic'][degree // 2]
s += (' Interpolation: tri-%s ' % name +
'(%d. degree polynomial)\n' % degree)
s += ' Poisson solver: %s' % self.poisson.get_description()
return s
def set_positions(self, spos_ac, rank_a):
Hamiltonian.set_positions(self, spos_ac, rank_a)
if self.vbar_g is None:
self.vbar_g = self.finegd.empty()
self.vbar_g[:] = 0.0
self.vbar.add(self.vbar_g)
def update_pseudo_potential(self, dens):
self.timer.start('vbar')
e_zero = self.finegd.integrate(self.vbar_g, dens.nt_g,
global_integral=False)
vt_g = self.vt_sg[0]
vt_g[:] = self.vbar_g
self.timer.stop('vbar')
e_external = 0.0
if self.vext is not None:
if self.vext.get_name() == 'CDFTPotential':
vext_g = self.vext.get_potential(self.finegd).copy()
e_external += self.vext.get_cdft_external_energy(
dens,
self.nspins, vext_g, vt_g, self.vbar_g, self.vt_sg)
else:
vext_g = self.vext.get_potential(self.finegd)
vt_g += vext_g
e_external = self.finegd.integrate(vext_g, dens.rhot_g,
global_integral=False)
if self.nspins == 2:
self.vt_sg[1] = vt_g
self.timer.start('XC 3D grid')
e_xc = self.xc.calculate(self.finegd, dens.nt_sg, self.vt_sg)
e_xc /= self.finegd.comm.size
self.timer.stop('XC 3D grid')
self.timer.start('Poisson')
# npoisson is the number of iterations:
self.npoisson = self.poisson.solve(self.vHt_g, dens.rhot_g,
charge=-dens.charge,
timer=self.timer)
self.timer.stop('Poisson')
self.timer.start('Hartree integrate/restrict')
e_coulomb = 0.5 * self.finegd.integrate(self.vHt_g, dens.rhot_g,
global_integral=False)
for vt_g in self.vt_sg:
vt_g += self.vHt_g
self.timer.stop('Hartree integrate/restrict')
energies = np.array([e_coulomb, e_zero, e_external, e_xc])
return energies
def calculate_kinetic_energy(self, density):
# XXX new timer item for kinetic energy?
self.timer.start('Hartree integrate/restrict')
self.restrict_and_collect(self.vt_sg, self.vt_sG)
e_kinetic = 0.0
s = 0
for vt_G, nt_G in zip(self.vt_sG, density.nt_sG):
if self.ref_vt_sG is not None:
vt_G += self.ref_vt_sG[s]
if s < self.nspins:
e_kinetic -= self.gd.integrate(vt_G, nt_G - density.nct_G,
global_integral=False)
else:
e_kinetic -= self.gd.integrate(vt_G, nt_G,
global_integral=False)
s += 1
self.timer.stop('Hartree integrate/restrict')
return e_kinetic
def calculate_atomic_hamiltonians(self, dens):
def getshape(a):
return sum(2 * l + 1 for l, _ in enumerate(self.setups[a].ghat_l)),
W_aL = ArrayDict(self.atomdist.aux_partition, getshape, float)
if self.vext:
if self.vext.get_name() != 'CDFTPotential':
vext_g = self.vext.get_potential(self.finegd)
dens.ghat.integrate(self.vHt_g + vext_g, W_aL)
else:
dens.ghat.integrate(self.vHt_g, W_aL)
else:
dens.ghat.integrate(self.vHt_g, W_aL)
return self.atomdist.to_work(self.atomdist.from_aux(W_aL))
def calculate_forces2(self, dens, ghat_aLv, nct_av, vbar_av):
if self.nspins == 2:
vt_G = self.vt_sG.mean(0)
else:
vt_G = self.vt_sG[0]
self.vbar.derivative(dens.nt_g, vbar_av)
if self.vext:
if self.vext.get_name() == 'CDFTPotential':
# CDFT force added in calculate_forces
dens.ghat.derivative(self.vHt_g, ghat_aLv)
else:
vext_g = self.vext.get_potential(self.finegd)
dens.ghat.derivative(self.vHt_g + vext_g, ghat_aLv)
else:
dens.ghat.derivative(self.vHt_g, ghat_aLv)
dens.nct.derivative(vt_G, nct_av)
def get_electrostatic_potential(self, dens):
self.update(dens)
v_g = self.finegd.collect(self.vHt_g, broadcast=True)
v_g = self.finegd.zero_pad(v_g)
if hasattr(self.poisson, 'correction'):
assert self.poisson.c == 2
v_g[:, :, 0] = self.poisson.correction
return v_g