Source code for gpaw.wavefunctions.fd

import numpy as np
from ase.units import Bohr

from gpaw.fd_operators import Laplace, Gradient
from gpaw.kpoint import KPoint
from gpaw.kpt_descriptor import KPointDescriptor
from gpaw.lfc import LocalizedFunctionsCollection as LFC
from gpaw.mpi import serial_comm
from gpaw.preconditioner import Preconditioner
from gpaw.projections import Projections
from gpaw.transformers import Transformer
from gpaw.utilities.blas import axpy
from gpaw.wavefunctions.arrays import UniformGridWaveFunctions
from gpaw.wavefunctions.fdpw import FDPWWaveFunctions
from gpaw.wavefunctions.mode import Mode
import _gpaw


class FD(Mode):
    name = 'fd'

    def __init__(self, nn=3, interpolation=3, force_complex_dtype=False):
        self.nn = nn
        self.interpolation = interpolation
        Mode.__init__(self, force_complex_dtype)

    def __call__(self, *args, **kwargs):
        return FDWaveFunctions(self.nn, *args, **kwargs)

    def todict(self):
        dct = Mode.todict(self)
        dct['nn'] = self.nn
        dct['interpolation'] = self.interpolation
        return dct


[docs]class FDWaveFunctions(FDPWWaveFunctions): mode = 'fd' def __init__(self, stencil, parallel, initksl, gd, nvalence, setups, bd, dtype, world, kd, kptband_comm, timer, reuse_wfs_method=None, collinear=True): 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) # Kinetic energy operator: self.kin = Laplace(self.gd, -0.5, stencil, self.dtype) 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, atom_partition=None): FDPWWaveFunctions.set_positions(self, spos_ac, atom_partition) def __str__(self): s = 'Wave functions: Uniform real-space grid\n' s += ' Kinetic energy operator: %s\n' % self.kin.description return s + FDPWWaveFunctions.__str__(self) def make_preconditioner(self, block=1): return Preconditioner(self.gd, self.kin, self.dtype, block) def apply_pseudo_hamiltonian(self, kpt, ham, psit_xG, Htpsit_xG): self.timer.start('Apply hamiltonian') self.kin.apply(psit_xG, Htpsit_xG, kpt.phase_cd) ham.apply_local_potential(psit_xG, Htpsit_xG, kpt.s) ham.xc.apply_orbital_dependent_hamiltonian( kpt, psit_xG, Htpsit_xG, ham.dH_asp) self.timer.stop('Apply hamiltonian') def get_pseudo_partial_waves(self): phit_aj = [setup.get_partial_waves_for_atomic_orbitals() for setup in self.setups] return LFC(self.gd, phit_aj, kd=self.kd, cut=True, dtype=self.dtype) 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] for f, psit_G in zip(f_n, kpt.psit_nG): # Same as nt_G += f * abs(psit_G)**2, but much faster: _gpaw.add_to_density(f, psit_G, 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 not isinstance(self.kpt_u[0].psit_nG, np.ndarray): return None 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.kptband_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)
[docs] 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_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_qs = [] kpt_u = [] for k in range(kd.nbzkpts): kpt_s = [] for s in range(self.nspins): # Index of symmetry related point in the IBZ ik = self.kd.bz2ibz_k[k] r, q = self.kd.get_rank_and_index(ik) assert r == 0 kpt = self.kpt_qs[q][s] phase_cd = np.exp(2j * np.pi * self.gd.sdisp_cd * kd.bzk_kc[k, :, np.newaxis]) # New k-point: kpt2 = KPoint(1.0 / kd.nbzkpts, 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 = UniformGridWaveFunctions( self.bd.nbands, self.gd, self.dtype, kpt=k, dist=(self.bd.comm, self.bd.comm.size), spin=kpt.s, collinear=True) self.gd.distribute(Psit_nG, kpt2.psit_nG) # Calculate PAW projections: nproj_a = [setup.ni for setup in self.setups] kpt2.projections = Projections( self.bd.nbands, nproj_a, kpt.projections.atom_partition, self.bd.comm, collinear=True, spin=s, dtype=self.dtype) kpt2.psit.matrix_elements(self.pt, out=kpt2.projections) kpt_s.append(kpt2) kpt_u.append(kpt2) kpt_qs.append(kpt_s) self.kd = kd self.kpt_qs = kpt_qs self.kpt_u = kpt_u
def _get_wave_function_array(self, u, n, realspace=True, periodic=False): assert realspace kpt = self.kpt_u[u] psit_G = kpt.psit_nG[n] if periodic and self.dtype == complex: k_c = self.kd.ibzk_kc[kpt.k] return self.gd.plane_wave(-k_c) * psit_G return psit_G def write(self, writer, write_wave_functions=False): FDPWWaveFunctions.write(self, writer) if not write_wave_functions: return writer.add_array( 'values', (self.nspins, self.kd.nibzkpts, self.bd.nbands) + tuple(self.gd.get_size_of_global_array()), self.dtype) 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 * Bohr**-1.5) def read(self, reader): FDPWWaveFunctions.read(self, reader) if 'values' not in reader.wave_functions: return c = reader.bohr**1.5 if reader.version < 0: c = 1 # old gpw file for kpt in self.kpt_u: # 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: psit_nG = reader.wave_functions.proxy('values', kpt.s, kpt.k) psit_nG.scale = c kpt.psit = UniformGridWaveFunctions( self.bd.nbands, self.gd, self.dtype, psit_nG, kpt=kpt.q, dist=(self.bd.comm, self.bd.comm.size), spin=kpt.s, collinear=True) if self.world.size > 1: # Read to memory: for kpt in self.kpt_u: kpt.psit.read_from_file() def initialize_from_lcao_coefficients(self, basis_functions): for kpt in self.kpt_u: kpt.psit = UniformGridWaveFunctions( self.bd.nbands, self.gd, self.dtype, kpt=kpt.q, dist=(self.bd.comm, self.bd.comm.size, 1), spin=kpt.s, collinear=True) kpt.psit_nG[:] = 0.0 mynbands = len(kpt.C_nM) basis_functions.lcao_to_grid(kpt.C_nM, kpt.psit_nG[:mynbands], kpt.q) kpt.C_nM = None
[docs] 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)