def allocate(self):
if self.allocated:
raise RuntimeError("Already allocated!")
# Calculate complex phase factors:
self.phase_ucd = np.ones((self.kd.mynks, 3, 2), complex)
if not self.gamma:
for myu, phase_cd in enumerate(self.phase_ucd):
u = self.kd.global_index(myu)
s, k = self.kd.what_is(u)
phase_cd[:] = np.exp(2j * np.pi * self.gd.sdisp_cd * self.ibzk_kc[k, :, np.newaxis])
assert self.dtype == complex, "Complex wavefunctions are required."
self.r_vG = self.gd.get_grid_point_coordinates()
self.wf_uG = self.gd.zeros(self.kd.mynks, dtype=self.dtype)
self.laplace0_uG = self.gd.zeros(self.kd.mynks, dtype=self.dtype)
buf_G = self.gd.empty(dtype=self.dtype)
sdisp_Ac = []
for a, spos_c in enumerate(self.atoms.get_scaled_positions() % 1.0):
for sdisp_x in range(-1 * self.gd.pbc_c[0], self.gd.pbc_c[0] + 1):
for sdisp_y in range(-1 * self.gd.pbc_c[1], self.gd.pbc_c[1] + 1):
for sdisp_z in range(-1 * self.gd.pbc_c[2], self.gd.pbc_c[2] + 1):
sdisp_c = np.array([sdisp_x, sdisp_y, sdisp_z])
if debug and world.rank == 0:
print "a=%d, spos=%s, sdisp_c=%s" % (a, spos_c, sdisp_c)
sdisp_Ac.append((a, spos_c, sdisp_c))
for a, spos_c, sdisp_c in sdisp_Ac:
if debug and world.rank == 0:
print "Adding gaussian at a=%d, spos=%s, sigma=%8.5f Ang" % (a, spos_c + sdisp_c, self.sigma * Bohr)
r0_v = np.dot(spos_c + sdisp_c, self.gd.cell_cv)
for myu in range(self.kd.mynks):
u = self.kd.global_index(myu)
s, k = self.kd.what_is(u)
ibzk_v = self.ibzk_kv[k]
# f(r) = sum_a A exp(-|r-R^a|^2 / 2sigma^2) exp(i k.r)
gaussian_wave(self.r_vG, r0_v, self.sigma, ibzk_v, A=1.0, dtype=self.dtype, out_G=buf_G)
self.wf_uG[myu] += buf_G
# d^2/dx^2 exp(ikx-(x-x0)^2/2sigma^2)
# ((ik-(x-x0)/sigma^2)^2 - 1/sigma^2) exp(ikx-(x-x0)^2/2sigma^2)
dr2_G = np.sum(
(
1j * ibzk_v[:, np.newaxis, np.newaxis, np.newaxis]
- (self.r_vG - r0_v[:, np.newaxis, np.newaxis, np.newaxis]) / self.sigma ** 2
)
** 2,
axis=0,
)
self.laplace0_uG[myu] += (dr2_G - 3 / self.sigma ** 2) * buf_G
self.allocated = True
def allocate(self):
if self.allocated:
raise RuntimeError('Already allocated!')
# Calculate complex phase factors:
self.phase_ucd = np.ones((self.kd.mynks, 3, 2), complex)
if not self.gamma:
for myu, phase_cd in enumerate(self.phase_ucd):
u = self.kd.global_index(myu)
s, k = self.kd.what_is(u)
phase_cd[:] = np.exp(2j * np.pi * self.gd.sdisp_cd * \
self.ibzk_kc[k,:,np.newaxis])
assert self.dtype == complex, 'Complex wavefunctions are required.'
self.r_vG = self.gd.get_grid_point_coordinates()
self.wf_uG = self.gd.zeros(self.kd.mynks, dtype=self.dtype)
self.laplace0_uG = self.gd.zeros(self.kd.mynks, dtype=self.dtype)
buf_G = self.gd.empty(dtype=self.dtype)
sdisp_Ac = []
for a,spos_c in enumerate(self.atoms.get_scaled_positions() % 1.0):
for sdisp_x in range(-1*self.gd.pbc_c[0],self.gd.pbc_c[0]+1):
for sdisp_y in range(-1*self.gd.pbc_c[1],self.gd.pbc_c[1]+1):
for sdisp_z in range(-1*self.gd.pbc_c[2],self.gd.pbc_c[2]+1):
sdisp_c = np.array([sdisp_x, sdisp_y, sdisp_z])
if debug and world.rank == 0:
print 'a=%d, spos=%s, sdisp_c=%s' % (a,spos_c,sdisp_c)
sdisp_Ac.append((a,spos_c,sdisp_c))
for a,spos_c,sdisp_c in sdisp_Ac:
if debug and world.rank == 0:
print 'Adding gaussian at a=%d, spos=%s, sigma=%8.5f Ang' % (a,spos_c+sdisp_c,self.sigma*Bohr)
r0_v = np.dot(spos_c+sdisp_c, self.gd.cell_cv)
for myu in range(self.kd.mynks):
u = self.kd.global_index(myu)
s, k = self.kd.what_is(u)
ibzk_v = self.ibzk_kv[k]
# f(r) = sum_a A exp(-|r-R^a|^2 / 2sigma^2) exp(i k.r)
gaussian_wave(self.r_vG, r0_v, self.sigma, ibzk_v, A=1.0,
dtype=self.dtype, out_G=buf_G)
self.wf_uG[myu] += buf_G
# d^2/dx^2 exp(ikx-(x-x0)^2/2sigma^2)
# ((ik-(x-x0)/sigma^2)^2 - 1/sigma^2) exp(ikx-(x-x0)^2/2sigma^2)
dr2_G = np.sum((1j*ibzk_v[:,np.newaxis,np.newaxis,np.newaxis] \
- (self.r_vG-r0_v[:,np.newaxis,np.newaxis,np.newaxis]) \
/ self.sigma**2)**2, axis=0)
self.laplace0_uG[myu] += (dr2_G - 3/self.sigma**2) * buf_G
self.allocated = True
def get_scaled_gaussian_wave(self, pos_c, sigma=None, k_c=None, scale=None):
if sigma is None:
sigma = self._sigma0
if k_c is None:
k_c = self._k0_c
if scale is None:
A = None
else:
# 4*pi*int(exp(-r^2/(2*w^2))^2*r^2, r=0...infinity)= w^3*pi^(3/2)
# = scale/A^2 -> A = scale*(sqrt(Pi)*w)^(-3/2) hence int -> scale^2
A = scale/(sigma*(np.pi)**0.5)**1.5
return gaussian_wave(self.r_cG, pos_c, sigma, k_c, A, self.dtype, self.buf_G)
def get_scaled_gaussian_wave(self, pos_c, sigma=None, k_c=None, scale=None):
if sigma is None:
sigma = self._sigma0
if k_c is None:
k_c = self._k0_c
if scale is None:
A = None
else:
# 4*pi*int(exp(-r^2/(2*w^2))^2*r^2, r=0...infinity)= w^3*pi^(3/2)
# = scale/A^2 -> A = scale*(sqrt(Pi)*w)^(-3/2) hence int -> scale^2
A = scale/(sigma*(np.pi)**0.5)**1.5
return gaussian_wave(self.r_cG, pos_c, sigma, k_c, A, self.dtype, self.buf_G)
k_c = np.random.normal(size=C)
A = np.random.uniform()*np.exp(1j*np.random.uniform(0,2*np.pi))
print('Allocation: %8.5f s' % (time.time()-t))
# -------------------------------------------------------------------
# Test case for real-part of gamma-point wave with normalized amplitude
_gaussRGN = lambda r_cG, r0_c, sigma: 1/(sigma*np.pi**0.5)**1.5 \
* np.exp(-np.sum((r_cG-r0_c[:,np.newaxis,np.newaxis,np.newaxis])**2, axis=0)/(2*sigma**2))
t = time.time()
gs0_G = _gaussRGN(r_cG, r0_c, sigma)
print('_gaussRGN: %8.5f s' % (time.time()-t))
t = time.time()
gs1_G = gaussian_wave(r_cG, r0_c, sigma)
print('+gaussRGN: %8.5f s' % (time.time()-t))
assert np.abs(gs0_G-gs1_G).max() < 1e-12, 'Max error %g' % np.abs(gs0_G-gs1_G).max()
del gs0_G, gs1_G
# Test case for real-part of gamma-point wave with complex amplitude
_gaussRGA = lambda r_cG, r0_c, sigma, A: np.real(A) \
* np.exp(-np.sum((r_cG-r0_c[:,np.newaxis,np.newaxis,np.newaxis])**2, axis=0)/(2*sigma**2))
t = time.time()
gs0_G = _gaussRGA(r_cG, r0_c, sigma, A)
print('_gaussRGA: %8.5f s' % (time.time()-t))
t = time.time()
gs1_G = gaussian_wave(r_cG, r0_c, sigma, None, A)
k_c = np.random.normal(size=C)
A = np.random.uniform()*np.exp(1j*np.random.uniform(0,2*np.pi))
print 'Allocation: %8.5f s' % (time.time()-t)
# -------------------------------------------------------------------
# Test case for real-part of gamma-point wave with normalized amplitude
_gaussRGN = lambda r_cG, r0_c, sigma: 1/(sigma*np.pi**0.5)**1.5 \
* np.exp(-np.sum((r_cG-r0_c[:,np.newaxis,np.newaxis,np.newaxis])**2, axis=0)/(2*sigma**2))
t = time.time()
gs0_G = _gaussRGN(r_cG, r0_c, sigma)
print '_gaussRGN: %8.5f s' % (time.time()-t)
t = time.time()
gs1_G = gaussian_wave(r_cG, r0_c, sigma)
print '+gaussRGN: %8.5f s' % (time.time()-t)
assert np.abs(gs0_G-gs1_G).max() < 1e-12, 'Max error %g' % np.abs(gs0_G-gs1_G).max()
del gs0_G, gs1_G
# Test case for real-part of gamma-point wave with complex amplitude
_gaussRGA = lambda r_cG, r0_c, sigma, A: np.real(A) \
* np.exp(-np.sum((r_cG-r0_c[:,np.newaxis,np.newaxis,np.newaxis])**2, axis=0)/(2*sigma**2))
t = time.time()
gs0_G = _gaussRGA(r_cG, r0_c, sigma, A)
print '_gaussRGA: %8.5f s' % (time.time()-t)
t = time.time()
gs1_G = gaussian_wave(r_cG, r0_c, sigma, None, A)
