def __init__(self,
ksl,
gd,
nvalence,
setups,
bd,
dtype,
world,
kd,
kptband_comm,
timer,
atomic_hamiltonian=None):
WaveFunctions.__init__(self, gd, nvalence, setups, bd, dtype, world,
kd, kptband_comm, timer)
self.ksl = ksl
self.S_qMM = None
self.T_qMM = None
self.P_aqMi = None
if atomic_hamiltonian is None:
if ksl.using_blacs:
atomic_hamiltonian = 'distributed'
else:
atomic_hamiltonian = 'dense'
if isinstance(atomic_hamiltonian, str):
atomic_hamiltonian = get_atomic_hamiltonian(atomic_hamiltonian)
self.atomic_hamiltonian = atomic_hamiltonian
self.timer.start('TCI: Evaluate splines')
self.tci = NewTCI(gd.cell_cv, gd.pbc_c, setups, kd.ibzk_qc, kd.gamma)
self.timer.stop('TCI: Evaluate splines')
self.basis_functions = BasisFunctions(
gd, [setup.phit_j for setup in setups], kd, dtype=dtype, cut=True)
def __init__(self, ksl, gd, nvalence, setups, bd,
dtype, world, kd, kptband_comm, timer,
atomic_correction=None, collinear=True):
WaveFunctions.__init__(self, gd, nvalence, setups, bd,
dtype, collinear, world, kd,
kptband_comm, timer)
self.ksl = ksl
self.S_qMM = None
self.T_qMM = None
self.P_aqMi = None
self.debug_tci = False
if atomic_correction is None:
if ksl.using_blacs:
atomic_correction = 'scipy'
else:
atomic_correction = 'dense'
if isinstance(atomic_correction, str):
atomic_correction = get_atomic_correction(atomic_correction)
self.atomic_correction = atomic_correction
#self.tci = NewTCI(gd.cell_cv, gd.pbc_c, setups, kd.ibzk_qc, kd.gamma)
with self.timer('TCI: Evaluate splines'):
self.tciexpansions = TCIExpansions.new_from_setups(setups)
self.basis_functions = BasisFunctions(gd,
[setup.phit_j
for setup in setups],
kd,
dtype=dtype,
cut=True)
def get_bfs(calc):
wfs = calc.wfs
bfs = BasisFunctions(wfs.gd, [setup.phit_j for setup in wfs.setups],
wfs.kd,
cut=True)
bfs.set_positions(calc.atoms.get_scaled_positions() % 1.)
return bfs
def get_bfs(calc):
wfs = calc.wfs
bfs = BasisFunctions(wfs.gd, [setup.phit_j for setup in wfs.setups],
wfs.kd,
cut=True)
bfs.set_positions(wfs.spos_ac)
return bfs
def initialize(self, density, hamiltonian, spos_ac):
if self.kpt_u[0].psit_nG is None:
basis_functions = BasisFunctions(
self.gd, [setup.phit_j for setup in self.setups], cut=True)
if not self.gamma:
basis_functions.set_k_points(self.kd.ibzk_qc)
basis_functions.set_positions(spos_ac)
elif isinstance(self.kpt_u[0].psit_nG, TarFileReference):
self.initialize_wave_functions_from_restart_file()
if self.kpt_u[0].psit_nG is not None:
density.initialize_from_wavefunctions(self)
elif density.nt_sG is None:
density.initialize_from_atomic_densities(basis_functions)
# Initialize GLLB-potential from basis function orbitals
if hamiltonian.xc.type == 'GLLB':
hamiltonian.xc.initialize_from_atomic_orbitals(basis_functions)
else: # XXX???
# We didn't even touch density, but some combinations in paw.set()
# will make it necessary to do this for some reason.
density.calculate_normalized_charges_and_mix()
hamiltonian.update(density)
if self.kpt_u[0].psit_nG is None:
self.initialize_wave_functions_from_basis_functions(
basis_functions, density, hamiltonian, spos_ac)
def get_bfs(calc):
wfs = calc.wfs
bfs = BasisFunctions(wfs.gd, [setup.phit_j for setup in wfs.setups],
wfs.kpt_comm, cut=True)
if not wfs.gamma:
bfs.set_k_points(wfs.ibzk_qc)
bfs.set_positions(calc.atoms.get_scaled_positions() % 1.)
return bfs
def get_density(self, atom_indices=None, gridrefinement=2):
"""Get sum of atomic densities from the given atom list.
All atoms are taken if the list is not given."""
all_atoms = self.calculator.get_atoms()
if atom_indices is None:
atom_indices = range(len(all_atoms))
density = self.calculator.density
spos_ac = all_atoms.get_scaled_positions()
rank_a = self.finegd.get_ranks_from_positions(spos_ac)
density.set_positions(all_atoms.get_scaled_positions(),
AtomPartition(self.finegd.comm, rank_a))
# select atoms
atoms = []
D_asp = {}
rank_a = []
all_D_asp = self.calculator.density.D_asp
all_rank_a = self.calculator.density.atom_partition.rank_a
for a in atom_indices:
if a in all_D_asp:
D_asp[len(atoms)] = all_D_asp.get(a)
atoms.append(all_atoms[a])
rank_a.append(all_rank_a[a])
atoms = Atoms(atoms,
cell=all_atoms.get_cell(),
pbc=all_atoms.get_pbc())
spos_ac = atoms.get_scaled_positions()
Z_a = atoms.get_atomic_numbers()
par = self.calculator.parameters
setups = Setups(Z_a, par.setups, par.basis, XC(par.xc),
self.calculator.wfs.world)
# initialize
self.initialize(setups, self.calculator.timer, np.zeros(len(atoms)),
False)
self.set_mixer(None)
# FIXME nparray causes partitionong.py test to fail
self.set_positions(spos_ac, AtomPartition(self.gd.comm, rank_a))
self.D_asp = D_asp
basis_functions = BasisFunctions(
self.gd, [setup.phit_j for setup in self.setups], cut=True)
basis_functions.set_positions(spos_ac)
self.initialize_from_atomic_densities(basis_functions)
aed_sg, gd = self.get_all_electron_density(atoms, gridrefinement)
return aed_sg.sum(axis=0), gd
def get_orbitals(calc):
"""Get LCAO orbitals on 3D grid by lcao_to_grid method."""
bfs_a = [setup.phit_j for setup in calc.wfs.setups]
from gpaw.lfc import BasisFunctions
bfs = BasisFunctions(calc.wfs.gd, bfs_a, calc.wfs.kd.comm, cut=True)
bfs.set_positions(calc.spos_ac)
nLCAO = calc.get_number_of_bands()
orb_MG = calc.wfs.gd.zeros(nLCAO)
C_M = np.identity(nLCAO)
bfs.lcao_to_grid(C_M, orb_MG, q=-1)
return orb_MG
def get_density(self, atom_indicees=None):
"""Get sum of atomic densities from the given atom list.
All atoms are taken if the list is not given."""
all_atoms = self.calculator.get_atoms()
if atom_indicees is None:
atom_indicees = range(len(all_atoms))
density = self.calculator.density
density.set_positions(all_atoms.get_scaled_positions() % 1.0,
self.calculator.wfs.rank_a)
# select atoms
atoms = []
D_asp = {}
rank_a = []
all_D_asp = self.calculator.density.D_asp
all_rank_a = self.calculator.density.rank_a
for a in atom_indicees:
if a in all_D_asp:
D_asp[len(atoms)] = all_D_asp.get(a)
atoms.append(all_atoms[a])
rank_a.append(all_rank_a[a])
atoms = Atoms(atoms, cell=all_atoms.get_cell())
spos_ac = atoms.get_scaled_positions() % 1.0
Z_a = atoms.get_atomic_numbers()
par = self.calculator.input_parameters
setups = Setups(Z_a, par.setups, par.basis, par.lmax, XC(par.xc),
self.calculator.wfs.world)
self.D_asp = D_asp
# initialize
self.initialize(setups, par.stencils[1], self.calculator.timer,
[0] * len(atoms), False)
self.set_mixer(None)
self.set_positions(spos_ac, rank_a)
basis_functions = BasisFunctions(
self.gd, [setup.phit_j for setup in self.setups], cut=True)
basis_functions.set_positions(spos_ac)
self.initialize_from_atomic_densities(basis_functions)
aed_sg, gd = Density.get_all_electron_density(self,
atoms,
gridrefinement=2)
return aed_sg[0], gd
def initialize(self, density, hamiltonian, spos_ac):
"""Initialize wave-functions, density and hamiltonian.
Return (nlcao, nrand) tuple with number of bands intialized from
LCAO and random numbers, respectively."""
if self.mykpts[0].psit is None:
basis_functions = BasisFunctions(
self.gd, [setup.phit_j for setup in self.setups],
self.kd,
dtype=self.dtype,
cut=True)
basis_functions.set_positions(spos_ac)
else:
self.initialize_wave_functions_from_restart_file()
if self.mykpts[0].psit is not None:
density.initialize_from_wavefunctions(self)
elif density.nt_sG is None:
density.initialize_from_atomic_densities(basis_functions)
# Initialize GLLB-potential from basis function orbitals
if hamiltonian.xc.type == 'GLLB':
hamiltonian.xc.initialize_from_atomic_orbitals(basis_functions)
else: # XXX???
# We didn't even touch density, but some combinations in paw.set()
# will make it necessary to do this for some reason.
density.calculate_normalized_charges_and_mix()
hamiltonian.update(density)
if self.mykpts[0].psit is None:
if 1: # self.collinear:
nlcao = self.initialize_wave_functions_from_basis_functions(
basis_functions, density, hamiltonian, spos_ac)
else:
self.random_wave_functions(0)
nlcao = 0
nrand = self.bd.nbands - nlcao
else:
# We got everything from file:
nlcao = 0
nrand = 0
return nlcao, nrand
def get_density(self, atom_indices=None, gridrefinement=2):
"""Get sum of atomic densities from the given atom list.
Parameters
----------
atom_indices : list_like
All atoms are taken if the list is not given.
gridrefinement : 1, 2, 4
Gridrefinement given to get_all_electron_density
Returns
-------
type
spin summed density, grid_descriptor
"""
all_atoms = self.calculator.get_atoms()
if atom_indices is None:
atom_indices = range(len(all_atoms))
# select atoms
atoms = self.calculator.get_atoms()[atom_indices]
spos_ac = atoms.get_scaled_positions()
Z_a = atoms.get_atomic_numbers()
par = self.calculator.parameters
setups = Setups(Z_a, par.setups, par.basis, XC(par.xc),
self.calculator.wfs.world)
# initialize
self.initialize(setups, self.calculator.timer, np.zeros(
(len(atoms), 3)), False)
self.set_mixer(None)
rank_a = self.gd.get_ranks_from_positions(spos_ac)
self.set_positions(spos_ac, AtomPartition(self.gd.comm, rank_a))
basis_functions = BasisFunctions(
self.gd, [setup.phit_j for setup in self.setups], cut=True)
basis_functions.set_positions(spos_ac)
self.initialize_from_atomic_densities(basis_functions)
aed_sg, gd = self.get_all_electron_density(atoms, gridrefinement)
return aed_sg.sum(axis=0), gd
def __init__(self, ksl, gd, nvalence, setups, bd,
dtype, world, kd, timer=None):
WaveFunctions.__init__(self, gd, nvalence, setups, bd,
dtype, world, kd, timer)
self.ksl = ksl
self.S_qMM = None
self.T_qMM = None
self.P_aqMi = None
self.timer.start('TCI: Evaluate splines')
self.tci = NewTCI(gd.cell_cv, gd.pbc_c, setups, kd.ibzk_qc, kd.gamma)
self.timer.stop('TCI: Evaluate splines')
self.basis_functions = BasisFunctions(gd,
[setup.phit_j
for setup in setups],
kd.comm,
cut=True)
if not kd.gamma:
self.basis_functions.set_k_points(kd.ibzk_qc)
def get_all_electron_density(self,
atoms=None,
gridrefinement=2,
spos_ac=None,
skip_core=False):
"""Return real all-electron density array.
Usage: Either get_all_electron_density(atoms) or
get_all_electron_density(spos_ac=spos_ac)
skip_core=True theoretically returns the
all-electron valence density (use with
care; will not in general integrate
to valence)
"""
if spos_ac is None:
spos_ac = atoms.get_scaled_positions() % 1.0
# Refinement of coarse grid, for representation of the AE-density
# XXXXXXXXXXXX think about distribution depending on gridrefinement!
if gridrefinement == 1:
gd = self.redistributor.aux_gd
n_sg = self.nt_sG.copy()
# This will get the density with the same distribution
# as finegd:
n_sg = self.redistributor.distribute(n_sg)
elif gridrefinement == 2:
gd = self.finegd
if self.nt_sg is None:
self.interpolate_pseudo_density()
n_sg = self.nt_sg.copy()
elif gridrefinement == 4:
# Extra fine grid
gd = self.finegd.refine()
# Interpolation function for the density:
interpolator = Transformer(self.finegd, gd, 3) # XXX grids!
# Transfer the pseudo-density to the fine grid:
n_sg = gd.empty(self.nspins)
if self.nt_sg is None:
self.interpolate_pseudo_density()
for s in range(self.nspins):
interpolator.apply(self.nt_sg[s], n_sg[s])
else:
raise NotImplementedError
# Add corrections to pseudo-density to get the AE-density
splines = {}
phi_aj = []
phit_aj = []
nc_a = []
nct_a = []
for a, id in enumerate(self.setups.id_a):
if id in splines:
phi_j, phit_j, nc, nct = splines[id]
else:
# Load splines:
phi_j, phit_j, nc, nct = self.setups[a].get_partial_waves()[:4]
splines[id] = (phi_j, phit_j, nc, nct)
phi_aj.append(phi_j)
phit_aj.append(phit_j)
nc_a.append([nc])
nct_a.append([nct])
# Create localized functions from splines
phi = BasisFunctions(gd, phi_aj)
phit = BasisFunctions(gd, phit_aj)
nc = LFC(gd, nc_a)
nct = LFC(gd, nct_a)
phi.set_positions(spos_ac)
phit.set_positions(spos_ac)
nc.set_positions(spos_ac)
nct.set_positions(spos_ac)
I_sa = np.zeros((self.nspins, len(spos_ac)))
a_W = np.empty(len(phi.M_W), np.intc)
W = 0
for a in phi.atom_indices:
nw = len(phi.sphere_a[a].M_w)
a_W[W:W + nw] = a
W += nw
x_W = phi.create_displacement_arrays()[0]
D_asp = self.D_asp # XXX really?
rho_MM = np.zeros((phi.Mmax, phi.Mmax))
for s, I_a in enumerate(I_sa):
M1 = 0
for a, setup in enumerate(self.setups):
ni = setup.ni
D_sp = D_asp.get(a)
if D_sp is None:
D_sp = np.empty((self.nspins, ni * (ni + 1) // 2))
else:
I_a[a] = (
(setup.Nct) / self.nspins -
sqrt(4 * pi) * np.dot(D_sp[s], setup.Delta_pL[:, 0]))
if not skip_core:
I_a[a] -= setup.Nc / self.nspins
if gd.comm.size > 1:
gd.comm.broadcast(D_sp, D_asp.partition.rank_a[a])
M2 = M1 + ni
rho_MM[M1:M2, M1:M2] = unpack2(D_sp[s])
M1 = M2
assert np.all(n_sg[s].shape == phi.gd.n_c)
phi.lfc.ae_valence_density_correction(rho_MM, n_sg[s], a_W, I_a,
x_W)
phit.lfc.ae_valence_density_correction(-rho_MM, n_sg[s], a_W, I_a,
x_W)
a_W = np.empty(len(nc.M_W), np.intc)
W = 0
for a in nc.atom_indices:
nw = len(nc.sphere_a[a].M_w)
a_W[W:W + nw] = a
W += nw
scale = 1.0 / self.nspins
for s, I_a in enumerate(I_sa):
if not skip_core:
nc.lfc.ae_core_density_correction(scale, n_sg[s], a_W, I_a)
nct.lfc.ae_core_density_correction(-scale, n_sg[s], a_W, I_a)
gd.comm.sum(I_a)
N_c = gd.N_c
g_ac = np.around(N_c * spos_ac).astype(int) % N_c - gd.beg_c
if not skip_core:
for I, g_c in zip(I_a, g_ac):
if (g_c >= 0).all() and (g_c < gd.n_c).all():
n_sg[s][tuple(g_c)] -= I / gd.dv
return n_sg, gd
def new_get_all_electron_density(self, atoms, gridrefinement=2):
"""Return real all-electron density array."""
# Refinement of coarse grid, for representation of the AE-density
if gridrefinement == 1:
gd = self.gd
n_sg = self.nt_sG.copy()
elif gridrefinement == 2:
gd = self.finegd
if self.nt_sg is None:
self.interpolate()
n_sg = self.nt_sg.copy()
elif gridrefinement == 4:
# Extra fine grid
gd = self.finegd.refine()
# Interpolation function for the density:
interpolator = Transformer(self.finegd, gd, 3)
# Transfer the pseudo-density to the fine grid:
n_sg = gd.empty(self.nspins)
if self.nt_sg is None:
self.interpolate()
for s in range(self.nspins):
interpolator.apply(self.nt_sg[s], n_sg[s])
else:
raise NotImplementedError
# Add corrections to pseudo-density to get the AE-density
splines = {}
phi_aj = []
phit_aj = []
nc_a = []
nct_a = []
for a, id in enumerate(self.setups.id_a):
if id in splines:
phi_j, phit_j, nc, nct = splines[id]
else:
# Load splines:
phi_j, phit_j, nc, nct = self.setups[a].get_partial_waves()[:4]
splines[id] = (phi_j, phit_j, nc, nct)
phi_aj.append(phi_j)
phit_aj.append(phit_j)
nc_a.append([nc])
nct_a.append([nct])
# Create localized functions from splines
phi = BasisFunctions(gd, phi_aj)
phit = BasisFunctions(gd, phit_aj)
nc = LFC(gd, nc_a)
nct = LFC(gd, nct_a)
spos_ac = atoms.get_scaled_positions() % 1.0
phi.set_positions(spos_ac)
phit.set_positions(spos_ac)
nc.set_positions(spos_ac)
nct.set_positions(spos_ac)
I_sa = np.zeros((self.nspins, len(atoms)))
a_W = np.empty(len(phi.M_W), np.int32)
W = 0
for a in phi.atom_indices:
nw = len(phi.sphere_a[a].M_w)
a_W[W:W + nw] = a
W += nw
rho_MM = np.zeros((phi.Mmax, phi.Mmax))
for s, I_a in enumerate(I_sa):
M1 = 0
for a, setup in enumerate(self.setups):
ni = setup.ni
D_sp = self.D_asp.get(a)
if D_sp is None:
D_sp = np.empty((self.nspins, ni * (ni + 1) // 2))
else:
I_a[a] = ((setup.Nct - setup.Nc) / self.nspins -
sqrt(4 * pi) *
np.dot(D_sp[s], setup.Delta_pL[:, 0]))
if gd.comm.size > 1:
gd.comm.broadcast(D_sp, self.rank_a[a])
M2 = M1 + ni
rho_MM[M1:M2, M1:M2] = unpack2(D_sp[s])
M1 = M2
phi.lfc.ae_valence_density_correction(rho_MM, n_sg[s], a_W, I_a)
phit.lfc.ae_valence_density_correction(-rho_MM, n_sg[s], a_W, I_a)
a_W = np.empty(len(nc.M_W), np.int32)
W = 0
for a in nc.atom_indices:
nw = len(nc.sphere_a[a].M_w)
a_W[W:W + nw] = a
W += nw
scale = 1.0 / self.nspins
for s, I_a in enumerate(I_sa):
nc.lfc.ae_core_density_correction(scale, n_sg[s], a_W, I_a)
nct.lfc.ae_core_density_correction(-scale, n_sg[s], a_W, I_a)
gd.comm.sum(I_a)
N_c = gd.N_c
g_ac = np.around(N_c * spos_ac).astype(int) % N_c - gd.beg_c
for I, g_c in zip(I_a, g_ac):
if (g_c >= 0).all() and (g_c < gd.n_c).all():
n_sg[s][tuple(g_c)] -= I / gd.dv
return n_sg, gd
# / X : \
# +--/-----------/-\--:--------\----------+
# | | | | : | |
# | | | | : | |
# | | x | | : x | |
# | | | | : | |
# | | | | : | |
# +--\-----------\-/--:--------/----------+
# \ X : /
# \ / \ : /
# --------- --:------
# :
#
# ':' is the domain wall if split on two cpu's
gd = GridDescriptor(N_c=[40, 8, 8], cell_cv=[10., 2., 2.], pbc_c=(0, 1, 1))
pos_ac = np.array([[.25, .5, .5], [.55, .5, .5]])
kpts_kc = None
spline = Spline(l=0,
rmax=2.0,
f_g=np.array([1, 0.9, 0.1, 0.0]),
r_g=None,
beta=None,
points=25)
spline_aj = [[spline] for pos_c in pos_ac]
bfs = BasisFunctions(gd, spline_aj)
if kpts_kc is not None:
bfs.set_k_points(kpts_kc)
bfs.set_positions(pos_ac)
def paw_corrections(self, gridrefinement=2):
Fn_wsg, gd = self.interpolate_pseudo_density(gridrefinement)
# Splines
splines = {}
phi_aj = []
phit_aj = []
for a, id in enumerate(self.setups.id_a):
if id in splines:
phi_j, phit_j = splines[id]
else:
# Load splines:
phi_j, phit_j = self.setups[a].get_partial_waves()[:2]
splines[id] = (phi_j, phit_j)
phi_aj.append(phi_j)
phit_aj.append(phit_j)
# Create localized functions from splines
phi = BasisFunctions(gd, phi_aj, dtype=float)
phit = BasisFunctions(gd, phit_aj, dtype=float)
# phi = BasisFunctions(gd, phi_aj, dtype=complex)
# phit = BasisFunctions(gd, phit_aj, dtype=complex)
spos_ac = self.atoms.get_scaled_positions()
phi.set_positions(spos_ac)
phit.set_positions(spos_ac)
tmp_g = gd.empty(dtype=float)
rho_MM = np.zeros((phi.Mmax, phi.Mmax), dtype=self.dtype)
rho2_MM = np.zeros_like(rho_MM)
for w in range(self.nw):
for s in range(self.nspins):
rho_MM[:] = 0
M1 = 0
for a, setup in enumerate(self.setups):
ni = setup.ni
FD_wsp = self.FD_awsp.get(a)
if FD_wsp is None:
FD_p = np.empty((ni * (ni + 1) // 2), dtype=self.dtype)
else:
FD_p = FD_wsp[w][s]
if gd.comm.size > 1:
gd.comm.broadcast(FD_p, self.rank_a[a])
D_ij = unpack2(FD_p)
# unpack does complex conjugation that we don't want so
# remove conjugation
D_ij = np.triu(D_ij, 1) + np.conj(np.tril(D_ij))
# if FD_wsp is None:
# FD_wsp = np.empty((self.nw, self.nspins,
# ni * (ni + 1) // 2),
# dtype=self.dtype)
# if gd.comm.size > 1:
# gd.comm.broadcast(FD_wsp, self.rank_a[a])
# D_ij = unpack2(FD_wsp[w][s])
# D_ij = np.triu(D_ij, 1) + np.conj(np.tril(D_ij))
M2 = M1 + ni
rho_MM[M1:M2, M1:M2] = D_ij
M1 = M2
# Add real part of AE corrections
tmp_g[:] = 0
rho2_MM[:] = rho_MM.real
# TODO: use ae_valence_density_correction
phi.construct_density(rho2_MM, tmp_g, q=-1)
phit.construct_density(-rho2_MM, tmp_g, q=-1)
# phi.lfc.ae_valence_density_correction(rho2_MM, tmp_g,
# np.zeros(len(phi.M_W),
# np.intc),
# np.zeros(self.na))
# phit.lfc.ae_valence_density_correction(-rho2_MM, tmp_g,
# np.zeros(len(phi.M_W),
# np.intc),
# np.zeros(self.na))
Fn_wsg[w][s] += tmp_g
# Add imag part of AE corrections
tmp_g[:] = 0
rho2_MM[:] = rho_MM.imag
# TODO: use ae_valence_density_correction
phi.construct_density(rho2_MM, tmp_g, q=-1)
phit.construct_density(-rho2_MM, tmp_g, q=-1)
# phi.lfc.ae_valence_density_correction(rho2_MM, tmp_g,
# np.zeros(len(phi.M_W),
# np.intc),
# np.zeros(self.na))
# phit.lfc.ae_valence_density_correction(-rho2_MM, tmp_g,
# np.zeros(len(phi.M_W),
# np.intc),
# np.zeros(self.na))
Fn_wsg[w][s] += 1.0j * tmp_g
return Fn_wsg, gd
def with_ae_corrections(self, finegrid=False):
"""Get pair density including the AE corrections"""
nij_g = self.get(finegrid)
# Generate the density matrix
D_ap = {}
# D_aii = {}
for a, P_ni in self.P_ani.items():
Pi_i = P_ni[self.i]
Pj_i = P_ni[self.j]
D_ii = np.outer(Pi_i.conj(), Pj_i)
# Note: D_ii is not symmetric but the products of partial waves are
# so that we can pack
D_ap[a] = pack(D_ii)
# D_aii[a] = D_ii
# Load partial waves if needed
if ((finegrid and (not hasattr(self, 'phi')))
or ((not finegrid) and (not hasattr(self, 'Phi')))):
# Splines
splines = {}
phi_aj = []
phit_aj = []
for a, id in enumerate(self.setups.id_a):
if id in splines:
phi_j, phit_j = splines[id]
else:
# Load splines:
phi_j, phit_j = self.setups[a].get_partial_waves()[:2]
splines[id] = (phi_j, phit_j)
phi_aj.append(phi_j)
phit_aj.append(phit_j)
# Store partial waves as class variables
if finegrid:
gd = self.density.finegd
self.__class__.phi = BasisFunctions(gd, phi_aj)
self.__class__.phit = BasisFunctions(gd, phit_aj)
self.__class__.phi.set_positions(self.spos_ac)
self.__class__.phit.set_positions(self.spos_ac)
else:
gd = self.density.gd
self.__class__.Phi = BasisFunctions(gd, phi_aj)
self.__class__.Phit = BasisFunctions(gd, phit_aj)
self.__class__.Phi.set_positions(self.spos_ac)
self.__class__.Phit.set_positions(self.spos_ac)
# Add AE corrections
if finegrid:
phi = self.phi
phit = self.phit
gd = self.density.finegd
else:
phi = self.Phi
phit = self.Phit
gd = self.density.gd
rho_MM = np.zeros((phi.Mmax, phi.Mmax))
M1 = 0
for a, setup in enumerate(self.setups):
ni = setup.ni
D_p = D_ap.get(a)
if D_p is None:
D_p = np.empty((ni * (ni + 1) // 2))
if gd.comm.size > 1:
gd.comm.broadcast(D_p, self.wfs.partition.rank_a[a])
D_ii = unpack2(D_p)
# D_ii = D_aii.get(a)
# if D_ii is None:
# D_ii = np.empty((ni, ni))
# if gd.comm.size > 1:
# gd.comm.broadcast(D_ii, self.wfs.atom_partition.rank_a[a])
M2 = M1 + ni
rho_MM[M1:M2, M1:M2] = D_ii
M1 = M2
# construct_density assumes symmetric rho_MM and
# takes only the upper half of it
phi.construct_density(rho_MM, nij_g, q=-1)
phit.construct_density(-rho_MM, nij_g, q=-1)
# TODO: use ae_valence_density_correction
# phi.lfc.ae_valence_density_correction(
# rho_MM, nij_g, np.zeros(len(phi.M_W), np.intc), np.zeros(self.na))
# phit.lfc.ae_valence_density_correction(
# -rho_MM, nij_g, np.zeros(len(phit.M_W), np.intc), np.zeros(self.na))
return nij_g
