Example #1
0
 def test_something(self):
     laplace_uG = np.empty_like(self.laplace0_uG)
     op = Laplace(self.gd, dtype=self.dtype)
     for myu, laplace_G in enumerate(laplace_uG):
         phase_cd = {float:None, complex:self.phase_ucd[myu]}[self.dtype]
         op.apply(self.wf_uG[myu], laplace_G, phase_cd)
         print 'myu:', myu, 'diff:', np.std(laplace_G-self.laplace0_uG[myu]), '/', np.abs(laplace_G-self.laplace0_uG[myu]).max()
Example #2
0
 def test_something(self):
     laplace_uG = np.empty_like(self.laplace0_uG)
     op = Laplace(self.gd, dtype=self.dtype)
     for myu, laplace_G in enumerate(laplace_uG):
         phase_cd = {float:None, complex:self.phase_ucd[myu]}[self.dtype]
         op.apply(self.wf_uG[myu], laplace_G, phase_cd)
         print 'myu:', myu, 'diff:', np.std(laplace_G-self.laplace0_uG[myu]), '/', np.abs(laplace_G-self.laplace0_uG[myu]).max()
Example #3
0
 def apply_t(self):
     """Apply kinetic energy operator and return new object."""
     p = 2  # padding
     newsize_c = self.size_c + 2 * p
     gd = GridDescriptor(N_c=newsize_c + 1, cell_cv=self.gd.h_c * (newsize_c + 1), pbc_c=False, comm=mpi.serial_comm)
     T = Laplace(gd, scale=1 / 2.0, n=p)
     f_ig = np.zeros((len(self.f_iG),) + tuple(newsize_c))
     f_ig[:, p:-p, p:-p, p:-p] = self.f_iG
     Tf_iG = np.empty_like(f_ig)
     T.apply(f_ig, Tf_iG)
     return LocalizedFunctions(self.gd, Tf_iG, self.corner_c - p, self.index)
Example #4
0
 def apply_t(self):
     """Apply kinetic energy operator and return new object."""
     p = 2  # padding
     newsize_c = self.size_c + 2 * p
     gd = GridDescriptor(N_c=newsize_c + 1,
                         cell_cv=self.gd.h_c * (newsize_c + 1),
                         pbc_c=False,
                         comm=mpi.serial_comm)
     T = Laplace(gd, scale =1/2., n=p)
     f_ig = np.zeros((len(self.f_iG),) + tuple(newsize_c))
     f_ig[:, p:-p, p:-p, p:-p] = self.f_iG
     Tf_iG = np.empty_like(f_ig)
     T.apply(f_ig, Tf_iG)
     return LocalizedFunctions(self.gd, Tf_iG, self.corner_c - p,
                               self.index)
Example #5
0
class Preconditioner:
    def __init__(self, gd0, kin0, dtype=float, block=1):
        gd1 = gd0.coarsen()
        gd2 = gd1.coarsen()
        self.kin0 = kin0
        self.kin1 = Laplace(gd1, -0.5, 1, dtype)
        self.kin2 = Laplace(gd2, -0.5, 1, dtype)
        self.scratch0 = gd0.zeros((2, block), dtype, False)
        self.scratch1 = gd1.zeros((3, block), dtype, False)
        self.scratch2 = gd2.zeros((3, block), dtype, False)
        self.step = 0.66666666 / kin0.get_diagonal_element()

        self.restrictor_object0 = Transformer(gd0, gd1, 1, dtype)
        self.restrictor_object1 = Transformer(gd1, gd2, 1, dtype)
        self.interpolator_object2 = Transformer(gd2, gd1, 1, dtype)
        self.interpolator_object1 = Transformer(gd1, gd0, 1, dtype)
        self.restrictor0 = self.restrictor_object0.apply
        self.restrictor1 = self.restrictor_object1.apply
        self.interpolator2 = self.interpolator_object2.apply
        self.interpolator1 = self.interpolator_object1.apply

    def calculate_kinetic_energy(self, psit_xG, kpt):
        return None

    def __call__(self, residuals, kpt, ekin=None, out=None):
        if residuals.ndim == 3:
            if out is None:
                return self.__call__(residuals[np.newaxis], kpt)[0]
            return self.__call__(residuals[np.newaxis], kpt,
                                 out=out[np.newaxis])[0]
        nb = len(residuals)  # number of bands
        phases = kpt.phase_cd
        step = self.step
        if out is None:
            d0, q0 = self.scratch0[:, :nb]
        else:
            d0 = out
            q0 = self.scratch0[0, :nb]
        r1, d1, q1 = self.scratch1[:, :nb]
        r2, d2, q2 = self.scratch2[:, :nb]
        self.restrictor0(-residuals, r1, phases)
        d1[:] = 4 * step * r1
        self.kin1.apply(d1, q1, phases)
        q1 -= r1
        self.restrictor1(q1, r2, phases)
        d2 = 16 * step * r2
        self.kin2.apply(d2, q2, phases)
        q2 -= r2
        d2 -= 16 * step * q2
        self.interpolator2(d2, q1, phases)
        d1 -= q1
        self.kin1.apply(d1, q1, phases)
        q1 -= r1
        d1 -= 4 * step * q1
        self.interpolator1(-d1, d0, phases)
        self.kin0.apply(d0, q0, phases)
        q0 -= residuals
        axpy(-step, q0, d0)  # d0 -= step * q0
        d0 *= -1.0
        return d0
