def test_compare_abtem_to_gpaw():
from gpaw import GPAW
from gpaw.utilities.ps2ae import PS2AE
from abtem.dft import GPAWPotential
atoms = Atoms('C', positions=[(0, 1, 2)], cell=(2, 2, 4), pbc=True)
calc = GPAW(h=.2, txt=None, kpts=(2, 2, 1))
atoms.calc = calc
atoms.get_potential_energy()
h = 0.01
t = PS2AE(calc, h=h)
ae = (-t.get_electrostatic_potential(rcgauss=.02 * units.Bohr) * h).sum(-1)
ae -= ae.min()
dft_pot = GPAWPotential(calc,
gpts=ae.shape,
core_size=.02,
slice_thickness=4)
dft_array = dft_pot.build(pbar=False)
abtem_ae = dft_array.array.sum(0)
abtem_ae -= abtem_ae.min()
valid = abtem_ae > 1
assert np.all(((abtem_ae[valid] - ae[valid]) / ae[valid]).max() < .001)
def get_ae_potential(atoms):
calc = GPAW(hund=True, eigensolver='cg', h=.2)
atoms.set_calculator(calc)
atoms.get_potential_energy()
ps2ae = PS2AE(atoms.calc)
return ps2ae.get_electrostatic_potential(ae=True)
def get_ae_pair_density_matrix(self, calc_A, calc_B, matrix_name=None):
if calc_A.wfs.kd.nibzkpts != 1:
raise ValueError(ae_ibz_error)
# <Psi_A|Psi_B> using the all-electron pair density
psi_A = PS2AE(calc_A, h=self.h)
psi_B = PS2AE(calc_B, h=self.h)
ns = calc_A.wfs.nspins
# total of filled a and b bands for each spin and kpt
n_occup_A = self.get_n_occupied_bands(calc_A)
n_occup_B = self.get_n_occupied_bands(calc_B)
# list to store k-dependent pair density
n_AB = []
w_k = np.zeros(calc_A.wfs.kd.nibzkpts) #store kpt weights
# get overlap at for each ij band at kpt and spin
# the resulting matrix is organized in alpha and beta blocks
# | | |
# | a | 0 | a:<psi_a|psi_a> != 0
# S =|_______|____| , <psi_a|psi_b> = 0
# | 0 | b | b:<psi_b|psi_b> != 0
# | | |
#
# a = naa x naa, b = nab x nab
for spin in range(ns):
for k in range(calc_A.wfs.kd.nibzkpts):
nAa, nAb, nBa, nBb = self.check_bands(n_occup_A, n_occup_B, k)
nas, n_occup, n_occup_s = self.check_spin_and_occupations(
nAa, nAb, nBa, nBb)
kd = calc_A.wfs.kd
w_kA = kd.weight_k[k]
kd = calc_B.wfs.kd
w_kB = kd.weight_k[k]
# check that a and b cDFT states have similar spin state
if np.sign(nAa - nAb) != np.sign(nBa - nBb):
warning = UserWarning(
'The cDFT wave functions have\n'
'different spin states! Similar\n'
'spin states are required for coupling constant\n'
'calculation!')
warnings.warn(warning)
# form overlap matrices of correct size for each kpt
if spin == 0:
n_AB.append(np.zeros((n_occup, n_occup), dtype=np.complex))
for i in range(n_occup_s[spin]):
for j in range(n_occup_s[spin]):
# take only the bands which contain electrons in spin-orbital
psi_kA = psi_A.get_wave_function(n=i,
k=k,
s=spin,
ae=True)
psi_kB = psi_B.get_wave_function(n=j,
k=k,
s=spin,
ae=True)
n_ij = psi_A.gd.integrate(psi_kA.conj(),
psi_kB,
global_integral=True)
if spin == 0 and i == 0 and j == 0:
w_k[k] = (w_kA + w_kB) / 2.
I = spin * nas + i
J = spin * nas + j
n_AB[k][I][J] = n_ij * Bohr**3
n_AB = np.asarray(n_AB)
self.w_k = w_k
self.n_ab = n_AB
if self.save_matrix:
if matrix_name is None:
np.save(self.S_matrix + 'final', self.n_ab)
else:
np.save('%s_final' % matrix_name, self.n_ab)
return self.n_ab, self.w_k
def get_ae_pair_weight_matrix(self):
if self.calc_A.wfs.kd.nibzkpts != 1:
raise ValueError(ae_ibz_error)
if hasattr(self, 'n_ab'):
pass
else:
raise ValueError(nab_missing_error)
# pseudo wfs to all-electron wfs
psi_A = PS2AE(self.calc_A, h=self.h)
psi_B = PS2AE(self.calc_B, h=self.h)
ns = self.calc_A.wfs.nspins
# weight functions sum_c VcWc
wa = []
wb = []
for a in range(len(self.Va)):
wa.append(
interpolate_weight(self.calc_A, self.fineweightA[a], h=self.h))
for b in range(len(self.Vb)):
wb.append(
interpolate_weight(self.calc_B, self.fineweightB[b], h=self.h))
wa = np.asarray(wa)
wb = np.asarray(wb)
# check number of occupied and total number of bands
n_occup_A = self.get_n_occupied_bands(self.calc_A)
n_occup_B = self.get_n_occupied_bands(self.calc_B)
# place to store <i_A(k)|w|j_B(k)>
w_kij_AB = []
w_kij_BA = []
# k-point weights
w_k = np.zeros(self.calc_A.wfs.kd.nibzkpts, dtype=np.float32)
# get weight matrix at for each ij band at kpt and spin
# the resulting matrix is organized in alpha and beta blocks
# | | |
# | a | 0 | a:<psi_a|Vb*wb|psi_a> != 0
# VW_AB =|_______|____| , <psi_a|w|psi_b> = 0
# | 0 | b | b:<psi_b|Vb*wb|psi_b> != 0
# | | |
#
# a = nAa x nAa, b = nAb x nAb
for spin in range(ns):
for k in range(self.calc_A.wfs.kd.nibzkpts):
# k-dependent overlap/pair density matrices
inv_S = np.linalg.inv(self.n_ab[k])
det_S = np.linalg.det(self.n_ab[k])
I = np.identity(inv_S.shape[0])
C_ab = np.transpose(np.dot(inv_S, (det_S * I)))
nAa, nAb, nBa, nBb = self.check_bands(n_occup_A, n_occup_B, k)
nas, n_occup, n_occup_s = self.check_spin_and_occupations(
nAa, nAb, nBa, nBb)
# check that a and b cDFT states have similar spin state
if np.sign(nAa - nAb) != np.sign(nBa - nBb):
warning = UserWarning(spin_state_error)
warnings.warn(warning)
# form overlap matrices of correct size for each kpt
if spin == 0:
w_kij_AB.append(
np.zeros((n_occup, n_occup), dtype=np.complex))
w_kij_BA.append(
np.zeros((n_occup, n_occup), dtype=np.complex))
# store k-point weights
kd = self.calc_A.wfs.kd
w_kA = kd.weight_k[k]
kd = self.calc_B.wfs.kd
w_kB = kd.weight_k[k]
for i in range(n_occup_s[spin]):
for j in range(n_occup_s[spin]):
I = spin * nas + i
J = spin * nas + j
psi_kA = psi_A.get_wave_function(n=i,
k=k,
s=spin,
ae=True)
psi_kB = psi_B.get_wave_function(n=j,
k=k,
s=spin,
ae=True)
w_ij_AB = []
w_ji_BA = []
for b in range(len(self.Vb)):
integral = psi_B.gd.integrate(
psi_kA.conj() * wb[b] * psi_kB,
global_integral=True) * C_ab[I][J]
if b >= self.n_charge_regionsB and spin == 1:
# for charge constraint w > 0
integral *= -1.
w_ij_AB.append(-integral)
for a in range(len(self.Va)):
integral = psi_A.gd.integrate(
psi_kB.conj() * wa[a] * psi_kA,
global_integral=True) * C_ab[J][I]
if a >= self.n_charge_regionsA and spin == 1:
integral *= -1.
w_ji_BA.append(-integral)
w_ij_AB = np.asarray(w_ij_AB) * Bohr**3
w_ji_BA = np.asarray(w_ji_BA) * Bohr**3
# collect kpt weight, only once per kpt
if spin == 0 and i == 0 and j == 0:
w_k[k] = (w_kA + w_kB) / 2.
w_kij_AB[k][I][J] += np.dot(self.Vb, w_ij_AB).sum()
w_kij_BA[k][J][I] += np.dot(self.Va, w_ji_BA).sum()
self.w_k = w_k
self.VW_AB = w_kij_AB
self.VW_BA = w_kij_BA
if self.save_matrix:
np.save(self.VW_matrix + 'final_AB', self.VW_AB)
np.save(self.VW_matrix + 'final_BA', self.VW_BA)
return self.VW_AB, self.VW_BA, self.w_k
# creates: hli.png
import matplotlib.pyplot as plt
from ase import Atoms
from ase.units import Bohr
from gpaw.utilities.ps2ae import PS2AE
from gpaw import GPAW
hli = Atoms('HLi', positions=[[0, 0, 0], [0, 0, 1.6]])
hli.center(vacuum=2.5)
hli.set_calculator(GPAW(txt='hli.txt', mode='fd'))
hli.get_potential_energy()
# Transformer:
t = PS2AE(hli.calc, h=0.05)
for n, color in enumerate(['green', 'red']):
ps = t.get_wave_function(n, ae=False)
ae = t.get_wave_function(n)
norm = t.gd.integrate(ae**2)
print('Norm:', norm)
assert abs(norm - 1) < 1e-2
i = ps.shape[0] // 2
x = t.gd.coords(2) * Bohr
# Interpolated PS and AE wfs:
plt.plot(x,
ps[i, i],
'--',
color=color,
label=r'$\tilde\psi_{}$'.format(n))
plt.plot(x, ae[i, i], '-', color=color, label=r'$\psi_{}$'.format(n))
import matplotlib.pyplot as plt
from ase.units import Bohr
from gpaw.utilities.ps2ae import PS2AE
from gpaw import GPAW
calc = GPAW('hli.gpw', txt=None)
# Transformer:
t = PS2AE(calc, h=0.05)
# Interpolated PS and AE potentials:
ps = t.get_electrostatic_potential(ae=False)
ae = t.get_electrostatic_potential()
i = ps.shape[0] // 2
x = t.gd.coords(2) * Bohr
plt.plot(x, ps[i, i], '--', label=r'$\tilde v$')
plt.plot(x, ae[i, i], '-', label=r'$v$')
# Raw PS wfs:
ps0 = calc.get_electrostatic_potential()
gd = calc.hamiltonian.finegd
i = ps0.shape[0] // 2
X = gd.coords(2) * Bohr
plt.plot(X, ps0[i, i], 'o')
plt.plot(x, 0 * x, 'k')
plt.xlabel('z [Ang]')
plt.ylabel('potential [eV]')
plt.ylim(bottom=-100, top=10)
plt.legend()