Source code for gpaw.hamiltonian

"""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.hubbard import hubbard
from gpaw.lfc import LFC
from gpaw.poisson import PoissonSolver
from gpaw.spinorbit import soc
from gpaw.transformers import Transformer
from gpaw.utilities import (pack2, pack_atomic_matrices, unpack,
                            unpack_atomic_matrices)
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(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(dH_xp[0]).T
            dH_vii = [unpack(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_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 {0}*{1}*{2} 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 += ' External potential:\n {0}\n'.format(self.vext) return s def summary(self, wfs, log): log('Energy contributions relative to reference atoms:', '(reference = {0:.6f})\n'.format(self.setups.Eref * Ha)) 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)] for name, e in energies: 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:] = pack2(dH_vii) else: dH_sp = np.zeros_like(D_sp) if setup.HubU is not None: # assert self.collinear for l, U, scale in zip(setup.Hubl, setup.HubU, setup.Hubs): eU, dHU_sp = hubbard(setup, D_sp, l, U, scale) e_xc += eU dH_sp += dHU_sp 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): """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) 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(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(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(e_kin) e_kin = wfs.kd.comm.sum(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(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(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(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