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