class PhononPerturbation(Perturbation): """Implementation of a phonon perturbation. This class implements the change in the effective potential due to a displacement of an atom ``a`` in direction ``v`` with wave-vector ``q``. The action of the perturbing potential on a state vector is implemented in the ``apply`` member function. """ def __init__(self, calc, kd, poisson_solver, dtype=float, **kwargs): """Store useful objects, e.g. lfc's for the various atomic functions. Depending on whether the system is periodic or finite, Poisson's equation is solved with FFT or multigrid techniques, respectively. Parameters ---------- calc: Calculator Ground-state calculation. kd: KPointDescriptor Descriptor for the q-vectors of the dynamical matrix. """ self.kd = kd self.dtype = dtype self.poisson = poisson_solver # Gamma wrt q-vector if self.kd.gamma: self.phase_cd = None else: assert self.kd.mynks == len(self.kd.ibzk_qc) self.phase_qcd = [] sdisp_cd = calc.wfs.gd.sdisp_cd for q in range(self.kd.mynks): phase_cd = np.exp(2j * np.pi * \ sdisp_cd * self.kd.ibzk_qc[q, :, np.newaxis]) self.phase_qcd.append(phase_cd) # Store grid-descriptors self.gd = calc.density.gd self.finegd = calc.density.finegd # Steal setups for the lfc's setups = calc.wfs.setups # Store projector coefficients self.dH_asp = calc.hamiltonian.dH_asp.copy() # Localized functions: # core corections self.nct = LFC(self.gd, [[setup.nct] for setup in setups], integral=[setup.Nct for setup in setups], dtype=self.dtype) # compensation charges #XXX what is the consequence of numerical errors in the integral ?? self.ghat = LFC(self.finegd, [setup.ghat_l for setup in setups], dtype=self.dtype) ## self.ghat = LFC(self.finegd, [setup.ghat_l for setup in setups], ## integral=sqrt(4 * pi), dtype=self.dtype) # vbar potential self.vbar = LFC(self.finegd, [[setup.vbar] for setup in setups], dtype=self.dtype) # Expansion coefficients for the compensation charges self.Q_aL = calc.density.Q_aL.copy() # Grid transformer -- convert array from fine to coarse grid self.restrictor = Transformer(self.finegd, self.gd, nn=3, dtype=self.dtype, allocate=False) # Atom, cartesian coordinate and q-vector of the perturbation self.a = None self.v = None # Local q-vector index of the perturbation if self.kd.gamma: self.q = -1 else: self.q = None def initialize(self, spos_ac): """Prepare the various attributes for a calculation.""" # Set positions on LFC's self.nct.set_positions(spos_ac) self.ghat.set_positions(spos_ac) self.vbar.set_positions(spos_ac) if not self.kd.gamma: # Set q-vectors and update self.ghat.set_k_points(self.kd.ibzk_qc) self.ghat._update(spos_ac) # Set q-vectors and update self.vbar.set_k_points(self.kd.ibzk_qc) self.vbar._update(spos_ac) # Phase factor exp(iq.r) needed to obtian the periodic part of lfcs coor_vg = self.finegd.get_grid_point_coordinates() cell_cv = self.finegd.cell_cv # Convert to scaled coordinates scoor_cg = np.dot(la.inv(cell_cv), coor_vg.swapaxes(0, -2)) scoor_cg = scoor_cg.swapaxes(1,-2) # Phase factor phase_qg = np.exp(2j * pi * np.dot(self.kd.ibzk_qc, scoor_cg.swapaxes(0,-2))) self.phase_qg = phase_qg.swapaxes(1, -2) #XXX To be removed from this class !! # Setup the Poisson solver -- to be used on the fine grid self.poisson.set_grid_descriptor(self.finegd) self.poisson.initialize() # Grid transformer self.restrictor.allocate() def set_q(self, q): """Set the index of the q-vector of the perturbation.""" assert not self.kd.gamma, "Gamma-point calculation" self.q = q # Update phases and Poisson solver self.phase_cd = self.phase_qcd[q] self.poisson.set_q(self.kd.ibzk_qc[q]) # Invalidate calculated quantities # - local part of perturbing potential self.v1_G = None def set_av(self, a, v): """Set atom and cartesian component of the perturbation. Parameters ---------- a: int Index of the atom. v: int Cartesian component (0, 1 or 2) of the atomic displacement. """ assert self.q is not None self.a = a self.v = v # Update derivative of local potential self.calculate_local_potential() def get_phase_cd(self): """Overwrite base class member function.""" return self.phase_cd def has_q(self): """Overwrite base class member function.""" return (not self.kd.gamma) def get_q(self): """Return q-vector.""" assert not self.kd.gamma, "Gamma-point calculation." return self.kd.ibzk_qc[self.q] def solve_poisson(self, phi_g, rho_g): """Solve Poisson's equation for a Bloch-type charge distribution. More to come here ... Parameters ---------- phi_g: GridDescriptor Grid for the solution of Poissons's equation. rho_g: GridDescriptor Grid with the charge distribution. """ #assert phi_g.shape == rho_g.shape == self.phase_qg.shape[-3:], \ # ("Arrays have incompatible shapes.") assert self.q is not None, ("q-vector not set") # Gamma point calculation wrt the q-vector -> rho_g periodic if self.kd.gamma: #XXX NOTICE: solve_neutral self.poisson.solve_neutral(phi_g, rho_g) else: # Divide out the phase factor to get the periodic part rhot_g = rho_g/self.phase_qg[self.q] # Solve Poisson's equation for the periodic part of the potential #XXX NOTICE: solve_neutral self.poisson.solve_neutral(phi_g, rhot_g) # Return to Bloch form phi_g *= self.phase_qg[self.q] def calculate_local_potential(self): """Derivate of the local potential wrt an atomic displacements. The local part of the PAW potential has contributions from the compensation charges (``ghat``) and a spherical symmetric atomic potential (``vbar``). """ assert self.a is not None assert self.v is not None assert self.q is not None a = self.a v = self.v # Expansion coefficients for the ghat functions Q_aL = self.ghat.dict(zero=True) # Remember sign convention for add_derivative method # And be sure not to change the dtype of the arrays by assigning values # to array elements. Q_aL[a][:] = -1 * self.Q_aL[a] # Grid for derivative of compensation charges ghat1_g = self.finegd.zeros(dtype=self.dtype) self.ghat.add_derivative(a, v, ghat1_g, c_axi=Q_aL, q=self.q) # Solve Poisson's eq. for the potential from the periodic part of the # compensation charge derivative v1_g = self.finegd.zeros(dtype=self.dtype) self.solve_poisson(v1_g, ghat1_g) # Store potential from the compensation charge self.vghat1_g = v1_g.copy() # Add derivative of vbar - sign convention in add_derivative method c_ai = self.vbar.dict(zero=True) c_ai[a][0] = -1. self.vbar.add_derivative(a, v, v1_g, c_axi=c_ai, q=self.q) # Store potential for the evaluation of the energy derivative self.v1_g = v1_g.copy() # Transfer to coarse grid v1_G = self.gd.zeros(dtype=self.dtype) self.restrictor.apply(v1_g, v1_G, phases=self.phase_cd) self.v1_G = v1_G def apply(self, psi_nG, y_nG, wfs, k, kplusq): """Apply perturbation to unperturbed wave-functions. Parameters ---------- psi_nG: ndarray Set of grid vectors to which the perturbation is applied. y_nG: ndarray Output vectors. wfs: WaveFunctions Instance of class ``WaveFunctions``. k: int Index of the k-point for the vectors. kplusq: int Index of the k+q vector. """ assert self.a is not None assert self.v is not None assert self.q is not None assert psi_nG.ndim in (3, 4) assert tuple(self.gd.n_c) == psi_nG.shape[-3:] if psi_nG.ndim == 3: y_nG += self.v1_G * psi_nG else: y_nG += self.v1_G[np.newaxis, :] * psi_nG self.apply_nonlocal_potential(psi_nG, y_nG, wfs, k, kplusq) def apply_nonlocal_potential(self, psi_nG, y_nG, wfs, k, kplusq): """Derivate of the non-local PAW potential wrt an atomic displacement. Parameters ---------- k: int Index of the k-point being operated on. kplusq: int Index of the k+q vector. """ assert self.a is not None assert self.v is not None assert psi_nG.ndim in (3, 4) assert tuple(self.gd.n_c) == psi_nG.shape[-3:] if psi_nG.ndim == 3: n = 1 else: n = psi_nG.shape[0] a = self.a v = self.v P_ani = wfs.kpt_u[k].P_ani dP_aniv = wfs.kpt_u[k].dP_aniv pt = wfs.pt # < p_a^i | Psi_nk > P_ni = P_ani[a] # < dp_av^i | Psi_nk > - remember the sign convention of the derivative dP_ni = -1 * dP_aniv[a][...,v] # Expansion coefficients for the projectors on atom a dH_ii = unpack(self.dH_asp[a][0]) # The derivative of the non-local PAW potential has two contributions # 1) Sum over projectors c_ni = np.dot(dP_ni, dH_ii) c_ani = pt.dict(shape=n, zero=True) c_ani[a] = c_ni # k+q !! pt.add(y_nG, c_ani, q=kplusq) # 2) Sum over derivatives of the projectors dc_ni = np.dot(P_ni, dH_ii) dc_ani = pt.dict(shape=n, zero=True) # Take care of sign of derivative in the coefficients dc_ani[a] = -1 * dc_ni # k+q !! pt.add_derivative(a, v, y_nG, dc_ani, q=kplusq)
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
import numpy as np from gpaw.lfc import LocalizedFunctionsCollection as LFC from gpaw.grid_descriptor import GridDescriptor from gpaw.spline import Spline 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
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)
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
from gpaw.setup import Setup rc = 2.0 a = 2.5 * rc n = 64 lmax = 2 b = 8.0 m = (lmax + 1)**2 gd = GridDescriptor([n, n, n], [a, a, a]) r = np.linspace(0, rc, 200) g = np.exp(-(r / rc * b)**2) splines = [Spline(l=l, rmax=rc, f_g=g) for l in range(lmax + 1)] c = LFC(gd, [splines]) c.set_positions([(0, 0, 0)]) psi = gd.zeros(m) d0 = c.dict(m) if 0 in d0: d0[0] = np.identity(m) c.add(psi, d0) d1 = c.dict(m, derivative=True) c.derivative(psi, d1) class TestSetup(Setup): l_j = range(lmax + 1) nj = lmax + 1 ni = m def __init__(self): pass rgd = EquidistantRadialGridDescriptor(r[1], len(r)) g = [np.exp(-(r / rc * b)**2) * r**l for l in range(lmax + 1)] d2 = TestSetup().get_derivative_integrals(rgd, g, np.zeros_like(g)) if 0 in d1:
# Initialize s, p, d (9 in total) wave and put them on grid rc = 2.0 a = 2.5 * rc n = 64 lmax = 2 b = 8.0 m = (lmax + 1)**2 gd = GridDescriptor([n, n, n], [a, a, a]) r = np.linspace(0, rc, 200) g = np.exp(-(r / rc * b)**2) splines = [Spline(l=l, rmax=rc, f_g=g) for l in range(lmax + 1)] c = LFC(gd, [splines]) c.set_positions([(0.5, 0.5, 0.5)]) psi = gd.zeros(m) d0 = c.dict(m) if 0 in d0: d0[0] = np.identity(m) c.add(psi, d0) # Calculate on 3d-grid < phi_i | e**(-ik.r) | phi_j > R_a = np.array([a / 2, a / 2, a / 2]) rr = gd.get_grid_point_coordinates() for dim in range(3): rr[dim] -= R_a[dim] k_G = np.array([[11., 0.2, 0.1], [10., 0., 10.]]) nkpt = k_G.shape[0] d0 = np.zeros((nkpt, m, m), dtype=complex) for i in range(m):
from __future__ import print_function import numpy as np from gpaw.lfc import LocalizedFunctionsCollection as LFC from gpaw.grid_descriptor import GridDescriptor from gpaw.spline import Spline gd = GridDescriptor([20, 16, 16], [(4, 2, 0), (0, 4, 0), (0, 0, 4)]) spos_ac = np.array([[0.252, 0.15, 0.35], [0.503, 0.5, 0.5]]) s = Spline(l=0, rmax=2.0, f_g=np.array([1, 0.9, 0.1, 0.0])) spline_aj = [[s], [s]] c = LFC(gd, spline_aj) c.set_positions(spos_ac) c_ai = c.dict(zero=True) if 1 in c_ai: c_ai[1][0] = 2.0 psi = gd.zeros() c.add(psi, c_ai) d_avv = dict([(a, np.zeros((3, 3))) for a in c.my_atom_indices]) c.second_derivative(psi, d_avv) if 0 in d_avv: print(d_avv[0]) eps = 0.000001 d_aiv = c.dict(derivative=True) pos_av = np.dot(spos_ac, gd.cell_cv) for v in range(3): pos_av[0, v] += eps c.set_positions(np.dot(pos_av, gd.icell_cv.T)) c.derivative(psi, d_aiv) if 0 in d_aiv:
from gpaw.setup import Setup rc = 2.0 a = 2.5 * rc n = 64 lmax = 2 b = 8.0 m = (lmax + 1)**2 gd = GridDescriptor([n, n, n], [a, a, a]) r = np.linspace(0, rc, 200) g = np.exp(-(r / rc * b)**2) splines = [Spline(l=l, rmax=rc, f_g=g) for l in range(lmax + 1)] c = LFC(gd, [splines]) c.set_positions([(0, 0, 0)]) psi = gd.zeros(m) d0 = c.dict(m) if 0 in d0: d0[0] = np.identity(m) c.add(psi, d0) d1 = c.dict(m, derivative=True) c.derivative(psi, d1) class TestSetup(Setup): l_j = range(lmax + 1) nj = lmax + 1 ni = m def __init__(self): pass
import numpy as np from gpaw.lfc import LocalizedFunctionsCollection as LFC from gpaw.grid_descriptor import GridDescriptor from gpaw.spline import Spline 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
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) self.wd = self.gd # wave function descriptor # Kinetic energy operator: self.kin = Laplace(self.gd, -0.5, stencil, self.dtype, allocate=False) self.matrixoperator = MatrixOperator(orthoksl) def set_setups(self, setups): self.pt = LFC(self.gd, [setup.pt_j for setup in setups], self.kpt_comm, dtype=self.dtype, forces=True) FDPWWaveFunctions.set_setups(self, setups) def set_positions(self, spos_ac): if not self.kin.is_allocated(): self.kin.allocate() FDPWWaveFunctions.set_positions(self, spos_ac) def summary(self, fd): fd.write('Mode: Finite-difference\n') 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.kin.apply(psit_xG, Htpsit_xG, kpt.phase_cd) hamiltonian.apply_local_potential(psit_xG, Htpsit_xG, kpt.s) 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, tauct, grad_v): assert not hasattr(self.kpt_u[0], 'c_on') if isinstance(self.kpt_u[0].psit_nG, TarFileReference): 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): grad_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 calculate_forces(self, hamiltonian, F_av): # Calculate force-contribution from k-points: F_av.fill(0.0) F_aniv = self.pt.dict(self.bd.mynbands, derivative=True) for kpt in self.kpt_u: self.pt.derivative(kpt.psit_nG, F_aniv, kpt.q) for a, F_niv in F_aniv.items(): F_niv = F_niv.conj() F_niv *= kpt.f_n[:, np.newaxis, np.newaxis] dH_ii = unpack(hamiltonian.dH_asp[a][kpt.s]) P_ni = kpt.P_ani[a] F_vii = np.dot(np.dot(F_niv.transpose(), P_ni), dH_ii) F_niv *= kpt.eps_n[:, np.newaxis, np.newaxis] dO_ii = hamiltonian.setups[a].dO_ii F_vii -= np.dot(np.dot(F_niv.transpose(), P_ni), dO_ii) F_av[a] += 2 * F_vii.real.trace(0, 1, 2) # Hack used in delta-scf calculations: if hasattr(kpt, 'c_on'): assert self.bd.comm.size == 1 self.pt.derivative(kpt.psit_nG, F_aniv, kpt.q) #XXX again 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 a, F_niv in F_aniv.items(): F_niv = F_niv.conj() dH_ii = unpack(hamiltonian.dH_asp[a][kpt.s]) Q_ni = np.dot(d_nn, kpt.P_ani[a]) F_vii = np.dot(np.dot(F_niv.transpose(), Q_ni), dH_ii) F_niv *= kpt.eps_n[:, np.newaxis, np.newaxis] dO_ii = hamiltonian.setups[a].dO_ii F_vii -= np.dot(np.dot(F_niv.transpose(), Q_ni), dO_ii) F_av[a] += 2 * F_vii.real.trace(0, 1, 2) self.bd.comm.sum(F_av, 0) if self.bd.comm.rank == 0: self.kpt_comm.sum(F_av, 0) def estimate_memory(self, mem): FDPWWaveFunctions.estimate_memory(self, mem) self.kin.estimate_memory(mem.subnode('Kinetic operator'))
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) self.wd = self.gd # wave function descriptor # Kinetic energy operator: self.kin = Laplace(self.gd, -0.5, stencil, self.dtype, allocate=False) self.matrixoperator = MatrixOperator(orthoksl) def set_setups(self, setups): self.pt = LFC(self.gd, [setup.pt_j for setup in setups], self.kpt_comm, dtype=self.dtype, forces=True) FDPWWaveFunctions.set_setups(self, setups) def set_positions(self, spos_ac): if not self.kin.is_allocated(): self.kin.allocate() FDPWWaveFunctions.set_positions(self, spos_ac) def summary(self, fd): fd.write('Mode: Finite-difference\n') 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.kin.apply(psit_xG, Htpsit_xG, kpt.phase_cd) hamiltonian.apply_local_potential(psit_xG, Htpsit_xG, kpt.s) 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, tauct, grad_v): assert not hasattr(self.kpt_u[0], 'c_on') if isinstance(self.kpt_u[0].psit_nG, TarFileReference): 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): grad_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 calculate_forces(self, hamiltonian, F_av): # Calculate force-contribution from k-points: F_av.fill(0.0) F_aniv = self.pt.dict(self.bd.mynbands, derivative=True) for kpt in self.kpt_u: self.pt.derivative(kpt.psit_nG, F_aniv, kpt.q) for a, F_niv in F_aniv.items(): F_niv = F_niv.conj() F_niv *= kpt.f_n[:, np.newaxis, np.newaxis] dH_ii = unpack(hamiltonian.dH_asp[a][kpt.s]) P_ni = kpt.P_ani[a] F_vii = np.dot(np.dot(F_niv.transpose(), P_ni), dH_ii) F_niv *= kpt.eps_n[:, np.newaxis, np.newaxis] dO_ii = hamiltonian.setups[a].dO_ii F_vii -= np.dot(np.dot(F_niv.transpose(), P_ni), dO_ii) F_av[a] += 2 * F_vii.real.trace(0, 1, 2) # Hack used in delta-scf calculations: if hasattr(kpt, 'c_on'): assert self.bd.comm.size == 1 self.pt.derivative(kpt.psit_nG, F_aniv, kpt.q) #XXX again 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 a, F_niv in F_aniv.items(): F_niv = F_niv.conj() dH_ii = unpack(hamiltonian.dH_asp[a][kpt.s]) Q_ni = np.dot(d_nn, kpt.P_ani[a]) F_vii = np.dot(np.dot(F_niv.transpose(), Q_ni), dH_ii) F_niv *= kpt.eps_n[:, np.newaxis, np.newaxis] dO_ii = hamiltonian.setups[a].dO_ii F_vii -= np.dot(np.dot(F_niv.transpose(), Q_ni), dO_ii) F_av[a] += 2 * F_vii.real.trace(0, 1, 2) self.bd.comm.sum(F_av, 0) if self.bd.comm.rank == 0: self.kpt_comm.sum(F_av, 0) def estimate_memory(self, mem): FDPWWaveFunctions.estimate_memory(self, mem) self.kin.estimate_memory(mem.subnode('Kinetic operator'))
df = DF(calc) df.spos_ac = spos_ac 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)
# Initialize s, p, d (9 in total) wave and put them on grid rc = 2.0 a = 2.5 * rc n = 64 lmax = 2 b = 8.0 m = (lmax + 1)**2 gd = GridDescriptor([n, n, n], [a, a, a]) r = np.linspace(0, rc, 200) g = np.exp(-(r / rc * b)**2) splines = [Spline(l=l, rmax=rc, f_g=g) for l in range(lmax + 1)] c = LFC(gd, [splines]) c.set_positions([(0.5, 0.5, 0.5)]) psi = gd.zeros(m) d0 = c.dict(m) if 0 in d0: d0[0] = np.identity(m) c.add(psi, d0) # Calculate on 3d-grid < phi_i | e**(-ik.r) | phi_j > R_a = np.array([a/2,a/2,a/2]) rr = gd.get_grid_point_coordinates() for dim in range(3): rr[dim] -= R_a[dim] k_G = np.array([[11.,0.2,0.1],[10., 0., 10.]]) nkpt = k_G.shape[0] d0 = np.zeros((nkpt,m,m), dtype=complex) for i in range(m):
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()
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)
import numpy as np from gpaw.lfc import LocalizedFunctionsCollection as LFC from gpaw.grid_descriptor import GridDescriptor from gpaw.spline import Spline gd = GridDescriptor([20, 16, 16], [(4, 2, 0), (0, 4, 0), (0, 0, 4)]) spos_ac = np.array([[0.252, 0.15, 0.35], [0.503, 0.5, 0.5]]) s = Spline(l=0, rmax=2.0, f_g=np.array([1, 0.9, 0.1, 0.0])) spline_aj = [[s], [s]] c = LFC(gd, spline_aj) c.set_positions(spos_ac) c_ai = c.dict(zero=True) if 1 in c_ai: c_ai[1][0] = 2.0 psi = gd.zeros() c.add(psi, c_ai) d_avv = dict([(a, np.zeros((3, 3))) for a in c.my_atom_indices]) c.second_derivative(psi, d_avv) if 0 in d_avv: print d_avv[0] eps = 0.000001 d_aiv = c.dict(derivative=True) pos_av = np.dot(spos_ac, gd.cell_cv) for v in range(3): pos_av[0, v] += eps c.set_positions(np.dot(pos_av, gd.icell_cv.T)) c.derivative(psi, d_aiv) if 0 in d_aiv: d0_v = d_aiv[0][0].copy()