import numpy as np
from gpaw.utilities import pack_density
def get_density(rho_MM, wfs, density, density_type='comp', u=0):
if wfs.ksl.using_blacs:
raise NotImplementedError('Scalapack is not supported')
rho_G = density.gd.zeros()
kpt = wfs.kpt_u[u]
assert kpt.q == 0
wfs.basis_functions.construct_density(rho_MM.astype(wfs.dtype),
rho_G, kpt.q)
# Uncomment this if you want to add the static part
# rho_G += density.nct_G
if density_type == 'pseudocoarse':
return rho_G
rho_g = density.finegd.zeros()
density.distribute_and_interpolate(rho_G, rho_g)
rho_G = None
if density_type == 'pseudo':
return rho_g
if density_type == 'comp':
D_asp = density.atom_partition.arraydict(density.D_asp.shapes_a)
Q_aL = {}
for a, D_sp in D_asp.items():
P_Mi = wfs.P_aqMi[a][kpt.q]
assert np.max(np.absolute(P_Mi.imag)) == 0
P_Mi = P_Mi.real
assert P_Mi.dtype == float
D_ii = np.dot(np.dot(P_Mi.T.conj(), rho_MM), P_Mi)
D_sp[:] = pack_density(D_ii)[np.newaxis, :]
Q_aL[a] = np.dot(D_sp.sum(axis=0), wfs.setups[a].Delta_pL)
density.ghat.add(rho_g, Q_aL)
return rho_g
raise RuntimeError('Unknown density type: %s' % density_type)
[docs]class DensityMatrix:
def __init__(self, paw):
self.wfs = paw.wfs
self.density = paw.density
self.using_blacs = self.wfs.ksl.using_blacs
self.tag = None
def zeros(self, dtype):
ksl = self.wfs.ksl
if self.using_blacs:
return ksl.mmdescriptor.zeros(dtype=dtype)
else:
return np.zeros((ksl.mynao, ksl.nao), dtype=dtype)
def _calculate_density_matrix(self, wfs, kpt):
if self.using_blacs:
ksl = wfs.ksl
rho_MM = ksl.calculate_blocked_density_matrix(kpt.f_n, kpt.C_nM)
else:
rho_MM = wfs.calculate_density_matrix(kpt.f_n, kpt.C_nM)
wfs.bd.comm.sum(rho_MM, root=0)
# TODO: should the sum over bands be moved to
# OrbitalLayouts.calculate_density_matrix()
return rho_MM
def get_density_matrix(self, tag=None):
if tag is None or self.tag != tag:
self.rho_uMM = []
for kpt in self.wfs.kpt_u:
rho_MM = self._calculate_density_matrix(self.wfs, kpt)
self.rho_uMM.append(rho_MM)
self.tag = tag
return self.rho_uMM
def get_density(self, rho_uMM=None, density_type='comp'):
assert len(self.wfs.kpt_u) == 1, 'K-points not implemented'
u = 0
if rho_uMM is None:
rho_uMM = self.rho_uMM
return get_density(rho_uMM[u], self.wfs, self.density, density_type, u)