def test_nr_rks(self):
cell = pbcgto.Cell()
cell.verbose = 5
cell.output = '/dev/null'
cell.a = np.eye(3) * 2.5
cell.mesh = [21] * 3
cell.atom = [
['He', (1., .8, 1.9)],
['He', (.1, .2, .3)],
]
cell.basis = 'ccpvdz'
cell.build(False, False)
grids = gen_grid.UniformGrids(cell)
grids.build()
nao = cell.nao_nr()
np.random.seed(1)
kpts = np.random.random((2, 3))
dms = np.random.random((2, nao, nao))
dms = (dms + dms.transpose(0, 2, 1)) * .5
ni = numint.NumInt()
with lib.temporary_env(pbcgto.eval_gto, EXTRA_PREC=1e-5):
ne, exc, vmat = ni.nr_rks(cell, grids, 'blyp', dms[0], 1, kpts[0])
self.assertAlmostEqual(ne, 5.0499199224525153, 8)
self.assertAlmostEqual(exc, -3.8870579114663886, 8)
self.assertAlmostEqual(lib.fp(vmat),
0.42538491159934377 + 0.14139753327162483j, 8)
ni = numint.KNumInt()
with lib.temporary_env(pbcgto.eval_gto, EXTRA_PREC=1e-5):
ne, exc, vmat = ni.nr_rks(cell, grids, 'blyp', dms, 1, kpts)
self.assertAlmostEqual(ne, 6.0923292346269742, 8)
self.assertAlmostEqual(exc, -3.9899423803106466, 8)
self.assertAlmostEqual(lib.fp(vmat[0]),
-2348.9577179701278 - 60.733087913116719j, 7)
self.assertAlmostEqual(lib.fp(vmat[1]),
-2353.0350086740673 - 117.74811536967495j, 7)
with lib.temporary_env(pbcgto.eval_gto, EXTRA_PREC=1e-5):
ne, exc, vmat = ni.nr_rks(cell, grids, 'blyp', [dms, dms], 1, kpts)
self.assertAlmostEqual(ne[1], 6.0923292346269742, 8)
self.assertAlmostEqual(exc[1], -3.9899423803106466, 8)
self.assertAlmostEqual(lib.fp(vmat[1][0]),
-2348.9577179701278 - 60.733087913116719j, 7)
self.assertAlmostEqual(lib.fp(vmat[1][1]),
-2353.0350086740673 - 117.74811536967495j, 7)
def test_1d_rho(self):
cell = pbcgto.Cell()
cell.a = '5 0 0; 0 1 0; 0 0 1'
cell.unit = 'B'
cell.atom = 'He 1. 0. 1.'
cell.basis = {'He': '321g'}
cell.dimension = 1
cell.verbose = 0
cell.mesh = [10, 30, 30]
cell.build()
grids = gen_grid.UniformGrids(cell)
grids.build()
numpy.random.seed(10)
nao = cell.nao_nr()
dm = numpy.random.random((nao, nao))
dm = dm + dm.T
ni = numint.NumInt()
rho = numint.get_rho(ni, cell, dm, grids)
self.assertAlmostEqual(lib.finger(rho), 1.1624587519868457, 9)
def test_2d_rho(self):
cell = pbcgto.Cell()
cell.a = '5 0 0; 0 5 0; 0 0 1'
cell.unit = 'B'
cell.atom = 'He 1. 0. 1.'
cell.basis = {'He': '321g'}
cell.dimension = 2
cell.low_dim_ft_type = 'inf_vacuum'
cell.verbose = 0
cell.mesh = [10,10,30]
cell.build()
grids = gen_grid.UniformGrids(cell)
grids.build()
numpy.random.seed(10)
nao = cell.nao_nr()
dm = numpy.random.random((nao,nao))
dm = dm + dm.T
ni = numint.NumInt()
rho = numint.get_rho(ni, cell, dm, grids)
self.assertAlmostEqual(lib.finger(rho), 7.2089907050590334, 9)
def test_eval_ao_kpt(self):
cell = pbcgto.Cell()
cell.verbose = 5
cell.output = '/dev/null'
cell.a = np.eye(3) * 2.5
cell.mesh = [21]*3
cell.atom = [['He', (1., .8, 1.9)],
['He', (.1, .2, .3)],]
cell.basis = 'ccpvdz'
cell.build(False, False)
grids = gen_grid.UniformGrids(cell)
grids.build()
np.random.seed(1)
kpt = np.random.random(3)
ni = numint.NumInt()
ao0 = eval_ao(cell, grids.coords, kpt)
ao1 = ni.eval_ao(cell, grids.coords, kpt)
self.assertTrue(numpy.allclose(ao0, ao1, atol=1e-9, rtol=1e-9))
self.assertAlmostEqual(finger(ao1), (-2.4066959390326477-0.98044994099240701j), 8)
def test_eval_ao_kpts(self):
cell = pbcgto.Cell()
cell.verbose = 5
cell.output = '/dev/null'
cell.a = np.eye(3) * 2.5
cell.mesh = [21]*3
cell.atom = [['He', (1., .8, 1.9)],
['He', (.1, .2, .3)],]
cell.basis = 'ccpvdz'
cell.build(False, False)
grids = gen_grid.UniformGrids(cell)
grids.build()
np.random.seed(1)
kpts = np.random.random((4,3))
ni = numint.KNumInt(kpts)
ao1 = ni.eval_ao(cell, grids.coords, kpts)
self.assertAlmostEqual(finger(ao1[0]), (-2.4066959390326477-0.98044994099240701j), 8)
self.assertAlmostEqual(finger(ao1[1]), (-0.30643153325360639+0.1571658820483913j), 8)
self.assertAlmostEqual(finger(ao1[2]), (-1.1937974302337684-0.39039259235266233j), 8)
self.assertAlmostEqual(finger(ao1[3]), (0.17701966968272009-0.20232879692603079j), 8)
def get_A_mat(self):
'''
Construct the projection matrix: A_{m,n}^{\mathbf{k}}
Equation (62) in MV, Phys. Rev. B 56, 12847 or equation (22) in SMV, Phys. Rev. B 65, 035109
'''
A_matrix_loc = np.empty([self.num_kpts_loc, self.num_wann_loc, self.num_bands_loc], dtype = np.complex)
if self.use_bloch_phases == True:
for k_id in range(self.num_kpts_loc):
Amn = np.zeros([self.num_wann_loc, self.num_bands_loc])
np.fill_diagonal(Amn, 1)
A_matrix_loc[k_id,:,:] = Amn
else:
grids = gen_grid.UniformGrids(self.cell)
grids.build()
for k_id in range(self.num_kpts_loc):
kpt = self.cell.get_abs_kpts(self.kpt_latt_loc[k_id])
ao = numint.eval_ao(self.cell, grids.coords, kpt = kpt)
for ith_wann in range(self.num_wann_loc):
frac_site = self.proj_site[ith_wann]
abs_site = frac_site.dot(self.real_lattice_loc) / param.BOHR
l = self.proj_l[ith_wann]
mr = self.proj_m[ith_wann]
r = self.proj_radial[ith_wann]
zona = self.proj_zona[ith_wann]
x_axis = self.proj_x[ith_wann]
z_axis = self.proj_z[ith_wann]
gr = g_r(grids.coords, abs_site, l, mr, r, zona, x_axis, z_axis, unit = 'B')
C = np.dot(self.U[k_id], self.mo_coeff_kpts[k_id])[:,self.band_included_list]
s_ao = np.einsum('i,iu,i->u', grids.weights, ao.conj(), gr)
A_matrix_loc[k_id,ith_wann,:] = np.einsum('um,u->m', C, s_ao)
return A_matrix_loc
def test_eval_ao(self):
cell = pbcgto.Cell()
cell.verbose = 5
cell.output = '/dev/null'
cell.a = np.eye(3) * 2.5
cell.mesh = [21] * 3
cell.atom = [
['C', (1., .8, 1.9)],
['C', (.1, .2, .3)],
]
cell.basis = 'ccpvdz'
cell.build(False, False)
grids = gen_grid.UniformGrids(cell)
grids.build()
ni = numint.NumInt()
ao10 = eval_ao(cell, grids.coords, deriv=1)
ao0 = ao10[0]
ao1 = ni.eval_ao(cell, grids.coords)
self.assertTrue(numpy.allclose(ao0, ao1, atol=1e-9, rtol=1e-9))
self.assertAlmostEqual(finger(ao1), -0.54069672246407219, 8)
ao11 = ni.eval_ao(cell, grids.coords, deriv=1)
self.assertTrue(numpy.allclose(ao10, ao11, atol=1e-9, rtol=1e-9))
self.assertAlmostEqual(finger(ao11), 8.8004405892746433, 8)
ni.non0tab = ni.make_mask(cell, grids.coords)
ao1 = ni.eval_ao(cell, grids.coords)
self.assertTrue(numpy.allclose(ao0, ao1, atol=1e-9, rtol=1e-9))
self.assertAlmostEqual(finger(ao1), -0.54069672246407219, 8)
ao11 = ni.eval_ao(cell, grids.coords, deriv=1)
self.assertTrue(numpy.allclose(ao10, ao11, atol=1e-9, rtol=1e-9))
self.assertAlmostEqual(finger(ao11), 8.8004405892746433, 8)
ao11 = ni.eval_ao(cell, grids.coords, deriv=1, shls_slice=(3, 7))
self.assertTrue(
numpy.allclose(ao10[:, :, 6:17], ao11, atol=1e-9, rtol=1e-9))
def multi_grids_tasks(cell, verbose=None):
log = lib.logger.new_logger(cell, verbose)
tasks = []
a = cell.lattice_vectors()
neighbour_images = lib.cartesian_prod(([0, -1, 1], [0, -1, 1], [0, -1, 1]))
# Remove the first one which is the unit cell itself
neighbour_images0 = neighbour_images
neighbour_images = neighbour_images[1:]
neighbour_images = neighbour_images.dot(a)
b = numpy.linalg.inv(a.T)
heights = 1. / numpy.linalg.norm(b, axis=1)
normal_vector = b * heights.reshape(-1, 1)
distance_to_edge = cell.atom_coords().dot(normal_vector.T)
#FIXME: if atoms out of unit cell
#assert(numpy.all(distance_to_edge >= 0))
distance_to_edge = numpy.hstack(
[distance_to_edge, heights - distance_to_edge])
min_distance_to_edge = distance_to_edge.min(axis=1)
# Split shells based on rcut
rcuts_pgto, kecuts_pgto = _primitive_gto_cutoff(cell)
ao_loc = cell.ao_loc_nr()
def make_cell_high_exp(shls_high, r0, r1):
cell_high = copy.copy(cell)
cell_high._bas = cell._bas.copy()
cell_high._env = cell._env.copy()
rcut_atom = [0] * cell.natm
ke_cutoff = 0
for ib in shls_high:
rc = rcuts_pgto[ib]
idx = numpy.where((r1 <= rc) & (rc < r0))[0]
np1 = len(idx)
cs = cell._libcint_ctr_coeff(ib)
np, nc = cs.shape
if np1 < np: # no pGTO splitting within the shell
pexp = cell._bas[ib, PTR_EXP]
pcoeff = cell._bas[ib, PTR_COEFF]
cs1 = cs[idx]
cell_high._env[pcoeff:pcoeff + cs1.size] = cs1.T.ravel()
cell_high._env[pexp:pexp + np1] = cell.bas_exp(ib)[idx]
cell_high._bas[ib, NPRIM_OF] = np1
ke_cutoff = max(ke_cutoff, kecuts_pgto[ib][idx].max())
ia = cell.bas_atom(ib)
rcut_atom[ia] = max(rcut_atom[ia], rc[idx].max())
cell_high._bas = cell_high._bas[shls_high]
ao_idx = numpy.hstack(
[numpy.arange(ao_loc[i], ao_loc[i + 1]) for i in shls_high])
return cell_high, ao_idx, ke_cutoff, rcut_atom
def make_cell_low_exp(shls_low, r0, r1):
cell_low = copy.copy(cell)
cell_low._bas = cell._bas.copy()
cell_low._env = cell._env.copy()
for ib in shls_low:
idx = numpy.where(r0 <= rcuts_pgto[ib])[0]
np1 = len(idx)
cs = cell._libcint_ctr_coeff(ib)
np, nc = cs.shape
if np1 < np: # no pGTO splitting within the shell
pexp = cell._bas[ib, PTR_EXP]
pcoeff = cell._bas[ib, PTR_COEFF]
cs1 = cs[idx]
cell_low._env[pcoeff:pcoeff + cs1.size] = cs1.T.ravel()
cell_low._env[pexp:pexp + np1] = cell.bas_exp(ib)[idx]
cell_low._bas[ib, NPRIM_OF] = np1
cell_low._bas = cell_low._bas[shls_low]
ao_idx = numpy.hstack(
[numpy.arange(ao_loc[i], ao_loc[i + 1]) for i in shls_low])
return cell_low, ao_idx
nao = cell.nao_nr()
rmax = a.max() * RMAX_FACTOR
n_delimeter = int(numpy.log(0.01 / rmax) / numpy.log(RMAX_RATIO))
rcut_delimeter = rmax * (RMAX_RATIO**numpy.arange(n_delimeter))
for r0, r1 in zip(numpy.append(1e9, rcut_delimeter),
numpy.append(rcut_delimeter, 0)):
# shells which have high exps (small rcut)
shls_high = [
ib for ib, rc in enumerate(rcuts_pgto)
if numpy.any((r1 <= rc) & (rc < r0))
]
if len(shls_high) == 0:
continue
cell_high, ao_idx_high, ke_cutoff, rcut_atom = \
make_cell_high_exp(shls_high, r0, r1)
# shells which have low exps (big rcut)
shls_low = [
ib for ib, rc in enumerate(rcuts_pgto) if numpy.any(r0 <= rc)
]
if len(shls_low) == 0:
cell_low = None
ao_idx_low = []
else:
cell_low, ao_idx_low = make_cell_low_exp(shls_low, r0, r1)
mesh = tools.cutoff_to_mesh(a, ke_cutoff)
if TO_EVEN_GRIDS:
mesh = (mesh + 1) // 2 * 2 # to the nearest even number
if numpy.all(mesh >= cell.mesh):
# Including all rest shells
shls_high = [
ib for ib, rc in enumerate(rcuts_pgto) if numpy.any(rc < r0)
]
cell_high, ao_idx_high = make_cell_high_exp(shls_high, r0, 0)[:2]
cell_high.mesh = mesh = numpy.min([mesh, cell.mesh], axis=0)
coords = cell.gen_uniform_grids(mesh)
coords_f4 = coords.astype(numpy.float32)
ngrids = coords_f4.shape[0]
coords_idx = numpy.zeros(ngrids, dtype=bool)
gg = numpy.einsum('px,px->p', coords_f4, coords_f4)
Lg = (2 * neighbour_images.astype(numpy.float32)).dot(coords_f4.T)
for ia in set(cell_high._bas[:, ATOM_OF]):
rcut = rcut_atom[ia]
log.debug1(' atom %d rcut %g', mesh, ia, rcut)
atom_coord = cell.atom_coord(ia)
#dr = coords_f4 - atom_coord.astype(numpy.float32)
#coords_idx |= numpy.einsum('px,px->p', dr, dr) <= rcut**2
#optimized to
gg_ag = gg - coords_f4.dot(2 * atom_coord.astype(numpy.float32))
coords_idx |= gg_ag <= rcut**2 - atom_coord.dot(atom_coord)
if min_distance_to_edge[ia] > rcut:
# atom + rcut is completely inside the unit cell
continue
atoms_in_neighbour = neighbour_images + atom_coord
distance_to_unit_cell = atoms_in_neighbour.dot(normal_vector.T)
distance_to_unit_cell = numpy.hstack([
abs(distance_to_unit_cell),
abs(heights - distance_to_unit_cell)
])
idx = distance_to_unit_cell.min(axis=1) <= rcut
#for r_atom in atoms_in_neighbour[idx]:
# dr = coords_f4 - r_atom.astype(numpy.float32)
# coords_idx |= numpy.einsum('px,px->p', dr, dr) <= rcut**2
#optimized to
for i in numpy.where(idx)[0]:
L_a = atoms_in_neighbour[i]
coords_idx |= gg_ag - Lg[i] <= rcut**2 - L_a.dot(L_a)
coords = coords[coords_idx]
grids_high = gen_grid.UniformGrids(cell_high)
grids_high.coords = coords
grids_high.non0tab = grids_high.make_mask(cell_high, coords)
grids_high.coords_idx = coords_idx
grids_high.ao_idx = ao_idx_high
if cell_low is None:
grids_low = None
else:
grids_low = gen_grid.UniformGrids(cell_low)
grids_low.coords = coords
grids_low.non0tab = grids_low.make_mask(cell_low, coords)
grids_low.coords_idx = coords_idx
grids_low.ao_idx = ao_idx_low
log.debug('mesh %s nao all/high/low %d %d %d, ngrids %d', mesh, nao,
len(ao_idx_high), len(ao_idx_low), coords.shape[0])
tasks.append([grids_high, grids_low])
if numpy.all(mesh >= cell.mesh):
break
return tasks