Example #6
0
class Preconditioner:
    def __init__(self, gd0, kin0, dtype=float, block=1):
        gd1 = gd0.coarsen()
        gd2 = gd1.coarsen()
        self.kin0 = kin0
        self.kin1 = Laplace(gd1, -0.5, 1, dtype)
        self.kin2 = Laplace(gd2, -0.5, 1, dtype)
        self.scratch0 = gd0.zeros((2, block), dtype, False)
        self.scratch1 = gd1.zeros((3, block),dtype, False)
        self.scratch2 = gd2.zeros((3, block), dtype, False)
        self.step = 0.66666666 / kin0.get_diagonal_element()

        self.restrictor_object0 = Transformer(gd0, gd1, 1,dtype, False)
        self.restrictor_object1 = Transformer(gd1, gd2, 1, dtype, False)
        self.interpolator_object2 = Transformer(gd2, gd1, 1, dtype, False)
        self.interpolator_object1 = Transformer(gd1, gd0, 1, dtype, False)
        self.restrictor0 = self.restrictor_object0.apply
        self.restrictor1 = self.restrictor_object1.apply
        self.interpolator2 = self.interpolator_object2.apply
        self.interpolator1 = self.interpolator_object1.apply
        self.allocated = False

    def allocate(self):
        assert not self.allocated
        for transformer in [self.restrictor_object0,
                            self.restrictor_object1,
                            self.interpolator_object2,
                            self.interpolator_object1]:
            transformer.allocate()
        self.allocated = True
        
    def __call__(self, residuals, kpt):
        nb = len(residuals) # number of bands
        phases = kpt.phase_cd
        step = self.step
        d0, q0 = self.scratch0[:,:nb]
        r1, d1, q1 = self.scratch1[:, :nb]
        r2, d2, q2 = self.scratch2[:, :nb]
        self.restrictor0(-residuals, r1, phases)
        d1[:] = 4 * step * r1
        self.kin1.apply(d1, q1, phases)
        q1 -= r1
        self.restrictor1(q1, r2, phases)
        d2 = 16 * step * r2
        self.kin2.apply(d2, q2, phases)
        q2 -= r2
        d2 -= 16 * step * q2
        self.interpolator2(d2, q1, phases)
        d1 -= q1
        self.kin1.apply(d1, q1, phases)
        q1 -= r1
        d1 -= 4 * step * q1
        self.interpolator1(-d1, d0, phases)
        self.kin0.apply(d0, q0, phases)
        q0 -= residuals
        axpy(-step, q0, d0)  # d0 -= step * q0
        d0 *= -1.0
        return d0
