def update_intraband(self, ind_m, kpt):
kd = self.calc.wfs.kd
gd = self.calc.wfs.gd
k_c = kd.bzk_kc[kpt.K] + kpt.shift_c
k_v = 2 * np.pi * np.dot(k_c, np.linalg.inv(gd.cell_cv).T)
atomdata_a = self.calc.wfs.setups
ut_mvR = self.calc.wfs.gd.zeros((len(ind_m), 3), complex)
for ind, ut_vR in zip(ind_m, ut_mvR):
ut_vR[:] = self.make_derivative(kpt.s, kpt.K,
kpt.n1 + ind,
kpt.n1 + ind + 1)[0]
npartocc = len(ind_m)
ut_mR = kpt.ut_nR[ind_m]
nabla0_mmv = np.zeros((npartocc, npartocc, 3), dtype=complex)
for m in range(npartocc):
ut_vR = ut_mvR[m]
C_avi = [np.dot(atomdata.nabla_iiv.T, P_mi[ind_m[m]])
for atomdata, P_mi in zip(atomdata_a, kpt.P_ani)]
nabla0_mv = -self.calc.wfs.gd.integrate(ut_vR, ut_mR).T
nt_m = self.calc.wfs.gd.integrate(ut_mR[m], ut_mR)
nabla0_mv += 1j * nt_m[:, np.newaxis] * k_v[np.newaxis, :]
for C_vi, P_mi in zip(C_avi, kpt.P_ani):
gemm(1.0, C_vi, P_mi[ind_m[0:npartocc]], 1.0, nabla0_mv, 'c')
nabla0_mmv[m] = nabla0_mv
return nabla0_mmv
def project(self, x_nG):
"""Project the vector onto the unoccupied states at k+q.
The projection operator is defined as follows::
-- --
P = > |Psi ><Psi | = 1 - > |Psi ><Psi |
c,k+q -- c c -- v v
c k+q k+q v k+q k+q
"""
# It might be a good idea to move this functionality to its own class
assert x_nG.ndim == 4
assert self.kplusq is not None
# k+q
kplusqpt = self.kpt_u[self.kplusq]
# Occupied wave function
psit_nG = kplusqpt.psit_nG
# Project out occupied sub-space using blas routine gemm
m = len(x_nG)
n = len(psit_nG)
proj_mn = np.zeros((m, n), dtype=self.dtype)
gemm(self.gd.dv, psit_nG, x_nG, 0.0, proj_mn, 'c')
gemm(-1.0, psit_nG, proj_mn, 1.0, x_nG)
def _pseudo_braket(self, bra_xG, ket_yG, A_yx, square=None):
"""Calculate matrix elements of braket pairs of pseudo wave functions.
Low-level helper function. Results will be put in the *A_yx* array::
/ ~ * ~
A = | dG bra (G) ket (G)
nn' / n n'
Parameters:
bra_xG: ndarray
Set of bra-like vectors in which the matrix elements are evaluated.
key_yG: ndarray
Set of ket-like vectors in which the matrix elements are evaluated.
A_yx: ndarray
Matrix in which to put calculated elements. Take care: Due to the
difference in Fortran/C array order and the inherent BLAS nature,
the matrix has to be filled in transposed (conjugated in future?).
"""
assert bra_xG.shape[1:] == ket_yG.shape[1:]
assert (ket_yG.shape[0], bra_xG.shape[0]) == A_yx.shape
if square is None:
square = (bra_xG.shape[0] == ket_yG.shape[0])
dv = self.gd.dv
if ket_yG is bra_xG:
rk(dv, bra_xG, 0.0, A_yx)
elif self.hermitian and square:
r2k(0.5 * dv, bra_xG, ket_yG, 0.0, A_yx)
else:
gemm(dv, bra_xG, ket_yG, 0.0, A_yx, 'c')
def calculate_sic_matrixelements(self):
# overlap of pseudo wavefunctions
Htphit_mG = self.vt_mG * self.phit_mG
V_mm = np.zeros((self.nocc, self.nocc), dtype=self.dtype)
gemm(self.gd.dv, self.phit_mG, Htphit_mG, 0.0, V_mm, 't')
#
# PAW
for a, P_mi in self.P_ami.items():
for m, dH_p in enumerate(self.dH_amp[a]):
dH_ii = unpack(dH_p)
V_mm[m,:] += np.dot(P_mi[m], np.dot(dH_ii, P_mi.T))
# accumulate over grid-domains
self.gd.comm.sum(V_mm)
self.V_mm = V_mm
# Symmetrization of V and kappa-matrix:
K_mm = 0.5 * (V_mm - V_mm.T.conj())
V_mm = 0.5 * (V_mm + V_mm.T.conj())
# evaluate the kinetic correction
self.ekin = -np.trace(V_mm) * (3 - self.nspins)
return V_mm, K_mm, np.vdot(K_mm, K_mm).real
def integrate(self,
a_xg,
b_yg=None,
global_integral=True,
hermitian=False,
_transposed_result=None):
"""Integrate function(s) over domain.
a_xg: ndarray
Function(s) to be integrated.
b_yg: ndarray
If present, integrate a_xg.conj() * b_yg.
global_integral: bool
If the array(s) are distributed over several domains, then the
total sum will be returned. To get the local contribution
only, use global_integral=False.
hermitian: bool
Result is hermitian.
_transposed_result: ndarray
Long story. Don't use this unless you are a method of the
MatrixOperator class ..."""
xshape = a_xg.shape[:-3]
if b_yg is None:
# Only one array:
result = a_xg.reshape(xshape + (-1, )).sum(axis=-1) * self.dv
if global_integral:
if result.ndim == 0:
result = self.comm.sum(result)
else:
self.comm.sum(result)
return result
A_xg = np.ascontiguousarray(a_xg.reshape((-1, ) + a_xg.shape[-3:]))
B_yg = np.ascontiguousarray(b_yg.reshape((-1, ) + b_yg.shape[-3:]))
if _transposed_result is None:
result_yx = np.zeros((len(B_yg), len(A_xg)), A_xg.dtype)
else:
result_yx = _transposed_result
global_integral = False
if a_xg is b_yg:
rk(self.dv, A_xg, 0.0, result_yx)
elif hermitian:
r2k(0.5 * self.dv, A_xg, B_yg, 0.0, result_yx)
else:
gemm(self.dv, A_xg, B_yg, 0.0, result_yx, 'c')
if global_integral:
self.comm.sum(result_yx)
yshape = b_yg.shape[:-3]
result = result_yx.T.reshape(xshape + yshape)
if result.ndim == 0:
return result.item()
else:
return result
def sqrt_matrix(a, preserve=False):
"""Get the sqrt of a symmetric matrix a (diagonalize is used).
The matrix is kept if preserve=True, a=sqrt(a) otherwise."""
n = len(a)
if debug:
assert a.flags.contiguous
assert a.dtype == float
assert a.shape == (n, n)
if preserve:
b = a.copy()
else:
b = a
# diagonalize to get the form b = Z * D * Z^T
# where D is diagonal
D = np.empty((n,))
diagonalize(b, D)
ZT = b.copy()
Z = np.transpose(b)
# c = Z * sqrt(D)
c = Z * np.sqrt(D)
# sqrt(b) = c * Z^T
gemm(1., ZT, np.ascontiguousarray(c), 0., b)
return b
def calculate_hamiltonian_matrix(self, hamiltonian, wfs, kpt, root=-1):
# XXX document parallel stuff, particularly root parameter
assert self.has_initialized
vt_G = hamiltonian.vt_sG[kpt.s]
H_MM = np.empty((wfs.ksl.mynao, wfs.ksl.nao), wfs.dtype)
wfs.timer.start('Potential matrix')
wfs.basis_functions.calculate_potential_matrix(vt_G, H_MM, kpt.q)
wfs.timer.stop('Potential matrix')
# Add atomic contribution
#
# -- a a a*
# H += > P dH P
# mu nu -- mu i ij nu j
# aij
#
wfs.timer.start('Atomic Hamiltonian')
Mstart = wfs.basis_functions.Mstart
Mstop = wfs.basis_functions.Mstop
for a, P_Mi in kpt.P_aMi.items():
dH_ii = np.asarray(unpack(hamiltonian.dH_asp[a][kpt.s]), wfs.dtype)
dHP_iM = np.zeros((dH_ii.shape[1], P_Mi.shape[0]), wfs.dtype)
# (ATLAS can't handle uninitialized output array)
gemm(1.0, P_Mi, dH_ii, 0.0, dHP_iM, 'c')
gemm(1.0, dHP_iM, P_Mi[Mstart:Mstop], 1.0, H_MM)
wfs.timer.stop('Atomic Hamiltonian')
wfs.timer.start('Distribute overlap matrix')
H_MM = wfs.ksl.distribute_overlap_matrix(H_MM, root)
wfs.timer.stop('Distribute overlap matrix')
H_MM += wfs.T_qMM[kpt.q]
return H_MM
def project(self, x_nG):
"""Project the vector onto the unoccupied states at k+q.
The projection operator is defined as follows::
-- --
P = > |Psi ><Psi | = 1 - > |Psi ><Psi |
c,k+q -- c c -- v v
c k+q k+q v k+q k+q
"""
# It might be a good idea to move this functionality to its own class
assert x_nG.ndim == 4
assert self.kplusq is not None
# k+q
kplusqpt = self.kpt_u[self.kplusq]
# Occupied wave function
psit_nG = kplusqpt.psit_nG
# Project out occupied sub-space using blas routine gemm
m = len(x_nG)
n = len(psit_nG)
proj_mn = np.zeros((m, n), dtype=self.dtype)
gemm(self.gd.dv, psit_nG, x_nG, 0.0, proj_mn, 'c')
gemm(-1.0, psit_nG, proj_mn, 1.0, x_nG)
def update_hilbert(self, n_mG, deps_m, df_m, chi0_wGG):
"""Update spectral function.
Updates spectral function A_wGG and saves it to chi0_wGG for
later hilbert-transform."""
self.timer.start('prep')
beta = (2**0.5 - 1) * self.domega0 / self.omega2
o_m = abs(deps_m)
w_m = (o_m / (self.domega0 + beta * o_m)).astype(int)
o1_m = self.omega_w[w_m]
o2_m = self.omega_w[w_m + 1]
p_m = self.prefactor * abs(df_m) / (o2_m - o1_m)**2 # XXX abs()?
p1_m = p_m * (o2_m - o_m)
p2_m = p_m * (o_m - o1_m)
self.timer.stop('prep')
if self.blockcomm.size > 1:
for p1, p2, n_G, w in zip(p1_m, p2_m, n_mG, w_m):
myn_G = n_G[self.Ga:self.Gb].reshape((-1, 1))
gemm(p1, n_G.reshape((-1, 1)), myn_G, 1.0, chi0_wGG[w], 'c')
gemm(p2, n_G.reshape((-1, 1)), myn_G, 1.0, chi0_wGG[w + 1],
'c')
return
for p1, p2, n_G, w in zip(p1_m, p2_m, n_mG, w_m):
czher(p1, n_G.conj(), chi0_wGG[w])
czher(p2, n_G.conj(), chi0_wGG[w + 1])
def update_intraband(self, ind_m, kpt):
kd = self.calc.wfs.kd
gd = self.calc.wfs.gd
k_c = kd.bzk_kc[kpt.K] + kpt.shift_c
k_v = 2 * np.pi * np.dot(k_c, np.linalg.inv(gd.cell_cv).T)
atomdata_a = self.calc.wfs.setups
ut_mvR = self.calc.wfs.gd.zeros((len(ind_m), 3), complex)
for ind, ut_vR in zip(ind_m, ut_mvR):
ut_vR[:] = self.make_derivative(kpt.s, kpt.K, kpt.na + ind,
kpt.na + ind + 1)[0]
npartocc = len(ind_m)
ut_mR = kpt.ut_nR[ind_m]
nabla0_mmv = np.zeros((npartocc, npartocc, 3), dtype=complex)
for m in range(npartocc):
ut_vR = ut_mvR[m]
C_avi = [
np.dot(atomdata.nabla_iiv.T, P_mi[ind_m[m]])
for atomdata, P_mi in zip(atomdata_a, kpt.P_ani)
]
nabla0_mv = -self.calc.wfs.gd.integrate(ut_vR, ut_mR).T
nt_m = self.calc.wfs.gd.integrate(ut_mR[m], ut_mR)
nabla0_mv += 1j * nt_m[:, np.newaxis] * k_v[np.newaxis, :]
for C_vi, P_mi in zip(C_avi, kpt.P_ani):
gemm(1.0, C_vi, P_mi[ind_m[0:npartocc]], 1.0, nabla0_mv, 'c')
nabla0_mmv[m] = nabla0_mv
return nabla0_mmv
def iterate_one_k_point(self, hamiltonian, wfs, kpt, Vt_xMM):
if wfs.bd.comm.size > 1 and wfs.bd.strided:
raise NotImplementedError
H_MM = self.calculate_hamiltonian_matrix(hamiltonian, wfs, kpt, Vt_xMM,
root=0)
S_MM = wfs.S_qMM[kpt.q]
if kpt.eps_n is None:
kpt.eps_n = np.empty(wfs.bd.mynbands)
diagonalization_string = repr(self.diagonalizer)
wfs.timer.start(diagonalization_string)
self.diagonalizer.diagonalize(H_MM, kpt.C_nM, kpt.eps_n, S_MM)
wfs.timer.stop(diagonalization_string)
wfs.timer.start('Calculate projections')
# P_ani are not strictly necessary as required quantities can be
# evaluated directly using P_aMi. We should probably get rid
# of the places in the LCAO code using P_ani directly
for a, P_ni in kpt.P_ani.items():
# ATLAS can't handle uninitialized output array:
P_ni.fill(117)
gemm(1.0, kpt.P_aMi[a], kpt.C_nM, 0.0, P_ni, 'n')
wfs.timer.stop('Calculate projections')
def integrate(self, a_xG, c_axi=None, q=-1):
c_xI = np.zeros(a_xG.shape[:-1] + (self.nI,), self.pd.dtype)
b_xI = c_xI.reshape((np.prod(c_xI.shape[:-1]), self.nI))
a_xG = a_xG.reshape((-1, a_xG.shape[-1]))
alpha = 1.0 / self.pd.gd.N_c.prod()
if self.pd.dtype == float:
alpha *= 2
a_xG = a_xG.view(float)
if c_axi is None:
c_axi = self.dict(a_xG.shape[:-1])
x = 0.0
for G1, G2 in self.block(q):
f_IG = self.expand(q, G1, G2)
if self.pd.dtype == float:
if G1 == 0:
f_IG[:, 0] *= 0.5
f_IG = f_IG.view(float)
G1 *= 2
G2 *= 2
gemm(alpha, f_IG, a_xG[:, G1:G2], x, b_xI, 'c')
x = 1.0
for a, I1, I2 in self.indices:
c_axi[a][:] = self.eikR_qa[q][a] * c_xI[..., I1:I2]
return c_axi
def iterate_one_k_point(self, hamiltonian, wfs, kpt, Vt_xMM):
if wfs.bd.comm.size > 1 and wfs.bd.strided:
raise NotImplementedError
H_MM = self.calculate_hamiltonian_matrix(hamiltonian,
wfs,
kpt,
Vt_xMM,
root=0)
S_MM = wfs.S_qMM[kpt.q]
if kpt.eps_n is None:
kpt.eps_n = np.empty(wfs.bd.mynbands)
diagonalization_string = repr(self.diagonalizer)
wfs.timer.start(diagonalization_string)
self.diagonalizer.diagonalize(H_MM, kpt.C_nM, kpt.eps_n, S_MM)
wfs.timer.stop(diagonalization_string)
wfs.timer.start('Calculate projections')
# P_ani are not strictly necessary as required quantities can be
# evaluated directly using P_aMi. We should probably get rid
# of the places in the LCAO code using P_ani directly
for a, P_ni in kpt.P_ani.items():
# ATLAS can't handle uninitialized output array:
P_ni.fill(117)
gemm(1.0, kpt.P_aMi[a], kpt.C_nM, 0.0, P_ni, 'n')
wfs.timer.stop('Calculate projections')
def add(self, a_xG, c_axi=1.0, q=-1, f0_IG=None):
if isinstance(c_axi, float):
assert q == -1, a_xG.dims == 1
a_xG += (c_axi / self.pd.gd.dv) * self.expand(-1).sum(0)
return
c_xI = np.empty(a_xG.shape[:-1] + (self.nI,), self.pd.dtype)
for a, I1, I2 in self.indices:
c_xI[..., I1:I2] = c_axi[a] * self.eikR_qa[q][a].conj()
c_xI = c_xI.reshape((np.prod(c_xI.shape[:-1]), self.nI))
a_xG = a_xG.reshape((-1, a_xG.shape[-1])).view(self.pd.dtype)
for G1, G2 in self.block(q):
if f0_IG is None:
f_IG = self.expand(q, G1, G2)
else:
f_IG = f0_IG
if self.pd.dtype == float:
f_IG = f_IG.view(float)
G1 *= 2
G2 *= 2
gemm(1.0 / self.pd.gd.dv, f_IG, c_xI, 1.0, a_xG[:, G1:G2])
def add(self, a_xG, c_axi=1.0, q=-1, f0_IG=None):
if isinstance(c_axi, float):
assert q == -1, a_xG.dims == 1
a_xG += (c_axi / self.pd.gd.dv) * self.expand(-1).sum(0)
return
c_xI = np.empty(a_xG.shape[:-1] + (self.nI,), self.pd.dtype)
for a, I1, I2 in self.indices:
c_xI[..., I1:I2] = c_axi[a] * self.eikR_qa[q][a].conj()
c_xI = c_xI.reshape((np.prod(c_xI.shape[:-1], dtype=int), self.nI))
a_xG = a_xG.reshape((-1, a_xG.shape[-1])).view(self.pd.dtype)
for G1, G2 in self.block(q):
if f0_IG is None:
f_IG = self.expand(q, G1, G2)
else:
f_IG = f0_IG
if self.pd.dtype == float:
f_IG = f_IG.view(float)
G1 *= 2
G2 *= 2
gemm(1.0 / self.pd.gd.dv, f_IG, c_xI, 1.0, a_xG[:, G1:G2])
def sqrt_matrix(a, preserve=False):
"""Get the sqrt of a symmetric matrix a (diagonalize is used).
The matrix is kept if preserve=True, a=sqrt(a) otherwise."""
n = len(a)
if debug:
assert a.flags.contiguous
assert a.dtype == float
assert a.shape == (n, n)
if preserve:
b = a.copy()
else:
b = a
# diagonalize to get the form b = Z * D * Z^T
# where D is diagonal
D = np.empty((n, ))
diagonalize(b, D)
ZT = b.copy()
Z = np.transpose(b)
# c = Z * sqrt(D)
c = Z * np.sqrt(D)
# sqrt(b) = c * Z^T
gemm(1., ZT, c, 0., b)
return b
def calculate_density_matrix(self, f_n, C_nM, rho_MM=None):
# ATLAS can't handle uninitialized output array:
#rho_MM.fill(42)
self.timer.start('Calculate density matrix')
rho_MM = self.ksl.calculate_density_matrix(f_n, C_nM, rho_MM)
self.timer.stop('Calculate density matrix')
return rho_MM
# ----------------------------
if 1:
# XXX Should not conjugate, but call gemm(..., 'c')
# Although that requires knowing C_Mn and not C_nM.
# that also conforms better to the usual conventions in literature
Cf_Mn = C_nM.T.conj() * f_n
self.timer.start('gemm')
gemm(1.0, C_nM, Cf_Mn, 0.0, rho_MM, 'n')
self.timer.stop('gemm')
self.timer.start('band comm sum')
self.bd.comm.sum(rho_MM)
self.timer.stop('band comm sum')
else:
# Alternative suggestion. Might be faster. Someone should test this
from gpaw.utilities.blas import r2k
C_Mn = C_nM.T.copy()
r2k(0.5, C_Mn, f_n * C_Mn, 0.0, rho_MM)
tri2full(rho_MM)
def calculate_sic_matrixelements(self):
# overlap of pseudo wavefunctions
Htphit_mG = self.vt_mG * self.phit_mG
V_mm = np.zeros((self.nocc, self.nocc), dtype=self.dtype)
gemm(self.gd.dv, self.phit_mG, Htphit_mG, 0.0, V_mm, 't')
#
# PAW
for a, P_mi in self.P_ami.items():
for m, dH_p in enumerate(self.dH_amp[a]):
dH_ii = unpack(dH_p)
V_mm[m, :] += np.dot(P_mi[m], np.dot(dH_ii, P_mi.T))
# accumulate over grid-domains
self.gd.comm.sum(V_mm)
self.V_mm = V_mm
# Symmetrization of V and kappa-matrix:
K_mm = 0.5 * (V_mm - V_mm.T.conj())
V_mm = 0.5 * (V_mm + V_mm.T.conj())
# evaluate the kinetic correction
self.ekin = -np.trace(V_mm) * (3 - self.nspins)
return V_mm, K_mm, np.vdot(K_mm, K_mm).real
def calculate_density_matrix(self, f_n, C_nM, rho_MM=None):
# ATLAS can't handle uninitialized output array:
#rho_MM.fill(42)
self.timer.start('Calculate density matrix')
rho_MM = self.ksl.calculate_density_matrix(f_n, C_nM, rho_MM)
self.timer.stop('Calculate density matrix')
return rho_MM
# ----------------------------
if 1:
# XXX Should not conjugate, but call gemm(..., 'c')
# Although that requires knowing C_Mn and not C_nM.
# that also conforms better to the usual conventions in literature
Cf_Mn = C_nM.T.conj() * f_n
self.timer.start('gemm')
gemm(1.0, C_nM, Cf_Mn, 0.0, rho_MM, 'n')
self.timer.stop('gemm')
self.timer.start('band comm sum')
self.bd.comm.sum(rho_MM)
self.timer.stop('band comm sum')
else:
# Alternative suggestion. Might be faster. Someone should test this
from gpaw.utilities.blas import r2k
C_Mn = C_nM.T.copy()
r2k(0.5, C_Mn, f_n * C_Mn, 0.0, rho_MM)
tri2full(rho_MM)
def update_optical_limit(self, n, m, kpt1, kpt2, deps_m, df_m, n_mG):
if self.ut_sKnvR is None or kpt1.K not in self.ut_sKnvR[kpt1.s]:
self.ut_sKnvR = self.calculate_derivatives(kpt1)
kd = self.calc.wfs.kd
gd = self.calc.wfs.gd
k_c = kd.bzk_kc[kpt1.K] + kpt1.shift_c
k_v = 2 * np.pi * np.dot(k_c, np.linalg.inv(gd.cell_cv).T)
ut_vR = self.ut_sKnvR[kpt1.s][kpt1.K][n]
atomdata_a = self.calc.wfs.setups
C_avi = [np.dot(atomdata.nabla_iiv.T, P_ni[n])
for atomdata, P_ni in zip(atomdata_a, kpt1.P_ani)]
n0_mv = -self.calc.wfs.gd.integrate(ut_vR, kpt2.ut_nR[m]).T
nt_m = self.calc.wfs.gd.integrate(kpt1.ut_nR[n], kpt2.ut_nR[m])
n0_mv += 1j * nt_m[:, np.newaxis] * k_v[np.newaxis, :]
for C_vi, P_mi in zip(C_avi, kpt2.P_ani):
P_mi = P_mi[m].copy()
gemm(1.0, C_vi, P_mi, 1.0, n0_mv, 'c')
deps_m = deps_m.copy()
deps_m[deps_m > -1e-3] = np.inf
n0_mv *= 1j / deps_m[:, np.newaxis]
n_mG[:, 0] = n0_mv[:, 0]
return n0_mv
def update_hilbert(self, n_mG, deps_m, wd, chi0_wGG):
"""Update spectral function.
Updates spectral function A_wGG and saves it to chi0_wGG for
later hilbert-transform."""
self.timer.start('prep')
omega_w = wd.get_data()
deps_m += self.eshift * np.sign(deps_m)
o_m = abs(deps_m)
w_m = wd.get_closest_index(o_m)
o1_m = omega_w[w_m]
o2_m = omega_w[w_m + 1]
p_m = np.abs(1 / (o2_m - o1_m)**2)
p1_m = p_m * (o2_m - o_m)
p2_m = p_m * (o_m - o1_m)
self.timer.stop('prep')
if self.blockcomm.size > 1:
for p1, p2, n_G, w in zip(p1_m, p2_m, n_mG, w_m):
if w + 1 < wd.wmax: # The last frequency is not reliable
myn_G = n_G[self.Ga:self.Gb].reshape((-1, 1))
gemm(p1, n_G.reshape((-1, 1)), myn_G,
1.0, chi0_wGG[w], 'c')
gemm(p2, n_G.reshape((-1, 1)), myn_G,
1.0, chi0_wGG[w + 1], 'c')
return
for p1, p2, n_G, w in zip(p1_m, p2_m, n_mG, w_m):
if w + 1 < wd.wmax: # The last frequency is not reliable
czher(p1, n_G.conj(), chi0_wGG[w])
czher(p2, n_G.conj(), chi0_wGG[w + 1])
def calculate_hamiltonian_matrix(self, hamiltonian, wfs, kpt, root=-1):
# XXX document parallel stuff, particularly root parameter
assert self.has_initialized
vt_G = hamiltonian.vt_sG[kpt.s]
H_MM = np.empty((wfs.ksl.mynao, wfs.ksl.nao), wfs.dtype)
wfs.timer.start('Potential matrix')
wfs.basis_functions.calculate_potential_matrix(vt_G, H_MM, kpt.q)
wfs.timer.stop('Potential matrix')
# Add atomic contribution
#
# -- a a a*
# H += > P dH P
# mu nu -- mu i ij nu j
# aij
#
wfs.timer.start('Atomic Hamiltonian')
Mstart = wfs.basis_functions.Mstart
Mstop = wfs.basis_functions.Mstop
for a, P_Mi in kpt.P_aMi.items():
dH_ii = np.asarray(unpack(hamiltonian.dH_asp[a][kpt.s]), wfs.dtype)
dHP_iM = np.zeros((dH_ii.shape[1], P_Mi.shape[0]), wfs.dtype)
# (ATLAS can't handle uninitialized output array)
gemm(1.0, P_Mi, dH_ii, 0.0, dHP_iM, 'c')
gemm(1.0, dHP_iM, P_Mi[Mstart:Mstop], 1.0, H_MM)
wfs.timer.stop('Atomic Hamiltonian')
wfs.timer.start('Distribute overlap matrix')
H_MM = wfs.ksl.distribute_overlap_matrix(H_MM, root)
wfs.timer.stop('Distribute overlap matrix')
H_MM += wfs.T_qMM[kpt.q]
return H_MM
def integrate(self, a_xG, c_axi=None, q=-1):
c_xI = np.zeros(a_xG.shape[:-1] + (self.nI,), self.pd.dtype)
b_xI = c_xI.reshape((np.prod(c_xI.shape[:-1], dtype=int), self.nI))
a_xG = a_xG.reshape((-1, a_xG.shape[-1]))
alpha = 1.0 / self.pd.gd.N_c.prod()
if self.pd.dtype == float:
alpha *= 2
a_xG = a_xG.view(float)
if c_axi is None:
c_axi = self.dict(a_xG.shape[:-1])
x = 0.0
for G1, G2 in self.block(q):
f_IG = self.expand(q, G1, G2)
if self.pd.dtype == float:
if G1 == 0:
f_IG[:, 0] *= 0.5
f_IG = f_IG.view(float)
G1 *= 2
G2 *= 2
gemm(alpha, f_IG, a_xG[:, G1:G2], x, b_xI, 'c')
x = 1.0
for a, I1, I2 in self.indices:
c_axi[a][:] = self.eikR_qa[q][a] * c_xI[..., I1:I2]
return c_axi
def derivative(self, a_xG, c_axiv, q=-1):
c_vxI = np.zeros((3,) + a_xG.shape[:-1] + (self.nI,), self.pd.dtype)
b_vxI = c_vxI.reshape((3, np.prod(c_vxI.shape[1:-1]), self.nI))
a_xG = a_xG.reshape((-1, a_xG.shape[-1])).view(self.pd.dtype)
alpha = 1.0 / self.pd.gd.N_c.prod()
K_v = self.pd.K_qv[q]
x = 0.0
for G1, G2 in self.block(q):
f_IG = self.expand(q, G1, G2)
G_Gv = self.pd.G_Qv[self.pd.Q_qG[q][G1:G2]]
if self.pd.dtype == float:
for v in range(3):
gemm(2 * alpha,
(f_IG * 1.0j * G_Gv[:, v]).view(float),
a_xG[:, 2 * G1:2 * G2],
x, b_vxI[v], 'c')
else:
for v in range(3):
gemm(-alpha,
f_IG * (G_Gv[:, v] + K_v[v]),
a_xG[:, G1:G2],
x, b_vxI[v], 'c')
x = 1.0
for v in range(3):
if self.pd.dtype == float:
for a, I1, I2 in self.indices:
c_axiv[a][..., v] = c_vxI[v, ..., I1:I2]
else:
for a, I1, I2 in self.indices:
c_axiv[a][..., v] = (1.0j * self.eikR_qa[q][a] *
c_vxI[v, ..., I1:I2])
def _pseudo_braket(self, bra_xG, ket_yG, A_yx, square=None):
"""Calculate matrix elements of braket pairs of pseudo wave functions.
Low-level helper function. Results will be put in the *A_yx* array::
/ ~ * ~
A = | dG bra (G) ket (G)
nn' / n n'
Parameters:
bra_xG: ndarray
Set of bra-like vectors in which the matrix elements are evaluated.
key_yG: ndarray
Set of ket-like vectors in which the matrix elements are evaluated.
A_yx: ndarray
Matrix in which to put calculated elements. Take care: Due to the
difference in Fortran/C array order and the inherent BLAS nature,
the matrix has to be filled in transposed (conjugated in future?).
"""
assert bra_xG.shape[1:] == ket_yG.shape[1:]
assert (ket_yG.shape[0], bra_xG.shape[0]) == A_yx.shape
if square is None:
square = (bra_xG.shape[0]==ket_yG.shape[0])
dv = self.gd.dv
if ket_yG is bra_xG:
rk(dv, bra_xG, 0.0, A_yx)
elif self.hermitian and square:
r2k(0.5 * dv, bra_xG, ket_yG, 0.0, A_yx)
else:
gemm(dv, bra_xG, ket_yG, 0.0, A_yx, 'c')
def update_projectors(self):
self.timer.start('LCAO update projectors')
# Loop over all k-points
for k, kpt in enumerate(self.wfs.kpt_u):
for a, P_ni in kpt.P_ani.items():
P_ni.fill(117)
gemm(1.0, self.wfs.P_aqMi[a][kpt.q], kpt.C_nM, 0.0, P_ni, 'n')
self.timer.stop('LCAO update projectors')
def calculate_atomic_density_matrices_k_point(self, D_sii, kpt, a, f_n):
P_Mi = kpt.P_aMi[a]
rhoP_Mi = np.zeros_like(P_Mi)
D0_ii = np.zeros(D_sii[0].shape, self.dtype)
for rho_MM, D_ii in zip(kpt.rho_sMM, D_sii):
gemm(1.0, P_Mi, rho_MM, 0.0, rhoP_Mi)
gemm(1.0, rhoP_Mi, P_Mi.T.conj().copy(), 0.0, D0_ii)
D_ii += 0.5 * (D0_ii.real + D0_ii.T.real)
def calculate_atomic_density_matrices_k_point(self, D_sii, kpt, a, f_n):
P_Mi = kpt.P_aMi[a]
rhoP_Mi = np.zeros_like(P_Mi)
D0_ii = np.zeros(D_sii[0].shape, self.dtype)
for rho_MM, D_ii in zip(kpt.rho_sMM, D_sii):
gemm(1.0, P_Mi, rho_MM, 0.0, rhoP_Mi)
gemm(1.0, rhoP_Mi, P_Mi.T.conj().copy(), 0.0, D0_ii)
D_ii += 0.5 * (D0_ii.real + D0_ii.T.real)
def update(self, n_mG, deps_m, df_m, chi0_wGG):
sign = 1 - 2 * self.timeordered
for w, omega in enumerate(self.omega_w):
x_m = df_m * (1.0 / (omega + deps_m + 1j * self.eta) -
1.0 / (omega - deps_m + 1j * self.eta * sign))
nx_mG = n_mG * x_m[:, np.newaxis]
gemm(self.prefactor, n_mG.conj(), np.ascontiguousarray(nx_mG.T),
1.0, chi0_wGG[w])
def integrate(self, a_xg, b_yg=None,
global_integral=True, hermitian=False,
_transposed_result=None):
"""Integrate function(s) over domain.
a_xg: ndarray
Function(s) to be integrated.
b_yg: ndarray
If present, integrate a_xg.conj() * b_yg.
global_integral: bool
If the array(s) are distributed over several domains, then the
total sum will be returned. To get the local contribution
only, use global_integral=False.
hermitian: bool
Result is hermitian.
_transposed_result: ndarray
Long story. Don't use this unless you are a method of the
MatrixOperator class ..."""
xshape = a_xg.shape[:-3]
if b_yg is None:
# Only one array:
result = a_xg.reshape(xshape + (-1,)).sum(axis=-1) * self.dv
if global_integral:
if result.ndim == 0:
result = self.comm.sum(result)
else:
self.comm.sum(result)
return result
A_xg = np.ascontiguousarray(a_xg.reshape((-1,) + a_xg.shape[-3:]))
B_yg = np.ascontiguousarray(b_yg.reshape((-1,) + b_yg.shape[-3:]))
if _transposed_result is None:
result_yx = np.zeros((len(B_yg), len(A_xg)), A_xg.dtype)
else:
result_yx = _transposed_result
global_integral = False
if a_xg is b_yg:
rk(self.dv, A_xg, 0.0, result_yx)
elif hermitian:
r2k(0.5 * self.dv, A_xg, B_yg, 0.0, result_yx)
else:
gemm(self.dv, A_xg, B_yg, 0.0, result_yx, 'c')
if global_integral:
self.comm.sum(result_yx)
yshape = b_yg.shape[:-3]
result = result_yx.T.reshape(xshape + yshape)
if result.ndim == 0:
return result.item()
else:
return result
def update_projectors(self):
self.timer.start('LCAO update projectors')
# Loop over all k-points
for k, kpt in enumerate(self.wfs.kpt_u):
for a, P_ni in kpt.P_ani.items():
print('Update projector: Rank:', world.rank, 'a', a)
P_ni.fill(117)
gemm(1.0, kpt.P_aMi[a], kpt.C_nM, 0.0, P_ni, 'n')
self.timer.stop('LCAO update projectors')
def __call__(self, S_wx):
"""Inplace transform"""
B_wx = S_wx.reshape((len(S_wx), -1))
nw, nx = B_wx.shape
tmp_wx = np.zeros((nw, min(nx, self.blocksize)), complex)
for x in range(0, nx, self.blocksize):
b_wx = B_wx[:, x:x + self.blocksize]
c_wx = tmp_wx[:, :b_wx.shape[1]]
gemm(1.0, b_wx, self.H_ww, 0.0, c_wx)
b_wx[:] = c_wx
def add_paw_correction_to_overlap(setups, P_aqMi, S_qMM, Mstart=0, Mstop=None):
if Mstop is None:
Mstop = setups.nao
for a, P_qMi in P_aqMi.items():
dO_ii = np.asarray(setups[a].dO_ii, S_qMM.dtype)
for S_MM, P_Mi in zip(S_qMM, P_qMi):
dOP_iM = np.zeros((dO_ii.shape[1], setups.nao), P_Mi.dtype)
# (ATLAS can't handle uninitialized output array)
gemm(1.0, P_Mi, dO_ii, 0.0, dOP_iM, 'c')
gemm(1.0, dOP_iM, P_Mi[Mstart:Mstop], 1.0, S_MM, 'n')
def __call__(self, S_wx):
"""Inplace transform"""
B_wx = S_wx.reshape((len(S_wx), -1))
nw, nx = B_wx.shape
tmp_wx = np.zeros((nw, min(nx, self.blocksize)), complex)
for x in range(0, nx, self.blocksize):
b_wx = B_wx[:, x : x + self.blocksize]
c_wx = tmp_wx[:, : b_wx.shape[1]]
gemm(1.0, b_wx, self.H_ww, 0.0, c_wx)
b_wx[:] = c_wx
def calculate(self, wfs, kpt, dH_asp, H_MM, y):
Mstart = wfs.ksl.Mstart
Mstop = wfs.ksl.Mstop
dtype = wfs.dtype
for a, P_Mi in kpt.P_aMi.items():
dH_ii = np.asarray(unpack(dH_asp[a][kpt.s]), dtype)
dHP_iM = np.zeros((dH_ii.shape[1], P_Mi.shape[0]), dtype)
# (ATLAS can't handle uninitialized output array)
gemm(1.0, P_Mi, dH_ii, 0.0, dHP_iM, 'c')
gemm(y, dHP_iM, P_Mi[Mstart:Mstop], 1.0, H_MM)
def update_optimal_states(self):
# pseudo wavefunctions
self.phit_mG = self.gd.zeros(self.nocc)
if self.nocc > 0:
gemm(1.0, self.kpt.psit_nG[:self.nocc], self.W_mn, 0.0,
self.phit_mG)
# PAW
self.P_ami = {}
for a, P_ni in self.kpt.P_ani.items():
self.P_ami[a] = np.dot(self.W_mn, P_ni[:self.nocc])
def update_optimal_states(self):
#
# pseudo wavefunctions
self.phit_mG = self.gd.zeros(self.nocc)
if self.nocc > 0:
gemm(1.0, self.kpt.psit_nG[: self.nocc], self.W_mn, 0.0, self.phit_mG)
#
# PAW
self.P_ami = {}
for a, P_ni in self.kpt.P_ani.items():
self.P_ami[a] = np.dot(self.W_mn, P_ni[: self.nocc])
def calculate_density_matrix(self, f_n, C_nM, rho_MM=None):
# Only a madman would use a non-transposed density matrix.
# Maybe we should use the get_transposed_density_matrix instead
if rho_MM is None:
rho_MM = np.zeros((self.mynao, self.nao), dtype=C_nM.dtype)
# XXX Should not conjugate, but call gemm(..., 'c')
# Although that requires knowing C_Mn and not C_nM.
# that also conforms better to the usual conventions in literature
Cf_Mn = C_nM.T.conj() * f_n
gemm(1.0, C_nM, Cf_Mn, 0.0, rho_MM, 'n')
self.bd.comm.sum(rho_MM)
return rho_MM
def rotate_function(self, psit_oG, bfs, q=-1, indices=None):
if indices is None:
U_ow = self.U_ow
U_Mw = self.U_Mw
else:
U_ow = self.U_ow[:, indices]
U_Mw = self.U_Mw[:, indices]
w_wG = np.zeros((U_ow.shape[1], ) + psit_oG.shape[1:])
if len(U_ow) > 0:
gemm(1., psit_oG, U_ow.T.copy(), 0., w_wG)
bfs.lcao_to_grid(U_Mw.T.copy(), w_wG, q)
return w_wG
def rotate_function(self, psit_oG, bfs, q=-1, indices=None):
if indices is None:
U_ow = self.U_ow
U_Mw = self.U_Mw
else:
U_ow = self.U_ow[:, indices]
U_Mw = self.U_Mw[:, indices]
w_wG = np.zeros((U_ow.shape[1],) + psit_oG.shape[1:])
if len(U_ow) > 0:
gemm(1.0, psit_oG, U_ow.T.copy(), 0.0, w_wG)
bfs.lcao_to_grid(U_Mw.T.copy(), w_wG, q)
return w_wG
def update(self, n_mG, deps_m, df_m, chi0_wGG):
if self.timeordered:
deps1_m = deps_m + 1j * self.eta * np.sign(deps_m)
deps2_m = deps1_m
else:
deps1_m = deps_m + 1j * self.eta
deps2_m = deps_m - 1j * self.eta
for omega, chi0_GG in zip(self.omega_w, chi0_wGG):
x_m = df_m * (1 / (omega + deps1_m) - 1 / (omega - deps2_m))
nx_mG = n_mG * x_m[:, np.newaxis]
gemm(self.prefactor, n_mG.conj(), np.ascontiguousarray(nx_mG.T), 1.0, chi0_GG)
def add_paw_correction_to_overlap(setups, P_aqMi, S_qMM, Mstart=0,
Mstop=None):
if Mstop is None:
Mstop = setups.nao
for a, P_qMi in P_aqMi.items():
dO_ii = np.asarray(setups[a].dO_ii, S_qMM.dtype)
for S_MM, P_Mi in zip(S_qMM, P_qMi):
dOP_iM = np.zeros((dO_ii.shape[1], setups.nao),
P_Mi.dtype)
# (ATLAS can't handle uninitialized output array)
gemm(1.0, P_Mi, dO_ii, 0.0, dOP_iM, 'c')
gemm(1.0, dOP_iM, P_Mi[Mstart:Mstop],
1.0, S_MM, 'n')
def update(self, n_mG, deps_m, df_m, chi0_wGG):
if self.timeordered:
deps1_m = deps_m + 1j * self.eta * np.sign(deps_m)
deps2_m = deps1_m
else:
deps1_m = deps_m + 1j * self.eta
deps2_m = deps_m - 1j * self.eta
for omega, chi0_GG in zip(self.omega_w, chi0_wGG):
x_m = df_m * (1 / (omega + deps1_m) - 1 / (omega - deps2_m))
nx_mG = n_mG * x_m[:, np.newaxis]
gemm(self.prefactor, n_mG.conj(), np.ascontiguousarray(nx_mG.T),
1.0, chi0_GG)
def calculate(self, wfs, q, dX_aii, X_MM):
dtype = X_MM.dtype
nops = 0
for a, dX_ii in dX_aii.items():
P_Mi = self.P_aqMi[a][q]
assert dtype == P_Mi.dtype
dXP_iM = np.zeros((dX_ii.shape[1], P_Mi.shape[0]), dtype)
# (ATLAS can't handle uninitialized output array)
gemm(1.0, P_Mi, dX_ii, 0.0, dXP_iM, 'c')
nops += dXP_iM.size * dX_ii.shape[0]
gemm(1.0, dXP_iM, P_Mi[self.Mstart:self.Mstop], 1.0, X_MM)
nops += X_MM.size * dXP_iM.shape[0]
self.nops = nops
def calculate_hamiltonian_matrix(self, hamiltonian, wfs, kpt, Vt_xMM=None,
root=-1):
# XXX document parallel stuff, particularly root parameter
assert self.has_initialized
bf = wfs.basis_functions
if Vt_xMM is None:
wfs.timer.start('Potential matrix')
vt_G = hamiltonian.vt_sG[kpt.s]
Vt_xMM = bf.calculate_potential_matrices(vt_G)
wfs.timer.stop('Potential matrix')
if bf.gamma:
y = 1.0
H_MM = Vt_xMM[0]
if wfs.dtype == complex:
H_MM = H_MM.astype(complex)
else:
wfs.timer.start('Sum over cells')
y = 0.5
k_c = wfs.kd.ibzk_qc[kpt.q]
H_MM = (0.5 + 0.0j) * Vt_xMM[0]
for sdisp_c, Vt_MM in zip(bf.sdisp_xc, Vt_xMM)[1:]:
H_MM += np.exp(2j * np.pi * np.dot(sdisp_c, k_c)) * Vt_MM
wfs.timer.stop('Sum over cells')
# Add atomic contribution
#
# -- a a a*
# H += > P dH P
# mu nu -- mu i ij nu j
# aij
#
wfs.timer.start('Atomic Hamiltonian')
Mstart = wfs.basis_functions.Mstart
Mstop = wfs.basis_functions.Mstop
for a, P_Mi in kpt.P_aMi.items():
dH_ii = np.asarray(unpack(hamiltonian.dH_asp[a][kpt.s]), wfs.dtype)
dHP_iM = np.zeros((dH_ii.shape[1], P_Mi.shape[0]), wfs.dtype)
# (ATLAS can't handle uninitialized output array)
gemm(1.0, P_Mi, dH_ii, 0.0, dHP_iM, 'c')
gemm(y, dHP_iM, P_Mi[Mstart:Mstop], 1.0, H_MM)
wfs.timer.stop('Atomic Hamiltonian')
wfs.timer.start('Distribute overlap matrix')
H_MM = wfs.ksl.distribute_overlap_matrix(
H_MM, root, add_hermitian_conjugate=(y == 0.5))
wfs.timer.stop('Distribute overlap matrix')
H_MM += wfs.T_qMM[kpt.q]
return H_MM
def make_optimized(qstart, qend):
g_qG = np.zeros((qend - qstart,) + w_wG.shape[1:], float)
P_aqp = {}
for a, P_wi in P_awi.items():
ni = P_wi.shape[1]
nii = ni * (ni + 1) // 2
P_aqp[a] = np.zeros((qend - qstart, nii), float)
for w1, w1_G in enumerate(w_wG):
U = Uisq_iqj[w1, qstart: qend].copy()
gemm(1., w1_G * w_wG, U, 1.0, g_qG)
for a, P_wi in P_awi.items():
P_wp = np.array([pack(np.outer(P_wi[w1], P_wi[w2]))
for w2 in range(Ni)])
gemm(1., P_wp, U, 1.0, P_aqp[a])
return g_qG, P_aqp
def matrix_multiply(C_nn, psit_mG, send_mG, recv_mG):
"""Calculate new linear compination of wave functions."""
rank = band_comm.rank
C_imim = C_nn.reshape((B, M, B, M))
send_mG[:] = psit_mG
psit_mG[:] = 0.0
beta = 0.0
for i in range(B - 1):
gemm(1.0, send_mG, C_imim[rank, :, (rank + i) % B], beta, psit_mG)
beta = 1.0
band_comm.sendreceive(send_mG, (rank - 1) % B, recv_mG, (rank + 1) % B, 117, 117)
send_mG, recv_mG = recv_mG, send_mG
gemm(1.0, send_mG, C_imim[rank, :, rank - 1], beta, psit_mG)
return psit_mG
def make_optimized(qstart, qend):
g_qG = np.zeros((qend - qstart,) + w_wG.shape[1:], float)
P_aqp = {}
for a, P_wi in P_awi.items():
ni = P_wi.shape[1]
nii = ni * (ni + 1) // 2
P_aqp[a] = np.zeros((qend - qstart, nii), float)
for w1, w1_G in enumerate(w_wG):
U = Uisq_iqj[w1, qstart: qend].copy()
gemm(1., w1_G * w_wG, U, 1.0, g_qG)
for a, P_wi in P_awi.items():
P_wp = np.array([pack(np.outer(P_wi[w1], P_wi[w2]))
for w2 in range(Ni)])
gemm(1., P_wp, U, 1.0, P_aqp[a])
return g_qG, P_aqp
def get_overlap_energydiff(self, k, spin, mynks):
"""Calculate overlap matrices and energy difference matrix
for requested k-point and spin with index mynks
k global k-point index
spin global spin index
mynks local k-point and spin index
overlap_qvnm overlap matrix of direction and bands
de_nm energy difference matrix of bands"""
nbands = self.calc.get_number_of_bands() # total number of bands
identity_nm = identity(nbands, dtype=self.calc.wfs.dtype) # identity
eigenvalues_n = self.calc.wfs.kpt_u[mynks].eps_n.copy() # eigenvalues
occupations_n = self.calc.wfs.kpt_u[mynks].f_n.copy() # occupations
coeff_nnu = self.calc.wfs.kpt_u[mynks].C_nM.copy() # LCAO coefficients
#Include the GLLBSC scissors correction, if implemented
eigenvalues_n[self.nocc:] += self.scissor
if self.singlet:
spin2 = self.swap_unoccupied(eigenvalues_n, occupations_n,
coeff_nnu, mynks)
else:
spin2 = spin
# Calculating PAW Corrections
de_nm = (outer(eigenvalues_n, ones(nbands)) -
outer(ones(nbands), eigenvalues_n))
df_nm = self.get_df_nm(occupations_n,
weight=self.calc.wfs.kd.weight_k[k],
number_of_spins=self.calc.get_number_of_spins())
gradcoeff_num = zeros(coeff_nnu.shape, dtype=self.calc.wfs.dtype)
# Overlap Matrix initialized with PAW overlaps
overlap_qvnm = self.get_paw_overlap(spin1=spin, spin2=spin2, k=k)
# Calculating Overlap Matrix
# Perform Matrix Multiplications for each direction
nkpts = self.grad_phi_kqvnumu.shape[0]
for qdir in range(3):
# Be careful that mynks points to the k-point index
gemm(1.0, self.grad_phi_kqvnumu[mynks % nkpts, qdir], coeff_nnu,
0.0, gradcoeff_num)
gemm(1.0, coeff_nnu, gradcoeff_num, 1.0, overlap_qvnm[qdir], 'c')
overlap_qvnm[qdir] /= de_nm + identity_nm
# Sum overlap_qvnm over domains
self.calc.comms['K'].sum(overlap_qvnm)
overlap_qvnm = self.prefactor * df_nm * (overlap_qvnm *
overlap_qvnm.conj())
return overlap_qvnm.real, de_nm
def calculate_atomic_density_matrices_k_point(self, D_sii, kpt, a, f_n):
if kpt.rho_MM is not None:
P_Mi = self.P_aqMi[a][kpt.q]
rhoP_Mi = np.zeros_like(P_Mi)
D_ii = np.zeros(D_sii[kpt.s].shape, kpt.rho_MM.dtype)
gemm(1.0, P_Mi, kpt.rho_MM, 0.0, rhoP_Mi)
gemm(1.0, rhoP_Mi, P_Mi.T.conj().copy(), 0.0, D_ii)
D_sii[kpt.s] += D_ii.real
else:
P_ni = kpt.P_ani[a]
D_sii[kpt.s] += np.dot(P_ni.T.conj() * f_n, P_ni).real
if hasattr(kpt, 'c_on'):
for ne, c_n in zip(kpt.ne_o, kpt.c_on):
d_nn = ne * np.outer(c_n.conj(), c_n)
D_sii[kpt.s] += np.dot(P_ni.T.conj(), np.dot(d_nn, P_ni)).real
def matrix_multiply(C_nn, psit_mG, send_mG, recv_mG):
"""Calculate new linear compination of wave functions."""
rank = band_comm.rank
C_imim = C_nn.reshape((B, M, B, M))
send_mG[:] = psit_mG
psit_mG[:] = 0.0
beta = 0.0
for i in range(B - 1):
gemm(1.0, send_mG, C_imim[rank, :, (rank + i) % B], beta, psit_mG)
beta = 1.0
band_comm.sendreceive(send_mG, (rank - 1) % B, recv_mG, (rank + 1) % B,
117, 117)
send_mG, recv_mG = recv_mG, send_mG
gemm(1.0, send_mG, C_imim[rank, :, rank - 1], beta, psit_mG)
return psit_mG
def calculate_residual_change(self, psit_xG, Htpsit_xG, P_axi, c_axi, n_x):
"""
"""
assert len(n_x) == 1
Htphit_mG = self.Htphit_mG
phit_mG = self.phit_mG
w_mx = np.zeros((self.nocc, 1), dtype=self.dtype)
v_mx = np.zeros((self.nocc, 1), dtype=self.dtype)
gemm(self.gd.dv, psit_xG, phit_mG, 0.0, w_mx, 't')
gemm(self.gd.dv, psit_xG, Htphit_mG, 0.0, v_mx, 't')
# PAW
for a, P_mi in self.P_ami.items():
P_xi = P_axi[a]
dO_ii = self.setups[a].dO_ii
#
w_mx += np.dot(P_mi, np.dot(dO_ii, P_xi.T))
for m, dH_p in enumerate(self.dH_amp[a]):
dH_ii = unpack(dH_p)
v_mx[m] += np.dot(P_mi[m], np.dot(dH_ii, P_xi.T))
# sum over grid-domains
self.finegd.comm.sum(w_mx)
self.finegd.comm.sum(v_mx)
V_mm = 0.5 * (self.V_mm + self.V_mm.T)
q_mx = v_mx - np.dot(V_mm, w_mx)
if self.stabpot != 0.0:
q_mx -= self.stabpot * w_mx
gemm(1.0, Htphit_mG, w_mx.T.copy(), 1.0, Htpsit_xG)
gemm(1.0, phit_mG, q_mx.T.copy(), 1.0, Htpsit_xG)
# PAW
for a, P_mi in self.P_ami.items():
c_xi = c_axi[a]
ct_mi = P_mi.copy()
dO_ii = self.setups[a].dO_ii
c_xi += np.dot(q_mx.T, np.dot(P_mi, dO_ii))
for m, dH_p in enumerate(self.dH_amp[a]):
dH_ii = unpack(dH_p)
ct_mi[m] = np.dot(P_mi[m], dH_ii)
c_xi += np.dot(w_mx.T, ct_mi)
if self.stabpot != 0.0:
Htphit_mG += self.stabpot * psit_xG
for a, P_xi in P_axi.items():
dO_ii = self.setups[a].dO_ii
c_axi[a] += self.stabpot * np.dot(P_xi, dO_ii)
def overlap(psit_mG, send_mG, recv_mG):
"""Calculate overlap matrix.
The master (rank=0) will return the matrix with only the lower
part filled in."""
Q = B // 2 + 1
rank = band_comm.rank
S_imm = np.empty((Q, M, M))
send_mG[:] = psit_mG
# Shift wave functions:
for i in range(Q - 1):
gemm(gd.dv, psit_mG, send_mG, 0.0, S_imm[i], 'c')
if (band_comm.rank % 2 == 0):
band_comm.send(send_mG, (rank - 1) % B, 42)
band_comm.receive(recv_mG, (rank + 1) % B, 42)
else:
band_comm.receive(recv_mG, (rank + 1) % B, 42)
band_comm.send(send_mG, (rank - 1) % B, 42)
send_mG, recv_mG = recv_mG, send_mG
gemm(gd.dv, psit_mG, send_mG, 0.0, S_imm[Q - 1], 'c')
domain_comm.sum(S_imm)
t1 = time()
if domain_comm.rank == 0:
if band_comm.rank == 0:
# Master collects submatrices:
S_nn = np.empty((N,N))
S_nn[:] = 11111.77777
S_imim = S_nn.reshape((B, M, B, M))
S_imim[:Q, :, 0] = S_imm
for i1 in range(1, B):
band_comm.receive(S_imm, i1, i1)
for i2 in range(Q):
if i1 + i2 < B:
S_imim[i1 + i2, :, i1] = S_imm[i2]
else:
S_imim[i1, :, i1 + i2 - B] = S_imm[i2].T
t2 = time()
print 'Collect submatrices time %f' % (t2-t1)
return S_nn
else:
# Slaves send their submatrices:
band_comm.send(S_imm, 0, band_comm.rank)
def calculate_residual(self, psit_nG, Htpsit_nG, P_ani, c_ani):
""" Calculate the action of the unified Hamiltonian on an
arbitrary state:
H_u|Psi> =
"""
nocc = self.nocc
nvirt = psit_nG.shape[0] - nocc
# constraint for unoccupied states
R_mk = np.zeros((nocc, nvirt), dtype=self.dtype)
if nvirt > 0:
gemm(self.gd.dv, psit_nG[nocc:], self.Htphit_mG, 0.0, R_mk, 't')
# PAW
for a, P_mi in self.P_ami.items():
P_ni = P_ani[a]
for m, dH_p in enumerate(self.dH_amp[a]):
dH_ii = unpack(dH_p)
R_mk[m] += np.dot(P_mi[m], np.dot(dH_ii, P_ni[nocc:].T))
self.finegd.comm.sum(R_mk)
# self.R_mk = R_mk
# R_mk = self.R_mk
W_mn = self.W_mn
Htphit_mG = self.Htphit_mG
phit_mG = self.phit_mG
K_mm = 0.5 * (self.V_mm - self.V_mm.T)
Q_mn = np.dot(K_mm, W_mn)
# Action of unified Hamiltonian on occupied states:
if nocc > 0:
gemm(1.0, Htphit_mG, W_mn.T.copy(), 1.0, Htpsit_nG[:nocc])
gemm(1.0, phit_mG, Q_mn.T.copy(), 1.0, Htpsit_nG[:nocc])
if nvirt > 0:
gemm(1.0, phit_mG, R_mk.T.copy(), 1.0, Htpsit_nG[nocc:])
if self.stabpot != 0.0:
Htpsit_nG[nocc:] += self.stabpot * psit_nG[nocc:]
# PAW
for a, P_mi in self.P_ami.items():
#
c_ni = c_ani[a]
ct_mi = P_mi.copy()
#
dO_ii = self.setups[a].dO_ii
c_ni[:nocc] += np.dot(Q_mn.T, np.dot(P_mi, dO_ii))
c_ni[nocc:] += np.dot(R_mk.T, np.dot(P_mi, dO_ii))
#
for m, dH_p in enumerate(self.dH_amp[a]):
dH_ii = unpack(dH_p)
ct_mi[m] = np.dot(P_mi[m], dH_ii)
c_ni[:nocc] += np.dot(W_mn.T, ct_mi)
c_ni[nocc:] += self.stabpot * np.dot(P_ani[a][nocc:], dO_ii)
def dot(a, b, transa='n'):
assert len(a.shape) == 2 and a.shape[1] == b.shape[0]
dtype = complex
if a.dtype == complex and b.dtype == complex:
c = a
d = b
elif a.dtype == float and b.dtype == complex:
c = np.array(a, complex)
d = b
elif a.dtype == complex and b.dtype == float:
d = np.array(b, complex)
c = a
else:
dtype = float
c = a
d = b
e = np.zeros([c.shape[0], d.shape[1]], dtype)
gemm(1.0, np.ascontiguousarray(d), np.ascontiguousarray(c), 0.0, e, transa)
return e
def get_component(self, wfs, s, vt_sG, dH_asp, kpt, H_MM):
wfs.basis_functions.calculate_potential_matrix(vt_sG[s], H_MM, kpt.q)
# Add atomic contribution
#
# -- a a a*
# H += > P dH P
# mu nu -- mu i ij nu j
# aij
#
wfs.timer.start('Atomic Hamiltonian')
Mstart = wfs.basis_functions.Mstart
Mstop = wfs.basis_functions.Mstop
wfs.timer.stop('Atomic Hamiltonian')
tri2full(H_MM)
for a, P_Mi in kpt.P_aMi.items():
dH_ii = np.asarray(unpack(dH_asp[a][s]), wfs.dtype)
dHP_iM = np.zeros((dH_ii.shape[1], P_Mi.shape[0]), wfs.dtype)
# (ATLAS can't handle uninitialized output array)
gemm(1.0, P_Mi, dH_ii, 0.0, dHP_iM, 'c')
gemm(1.0, dHP_iM, P_Mi[Mstart:Mstop], 1.0, H_MM)
def calculate_atomic_density_matrices_k_point(self, D_sii, kpt, a, f_n):
if kpt.rho_MM is not None:
P_Mi = kpt.P_aMi[a]
#P_Mi = kpt.P_aMi_sparse[a]
#ind = get_matrix_index(kpt.P_aMi_index[a])
#D_sii[kpt.s] += np.dot(np.dot(P_Mi.T.conj(), kpt.rho_MM),
# P_Mi).real
rhoP_Mi = np.zeros_like(P_Mi)
D_ii = np.zeros(D_sii[kpt.s].shape, kpt.rho_MM.dtype)
#gemm(1.0, P_Mi, kpt.rho_MM[ind.T, ind], 0.0, tmp)
gemm(1.0, P_Mi, kpt.rho_MM, 0.0, rhoP_Mi)
gemm(1.0, rhoP_Mi, P_Mi.T.conj().copy(), 0.0, D_ii)
D_sii[kpt.s] += D_ii.real
#D_sii[kpt.s] += dot(dot(P_Mi.T.conj().copy(),
# kpt.rho_MM[ind.T, ind]), P_Mi).real
else:
P_ni = kpt.P_ani[a]
D_sii[kpt.s] += np.dot(P_ni.T.conj() * f_n, P_ni).real
if hasattr(kpt, 'c_on'):
for ne, c_n in zip(kpt.ne_o, kpt.c_on):
d_nn = ne * np.outer(c_n.conj(), c_n)
D_sii[kpt.s] += np.dot(P_ni.T.conj(), np.dot(d_nn, P_ni)).real
def update_hilbert(self, n_mG, deps_m, df_m, chi0_wGG):
self.timer.start("prep")
beta = (2 ** 0.5 - 1) * self.domega0 / self.omega2
o_m = abs(deps_m)
w_m = (o_m / (self.domega0 + beta * o_m)).astype(int)
o1_m = self.omega_w[w_m]
o2_m = self.omega_w[w_m + 1]
p_m = self.prefactor * abs(df_m) / (o2_m - o1_m) ** 2 # XXX abs()?
p1_m = p_m * (o2_m - o_m)
p2_m = p_m * (o_m - o1_m)
self.timer.stop("prep")
if self.blockcomm.size > 1:
for p1, p2, n_G, w in zip(p1_m, p2_m, n_mG, w_m):
myn_G = n_G[self.Ga : self.Gb].reshape((-1, 1))
gemm(p1, n_G.reshape((-1, 1)), myn_G, 1.0, chi0_wGG[w], "c")
gemm(p2, n_G.reshape((-1, 1)), myn_G, 1.0, chi0_wGG[w + 1], "c")
# chi0_wGG[w + 1] += p2 * np.outer(myn_G, n_G.conj())
return
for p1, p2, n_G, w in zip(p1_m, p2_m, n_mG, w_m):
czher(p1, n_G.conj(), chi0_wGG[w])
czher(p2, n_G.conj(), chi0_wGG[w + 1])
