Source code for gpaw.wavefunctions.pw

import numbers
from math import pi

import numpy as np
from ase.units import Bohr, Ha
from ase.utils.timing import timer

import _gpaw
import gpaw
import gpaw.fftw as fftw
from gpaw.band_descriptor import BandDescriptor
from gpaw.blacs import BlacsDescriptor, BlacsGrid, Redistributor
from gpaw.lfc import BasisFunctions
from gpaw.matrix_descriptor import MatrixDescriptor
from gpaw.pw.descriptor import PWDescriptor
from gpaw.pw.lfc import PWLFC
from gpaw.utilities import unpack
from gpaw.utilities.blas import axpy
from gpaw.utilities.progressbar import ProgressBar
from gpaw.wavefunctions.arrays import PlaneWaveExpansionWaveFunctions
from gpaw.wavefunctions.fdpw import FDPWWaveFunctions
from gpaw.wavefunctions.mode import Mode


[docs]class PW(Mode): name = 'pw' def __init__(self, ecut=340, fftwflags=fftw.MEASURE, cell=None, gammacentered=False, pulay_stress=None, dedecut=None, force_complex_dtype=False): """Plane-wave basis mode. ecut: float Plane-wave cutoff in eV. gammacentered: bool Center the grid of chosen plane waves around the gamma point or q/k-vector dedecut: float or None or 'estimate' Estimate of derivative of total energy with respect to plane-wave cutoff. Used to calculate pulay_stress. pulay_stress: float or None Pulay-stress correction. fftwflags: int Flags for making an FFTW plan. There are 4 possibilities (default is MEASURE):: from gpaw.fftw import ESTIMATE, MEASURE, PATIENT, EXHAUSTIVE cell: 3x3 ndarray Use this unit cell to chose the planewaves. Only one of dedecut and pulay_stress can be used. """ assert not gammacentered self.ecut = ecut / Ha # Don't do expensive planning in dry-run mode: self.fftwflags = fftwflags if not gpaw.dry_run else fftw.MEASURE self.dedecut = dedecut self.pulay_stress = (None if pulay_stress is None else pulay_stress * Bohr**3 / Ha) assert pulay_stress is None or dedecut is None if cell is None: self.cell_cv = None else: self.cell_cv = cell / Bohr Mode.__init__(self, force_complex_dtype) def __call__(self, parallel, initksl, gd, **kwargs): dedepsilon = 0.0 if self.cell_cv is None: ecut = self.ecut else: volume0 = abs(np.linalg.det(self.cell_cv)) ecut = self.ecut * (volume0 / gd.volume)**(2 / 3.0) if self.pulay_stress is not None: dedepsilon = self.pulay_stress * gd.volume elif self.dedecut is not None: if self.dedecut == 'estimate': dedepsilon = 'estimate' else: dedepsilon = self.dedecut * 2 / 3 * ecut wfs = PWWaveFunctions(ecut, self.fftwflags, dedepsilon, parallel, initksl, gd=gd, **kwargs) return wfs def todict(self): dct = Mode.todict(self) dct['ecut'] = self.ecut * Ha if self.cell_cv is not None: dct['cell'] = self.cell_cv * Bohr if self.pulay_stress is not None: dct['pulay_stress'] = self.pulay_stress * Ha / Bohr**3 if self.dedecut is not None: dct['dedecut'] = self.dedecut return dct
class Preconditioner: """Preconditioner for KS equation. From: Teter, Payne and Allen, Phys. Rev. B 40, 12255 (1989) as modified by: Kresse and Furthm├╝ller, Phys. Rev. B 54, 11169 (1996) """ def __init__(self, G2_qG, pd): self.G2_qG = G2_qG self.pd = pd def calculate_kinetic_energy(self, psit_xG, kpt): if psit_xG.ndim == 1: return self.calculate_kinetic_energy(psit_xG[np.newaxis], kpt)[0] G2_G = self.G2_qG[kpt.q] return np.array([self.pd.integrate(0.5 * G2_G * psit_G, psit_G).real for psit_G in psit_xG]) def __call__(self, R_xG, kpt, ekin_x, out=None): if out is None: out = np.empty_like(R_xG) G2_G = self.G2_qG[kpt.q] if R_xG.ndim == 1: _gpaw.pw_precond(G2_G, R_xG, ekin_x, out) else: for PR_G, R_G, ekin in zip(out, R_xG, ekin_x): _gpaw.pw_precond(G2_G, R_G, ekin, PR_G) return out class NonCollinearPreconditioner(Preconditioner): def calculate_kinetic_energy(self, psit_xsG, kpt): shape = psit_xsG.shape ekin_xs = Preconditioner.calculate_kinetic_energy( self, psit_xsG.reshape((-1, shape[-1])), kpt) return ekin_xs.reshape(shape[:-1]).sum(-1) def __call__(self, R_sG, kpt, ekin, out=None): return Preconditioner.__call__(self, R_sG, kpt, [ekin, ekin], out)
[docs]class PWWaveFunctions(FDPWWaveFunctions): mode = 'pw' def __init__(self, ecut, fftwflags, dedepsilon, parallel, initksl, reuse_wfs_method, collinear, gd, nvalence, setups, bd, dtype, world, kd, kptband_comm, timer): self.ecut = ecut self.fftwflags = fftwflags self.dedepsilon = dedepsilon # Pulay correction for stress tensor self.ng_k = None # number of G-vectors for all IBZ k-points FDPWWaveFunctions.__init__(self, parallel, initksl, reuse_wfs_method=reuse_wfs_method, collinear=collinear, gd=gd, nvalence=nvalence, setups=setups, bd=bd, dtype=dtype, world=world, kd=kd, kptband_comm=kptband_comm, timer=timer) def empty(self, n=(), global_array=False, realspace=False, q=None): if isinstance(n, numbers.Integral): n = (n,) if realspace: return self.gd.empty(n, self.dtype, global_array) elif global_array: return np.zeros(n + (self.pd.ngmax,), complex) elif q is None: return np.zeros(n + (self.pd.maxmyng,), complex) else: return self.pd.empty(n, self.dtype, q) def integrate(self, a_xg, b_yg=None, global_integral=True): return self.pd.integrate(a_xg, b_yg, global_integral) def bytes_per_wave_function(self): return 16 * self.pd.maxmyng def set_setups(self, setups): self.timer.start('PWDescriptor') self.pd = PWDescriptor(self.ecut, self.gd, self.dtype, self.kd, self.fftwflags) self.timer.stop('PWDescriptor') # Build array of number of plane wave coefficiants for all k-points # in the IBZ: self.ng_k = np.zeros(self.kd.nibzkpts, dtype=int) for kpt in self.kpt_u: if kpt.s != 1: # avoid double counting (only sum over s=0 or None) self.ng_k[kpt.k] = len(self.pd.Q_qG[kpt.q]) self.kd.comm.sum(self.ng_k) self.pt = PWLFC([setup.pt_j for setup in setups], self.pd) FDPWWaveFunctions.set_setups(self, setups) if self.dedepsilon == 'estimate': dedecut = self.setups.estimate_dedecut(self.ecut) self.dedepsilon = dedecut * 2 / 3 * self.ecut def get_pseudo_partial_waves(self): return PWLFC([setup.get_partial_waves_for_atomic_orbitals() for setup in self.setups], self.pd) def __str__(self): s = 'Wave functions: Plane wave expansion\n' s += ' Cutoff energy: %.3f eV\n' % (self.pd.ecut * Ha) if self.dtype == float: s += (' Number of coefficients: %d (reduced to %d)\n' % (self.pd.ngmax * 2 - 1, self.pd.ngmax)) else: s += (' Number of coefficients (min, max): %d, %d\n' % (self.pd.ngmin, self.pd.ngmax)) stress = self.dedepsilon / self.gd.volume * Ha / Bohr**3 dedecut = 1.5 * self.dedepsilon / self.ecut s += (' Pulay-stress correction: {:.6f} eV/Ang^3 ' '(de/decut={:.6f})\n'.format(stress, dedecut)) if fftw.have_fftw(): s += ' Using FFTW library\n' else: s += " Using Numpy's FFT\n" return s + FDPWWaveFunctions.__str__(self) def make_preconditioner(self, block=1): if self.collinear: return Preconditioner(self.pd.G2_qG, self.pd) return NonCollinearPreconditioner(self.pd.G2_qG, self.pd)
[docs] @timer('Apply H') def apply_pseudo_hamiltonian(self, kpt, ham, psit_xG, Htpsit_xG): """Apply the pseudo Hamiltonian i.e. without PAW corrections.""" if not self.collinear: self.apply_pseudo_hamiltonian_nc(kpt, ham, psit_xG, Htpsit_xG) return N = len(psit_xG) S = self.gd.comm.size vt_R = self.gd.collect(ham.vt_sG[kpt.s], broadcast=True) Q_G = self.pd.Q_qG[kpt.q] T_G = 0.5 * self.pd.G2_qG[kpt.q] for n1 in range(0, N, S): n2 = min(n1 + S, N) psit_G = self.pd.alltoall1(psit_xG[n1:n2], kpt.q) with self.timer('HMM T'): np.multiply(T_G, psit_xG[n1:n2], Htpsit_xG[n1:n2]) if psit_G is not None: psit_R = self.pd.ifft(psit_G, kpt.q, local=True, safe=False) psit_R *= vt_R self.pd.fftplan.execute() vtpsit_G = self.pd.tmp_Q.ravel()[Q_G] else: vtpsit_G = self.pd.tmp_G self.pd.alltoall2(vtpsit_G, kpt.q, Htpsit_xG[n1:n2]) ham.xc.apply_orbital_dependent_hamiltonian( kpt, psit_xG, Htpsit_xG, ham.dH_asp)
def apply_pseudo_hamiltonian_nc(self, kpt, ham, psit_xG, Htpsit_xG): Htpsit_xG[:] = 0.5 * self.pd.G2_qG[kpt.q] * psit_xG v, x, y, z = ham.vt_xG iy = y * 1j for psit_sG, Htpsit_sG in zip(psit_xG, Htpsit_xG): a = self.pd.ifft(psit_sG[0], kpt.q) b = self.pd.ifft(psit_sG[1], kpt.q) Htpsit_sG[0] += self.pd.fft(a * (v + z) + b * (x - iy), kpt.q) Htpsit_sG[1] += self.pd.fft(a * (x + iy) + b * (v - z), kpt.q) def add_orbital_density(self, nt_G, kpt, n): axpy(1.0, abs(self.pd.ifft(kpt.psit_nG[n], kpt.q))**2, nt_G) def add_to_density_from_k_point_with_occupation(self, nt_xR, kpt, f_n): if not self.collinear: self.add_to_density_from_k_point_with_occupation_nc( nt_xR, kpt, f_n) return comm = self.gd.comm nt_R = self.gd.zeros(global_array=True) for n1 in range(0, self.bd.mynbands, comm.size): n2 = min(n1 + comm.size, self.bd.mynbands) psit_G = self.pd.alltoall1(kpt.psit.array[n1:n2], kpt.q) if psit_G is not None: f = f_n[n1 + comm.rank] psit_R = self.pd.ifft(psit_G, kpt.q, local=True, safe=False) # Same as nt_R += f * abs(psit_R)**2, but much faster: _gpaw.add_to_density(f, psit_R, nt_R) comm.sum(nt_R) nt_R = self.gd.distribute(nt_R) nt_xR[kpt.s] += nt_R def add_to_density_from_k_point_with_occupation_nc(self, nt_xR, kpt, f_n): for f, psit_sG in zip(f_n, kpt.psit.array): p1 = self.pd.ifft(psit_sG[0], kpt.q) p2 = self.pd.ifft(psit_sG[1], kpt.q) p11 = p1.real**2 + p1.imag**2 p22 = p2.real**2 + p2.imag**2 p12 = p1.conj() * p2 nt_xR[0] += f * (p11 + p22) nt_xR[1] += 2 * f * p12.real nt_xR[2] += 2 * f * p12.imag nt_xR[3] += f * (p11 - p22) def add_to_kinetic_energy_density_kpt(self, kpt, psit_xG, taut_xR): N = self.bd.mynbands S = self.gd.comm.size Gpsit_xG = np.empty((S,) + psit_xG.shape[1:], dtype=psit_xG.dtype) taut_R = self.gd.zeros(global_array=True) G_Gv = self.pd.get_reciprocal_vectors(q=kpt.q) comm = self.gd.comm for v in range(3): for n1 in range(0, N, S): n2 = min(n1 + S, N) dn = n2 - n1 Gpsit_xG[:dn] = 1j * G_Gv[:, v] * psit_xG[n1:n2] Gpsit_G = self.pd.alltoall1(Gpsit_xG[:dn], kpt.q) if Gpsit_G is not None: f = kpt.f_n[n1 + comm.rank] a_R = self.pd.ifft(Gpsit_G, kpt.q, local=True, safe=False) _gpaw.add_to_density(0.5 * f, a_R, taut_R) comm.sum(taut_R) taut_R = self.gd.distribute(taut_R) taut_xR[kpt.s] += taut_R def add_to_ke_crossterms_kpt(self, kpt, psit_xG, taut_xR): N = self.bd.mynbands S = self.gd.comm.size Gpsit_xG = np.empty((S,) + psit_xG.shape[1:], dtype=psit_xG.dtype) taut_vvR = self.gd.zeros((3, 3), global_array=True) G_Gv = self.pd.get_reciprocal_vectors(q=kpt.q) comm = self.gd.comm for n1 in range(0, N, S): n2 = min(n1 + S, N) dn = n2 - n1 a_vR = {} for v in range(3): Gpsit_xG[:dn] = 1j * G_Gv[:, v] * psit_xG[n1:n2] Gpsit_G = self.pd.alltoall1(Gpsit_xG[:dn], kpt.q) if Gpsit_G is not None: f = kpt.f_n[n1 + comm.rank] a_vR[v] = self.pd.ifft(Gpsit_G, kpt.q, local=True, safe=True) if len(a_vR) == 3: f = kpt.f_n[n1 + comm.rank] for v1 in range(3): for v2 in range(3): # imaginary parts should cancel taut_vvR[v1, v2] += f * (a_vR[v1].conj() * a_vR[v2]).real elif len(a_vR) == 0: pass else: raise RuntimeError('Parallelization issue') comm.sum(taut_vvR) taut_vvR = self.gd.distribute(taut_vvR) taut_xR[kpt.s, :, :] += taut_vvR def calculate_kinetic_energy_density(self): if self.kpt_u[0].f_n is None: return None taut_sR = self.gd.zeros(self.nspins) for kpt in self.kpt_u: self.add_to_kinetic_energy_density_kpt(kpt, kpt.psit_nG, taut_sR) self.kptband_comm.sum(taut_sR) for taut_R in taut_sR: self.kd.symmetry.symmetrize(taut_R, self.gd) return taut_sR def calculate_kinetic_energy_density_crossterms(self): if self.kpt_u[0].f_n is None: return None taut_svvR = self.gd.zeros((self.nspins, 3, 3)) for kpt in self.kpt_u: self.add_to_ke_crossterms_kpt(kpt, kpt.psit_nG, taut_svvR) self.kptband_comm.sum(taut_svvR) for taut_R in taut_svvR.reshape(-1, *taut_svvR.shape[-3:]): self.kd.symmetry.symmetrize(taut_R, self.gd) return taut_svvR def apply_mgga_orbital_dependent_hamiltonian(self, kpt, psit_xG, Htpsit_xG, dH_asp, dedtaut_R): N = len(psit_xG) S = self.gd.comm.size Q_G = self.pd.Q_qG[kpt.q] Gpsit_xG = np.empty((S,) + psit_xG.shape[1:], dtype=psit_xG.dtype) tmp_xG = np.empty((S,) + psit_xG.shape[1:], dtype=Htpsit_xG.dtype) G_Gv = self.pd.get_reciprocal_vectors(q=kpt.q) dedtaut_R = self.gd.collect(dedtaut_R, broadcast=True) for v in range(3): for n1 in range(0, N, S): n2 = min(n1 + S, N) dn = n2 - n1 Gpsit_xG[:dn] = 1j * G_Gv[:, v] * psit_xG[n1:n2] tmp_xG[:] = 0 Gpsit_G = self.pd.alltoall1(Gpsit_xG[:dn], kpt.q) if Gpsit_G is not None: a_R = self.pd.ifft(Gpsit_G, kpt.q, local=True, safe=False) a_R *= dedtaut_R self.pd.fftplan.execute() a_R = self.pd.tmp_Q.ravel()[Q_G] else: a_R = self.pd.tmp_G self.pd.alltoall2(a_R, kpt.q, tmp_xG[:dn]) axpy(-0.5, (1j * G_Gv[:, v] * tmp_xG[:dn]).ravel(), Htpsit_xG[n1:n2].ravel()) def _get_wave_function_array(self, u, n, realspace=True, periodic=False): kpt = self.kpt_u[u] psit_G = kpt.psit_nG[n] if realspace: psit_R = self.pd.ifft(psit_G, kpt.q) if self.kd.gamma or periodic: return psit_R k_c = self.kd.ibzk_kc[kpt.k] eikr_R = self.gd.plane_wave(k_c) return psit_R * eikr_R return psit_G
[docs] def get_wave_function_array(self, n, k, s, realspace=True, cut=True, periodic=False): kpt_rank, q = self.kd.get_rank_and_index(k) u = q * self.nspins + s band_rank, myn = self.bd.who_has(n) rank = self.world.rank if (self.kd.comm.rank == kpt_rank and self.bd.comm.rank == band_rank): psit_G = self._get_wave_function_array(u, myn, realspace, periodic) if realspace: psit_G = self.gd.collect(psit_G) else: assert not cut tmp_G = self.pd.gather(psit_G, self.kpt_u[u].q) if tmp_G is not None: ng = self.pd.ngmax if self.collinear: psit_G = np.zeros(ng, complex) else: psit_G = np.zeros((2, ng), complex) psit_G[..., :tmp_G.shape[-1]] = tmp_G if rank == 0: return psit_G # Domain master send this to the global master if self.gd.comm.rank == 0: self.world.ssend(psit_G, 0, 1398) if rank == 0: # allocate full wave function and receive shape = () if self.collinear else (2,) psit_G = self.empty(shape, global_array=True, realspace=realspace) # XXX this will fail when using non-standard nesting # of communicators. world_rank = (kpt_rank * self.gd.comm.size * self.bd.comm.size + band_rank * self.gd.comm.size) self.world.receive(psit_G, world_rank, 1398) return psit_G # We return a number instead of None on all the slaves. Most of # the time the return value will be ignored on the slaves, but # in some cases it will be multiplied by some other number and # then ignored. Allowing for this will simplify some code here # and there. return np.nan
def write(self, writer, write_wave_functions=False): FDPWWaveFunctions.write(self, writer) if not write_wave_functions: return if self.collinear: shape = (self.nspins, self.kd.nibzkpts, self.bd.nbands, self.pd.ngmax) else: shape = (self.kd.nibzkpts, self.bd.nbands, 2, self.pd.ngmax) writer.add_array('coefficients', shape, complex) c = Bohr**-1.5 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, realspace=False, cut=False) writer.fill(psit_G * c) writer.add_array('indices', (self.kd.nibzkpts, self.pd.ngmax), np.int32) if self.bd.comm.rank > 0: return Q_G = np.empty(self.pd.ngmax, np.int32) kk = 0 for r in range(self.kd.comm.size): for q, k in enumerate(self.kd.get_indices(r)): ng = self.ng_k[k] if r == self.kd.comm.rank: Q_G[:ng] = self.pd.Q_qG[q] if r > 0: self.kd.comm.send(Q_G, 0) if self.kd.comm.rank == 0: if r > 0: self.kd.comm.receive(Q_G, r) Q_G[ng:] = -1 writer.fill(Q_G) assert k == kk kk += 1 def read(self, reader): FDPWWaveFunctions.read(self, reader) if 'coefficients' not in reader.wave_functions: return Q_kG = reader.wave_functions.indices for kpt in self.kpt_u: if kpt.s == 0: Q_G = Q_kG[kpt.k] ng = self.ng_k[kpt.k] assert (Q_G[:ng] == self.pd.Q_qG[kpt.q]).all() assert (Q_G[ng:] == -1).all() c = reader.bohr**1.5 if reader.version < 0: c = 1 # old gpw file for kpt in self.kpt_u: ng = self.ng_k[kpt.k] index = (kpt.s, kpt.k) if self.collinear else (kpt.k,) psit_nG = reader.wave_functions.proxy('coefficients', *index) psit_nG.scale = c psit_nG.length_of_last_dimension = ng kpt.psit = PlaneWaveExpansionWaveFunctions( self.bd.nbands, self.pd, self.dtype, psit_nG, kpt=kpt.q, dist=(self.bd.comm, self.bd.comm.size), spin=kpt.s, collinear=self.collinear) if self.world.size > 1: # Read to memory: for kpt in self.kpt_u: kpt.psit.read_from_file() def hs(self, ham, q=-1, s=0, md=None): npw = len(self.pd.Q_qG[q]) N = self.pd.tmp_R.size if md is None: H_GG = np.zeros((npw, npw), complex) S_GG = np.zeros((npw, npw), complex) G1 = 0 G2 = npw else: H_GG = md.zeros(dtype=complex) S_GG = md.zeros(dtype=complex) if S_GG.size == 0: return H_GG, S_GG G1, G2 = next(md.my_blocks(S_GG))[:2] H_GG.ravel()[G1::npw + 1] = (0.5 * self.pd.gd.dv / N * self.pd.G2_qG[q][G1:G2]) for G in range(G1, G2): x_G = self.pd.zeros(q=q) x_G[G] = 1.0 H_GG[G - G1] += (self.pd.gd.dv / N * self.pd.fft(ham.vt_sG[s] * self.pd.ifft(x_G, q), q)) if ham.xc.type == 'MGGA': G_Gv = self.pd.get_reciprocal_vectors(q=q) for G in range(G1, G2): x_G = self.pd.zeros(q=q) x_G[G] = 1.0 for v in range(3): a_R = self.pd.ifft(1j * G_Gv[:, v] * x_G, q) H_GG[G - G1] += (self.pd.gd.dv / N * (-0.5) * 1j * G_Gv[:, v] * self.pd.fft(ham.xc.dedtaut_sG[s] * a_R, q)) S_GG.ravel()[G1::npw + 1] = self.pd.gd.dv / N f_GI = self.pt.expand(q) nI = f_GI.shape[1] dH_II = np.zeros((nI, nI)) dS_II = np.zeros((nI, nI)) I1 = 0 for a in self.pt.my_atom_indices: dH_ii = unpack(ham.dH_asp[a][s]) dS_ii = self.setups[a].dO_ii I2 = I1 + len(dS_ii) dH_II[I1:I2, I1:I2] = dH_ii / N**2 dS_II[I1:I2, I1:I2] = dS_ii / N**2 I1 = I2 H_GG += np.dot(f_GI[G1:G2].conj(), np.dot(dH_II, f_GI.T)) S_GG += np.dot(f_GI[G1:G2].conj(), np.dot(dS_II, f_GI.T)) return H_GG, S_GG @timer('Full diag') def diagonalize_full_hamiltonian(self, ham, atoms, log, nbands=None, ecut=None, scalapack=None, expert=False): if self.dtype != complex: raise ValueError( 'Please use mode=PW(..., force_complex_dtype=True)') if self.gd.comm.size > 1: raise ValueError( "Please use parallel={'domain': 1}") S = self.bd.comm.size if nbands is None and ecut is None: nbands = self.pd.ngmin // S * S elif nbands is None: ecut /= Ha # XXX I have seen this nbands expression elsewhere, # extract to function! nbands = int(self.gd.volume * ecut**1.5 * 2**0.5 / 3 / pi**2) if nbands % S != 0: nbands += S - nbands % S assert nbands <= self.pd.ngmin if expert: iu = nbands else: iu = None self.bd = bd = BandDescriptor(nbands, self.bd.comm) self.occupations.bd = bd log('Diagonalizing full Hamiltonian ({} lowest bands)'.format(nbands)) log('Matrix size (min, max): {}, {}'.format(self.pd.ngmin, self.pd.ngmax)) mem = 3 * self.pd.ngmax**2 * 16 / S / 1024**2 log('Approximate memory used per core to store H_GG, S_GG: {:.3f} MB' .format(mem)) log('Notice: Up to twice the amount of memory might be allocated\n' 'during diagonalization algorithm.') log('The least memory is required when the parallelization is purely\n' 'over states (bands) and not k-points, set ' "GPAW(..., parallel={'kpt': 1}, ...).") if S > 1: if isinstance(scalapack, (list, tuple)): nprow, npcol, b = scalapack assert nprow * npcol == S, (nprow, npcol, S) else: nprow = int(round(S**0.5)) while S % nprow != 0: nprow -= 1 npcol = S // nprow b = 64 log('ScaLapack grid: {}x{},'.format(nprow, npcol), 'block-size:', b) bg = BlacsGrid(bd.comm, S, 1) bg2 = BlacsGrid(bd.comm, nprow, npcol) scalapack = True else: scalapack = False self.set_positions(atoms.get_scaled_positions()) self.kpt_u[0].projections = None self.allocate_arrays_for_projections(self.pt.my_atom_indices) myslice = bd.get_slice() pb = ProgressBar(log.fd) nkpt = len(self.kpt_u) for u, kpt in enumerate(self.kpt_u): pb.update(u / nkpt) npw = len(self.pd.Q_qG[kpt.q]) if scalapack: mynpw = -(-npw // S) md = BlacsDescriptor(bg, npw, npw, mynpw, npw) md2 = BlacsDescriptor(bg2, npw, npw, b, b) else: md = md2 = MatrixDescriptor(npw, npw) with self.timer('Build H and S'): H_GG, S_GG = self.hs(ham, kpt.q, kpt.s, md) if scalapack: r = Redistributor(bd.comm, md, md2) H_GG = r.redistribute(H_GG) S_GG = r.redistribute(S_GG) psit_nG = md2.empty(dtype=complex) eps_n = np.empty(npw) with self.timer('Diagonalize'): if not scalapack: md2.general_diagonalize_dc(H_GG, S_GG, psit_nG, eps_n, iu=iu) else: md2.general_diagonalize_dc(H_GG, S_GG, psit_nG, eps_n) if eps_n[0] < -1000: msg = f"""Lowest eigenvalue is {eps_n[0]} Hartree. You might be suffering from MKL library bug MKLD-11440. See issue #241 in GPAW. Creashing to prevent corrupted results.""" raise RuntimeError(msg) del H_GG, S_GG kpt.eps_n = eps_n[myslice].copy() if scalapack: md3 = BlacsDescriptor(bg, npw, npw, bd.maxmynbands, npw) r = Redistributor(bd.comm, md2, md3) psit_nG = r.redistribute(psit_nG) kpt.psit = PlaneWaveExpansionWaveFunctions( self.bd.nbands, self.pd, self.dtype, psit_nG[:bd.mynbands].copy(), kpt=kpt.q, dist=(self.bd.comm, self.bd.comm.size), spin=kpt.s, collinear=self.collinear) del psit_nG with self.timer('Projections'): self.pt.integrate(kpt.psit_nG, kpt.P_ani, kpt.q) kpt.f_n = None pb.finish() self.calculate_occupation_numbers() return nbands
[docs] def initialize_from_lcao_coefficients(self, basis_functions: BasisFunctions, block_size: int = 10) -> None: """Convert from LCAO to PW coefficients.""" nlcao = len(self.kpt_qs[0][0].C_nM) # We go from LCAO to real-space and then to PW's. # It's too expensive to allocate one big real-space array: block_size = min(max(nlcao, 1), block_size) psit_nR = self.gd.empty(block_size, self.dtype) for kpt in self.kpt_u: if self.kd.gamma: emikr_R = 1.0 else: k_c = self.kd.ibzk_kc[kpt.k] emikr_R = self.gd.plane_wave(-k_c) kpt.psit = PlaneWaveExpansionWaveFunctions( self.bd.nbands, self.pd, self.dtype, kpt=kpt.q, dist=(self.bd.comm, -1, 1), spin=kpt.s, collinear=self.collinear) psit_nG = kpt.psit.array if psit_nG.ndim == 3: # non-collinear calculation N, S, G = psit_nG.shape psit_nG = psit_nG.reshape((N * S, G)) for n1 in range(0, nlcao, block_size): n2 = min(n1 + block_size, nlcao) psit_nR[:] = 0.0 basis_functions.lcao_to_grid(kpt.C_nM[n1:n2], psit_nR[:n2 - n1], kpt.q, block_size) for psit_R, psit_G in zip(psit_nR, psit_nG[n1:n2]): psit_G[:] = self.pd.fft(psit_R * emikr_R, kpt.q) kpt.C_nM = None
def random_wave_functions(self, mynao): rs = np.random.RandomState(self.world.rank) for kpt in self.kpt_u: if kpt.psit is None: kpt.psit = PlaneWaveExpansionWaveFunctions( self.bd.nbands, self.pd, self.dtype, kpt=kpt.q, dist=(self.bd.comm, -1, 1), spin=kpt.s, collinear=self.collinear) array = kpt.psit.array[mynao:] weight_G = 1.0 / (1.0 + self.pd.G2_qG[kpt.q]) array.real = rs.uniform(-1, 1, array.shape) * weight_G array.imag = rs.uniform(-1, 1, array.shape) * weight_G def estimate_memory(self, mem): FDPWWaveFunctions.estimate_memory(self, mem) self.pd.estimate_memory(mem.subnode('PW-descriptor')) def get_kinetic_stress(self): sigma_vv = np.zeros((3, 3), dtype=complex) pd = self.pd dOmega = pd.gd.dv / pd.gd.N_c.prod() if pd.dtype == float: dOmega *= 2 K_qv = self.pd.K_qv for kpt in self.kpt_u: G_Gv = pd.get_reciprocal_vectors(q=kpt.q, add_q=False) psit2_G = 0.0 for n, f in enumerate(kpt.f_n): psit2_G += f * np.abs(kpt.psit_nG[n])**2 for alpha in range(3): Ga_G = G_Gv[:, alpha] + K_qv[kpt.q, alpha] for beta in range(3): Gb_G = G_Gv[:, beta] + K_qv[kpt.q, beta] sigma_vv[alpha, beta] += (psit2_G * Ga_G * Gb_G).sum() sigma_vv *= -dOmega self.world.sum(sigma_vv) return sigma_vv