Example #7
0
class Preconditioner:
    def __init__(self, gd0, kin0, dtype=float, block=1):
        gd1 = gd0.coarsen()
        gd2 = gd1.coarsen()
        self.kin0 = kin0
        self.kin1 = Laplace(gd1, -0.5, 1, dtype)
        self.kin2 = Laplace(gd2, -0.5, 1, dtype)
        self.scratch0 = gd0.zeros((2, block), dtype, False)
        self.scratch1 = gd1.zeros((3, block), dtype, False)
        self.scratch2 = gd2.zeros((3, block), dtype, False)
        self.step = 0.66666666 / kin0.get_diagonal_element()

        self.restrictor_object0 = Transformer(gd0, gd1, 1, dtype, False)
        self.restrictor_object1 = Transformer(gd1, gd2, 1, dtype, False)
        self.interpolator_object2 = Transformer(gd2, gd1, 1, dtype, False)
        self.interpolator_object1 = Transformer(gd1, gd0, 1, dtype, False)
        self.restrictor0 = self.restrictor_object0.apply
        self.restrictor1 = self.restrictor_object1.apply
        self.interpolator2 = self.interpolator_object2.apply
        self.interpolator1 = self.interpolator_object1.apply
        self.allocated = False

    def allocate(self):
        assert not self.allocated
        for transformer in [
                self.restrictor_object0, self.restrictor_object1,
                self.interpolator_object2, self.interpolator_object1
        ]:
            transformer.allocate()
        self.allocated = True

    def __call__(self, residuals, kpt):
        nb = len(residuals)  # number of bands
        phases = kpt.phase_cd
        step = self.step
        d0, q0 = self.scratch0[:, :nb]
        r1, d1, q1 = self.scratch1[:, :nb]
        r2, d2, q2 = self.scratch2[:, :nb]
        self.restrictor0(-residuals, r1, phases)
        d1[:] = 4 * step * r1
        self.kin1.apply(d1, q1, phases)
        q1 -= r1
        self.restrictor1(q1, r2, phases)
        d2 = 16 * step * r2
        self.kin2.apply(d2, q2, phases)
        q2 -= r2
        d2 -= 16 * step * q2
        self.interpolator2(d2, q1, phases)
        d1 -= q1
        self.kin1.apply(d1, q1, phases)
        q1 -= r1
        d1 -= 4 * step * q1
        self.interpolator1(-d1, d0, phases)
        self.kin0.apply(d0, q0, phases)
        q0 -= residuals
        axpy(-step, q0, d0)  # d0 -= step * q0
        d0 *= -1.0
        return d0
