def get_projections(self, locfun, spin=0): """Project wave functions onto localized functions Determine the projections of the Kohn-Sham eigenstates onto specified localized functions of the format:: locfun = [[spos_c, l, sigma], [...]] spos_c can be an atom index, or a scaled position vector. l is the angular momentum, and sigma is the (half-) width of the radial gaussian. Return format is:: f_kni = <psi_kn | f_i> where psi_kn are the wave functions, and f_i are the specified localized functions. As a special case, locfun can be the string 'projectors', in which case the bound state projectors are used as localized functions. """ wfs = self.wfs if locfun == 'projectors': f_kin = [] for kpt in wfs.kpt_u: if kpt.s == spin: f_in = [] for a, P_ni in kpt.P_ani.items(): i = 0 setup = wfs.setups[a] for l, n in zip(setup.l_j, setup.n_j): if n >= 0: for j in range(i, i + 2 * l + 1): f_in.append(P_ni[:, j]) i += 2 * l + 1 f_kin.append(f_in) f_kni = np.array(f_kin).transpose(0, 2, 1) return f_kni.conj() from gpaw.lfc import LocalizedFunctionsCollection as LFC from gpaw.spline import Spline from gpaw.utilities import _fact nkpts = len(wfs.ibzk_kc) nbf = np.sum([2 * l + 1 for pos, l, a in locfun]) f_kni = np.zeros((nkpts, wfs.nbands, nbf), wfs.dtype) spos_ac = self.atoms.get_scaled_positions() % 1.0 spos_xc = [] splines_x = [] for spos_c, l, sigma in locfun: if isinstance(spos_c, int): spos_c = spos_ac[spos_c] spos_xc.append(spos_c) alpha = .5 * Bohr**2 / sigma**2 r = np.linspace(0, 6. * sigma, 500) f_g = (_fact[l] * (4 * alpha)**(l + 3 / 2.) * np.exp(-alpha * r**2) / (np.sqrt(4 * np.pi) * _fact[2 * l + 1])) splines_x.append([Spline(l, rmax=r[-1], f_g=f_g, points=61)]) lf = LFC(wfs.gd, splines_x, wfs.kpt_comm, dtype=wfs.dtype) if not wfs.gamma: lf.set_k_points(wfs.ibzk_qc) lf.set_positions(spos_xc) k = 0 f_ani = lf.dict(wfs.nbands) for kpt in wfs.kpt_u: if kpt.s != spin: continue lf.integrate(kpt.psit_nG[:], f_ani, kpt.q) i1 = 0 for x, f_ni in f_ani.items(): i2 = i1 + f_ni.shape[1] f_kni[k, :, i1:i2] = f_ni i1 = i2 k += 1 return f_kni.conj()
class FDWaveFunctions(FDPWWaveFunctions): def __init__(self, stencil, diagksl, orthoksl, initksl, gd, nvalence, setups, bd, dtype, world, kd, timer=None): FDPWWaveFunctions.__init__(self, diagksl, orthoksl, initksl, gd, nvalence, setups, bd, dtype, world, kd, timer) # Kinetic energy operator: self.kin = Laplace(self.gd, -0.5, stencil, self.dtype) self.matrixoperator = MatrixOperator(self.orthoksl) self.taugrad_v = None # initialized by MGGA functional def empty(self, n=(), global_array=False, realspace=False, q=-1): return self.gd.empty(n, self.dtype, global_array) def integrate(self, a_xg, b_yg=None, global_integral=True): return self.gd.integrate(a_xg, b_yg, global_integral) def bytes_per_wave_function(self): return self.gd.bytecount(self.dtype) def set_setups(self, setups): self.pt = LFC(self.gd, [setup.pt_j for setup in setups], self.kd, dtype=self.dtype, forces=True) FDPWWaveFunctions.set_setups(self, setups) def set_positions(self, spos_ac): FDPWWaveFunctions.set_positions(self, spos_ac) def summary(self, fd): fd.write('Wave functions: Uniform real-space grid\n') fd.write('Kinetic energy operator: %s\n' % self.kin.description) def make_preconditioner(self, block=1): return Preconditioner(self.gd, self.kin, self.dtype, block) def apply_pseudo_hamiltonian(self, kpt, hamiltonian, psit_xG, Htpsit_xG): self.timer.start('Apply hamiltonian') self.kin.apply(psit_xG, Htpsit_xG, kpt.phase_cd) hamiltonian.apply_local_potential(psit_xG, Htpsit_xG, kpt.s) self.timer.stop('Apply hamiltonian') def add_orbital_density(self, nt_G, kpt, n): if self.dtype == float: axpy(1.0, kpt.psit_nG[n]**2, nt_G) else: axpy(1.0, kpt.psit_nG[n].real**2, nt_G) axpy(1.0, kpt.psit_nG[n].imag**2, nt_G) def add_to_density_from_k_point_with_occupation(self, nt_sG, kpt, f_n): # Used in calculation of response part of GLLB-potential nt_G = nt_sG[kpt.s] if self.dtype == float: for f, psit_G in zip(f_n, kpt.psit_nG): axpy(f, psit_G**2, nt_G) else: for f, psit_G in zip(f_n, kpt.psit_nG): axpy(f, psit_G.real**2, nt_G) axpy(f, psit_G.imag**2, nt_G) # Hack used in delta-scf calculations: if hasattr(kpt, 'c_on'): assert self.bd.comm.size == 1 d_nn = np.zeros((self.bd.mynbands, self.bd.mynbands), dtype=complex) for ne, c_n in zip(kpt.ne_o, kpt.c_on): d_nn += ne * np.outer(c_n.conj(), c_n) for d_n, psi0_G in zip(d_nn, kpt.psit_nG): for d, psi_G in zip(d_n, kpt.psit_nG): if abs(d) > 1.e-12: nt_G += (psi0_G.conj() * d * psi_G).real def calculate_kinetic_energy_density(self): if self.taugrad_v is None: self.taugrad_v = [ Gradient(self.gd, v, n=3, dtype=self.dtype).apply for v in range(3)] assert not hasattr(self.kpt_u[0], 'c_on') if self.kpt_u[0].psit_nG is None: raise RuntimeError('No wavefunctions yet') if isinstance(self.kpt_u[0].psit_nG, FileReference): # XXX initialize raise RuntimeError('Wavefunctions have not been initialized.') taut_sG = self.gd.zeros(self.nspins) dpsit_G = self.gd.empty(dtype=self.dtype) for kpt in self.kpt_u: for f, psit_G in zip(kpt.f_n, kpt.psit_nG): for v in range(3): self.taugrad_v[v](psit_G, dpsit_G, kpt.phase_cd) axpy(0.5 * f, abs(dpsit_G)**2, taut_sG[kpt.s]) self.kpt_comm.sum(taut_sG) self.band_comm.sum(taut_sG) return taut_sG def apply_mgga_orbital_dependent_hamiltonian(self, kpt, psit_xG, Htpsit_xG, dH_asp, dedtaut_G): a_G = self.gd.empty(dtype=psit_xG.dtype) for psit_G, Htpsit_G in zip(psit_xG, Htpsit_xG): for v in range(3): self.taugrad_v[v](psit_G, a_G, kpt.phase_cd) self.taugrad_v[v](dedtaut_G * a_G, a_G, kpt.phase_cd) axpy(-0.5, a_G, Htpsit_G) def ibz2bz(self, atoms): """Transform wave functions in IBZ to the full BZ.""" assert self.kd.comm.size == 1 # New k-point descriptor for full BZ: kd = KPointDescriptor(self.kd.bzk_kc, nspins=self.nspins) kd.set_symmetry(atoms, self.setups, usesymm=None) kd.set_communicator(serial_comm) self.pt = LFC(self.gd, [setup.pt_j for setup in self.setups], kd, dtype=self.dtype) self.pt.set_positions(atoms.get_scaled_positions()) self.initialize_wave_functions_from_restart_file() weight = 2.0 / kd.nspins / kd.nbzkpts # Build new list of k-points: kpt_u = [] for s in range(self.nspins): for k in range(kd.nbzkpts): # Index of symmetry related point in the IBZ ik = self.kd.bz2ibz_k[k] r, u = self.kd.get_rank_and_index(s, ik) assert r == 0 kpt = self.kpt_u[u] phase_cd = np.exp(2j * np.pi * self.gd.sdisp_cd * kd.bzk_kc[k, :, np.newaxis]) # New k-point: kpt2 = KPoint(weight, s, k, k, phase_cd) kpt2.f_n = kpt.f_n / kpt.weight / kd.nbzkpts * 2 / self.nspins kpt2.eps_n = kpt.eps_n.copy() # Transform wave functions using symmetry operation: Psit_nG = self.gd.collect(kpt.psit_nG) if Psit_nG is not None: Psit_nG = Psit_nG.copy() for Psit_G in Psit_nG: Psit_G[:] = self.kd.transform_wave_function(Psit_G, k) kpt2.psit_nG = self.gd.empty(self.bd.nbands, dtype=self.dtype) self.gd.distribute(Psit_nG, kpt2.psit_nG) # Calculate PAW projections: kpt2.P_ani = self.pt.dict(len(kpt.psit_nG)) self.pt.integrate(kpt2.psit_nG, kpt2.P_ani, k) kpt_u.append(kpt2) self.kd = kd self.kpt_u = kpt_u def write(self, writer, write_wave_functions=False): writer['Mode'] = 'fd' if not write_wave_functions: return writer.add('PseudoWaveFunctions', ('nspins', 'nibzkpts', 'nbands', 'ngptsx', 'ngptsy', 'ngptsz'), dtype=self.dtype) if hasattr(writer, 'hdf5'): parallel = (self.world.size > 1) for kpt in self.kpt_u: indices = [kpt.s, kpt.k] indices.append(self.bd.get_slice()) indices += self.gd.get_slice() writer.fill(kpt.psit_nG, parallel=parallel, *indices) else: for s in range(self.nspins): for k in range(self.nibzkpts): for n in range(self.bd.nbands): psit_G = self.get_wave_function_array(n, k, s) writer.fill(psit_G, s, k, n) def read(self, reader, hdf5): if ((not hdf5 and self.bd.comm.size == 1) or (hdf5 and self.world.size == 1)): # We may not be able to keep all the wave # functions in memory - so psit_nG will be a special type of # array that is really just a reference to a file: for kpt in self.kpt_u: kpt.psit_nG = reader.get_reference('PseudoWaveFunctions', (kpt.s, kpt.k)) else: for kpt in self.kpt_u: kpt.psit_nG = self.empty(self.bd.mynbands) if hdf5: indices = [kpt.s, kpt.k] indices.append(self.bd.get_slice()) indices += self.gd.get_slice() reader.get('PseudoWaveFunctions', out=kpt.psit_nG, parallel=(self.world.size > 1), *indices) else: # Read band by band to save memory for myn, psit_G in enumerate(kpt.psit_nG): n = self.bd.global_index(myn) if self.gd.comm.rank == 0: big_psit_G = np.array( reader.get('PseudoWaveFunctions', kpt.s, kpt.k, n), self.dtype) else: big_psit_G = None self.gd.distribute(big_psit_G, psit_G) def initialize_from_lcao_coefficients(self, basis_functions, mynbands): for kpt in self.kpt_u: kpt.psit_nG = self.gd.zeros(self.bd.mynbands, self.dtype) basis_functions.lcao_to_grid(kpt.C_nM, kpt.psit_nG[:mynbands], kpt.q) kpt.C_nM = None def random_wave_functions(self, nao): """Generate random wave functions.""" gpts = self.gd.N_c[0] * self.gd.N_c[1] * self.gd.N_c[2] if self.bd.nbands < gpts / 64: gd1 = self.gd.coarsen() gd2 = gd1.coarsen() psit_G1 = gd1.empty(dtype=self.dtype) psit_G2 = gd2.empty(dtype=self.dtype) interpolate2 = Transformer(gd2, gd1, 1, self.dtype).apply interpolate1 = Transformer(gd1, self.gd, 1, self.dtype).apply shape = tuple(gd2.n_c) scale = np.sqrt(12 / abs(np.linalg.det(gd2.cell_cv))) old_state = np.random.get_state() np.random.seed(4 + self.world.rank) for kpt in self.kpt_u: for psit_G in kpt.psit_nG[nao:]: if self.dtype == float: psit_G2[:] = (np.random.random(shape) - 0.5) * scale else: psit_G2.real = (np.random.random(shape) - 0.5) * scale psit_G2.imag = (np.random.random(shape) - 0.5) * scale interpolate2(psit_G2, psit_G1, kpt.phase_cd) interpolate1(psit_G1, psit_G, kpt.phase_cd) np.random.set_state(old_state) elif gpts / 64 <= self.bd.nbands < gpts / 8: gd1 = self.gd.coarsen() psit_G1 = gd1.empty(dtype=self.dtype) interpolate1 = Transformer(gd1, self.gd, 1, self.dtype).apply shape = tuple(gd1.n_c) scale = np.sqrt(12 / abs(np.linalg.det(gd1.cell_cv))) old_state = np.random.get_state() np.random.seed(4 + self.world.rank) for kpt in self.kpt_u: for psit_G in kpt.psit_nG[nao:]: if self.dtype == float: psit_G1[:] = (np.random.random(shape) - 0.5) * scale else: psit_G1.real = (np.random.random(shape) - 0.5) * scale psit_G1.imag = (np.random.random(shape) - 0.5) * scale interpolate1(psit_G1, psit_G, kpt.phase_cd) np.random.set_state(old_state) else: shape = tuple(self.gd.n_c) scale = np.sqrt(12 / abs(np.linalg.det(self.gd.cell_cv))) old_state = np.random.get_state() np.random.seed(4 + self.world.rank) for kpt in self.kpt_u: for psit_G in kpt.psit_nG[nao:]: if self.dtype == float: psit_G[:] = (np.random.random(shape) - 0.5) * scale else: psit_G.real = (np.random.random(shape) - 0.5) * scale psit_G.imag = (np.random.random(shape) - 0.5) * scale np.random.set_state(old_state) def estimate_memory(self, mem): FDPWWaveFunctions.estimate_memory(self, mem)
class WaveFunctions: """Class for wave-function related stuff (e.g. projectors).""" def __init__(self, nbands, kpt_u, setups, kd, gd, dtype=float): """Store and initialize required attributes. Parameters ---------- nbands: int Number of occupied bands. kpt_u: list of KPoints List of KPoint instances from a ground-state calculation (i.e. the attribute ``calc.wfs.kpt_u``). setups: Setups LocalizedFunctionsCollection setups. kd: KPointDescriptor K-point and symmetry related stuff. gd: GridDescriptor Descriptor for the coarse grid. dtype: dtype This is the ``dtype`` for the wave-function derivatives (same as the ``dtype`` for the ground-state wave-functions). """ self.dtype = dtype # K-point related attributes self.kd = kd # Number of occupied bands self.nbands = nbands # Projectors self.pt = LFC(gd, [setup.pt_j for setup in setups], dtype=self.dtype) # Store grid self.gd = gd # Unfold the irreducible BZ to the full BZ # List of KPointContainers for the k-points in the full BZ self.kpt_u = [] # No symmetries or only time-reversal symmetry used if kd.symmetry is None: # For now, time-reversal symmetry not allowed assert len(kpt_u) == kd.nbzkpts for k in range(kd.nbzkpts): kpt_ = kpt_u[k] psit_nG = gd.empty(nbands, dtype=self.dtype) for n, psit_G in enumerate(psit_nG): psit_G[:] = kpt_.psit_nG[n] # psit_0 = psit_G[0, 0, 0] # psit_G *= psit_0.conj() / (abs(psit_0)) # Strip off KPoint attributes and store in the KPointContainer # Note, only the occupied GS wave-functions are retained here ! kpt = KPointContainer(weight=kpt_.weight, k=kpt_.k, s=kpt_.s, phase_cd=kpt_.phase_cd, eps_n=kpt_.eps_n[:nbands], psit_nG=psit_nG, psit1_nG=None, P_ani=None, dP_aniv=None) # q=kpt.q, # f_n=kpt.f_n[:nbands]) self.kpt_u.append(kpt) else: assert len(kpt_u) == kd.nibzkpts for k, k_c in enumerate(kd.bzk_kc): # Index of symmetry related point in the irreducible BZ ik = kd.kibz_k[k] # Index of point group operation s = kd.sym_k[k] # Time-reversal symmetry used time_reversal = kd.time_reversal_k[k] # Coordinates of symmetry related point in the irreducible BZ ik_c = kd.ibzk_kc[ik] # Point group operation op_cc = kd.symmetry.op_scc[s] # KPoint from ground-state calculation kpt_ = kpt_u[ik] weight = 1. / kd.nbzkpts * (2 - kpt_.s) phase_cd = np.exp(2j * pi * gd.sdisp_cd * k_c[:, np.newaxis]) psit_nG = gd.empty(nbands, dtype=self.dtype) for n, psit_G in enumerate(psit_nG): #XXX Seems to corrupt my memory somehow ??? psit_G[:] = kd.symmetry.symmetrize_wavefunction( kpt_.psit_nG[n], ik_c, k_c, op_cc, time_reversal) # Choose gauge # psit_0 = psit_G[0, 0, 0] # psit_G *= psit_0.conj() / (abs(psit_0)) kpt = KPointContainer(weight=weight, k=k, s=kpt_.s, phase_cd=phase_cd, eps_n=kpt_.eps_n[:nbands], psit_nG=psit_nG, psit1_nG=None, P_ani=None, dP_aniv=None) self.kpt_u.append(kpt) def initialize(self, spos_ac): """Initialize projectors according to the ``gamma`` attribute.""" # Set positions on LFC's self.pt.set_positions(spos_ac) if not self.kd.gamma: # Set k-vectors and update self.pt.set_k_points(self.kd.ibzk_kc) self.pt._update(spos_ac) # Calculate projector coefficients for the GS wave-functions self.calculate_projector_coef() def reset(self): """Make fresh arrays for wave-function derivatives.""" for kpt in self.kpt_u: kpt.psit1_nG = self.gd.zeros(n=self.nbands, dtype=self.dtype) def calculate_projector_coef(self): """Coefficients for the derivative of the non-local part of the PP. Parameters ---------- k: int Index of the k-point of the Bloch state on which the non-local potential operates on. The calculated coefficients are the following (except for an overall sign of -1; see ``derivative`` member function of class ``LFC``): 1. Coefficients from the projector functions:: / a P_ani = | dG p (G) Psi (G) , / i n 2. Coefficients from the derivative of the projector functions:: / a dP_aniv = | dG dp (G) Psi (G) , / iv n where:: a d a dp (G) = --- Phi (G) . iv a i dR """ n = self.nbands for kpt in self.kpt_u: # K-point index and wave-functions k = kpt.k psit_nG = kpt.psit_nG # Integration dicts P_ani = self.pt.dict(shape=n) dP_aniv = self.pt.dict(shape=n, derivative=True) # 1) Integrate with projectors self.pt.integrate(psit_nG, P_ani, q=k) kpt.P_ani = P_ani # 2) Integrate with derivative of projectors self.pt.derivative(psit_nG, dP_aniv, q=k) kpt.dP_aniv = dP_aniv
a = 4.0 gd = GridDescriptor(N_c=[16, 20, 20], cell_cv=[a, a + 1, a + 2], pbc_c=(0, 1, 1)) spos_ac = np.array([[0.25, 0.15, 0.35], [0.5, 0.5, 0.5]]) kpts_kc = None s = Spline(l=0, rmax=2.0, f_g=np.array([1, 0.9, 0.1, 0.0])) p = Spline(l=1, rmax=2.0, f_g=np.array([1, 0.9, 0.1, 0.0])) spline_aj = [[s], [s, p]] c = LFC(gd, spline_aj, cut=True, forces=True) c.set_positions(spos_ac) C_ani = c.dict(3, zero=True) if 1 in C_ani: C_ani[1][:, 1:] = np.eye(3) psi = gd.zeros(3) c.add(psi, C_ani) c.integrate(psi, C_ani) if 1 in C_ani: d = C_ani[1][:, 1:].diagonal() assert d.ptp() < 4e-6 C_ani[1][:, 1:] -= np.diag(d) assert abs(C_ani[1]).max() < 5e-17 d_aniv = c.dict(3, derivative=True) c.derivative(psi, d_aniv) if 1 in d_aniv: for v in range(3): assert abs(d_aniv[1][v - 1, 0, v] + 0.2144) < 5e-5 d_aniv[1][v - 1, 0, v] = 0 assert abs(d_aniv[1]).max() < 3e-16 eps = 0.0001 pos_av = np.dot(spos_ac, gd.cell_cv) for v in range(3):
class WaveFunctions: """Class for wave-function related stuff (e.g. projectors).""" def __init__(self, nbands, kpt_u, setups, kd, gd, dtype=float): """Store and initialize required attributes. Parameters ---------- nbands: int Number of occupied bands. kpt_u: list of KPoints List of KPoint instances from a ground-state calculation (i.e. the attribute ``calc.wfs.kpt_u``). setups: Setups LocalizedFunctionsCollection setups. kd: KPointDescriptor K-point and symmetry related stuff. gd: GridDescriptor Descriptor for the coarse grid. dtype: dtype This is the ``dtype`` for the wave-function derivatives (same as the ``dtype`` for the ground-state wave-functions). """ self.dtype = dtype # K-point related attributes self.kd = kd # Number of occupied bands self.nbands = nbands # Projectors self.pt = LFC(gd, [setup.pt_j for setup in setups], kd, dtype=self.dtype) # Store grid self.gd = gd # Unfold the irreducible BZ to the full BZ # List of KPointContainers for the k-points in the full BZ self.kpt_u = [] # No symmetries or only time-reversal symmetry used assert kd.symmetry.point_group == False if kd.symmetry.time_reversal == False: # For now, time-reversal symmetry not allowed assert len(kpt_u) == kd.nbzkpts for k in range(kd.nbzkpts): kpt_ = kpt_u[k] psit_nG = gd.empty(nbands, dtype=self.dtype) for n, psit_G in enumerate(psit_nG): psit_G[:] = kpt_.psit_nG[n] # psit_0 = psit_G[0, 0, 0] # psit_G *= psit_0.conj() / (abs(psit_0)) # Strip off KPoint attributes and store in the KPointContainer # Note, only the occupied GS wave-functions are retained here ! kpt = KPointContainer(weight=kpt_.weight, k=kpt_.k, s=kpt_.s, phase_cd=kpt_.phase_cd, eps_n=kpt_.eps_n[:nbands], psit_nG=psit_nG, psit1_nG=None, P_ani=None, dP_aniv=None) # q=kpt.q, # f_n=kpt.f_n[:nbands]) self.kpt_u.append(kpt) else: assert len(kpt_u) == kd.nibzkpts for k, k_c in enumerate(kd.bzk_kc): # Index of symmetry related point in the irreducible BZ ik = kd.bz2ibz_k[k] # Index of point group operation s = kd.sym_k[k] # Time-reversal symmetry used time_reversal = kd.time_reversal_k[k] # Coordinates of symmetry related point in the irreducible BZ ik_c = kd.ibzk_kc[ik] # Point group operation op_cc = kd.symmetry.op_scc[s] # KPoint from ground-state calculation kpt_ = kpt_u[ik] weight = 1. / kd.nbzkpts * (2 - kpt_.s) phase_cd = np.exp(2j * pi * gd.sdisp_cd * k_c[:, np.newaxis]) psit_nG = gd.empty(nbands, dtype=self.dtype) for n, psit_G in enumerate(psit_nG): #XXX Seems to corrupt my memory somehow ??? psit_G[:] = kd.symmetry.symmetrize_wavefunction( kpt_.psit_nG[n], ik_c, k_c, op_cc, time_reversal) # Choose gauge # psit_0 = psit_G[0, 0, 0] # psit_G *= psit_0.conj() / (abs(psit_0)) kpt = KPointContainer(weight=weight, k=k, s=kpt_.s, phase_cd=phase_cd, eps_n=kpt_.eps_n[:nbands], psit_nG=psit_nG, psit1_nG=None, P_ani=None, dP_aniv=None) self.kpt_u.append(kpt) def initialize(self, spos_ac): """Initialize projectors according to the ``gamma`` attribute.""" # Set positions on LFC's self.pt.set_positions(spos_ac) # Calculate projector coefficients for the GS wave-functions self.calculate_projector_coef() def reset(self): """Make fresh arrays for wave-function derivatives.""" for kpt in self.kpt_u: kpt.psit1_nG = self.gd.zeros(n=self.nbands, dtype=self.dtype) def calculate_projector_coef(self): """Coefficients for the derivative of the non-local part of the PP. Parameters ---------- k: int Index of the k-point of the Bloch state on which the non-local potential operates on. The calculated coefficients are the following (except for an overall sign of -1; see ``derivative`` member function of class ``LFC``): 1. Coefficients from the projector functions:: / a P_ani = | dG p (G) Psi (G) , / i n 2. Coefficients from the derivative of the projector functions:: / a dP_aniv = | dG dp (G) Psi (G) , / iv n where:: a d a dp (G) = --- Phi (G) . iv a i dR """ n = self.nbands for kpt in self.kpt_u: # K-point index and wave-functions k = kpt.k psit_nG = kpt.psit_nG # Integration dicts P_ani = self.pt.dict(shape=n) dP_aniv = self.pt.dict(shape=n, derivative=True) # 1) Integrate with projectors self.pt.integrate(psit_nG, P_ani, q=k) kpt.P_ani = P_ani # 2) Integrate with derivative of projectors self.pt.derivative(psit_nG, dP_aniv, q=k) kpt.dP_aniv = dP_aniv
class FDWaveFunctions(FDPWWaveFunctions): mode = 'fd' def __init__(self, stencil, diagksl, orthoksl, initksl, gd, nvalence, setups, bd, dtype, world, kd, kptband_comm, timer=None): FDPWWaveFunctions.__init__(self, diagksl, orthoksl, initksl, gd, nvalence, setups, bd, dtype, world, kd, kptband_comm, timer) # Kinetic energy operator: self.kin = Laplace(self.gd, -0.5, stencil, self.dtype) self.matrixoperator = MatrixOperator(self.orthoksl) self.taugrad_v = None # initialized by MGGA functional def empty(self, n=(), global_array=False, realspace=False, q=-1): return self.gd.empty(n, self.dtype, global_array) def integrate(self, a_xg, b_yg=None, global_integral=True): return self.gd.integrate(a_xg, b_yg, global_integral) def bytes_per_wave_function(self): return self.gd.bytecount(self.dtype) def set_setups(self, setups): self.pt = LFC(self.gd, [setup.pt_j for setup in setups], self.kd, dtype=self.dtype, forces=True) FDPWWaveFunctions.set_setups(self, setups) def set_positions(self, spos_ac): FDPWWaveFunctions.set_positions(self, spos_ac) def summary(self, fd): fd.write('Wave functions: Uniform real-space grid\n') fd.write('Kinetic energy operator: %s\n' % self.kin.description) def make_preconditioner(self, block=1): return Preconditioner(self.gd, self.kin, self.dtype, block) def apply_pseudo_hamiltonian(self, kpt, hamiltonian, psit_xG, Htpsit_xG): self.timer.start('Apply hamiltonian') self.kin.apply(psit_xG, Htpsit_xG, kpt.phase_cd) hamiltonian.apply_local_potential(psit_xG, Htpsit_xG, kpt.s) self.timer.stop('Apply hamiltonian') def add_orbital_density(self, nt_G, kpt, n): if self.dtype == float: axpy(1.0, kpt.psit_nG[n]**2, nt_G) else: axpy(1.0, kpt.psit_nG[n].real**2, nt_G) axpy(1.0, kpt.psit_nG[n].imag**2, nt_G) def add_to_density_from_k_point_with_occupation(self, nt_sG, kpt, f_n): # Used in calculation of response part of GLLB-potential nt_G = nt_sG[kpt.s] if self.dtype == float: for f, psit_G in zip(f_n, kpt.psit_nG): axpy(f, psit_G**2, nt_G) else: for f, psit_G in zip(f_n, kpt.psit_nG): axpy(f, psit_G.real**2, nt_G) axpy(f, psit_G.imag**2, nt_G) # Hack used in delta-scf calculations: if hasattr(kpt, 'c_on'): assert self.bd.comm.size == 1 d_nn = np.zeros((self.bd.mynbands, self.bd.mynbands), dtype=complex) for ne, c_n in zip(kpt.ne_o, kpt.c_on): d_nn += ne * np.outer(c_n.conj(), c_n) for d_n, psi0_G in zip(d_nn, kpt.psit_nG): for d, psi_G in zip(d_n, kpt.psit_nG): if abs(d) > 1.e-12: nt_G += (psi0_G.conj() * d * psi_G).real def calculate_kinetic_energy_density(self): if self.taugrad_v is None: self.taugrad_v = [ Gradient(self.gd, v, n=3, dtype=self.dtype).apply for v in range(3) ] assert not hasattr(self.kpt_u[0], 'c_on') if self.kpt_u[0].psit_nG is None: raise RuntimeError('No wavefunctions yet') if isinstance(self.kpt_u[0].psit_nG, FileReference): # XXX initialize raise RuntimeError('Wavefunctions have not been initialized.') taut_sG = self.gd.zeros(self.nspins) dpsit_G = self.gd.empty(dtype=self.dtype) for kpt in self.kpt_u: for f, psit_G in zip(kpt.f_n, kpt.psit_nG): for v in range(3): self.taugrad_v[v](psit_G, dpsit_G, kpt.phase_cd) axpy(0.5 * f, abs(dpsit_G)**2, taut_sG[kpt.s]) self.kd.comm.sum(taut_sG) self.band_comm.sum(taut_sG) return taut_sG def apply_mgga_orbital_dependent_hamiltonian(self, kpt, psit_xG, Htpsit_xG, dH_asp, dedtaut_G): a_G = self.gd.empty(dtype=psit_xG.dtype) for psit_G, Htpsit_G in zip(psit_xG, Htpsit_xG): for v in range(3): self.taugrad_v[v](psit_G, a_G, kpt.phase_cd) self.taugrad_v[v](dedtaut_G * a_G, a_G, kpt.phase_cd) axpy(-0.5, a_G, Htpsit_G) def ibz2bz(self, atoms): """Transform wave functions in IBZ to the full BZ.""" assert self.kd.comm.size == 1 # New k-point descriptor for full BZ: kd = KPointDescriptor(self.kd.bzk_kc, nspins=self.nspins) #kd.set_symmetry(atoms, self.setups, enabled=False) kd.set_communicator(serial_comm) self.pt = LFC(self.gd, [setup.pt_j for setup in self.setups], kd, dtype=self.dtype) self.pt.set_positions(atoms.get_scaled_positions()) self.initialize_wave_functions_from_restart_file() weight = 2.0 / kd.nspins / kd.nbzkpts # Build new list of k-points: kpt_u = [] for s in range(self.nspins): for k in range(kd.nbzkpts): # Index of symmetry related point in the IBZ ik = self.kd.bz2ibz_k[k] r, u = self.kd.get_rank_and_index(s, ik) assert r == 0 kpt = self.kpt_u[u] phase_cd = np.exp(2j * np.pi * self.gd.sdisp_cd * kd.bzk_kc[k, :, np.newaxis]) # New k-point: kpt2 = KPoint(weight, s, k, k, phase_cd) kpt2.f_n = kpt.f_n / kpt.weight / kd.nbzkpts * 2 / self.nspins kpt2.eps_n = kpt.eps_n.copy() # Transform wave functions using symmetry operation: Psit_nG = self.gd.collect(kpt.psit_nG) if Psit_nG is not None: Psit_nG = Psit_nG.copy() for Psit_G in Psit_nG: Psit_G[:] = self.kd.transform_wave_function(Psit_G, k) kpt2.psit_nG = self.gd.empty(self.bd.nbands, dtype=self.dtype) self.gd.distribute(Psit_nG, kpt2.psit_nG) # Calculate PAW projections: kpt2.P_ani = self.pt.dict(len(kpt.psit_nG)) self.pt.integrate(kpt2.psit_nG, kpt2.P_ani, k) kpt_u.append(kpt2) self.kd = kd self.kpt_u = kpt_u def write(self, writer, write_wave_functions=False): writer['Mode'] = 'fd' if not write_wave_functions: return writer.add( 'PseudoWaveFunctions', ('nspins', 'nibzkpts', 'nbands', 'ngptsx', 'ngptsy', 'ngptsz'), dtype=self.dtype) if hasattr(writer, 'hdf5'): parallel = (self.world.size > 1) for kpt in self.kpt_u: indices = [kpt.s, kpt.k] indices.append(self.bd.get_slice()) indices += self.gd.get_slice() writer.fill(kpt.psit_nG, parallel=parallel, *indices) else: for s in range(self.nspins): for k in range(self.kd.nibzkpts): for n in range(self.bd.nbands): psit_G = self.get_wave_function_array(n, k, s) writer.fill(psit_G, s, k, n) def read(self, reader, hdf5): if ((not hdf5 and self.bd.comm.size == 1) or (hdf5 and self.world.size == 1)): # We may not be able to keep all the wave # functions in memory - so psit_nG will be a special type of # array that is really just a reference to a file: for kpt in self.kpt_u: kpt.psit_nG = reader.get_reference('PseudoWaveFunctions', (kpt.s, kpt.k)) else: for kpt in self.kpt_u: kpt.psit_nG = self.empty(self.bd.mynbands) if hdf5: indices = [kpt.s, kpt.k] indices.append(self.bd.get_slice()) indices += self.gd.get_slice() reader.get('PseudoWaveFunctions', out=kpt.psit_nG, parallel=(self.world.size > 1), *indices) else: # Read band by band to save memory for myn, psit_G in enumerate(kpt.psit_nG): n = self.bd.global_index(myn) if self.gd.comm.rank == 0: big_psit_G = np.array( reader.get('PseudoWaveFunctions', kpt.s, kpt.k, n), self.dtype) else: big_psit_G = None self.gd.distribute(big_psit_G, psit_G) def initialize_from_lcao_coefficients(self, basis_functions, mynbands): for kpt in self.kpt_u: kpt.psit_nG = self.gd.zeros(self.bd.mynbands, self.dtype) basis_functions.lcao_to_grid(kpt.C_nM, kpt.psit_nG[:mynbands], kpt.q) kpt.C_nM = None if use_mic: kpt.psit_nG_mic = stream.bind(kpt.psit_nG) stream.sync() def random_wave_functions(self, nao): """Generate random wave functions.""" gpts = self.gd.N_c[0] * self.gd.N_c[1] * self.gd.N_c[2] if self.bd.nbands < gpts / 64: gd1 = self.gd.coarsen() gd2 = gd1.coarsen() psit_G1 = gd1.empty(dtype=self.dtype) psit_G2 = gd2.empty(dtype=self.dtype) interpolate2 = Transformer(gd2, gd1, 1, self.dtype).apply interpolate1 = Transformer(gd1, self.gd, 1, self.dtype).apply shape = tuple(gd2.n_c) scale = np.sqrt(12 / abs(np.linalg.det(gd2.cell_cv))) old_state = np.random.get_state() np.random.seed(4 + self.world.rank) for kpt in self.kpt_u: for psit_G in kpt.psit_nG[nao:]: if self.dtype == float: psit_G2[:] = (np.random.random(shape) - 0.5) * scale else: psit_G2.real = (np.random.random(shape) - 0.5) * scale psit_G2.imag = (np.random.random(shape) - 0.5) * scale interpolate2(psit_G2, psit_G1, kpt.phase_cd) interpolate1(psit_G1, psit_G, kpt.phase_cd) np.random.set_state(old_state) elif gpts / 64 <= self.bd.nbands < gpts / 8: gd1 = self.gd.coarsen() psit_G1 = gd1.empty(dtype=self.dtype) interpolate1 = Transformer(gd1, self.gd, 1, self.dtype).apply shape = tuple(gd1.n_c) scale = np.sqrt(12 / abs(np.linalg.det(gd1.cell_cv))) old_state = np.random.get_state() np.random.seed(4 + self.world.rank) for kpt in self.kpt_u: for psit_G in kpt.psit_nG[nao:]: if self.dtype == float: psit_G1[:] = (np.random.random(shape) - 0.5) * scale else: psit_G1.real = (np.random.random(shape) - 0.5) * scale psit_G1.imag = (np.random.random(shape) - 0.5) * scale interpolate1(psit_G1, psit_G, kpt.phase_cd) np.random.set_state(old_state) else: shape = tuple(self.gd.n_c) scale = np.sqrt(12 / abs(np.linalg.det(self.gd.cell_cv))) old_state = np.random.get_state() np.random.seed(4 + self.world.rank) for kpt in self.kpt_u: for psit_G in kpt.psit_nG[nao:]: if self.dtype == float: psit_G[:] = (np.random.random(shape) - 0.5) * scale else: psit_G.real = (np.random.random(shape) - 0.5) * scale psit_G.imag = (np.random.random(shape) - 0.5) * scale np.random.set_state(old_state) def estimate_memory(self, mem): FDPWWaveFunctions.estimate_memory(self, mem)
def get_projections(self, locfun, spin=0): """Project wave functions onto localized functions Determine the projections of the Kohn-Sham eigenstates onto specified localized functions of the format:: locfun = [[spos_c, l, sigma], [...]] spos_c can be an atom index, or a scaled position vector. l is the angular momentum, and sigma is the (half-) width of the radial gaussian. Return format is:: f_kni = <psi_kn | f_i> where psi_kn are the wave functions, and f_i are the specified localized functions. As a special case, locfun can be the string 'projectors', in which case the bound state projectors are used as localized functions. """ wfs = self.wfs if locfun == 'projectors': f_kin = [] for kpt in wfs.kpt_u: if kpt.s == spin: f_in = [] for a, P_ni in kpt.P_ani.items(): i = 0 setup = wfs.setups[a] for l, n in zip(setup.l_j, setup.n_j): if n >= 0: for j in range(i, i + 2 * l + 1): f_in.append(P_ni[:, j]) i += 2 * l + 1 f_kin.append(f_in) f_kni = np.array(f_kin).transpose(0, 2, 1) return f_kni.conj() from gpaw.lfc import LocalizedFunctionsCollection as LFC from gpaw.spline import Spline from gpaw.utilities import _fact nkpts = len(wfs.ibzk_kc) nbf = np.sum([2 * l + 1 for pos, l, a in locfun]) f_kni = np.zeros((nkpts, wfs.nbands, nbf), wfs.dtype) spos_ac = self.atoms.get_scaled_positions() % 1.0 spos_xc = [] splines_x = [] for spos_c, l, sigma in locfun: if isinstance(spos_c, int): spos_c = spos_ac[spos_c] spos_xc.append(spos_c) alpha = .5 * Bohr**2 / sigma**2 r = np.linspace(0, 6. * sigma, 500) f_g = (_fact[l] * (4 * alpha)**(l + 3 / 2.) * np.exp(-alpha * r**2) / (np.sqrt(4 * np.pi) * _fact[2 * l + 1])) splines_x.append([Spline(l, rmax=r[-1], f_g=f_g, points=61)]) lf = LFC(wfs.gd, splines_x, wfs.kpt_comm, dtype=wfs.dtype) if not wfs.gamma: lf.set_k_points(wfs.ibzk_qc) lf.set_positions(spos_xc) k = 0 f_ani = lf.dict(wfs.nbands) for kpt in wfs.kpt_u: if kpt.s != spin: continue lf.integrate(kpt.psit_nG[:], f_ani, kpt.q) i1 = 0 for x, f_ni in f_ani.items(): i2 = i1 + f_ni.shape[1] f_kni[k, :, i1:i2] = f_ni i1 = i2 k += 1 return f_kni.conj()
gd = GridDescriptor(N_c=[16, 20, 20], cell_cv=[a, a + 1, a + 2], pbc_c=(0, 1, 1)) spos_ac = np.array([[0.25, 0.15, 0.35], [0.5, 0.5, 0.5]]) kpts_kc = None s = Spline(l=0, rmax=2.0, f_g=np.array([1, 0.9, 0.1, 0.0])) p = Spline(l=1, rmax=2.0, f_g=np.array([1, 0.9, 0.1, 0.0])) spline_aj = [[s], [s, p]] c = LFC(gd, spline_aj, cut=True, forces=True) c.set_positions(spos_ac) C_ani = c.dict(3, zero=True) if 1 in C_ani: C_ani[1][:, 1:] = np.eye(3) psi = gd.zeros(3) c.add(psi, C_ani) c.integrate(psi, C_ani) if 1 in C_ani: d = C_ani[1][:, 1:].diagonal() assert d.ptp() < 4e-6 C_ani[1][:, 1:] -= np.diag(d) assert abs(C_ani[1]).max() < 5e-17 d_aniv = c.dict(3, derivative=True) c.derivative(psi, d_aniv) if 1 in d_aniv: for v in range(3): assert abs(d_aniv[1][v - 1, 0, v] + 0.2144) < 5e-5 d_aniv[1][v - 1, 0, v] = 0 assert abs(d_aniv[1]).max() < 3e-16 eps = 0.0001 pos_av = np.dot(spos_ac, gd.cell_cv) for v in range(3):
if mode == 'fd': pt = LFC(gd, [setup.pt_j for setup in setups], KPointDescriptor(bzk_kc), dtype=calc.wfs.dtype) pt.set_positions(spos_ac) for spin in range(nspins): for k in range(len(bzk_kc)): ibzk = k # since no symmetry u = kd.get_rank_and_index(spin, ibzk)[1] kpt = calc.wfs.kpt_u[u] for n in range(nbands): P_ai = pt.dict() psit_G = calc.wfs.get_wave_function_array(n, ibzk, spin) pt.integrate(psit_G, P_ai, ibzk) for a in range(len(P_ai)): assert np.abs( P_ai[a] - calc.wfs.kpt_u[u].P_ani[a][n]).sum() < 1e-8 assert np.abs( P_ai[a] - df.get_P_ai(k, n, spin)[a]).sum() < 1e-8 else: pt = PWLFC([setup.pt_j for setup in setups], calc.wfs.pd) pt.set_positions(spos_ac) for spin in range(nspins): for k in range(len(bzk_kc)):