def solve(self,
phi,
rho,
charge=None,
eps=None,
maxcharge=1e-6,
zero_initial_phi=False):
if self.phi_tilde is None:
self.phi_tilde = self.gd.zeros()
phi_tilde = self.phi_tilde
niter_tilde = FDPoissonSolver.solve(self, phi_tilde, rho, None,
self.eps, maxcharge, False)
epsr, dx_epsr, dy_epsr, dz_epsr = self.dielectric.eps_gradeps
dx_phi_tilde = self.gd.empty()
dy_phi_tilde = self.gd.empty()
dz_phi_tilde = self.gd.empty()
Gradient(self.gd, 0, 1.0, self.nn).apply(phi_tilde, dx_phi_tilde)
Gradient(self.gd, 1, 1.0, self.nn).apply(phi_tilde, dy_phi_tilde)
Gradient(self.gd, 2, 1.0, self.nn).apply(phi_tilde, dz_phi_tilde)
scalar_product = (dx_epsr * dx_phi_tilde + dy_epsr * dy_phi_tilde +
dz_epsr * dz_phi_tilde)
rho_and_pol = (rho / epsr + scalar_product / (4. * np.pi * epsr**2))
niter = FDPoissonSolver.solve(self, phi, rho_and_pol, None, eps,
maxcharge, zero_initial_phi)
return niter_tilde + niter
def calculate_kinetic_energy_density(self):
if self.taugrad_v is None:
self.taugrad_v = [
Gradient(self.gd, v, n=3, dtype=self.dtype).apply
for v in range(3)
]
assert not hasattr(self.kpt_u[0], 'c_on')
if self.kpt_u[0].psit_nG is None:
raise RuntimeError('No wavefunctions yet')
if isinstance(self.kpt_u[0].psit_nG, FileReference):
# XXX initialize
raise RuntimeError('Wavefunctions have not been initialized.')
taut_sG = self.gd.zeros(self.nspins)
dpsit_G = self.gd.empty(dtype=self.dtype)
for kpt in self.kpt_u:
for f, psit_G in zip(kpt.f_n, kpt.psit_nG):
for v in range(3):
self.taugrad_v[v](psit_G, dpsit_G, kpt.phase_cd)
axpy(0.5 * f, abs(dpsit_G)**2, taut_sG[kpt.s])
self.kd.comm.sum(taut_sG)
self.band_comm.sum(taut_sG)
return taut_sG
def __init__(self, calc, xc, ibzq_qc, fd, unit_cells, density_cut, ecut,
tag, timer):
self.calc = calc
self.gd = calc.density.gd
self.xc = xc
self.ibzq_qc = ibzq_qc
self.fd = fd
self.unit_cells = unit_cells
self.density_cut = density_cut
self.ecut = ecut
self.tag = tag
self.timer = timer
self.A_x = -(3 / 4.) * (3 / np.pi)**(1 / 3.)
self.n_g = calc.get_all_electron_density(gridrefinement=1)
self.n_g *= Bohr**3
if xc[-3:] == 'PBE':
nf_g = calc.get_all_electron_density(gridrefinement=2)
nf_g *= Bohr**3
gdf = self.gd.refine()
grad_v = [Gradient(gdf, v, n=1).apply for v in range(3)]
gradnf_vg = gdf.empty(3)
for v in range(3):
grad_v[v](nf_g, gradnf_vg[v])
self.gradn_vg = gradnf_vg[:, ::2, ::2, ::2]
qd = KPointDescriptor(self.ibzq_qc)
self.pd = PWDescriptor(ecut / Hartree, self.gd, complex, qd)
def set_absorbing_boundary(self, absorbing_boundary):
""" Sets up the absorbing boundary.
Parameters:
absorbing_boundary: absorbing boundary object of any kind.
"""
self.absorbing_boundary = absorbing_boundary
self.absorbing_boundary.set_up(self.hamiltonian.gd)
if self.absorbing_boundary.type == 'PML':
gd = self.hamiltonian.gd
self.laplace = Laplace(gd, n=2, dtype=complex)
self.gradient = np.array(
(Gradient(gd, 0, n=2,
dtype=complex), Gradient(gd, 1, n=2, dtype=complex),
Gradient(gd, 2, n=2, dtype=complex)))
self.lpsit = None
def make_derivative(self, s, K, n1, n2):
wfs = self.calc.wfs
if self.real_space_derivatives:
grad_v = [
Gradient(wfs.gd, v, 1.0, 4, complex).apply for v in range(3)
]
U_cc, T, a_a, U_aii, shift_c, time_reversal = \
self.construct_symmetry_operators(K)
A_cv = wfs.gd.cell_cv
M_vv = np.dot(np.dot(A_cv.T, U_cc.T), np.linalg.inv(A_cv).T)
ik = wfs.kd.bz2ibz_k[K]
kpt = wfs.kpt_u[s * wfs.kd.nibzkpts + ik]
psit_nG = kpt.psit_nG
iG_Gv = 1j * wfs.pd.get_reciprocal_vectors(q=ik, add_q=False)
ut_nvR = wfs.gd.zeros((n2 - n1, 3), complex)
for n in range(n1, n2):
for v in range(3):
if self.real_space_derivatives:
ut_R = T(wfs.pd.ifft(psit_nG[n], ik))
grad_v[v](ut_R, ut_nvR[n - n1, v], np.ones((3, 2),
complex))
else:
ut_R = T(wfs.pd.ifft(iG_Gv[:, v] * psit_nG[n], ik))
for v2 in range(3):
ut_nvR[n - n1, v2] += ut_R * M_vv[v, v2]
return ut_nvR
def calculate_field(
gd,
rho_g,
bgef_v,
phi_g,
ef_vg,
fe_g, # preallocated numpy arrays
nv=3,
poisson_nn=3,
poisson_relax='J',
gradient_n=3,
poisson_eps=1e-20):
dtype = rho_g.dtype
yes_complex = dtype == complex
phi_g[:] = 0.0
ef_vg[:] = 0.0
fe_g[:] = 0.0
tmp_g = gd.zeros(dtype=float)
# Poissonsolver
poissonsolver = PoissonSolver(nn=poisson_nn,
relax=poisson_relax,
eps=poisson_eps)
poissonsolver.set_grid_descriptor(gd)
poissonsolver.initialize()
# Potential, real part
poissonsolver.solve(tmp_g, rho_g.real.copy())
phi_g += tmp_g
# Potential, imag part
if yes_complex:
tmp_g[:] = 0.0
poissonsolver.solve(tmp_g, rho_g.imag.copy())
phi_g += 1.0j * tmp_g
# Gradient
gradient = [Gradient(gd, v, scale=1.0, n=gradient_n) for v in range(nv)]
for v in range(nv):
# Electric field, real part
gradient[v].apply(-phi_g.real, tmp_g)
ef_vg[v] += tmp_g
# Electric field, imag part
if yes_complex:
gradient[v].apply(-phi_g.imag, tmp_g)
ef_vg[v] += 1.0j * tmp_g
# Electric field enhancement
tmp_g[:] = 0.0 # total electric field norm
bgefnorm = 0.0 # background electric field norm
for v in range(nv):
tmp_g += np.absolute(bgef_v[v] + ef_vg[v])**2
bgefnorm += np.absolute(bgef_v[v])**2
tmp_g = np.sqrt(tmp_g)
bgefnorm = np.sqrt(bgefnorm)
fe_g[:] = tmp_g / bgefnorm
def density_matrix(self, n, m, k, Gspace=True):
ibzk_kc = self.ibzk_kc
bzk_kc = self.bzk_kc
kq_k = self.kq_k
gd = self.gd
kd = self.kd
ibzkpt1 = kd.kibz_k[k]
ibzkpt2 = kd.kibz_k[kq_k[k]]
psitold_g = self.get_wavefunction(ibzkpt1, n, True)
psit1_g = kd.transform_wave_function(psitold_g, k)
psitold_g = self.get_wavefunction(ibzkpt2, m, True)
psit2_g = kd.transform_wave_function(psitold_g, kq_k[k])
if Gspace is False:
return psit1_g, psit2_g
else:
# FFT
tmp_g = psit1_g.conj() * psit2_g * self.expqr_g
rho_g = np.fft.fftn(tmp_g) * self.vol / self.nG0
# Here, planewave cutoff is applied
rho_G = np.zeros(self.npw, dtype=complex)
for iG in range(self.npw):
index = self.Gindex_G[iG]
rho_G[iG] = rho_g[index[0], index[1], index[2]]
if self.optical_limit:
d_c = [
Gradient(gd, i, n=4, dtype=complex).apply for i in range(3)
]
dpsit_g = gd.empty(dtype=complex)
tmp = np.zeros((3), dtype=complex)
phase_cd = np.exp(2j * pi * gd.sdisp_cd *
bzk_kc[kq_k[k], :, np.newaxis])
for ix in range(3):
d_c[ix](psit2_g, dpsit_g, phase_cd)
tmp[ix] = gd.integrate(psit1_g.conj() * dpsit_g)
rho_G[0] = -1j * np.dot(self.qq_v, tmp)
# PAW correction
pt = self.pt
P1_ai = pt.dict()
pt.integrate(psit1_g, P1_ai, k)
P2_ai = pt.dict()
pt.integrate(psit2_g, P2_ai, kq_k[k])
for a, id in enumerate(self.calc.wfs.setups.id_a):
P_p = np.outer(P1_ai[a].conj(), P2_ai[a]).ravel()
gemv(1.0, self.phi_aGp[a], P_p, 1.0, rho_G)
if self.optical_limit:
rho_G[0] /= self.e_kn[ibzkpt2, m] - self.e_kn[ibzkpt1, n]
return rho_G
def allocate(self):
SurfaceCalculator.allocate(self)
self.gradient = [
Gradient(self.gd, i, 1.0, self.nn) for i in (0, 1, 2)
]
self.gradient_out = self.gd.empty(3)
self.norm_grad_out = self.gd.empty()
self.div_tmp = self.gd.empty()
def initialize(self):
if self.dipcorr:
self.gradient = [
Gradient(self.finegd, i, 1.0, self.poisson.poissonsolver.nn)
for i in (0, 1, 2)
]
else:
self.gradient = [
Gradient(self.finegd, i, 1.0, self.poisson.nn)
for i in (0, 1, 2)
]
self.vt_ia_g = self.finegd.zeros()
self.cavity.allocate()
self.dielectric.allocate()
for ia in self.interactions:
ia.allocate()
RealSpaceHamiltonian.initialize(self)
def interpolate(self):
self.density.interpolate_pseudo_density()
if self.taut_sg is None:
self.taut_sg = self.finegd.empty(self.nspins)
self.nt_grad2_sg = self.finegd.empty(self.nspins)
ddr_v = [Gradient(self.finegd, v, n=3).apply for v in range(3)]
self.nt_grad2_sg[:] = 0.0
d_g = self.finegd.empty()
# Transfer the densities from the coarse to the fine grid
for s in range(self.nspins):
self.density.interpolator.apply(self.taut_sG[s], self.taut_sg[s])
#self.density.interpolator.apply(self.nt_grad2_sG[s],
# self.nt_grad2_sg[s])
for v in range(3):
ddr_v[v](self.density.nt_sg[s], d_g)
self.nt_grad2_sg[s] += d_g**2.0
def calculate_kinetic_energy_density(self):
if self.taugrad_v is None:
self.taugrad_v = [
Gradient(self.gd, v, n=3, dtype=self.dtype).apply
for v in range(3)
]
assert not hasattr(self.kpt_u[0], 'c_on')
if not isinstance(self.kpt_u[0].psit_nG, np.ndarray):
return None
taut_sG = self.gd.zeros(self.nspins)
dpsit_G = self.gd.empty(dtype=self.dtype)
for kpt in self.kpt_u:
for f, psit_G in zip(kpt.f_n, kpt.psit_nG):
for v in range(3):
self.taugrad_v[v](psit_G, dpsit_G, kpt.phase_cd)
axpy(0.5 * f, abs(dpsit_G)**2, taut_sG[kpt.s])
self.kptband_comm.sum(taut_sG)
return taut_sG
def initialize(self, load_gauss=False):
self.presmooths[self.levels] = 8
self.postsmooths[self.levels] = 8
self.phis = [None] + [gd.zeros() for gd in self.gds[1:]]
self.residuals = [gd.zeros() for gd in self.gds]
self.rhos = [gd.zeros() for gd in self.gds]
self.op_coarse_weights = [[g.empty() for g in (gd, ) * 4]
for gd in self.gds[1:]]
scale = -0.25 / np.pi
for i, gd in enumerate(self.gds):
if i == 0:
nn = self.nn
weights = self.dielectric.eps_gradeps
else:
nn = 1
weights = self.op_coarse_weights[i - 1]
operators = [Laplace(gd, scale, nn)] + \
[Gradient(gd, j, scale, nn) for j in (0, 1, 2)]
self.operators.append(WeightedFDOperator(weights, operators))
if load_gauss:
self.load_gauss()
def update(self, wfs):
ddr_v = [Gradient(self.gd, v).apply for v in range(3)]
assert self.nspins == 1
self.taut_sG[:] = wfs.calculate_kinetic_energy_density(
self.taut_sG[:1], ddr_v)
# Add the pseudo core kinetic array
self.tauct.add(self.taut_sG[0])
# For periodic boundary conditions
if wfs.symmetry is not None:
wfs.symmetry.symmetrize(self.taut_sG[0], wfs.gd)
self.nt_grad2_sG[:] = 0.0
d_G = self.gd.empty()
for s in range(self.nspins):
for v in range(3):
ddr_v[v](self.density.nt_sG[s], d_G)
self.nt_grad2_sG[s] += d_G**2.0
def set_grid_descriptor(self, gd):
LDA.set_grid_descriptor(self, gd)
self.grad_v = [Gradient(gd, v).apply for v in range(3)]
def gradient(gd, x, nn):
out = gd.empty(3)
for i in (0, 1, 2):
Gradient(gd, i, 1.0, nn).apply(x, out[i])
return out
from gpaw.mpi import world
if world.size > 4:
# Grid is so small that domain decomposition cannot exceed 4 domains
assert world.size % 4 == 0
group, other = divmod(world.rank, 4)
ranks = np.arange(4 * group, 4 * (group + 1))
domain_comm = world.new_communicator(ranks)
else:
domain_comm = world
gd = GridDescriptor((8, 1, 1), (8.0, 1.0, 1.0), comm=domain_comm)
a = gd.zeros()
dadx = gd.zeros()
a[:, 0, 0] = np.arange(gd.beg_c[0], gd.end_c[0])
gradx = Gradient(gd, v=0)
print a.itemsize, a.dtype, a.shape
print dadx.itemsize, dadx.dtype, dadx.shape
gradx.apply(a, dadx)
# a = [ 0. 1. 2. 3. 4. 5. 6. 7.]
#
# da
# -- = [-2.5 1. 1. 1. 1. 1. 1. -2.5]
# dx
dadx = gd.collect(dadx, broadcast=True)
assert dadx[3, 0, 0] == 1.0 and np.sum(dadx[:, 0, 0]) == 0.0
gd = GridDescriptor((1, 8, 1), (1.0, 8.0, 1.0), (1, 0, 1), comm=domain_comm)
dady = gd.zeros()
def calculate(self, spin=0):
"""Calculate the non-interacting density response function. """
calc = self.calc
kd = self.kd
gd = self.gd
sdisp_cd = gd.sdisp_cd
ibzk_kc = self.ibzk_kc
bzk_kc = self.bzk_kc
kq_k = self.kq_k
pt = self.pt
f_kn = self.f_kn
e_kn = self.e_kn
# Matrix init
chi0_wGG = np.zeros((self.Nw_local, self.npw, self.npw), dtype=complex)
if not (f_kn > self.ftol).any():
self.chi0_wGG = chi0_wGG
return
if self.hilbert_trans:
specfunc_wGG = np.zeros((self.NwS_local, self.npw, self.npw),
dtype=complex)
# Prepare for the derivative of pseudo-wavefunction
if self.optical_limit:
d_c = [Gradient(gd, i, n=4, dtype=complex).apply for i in range(3)]
dpsit_g = gd.empty(dtype=complex)
tmp = np.zeros((3), dtype=complex)
rho_G = np.zeros(self.npw, dtype=complex)
t0 = time()
t_get_wfs = 0
for k in range(self.kstart, self.kend):
# Find corresponding kpoint in IBZ
ibzkpt1 = kd.kibz_k[k]
if self.optical_limit:
ibzkpt2 = ibzkpt1
else:
ibzkpt2 = kd.kibz_k[kq_k[k]]
for n in range(self.nstart, self.nend):
# print >> self.txt, k, n, t_get_wfs, time() - t0
t1 = time()
psitold_g = self.get_wavefunction(ibzkpt1, n, True, spin=spin)
t_get_wfs += time() - t1
psit1new_g = kd.transform_wave_function(psitold_g, k)
P1_ai = pt.dict()
pt.integrate(psit1new_g, P1_ai, k)
psit1_g = psit1new_g.conj() * self.expqr_g
for m in range(self.nbands):
if self.hilbert_trans:
check_focc = (f_kn[ibzkpt1, n] -
f_kn[ibzkpt2, m]) > self.ftol
else:
check_focc = np.abs(f_kn[ibzkpt1, n] -
f_kn[ibzkpt2, m]) > self.ftol
t1 = time()
psitold_g = self.get_wavefunction(ibzkpt2,
m,
check_focc,
spin=spin)
t_get_wfs += time() - t1
if check_focc:
psit2_g = kd.transform_wave_function(
psitold_g, kq_k[k])
P2_ai = pt.dict()
pt.integrate(psit2_g, P2_ai, kq_k[k])
# fft
tmp_g = np.fft.fftn(
psit2_g * psit1_g) * self.vol / self.nG0
for iG in range(self.npw):
index = self.Gindex_G[iG]
rho_G[iG] = tmp_g[index[0], index[1], index[2]]
if self.optical_limit:
phase_cd = np.exp(2j * pi * sdisp_cd *
bzk_kc[kq_k[k], :, np.newaxis])
for ix in range(3):
d_c[ix](psit2_g, dpsit_g, phase_cd)
tmp[ix] = gd.integrate(psit1_g * dpsit_g)
rho_G[0] = -1j * np.dot(self.qq_v, tmp)
# PAW correction
for a, id in enumerate(calc.wfs.setups.id_a):
P_p = np.outer(P1_ai[a].conj(), P2_ai[a]).ravel()
gemv(1.0, self.phi_aGp[a], P_p, 1.0, rho_G)
if self.optical_limit:
rho_G[0] /= e_kn[ibzkpt2, m] - e_kn[ibzkpt1, n]
rho_GG = np.outer(rho_G, rho_G.conj())
if not self.hilbert_trans:
for iw in range(self.Nw_local):
w = self.w_w[iw + self.wstart] / Hartree
C = (f_kn[ibzkpt1, n] - f_kn[ibzkpt2, m]) / (
w + e_kn[ibzkpt1, n] - e_kn[ibzkpt2, m] +
1j * self.eta)
axpy(C, rho_GG, chi0_wGG[iw])
else:
focc = f_kn[ibzkpt1, n] - f_kn[ibzkpt2, m]
w0 = e_kn[ibzkpt2, m] - e_kn[ibzkpt1, n]
scal(focc, rho_GG)
# calculate delta function
w0_id = int(w0 / self.dw)
if w0_id + 1 < self.NwS:
# rely on the self.NwS_local is equal in each node!
if self.wScomm.rank == w0_id // self.NwS_local:
alpha = (w0_id + 1 -
w0 / self.dw) / self.dw
axpy(alpha, rho_GG,
specfunc_wGG[w0_id % self.NwS_local])
if self.wScomm.rank == (w0_id +
1) // self.NwS_local:
alpha = (w0 / self.dw - w0_id) / self.dw
axpy(
alpha, rho_GG,
specfunc_wGG[(w0_id + 1) %
self.NwS_local])
# deltaw = delta_function(w0, self.dw, self.NwS, self.sigma)
# for wi in range(self.NwS_local):
# if deltaw[wi + self.wS1] > 1e-8:
# specfunc_wGG[wi] += tmp_GG * deltaw[wi + self.wS1]
if self.nkpt == 1:
if n == 0:
dt = time() - t0
totaltime = dt * self.nband_local
self.printtxt(
'Finished n 0 in %f seconds, estimated %f seconds left.'
% (dt, totaltime))
if rank == 0 and self.nband_local // 5 > 0:
if n > 0 and n % (self.nband_local // 5) == 0:
dt = time() - t0
self.printtxt(
'Finished n %d in %f seconds, estimated %f seconds left.'
% (n, dt, totaltime - dt))
if calc.wfs.world.size != 1:
self.kcomm.barrier()
if k == 0:
dt = time() - t0
totaltime = dt * self.nkpt_local
self.printtxt(
'Finished k 0 in %f seconds, estimated %f seconds left.' %
(dt, totaltime))
if rank == 0 and self.nkpt_local // 5 > 0:
if k > 0 and k % (self.nkpt_local // 5) == 0:
dt = time() - t0
self.printtxt(
'Finished k %d in %f seconds, estimated %f seconds left. '
% (k, dt, totaltime - dt))
self.printtxt('Finished summation over k')
self.kcomm.barrier()
del rho_GG, rho_G
# Hilbert Transform
if not self.hilbert_trans:
self.kcomm.sum(chi0_wGG)
else:
self.kcomm.sum(specfunc_wGG)
if self.wScomm.size == 1:
if not self.full_hilbert_trans:
chi0_wGG = hilbert_transform(
specfunc_wGG, self.Nw, self.dw,
self.eta)[self.wstart:self.wend]
else:
chi0_wGG = full_hilbert_transform(
specfunc_wGG, self.Nw, self.dw,
self.eta)[self.wstart:self.wend]
self.printtxt('Finished hilbert transform !')
del specfunc_wGG
else:
# redistribute specfunc_wGG to all nodes
assert self.NwS % size == 0
NwStmp1 = (rank % self.kcomm.size) * self.NwS // size
NwStmp2 = (rank % self.kcomm.size + 1) * self.NwS // size
specfuncnew_wGG = specfunc_wGG[NwStmp1:NwStmp2]
del specfunc_wGG
coords = np.zeros(self.wcomm.size, dtype=int)
nG_local = self.npw**2 // self.wcomm.size
if self.wcomm.rank == self.wcomm.size - 1:
nG_local = self.npw**2 - (self.wcomm.size - 1) * nG_local
self.wcomm.all_gather(np.array([nG_local]), coords)
specfunc_Wg = SliceAlongFrequency(specfuncnew_wGG, coords,
self.wcomm)
self.printtxt('Finished Slice Along Frequency !')
if not self.full_hilbert_trans:
chi0_Wg = hilbert_transform(specfunc_Wg, self.Nw, self.dw,
self.eta)[:self.Nw]
else:
chi0_Wg = full_hilbert_transform(specfunc_Wg, self.Nw,
self.dw,
self.eta)[:self.Nw]
self.printtxt('Finished hilbert transform !')
self.comm.barrier()
del specfunc_Wg
chi0_wGG = SliceAlongOrbitals(chi0_Wg, coords, self.wcomm)
self.printtxt('Finished Slice along orbitals !')
self.comm.barrier()
del chi0_Wg
self.chi0_wGG = chi0_wGG / self.vol
self.printtxt('')
self.printtxt('Finished chi0 !')
return
def solve_electric_field(self, phi):
for v in range(3):
Gradient(self.gd, v, n=3).apply(-1.0 * self.sign * phi, self.electric_field[v])
def get_gradient_ops(gd, nn):
return [Gradient(gd, v, n=nn).apply for v in range(3)]
def set_grid_descriptor(self, gd):
SolvationPoissonSolver.set_grid_descriptor(self, gd)
self.gradx = Gradient(gd, 0, 1.0, self.nn)
self.grady = Gradient(gd, 1, 1.0, self.nn)
self.gradz = Gradient(gd, 2, 1.0, self.nn)
def __init__(self,
iidx=None,
jidx=None,
pspin=None,
kpt=None,
paw=None,
string=None,
fijscale=1,
dtype=float):
if string is not None:
self.fromstring(string, dtype)
return None
# normal entry
PairDensity.__init__(self, paw)
PairDensity.initialize(self, kpt, iidx, jidx)
self.pspin = pspin
self.energy = 0.0
self.fij = 0.0
self.me = np.zeros((3), dtype=dtype)
self.mur = np.zeros((3), dtype=dtype)
self.muv = np.zeros((3), dtype=dtype)
self.magn = np.zeros((3), dtype=dtype)
self.kpt_comm = paw.wfs.kd.comm
# leave empty if not my kpt
if kpt is None:
return
wfs = paw.wfs
gd = wfs.gd
self.energy = kpt.eps_n[jidx] - kpt.eps_n[iidx]
self.fij = (kpt.f_n[iidx] - kpt.f_n[jidx]) * fijscale
# calculate matrix elements -----------
# length form ..........................
# course grid contribution
# <i|r|j> is the negative of the dipole moment (because of negative
# e- charge)
me = -gd.calculate_dipole_moment(self.get())
# augmentation contributions
ma = np.zeros(me.shape, dtype=dtype)
pos_av = paw.atoms.get_positions() / Bohr
for a, P_ni in kpt.P_ani.items():
Ra = pos_av[a]
Pi_i = P_ni[self.i].conj()
Pj_i = P_ni[self.j]
Delta_pL = wfs.setups[a].Delta_pL
ni = len(Pi_i)
ma0 = 0
ma1 = np.zeros(me.shape, dtype=me.dtype)
for i in range(ni):
for j in range(ni):
pij = Pi_i[i] * Pj_i[j]
ij = packed_index(i, j, ni)
# L=0 term
ma0 += Delta_pL[ij, 0] * pij
# L=1 terms
if wfs.setups[a].lmax >= 1:
# see spherical_harmonics.py for
# L=1:y L=2:z; L=3:x
ma1 += np.array([
Delta_pL[ij, 3], Delta_pL[ij, 1], Delta_pL[ij, 2]
]) * pij
ma += sqrt(4 * pi / 3) * ma1 + Ra * sqrt(4 * pi) * ma0
gd.comm.sum(ma)
self.me = sqrt(self.energy * self.fij) * (me + ma)
self.mur = -(me + ma)
# velocity form .............................
if self.lcao:
# XXX Velocity form not supported in LCAO
return
me = np.zeros(self.mur.shape, dtype=dtype)
# get derivatives
dtype = self.wfj.dtype
dwfj_cg = gd.empty((3), dtype=dtype)
if not hasattr(gd, 'ddr'):
gd.ddr = [Gradient(gd, c, dtype=dtype).apply for c in range(3)]
for c in range(3):
gd.ddr[c](self.wfj, dwfj_cg[c], kpt.phase_cd)
me[c] = gd.integrate(self.wfi.conj() * dwfj_cg[c])
if 0:
me2 = np.zeros(self.mur.shape)
for c in range(3):
gd.ddr[c](self.wfi, dwfj_cg[c], kpt.phase_cd)
me2[c] = gd.integrate(self.wfj * dwfj_cg[c])
print(me, -me2, me2 + me)
# augmentation contributions
ma = np.zeros(me.shape, dtype=me.dtype)
for a, P_ni in kpt.P_ani.items():
Pi_i = P_ni[self.i].conj()
Pj_i = P_ni[self.j]
nabla_iiv = paw.wfs.setups[a].nabla_iiv
for c in range(3):
for i1, Pi in enumerate(Pi_i):
for i2, Pj in enumerate(Pj_i):
ma[c] += Pi * Pj * nabla_iiv[i1, i2, c]
gd.comm.sum(ma)
self.muv = -(me + ma) / self.energy
# magnetic transition dipole ................
r_cg, r2_g = coordinates(gd)
magn = np.zeros(me.shape, dtype=dtype)
wfi_g = self.wfi.conj()
for ci in range(3):
cj = (ci + 1) % 3
ck = (ci + 2) % 3
magn[ci] = gd.integrate(wfi_g * r_cg[cj] * dwfj_cg[ck] -
wfi_g * r_cg[ck] * dwfj_cg[cj])
# augmentation contributions
ma = np.zeros(magn.shape, dtype=magn.dtype)
for a, P_ni in kpt.P_ani.items():
Pi_i = P_ni[self.i].conj()
Pj_i = P_ni[self.j]
rnabla_iiv = paw.wfs.setups[a].rnabla_iiv
for c in range(3):
for i1, Pi in enumerate(Pi_i):
for i2, Pj in enumerate(Pj_i):
ma[c] += Pi * Pj * rnabla_iiv[i1, i2, c]
gd.comm.sum(ma)
self.magn = -alpha / 2. * (magn + ma)
def integrand(self):
# polarisation in the direction of vk
costh = self.angle['x']
sinth = sqrt(1. - costh**2)
sinphi = sin(self.angle['phi'])
cosphi = cos(self.angle['phi'])
eps0 = np.array([sinth * cosphi, sinth * sinphi, costh])
vk = self.k * eps0
# polarisation at the magic angle
costhm = self.costhm
sinthm = self.sinthm
sinpsi = sin(self.angle['psi'])
cospsi = cos(self.angle['psi'])
epsm = np.array([
sinthm * (cosphi * sinpsi * costh + sinphi * cospsi) +
costhm * cosphi * sinth,
sinthm * (sinphi * sinpsi * costh - cosphi * cospsi) +
costhm * sinphi * sinth, costhm * costh - sinthm * sinth * sinpsi
])
# initial and final state on the grid
initial_G = self.initial.get_grid()
final_G = self.final.get_grid(vk, self.r0)
ini_analyt = H1s(self.initial.gd, self.r0)
gd = self.initial.gd
if self.form == 'L':
if_G = initial_G * final_G
omega = self.get_omega()
if 0:
me_c = []
for c in range(3):
xyz_G = ((np.arange(gd.n_c[c]) + gd.beg_c[c]) * gd.h_c[c] -
self.r0[c])
shape = [1, 1, 1]
shape[c] = -1
xyz_G.shape = tuple(shape)
np.resize(xyz_G, gd.n_c)
me_c.append(gd.integrate(if_G * xyz_G))
me_c = np.array(me_c) * omega
else:
me_c = gd.calculate_dipole_moment(if_G)
me_c += self.r0 * gd.integrate(if_G)
me_c *= -omega
elif self.form == 'V':
dtype = final_G.dtype
phase_cd = np.ones((3, 2), dtype)
if not hasattr(gd, 'ddr'):
gd.ddr = [Gradient(gd, c, dtype=dtype).apply for c in range(3)]
dfinal_G = gd.empty(dtype=dtype)
me_c = np.empty(3, dtype=dtype)
for c in range(3):
# print "rank, c, apply", rank, c, dtype, final_G.shape, dfinal_G.shape
gd.ddr[c](final_G, dfinal_G, phase_cd)
me_c[c] = gd.integrate(initial_G * dfinal_G)
else:
raise NotImplementedError
# print self.k, self.initial.get_me_c(vk, self.form)[0].imag, me_c[0].imag
if 0:
omega = self.get_omega()
me_analyt = ini_analyt.get_me_c(vk, self.form)[0].imag
me = me_c[0].imag
def ds(me):
return self.k / omega * me**2
print omega, ds(me_analyt), ds(me), me_analyt, me
# print 'analyt', self.initial.get_me_c(vk, self.form)
# print 'num', me_c
# print 'analyt/num', self.initial.get_me_c(vk, self.form) / me_c
# return the squared matrix elements
T2 = []
for eps in [eps0, epsm]:
me = np.dot(eps, me_c)
# print "eps, T2:", eps, (me * me.conj()).real
T2.append((me * me.conj()).real)
# print vk, T2
return T2[0], T2[1]
def get_dipole_transitions(calc, momname=None, basename=None):
"""
Finds the dipole matrix elements:
<\psi_n|\nabla|\psi_m> = <u_n|nabla|u_m> + ik<u_n|u_m> where psi_n = u_n(r)*exp(ikr).
Input:
atoms Relevant ASE atoms object
momname Suffix for the dipole transition file
basename Suffix used for the gs.gpw file
Output:
dip_vknm.npy Array with dipole matrix elements
"""
# par = MPI4PY() # FIXME: use a comm from gpaw
#calc = atoms.calc
bzk_kc = calc.get_ibz_k_points()
n = calc.wfs.bd.nbands
nk = np.shape(bzk_kc)[0]
wfs = {}
parprint("Distributing wavefunctions.")
kcomm = calc.wfs.kd.comm
world = calc.wfs.world
if not calc.wfs.positions_set:
calc.initialize_positions()
for k in range(nk):
# Collects the wavefunctions and the projections to rank 0. Periodic -> u_n(r)
spin = 0 # FIXME
wf = np.array([
calc.wfs.get_wave_function_array(
i, k, spin, realspace=True, periodic=True) for i in range(n)
],
dtype=complex)
P_nI = calc.wfs.collect_projections(k, spin)
# Distributes the information to rank k % size.
if kcomm.rank == 0:
if k % kcomm.size == kcomm.rank:
wfs[k] = wf, P_nI
else:
kcomm.send(P_nI, dest=k % kcomm.size, tag=nk + k)
kcomm.send(wf, dest=k % kcomm.size, tag=k)
else:
if k % kcomm.size == kcomm.rank:
nproj = sum(setup.ni for setup in calc.wfs.setups)
if not calc.wfs.collinear:
nproj *= 2
P_nI = np.empty((calc.wfs.bd.nbands, nproj), calc.wfs.dtype)
shape = () if calc.wfs.collinear else (2, )
wf = np.tile(
calc.wfs.empty(shape, global_array=True, realspace=True),
(n, 1, 1, 1))
kcomm.receive(P_nI, src=0, tag=nk + k)
kcomm.receive(wf, src=0, tag=k)
wfs[k] = wf, P_nI
parprint("Evaluating dipole transition matrix elements.")
dip_vknm = np.zeros((3, nk, n, n), dtype=complex)
overlap_knm = np.zeros((nk, n, n), dtype=complex)
nabla_v = [
Gradient(calc.wfs.gd, v, 1.0, 4, complex).apply for v in range(3)
]
phases = np.ones((3, 2), dtype=complex)
grad_nv = calc.wfs.gd.zeros((n, 3), complex)
for k, (wf, P_nI) in wfs.items():
# Calculate <phit|nabla|phit> for the pseudo wavefunction
for v in range(3):
for i in range(n):
nabla_v[v](wf[i], grad_nv[i, v], phases)
dip_vknm[:, k] = np.transpose(calc.wfs.gd.integrate(wf, grad_nv),
(2, 0, 1))
overlap_knm[k] = [calc.wfs.gd.integrate(wf[i], wf) for i in range(n)]
k_v = np.dot(calc.wfs.kd.ibzk_kc[k], calc.wfs.gd.icell_cv) * 2 * pi
dip_vknm[:, k] += 1j * k_v[:, None, None] * overlap_knm[None, k, :, :]
# The PAW corrections are added - see https://wiki.fysik.dtu.dk/gpaw/dev/documentation/tddft/dielectric_response.html#paw-terms
I1 = 0
# np.einsum is slow but very memory efficient.
for a, setup in enumerate(calc.wfs.setups):
I2 = I1 + setup.ni
P_ni = P_nI[:, I1:I2]
dip_vknm[:, k, :, :] += np.einsum('ni,ijv,mj->vnm', P_ni.conj(),
setup.nabla_iiv, P_ni)
I1 = I2
world.sum(dip_vknm)
if world.rank == 0:
np.save('dip_vknm{}.npy'.format(make_suffix(momname)), dip_vknm)
def solve_rho(self):
self.charge_density *= 0.0
dmy = self.gd.empty()
for v in range(3):
Gradient(self.gd, v, n=3).apply(self.polarization_total[v], dmy)
self.charge_density -= dmy
from gpaw.mpi import world
if world.size > 4:
# Grid is so small that domain decomposition cannot exceed 4 domains
assert world.size % 4 == 0
group, other = divmod(world.rank, 4)
ranks = np.arange(4*group, 4*(group+1))
domain_comm = world.new_communicator(ranks)
else:
domain_comm = world
gd = GridDescriptor((8, 1, 1), (8.0, 1.0, 1.0), comm=domain_comm)
a = gd.zeros()
dadx = gd.zeros()
a[:, 0, 0] = np.arange(gd.beg_c[0], gd.end_c[0])
gradx = Gradient(gd, v=0)
print a.itemsize, a.dtype, a.shape
print dadx.itemsize, dadx.dtype, dadx.shape
gradx.apply(a, dadx)
# a = [ 0. 1. 2. 3. 4. 5. 6. 7.]
#
# da
# -- = [-2.5 1. 1. 1. 1. 1. 1. -2.5]
# dx
dadx = gd.collect(dadx, broadcast=True)
assert dadx[3, 0, 0] == 1.0 and np.sum(dadx[:, 0, 0]) == 0.0
gd = GridDescriptor((1, 8, 1), (1.0, 8.0, 1.0), (1, 0, 1), comm=domain_comm)
dady = gd.zeros()
class ADM12PoissonSolver(SolvationPoissonSolver):
"""Poisson solver with dielectric.
Following O. Andreussi, I. Dabo, and N. Marzari,
J. Chem. Phys. 136, 064102 (2012).
Warning: Not intended for production use, as it is not tested
thouroughly!
XXX TODO : * Correction for charged systems???
* Check: Can the polarization charge introduce a monopole?
* Convergence problems depending on eta. Apparently this
method works best with FFT as in the original Paper.
* Optimize numerics.
"""
def __init__(self,
nn=3,
relax='J',
eps=2e-10,
maxiter=1000,
remove_moment=None,
eta=.6,
use_charge_center=False):
"""Constructor for ADM12PoissonSolver.
Additional arguments not present in SolvationPoissonSolver:
eta -- linear mixing parameter
"""
adm12_warning = UserWarning(
'ADM12PoissonSolver is not tested thoroughly'
' and therefore not recommended for production code!')
warnings.warn(adm12_warning)
self.eta = eta
SolvationPoissonSolver.__init__(self,
nn,
relax,
eps,
maxiter,
remove_moment,
use_charge_center=use_charge_center)
def set_grid_descriptor(self, gd):
SolvationPoissonSolver.set_grid_descriptor(self, gd)
self.gradx = Gradient(gd, 0, 1.0, self.nn)
self.grady = Gradient(gd, 1, 1.0, self.nn)
self.gradz = Gradient(gd, 2, 1.0, self.nn)
def get_description(self):
if len(self.operators) == 0:
return 'uninitialized ADM12PoissonSolver'
else:
description = SolvationPoissonSolver.get_description(self)
return description.replace('solver with',
'ADM12 solver with dielectric and')
def _init(self):
if self._initialized:
return
self.rho_iter = self.gd.zeros()
self.d_phi = self.gd.empty()
return SolvationPoissonSolver._init(self)
def solve(self,
phi,
rho,
charge=None,
eps=None,
maxcharge=1e-6,
zero_initial_phi=False,
timer=None):
self._init()
if self.gd.pbc_c.all():
actual_charge = self.gd.integrate(rho)
if abs(actual_charge) > maxcharge:
raise NotImplementedError(
'charged periodic systems are not implemented')
return FDPoissonSolver.solve(self,
phi,
rho,
charge,
eps,
maxcharge,
zero_initial_phi,
timer=timer)
def solve_neutral(self, phi, rho, eps=2e-10, timer=None):
self._init()
self.rho = rho
return SolvationPoissonSolver.solve_neutral(self, phi, rho, eps)
def iterate2(self, step, level=0):
self._init()
if level == 0:
epsr, dx_epsr, dy_epsr, dz_epsr = self.dielectric.eps_gradeps
self.gradx.apply(self.phis[0], self.d_phi)
sp = dx_epsr * self.d_phi
self.grady.apply(self.phis[0], self.d_phi)
sp += dy_epsr * self.d_phi
self.gradz.apply(self.phis[0], self.d_phi)
sp += dz_epsr * self.d_phi
self.rho_iter = self.eta / (4. * np.pi) * sp + \
(1. - self.eta) * self.rho_iter
self.rhos[0][:] = (self.rho_iter + self.rho) / epsr
return SolvationPoissonSolver.iterate2(self, step, level)
def initialize(self):
self.eta /= Hartree
self.ecut /= Hartree
calc = self.calc
self.nspins = self.calc.wfs.nspins
# kpoint init
self.kd = kd = calc.wfs.kd
self.nikpt = kd.nibzkpts
self.ftol /= kd.nbzkpts
# cell init
self.acell_cv = calc.wfs.gd.cell_cv
self.acell_cv, self.bcell_cv, self.vol, self.BZvol = \
get_primitive_cell(self.acell_cv,rpad=self.rpad)
# grid init
gd = calc.wfs.gd.new_descriptor(comm=serial_comm)
self.pbc = gd.pbc_c
self.gd = gd
self.nG0 = np.prod(gd.N_c)
# Number of grid points and volume including zero padding
self.nGrpad = gd.N_c * self.rpad
self.nG0rpad = np.prod(self.nGrpad)
self.d_c = [
Gradient(gd, i, n=4, dtype=complex).apply for i in range(3)
]
# obtain eigenvalues, occupations
nibzkpt = kd.nibzkpts
kweight_k = kd.weight_k
self.eFermi = self.calc.occupations.get_fermi_level()
try:
self.e_skn
self.printtxt('Use eigenvalues from user.')
except:
self.printtxt('Use eigenvalues from the calculator.')
self.e_skn = {}
self.f_skn = {}
for ispin in range(self.nspins):
self.e_skn[ispin] = np.array([
calc.get_eigenvalues(kpt=k, spin=ispin)
for k in range(nibzkpt)
]) / Hartree
self.f_skn[ispin] = np.array([
calc.get_occupation_numbers(kpt=k, spin=ispin) /
kweight_k[k] for k in range(nibzkpt)
]) / kd.nbzkpts
#self.printtxt('Eigenvalues(k=0) are:')
#print >> self.txt, self.e_skn[0][0] * Hartree
self.enoshift_skn = {}
for ispin in range(self.nspins):
self.enoshift_skn[ispin] = self.e_skn[ispin].copy()
if self.eshift is not None:
self.add_discontinuity(self.eshift)
self.printtxt('Shift unoccupied bands by %f eV' % (self.eshift))
# k + q init
if self.q_c is not None:
self.qq_v = np.dot(self.q_c, self.bcell_cv) # summation over c
if self.optical_limit:
kq_k = np.arange(kd.nbzkpts)
self.expqr_g = 1.
else:
r_vg = gd.get_grid_point_coordinates() # (3, nG)
qr_g = gemmdot(self.qq_v, r_vg, beta=0.0)
self.expqr_g = np.exp(-1j * qr_g)
del r_vg, qr_g
kq_k = kd.find_k_plus_q(self.q_c)
self.kq_k = kq_k
# Plane wave init
if self.G_plus_q:
self.npw, self.Gvec_Gc, self.Gindex_G = set_Gvectors(self.acell_cv,
self.bcell_cv,
self.gd.N_c,
self.ecut,
q=self.q_c)
else:
self.npw, self.Gvec_Gc, self.Gindex_G = set_Gvectors(
self.acell_cv, self.bcell_cv, self.gd.N_c, self.ecut)
# band init
if self.nbands is None:
self.nbands = calc.wfs.bd.nbands
self.nvalence = calc.wfs.nvalence
# Projectors init
setups = calc.wfs.setups
self.spos_ac = calc.atoms.get_scaled_positions()
if self.pwmode:
self.pt = PWLFC([setup.pt_j for setup in setups], self.calc.wfs.pd)
self.pt.set_positions(self.spos_ac)
else:
self.pt = LFC(gd, [setup.pt_j for setup in setups],
KPointDescriptor(self.kd.bzk_kc),
dtype=complex,
forces=True)
self.pt.set_positions(self.spos_ac)
# Printing calculation information
self.print_stuff()
return
def create_laplace(self, gd, scale=1.0, n=1, dtype=float):
operators = [Laplace(gd, scale, n, dtype)]
operators += [Gradient(gd, j, scale, n, dtype) for j in (0, 1, 2)]
return WeightedFDOperator(operators)
def calculate(self, seperate_spin=None):
"""Calculate the non-interacting density response function. """
calc = self.calc
kd = self.kd
gd = self.gd
sdisp_cd = gd.sdisp_cd
ibzk_kc = kd.ibzk_kc
bzk_kc = kd.bzk_kc
kq_k = self.kq_k
f_skn = self.f_skn
e_skn = self.e_skn
# Matrix init
chi0_wGG = np.zeros((self.Nw_local, self.npw, self.npw), dtype=complex)
if self.hilbert_trans:
specfunc_wGG = np.zeros((self.NwS_local, self.npw, self.npw), dtype = complex)
# Prepare for the derivative of pseudo-wavefunction
if self.optical_limit:
d_c = [Gradient(gd, i, n=4, dtype=complex).apply for i in range(3)]
dpsit_g = gd.empty(dtype=complex)
tmp = np.zeros((3), dtype=complex)
rhoG0_v = np.zeros(3, dtype=complex)
self.chi0G0_wGv = np.zeros((self.Nw_local, self.npw, 3), dtype=complex)
self.chi00G_wGv = np.zeros((self.Nw_local, self.npw, 3), dtype=complex)
specfuncG0_wGv = np.zeros((self.NwS_local, self.npw, 3), dtype=complex)
specfunc0G_wGv = np.zeros((self.NwS_local, self.npw, 3), dtype=complex)
use_zher = False
if self.eta < 1e-5:
use_zher = True
rho_G = np.zeros(self.npw, dtype=complex)
t0 = time()
if seperate_spin is None:
spinlist = np.arange(self.nspins)
else:
spinlist = [seperate_spin]
for spin in spinlist:
if not (f_skn[spin] > self.ftol).any():
self.chi0_wGG = chi0_wGG
continue
for k in range(self.kstart, self.kend):
k_pad = False
if k >= self.kd.nbzkpts:
k = 0
k_pad = True
# Find corresponding kpoint in IBZ
ibzkpt1 = kd.bz2ibz_k[k]
if self.optical_limit:
ibzkpt2 = ibzkpt1
else:
ibzkpt2 = kd.bz2ibz_k[kq_k[k]]
if self.pwmode:
N_c = self.gd.N_c
k_c = self.kd.ibzk_kc[ibzkpt1]
eikr1_R = np.exp(2j * pi * np.dot(np.indices(N_c).T, k_c / N_c).T)
k_c = self.kd.ibzk_kc[ibzkpt2]
eikr2_R = np.exp(2j * pi * np.dot(np.indices(N_c).T, k_c / N_c).T)
index1_g, phase1_g = kd.get_transform_wavefunction_index(self.gd.N_c - (self.pbc == False), k)
index2_g, phase2_g = kd.get_transform_wavefunction_index(self.gd.N_c - (self.pbc == False), kq_k[k])
for n in range(self.nvalbands):
if self.calc.wfs.world.size == 1:
if (self.f_skn[spin][ibzkpt1, n] - self.ftol < 0):
continue
t1 = time()
if self.pwmode:
u = self.kd.get_rank_and_index(spin, ibzkpt1)[1]
psitold_g = calc.wfs._get_wave_function_array(u, n, realspace=True, phase=eikr1_R)
else:
u = None
psitold_g = self.get_wavefunction(ibzkpt1, n, True, spin=spin)
psit1new_g = kd.transform_wave_function(psitold_g,k,index1_g,phase1_g)
P1_ai = self.pawstuff(psit1new_g, k, n, spin, u, ibzkpt1)
psit1_g = psit1new_g.conj() * self.expqr_g
for m in self.mlist:
if self.nbands > 1000 and m % 200 == 0:
print(' ', k, n, m, time() - t0, file=self.txt)
check_focc = (f_skn[spin][ibzkpt1, n] - f_skn[spin][ibzkpt2, m]) > self.ftol
if not self.pwmode:
psitold_g = self.get_wavefunction(ibzkpt2, m, check_focc, spin=spin)
if check_focc:
if self.pwmode:
u = self.kd.get_rank_and_index(spin, ibzkpt2)[1]
psitold_g = calc.wfs._get_wave_function_array(u, m, realspace=True, phase=eikr2_R)
psit2_g = kd.transform_wave_function(psitold_g, kq_k[k], index2_g, phase2_g)
# zero padding is included through the FFT
rho_g = np.fft.fftn(psit2_g * psit1_g, s=self.nGrpad) * self.vol / self.nG0rpad
# Here, planewave cutoff is applied
rho_G = rho_g.ravel()[self.Gindex_G]
if self.optical_limit:
phase_cd = np.exp(2j * pi * sdisp_cd * kd.bzk_kc[kq_k[k], :, np.newaxis])
for ix in range(3):
d_c[ix](psit2_g, dpsit_g, phase_cd)
tmp[ix] = gd.integrate(psit1_g * dpsit_g)
rho_G[0] = -1j * np.dot(self.qq_v, tmp)
for ix in range(3):
q2_c = np.diag((1,1,1))[ix] * self.qopt
qq2_v = np.dot(q2_c, self.bcell_cv) # summation over c
rhoG0_v[ix] = -1j * np.dot(qq2_v, tmp)
P2_ai = self.pawstuff(psit2_g, kq_k[k], m, spin, u, ibzkpt2)
for a, id in enumerate(calc.wfs.setups.id_a):
P_p = np.outer(P1_ai[a].conj(), P2_ai[a]).ravel()
gemv(1.0, self.phi_aGp[a], P_p, 1.0, rho_G)
if self.optical_limit:
gemv(1.0, self.phiG0_avp[a], P_p, 1.0, rhoG0_v)
if self.optical_limit:
if np.abs(self.enoshift_skn[spin][ibzkpt2, m] -
self.enoshift_skn[spin][ibzkpt1, n]) > 0.1/Hartree:
rho_G[0] /= self.enoshift_skn[spin][ibzkpt2, m] \
- self.enoshift_skn[spin][ibzkpt1, n]
rhoG0_v /= self.enoshift_skn[spin][ibzkpt2, m] \
- self.enoshift_skn[spin][ibzkpt1, n]
else:
rho_G[0] = 0.
rhoG0_v[:] = 0.
if k_pad:
rho_G[:] = 0.
if self.optical_limit:
rho0G_Gv = np.outer(rho_G.conj(), rhoG0_v)
rhoG0_Gv = np.outer(rho_G, rhoG0_v.conj())
rho0G_Gv[0,:] = rhoG0_v * rhoG0_v.conj()
rhoG0_Gv[0,:] = rhoG0_v * rhoG0_v.conj()
if not self.hilbert_trans:
if not use_zher:
rho_GG = np.outer(rho_G, rho_G.conj())
for iw in range(self.Nw_local):
w = self.w_w[iw + self.wstart] / Hartree
coef = ( 1. / (w + e_skn[spin][ibzkpt1, n] - e_skn[spin][ibzkpt2, m]
+ 1j * self.eta)
- 1. / (w - e_skn[spin][ibzkpt1, n] + e_skn[spin][ibzkpt2, m]
+ 1j * self.eta) )
C = (f_skn[spin][ibzkpt1, n] - f_skn[spin][ibzkpt2, m]) * coef
if use_zher:
czher(C.real, rho_G.conj(), chi0_wGG[iw])
else:
axpy(C, rho_GG, chi0_wGG[iw])
if self.optical_limit:
axpy(C, rho0G_Gv, self.chi00G_wGv[iw])
axpy(C, rhoG0_Gv, self.chi0G0_wGv[iw])
else:
rho_GG = np.outer(rho_G, rho_G.conj())
focc = f_skn[spin][ibzkpt1,n] - f_skn[spin][ibzkpt2,m]
w0 = e_skn[spin][ibzkpt2,m] - e_skn[spin][ibzkpt1,n]
scal(focc, rho_GG)
if self.optical_limit:
scal(focc, rhoG0_Gv)
scal(focc, rho0G_Gv)
# calculate delta function
w0_id = int(w0 / self.dw)
if w0_id + 1 < self.NwS:
# rely on the self.NwS_local is equal in each node!
if self.wScomm.rank == w0_id // self.NwS_local:
alpha = (w0_id + 1 - w0/self.dw) / self.dw
axpy(alpha, rho_GG, specfunc_wGG[w0_id % self.NwS_local] )
if self.optical_limit:
axpy(alpha, rho0G_Gv, specfunc0G_wGv[w0_id % self.NwS_local] )
axpy(alpha, rhoG0_Gv, specfuncG0_wGv[w0_id % self.NwS_local] )
if self.wScomm.rank == (w0_id+1) // self.NwS_local:
alpha = (w0 / self.dw - w0_id) / self.dw
axpy(alpha, rho_GG, specfunc_wGG[(w0_id+1) % self.NwS_local] )
if self.optical_limit:
axpy(alpha, rho0G_Gv, specfunc0G_wGv[(w0_id+1) % self.NwS_local] )
axpy(alpha, rhoG0_Gv, specfuncG0_wGv[(w0_id+1) % self.NwS_local] )
# deltaw = delta_function(w0, self.dw, self.NwS, self.sigma)
# for wi in range(self.NwS_local):
# if deltaw[wi + self.wS1] > 1e-8:
# specfunc_wGG[wi] += tmp_GG * deltaw[wi + self.wS1]
if self.kd.nbzkpts == 1:
if n == 0:
dt = time() - t0
totaltime = dt * self.nvalbands * self.nspins
self.printtxt('Finished n 0 in %d seconds, estimate %d seconds left.' %(dt, totaltime) )
if rank == 0 and self.nvalbands // 5 > 0:
if n > 0 and n % (self.nvalbands // 5) == 0:
dt = time() - t0
self.printtxt('Finished n %d in %d seconds, estimate %d seconds left.'%(n, dt, totaltime-dt))
if calc.wfs.world.size != 1:
self.kcomm.barrier()
if k == 0:
dt = time() - t0
totaltime = dt * self.nkpt_local * self.nspins
self.printtxt('Finished k 0 in %d seconds, estimate %d seconds left.' %(dt, totaltime))
if rank == 0 and self.nkpt_local // 5 > 0:
if k > 0 and k % (self.nkpt_local // 5) == 0:
dt = time() - t0
self.printtxt('Finished k %d in %d seconds, estimate %d seconds left. '%(k, dt, totaltime - dt) )
self.printtxt('Finished summation over k')
self.kcomm.barrier()
# Hilbert Transform
if not self.hilbert_trans:
for iw in range(self.Nw_local):
self.kcomm.sum(chi0_wGG[iw])
if self.optical_limit:
self.kcomm.sum(self.chi0G0_wGv[iw])
self.kcomm.sum(self.chi00G_wGv[iw])
if use_zher:
assert (np.abs(chi0_wGG[0,1:,0]) < 1e-10).all()
for iw in range(self.Nw_local):
chi0_wGG[iw] += chi0_wGG[iw].conj().T
for iG in range(self.npw):
chi0_wGG[iw, iG, iG] /= 2.
assert np.abs(np.imag(chi0_wGG[iw, iG, iG])) < 1e-10
else:
for iw in range(self.NwS_local):
self.kcomm.sum(specfunc_wGG[iw])
if self.optical_limit:
self.kcomm.sum(specfuncG0_wGv[iw])
self.kcomm.sum(specfunc0G_wGv[iw])
if self.wScomm.size == 1:
chi0_wGG = hilbert_transform(specfunc_wGG, self.w_w, self.Nw, self.dw, self.eta,
self.full_hilbert_trans)[self.wstart:self.wend]
self.printtxt('Finished hilbert transform !')
del specfunc_wGG
else:
# redistribute specfunc_wGG to all nodes
size = self.comm.size
assert self.NwS % size == 0
NwStmp1 = (rank % self.kcomm.size) * self.NwS // size
NwStmp2 = (rank % self.kcomm.size + 1) * self.NwS // size
specfuncnew_wGG = specfunc_wGG[NwStmp1:NwStmp2]
del specfunc_wGG
coords = np.zeros(self.wcomm.size, dtype=int)
nG_local = self.npw**2 // self.wcomm.size
if self.wcomm.rank == self.wcomm.size - 1:
nG_local = self.npw**2 - (self.wcomm.size - 1) * nG_local
self.wcomm.all_gather(np.array([nG_local]), coords)
specfunc_Wg = SliceAlongFrequency(specfuncnew_wGG, coords, self.wcomm)
self.printtxt('Finished Slice Along Frequency !')
chi0_Wg = hilbert_transform(specfunc_Wg, self.w_w, self.Nw, self.dw, self.eta,
self.full_hilbert_trans)[:self.Nw]
self.printtxt('Finished hilbert transform !')
self.comm.barrier()
del specfunc_Wg
chi0_wGG = SliceAlongOrbitals(chi0_Wg, coords, self.wcomm)
self.printtxt('Finished Slice along orbitals !')
self.comm.barrier()
del chi0_Wg
if self.optical_limit:
specfuncG0_WGv = np.zeros((self.NwS, self.npw, 3), dtype=complex)
specfunc0G_WGv = np.zeros((self.NwS, self.npw, 3), dtype=complex)
self.wScomm.all_gather(specfunc0G_wGv, specfunc0G_WGv)
self.wScomm.all_gather(specfuncG0_wGv, specfuncG0_WGv)
specfunc0G_wGv = specfunc0G_WGv
specfuncG0_wGv = specfuncG0_WGv
if self.optical_limit:
self.chi00G_wGv = hilbert_transform(specfunc0G_wGv, self.w_w, self.Nw, self.dw, self.eta,
self.full_hilbert_trans)[self.wstart:self.wend]
self.chi0G0_wGv = hilbert_transform(specfuncG0_wGv, self.w_w, self.Nw, self.dw, self.eta,
self.full_hilbert_trans)[self.wstart:self.wend]
if self.optical_limit:
self.chi00G_wGv /= self.vol
self.chi0G0_wGv /= self.vol
self.chi0_wGG = chi0_wGG
self.chi0_wGG /= self.vol
self.printtxt('')
self.printtxt('Finished chi0 !')
continue
print('------------------')
print(name, D)
print(cell[0])
print(cell[1])
print(cell[2])
for n in range(1, 6):
N = 2 * n + 2
gd = GridDescriptor((N, N, N), cell)
b_g = gd.zeros()
r_gv = gd.get_grid_point_coordinates().transpose((1, 2, 3, 0))
c_v = gd.cell_cv.sum(0) / 2
r_gv -= c_v
lap = Laplace(gd, n=n)
grad_v = [Gradient(gd, v, n=n) for v in range(3)]
assert lap.npoints == D * 2 * n + 1
for m in range(0, 2 * n + 1):
for ix in range(m + 1):
for iy in range(m - ix + 1):
iz = m - ix - iy
a_g = (r_gv**(ix, iy, iz)).prod(3)
if ix + iy + iz == 2 and max(ix, iy, iz) == 2:
r = 2.0
else:
r = 0.0
lap.apply(a_g, b_g)
e = b_g[n + 1, n + 1, n + 1] - r
assert abs(e) < 2e-12, e
for v in range(3):
grad_v[v].apply(a_g, b_g)
def allocate(self):
Density.allocate(self)
self.gradient = [Gradient(self.gd, i, 1.0, self.nn) for i in (0, 1, 2)]