Example #8
0
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)
Example #9
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)
Example #10
0
class SternheimerOperator:
    """Class implementing the linear operator in the Sternheimer equation.
    
    The main purpose of this class is to implement the multiplication of the
    linear operator (H - eps_nk) with a vector in the ``apply`` member
    function. An additional ``matvec`` member function has been defined so that
    instances of this class can be passed as a linear operator to scipy's
    iterative Krylov solvers in ``scipy.sparse.linalg``.

    """
    
    def __init__(self, hamiltonian, wfs, gd, dtype=float):
        """Save useful objects for the Sternheimer operator.

        Parameters
        ----------
        hamiltonian: Hamiltonian
            Hamiltonian for a ground-state calculation.
        wfs: WaveFunctions
            Ground-state wave-functions.
        gd: GridDescriptor
            Grid on which the operator is defined.
        dtype: dtype
            dtype of the wave-function being operated on.

        """

        self.hamiltonian = hamiltonian
        self.kin = Laplace(gd, scale=-0.5, n=3, dtype=dtype)
        self.kpt_u = wfs.kpt_u
        self.pt = wfs.pt
        self.gd = gd

        # Variables for k-point and band index
        self.k = None
        self.n = None
        self.kplusq = None
        
        # For scipy's linear solver
        N = np.prod(gd.n_c)
        self.shape = (N, N)
        self.dtype = dtype

    def set_k(self, k):
        """Set k-vector for the Bloch-state in consideration."""

        self.k = k
        
    def set_band(self, n):
        """Set band index for the Bloch-state in consideration.

        Note that n different from None has implications for the interpretation
        of the input vector in the ``apply`` member function.

        """

        self.n = n

    def set_kplusq(self, kplusq):
        """Set q-vector of the perturbing potential.

        Parameters
        ----------
        kplusq: int
           Index of the k+q vector.

        """

        self.kplusq = kplusq

    def apply(self, x_nG, y_nG):
        """Apply the linear operator in the Sternheimer equation to a vector.

        For the eigenvalue term the k-point is the one of the state.
        For the other terms the k-point to be used is the one given by the k+q
        phase of the first-order of the state. Only for q=0 do the two coincide.
        
        Parameters
        ----------
        x_nG: ndarray
            Vector(s) to which the Sternheimer operator is applied.
        y_nG: ndarray
            Resulting vector(s).
            
        """

        assert x_nG.ndim in (3, 4)
        assert x_nG.shape == y_nG.shape
        assert tuple(self.gd.n_c) == x_nG.shape[-3:]
        assert self.k is not None

        # k
        kpt = self.kpt_u[self.k]
        # k+q
        kplusqpt = self.kpt_u[self.kplusq]
        
        # Kintetic energy
        # k+q
        self.kin.apply(x_nG, y_nG, phase_cd=kplusqpt.phase_cd)

        # Local part of effective potential - no phase !!
        self.hamiltonian.apply_local_potential(x_nG, y_nG, kpt.s)
        
        # Non-local part from projectors (coefficients can not be reused)
        shape = x_nG.shape[:-3]
        P_ani = self.pt.dict(shape=shape)

        # k+q 
        self.pt.integrate(x_nG, P_ani, q=kplusqpt.k)

        for a, P_ni in P_ani.items():
            dH_ii = unpack(self.hamiltonian.dH_asp[a][kpt.s])
            P_ani[a] = np.dot(P_ni, dH_ii)
        # k+q
        self.pt.add(y_nG, P_ani, q=kplusqpt.k)

        # XXX Eigenvalue term
        if self.n is not None:
            # k
            y_nG -= kpt.eps_n[self.n] * x_nG
        else:
            for n, x_G in enumerate(x_nG):
                # k
                y_nG[n] -= kpt.eps_n[n] * x_G

        # Project out undesired (numerical) components
        # k+q
        self.project(y_nG)

    def project(self, x_nG):
        """Project the vector onto the unoccupied states at k+q.

        The projection operator is defined as follows::

                      --                    --             
             P      = >  |Psi ><Psi | = 1 - >  |Psi ><Psi |
              c,k+q   --     c     c        --     v     v 
                       c    k+q   k+q        v    k+q   k+q
        
        """

        # It might be a good idea to move this functionality to its own class
        assert x_nG.ndim == 4
        assert self.kplusq is not None

        # k+q
        kplusqpt = self.kpt_u[self.kplusq]
        # Occupied wave function
        psit_nG = kplusqpt.psit_nG

        # Project out occupied sub-space using blas routine gemm
        m = len(x_nG)
        n = len(psit_nG)
        proj_mn = np.zeros((m, n), dtype=self.dtype)
        gemm(self.gd.dv, psit_nG, x_nG, 0.0, proj_mn, 'c')
        gemm(-1.0, psit_nG, proj_mn, 1.0, x_nG)

        # Project out one orbital at a time
        ## for n, psit_G in enumerate(psit_nG):
        ## 
        ##     proj = self.gd.integrate(psit_G.conjugate() * x_nG)
        ##     x_nG -= proj * psit_G
        
    def matvec(self, x):
        """Matrix-vector multiplication for scipy's Krylov solvers.

        This is a wrapper around the ``apply`` member function above. It allows
        to multiply the sternheimer operator onto the 1-dimensional scipy
        representation of a gpaw grid vector(s).
        
        Parameters
        ----------
        a: ndarray
            1-dimensional array holding the representation of a gpaw grid
            vector (possibly a set of vectors).

        """

        # Check that a is 1-dimensional
        assert len(x.shape) == 1
        
        # Find the number of states in x
        grid_shape = tuple(self.gd.n_c)
        assert ((x.size % np.prod(grid_shape)) == 0), ("Incompatible array " +
                                                       "shapes")
        # Number of states in the vector
        N = x.size / np.prod(grid_shape)
        
        # Output array
        y_nG = self.gd.zeros(n=N, dtype=self.dtype)
        shape = y_nG.shape

        assert x.size == np.prod(y_nG.shape)
        
        x_nG = x.reshape(shape)

        self.apply(x_nG, y_nG)
        
        y = y_nG.ravel()
        
        return y
