from gpaw.symmetry import Symmetry from ase.dft.kpoints import monkhorst_pack # Primitive diamond lattice, with Si lattice parameter a = 5.475 cell_cv = .5 * a * np.array([(1, 1, 0), (1, 0, 1), (0, 1, 1)]) spos_ac = np.array([(.00, .00, .00), (.25, .25, .25)]) id_a = [1, 1] # Two identical atoms pbc_c = np.ones(3, bool) bzk_kc = monkhorst_pack((4, 4, 4)) # Do check symm = Symmetry(id_a, cell_cv, pbc_c, fractrans=False) symm.analyze(spos_ac) ibzk_kc, w_k = symm.reduce(bzk_kc)[:2] assert len(symm.op_scc) == 24 assert len(w_k) == 10 a = 3 / 32.; b = 1 / 32.; c = 6 / 32. assert np.all(w_k == [a, b, a, c, c, a, a, a, a, b]) assert not symm.op_scc.sum(0).any() # Rotate unit cell and check again: cell_cv = a / sqrt(2) * np.array([(1, 0, 0), (0.5, sqrt(3) / 2, 0), (0.5, sqrt(3) / 6, sqrt(2.0 / 3))]) symm = Symmetry(id_a, cell_cv, pbc_c, fractrans=False) symm.analyze(spos_ac) ibzkb_kc, wb_k = symm.reduce(bzk_kc)[:2] assert len(symm.op_scc) == 24 assert abs(w_k - wb_k).sum() < 1e-14
def get_realspace_hs(h_skmm, s_kmm, bzk_kc, weight_k, R_c=(0, 0, 0), direction='x', usesymm=None): from gpaw.symmetry import Symmetry from ase.dft.kpoints import get_monkhorst_pack_size_and_offset, \ monkhorst_pack if usesymm is True: raise NotImplementedError, 'Only None and False have been implemented' nspins, nk, nbf = h_skmm.shape[:3] dir = 'xyz'.index(direction) transverse_dirs = np.delete([0, 1, 2], [dir]) dtype = float if len(bzk_kc) > 1 or np.any(bzk_kc[0] != [0, 0, 0]): dtype = complex kpts_grid = get_monkhorst_pack_size_and_offset(bzk_kc)[0] # kpts in the transport direction nkpts_p = kpts_grid[dir] bzk_p_kc = monkhorst_pack((nkpts_p,1,1))[:, 0] weight_p_k = 1. / nkpts_p # kpts in the transverse directions bzk_t_kc = monkhorst_pack(tuple(kpts_grid[transverse_dirs]) + (1, )) if usesymm == False: #XXX a somewhat ugly hack: # By default GPAW reduces inversion sym in the z direction. The steps # below assure reduction in the transverse dirs. # For now this part seems to do the job, but it may be written # in a smarter way in the future. symmetry = Symmetry([1], np.eye(3)) symmetry.prune_symmetries(np.zeros((1, 3))) ibzk_kc, ibzweight_k = symmetry.reduce(bzk_kc)[:2] ibzk_t_kc, weights_t_k = symmetry.reduce(bzk_t_kc)[:2] ibzk_t_kc = ibzk_t_kc[:, :2] nkpts_t = len(ibzk_t_kc) else: # usesymm = None ibzk_kc = bzk_kc.copy() ibzk_t_kc = bzk_t_kc nkpts_t = len(bzk_t_kc) weights_t_k = [1. / nkpts_t for k in range(nkpts_t)] h_skii = np.zeros((nspins, nkpts_t, nbf, nbf), dtype) if s_kmm is not None: s_kii = np.zeros((nkpts_t, nbf, nbf), dtype) tol = 7 for j, k_t in enumerate(ibzk_t_kc): for k_p in bzk_p_kc: k = np.zeros((3,)) k[dir] = k_p k[transverse_dirs] = k_t kpoint_list = [list(np.round(k_kc, tol)) for k_kc in ibzk_kc] if list(np.round(k, tol)) not in kpoint_list: k = -k # inversion index = kpoint_list.index(list(np.round(k,tol))) h = h_skmm[:, index].conjugate() if s_kmm is not None: s = s_kmm[index].conjugate() k=-k else: # kpoint in the ibz index = kpoint_list.index(list(np.round(k, tol))) h = h_skmm[:, index] if s_kmm is not None: s = s_kmm[index] c_k = np.exp(2.j * np.pi * np.dot(k, R_c)) * weight_p_k h_skii[:, j] += c_k * h if s_kmm is not None: s_kii[j] += c_k * s if s_kmm is None: return ibzk_t_kc, weights_t_k, h_skii else: return ibzk_t_kc, weights_t_k, h_skii, s_kii
def get_realspace_hs(h_skmm, s_kmm, bzk_kc, weight_k, R_c=(0, 0, 0), direction='x', symmetry={'enabled': False}): from gpaw.symmetry import Symmetry from ase.dft.kpoints import get_monkhorst_pack_size_and_offset, \ monkhorst_pack if symmetry['point_group'] is True: raise NotImplementedError, 'Point group symmetry not implemented' nspins, nk, nbf = h_skmm.shape[:3] dir = 'xyz'.index(direction) transverse_dirs = np.delete([0, 1, 2], [dir]) dtype = float if len(bzk_kc) > 1 or np.any(bzk_kc[0] != [0, 0, 0]): dtype = complex kpts_grid = get_monkhorst_pack_size_and_offset(bzk_kc)[0] # kpts in the transport direction nkpts_p = kpts_grid[dir] bzk_p_kc = monkhorst_pack((nkpts_p,1,1))[:, 0] weight_p_k = 1. / nkpts_p # kpts in the transverse directions bzk_t_kc = monkhorst_pack(tuple(kpts_grid[transverse_dirs]) + (1, )) if not 'time_reversal' in symmetry: symmetry['time_reversal'] = True if symmetry['time_reversal'] is True: #XXX a somewhat ugly hack: # By default GPAW reduces inversion sym in the z direction. The steps # below assure reduction in the transverse dirs. # For now this part seems to do the job, but it may be written # in a smarter way in the future. symmetry = Symmetry([1], np.eye(3)) symmetry.prune_symmetries_atoms(np.zeros((1, 3))) ibzk_kc, ibzweight_k = symmetry.reduce(bzk_kc)[:2] ibzk_t_kc, weights_t_k = symmetry.reduce(bzk_t_kc)[:2] ibzk_t_kc = ibzk_t_kc[:, :2] nkpts_t = len(ibzk_t_kc) else: ibzk_kc = bzk_kc.copy() ibzk_t_kc = bzk_t_kc nkpts_t = len(bzk_t_kc) weights_t_k = [1. / nkpts_t for k in range(nkpts_t)] h_skii = np.zeros((nspins, nkpts_t, nbf, nbf), dtype) if s_kmm is not None: s_kii = np.zeros((nkpts_t, nbf, nbf), dtype) tol = 7 for j, k_t in enumerate(ibzk_t_kc): for k_p in bzk_p_kc: k = np.zeros((3,)) k[dir] = k_p k[transverse_dirs] = k_t kpoint_list = [list(np.round(k_kc, tol)) for k_kc in ibzk_kc] if list(np.round(k, tol)) not in kpoint_list: k = -k # inversion index = kpoint_list.index(list(np.round(k,tol))) h = h_skmm[:, index].conjugate() if s_kmm is not None: s = s_kmm[index].conjugate() k=-k else: # kpoint in the ibz index = kpoint_list.index(list(np.round(k, tol))) h = h_skmm[:, index] if s_kmm is not None: s = s_kmm[index] c_k = np.exp(2.j * np.pi * np.dot(k, R_c)) * weight_p_k h_skii[:, j] += c_k * h if s_kmm is not None: s_kii[j] += c_k * s if s_kmm is None: return ibzk_t_kc, weights_t_k, h_skii else: return ibzk_t_kc, weights_t_k, h_skii, s_kii
from gpaw.symmetry import Symmetry from ase.dft.kpoints import monkhorst_pack # Primitive diamond lattice, with Si lattice parameter a = 5.475 cell_cv = .5 * a * np.array([(1, 1, 0), (1, 0, 1), (0, 1, 1)]) spos_ac = np.array([(.00, .00, .00), (.25, .25, .25)]) id_a = [1, 1] # Two identical atoms pbc_c = np.ones(3, bool) bzk_kc = monkhorst_pack((4, 4, 4)) # Do check symm = Symmetry(id_a, cell_cv, pbc_c, symmorphic=False) symm.analyze(spos_ac) ibzk_kc, w_k = symm.reduce(bzk_kc)[:2] assert len(symm.op_scc) == 48 assert len(w_k) == 10 a = 3 / 32. b = 1 / 32. c = 6 / 32. assert np.all(w_k == [a, b, a, c, c, a, a, a, a, b]) assert not symm.op_scc.sum(0).any() # Rotate unit cell and check again: cell_cv = a / sqrt(2) * np.array([(1, 0, 0), (0.5, sqrt(3) / 2, 0), (0.5, sqrt(3) / 6, sqrt(2.0 / 3))]) symm = Symmetry(id_a, cell_cv, pbc_c, symmorphic=False) symm.analyze(spos_ac) ibzkb_kc, wb_k = symm.reduce(bzk_kc)[:2] assert len(symm.op_scc) == 48
class KPointDescriptor: """Descriptor-class for k-points.""" def __init__(self, kpts, nspins=1, collinear=True, usefractrans=False): """Construct descriptor object for kpoint/spin combinations (ks-pair). Parameters ---------- kpts: None, sequence of 3 ints, or (n,3)-shaped array Specification of the k-point grid. None=Gamma, list of ints=Monkhorst-Pack, ndarray=user specified. nspins: int Number of spins. usefractrans: bool Switch for the use of non-symmorphic symmetries aka: symmetries with fractional translations. False by default (experimental!!!) Attributes =================== ================================================= ``N_c`` Number of k-points in the different directions. ``nspins`` Number of spins in total. ``mynspins`` Number of spins on this CPU. ``nibzkpts`` Number of irreducible kpoints in 1st BZ. ``nks`` Number of k-point/spin combinations in total. ``mynks`` Number of k-point/spin combinations on this CPU. ``gamma`` Boolean indicator for gamma point calculation. ``comm`` MPI-communicator for kpoint distribution. ``weight_k`` Weights of each k-point ``ibzk_kc`` Unknown ``sym_k`` Unknown ``time_reversal_k`` Unknown ``bz2ibz_k`` Unknown ``ibz2bz_k`` Unknown ``bz2bz_ks`` Unknown ``symmetry`` Object representing symmetries =================== ================================================= """ if kpts is None: self.bzk_kc = np.zeros((1, 3)) self.N_c = np.array((1, 1, 1), dtype=int) self.offset_c = np.zeros(3) elif isinstance(kpts[0], int): self.bzk_kc = monkhorst_pack(kpts) self.N_c = np.array(kpts, dtype=int) self.offset_c = np.zeros(3) else: self.bzk_kc = np.array(kpts, float) try: self.N_c, self.offset_c = \ get_monkhorst_pack_size_and_offset(self.bzk_kc) except ValueError: self.N_c = None self.offset_c = None self.collinear = collinear self.nspins = nspins self.nbzkpts = len(self.bzk_kc) # Gamma-point calculation? self.usefractrans = usefractrans self.gamma = (self.nbzkpts == 1 and np.allclose(self.bzk_kc[0], 0.0)) self.set_symmetry(None, None, usesymm=None) self.set_communicator(mpi.serial_comm) if self.gamma: self.description = '1 k-point (Gamma)' else: self.description = '%d k-points' % self.nbzkpts if self.N_c is not None: self.description += (': %d x %d x %d Monkhorst-Pack grid' % tuple(self.N_c)) if self.offset_c.any(): self.description += ' + [' for x in self.offset_c: if x != 0 and abs(round(1 / x) - 1 / x) < 1e-12: self.description += '1/%d,' % round(1 / x) else: self.description += '%f,' % x self.description = self.description[:-1] + ']' def __len__(self): """Return number of k-point/spin combinations of local CPU.""" return self.mynks def set_symmetry(self, atoms, setups, magmom_av=None, usesymm=False, N_c=None, comm=None): """Create symmetry object and construct irreducible Brillouin zone. atoms: Atoms object Defines atom positions and types and also unit cell and boundary conditions. setups: instance of class Setups PAW setups for the atoms. magmom_av: ndarray Initial magnetic moments. usesymm: bool Symmetry flag. N_c: three int's or None If not None: Check also symmetry of grid. """ if atoms is not None: for c, periodic in enumerate(atoms.pbc): if not periodic and not np.allclose(self.bzk_kc[:, c], 0.0): raise ValueError('K-points can only be used with PBCs!') self.cell_cv = atoms.cell / Bohr if magmom_av is None: magmom_av = np.zeros((len(atoms), 3)) magmom_av[:, 2] = atoms.get_initial_magnetic_moments() magmom_av = magmom_av.round(decimals=3) # round off id_a = zip(setups.id_a, *magmom_av.T) # Construct a Symmetry instance containing the identity operation # only self.symmetry = Symmetry(id_a, atoms.cell / Bohr, atoms.pbc, fractrans=self.usefractrans) self.usefractrans = self.symmetry.usefractrans else: self.symmetry = None if self.gamma or usesymm is None: # Point group and time-reversal symmetry neglected self.weight_k = np.ones(self.nbzkpts) / self.nbzkpts self.ibzk_kc = self.bzk_kc.copy() self.sym_k = np.zeros(self.nbzkpts, int) self.time_reversal_k = np.zeros(self.nbzkpts, bool) self.bz2ibz_k = np.arange(self.nbzkpts) self.ibz2bz_k = np.arange(self.nbzkpts) self.bz2bz_ks = np.arange(self.nbzkpts)[:, np.newaxis] else: if usesymm: # Find symmetry operations of atoms self.symmetry.analyze(atoms.get_scaled_positions()) if N_c is not None: if self.usefractrans: ## adjust N_c to symmetries # the factor (denominator) the grid must follow factor = np.ones(3, float) indexes = np.where(np.abs(self.symmetry.ft_sc) > 1e-3) for i in range(len(indexes[0])): # find smallest common denominator a = factor[indexes[1][i]] b = np.rint(1. / self.symmetry.ft_sc[indexes[0][i]][indexes[1][i]]) factor[indexes[1][i]] = a * b while b != 0: rem = a % b a = b b = rem factor[indexes[1][i]] /= a Nnew_c = np.array(np.rint(N_c / factor) * factor, int) # make sure new grid is not less dense Nnew_c = np.array(np.where(Nnew_c >= N_c, Nnew_c, Nnew_c + factor), int) N_c = Nnew_c else: ## adjust symmetries to grid self.symmetry.prune_symmetries_grid(N_c) (self.ibzk_kc, self.weight_k, self.sym_k, self.time_reversal_k, self.bz2ibz_k, self.ibz2bz_k, self.bz2bz_ks) = self.symmetry.reduce(self.bzk_kc, comm) if setups is not None: setups.set_symmetry(self.symmetry) # Number of irreducible k-points and k-point/spin combinations. self.nibzkpts = len(self.ibzk_kc) if self.collinear: self.nks = self.nibzkpts * self.nspins else: self.nks = self.nibzkpts return N_c def set_communicator(self, comm): """Set k-point communicator.""" # Ranks < self.rank0 have mynks0 k-point/spin combinations and # ranks >= self.rank0 have mynks0+1 k-point/spin combinations. mynks0, x = divmod(self.nks, comm.size) self.rank0 = comm.size - x self.comm = comm # My number and offset of k-point/spin combinations self.mynks = self.get_count() self.ks0 = self.get_offset() if self.nspins == 2 and comm.size == 1: # NCXXXXXXXX # Avoid duplicating k-points in local list of k-points. self.ibzk_qc = self.ibzk_kc.copy() self.weight_q = self.weight_k else: self.ibzk_qc = np.vstack((self.ibzk_kc, self.ibzk_kc))[self.get_slice()] self.weight_q = np.hstack((self.weight_k, self.weight_k))[self.get_slice()] def copy(self, comm=mpi.serial_comm): """Create a copy with shared symmetry object.""" kd = KPointDescriptor(self.bzk_kc, self.nspins) kd.weight_k = self.weight_k kd.ibzk_kc = self.ibzk_kc kd.sym_k = self.sym_k kd.time_reversal_k = self.time_reversal_k kd.bz2ibz_k = self.bz2ibz_k kd.ibz2bz_k = self.ibz2bz_k kd.bz2bz_ks = self.bz2bz_ks kd.symmetry = self.symmetry kd.nibzkpts = self.nibzkpts kd.nks = self.nks kd.set_communicator(comm) return kd def create_k_points(self, gd): """Return a list of KPoints.""" sdisp_cd = gd.sdisp_cd kpt_u = [] for ks in range(self.ks0, self.ks0 + self.mynks): s, k = divmod(ks, self.nibzkpts) q = (ks - self.ks0) % self.nibzkpts if self.collinear: weight = self.weight_k[k] * 2 / self.nspins else: weight = self.weight_k[k] if self.gamma: phase_cd = np.ones((3, 2), complex) else: phase_cd = np.exp(2j * np.pi * sdisp_cd * self.ibzk_kc[k, :, np.newaxis]) kpt_u.append(KPoint(weight, s, k, q, phase_cd)) return kpt_u def collect(self, a_ux, broadcast=True): """Collect distributed data to all.""" if self.comm.rank == 0 or broadcast: xshape = a_ux.shape[1:] a_skx = np.empty((self.nspins, self.nibzkpts) + xshape, a_ux.dtype) a_Ux = a_skx.reshape((-1,) + xshape) else: a_skx = None if self.comm.rank > 0: self.comm.send(a_ux, 0) else: u1 = self.get_count(0) a_Ux[0:u1] = a_ux requests = [] for rank in range(1, self.comm.size): u2 = u1 + self.get_count(rank) requests.append(self.comm.receive(a_Ux[u1:u2], rank, block=False)) u1 = u2 assert u1 == len(a_Ux) self.comm.waitall(requests) if broadcast: self.comm.broadcast(a_Ux, 0) return a_skx def transform_wave_function(self, psit_G, k, index_G=None, phase_G=None): """Transform wave function from IBZ to BZ. k is the index of the desired k-point in the full BZ. """ s = self.sym_k[k] time_reversal = self.time_reversal_k[k] op_cc = np.linalg.inv(self.symmetry.op_scc[s]).round().astype(int) # Identity if (np.abs(op_cc - np.eye(3, dtype=int)) < 1e-10).all(): if time_reversal: return psit_G.conj() else: return psit_G # General point group symmetry else: ik = self.bz2ibz_k[k] kibz_c = self.ibzk_kc[ik] b_g = np.zeros_like(psit_G) kbz_c = np.dot(self.symmetry.op_scc[s], kibz_c) if index_G is not None: assert index_G.shape == psit_G.shape == phase_G.shape,\ 'Shape mismatch %s vs %s vs %s' % (index_G.shape, psit_G.shape, phase_G.shape) _gpaw.symmetrize_with_index(psit_G, b_g, index_G, phase_G) else: _gpaw.symmetrize_wavefunction(psit_G, b_g, op_cc.copy(), np.ascontiguousarray(kibz_c), kbz_c) if time_reversal: return b_g.conj() else: return b_g def get_transform_wavefunction_index(self, nG, k): """Get the "wavefunction transform index". This is a permutation of the numbers 1, 2, .. N which associates k + q to some k, and where N is the total number of grid points as specified by nG which is a 3D tuple. Returns index_G and phase_G which are one-dimensional arrays on the grid.""" s = self.sym_k[k] op_cc = np.linalg.inv(self.symmetry.op_scc[s]).round().astype(int) # General point group symmetry if (np.abs(op_cc - np.eye(3, dtype=int)) < 1e-10).all(): nG0 = np.prod(nG) index_G = np.arange(nG0).reshape(nG) phase_G = np.ones(nG) else: ik = self.bz2ibz_k[k] kibz_c = self.ibzk_kc[ik] index_G = np.zeros(nG, dtype=int) phase_G = np.zeros(nG, dtype=complex) kbz_c = np.dot(self.symmetry.op_scc[s], kibz_c) _gpaw.symmetrize_return_index(index_G, phase_G, op_cc.copy(), np.ascontiguousarray(kibz_c), kbz_c) return index_G, phase_G #def find_k_plus_q(self, q_c, k_x=None): def find_k_plus_q(self, q_c, kpts_k=None): """Find the indices of k+q for all kpoints in the Brillouin zone. In case that k+q is outside the BZ, the k-point inside the BZ corresponding to k+q is given. Parameters ---------- q_c: ndarray Coordinates for the q-vector in units of the reciprocal lattice vectors. kpts_k: list of ints Restrict search to specified k-points. """ k_x = kpts_k if k_x is None: return self.find_k_plus_q(q_c, range(self.nbzkpts)) i_x = [] for k in k_x: kpt_c = self.bzk_kc[k] + q_c d_kc = kpt_c - self.bzk_kc d_k = abs(d_kc - d_kc.round()).sum(1) i = d_k.argmin() if d_k[i] > 1e-8: raise RuntimeError('Could not find k+q!') i_x.append(i) return i_x def get_bz_q_points(self, first=False): """Return the q=k1-k2. q-mesh is always Gamma-centered.""" shift_c = 0.5 * ((self.N_c + 1) % 2) / self.N_c bzq_qc = monkhorst_pack(self.N_c) + shift_c if first: return to1bz(bzq_qc, self.cell_cv) else: return bzq_qc def get_ibz_q_points(self, bzq_qc, op_scc): """Return ibz q points and the corresponding symmetry operations that work for k-mesh as well.""" ibzq_qc_tmp = [] ibzq_qc_tmp.append(bzq_qc[-1]) weight_tmp = [0] for i, op_cc in enumerate(op_scc): if np.abs(op_cc - np.eye(3)).sum() < 1e-8: identity_iop = i break ibzq_q_tmp = {} iop_q = {} timerev_q = {} diff_qc = {} for i in range(len(bzq_qc) - 1, -1, -1): # loop opposite to kpoint try: ibzk, iop, timerev, diff_c = self.find_ibzkpt( op_scc, ibzq_qc_tmp, bzq_qc[i]) find = False for ii, iop1 in enumerate(self.sym_k): if iop1 == iop and self.time_reversal_k[ii] == timerev: find = True break if find is False: raise ValueError('cant find k!') ibzq_q_tmp[i] = ibzk weight_tmp[ibzk] += 1. iop_q[i] = iop timerev_q[i] = timerev diff_qc[i] = diff_c except ValueError: ibzq_qc_tmp.append(bzq_qc[i]) weight_tmp.append(1.) ibzq_q_tmp[i] = len(ibzq_qc_tmp) - 1 iop_q[i] = identity_iop timerev_q[i] = False diff_qc[i] = np.zeros(3) # reverse the order. nq = len(ibzq_qc_tmp) ibzq_qc = np.zeros((nq, 3)) ibzq_q = np.zeros(len(bzq_qc), dtype=int) for i in range(nq): ibzq_qc[i] = ibzq_qc_tmp[nq - i - 1] for i in range(len(bzq_qc)): ibzq_q[i] = nq - ibzq_q_tmp[i] - 1 self.q_weights = np.array(weight_tmp[::-1]) / len(bzq_qc) return ibzq_qc, ibzq_q, iop_q, timerev_q, diff_qc def find_ibzkpt(self, symrel, ibzk_kc, bzk_c): """Find index in IBZ and related symmetry operations.""" find = False ibzkpt = 0 iop = 0 timerev = False for sign in (1, -1): for ioptmp, op in enumerate(symrel): for i, ibzk in enumerate(ibzk_kc): diff_c = bzk_c - sign * np.dot(op, ibzk) if (np.abs(diff_c - diff_c.round()) < 1e-8).all(): ibzkpt = i iop = ioptmp find = True if sign == -1: timerev = True break if find == True: break if find == True: break if find == False: raise ValueError('Cant find corresponding IBZ kpoint!') return ibzkpt, iop, timerev, diff_c.round() def where_is_q(self, q_c, bzq_qc): """Find the index of q points in BZ.""" d_qc = q_c - bzq_qc d_q = abs(d_qc - d_qc.round()).sum(1) q = d_q.argmin() if d_q[q] > 1e-8: raise RuntimeError('Could not find q!') return q def get_count(self, rank=None): """Return the number of ks-pairs which belong to a given rank.""" if rank is None: rank = self.comm.rank assert rank in xrange(self.comm.size) mynks0 = self.nks // self.comm.size mynks = mynks0 if rank >= self.rank0: mynks += 1 return mynks def get_offset(self, rank=None): """Return the offset of the first ks-pair on a given rank.""" if rank is None: rank = self.comm.rank assert rank in xrange(self.comm.size) mynks0 = self.nks // self.comm.size ks0 = rank * mynks0 if rank >= self.rank0: ks0 += rank - self.rank0 return ks0 def get_rank_and_index(self, s, k): """Find rank and local index of k-point/spin combination.""" u = self.where_is(s, k) rank, myu = self.who_has(u) return rank, myu def get_slice(self, rank=None): """Return the slice of global ks-pairs which belong to a given rank.""" if rank is None: rank = self.comm.rank assert rank in xrange(self.comm.size) mynks, ks0 = self.get_count(rank), self.get_offset(rank) uslice = slice(ks0, ks0 + mynks) return uslice def get_indices(self, rank=None): """Return the global ks-pair indices which belong to a given rank.""" uslice = self.get_slice(rank) return np.arange(*uslice.indices(self.nks)) def get_ranks(self): """Return array of ranks as a function of global ks-pair indices.""" ranks = np.empty(self.nks, dtype=int) for rank in range(self.comm.size): uslice = self.get_slice(rank) ranks[uslice] = rank assert (ranks >= 0).all() and (ranks < self.comm.size).all() return ranks def who_has(self, u): """Convert global index to rank information and local index.""" mynks0 = self.nks // self.comm.size if u < mynks0 * self.rank0: rank, myu = divmod(u, mynks0) else: rank, myu = divmod(u - mynks0 * self.rank0, mynks0 + 1) rank += self.rank0 return rank, myu def global_index(self, myu, rank=None): """Convert rank information and local index to global index.""" if rank is None: rank = self.comm.rank assert rank in xrange(self.comm.size) ks0 = self.get_offset(rank) u = ks0 + myu return u def what_is(self, u): """Convert global index to corresponding kpoint/spin combination.""" s, k = divmod(u, self.nibzkpts) return s, k def where_is(self, s, k): """Convert kpoint/spin combination to the global index thereof.""" u = k + self.nibzkpts * s return u
class KPointDescriptor: """Descriptor-class for k-points.""" def __init__(self, kpts, nspins): """Construct descriptor object for kpoint/spin combinations (ks-pair). Parameters ---------- kpts: None, list of ints, or ndarray Specification of the k-point grid. None=Gamma, list of ints=Monkhorst-Pack, ndarray=user specified. nspins: int Number of spins. Attributes ============ ====================================================== ``N_c`` Number of k-points in the different directions. ``nspins`` Number of spins. ``nibzkpts`` Number of irreducible kpoints in 1st Brillouin zone. ``nks`` Number of k-point/spin combinations in total. ``mynks`` Number of k-point/spin combinations on this CPU. ``gamma`` Boolean indicator for gamma point calculation. ``comm`` MPI-communicator for kpoint distribution. ============ ====================================================== """ if kpts is None: self.bzk_kc = np.zeros((1, 3)) self.N_c = np.array((1, 1, 1), dtype=int) elif isinstance(kpts[0], int): self.bzk_kc = monkhorst_pack(kpts) self.N_c = np.array(kpts, dtype=int) else: self.bzk_kc = np.array(kpts) self.N_c = None self.nspins = nspins self.nbzkpts = len(self.bzk_kc) # Gamma-point calculation self.gamma = self.nbzkpts == 1 and not self.bzk_kc[0].any() self.symmetry = None self.comm = None self.ibzk_kc = None self.weight_k = None self.nibzkpts = None self.rank0 = None self.mynks = None self.ks0 = None self.ibzk_qc = None def __len__(self): """Return number of k-point/spin combinations of local CPU.""" return self.mynks def set_symmetry(self, atoms, setups, usesymm, N_c=None): """Create symmetry object and construct irreducible Brillouin zone. Parameters ---------- atoms: Atoms object Defines atom positions and types and also unit cell and boundary conditions. setups: instance of class Setups PAW setups for the atoms. usesymm: bool Symmetry flag. N_c: three int's or None If not None: Check also symmetry of grid. """ if (~atoms.pbc & self.bzk_kc.any(0)).any(): raise ValueError('K-points can only be used with PBCs!') # Construct a Symmetry instance containing the identity operation # only # Round off magmom_a = atoms.get_initial_magnetic_moments().round(decimals=3) id_a = zip(magmom_a, setups.id_a) self.symmetry = Symmetry(id_a, atoms.cell / Bohr, atoms.pbc) if self.gamma or usesymm is None: # Point group and time-reversal symmetry neglected nkpts = len(self.bzk_kc) self.weight_k = np.ones(nkpts) / nkpts self.ibzk_kc = self.bzk_kc.copy() self.sym_k = np.zeros(nkpts) self.time_reversal_k = np.zeros(nkpts, bool) self.kibz_k = np.arange(nkpts) else: if usesymm: # Find symmetry operations of atoms self.symmetry.analyze(atoms.get_scaled_positions()) if N_c is not None: self.symmetry.prune_symmetries_grid(N_c) (self.ibzk_kc, self.weight_k, self.sym_k, self.time_reversal_k, self.kibz_k) = self.symmetry.reduce(self.bzk_kc) setups.set_symmetry(self.symmetry) # Number of irreducible k-points and k-point/spin combinations. self.nibzkpts = len(self.ibzk_kc) self.nks = self.nibzkpts * self.nspins def set_communicator(self, comm): """Set k-point communicator.""" # Ranks < self.rank0 have mynks0 k-point/spin combinations and # ranks >= self.rank0 have mynks0+1 k-point/spin combinations. mynks0, x = divmod(self.nks, comm.size) self.rank0 = comm.size - x self.comm = comm # My number and offset of k-point/spin combinations self.mynks, self.ks0 = self.get_count(), self.get_offset() if self.nspins == 2 and comm.size == 1: # Avoid duplicating k-points in local list of k-points. self.ibzk_qc = self.ibzk_kc.copy() else: self.ibzk_qc = np.vstack((self.ibzk_kc, self.ibzk_kc))[self.get_slice()] def create_k_points(self, gd): """Return a list of KPoints.""" sdisp_cd = gd.sdisp_cd kpt_u = [] for ks in range(self.ks0, self.ks0 + self.mynks): s, k = divmod(ks, self.nibzkpts) q = (ks - self.ks0) % self.nibzkpts weight = self.weight_k[k] * 2 / self.nspins if self.gamma: phase_cd = np.ones((3, 2), complex) else: phase_cd = np.exp(2j * np.pi * sdisp_cd * self.ibzk_kc[k, :, np.newaxis]) kpt_u.append(KPoint(weight, s, k, q, phase_cd)) return kpt_u def transform_wave_function(self, psit_G, k): """Transform wave function from IBZ to BZ. k is the index of the desired k-point in the full BZ.""" s = self.sym_k[k] time_reversal = self.time_reversal_k[k] op_cc = np.linalg.inv(self.symmetry.op_scc[s]).round().astype(int) # Identity if (np.abs(op_cc - np.eye(3, dtype=int)) < 1e-10).all(): if time_reversal: return psit_G.conj() else: return psit_G # Inversion symmetry elif (np.abs(op_cc + np.eye(3, dtype=int)) < 1e-10).all(): return psit_G.conj() # General point group symmetry else: ik = self.kibz_k[k] kibz_c = self.ibzk_kc[ik] kbz_c = self.bzk_kc[k] import _gpaw b_g = np.zeros_like(psit_G) if time_reversal: # assert abs(np.dot(op_cc, kibz_c) - -kbz_c) < tol _gpaw.symmetrize_wavefunction(psit_G, b_g, op_cc.copy(), kibz_c, -kbz_c) return b_g.conj() else: # assert abs(np.dot(op_cc, kibz_c) - kbz_c) < tol _gpaw.symmetrize_wavefunction(psit_G, b_g, op_cc.copy(), kibz_c, kbz_c) return b_g def find_k_plus_q(self, q_c): """Find the indices of k+q for all kpoints in the Brillouin zone. In case that k+q is outside the BZ, the k-point inside the BZ corresponding to k+q is given. Parameters ---------- q_c: ndarray Coordinates for the q-vector in units of the reciprocal lattice vectors. """ # Monkhorst-pack grid if self.N_c is not None: N_c = self.N_c dk_c = 1. / N_c kmax_c = (N_c - 1) * dk_c / 2. N = np.zeros(3, dtype=int) # k+q vectors kplusq_kc = self.bzk_kc + q_c # Translate back into the first BZ kplusq_kc[np.where(kplusq_kc > 0.5)] -= 1. kplusq_kc[np.where(kplusq_kc <= -0.5)] += 1. # List of k+q indices kplusq_k = [] # Find index of k+q vector in the bzk_kc attribute for kplusq, kplusq_c in enumerate(kplusq_kc): # Calculate index for Monkhorst-Pack grids if self.N_c is not None: N = np.asarray(np.round((kplusq_c + kmax_c) / dk_c), dtype=int) kplusq_k.append(N[2] + N[1] * N_c[2] + N[0] * N_c[2] * N_c[1]) else: k = np.argmin(np.sum(np.abs(self.bzk_kc - kplusq_c), axis=1)) kplusq_k.append(k) # Check the k+q vector index k_c = self.bzk_kc[kplusq_k[kplusq]] assert abs(kplusq_c - k_c).sum() < 1e-8, "Could not find k+q!" return kplusq_k def get_bz_q_points(self): """Return the q=k1-k2.""" bzk_kc = self.bzk_kc # Get all q-points all_qs = [] for k1 in bzk_kc: for k2 in bzk_kc: all_qs.append(k1-k2) all_qs = np.array(all_qs) # Fold q-points into Brillouin zone all_qs[np.where(all_qs > 0.501)] -= 1. all_qs[np.where(all_qs < -0.499)] += 1. # Make list of non-identical q-points in full BZ bz_qs = [all_qs[0]] for q_a in all_qs: q_in_list = False for q_b in bz_qs: if (abs(q_a[0]-q_b[0]) < 0.01 and abs(q_a[1]-q_b[1]) < 0.01 and abs(q_a[2]-q_b[2]) < 0.01): q_in_list = True break if q_in_list == False: bz_qs.append(q_a) self.bzq_kc = bz_qs return def where_is_q(self, q_c): """Find the index of q points.""" q_c[np.where(q_c>0.499)] -= 1 q_c[np.where(q_c<-0.499)] += 1 found = False for ik in range(self.nbzkpts): if (np.abs(self.bzq_kc[ik] - q_c) < 1e-8).all(): found = True return ik break if found is False: print self.bzq_kc, q_c raise ValueError('q-points can not be found!') def get_count(self, rank=None): """Return the number of ks-pairs which belong to a given rank.""" if rank is None: rank = self.comm.rank assert rank in xrange(self.comm.size) mynks0 = self.nks // self.comm.size mynks = mynks0 if rank >= self.rank0: mynks += 1 return mynks def get_offset(self, rank=None): """Return the offset of the first ks-pair on a given rank.""" if rank is None: rank = self.comm.rank assert rank in xrange(self.comm.size) mynks0 = self.nks // self.comm.size ks0 = rank * mynks0 if rank >= self.rank0: ks0 += rank - self.rank0 return ks0 def get_rank_and_index(self, s, k): """Find rank and local index of k-point/spin combination.""" u = self.where_is(s, k) rank, myu = self.who_has(u) return rank, myu def get_slice(self, rank=None): """Return the slice of global ks-pairs which belong to a given rank.""" if rank is None: rank = self.comm.rank assert rank in xrange(self.comm.size) mynks, ks0 = self.get_count(rank), self.get_offset(rank) uslice = slice(ks0, ks0 + mynks) return uslice def get_indices(self, rank=None): """Return the global ks-pair indices which belong to a given rank.""" uslice = self.get_slice(rank) return np.arange(*uslice.indices(self.nks)) def get_ranks(self): """Return array of ranks as a function of global ks-pair indices.""" ranks = np.empty(self.nks, dtype=int) for rank in range(self.comm.size): uslice = self.get_slice(rank) ranks[uslice] = rank assert (ranks >= 0).all() and (ranks < self.comm.size).all() return ranks def who_has(self, u): """Convert global index to rank information and local index.""" mynks0 = self.nks // self.comm.size if u < mynks0 * self.rank0: rank, myu = divmod(u, mynks0) else: rank, myu = divmod(u - mynks0 * self.rank0, mynks0 + 1) rank += self.rank0 return rank, myu def global_index(self, myu, rank=None): """Convert rank information and local index to global index.""" if rank is None: rank = self.comm.rank assert rank in xrange(self.comm.size) ks0 = self.get_offset(rank) u = ks0 + myu return u def what_is(self, u): """Convert global index to corresponding kpoint/spin combination.""" s, k = divmod(u, self.nibzkpts) return s, k def where_is(self, s, k): """Convert kpoint/spin combination to the global index thereof.""" u = k + self.nibzkpts * s return u