Example #11
0
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'))
Example #12
0
File: fd.py Project: qsnake/gpaw
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'))
Example #13
0
class SternheimerOperator:
    """Class implementing the linear operator in the Sternheimer equation.
    
    The main purpose of this class is to implement the multiplication of the
    linear operator (H - eps_nk) with a vector in the ``apply`` member
    function. An additional ``matvec`` member function has been defined so that
    instances of this class can be passed as a linear operator to scipy's
    iterative Krylov solvers in ``scipy.sparse.linalg``.

    """
    def __init__(self, hamiltonian, wfs, gd, dtype=float):
        """Save useful objects for the Sternheimer operator.

        Parameters
        ----------
        hamiltonian: Hamiltonian
            Hamiltonian for a ground-state calculation.
        wfs: WaveFunctions
            Ground-state wave-functions.
        gd: GridDescriptor
            Grid on which the operator is defined.
        dtype: dtype
            dtype of the wave-function being operated on.

        """

        self.hamiltonian = hamiltonian
        self.kin = Laplace(gd, scale=-0.5, n=3, dtype=dtype)
        self.kpt_u = wfs.kpt_u
        self.pt = wfs.pt
        self.gd = gd

        # Variables for k-point and band index
        self.k = None
        self.n = None
        self.kplusq = None

        # For scipy's linear solver
        N = np.prod(gd.n_c)
        self.shape = (N, N)
        self.dtype = dtype

    def set_k(self, k):
        """Set k-vector for the Bloch-state in consideration."""

        self.k = k

    def set_band(self, n):
        """Set band index for the Bloch-state in consideration.

        Note that n different from None has implications for the interpretation
        of the input vector in the ``apply`` member function.

        """

        self.n = n

    def set_kplusq(self, kplusq):
        """Set q-vector of the perturbing potential.

        Parameters
        ----------
        kplusq: int
           Index of the k+q vector.

        """

        self.kplusq = kplusq

    def apply(self, x_nG, y_nG):
        """Apply the linear operator in the Sternheimer equation to a vector.

        For the eigenvalue term the k-point is the one of the state.
        For the other terms the k-point to be used is the one given by the k+q
        phase of the first-order of the state. Only for q=0 do the two coincide.
        
        Parameters
        ----------
        x_nG: ndarray
            Vector(s) to which the Sternheimer operator is applied.
        y_nG: ndarray
            Resulting vector(s).
            
        """

        assert x_nG.ndim in (3, 4)
        assert x_nG.shape == y_nG.shape
        assert tuple(self.gd.n_c) == x_nG.shape[-3:]
        assert self.k is not None

        # k
        kpt = self.kpt_u[self.k]
        # k+q
        kplusqpt = self.kpt_u[self.kplusq]

        # Kintetic energy
        # k+q
        self.kin.apply(x_nG, y_nG, phase_cd=kplusqpt.phase_cd)

        # Local part of effective potential - no phase !!
        self.hamiltonian.apply_local_potential(x_nG, y_nG, kpt.s)

        # Non-local part from projectors (coefficients can not be reused)
        shape = x_nG.shape[:-3]
        P_ani = self.pt.dict(shape=shape)

        # k+q
        self.pt.integrate(x_nG, P_ani, q=kplusqpt.k)

        for a, P_ni in P_ani.items():
            dH_ii = unpack(self.hamiltonian.dH_asp[a][kpt.s])
            P_ani[a] = np.dot(P_ni, dH_ii)
        # k+q
        self.pt.add(y_nG, P_ani, q=kplusqpt.k)

        # XXX Eigenvalue term
        if self.n is not None:
            # k
            y_nG -= kpt.eps_n[self.n] * x_nG
        else:
            for n, x_G in enumerate(x_nG):
                # k
                y_nG[n] -= kpt.eps_n[n] * x_G

        # Project out undesired (numerical) components
        # k+q
        self.project(y_nG)

    def project(self, x_nG):
        """Project the vector onto the unoccupied states at k+q.

        The projection operator is defined as follows::

                      --                    --             
             P      = >  |Psi ><Psi | = 1 - >  |Psi ><Psi |
              c,k+q   --     c     c        --     v     v 
                       c    k+q   k+q        v    k+q   k+q
        
        """

        # It might be a good idea to move this functionality to its own class
        assert x_nG.ndim == 4
        assert self.kplusq is not None

        # k+q
        kplusqpt = self.kpt_u[self.kplusq]
        # Occupied wave function
        psit_nG = kplusqpt.psit_nG

        # Project out occupied sub-space using blas routine gemm
        m = len(x_nG)
        n = len(psit_nG)
        proj_mn = np.zeros((m, n), dtype=self.dtype)
        gemm(self.gd.dv, psit_nG, x_nG, 0.0, proj_mn, 'c')
        gemm(-1.0, psit_nG, proj_mn, 1.0, x_nG)

        # Project out one orbital at a time
        ## for n, psit_G in enumerate(psit_nG):
        ##
        ##     proj = self.gd.integrate(psit_G.conjugate() * x_nG)
        ##     x_nG -= proj * psit_G

    def matvec(self, x):
        """Matrix-vector multiplication for scipy's Krylov solvers.

        This is a wrapper around the ``apply`` member function above. It allows
        to multiply the sternheimer operator onto the 1-dimensional scipy
        representation of a gpaw grid vector(s).
        
        Parameters
        ----------
        a: ndarray
            1-dimensional array holding the representation of a gpaw grid
            vector (possibly a set of vectors).

        """

        # Check that a is 1-dimensional
        assert len(x.shape) == 1

        # Find the number of states in x
        grid_shape = tuple(self.gd.n_c)
        assert ((x.size % np.prod(grid_shape)) == 0), ("Incompatible array " +
                                                       "shapes")
        # Number of states in the vector
        N = x.size / np.prod(grid_shape)

        # Output array
        y_nG = self.gd.zeros(n=N, dtype=self.dtype)
        shape = y_nG.shape

        assert x.size == np.prod(y_nG.shape)

        x_nG = x.reshape(shape)

        self.apply(x_nG, y_nG)

        y = y_nG.ravel()

        return y
Example #14
0
File: fd.py Project: thonmaker/gpaw
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)

    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_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.mykpts[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 = 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_u.append(kpt2)

        self.kd = kd
        self.mykpts = 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.mykpts:
            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

    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)
Example #15
0
class ExtraVacuumPoissonSolver(_PoissonSolver):
    """Wrapper around PoissonSolver extending the vacuum size.

       Poisson equation is solved on the large grid defined by `gpts`
       using `poissonsolver_large`.

       If `coarses` is given, then the large grid is first coarsed
       `coarses` times from to the original grid. Then, the coarse
       potential is used to correct the boundary conditions
       of the potential calculated on the original (small, fine)
       grid by `poissonsolver_small`.

       The parameters `nn_*` control the finite difference stencils
       used in the coarsening, refining, and Laplace.

       If the parameter `use_aux_grid` is `True`, an auxiliary
       medium-sized grid is used between the large and small grids.
       The parameter does not affect the result but can be used to
       achieve speed-up depending on the grid sizes.
       """

    def __init__(self, gpts, poissonsolver_large,
                 poissonsolver_small=None, coarses=0,
                 nn_coarse=3, nn_refine=3, nn_laplace=3,
                 use_aux_grid=True):
        # TODO: Alternative options: vacuum size and h
        self.N_large_fine_c = np.array(gpts, dtype=int)
        self.Ncoar = coarses  # coar == coarse
        if self.Ncoar > 0:
            self.use_coarse = True
        else:
            self.use_coarse = False
        self.ps_large_coar = create_poisson_solver(poissonsolver_large)
        if self.use_coarse:
            self.ps_small_fine = create_poisson_solver(poissonsolver_small)
        else:
            assert poissonsolver_small is None
        self.nn_coarse = nn_coarse
        self.nn_refine = nn_refine
        self.nn_laplace = nn_laplace
        self.use_aux_grid = use_aux_grid
        self._initialized = False

    def set_grid_descriptor(self, gd):
        # If non-periodic boundary conditions is used,
        # there is problems with auxiliary grid.
        # Maybe with use_aux_grid=False it would work?
        if gd.pbc_c.any():
            raise NotImplementedError('Only non-periodic boundary '
                                      'conditions are tested')

        self.gd_small_fine = gd
        assert np.all(self.gd_small_fine.N_c <= self.N_large_fine_c), \
            'extended grid has to be larger than the original one'

        if self.use_coarse:
            # 1.1. Construct coarse chain on the small grid
            self.coarser_i = []
            gd = self.gd_small_fine
            N_c = self.N_large_fine_c.copy()
            for i in range(self.Ncoar):
                gd2 = gd.coarsen()
                self.coarser_i.append(Transformer(gd, gd2, self.nn_coarse))
                N_c //= 2
                gd = gd2
            self.gd_small_coar = gd
        else:
            self.gd_small_coar = self.gd_small_fine
            N_c = self.N_large_fine_c

        # 1.2. Construct coarse extended grid
        self.gd_large_coar = ext_gd(self.gd_small_coar, N_c=N_c)

        # Initialize poissonsolvers
        self.ps_large_coar.set_grid_descriptor(self.gd_large_coar)
        if not self.use_coarse:
            return
        self.ps_small_fine.set_grid_descriptor(self.gd_small_fine)

        if self.use_aux_grid:
            # 2.1. Construct an auxiliary grid that is the small grid plus
            # a buffer region allowing Laplace and refining
            # with the used stencils
            buf = self.nn_refine
            for i in range(self.Ncoar):
                buf = 2 * buf + self.nn_refine
            buf += self.nn_laplace
            div = 2**self.Ncoar
            if buf % div != 0:
                buf += div - buf % div
            N_c = self.gd_small_fine.N_c + 2 * buf
            if np.any(N_c > self.N_large_fine_c):
                self.use_aux_grid = False
                N_c = self.N_large_fine_c
            self.gd_aux_fine = ext_gd(self.gd_small_fine, N_c=N_c)
        else:
            self.gd_aux_fine = ext_gd(self.gd_small_fine,
                                      N_c=self.N_large_fine_c)

        # 2.2 Construct Laplace on the aux grid
        self.laplace_aux_fine = Laplace(self.gd_aux_fine, - 0.25 / np.pi,
                                        self.nn_laplace)

        # 2.3 Construct refine chain
        self.refiner_i = []
        gd = self.gd_aux_fine
        N_c = gd.N_c.copy()
        for i in range(self.Ncoar):
            gd2 = gd.coarsen()
            self.refiner_i.append(Transformer(gd2, gd, self.nn_refine))
            N_c //= 2
            gd = gd2
        self.refiner_i = self.refiner_i[::-1]
        self.gd_aux_coar = gd

        if self.use_aux_grid:
            # 2.4 Construct large coarse grid from aux grid
            self.gd_large_coar_from_aux = ext_gd(self.gd_aux_coar,
                                                 N_c=self.gd_large_coar.N_c)
            # Check the consistency of the grids
            gd1 = self.gd_large_coar
            gd2 = self.gd_large_coar_from_aux
            assert np.all(gd1.N_c == gd2.N_c) and np.all(gd1.h_cv == gd2.h_cv)

        self._initialized = False

    def _init(self):
        if self._initialized:
            return
        # Allocate arrays
        self.phi_large_coar_g = self.gd_large_coar.zeros()
        self._initialized = True

        # Initialize poissonsolvers
        #self.ps_large_coar._init()
        #if not self.use_coarse:
        #    return
        #self.ps_small_fine._init()

    def solve(self, phi, rho, **kwargs):
        self._init()
        phi_small_fine_g = phi
        rho_small_fine_g = rho.copy()

        if self.use_coarse:
            # 1.1. Coarse rho on the small grid
            tmp_g = rho_small_fine_g
            for coarser in self.coarser_i:
                tmp_g = coarser.apply(tmp_g)
            rho_small_coar_g = tmp_g
        else:
            rho_small_coar_g = rho_small_fine_g

        # 1.2. Extend rho to the large grid
        rho_large_coar_g = self.gd_large_coar.zeros()
        extend_array(self.gd_small_coar, self.gd_large_coar,
                     rho_small_coar_g, rho_large_coar_g)

        # 1.3 Solve potential on the large coarse grid
        niter_large = self.ps_large_coar.solve(self.phi_large_coar_g,
                                               rho_large_coar_g, **kwargs)
        rho_large_coar_g = None

        if not self.use_coarse:
            deextend_array(self.gd_small_fine, self.gd_large_coar,
                           phi_small_fine_g, self.phi_large_coar_g)
            return niter_large

        if self.use_aux_grid:
            # 2.1 De-extend the potential to the auxiliary grid
            phi_aux_coar_g = self.gd_aux_coar.empty()
            deextend_array(self.gd_aux_coar, self.gd_large_coar_from_aux,
                           phi_aux_coar_g, self.phi_large_coar_g)
        else:
            phi_aux_coar_g = self.phi_large_coar_g

        # 3.1 Refine the potential
        tmp_g = phi_aux_coar_g
        for refiner in self.refiner_i:
            tmp_g = refiner.apply(tmp_g)
        phi_aux_coar_g = None
        phi_aux_fine_g = tmp_g

        # 3.2 Calculate the corresponding density with Laplace
        # (the refined coarse density would not accurately match with
        # the potential)
        rho_aux_fine_g = self.gd_aux_fine.empty()
        self.laplace_aux_fine.apply(phi_aux_fine_g, rho_aux_fine_g)

        # 3.3 De-extend the potential and density to the small grid
        cor_phi_small_fine_g = self.gd_small_fine.empty()
        deextend_array(self.gd_small_fine, self.gd_aux_fine,
                       cor_phi_small_fine_g, phi_aux_fine_g)
        phi_aux_fine_g = None
        cor_rho_small_fine_g = self.gd_small_fine.empty()
        deextend_array(self.gd_small_fine, self.gd_aux_fine,
                       cor_rho_small_fine_g, rho_aux_fine_g)
        rho_aux_fine_g = None

        # 3.4 Remove the correcting density and potential
        rho_small_fine_g -= cor_rho_small_fine_g
        phi_small_fine_g -= cor_phi_small_fine_g

        # 3.5 Solve potential on the small grid
        niter_small = self.ps_small_fine.solve(phi_small_fine_g,
                                               rho_small_fine_g, **kwargs)

        # 3.6 Correct potential and density
        phi_small_fine_g += cor_phi_small_fine_g
        # rho_small_fine_g += cor_rho_small_fine_g

        return (niter_large, niter_small)

    def estimate_memory(self, mem):
        self.ps_large_coar.estimate_memory(mem.subnode('Large grid Poisson'))
        if self.use_coarse:
            ps = self.ps_small_fine
            ps.estimate_memory(mem.subnode('Small grid Poisson'))
        mem.subnode('Large coarse phi', self.gd_large_coar.bytecount())
        tmp = max(self.gd_small_fine.bytecount(),
                  self.gd_large_coar.bytecount())
        if self.use_coarse:
            tmp = max(tmp,
                      self.gd_aux_coar.bytecount(),
                      self.gd_aux_fine.bytecount() * 2 +
                      self.gd_small_fine.bytecount(),
                      self.gd_aux_fine.bytecount() +
                      self.gd_small_fine.bytecount() * 2)
        mem.subnode('Temporary arrays', tmp)

    def get_description(self):
        line = '%s with ' % self.__class__.__name__
        if self.use_coarse:
            line += 'large and small grids'
        else:
            line += 'large grid'
        lines = [line]

        def add_line(line, pad=0):
            lines.extend(['%s%s' % (' ' * pad, line)])

        def get_cell(ps):
            descr = ps.get_description().replace('\n', '\n%s' % (' ' * 8))
            add_line('Poisson solver: %s' % descr, 8)
            if hasattr(ps, 'gd'):
                gd = ps.gd
                par = cell_to_cellpar(gd.cell_cv * Bohr)
                h_eff = gd.dv**(1.0 / 3.0) * Bohr
                l1 = '{:8d} x {:8d} x {:8d} points'.format(*gd.N_c)
                l2 = '{:8.2f} x {:8.2f} x {:8.2f} AA'.format(*par[:3])
                l3 = 'Effective grid spacing dv^(1/3) = {:.4f}'.format(h_eff)
                add_line('Grid: %s' % l1, 8)
                add_line('      %s' % l2, 8)
                add_line('      %s' % l3, 8)

        add_line('Large grid:', 4)
        get_cell(self.ps_large_coar)

        if self.use_coarse:
            add_line('Small grid:', 4)
            get_cell(self.ps_small_fine)

        return '\n'.join(lines)

    def todict(self):
        d = {'name': self.__class__.__name__}
        d['gpts'] = self.N_large_fine_c
        d['coarses'] = self.Ncoar
        d['nn_coarse'] = self.nn_coarse
        d['nn_refine'] = self.nn_refine
        d['nn_laplace'] = self.nn_laplace
        d['use_aux_grid'] = self.use_aux_grid
        d['poissonsolver_large'] = self.ps_large_coar.todict()
        if self.use_coarse:
            d['poissonsolver_small'] = self.ps_small_fine.todict()
        else:
            d['poissonsolver_small'] = None
        return d