def interpolate_pseudo_density(self, gridrefinement=2):
        
        gd = self.gd
        Fnt_wsg = self.Fnt_wsG.copy()
        
        # Find m for
        # gridrefinement = 2**m
        m1 = np.log(gridrefinement) / np.log(2.)
        m = int(np.round(m1))
        
        # Check if m is really integer
        if np.absolute(m - m1) < 1e-8:
            for i in range(m):
                gd2 = gd.refine()
                
                # Interpolate
                interpolator = Transformer(gd, gd2, self.stencil,
                                           dtype=self.dtype)
                Fnt2_wsg = gd2.empty((self.nw, self.nspins), dtype=self.dtype)
                for w in range(self.nw):
                    for s in range(self.nspins):
                        interpolator.apply(Fnt_wsg[w][s], Fnt2_wsg[w][s],
                                           np.ones((3, 2), dtype=complex))

                gd = gd2
                Fnt_wsg = Fnt2_wsg
        else:
            raise NotImplementedError
        
        return Fnt_wsg, gd
Example #2
0
    def interpolate_pseudo_density(self, gridrefinement=2):

        gd = self.gd
        Fnt_wsg = self.Fnt_wsG.copy()

        # Find m for
        # gridrefinement = 2**m
        m1 = np.log(gridrefinement) / np.log(2.)
        m = int(np.round(m1))

        # Check if m is really integer
        if np.absolute(m - m1) < 1e-8:
            for i in range(m):
                gd2 = gd.refine()

                # Interpolate
                interpolator = Transformer(gd,
                                           gd2,
                                           self.stencil,
                                           dtype=self.dtype)
                Fnt2_wsg = gd2.empty((self.nw, self.nspins), dtype=self.dtype)
                for w in range(self.nw):
                    for s in range(self.nspins):
                        interpolator.apply(Fnt_wsg[w][s], Fnt2_wsg[w][s],
                                           np.ones((3, 2), dtype=complex))

                gd = gd2
                Fnt_wsg = Fnt2_wsg
        else:
            raise NotImplementedError

        return Fnt_wsg, gd
Example #3
0
def interpolate_2d(mat):
    from gpaw.grid_descriptor import GridDescriptor
    from gpaw.transformers import Transformer
    nn = 10
    N_c = np.zeros([3], dtype=int)
    N_c[1:] = mat.shape[:2]
    N_c[0] = nn
    bmat = np.resize(mat, N_c)
    gd = GridDescriptor(N_c, N_c)
    finegd = GridDescriptor(N_c * 2, N_c)
    interpolator = Transformer(gd, finegd, 3)
    fine_bmat = finegd.zeros()
    interpolator.apply(bmat, fine_bmat)
    return fine_bmat[0]
Example #4
0
    def new_get_all_electron_density(self, atoms, gridrefinement=2):
        """Return real all-electron density array."""

        # Refinement of coarse grid, for representation of the AE-density
        if gridrefinement == 1:
            gd = self.gd
            n_sg = self.nt_sG.copy()
        elif gridrefinement == 2:
            gd = self.finegd
            if self.nt_sg is None:
                self.interpolate()
            n_sg = self.nt_sg.copy()
        elif gridrefinement == 4:
            # Extra fine grid
            gd = self.finegd.refine()
            
            # Interpolation function for the density:
            interpolator = Transformer(self.finegd, gd, 3)

            # Transfer the pseudo-density to the fine grid:
            n_sg = gd.empty(self.nspins)
            if self.nt_sg is None:
                self.interpolate()
            for s in range(self.nspins):
                interpolator.apply(self.nt_sg[s], n_sg[s])
        else:
            raise NotImplementedError

        # Add corrections to pseudo-density to get the AE-density
        splines = {}
        phi_aj = []
        phit_aj = []
        nc_a = []
        nct_a = []
        for a, id in enumerate(self.setups.id_a):
            if id in splines:
                phi_j, phit_j, nc, nct = splines[id]
            else:
                # Load splines:
                phi_j, phit_j, nc, nct = self.setups[a].get_partial_waves()[:4]
                splines[id] = (phi_j, phit_j, nc, nct)
            phi_aj.append(phi_j)
            phit_aj.append(phit_j)
            nc_a.append([nc])
            nct_a.append([nct])

        # Create localized functions from splines
        phi = BasisFunctions(gd, phi_aj)
        phit = BasisFunctions(gd, phit_aj)
        nc = LFC(gd, nc_a)
        nct = LFC(gd, nct_a)
        spos_ac = atoms.get_scaled_positions() % 1.0
        phi.set_positions(spos_ac)
        phit.set_positions(spos_ac)
        nc.set_positions(spos_ac)
        nct.set_positions(spos_ac)

        I_sa = np.zeros((self.nspins, len(atoms)))
        a_W =  np.empty(len(phi.M_W), np.int32)
        W = 0
        for a in phi.atom_indices:
            nw = len(phi.sphere_a[a].M_w)
            a_W[W:W + nw] = a
            W += nw
        rho_MM = np.zeros((phi.Mmax, phi.Mmax))
        for s, I_a in enumerate(I_sa):
            M1 = 0
            for a, setup in enumerate(self.setups):
                ni = setup.ni
                D_sp = self.D_asp.get(a)
                if D_sp is None:
                    D_sp = np.empty((self.nspins, ni * (ni + 1) // 2))
                else:
                    I_a[a] = ((setup.Nct - setup.Nc) / self.nspins -
                              sqrt(4 * pi) *
                              np.dot(D_sp[s], setup.Delta_pL[:, 0]))
                if gd.comm.size > 1:
                    gd.comm.broadcast(D_sp, self.rank_a[a])
                M2 = M1 + ni
                rho_MM[M1:M2, M1:M2] = unpack2(D_sp[s])
                M1 = M2

            phi.lfc.ae_valence_density_correction(rho_MM, n_sg[s], a_W, I_a)
            phit.lfc.ae_valence_density_correction(-rho_MM, n_sg[s], a_W, I_a)

        a_W =  np.empty(len(nc.M_W), np.int32)
        W = 0
        for a in nc.atom_indices:
            nw = len(nc.sphere_a[a].M_w)
            a_W[W:W + nw] = a
            W += nw
        scale = 1.0 / self.nspins
        for s, I_a in enumerate(I_sa):
            nc.lfc.ae_core_density_correction(scale, n_sg[s], a_W, I_a)
            nct.lfc.ae_core_density_correction(-scale, n_sg[s], a_W, I_a)
            gd.comm.sum(I_a)
            N_c = gd.N_c
            g_ac = np.around(N_c * spos_ac).astype(int) % N_c - gd.beg_c
            for I, g_c in zip(I_a, g_ac):
                if (g_c >= 0).all() and (g_c < gd.n_c).all():
                    n_sg[s][tuple(g_c)] -= I / gd.dv
        return n_sg, gd
Example #5
0
    def get_all_electron_density(self, atoms, gridrefinement=2):
        """Return real all-electron density array."""

        # Refinement of coarse grid, for representation of the AE-density
        if gridrefinement == 1:
            gd = self.gd
            n_sg = self.nt_sG.copy()
        elif gridrefinement == 2:
            gd = self.finegd
            if self.nt_sg is None:
                self.interpolate()
            n_sg = self.nt_sg.copy()
        elif gridrefinement == 4:
            # Extra fine grid
            gd = self.finegd.refine()
            
            # Interpolation function for the density:
            interpolator = Transformer(self.finegd, gd, 3)

            # Transfer the pseudo-density to the fine grid:
            n_sg = gd.empty(self.nspins)
            if self.nt_sg is None:
                self.interpolate()
            for s in range(self.nspins):
                interpolator.apply(self.nt_sg[s], n_sg[s])
        else:
            raise NotImplementedError

        # Add corrections to pseudo-density to get the AE-density
        splines = {}
        phi_aj = []
        phit_aj = []
        nc_a = []
        nct_a = []
        for a, id in enumerate(self.setups.id_a):
            if id in splines:
                phi_j, phit_j, nc, nct = splines[id]
            else:
                # Load splines:
                phi_j, phit_j, nc, nct = self.setups[a].get_partial_waves()[:4]
                splines[id] = (phi_j, phit_j, nc, nct)
            phi_aj.append(phi_j)
            phit_aj.append(phit_j)
            nc_a.append([nc])
            nct_a.append([nct])

        # Create localized functions from splines
        phi = LFC(gd, phi_aj)
        phit = LFC(gd, phit_aj)
        nc = LFC(gd, nc_a)
        nct = LFC(gd, nct_a)
        spos_ac = atoms.get_scaled_positions() % 1.0
        phi.set_positions(spos_ac)
        phit.set_positions(spos_ac)
        nc.set_positions(spos_ac)
        nct.set_positions(spos_ac)

        all_D_asp = []
        for a, setup in enumerate(self.setups):
            D_sp = self.D_asp.get(a)
            if D_sp is None:
                ni = setup.ni
                D_sp = np.empty((self.nspins, ni * (ni + 1) // 2))
            if gd.comm.size > 1:
                gd.comm.broadcast(D_sp, self.rank_a[a])
            all_D_asp.append(D_sp)

        for s in range(self.nspins):
            I_a = np.zeros(len(atoms))
            nc.add1(n_sg[s], 1.0 / self.nspins, I_a)
            nct.add1(n_sg[s], -1.0 / self.nspins, I_a)
            phi.add2(n_sg[s], all_D_asp, s, 1.0, I_a)
            phit.add2(n_sg[s], all_D_asp, s, -1.0, I_a)
            for a, D_sp in self.D_asp.items():
                setup = self.setups[a]
                I_a[a] -= ((setup.Nc - setup.Nct) / self.nspins +
                           sqrt(4 * pi) *
                           np.dot(D_sp[s], setup.Delta_pL[:, 0]))
            gd.comm.sum(I_a)
            N_c = gd.N_c
            g_ac = np.around(N_c * spos_ac).astype(int) % N_c - gd.beg_c
            for I, g_c in zip(I_a, g_ac):
                if (g_c >= 0).all() and (g_c < gd.n_c).all():
                    n_sg[s][tuple(g_c)] -= I / gd.dv

        return n_sg, gd
Example #6
0
class Density:
    """Density object.
    
    Attributes:
     =============== =====================================================
     ``gd``          Grid descriptor for coarse grids.
     ``finegd``      Grid descriptor for fine grids.
     ``interpolate`` Function for interpolating the electron density.
     ``mixer``       ``DensityMixer`` object.
     =============== =====================================================

    Soft and smooth pseudo functions on uniform 3D grids:
     ========== =========================================
     ``nt_sG``  Electron density on the coarse grid.
     ``nt_sg``  Electron density on the fine grid.
     ``nt_g``   Electron density on the fine grid.
     ``rhot_g`` Charge density on the fine grid.
     ``nct_G``  Core electron-density on the coarse grid.
     ========== =========================================
    """
    
    def __init__(self, gd, finegd, nspins, charge):
        """Create the Density object."""

        self.gd = gd
        self.finegd = finegd
        self.nspins = nspins
        self.charge = float(charge)

        self.charge_eps = 1e-7
        
        self.D_asp = None
        self.Q_aL = None

        self.nct_G = None
        self.nt_sG = None
        self.rhot_g = None
        self.nt_sg = None
        self.nt_g = None

        self.rank_a = None

        self.mixer = BaseMixer()
        self.timer = nulltimer
        self.allocated = False
        
    def initialize(self, setups, stencil, timer, magmom_a, hund):
        self.timer = timer
        self.setups = setups
        self.hund = hund
        self.magmom_a = magmom_a
        
        # Interpolation function for the density:
        self.interpolator = Transformer(self.gd, self.finegd, stencil,
                                        allocate=False)
        
        spline_aj = []
        for setup in setups:
            if setup.nct is None:
                spline_aj.append([])
            else:
                spline_aj.append([setup.nct])
        self.nct = LFC(self.gd, spline_aj,
                       integral=[setup.Nct for setup in setups],
                       forces=True, cut=True)
        self.ghat = LFC(self.finegd, [setup.ghat_l for setup in setups],
                        integral=sqrt(4 * pi), forces=True)
        if self.allocated:
            self.allocated = False
            self.allocate()

    def allocate(self):
        assert not self.allocated
        self.interpolator.allocate()
        self.allocated = True

    def reset(self):
        # TODO: reset other parameters?
        self.nt_sG = None

    def set_positions(self, spos_ac, rank_a=None):
        if not self.allocated:
            self.allocate()
        self.nct.set_positions(spos_ac)
        self.ghat.set_positions(spos_ac)
        self.mixer.reset()

        self.nct_G = self.gd.zeros()
        self.nct.add(self.nct_G, 1.0 / self.nspins)
        #self.nt_sG = None
        self.nt_sg = None
        self.nt_g = None
        self.rhot_g = None
        self.Q_aL = None

        # If both old and new atomic ranks are present, start a blank dict if
        # it previously didn't exist but it will needed for the new atoms.
        if (self.rank_a is not None and rank_a is not None and
            self.D_asp is None and (rank_a == self.gd.comm.rank).any()):
            self.D_asp = {}

        if self.rank_a is not None and self.D_asp is not None:
            self.timer.start('Redistribute')
            requests = []
            flags = (self.rank_a != rank_a)
            my_incoming_atom_indices = np.argwhere(np.bitwise_and(flags, \
                rank_a == self.gd.comm.rank)).ravel()
            my_outgoing_atom_indices = np.argwhere(np.bitwise_and(flags, \
                self.rank_a == self.gd.comm.rank)).ravel()

            for a in my_incoming_atom_indices:
                # Get matrix from old domain:
                ni = self.setups[a].ni
                D_sp = np.empty((self.nspins, ni * (ni + 1) // 2))
                requests.append(self.gd.comm.receive(D_sp, self.rank_a[a],
                                                     tag=a, block=False))
                assert a not in self.D_asp
                self.D_asp[a] = D_sp

            for a in my_outgoing_atom_indices:
                # Send matrix to new domain:
                D_sp = self.D_asp.pop(a)
                requests.append(self.gd.comm.send(D_sp, rank_a[a],
                                                  tag=a, block=False))
            self.gd.comm.waitall(requests)
            self.timer.stop('Redistribute')

        self.rank_a = rank_a

    def calculate_pseudo_density(self, wfs):
        """Calculate nt_sG from scratch.

        nt_sG will be equal to nct_G plus the contribution from
        wfs.add_to_density().
        """
        wfs.calculate_density_contribution(self.nt_sG)
        self.nt_sG += self.nct_G

    def update(self, wfs):
        self.timer.start('Density')
        self.timer.start('Pseudo density')
        self.calculate_pseudo_density(wfs)
        self.timer.stop('Pseudo density')
        self.timer.start('Atomic density matrices')
        wfs.calculate_atomic_density_matrices(self.D_asp)
        self.timer.stop('Atomic density matrices')
        self.timer.start('Multipole moments')
        comp_charge = self.calculate_multipole_moments()
        self.timer.stop('Multipole moments')
        
        if isinstance(wfs, LCAOWaveFunctions):
            self.timer.start('Normalize')
            self.normalize(comp_charge)
            self.timer.stop('Normalize')

        self.timer.start('Mix')
        self.mix(comp_charge)
        self.timer.stop('Mix')
        self.timer.stop('Density')

    def normalize(self, comp_charge=None):
        """Normalize pseudo density."""
        if comp_charge is None:
            comp_charge = self.calculate_multipole_moments()
        
        pseudo_charge = self.gd.integrate(self.nt_sG).sum()

        if pseudo_charge + self.charge + comp_charge != 0:
            if pseudo_charge != 0:
                x = -(self.charge + comp_charge) / pseudo_charge
                self.nt_sG *= x
            else:
                # Use homogeneous background:
                self.nt_sG[:] = (self.charge + comp_charge) * self.gd.dv

    def calculate_pseudo_charge(self, comp_charge):
        self.nt_g = self.nt_sg.sum(axis=0)
        self.rhot_g = self.nt_g.copy()
        self.ghat.add(self.rhot_g, self.Q_aL)

        if debug:
            charge = self.finegd.integrate(self.rhot_g) + self.charge
            if abs(charge) > self.charge_eps:
                raise RuntimeError('Charge not conserved: excess=%.9f' %
                                   charge)

    def mix(self, comp_charge):
        if not self.mixer.mix_rho:
            self.mixer.mix(self)
            comp_charge = None
          
        self.interpolate(comp_charge)
        self.calculate_pseudo_charge(comp_charge)

        if self.mixer.mix_rho:
            self.mixer.mix(self)

    def interpolate(self, comp_charge=None):
        """Interpolate pseudo density to fine grid."""
        if comp_charge is None:
            comp_charge = self.calculate_multipole_moments()

        if self.nt_sg is None:
            self.nt_sg = self.finegd.empty(self.nspins)

        for s in range(self.nspins):
            self.interpolator.apply(self.nt_sG[s], self.nt_sg[s])

        # With periodic boundary conditions, the interpolation will
        # conserve the number of electrons.
        if not self.gd.pbc_c.all():
            # With zero-boundary conditions in one or more directions,
            # this is not the case.
            pseudo_charge = -(self.charge + comp_charge)
            if abs(pseudo_charge) > 1.0e-14:
                x = pseudo_charge / self.finegd.integrate(self.nt_sg).sum()
                self.nt_sg *= x

    def calculate_multipole_moments(self):
        """Calculate multipole moments of compensation charges.

        Returns the total compensation charge in units of electron
        charge, so the number will be negative because of the
        dominating contribution from the nuclear charge."""

        comp_charge = 0.0
        self.Q_aL = {}
        for a, D_sp in self.D_asp.items():
            Q_L = self.Q_aL[a] = np.dot(D_sp.sum(0), self.setups[a].Delta_pL)
            Q_L[0] += self.setups[a].Delta0
            comp_charge += Q_L[0]
        return self.gd.comm.sum(comp_charge) * sqrt(4 * pi)

    def initialize_from_atomic_densities(self, basis_functions):
        """Initialize D_asp, nt_sG and Q_aL from atomic densities.

        nt_sG is initialized from atomic orbitals, and will
        be constructed with the specified magnetic moments and
        obeying Hund's rules if ``hund`` is true."""

        # XXX does this work with blacs?  What should be distributed?
        # Apparently this doesn't use blacs at all, so it's serial
        # with respect to the blacs distribution.  That means it works
        # but is not particularly efficient (not that this is a time
        # consuming step)

        f_sM = np.empty((self.nspins, basis_functions.Mmax))
        self.D_asp = {}
        f_asi = {}
        for a in basis_functions.atom_indices:
            c = self.charge / len(self.setups)  # distribute on all atoms
            f_si = self.setups[a].calculate_initial_occupation_numbers(
                    self.magmom_a[a], self.hund, charge=c, nspins=self.nspins)
            if a in basis_functions.my_atom_indices:
                self.D_asp[a] = self.setups[a].initialize_density_matrix(f_si)
            f_asi[a] = f_si

        self.nt_sG = self.gd.zeros(self.nspins)
        basis_functions.add_to_density(self.nt_sG, f_asi)
        self.nt_sG += self.nct_G
        self.calculate_normalized_charges_and_mix()

    def initialize_from_wavefunctions(self, wfs):
        """Initialize D_asp, nt_sG and Q_aL from wave functions."""
        self.nt_sG = self.gd.empty(self.nspins)
        self.calculate_pseudo_density(wfs)
        self.D_asp = {}
        my_atom_indices = np.argwhere(wfs.rank_a == self.gd.comm.rank).ravel()
        for a in my_atom_indices:
            ni = self.setups[a].ni
            self.D_asp[a] = np.empty((self.nspins, ni * (ni + 1) // 2))
        wfs.calculate_atomic_density_matrices(self.D_asp)
        self.calculate_normalized_charges_and_mix()

    def initialize_directly_from_arrays(self, nt_sG, D_asp):
        """Set D_asp and nt_sG directly."""
        self.nt_sG = nt_sG
        self.D_asp = D_asp
        #self.calculate_normalized_charges_and_mix()
        # No calculate multipole moments?  Tests will fail because of
        # improperly initialized mixer

    def calculate_normalized_charges_and_mix(self):
        comp_charge = self.calculate_multipole_moments()
        self.normalize(comp_charge)
        self.mix(comp_charge)

    def set_mixer(self, mixer):
        if mixer is not None:
            if self.nspins == 1 and isinstance(mixer, MixerSum):
                raise RuntimeError('Cannot use MixerSum with nspins==1')
            self.mixer = mixer
        else:
            if self.gd.pbc_c.any():
                beta = 0.1
                weight = 50.0
            else:
                beta = 0.25
                weight = 1.0
                
            if self.nspins == 2:
                self.mixer = MixerSum(beta=beta, weight=weight)
            else:
                self.mixer = Mixer(beta=beta, weight=weight)

        self.mixer.initialize(self)
        
    def estimate_magnetic_moments(self):
        magmom_a = np.zeros_like(self.magmom_a)
        if self.nspins == 2:
            for a, D_sp in self.D_asp.items():
                magmom_a[a] = np.dot(D_sp[0] - D_sp[1], self.setups[a].N0_p)
            self.gd.comm.sum(magmom_a)
        return magmom_a

    def get_correction(self, a, spin):
        """Integrated atomic density correction.

        Get the integrated correction to the pseuso density relative to
        the all-electron density.
        """
        setup = self.setups[a]
        return sqrt(4 * pi) * (
            np.dot(self.D_asp[a][spin], setup.Delta_pL[:, 0])
            + setup.Delta0 / self.nspins)

    def get_density_array(self):
        XXX
        # XXX why not replace with get_spin_density and get_total_density?
        """Return pseudo-density array."""
        if self.nspins == 2:
            return self.nt_sG
        else:
            return self.nt_sG[0]
    
    def get_all_electron_density(self, atoms, gridrefinement=2):
        """Return real all-electron density array."""

        # Refinement of coarse grid, for representation of the AE-density
        if gridrefinement == 1:
            gd = self.gd
            n_sg = self.nt_sG.copy()
        elif gridrefinement == 2:
            gd = self.finegd
            if self.nt_sg is None:
                self.interpolate()
            n_sg = self.nt_sg.copy()
        elif gridrefinement == 4:
            # Extra fine grid
            gd = self.finegd.refine()
            
            # Interpolation function for the density:
            interpolator = Transformer(self.finegd, gd, 3)

            # Transfer the pseudo-density to the fine grid:
            n_sg = gd.empty(self.nspins)
            if self.nt_sg is None:
                self.interpolate()
            for s in range(self.nspins):
                interpolator.apply(self.nt_sg[s], n_sg[s])
        else:
            raise NotImplementedError

        # Add corrections to pseudo-density to get the AE-density
        splines = {}
        phi_aj = []
        phit_aj = []
        nc_a = []
        nct_a = []
        for a, id in enumerate(self.setups.id_a):
            if id in splines:
                phi_j, phit_j, nc, nct = splines[id]
            else:
                # Load splines:
                phi_j, phit_j, nc, nct = self.setups[a].get_partial_waves()[:4]
                splines[id] = (phi_j, phit_j, nc, nct)
            phi_aj.append(phi_j)
            phit_aj.append(phit_j)
            nc_a.append([nc])
            nct_a.append([nct])

        # Create localized functions from splines
        phi = LFC(gd, phi_aj)
        phit = LFC(gd, phit_aj)
        nc = LFC(gd, nc_a)
        nct = LFC(gd, nct_a)
        spos_ac = atoms.get_scaled_positions() % 1.0
        phi.set_positions(spos_ac)
        phit.set_positions(spos_ac)
        nc.set_positions(spos_ac)
        nct.set_positions(spos_ac)

        all_D_asp = []
        for a, setup in enumerate(self.setups):
            D_sp = self.D_asp.get(a)
            if D_sp is None:
                ni = setup.ni
                D_sp = np.empty((self.nspins, ni * (ni + 1) // 2))
            if gd.comm.size > 1:
                gd.comm.broadcast(D_sp, self.rank_a[a])
            all_D_asp.append(D_sp)

        for s in range(self.nspins):
            I_a = np.zeros(len(atoms))
            nc.add1(n_sg[s], 1.0 / self.nspins, I_a)
            nct.add1(n_sg[s], -1.0 / self.nspins, I_a)
            phi.add2(n_sg[s], all_D_asp, s, 1.0, I_a)
            phit.add2(n_sg[s], all_D_asp, s, -1.0, I_a)
            for a, D_sp in self.D_asp.items():
                setup = self.setups[a]
                I_a[a] -= ((setup.Nc - setup.Nct) / self.nspins +
                           sqrt(4 * pi) *
                           np.dot(D_sp[s], setup.Delta_pL[:, 0]))
            gd.comm.sum(I_a)
            N_c = gd.N_c
            g_ac = np.around(N_c * spos_ac).astype(int) % N_c - gd.beg_c
            for I, g_c in zip(I_a, g_ac):
                if (g_c >= 0).all() and (g_c < gd.n_c).all():
                    n_sg[s][tuple(g_c)] -= I / gd.dv

        return n_sg, gd

    def new_get_all_electron_density(self, atoms, gridrefinement=2):
        """Return real all-electron density array."""

        # Refinement of coarse grid, for representation of the AE-density
        if gridrefinement == 1:
            gd = self.gd
            n_sg = self.nt_sG.copy()
        elif gridrefinement == 2:
            gd = self.finegd
            if self.nt_sg is None:
                self.interpolate()
            n_sg = self.nt_sg.copy()
        elif gridrefinement == 4:
            # Extra fine grid
            gd = self.finegd.refine()
            
            # Interpolation function for the density:
            interpolator = Transformer(self.finegd, gd, 3)

            # Transfer the pseudo-density to the fine grid:
            n_sg = gd.empty(self.nspins)
            if self.nt_sg is None:
                self.interpolate()
            for s in range(self.nspins):
                interpolator.apply(self.nt_sg[s], n_sg[s])
        else:
            raise NotImplementedError

        # Add corrections to pseudo-density to get the AE-density
        splines = {}
        phi_aj = []
        phit_aj = []
        nc_a = []
        nct_a = []
        for a, id in enumerate(self.setups.id_a):
            if id in splines:
                phi_j, phit_j, nc, nct = splines[id]
            else:
                # Load splines:
                phi_j, phit_j, nc, nct = self.setups[a].get_partial_waves()[:4]
                splines[id] = (phi_j, phit_j, nc, nct)
            phi_aj.append(phi_j)
            phit_aj.append(phit_j)
            nc_a.append([nc])
            nct_a.append([nct])

        # Create localized functions from splines
        phi = BasisFunctions(gd, phi_aj)
        phit = BasisFunctions(gd, phit_aj)
        nc = LFC(gd, nc_a)
        nct = LFC(gd, nct_a)
        spos_ac = atoms.get_scaled_positions() % 1.0
        phi.set_positions(spos_ac)
        phit.set_positions(spos_ac)
        nc.set_positions(spos_ac)
        nct.set_positions(spos_ac)

        I_sa = np.zeros((self.nspins, len(atoms)))
        a_W =  np.empty(len(phi.M_W), np.int32)
        W = 0
        for a in phi.atom_indices:
            nw = len(phi.sphere_a[a].M_w)
            a_W[W:W + nw] = a
            W += nw
        rho_MM = np.zeros((phi.Mmax, phi.Mmax))
        for s, I_a in enumerate(I_sa):
            M1 = 0
            for a, setup in enumerate(self.setups):
                ni = setup.ni
                D_sp = self.D_asp.get(a)
                if D_sp is None:
                    D_sp = np.empty((self.nspins, ni * (ni + 1) // 2))
                else:
                    I_a[a] = ((setup.Nct - setup.Nc) / self.nspins -
                              sqrt(4 * pi) *
                              np.dot(D_sp[s], setup.Delta_pL[:, 0]))
                if gd.comm.size > 1:
                    gd.comm.broadcast(D_sp, self.rank_a[a])
                M2 = M1 + ni
                rho_MM[M1:M2, M1:M2] = unpack2(D_sp[s])
                M1 = M2

            phi.lfc.ae_valence_density_correction(rho_MM, n_sg[s], a_W, I_a)
            phit.lfc.ae_valence_density_correction(-rho_MM, n_sg[s], a_W, I_a)

        a_W =  np.empty(len(nc.M_W), np.int32)
        W = 0
        for a in nc.atom_indices:
            nw = len(nc.sphere_a[a].M_w)
            a_W[W:W + nw] = a
            W += nw
        scale = 1.0 / self.nspins
        for s, I_a in enumerate(I_sa):
            nc.lfc.ae_core_density_correction(scale, n_sg[s], a_W, I_a)
            nct.lfc.ae_core_density_correction(-scale, n_sg[s], a_W, I_a)
            gd.comm.sum(I_a)
            N_c = gd.N_c
            g_ac = np.around(N_c * spos_ac).astype(int) % N_c - gd.beg_c
            for I, g_c in zip(I_a, g_ac):
                if (g_c >= 0).all() and (g_c < gd.n_c).all():
                    n_sg[s][tuple(g_c)] -= I / gd.dv
        return n_sg, gd

    if extra_parameters.get('usenewlfc', True):
        get_all_electron_density = new_get_all_electron_density
        
    def estimate_memory(self, mem):
        nspins = self.nspins
        nbytes = self.gd.bytecount()
        nfinebytes = self.finegd.bytecount()

        arrays = mem.subnode('Arrays')
        for name, size in [('nt_sG', nbytes * nspins),
                           ('nt_sg', nfinebytes * nspins),
                           ('nt_g', nfinebytes),
                           ('rhot_g', nfinebytes),
                           ('nct_G', nbytes)]:
            arrays.subnode(name, size)

        lfs = mem.subnode('Localized functions')
        for name, obj in [('nct', self.nct),
                          ('ghat', self.ghat)]:
            obj.estimate_memory(lfs.subnode(name))
        self.mixer.estimate_memory(mem.subnode('Mixer'), self.gd)

        # TODO
        # The implementation of interpolator memory use is not very
        # accurate; 20 MiB vs 13 MiB estimated in one example, probably
        # worse for parallel calculations.
        
        self.interpolator.estimate_memory(mem.subnode('Interpolator'))

    def get_spin_contamination(self, atoms, majority_spin=0):
        """Calculate the spin contamination.

        Spin contamination is defined as the integral over the
        spin density difference, where it is negative (i.e. the
        minority spin density is larger than the majority spin density.
        """

        if majority_spin == 0:
            smaj = 0
            smin = 1
        else:
            smaj = 1
            smin = 0
        nt_sg, gd = self.get_all_electron_density(atoms)
        dt_sg = nt_sg[smin] - nt_sg[smaj]
        dt_sg = np.where(dt_sg > 0, dt_sg, 0.0)
        return gd.integrate(dt_sg)
Example #7
0
class ResponseCalculator:
    """This class is a calculator for the sc density variation.

    From the given perturbation, the set of coupled equations for the
    first-order density response is solved self-consistently.

    Parameters
    ----------
    max_iter: int
        Maximum number of iterations in the self-consistent evaluation of
        the density variation.
    tolerance_sc: float
        Tolerance for the self-consistent loop measured in terms of
        integrated absolute change of the density derivative between two
        iterations.
    tolerance_sternheimer: float
        Tolerance for the solution of the Sternheimer equation -- passed to
        the ``LinearSolver``.
    beta: float (0 < beta < 1)
        Mixing coefficient.
    nmaxold: int
        Length of history for the mixer.
    weight: int
        Weight for the mixer metric (=1 -> no metric used).
        
    """

    parameters = {'verbose':               False,
                  'max_iter':              100,
                  'max_iter_krylov':       1000,
                  'krylov_solver':         'cg',
                  'tolerance_sc':          1.0e-5,
                  'tolerance_sternheimer': 1.0e-4,
                  'use_pc':                True,
                  'beta':                  0.1,
                  'nmaxold':               6,
                  'weight':                50
                  }
    
    def __init__(self, calc, wfs, poisson_solver=None, dtype=float, **kwargs):
        """Store calculator etc.

        Parameters
        ----------
        calc: Calculator
            Calculator instance containing a ground-state calculation
            (calc.set_positions must have been called before this point!).
        wfs: WaveFunctions
            Class taking care of wave-functions, projectors, k-point related
            quantities and symmetries.
        poisson_solver: PoissonSolver
            Multigrid or FFT poisson solver (not required if the
            ``Perturbation`` to be solved for has a ``solve_poisson`` member
            function). 
        dtype: ...
            dtype of the density response.
            
        """
        
        # Store ground-state quantities
        self.hamiltonian = calc.hamiltonian
        self.density = calc.density

        self.wfs = wfs

        # Get list of k-point containers
        self.kpt_u = wfs.kpt_u

        # Poisson solver
        if poisson_solver is None:
            # Solver must be provided by the perturbation
            self.poisson = None
            self.solve_poisson = None
        else:
            self.poisson = poisson_solver
            self.solve_poisson = self.poisson.solve_neutral
       
        # Store grid-descriptors
        self.gd = calc.density.gd
        self.finegd = calc.density.finegd

        # dtype for ground-state wave-functions
        self.gs_dtype = calc.wfs.dtype
        # dtype for the perturbing potential and density
        self.dtype = dtype
        
        # Grid transformer -- convert array from coarse to fine grid
        self.interpolator = Transformer(self.gd, self.finegd, nn=3,
                                        dtype=self.dtype, allocate=False)
        # Grid transformer -- convert array from fine to coarse grid
        self.restrictor = Transformer(self.finegd, self.gd, nn=3,
                                      dtype=self.dtype, allocate=False)

        # Sternheimer operator
        self.sternheimer_operator = None
        # Krylov solver
        self.linear_solver = None

        # Phases for transformer objects - since the perturbation determines
        # the form of the density response this is obtained from the
        # perturbation in the ``__call__`` member function below.
        self.phase_cd = None

        # Array attributes
        self.nt1_G = None
        self.vHXC1_G = None        
        self.nt1_g = None
        self.vH1_g = None

        # Perturbation
        self.perturbation = None
        
        # Number of occupied bands
        nvalence = calc.wfs.nvalence
        self.nbands = nvalence/2 + nvalence%2
        assert self.nbands <= calc.wfs.nbands
                                  
        self.initialized = False

        self.parameters = {}
        self.set(**kwargs)

    def clean(self):
        """Cleanup before call to ``__call__``."""

        self.perturbation = None
        self.solve_poisson = None

        self.nt1_G = None
        self.vHXC1_G = None        
        self.nt1_g = None
        self.vH1_g = None
        
    def __call__(self, perturbation):
        """Calculate density response (derivative) to perturbation.

        Parameters
        ----------
        perturbation: Perturbation
            Class implementing the perturbing potential. Must provide an
            ``apply`` member function implementing the multiplication of the
            perturbing potential to a (set of) state vector(s).
            
        """
        
        assert self.initialized, ("Linear response calculator "
                                  "not initizalized.")
        self.clean()
        
        if self.poisson is None:
            assert hasattr(perturbation, 'solve_poisson')
            self.solve_poisson = perturbation.solve_poisson

        # Store perturbation - used in other member functions
        self.perturbation = perturbation
        # Reset mixer
        self.mixer.reset()
        # Reset wave-functions
        self.wfs.reset()

        # Set phase attribute for Transformer objects
        self.phase_cd = self.perturbation.get_phase_cd()

        # Parameters for the SC-loop
        p = self.parameters
        max_iter = p['max_iter']
        tolerance = p['tolerance_sc']
        
        for iter in range(max_iter):

            if iter == 0:
                self.first_iteration()
            else:
                print "iter:%3i\t" % iter,
                norm = self.iteration()
                print "abs-norm: %6.3e\t" % norm,
                print ("integrated density response (abs): % 5.2e (%5.2e) "
                       % (self.gd.integrate(self.nt1_G.real),
                          self.gd.integrate(np.absolute(self.nt1_G))))
                       
                if norm < tolerance:
                    print ("self-consistent loop converged in %i iterations"
                           % iter)
                    break
                
            if iter == max_iter:
                raise RuntimeError, ("self-consistent loop did not converge "
                                     "in %i iterations" % iter)
   
    def set(self, **kwargs):
        """Set parameters for calculation."""

        # Check for legal input parameters
        for key, value in kwargs.items():
            if not key in ResponseCalculator.parameters:
                raise TypeError("Unknown keyword argument: '%s'" % key)

        # Insert default values if not given
        for key, value in ResponseCalculator.parameters.items():
            if key not in kwargs:
                kwargs[key] = value

        self.parameters.update(kwargs)
            
    def initialize(self, spos_ac):
        """Make the object ready for a calculation."""

        # Parameters
        p = self.parameters
        beta = p['beta']
        nmaxold = p['nmaxold']
        weight = p['weight']
        use_pc = p['use_pc']
        tolerance_sternheimer = p['tolerance_sternheimer']
        max_iter_krylov = p['max_iter_krylov']
        krylov_solver = p['krylov_solver']
                
        # Initialize WaveFunctions attribute
        self.wfs.initialize(spos_ac)
        
        # Initialize interpolator and restrictor
        self.interpolator.allocate()
        self.restrictor.allocate()
        
        # Initialize mixer
        # weight = 1 -> no metric is used
        self.mixer = BaseMixer(beta=beta, nmaxold=nmaxold,
                               weight=weight, dtype=self.dtype)
        self.mixer.initialize_metric(self.gd)
        
        # Linear operator in the Sternheimer equation
        self.sternheimer_operator = \
            SternheimerOperator(self.hamiltonian, self.wfs, self.gd,
                                dtype=self.gs_dtype)

        # Preconditioner for the Sternheimer equation
        if p['use_pc']:
            pc = ScipyPreconditioner(self.gd,
                                     self.sternheimer_operator.project,
                                     dtype=self.gs_dtype)
        else:
            pc = None

        #XXX K-point of the pc must be set in the k-point loop -> store a ref.
        self.pc = pc
        # Linear solver for the solution of Sternheimer equation            
        self.linear_solver = ScipyLinearSolver(method=krylov_solver,
                                               preconditioner=pc,
                                               tolerance=tolerance_sternheimer,
                                               max_iter=max_iter_krylov)

        self.initialized = True

    def first_iteration(self):
        """Perform first iteration of sc-loop."""

        self.wave_function_variations()
        self.density_response()
        self.mixer.mix(self.nt1_G, [], phase_cd=self.phase_cd)
        self.interpolate_density()
        
    def iteration(self):
        """Perform iteration."""

        # Update variation in the effective potential
        self.effective_potential_variation()
        # Update wave function variations
        self.wave_function_variations()
        # Update density
        self.density_response()
        # Mix - supply phase_cd here for metric inside the mixer
        self.mixer.mix(self.nt1_G, [], phase_cd=self.phase_cd)
        norm = self.mixer.get_charge_sloshing()

        self.interpolate_density()
       
        return norm

    def interpolate_density(self):
        """Interpolate density derivative onto the fine grid."""

        self.nt1_g = self.finegd.zeros(dtype=self.dtype)
        self.interpolator.apply(self.nt1_G, self.nt1_g, phases=self.phase_cd)
        
    def effective_potential_variation(self):
        """Calculate derivative of the effective potential (Hartree + XC)."""

        # Hartree part
        vHXC1_g = self.finegd.zeros(dtype=self.dtype)
        self.solve_poisson(vHXC1_g, self.nt1_g)
        # Store for evaluation of second order derivative
        self.vH1_g = vHXC1_g.copy()
        
        # XC part
        nt_sg = self.density.nt_sg
        fxct_sg = np.zeros_like(nt_sg)
        self.hamiltonian.xc.calculate_fxc(self.finegd, nt_sg, fxct_sg)
        vHXC1_g += fxct_sg[0] * self.nt1_g

        # Transfer to coarse grid
        self.vHXC1_G = self.gd.zeros(dtype=self.dtype)
        self.restrictor.apply(vHXC1_g, self.vHXC1_G, phases=self.phase_cd)
    
    def wave_function_variations(self):
        """Calculate variation in the wave-functions.

        Parameters
        ----------
        v1_G: ndarray
            Variation of the local effective potential (Hartree + XC).

        """

        verbose = self.parameters['verbose']

        if verbose:
            print "Calculating wave function variations"

        if self.perturbation.has_q():
            q_c = self.perturbation.get_q()
            kplusq_k = self.wfs.kd.find_k_plus_q(q_c)
        else:
            kplusq_k = None

        # Calculate wave-function variations for all k-points.
        for kpt in self.kpt_u:

            k = kpt.k

            if verbose:
                print "k-point %2.1i" % k
            
            # Index of k+q vector
            if kplusq_k is None:
                kplusq = k
                kplusqpt = kpt
            else:
                kplusq = kplusq_k[k]
                kplusqpt = self.kpt_u[kplusq]

            # Ground-state and first-order wave-functions
            psit_nG = kpt.psit_nG
            psit1_nG = kpt.psit1_nG
            # Update the SternheimerOperator
            self.sternheimer_operator.set_k(k)
            self.sternheimer_operator.set_kplusq(kplusq)
            # Update preconditioner
            if self.pc is not None:
                # k+q
                self.pc.set_kpt(kplusqpt)
                
            # Right-hand side of Sternheimer equations
            # k and k+q
            # XXX should only be done once for all k-points but maybe too cheap
            # to bother ??
            rhs_nG = self.gd.zeros(n=self.nbands, dtype=self.gs_dtype)            
            self.perturbation.apply(psit_nG, rhs_nG, self.wfs, k, kplusq)
            if self.vHXC1_G is not None:
                rhs_nG += self.vHXC1_G * psit_nG
            # Project out occupied subspace
            self.sternheimer_operator.project(rhs_nG)
              
            # Loop over occupied bands
            for n in range(self.nbands):

                # Update band index in SternheimerOperator
                self.sternheimer_operator.set_band(n)
                # Get view of the Bloch function derivative
                psit1_G = psit1_nG[n]
                # Rhs of Sternheimer equation                
                rhs_G = -1 * rhs_nG[n]
               
                # Solve Sternheimer equation
                iter, info = self.linear_solver.solve(self.sternheimer_operator,
                                                      psit1_G, rhs_G)
                
                if verbose:
                    print "\tBand %2.1i -" % n,
                    
                if info == 0:
                    if verbose:
                        print "linear solver converged in %i iterations" % iter
                elif info > 0:
                    assert False, ("linear solver did not converge in maximum "
                                   "number (=%i) of iterations for "
                                   "k-point number %d" % (iter, k))
                else:
                    assert False, ("linear solver failed to converge")

    def density_response(self):
        """Calculate density response from variation in the wave-functions."""

        # Density might be complex
        self.nt1_G = self.gd.zeros(dtype=self.dtype)

        for kpt in self.kpt_u:
            # The weight of the k-points includes spin-degeneracy
            w = kpt.weight
            # Wave functions
            psit_nG = kpt.psit_nG
            psit1_nG = kpt.psit1_nG

            for psit_G, psit1_G in zip(psit_nG, psit1_nG):
                # NOTICE: this relies on the automatic down-cast of the complex
                # array on the rhs to a real array when the lhs is real !!
                # Factor 2 for time-reversal symmetry
                self.nt1_G += 2 * w * psit_G.conj() * psit1_G
Example #8
0
class HybridXC(HybridXCBase):
    def __init__(self,
                 name,
                 hybrid=None,
                 xc=None,
                 finegrid=False,
                 unocc=False,
                 omega=None,
                 excitation=None,
                 excited=0,
                 stencil=2):
        """Mix standard functionals with exact exchange.

        finegrid: boolean
            Use fine grid for energy functional evaluations ?
        unocc: boolean
            Apply vxx also to unoccupied states ?
        omega: float
            RSF mixing parameter
        excitation: string:
            Apply operator for improved virtual orbitals
            to unocc states? Possible modes:
                singlet: excitations to singlets
                triplet: excitations to triplets
                average: average between singlets and tripletts
                see f.e. http://dx.doi.org/10.1021/acs.jctc.8b00238
        excited: number
            Band to excite from - counted from H**O downwards

        """
        self.finegrid = finegrid
        self.unocc = unocc
        self.excitation = excitation
        self.excited = excited
        HybridXCBase.__init__(self,
                              name,
                              hybrid=hybrid,
                              xc=xc,
                              omega=omega,
                              stencil=stencil)

    def calculate_paw_correction(self,
                                 setup,
                                 D_sp,
                                 dEdD_sp=None,
                                 addcoredensity=True,
                                 a=None):
        return self.xc.calculate_paw_correction(setup, D_sp, dEdD_sp,
                                                addcoredensity, a)

    def initialize(self, density, hamiltonian, wfs, occupations):
        assert wfs.kd.gamma
        self.xc.initialize(density, hamiltonian, wfs, occupations)
        self.kpt_comm = wfs.kd.comm
        self.nspins = wfs.nspins
        self.setups = wfs.setups
        self.density = density
        self.kpt_u = wfs.kpt_u
        self.exx_s = np.zeros(self.nspins)
        self.ekin_s = np.zeros(self.nspins)
        self.nocc_s = np.empty(self.nspins, int)

        self.gd = density.gd
        self.redistributor = density.redistributor

        use_charge_center = hamiltonian.poisson.use_charge_center
        # XXX How do we construct a copy of the Poisson solver of the
        # Hamiltonian?  We don't know what class it is, etc., but gd
        # may differ.
        # XXX One might consider using a charged centered compensation
        # charge for the PoissonSolver in the case of EXX as standard
        self.poissonsolver = PoissonSolver('fd',
                                           eps=1e-11,
                                           use_charge_center=use_charge_center)
        # self.poissonsolver = hamiltonian.poisson

        if self.finegrid:
            self.finegd = self.gd.refine()
            # XXX Taking restrictor from Hamiltonian will not work in PW mode,
            # will it?  I think this supports only real-space mode.
            # self.restrictor = hamiltonian.restrictor
            self.restrictor = Transformer(self.finegd, self.gd, 3)
            self.interpolator = Transformer(self.gd, self.finegd, 3)
        else:
            self.finegd = self.gd

        self.ghat = LFC(self.finegd,
                        [setup.ghat_l for setup in density.setups],
                        integral=np.sqrt(4 * np.pi),
                        forces=True)
        self.poissonsolver.set_grid_descriptor(self.finegd)
        if self.rsf == 'Yukawa':
            omega2 = self.omega**2
            self.screened_poissonsolver = HelmholtzSolver(
                k2=-omega2,
                eps=1e-11,
                nn=3,
                use_charge_center=use_charge_center)
            self.screened_poissonsolver.set_grid_descriptor(self.finegd)

    def set_positions(self, spos_ac):
        self.ghat.set_positions(spos_ac)

    def calculate(self, gd, n_sg, v_sg=None, e_g=None):
        # Normal XC contribution:
        exc = self.xc.calculate(gd, n_sg, v_sg, e_g)
        # Note that the quantities passed are on the
        # density/Hamiltonian grids!
        # They may be distributed differently from own quantities.
        self.ekin = self.kpt_comm.sum(self.ekin_s.sum())
        return exc + self.kpt_comm.sum(self.exx_s.sum())

    def calculate_exx(self):
        for kpt in self.kpt_u:
            self.apply_orbital_dependent_hamiltonian(kpt, kpt.psit_nG)

    def apply_orbital_dependent_hamiltonian(self,
                                            kpt,
                                            psit_nG,
                                            Htpsit_nG=None,
                                            dH_asp=None):
        if kpt.f_n is None:
            return

        deg = 2 // self.nspins  # Spin degeneracy
        hybrid = self.hybrid
        P_ani = kpt.P_ani
        setups = self.setups
        is_cam = self.is_cam

        vt_g = self.finegd.empty()
        if self.gd is not self.finegd:
            vt_G = self.gd.empty()
        if self.rsf == 'Yukawa':
            y_vt_g = self.finegd.empty()
            # if self.gd is not self.finegd:
            #     y_vt_G = self.gd.empty()

        nocc = int(ceil(kpt.f_n.sum())) // (3 - self.nspins)
        if self.excitation is not None:
            ex_band = nocc - self.excited - 1
            if self.excitation == 'singlet':
                ex_weight = -1
            elif self.excitation == 'triplet':
                ex_weight = +1
            else:
                ex_weight = 0

        if self.unocc or self.excitation is not None:
            nbands = len(kpt.f_n)
        else:
            nbands = nocc
        self.nocc_s[kpt.s] = nocc

        if Htpsit_nG is not None:
            kpt.vt_nG = self.gd.empty(nbands)
            kpt.vxx_ani = {}
            kpt.vxx_anii = {}
            for a, P_ni in P_ani.items():
                I = P_ni.shape[1]
                kpt.vxx_ani[a] = np.zeros((nbands, I))
                kpt.vxx_anii[a] = np.zeros((nbands, I, I))

        exx = 0.0
        ekin = 0.0

        # XXXX nbands can be different numbers on different cpus!
        # That means some will execute the loop and others not.
        # And deadlocks with augment-grids.

        # Determine pseudo-exchange
        for n1 in range(nbands):
            psit1_G = psit_nG[n1]
            f1 = kpt.f_n[n1] / deg
            for n2 in range(n1, nbands):
                psit2_G = psit_nG[n2]
                f2 = kpt.f_n[n2] / deg
                if n1 != n2 and f1 == 0 and f1 == f2:
                    continue  # Don't work on double unocc. bands
                # Double count factor:
                dc = (1 + (n1 != n2)) * deg
                nt_G, rhot_g = self.calculate_pair_density(
                    n1, n2, psit_nG, P_ani)
                vt_g[:] = 0.0
                # XXXXX This will go wrong because we are solving the
                # Poisson equation on the distribution of gd, not finegd
                # Or maybe it's fixed now

                self.poissonsolver.solve(vt_g,
                                         -rhot_g,
                                         charge=-float(n1 == n2),
                                         eps=1e-12,
                                         zero_initial_phi=True)
                vt_g *= hybrid
                if self.rsf == 'Yukawa':
                    y_vt_g[:] = 0.0
                    self.screened_poissonsolver.solve(y_vt_g,
                                                      -rhot_g,
                                                      charge=-float(n1 == n2),
                                                      eps=1e-12,
                                                      zero_initial_phi=True)
                    if is_cam:  # Cam like correction
                        y_vt_g *= self.cam_beta
                    else:
                        y_vt_g *= hybrid
                    vt_g -= y_vt_g
                if self.gd is self.finegd:
                    vt_G = vt_g
                else:
                    self.restrictor.apply(vt_g, vt_G)

                # Integrate the potential on fine and coarse grids
                int_fine = self.finegd.integrate(vt_g * rhot_g)
                int_coarse = self.gd.integrate(vt_G * nt_G)
                if self.gd.comm.rank == 0:  # only add to energy on master CPU
                    exx += 0.5 * dc * f1 * f2 * int_fine
                    ekin -= dc * f1 * f2 * int_coarse
                if Htpsit_nG is not None:
                    Htpsit_nG[n1] += f2 * vt_G * psit2_G
                    if n1 == n2:
                        kpt.vt_nG[n1] = f1 * vt_G
                        if self.excitation is not None and n1 == ex_band:
                            Htpsit_nG[nocc:] += f1 * vt_G * psit_nG[nocc:]
                    else:
                        if self.excitation is None or n1 != ex_band \
                                or n2 < nocc:
                            Htpsit_nG[n2] += f1 * vt_G * psit1_G
                        else:
                            Htpsit_nG[n2] += f1 * ex_weight * vt_G * psit1_G

                    # Update the vxx_uni and vxx_unii vectors of the nuclei,
                    # used to determine the atomic hamiltonian, and the
                    # residuals
                    v_aL = self.ghat.dict()
                    self.ghat.integrate(vt_g, v_aL)
                    for a, v_L in v_aL.items():
                        v_ii = unpack(np.dot(setups[a].Delta_pL, v_L))
                        v_ni = kpt.vxx_ani[a]
                        v_nii = kpt.vxx_anii[a]
                        P_ni = P_ani[a]
                        v_ni[n1] += f2 * np.dot(v_ii, P_ni[n2])
                        if n1 != n2:
                            if self.excitation is None or n1 != ex_band or \
                                    n2 < nocc:
                                v_ni[n2] += f1 * np.dot(v_ii, P_ni[n1])
                            else:
                                v_ni[n2] += f1 * ex_weight * \
                                    np.dot(v_ii, P_ni[n1])
                        else:
                            # XXX Check this:
                            v_nii[n1] = f1 * v_ii
                            if self.excitation is not None and n1 == ex_band:
                                for nuoc in range(nocc, nbands):
                                    v_ni[nuoc] += f1 * \
                                        np.dot(v_ii, P_ni[nuoc])

        def calculate_vv(ni, D_ii, M_pp, weight, addme=False):
            """Calculate the local corrections depending on Mpp."""
            dexx = 0
            dekin = 0
            if not addme:
                addsign = -2.0
            else:
                addsign = 2.0
            for i1 in range(ni):
                for i2 in range(ni):
                    A = 0.0
                    for i3 in range(ni):
                        p13 = packed_index(i1, i3, ni)
                        for i4 in range(ni):
                            p24 = packed_index(i2, i4, ni)
                            A += M_pp[p13, p24] * D_ii[i3, i4]
                    p12 = packed_index(i1, i2, ni)
                    if Htpsit_nG is not None:
                        dH_p[p12] += addsign * weight / \
                            deg * A / ((i1 != i2) + 1)
                    dekin += 2 * weight / deg * D_ii[i1, i2] * A
                    dexx -= weight / deg * D_ii[i1, i2] * A
            return (dexx, dekin)

        # Apply the atomic corrections to the energy and the Hamiltonian
        # matrix
        for a, P_ni in P_ani.items():
            setup = setups[a]

            if Htpsit_nG is not None:
                # Add non-trivial corrections the Hamiltonian matrix
                h_nn = symmetrize(
                    np.inner(P_ni[:nbands], kpt.vxx_ani[a][:nbands]))
                ekin -= np.dot(kpt.f_n[:nbands], h_nn.diagonal())

                dH_p = dH_asp[a][kpt.s]

            # Get atomic density and Hamiltonian matrices
            D_p = self.density.D_asp[a][kpt.s]
            D_ii = unpack2(D_p)
            ni = len(D_ii)

            # Add atomic corrections to the valence-valence exchange energy
            # --
            # >  D   C     D
            # --  ii  iiii  ii
            (dexx, dekin) = calculate_vv(ni, D_ii, setup.M_pp, hybrid)
            ekin += dekin
            exx += dexx
            if self.rsf is not None:
                Mg_pp = setup.calculate_yukawa_interaction(self.omega)
                if is_cam:
                    (dexx, dekin) = calculate_vv(ni,
                                                 D_ii,
                                                 Mg_pp,
                                                 self.cam_beta,
                                                 addme=True)
                else:
                    (dexx, dekin) = calculate_vv(ni,
                                                 D_ii,
                                                 Mg_pp,
                                                 hybrid,
                                                 addme=True)
                ekin -= dekin
                exx -= dexx
            # Add valence-core exchange energy
            # --
            # >  X   D
            # --  ii  ii
            if setup.X_p is not None:
                exx -= hybrid * np.dot(D_p, setup.X_p)
                if Htpsit_nG is not None:
                    dH_p -= hybrid * setup.X_p
                    ekin += hybrid * np.dot(D_p, setup.X_p)

                if self.rsf == 'Yukawa' and setup.X_pg is not None:
                    if is_cam:
                        thybrid = self.cam_beta  # 0th order
                    else:
                        thybrid = hybrid
                    exx += thybrid * np.dot(D_p, setup.X_pg)
                    if Htpsit_nG is not None:
                        dH_p += thybrid * setup.X_pg
                        ekin -= thybrid * np.dot(D_p, setup.X_pg)
                elif self.rsf == 'Yukawa' and setup.X_pg is None:
                    thybrid = exp(-3.62e-2 * self.omega)  # educated guess
                    if is_cam:
                        thybrid *= self.cam_beta
                    else:
                        thybrid *= hybrid
                    exx += thybrid * np.dot(D_p, setup.X_p)
                    if Htpsit_nG is not None:
                        dH_p += thybrid * setup.X_p
                        ekin -= thybrid * np.dot(D_p, setup.X_p)
                # Add core-core exchange energy
                if kpt.s == 0:
                    if self.rsf is None or is_cam:
                        if is_cam:
                            exx += self.cam_alpha * setup.ExxC
                        else:
                            exx += hybrid * setup.ExxC

        self.exx_s[kpt.s] = self.gd.comm.sum(exx)
        self.ekin_s[kpt.s] = self.gd.comm.sum(ekin)

    def correct_hamiltonian_matrix(self, kpt, H_nn):
        if not hasattr(kpt, 'vxx_ani'):
            return

        # if self.gd.comm.rank > 0:
        #    H_nn[:] = 0.0

        nocc = self.nocc_s[kpt.s]
        nbands = len(kpt.vt_nG)
        for a, P_ni in kpt.P_ani.items():
            H_nn[:nbands, :nbands] += symmetrize(
                np.inner(P_ni[:nbands], kpt.vxx_ani[a]))
        # self.gd.comm.sum(H_nn)

        if not self.unocc or self.excitation is not None:
            H_nn[:nocc, nocc:] = 0.0
            H_nn[nocc:, :nocc] = 0.0

    def calculate_pair_density(self, n1, n2, psit_nG, P_ani):
        Q_aL = {}
        for a, P_ni in P_ani.items():
            P1_i = P_ni[n1]
            P2_i = P_ni[n2]
            D_ii = np.outer(P1_i, P2_i.conj()).real
            D_p = pack(D_ii)
            Q_aL[a] = np.dot(D_p, self.setups[a].Delta_pL)

        nt_G = psit_nG[n1] * psit_nG[n2]

        if self.finegd is self.gd:
            nt_g = nt_G
        else:
            nt_g = self.finegd.empty()
            self.interpolator.apply(nt_G, nt_g)

        rhot_g = nt_g.copy()
        self.ghat.add(rhot_g, Q_aL)

        return nt_G, rhot_g

    def add_correction(self,
                       kpt,
                       psit_xG,
                       Htpsit_xG,
                       P_axi,
                       c_axi,
                       n_x,
                       calculate_change=False):
        if kpt.f_n is None:
            return

        if self.unocc or self.excitation is not None:
            nocc = len(kpt.vt_nG)
        else:
            nocc = self.nocc_s[kpt.s]

        if calculate_change:
            for x, n in enumerate(n_x):
                if n < nocc:
                    Htpsit_xG[x] += kpt.vt_nG[n] * psit_xG[x]
                    for a, P_xi in P_axi.items():
                        c_axi[a][x] += np.dot(kpt.vxx_anii[a][n], P_xi[x])
        else:
            for a, c_xi in c_axi.items():
                c_xi[:nocc] += kpt.vxx_ani[a][:nocc]

    def rotate(self, kpt, U_nn):
        if kpt.f_n is None:
            return

        U_nn = U_nn.T.copy()
        nocc = self.nocc_s[kpt.s]
        if len(kpt.vt_nG) == nocc:
            U_nn = U_nn[:nocc, :nocc]
        gemm(1.0, kpt.vt_nG.copy(), U_nn, 0.0, kpt.vt_nG)
        for v_ni in kpt.vxx_ani.values():
            gemm(1.0, v_ni.copy(), U_nn, 0.0, v_ni)
        for v_nii in kpt.vxx_anii.values():
            gemm(1.0, v_nii.copy(), U_nn, 0.0, v_nii)
class TimeDependentHamiltonian():
    def __init__(self, calc):

        #initialization
        self.calc = calc
        self.wfs = calc.wfs
        self.ham = calc.hamiltonian
        self.den = calc.density
        self.occ = calc.occupations

        #initialization plane and grid descriptors from GPAW calculation
        self.pd = calc.wfs.pd
        self.gd = calc.wfs.gd
        self.volume = np.abs(np.linalg.det(self.gd.cell_cv))

        #number of k-points
        self.nq = len(calc.wfs.kpt_u)
        #number of bands
        self.nbands = calc.get_number_of_bands()
        #number of electrons
        self.nelectrons = calc.get_number_of_electrons()

        #kinetic operator
        self.kinetic = np.zeros((self.nq, self.nbands, self.nbands),
                                dtype=complex)
        #overlap operator
        self.overlap = np.zeros((self.nq, self.nbands, self.nbands),
                                dtype=complex)
        #local momentum operator
        self.local_moment = np.zeros((3, self.nq, self.nbands, self.nbands),
                                     dtype=complex)
        #nonlocal momentum operator
        self.nonlocal_moment = np.zeros((3, self.nq, self.nbands, self.nbands),
                                        dtype=complex)
        #Fermi-Dirac occupation
        self.f_n = np.zeros((self.nq, self.nbands), dtype=float)

        #ground state Kohn-Sham orbitals wavefunctions
        psi_gs = []
        #ground state Kohn-Sham orbitals density
        den_gs = []

        for kpt in self.wfs.kpt_u:
            self.overlap[kpt.q] = np.eye(self.nbands)
            self.f_n[kpt.q] = kpt.f_n
            kinetic = 0.5 * self.pd.G2_qG[kpt.q]  # |G+q|^2/2
            gradient = self.pd.get_reciprocal_vectors(kpt.q)
            psi = []
            den = []
            for n in range(self.nbands):
                psi.append(self.pd.ifft(kpt.psit_nG[n], kpt.q))
                den.append(np.abs(psi[-1])**2)
                for m in range(self.nbands):
                    self.kinetic[kpt.q, n, m] = self.pd.integrate(
                        kpt.psit_nG[n], kinetic * kpt.psit_nG[m])
                    #calculation local momentum
                    #<psi_qn|\nabla|psi_qm>
                    for i in range(3):
                        self.local_moment[i, kpt.q, n, m] = self.pd.integrate(
                            kpt.psit_nG[n], gradient[:, i] * kpt.psit_nG[m])
            psi_gs.append(psi)
            den_gs.append(den)

        self.psi_gs = np.array(psi_gs)
        self.den_gs = np.array(den_gs, dtype=float)

        #real space grid points
        self.r = self.gd.get_grid_point_coordinates()

        #initialization local and nonlocal part of pseudopotential
        self.init_potential()

        self.proj = np.zeros((self.nq, self.nbands, self.norb), dtype=complex)
        self.proj_r = np.zeros((3, self.nq, self.nbands, self.norb),
                               dtype=complex)

        self.density = self.den.nt_sG.copy()

        #initialization charge density (ion+electrons) for Hartree potential
        self.ion_density = calc.hamiltonian.poisson.pd.ifft(
            calc.density.rhot_q) - calc.density.nt_sg[0]
        #plane wave descriptor for Hartree potential
        self.pd0 = self.ham.poisson.pd
        #reciprocal |G|^2 vectors for Hartree potential V(G)=4pi/|G|^2
        self.G2 = self.ham.poisson.G2_q
        self.G = self.pd0.get_reciprocal_vectors()

        #fine to coarse and coarse to fine grids transformers (for correct calculation local potential)
        self.fine_to_coarse = Transformer(calc.density.finegd, calc.density.gd,
                                          3)
        self.coarse_to_fine = Transformer(calc.density.gd, calc.density.finegd,
                                          3)

        #initialization local potenital from ground state density
        self.update_local_potential()
        self.update_gauge([0, 0, 0])
#----------------------------------------------------------------------------------------------------

    def init_potential(self):
        #initialization of nonlocal part of pseudopotential
        # V_NL=|chi_i> V_i <chi_i|
        spline_aj = []
        for setup in self.wfs.setups:
            spline_aj.append(setup.pt_j)
        self.lfc = PWLFC(spline_aj, self.pd)
        self.lfc.set_positions(self.calc.spos_ac)  #set position of atoms
        proj_G = []
        proj_r = []  #collect chi_i in real space using FFT

        for kpt in self.wfs.kpt_u:
            proj_G.append(self.lfc.expand(kpt.q))
            proj = []
            n_i = proj_G[-1].shape[1]
            for i in range(n_i):
                proj.append(self.pd.ifft(proj_G[-1][:, i].copy(), kpt.q))
            proj_r.append(proj)
        self.chi = np.array(proj_r) / self.gd.dv

        s = 0  # s=0 because we perform spin-paired calculation
        V = []  #collect V
        for a in range(len(self.wfs.setups)):
            dH_ii = unpack(self.ham.dH_asp[a][s])
            V.append(dH_ii.diagonal())

        V = np.array(V)
        self.V = V.ravel()

        self.norb = self.V.size  # number of orbitals in nonlocal potential

        #initialization of local part of pseudopotential
        V = self.ham.vbar.pd.zeros()
        self.ham.vbar.add(V)
        self.Vloc = self.ham.vbar.pd.ifft(V)

#----------------------------------------------------------------------------------------------------

    def update_density(self, wfn):
        self.density[0] = np.zeros_like(self.density[0])
        fast_density(self.density[0], wfn, self.den_gs)
        self.occupation = np.sum(np.abs(wfn)**2, axis=2)
        self.update_local_potential()

#----------------------------------------------------------------------------------------------------

    def update_gauge(self, A):
        #update projections p_qno=<chi_qo|psi_qn>
        phase = np.exp(1j * np.einsum('ixyz,i->xyz', self.r, A))
        fast_projections(self.proj, self.chi, phase, self.psi_gs, self.gd.dv)

        #update interaction hamiltonian
        #H_I=\nabla*A+0.5*A^2
        self.interaction = np.einsum(
            'i,iqnm->qnm', A,
            self.local_moment) + 0.5 * np.linalg.norm(A)**2 * self.overlap

        #calculation nonlocal momentum operator
        # I_NL=-i[r,V_NL]=-i sum_o V_o (r|chi_o><chi_o - |chi_o><chi_o|r)
        for i in range(3):
            fast_projections(self.proj_r[0], self.chi, self.r[i] * phase,
                             self.psi_gs, self.gd.dv)
        self.nonlocal_moment = -1j * (np.einsum(
            'o,iqno,qmo->iqnm', self.V, self.proj_r.conj(), self.proj) -
                                      np.einsum('o,qno,iqmo->iqnm', self.V,
                                                self.proj.conj(), self.proj_r))
        self.moment = self.local_moment + self.nonlocal_moment

#----------------------------------------------------------------------------------------------------

    def update_local_potential(self):

        #tranform density from coarse to fine grids
        density = self.coarse_to_fine.apply(self.density.copy())

        #calculate XC potential
        VXC = np.zeros_like(density)
        self.ham.xc.calculate(self.pd0.gd, density, VXC)

        # calculate Hartree potential
        charge_density = density + self.ion_density
        VH = 4 * np.pi * self.pd0.fft(charge_density) / self.G2
        VH = self.pd0.ifft(VH)

        #transform Hartree and XC potential from fine to coarse grids
        self.VXC = self.fine_to_coarse.apply(VXC[0])
        self.VH = self.fine_to_coarse.apply(VH)

#----------------------------------------------------------------------------------------------------

    def calculate_nonlocal(self):
        #calculation nonlocal part of Hamiltonian
        return np.einsum('o,qno,qmo->qnm', self.V, self.proj.conj(), self.proj)

#----------------------------------------------------------------------------------------------------

    def calculate_local(self):
        #calculation local part of Hamiltonian
        local = np.zeros((self.nq, self.nbands, self.nbands), dtype=complex)
        return fast_local(local, self.Vloc + self.VXC + self.VH, self.psi_gs,
                          self.gd.dv)

#----------------------------------------------------------------------------------------------------

    def calculate_kinetic(self):
        #calculation kinetic part of Hamiltonian
        return self.kinetic

#----------------------------------------------------------------------------------------------------

    def hamiltonian(self):
        #calculation hamiltonian
        return self.calculate_kinetic() + self.calculate_local(
        ) + self.calculate_nonlocal() + self.interaction

#----------------------------------------------------------------------------------------------------

    def calculate_current(self, wfn):
        return np.einsum('iqnm,qne,qme->i', self.moment, wfn.conj(),
                         wfn) / self.volume / self.nq
Example #10
0
class HybridXC(HybridXCBase):
    def __init__(self,
                 name,
                 hybrid=None,
                 xc=None,
                 finegrid=False,
                 unocc=False):
        """Mix standard functionals with exact exchange.

        finegrid: boolean
            Use fine grid for energy functional evaluations ?
        unocc: boolean
            Apply vxx also to unoccupied states ?
        """
        self.finegrid = finegrid
        self.unocc = unocc
        HybridXCBase.__init__(self, name, hybrid, xc)

    def calculate_paw_correction(self,
                                 setup,
                                 D_sp,
                                 dEdD_sp=None,
                                 addcoredensity=True,
                                 a=None):
        return self.xc.calculate_paw_correction(setup, D_sp, dEdD_sp,
                                                addcoredensity, a)

    def initialize(self, density, hamiltonian, wfs, occupations):
        assert wfs.kd.gamma
        self.xc.initialize(density, hamiltonian, wfs, occupations)
        self.kpt_comm = wfs.kd.comm
        self.nspins = wfs.nspins
        self.setups = wfs.setups
        self.density = density
        self.kpt_u = wfs.kpt_u
        self.exx_s = np.zeros(self.nspins)
        self.ekin_s = np.zeros(self.nspins)
        self.nocc_s = np.empty(self.nspins, int)

        self.gd = density.gd
        self.redistributor = density.redistributor

        # XXX How do we construct a copy of the Poisson solver of the
        # Hamiltonian?  We don't know what class it is, etc., but gd
        # may differ.
        self.poissonsolver = PoissonSolver(eps=1e-11)
        #self.poissonsolver = hamiltonian.poisson

        if self.finegrid:
            self.finegd = self.gd.refine()
            # XXX Taking restrictor from Hamiltonian will not work in PW mode,
            # will it?  I think this supports only real-space mode.
            #self.restrictor = hamiltonian.restrictor
            self.restrictor = Transformer(self.finegd, self.gd, 3)
            self.interpolator = Transformer(self.gd, self.finegd, 3)
        else:
            self.finegd = self.gd

        self.ghat = LFC(self.finegd,
                        [setup.ghat_l for setup in density.setups],
                        integral=np.sqrt(4 * np.pi),
                        forces=True)
        self.poissonsolver.set_grid_descriptor(self.finegd)
        self.poissonsolver.initialize()

    def set_positions(self, spos_ac):
        self.ghat.set_positions(spos_ac)

    def calculate(self, gd, n_sg, v_sg=None, e_g=None):
        # Normal XC contribution:
        exc = self.xc.calculate(gd, n_sg, v_sg, e_g)
        # Note that the quantities passed are on the density/Hamiltonian grids!
        # They may be distributed differently from own quantities.
        self.ekin = self.kpt_comm.sum(self.ekin_s.sum())
        return exc + self.kpt_comm.sum(self.exx_s.sum())

    def calculate_exx(self):
        for kpt in self.kpt_u:
            self.apply_orbital_dependent_hamiltonian(kpt, kpt.psit_nG)

    def apply_orbital_dependent_hamiltonian(self,
                                            kpt,
                                            psit_nG,
                                            Htpsit_nG=None,
                                            dH_asp=None):
        if kpt.f_n is None:
            return

        deg = 2 // self.nspins  # Spin degeneracy
        hybrid = self.hybrid
        P_ani = kpt.P_ani
        setups = self.setups

        vt_g = self.finegd.empty()
        if self.gd is not self.finegd:
            vt_G = self.gd.empty()

        nocc = int(kpt.f_n.sum()) // (3 - self.nspins)
        if self.unocc:
            nbands = len(kpt.f_n)
        else:
            nbands = nocc
        self.nocc_s[kpt.s] = nocc

        if Htpsit_nG is not None:
            kpt.vt_nG = self.gd.empty(nbands)
            kpt.vxx_ani = {}
            kpt.vxx_anii = {}
            for a, P_ni in P_ani.items():
                I = P_ni.shape[1]
                kpt.vxx_ani[a] = np.zeros((nbands, I))
                kpt.vxx_anii[a] = np.zeros((nbands, I, I))

        exx = 0.0
        ekin = 0.0

        # XXXX nbands can be different numbers on different cpus!
        # That means some will execute the loop and others not.
        # And deadlocks with augment-grids.

        # Determine pseudo-exchange
        for n1 in range(nbands):
            psit1_G = psit_nG[n1]
            f1 = kpt.f_n[n1] / deg
            for n2 in range(n1, nbands):
                psit2_G = psit_nG[n2]
                f2 = kpt.f_n[n2] / deg

                # Double count factor:
                dc = (1 + (n1 != n2)) * deg
                nt_G, rhot_g = self.calculate_pair_density(
                    n1, n2, psit_nG, P_ani)
                vt_g[:] = 0.0
                # XXXXX This will go wrong because we are solving the
                # Poisson equation on the distribution of gd, not finegd
                # Or maybe it's fixed now

                self.poissonsolver.solve(vt_g,
                                         -rhot_g,
                                         charge=-float(n1 == n2),
                                         eps=1e-12,
                                         zero_initial_phi=True)
                vt_g *= hybrid

                if self.gd is self.finegd:
                    vt_G = vt_g
                else:
                    self.restrictor.apply(vt_g, vt_G)

                # Integrate the potential on fine and coarse grids
                int_fine = self.finegd.integrate(vt_g * rhot_g)
                int_coarse = self.gd.integrate(vt_G * nt_G)
                if self.gd.comm.rank == 0:  # only add to energy on master CPU
                    exx += 0.5 * dc * f1 * f2 * int_fine
                    ekin -= dc * f1 * f2 * int_coarse
                if Htpsit_nG is not None:
                    Htpsit_nG[n1] += f2 * vt_G * psit2_G
                    if n1 == n2:
                        kpt.vt_nG[n1] = f1 * vt_G
                    else:
                        Htpsit_nG[n2] += f1 * vt_G * psit1_G

                    # Update the vxx_uni and vxx_unii vectors of the nuclei,
                    # used to determine the atomic hamiltonian, and the
                    # residuals
                    v_aL = self.ghat.dict()
                    self.ghat.integrate(vt_g, v_aL)
                    for a, v_L in v_aL.items():
                        v_ii = unpack(np.dot(setups[a].Delta_pL, v_L))
                        v_ni = kpt.vxx_ani[a]
                        v_nii = kpt.vxx_anii[a]
                        P_ni = P_ani[a]
                        v_ni[n1] += f2 * np.dot(v_ii, P_ni[n2])
                        if n1 != n2:
                            v_ni[n2] += f1 * np.dot(v_ii, P_ni[n1])
                        else:
                            # XXX Check this:
                            v_nii[n1] = f1 * v_ii

        # Apply the atomic corrections to the energy and the Hamiltonian matrix
        for a, P_ni in P_ani.items():
            setup = setups[a]

            if Htpsit_nG is not None:
                # Add non-trivial corrections the Hamiltonian matrix
                h_nn = symmetrize(
                    np.inner(P_ni[:nbands], kpt.vxx_ani[a][:nbands]))
                ekin -= np.dot(kpt.f_n[:nbands], h_nn.diagonal())

                dH_p = dH_asp[a][kpt.s]

            # Get atomic density and Hamiltonian matrices
            D_p = self.density.D_asp[a][kpt.s]
            D_ii = unpack2(D_p)
            ni = len(D_ii)

            # Add atomic corrections to the valence-valence exchange energy
            # --
            # >  D   C     D
            # --  ii  iiii  ii
            for i1 in range(ni):
                for i2 in range(ni):
                    A = 0.0
                    for i3 in range(ni):
                        p13 = packed_index(i1, i3, ni)
                        for i4 in range(ni):
                            p24 = packed_index(i2, i4, ni)
                            A += setup.M_pp[p13, p24] * D_ii[i3, i4]
                    p12 = packed_index(i1, i2, ni)
                    if Htpsit_nG is not None:
                        dH_p[p12] -= 2 * hybrid / deg * A / ((i1 != i2) + 1)
                    ekin += 2 * hybrid / deg * D_ii[i1, i2] * A
                    exx -= hybrid / deg * D_ii[i1, i2] * A

            # Add valence-core exchange energy
            # --
            # >  X   D
            # --  ii  ii
            if setup.X_p is not None:
                exx -= hybrid * np.dot(D_p, setup.X_p)
                if Htpsit_nG is not None:
                    dH_p -= hybrid * setup.X_p
                    ekin += hybrid * np.dot(D_p, setup.X_p)

                # Add core-core exchange energy
                if kpt.s == 0:
                    exx += hybrid * setup.ExxC

        self.exx_s[kpt.s] = self.gd.comm.sum(exx)
        self.ekin_s[kpt.s] = self.gd.comm.sum(ekin)

    def correct_hamiltonian_matrix(self, kpt, H_nn):
        if not hasattr(kpt, 'vxx_ani'):
            return

        if self.gd.comm.rank > 0:
            H_nn[:] = 0.0

        nocc = self.nocc_s[kpt.s]
        nbands = len(kpt.vt_nG)
        for a, P_ni in kpt.P_ani.items():
            H_nn[:nbands, :nbands] += symmetrize(
                np.inner(P_ni[:nbands], kpt.vxx_ani[a]))
        self.gd.comm.sum(H_nn)

        H_nn[:nocc, nocc:] = 0.0
        H_nn[nocc:, :nocc] = 0.0

    def calculate_pair_density(self, n1, n2, psit_nG, P_ani):
        Q_aL = {}
        for a, P_ni in P_ani.items():
            P1_i = P_ni[n1]
            P2_i = P_ni[n2]
            D_ii = np.outer(P1_i, P2_i.conj()).real
            D_p = pack(D_ii)
            Q_aL[a] = np.dot(D_p, self.setups[a].Delta_pL)

        nt_G = psit_nG[n1] * psit_nG[n2]

        if self.finegd is self.gd:
            nt_g = nt_G
        else:
            nt_g = self.finegd.empty()
            self.interpolator.apply(nt_G, nt_g)

        rhot_g = nt_g.copy()
        self.ghat.add(rhot_g, Q_aL)

        return nt_G, rhot_g

    def add_correction(self,
                       kpt,
                       psit_xG,
                       Htpsit_xG,
                       P_axi,
                       c_axi,
                       n_x,
                       calculate_change=False):
        if kpt.f_n is None:
            return

        nocc = self.nocc_s[kpt.s]

        if calculate_change:
            for x, n in enumerate(n_x):
                if n < nocc:
                    Htpsit_xG[x] += kpt.vt_nG[n] * psit_xG[x]
                    for a, P_xi in P_axi.items():
                        c_axi[a][x] += np.dot(kpt.vxx_anii[a][n], P_xi[x])
        else:
            for a, c_xi in c_axi.items():
                c_xi[:nocc] += kpt.vxx_ani[a][:nocc]

    def rotate(self, kpt, U_nn):
        if kpt.f_n is None:
            return

        nocc = self.nocc_s[kpt.s]
        if len(kpt.vt_nG) == nocc:
            U_nn = U_nn[:nocc, :nocc]
        gemm(1.0, kpt.vt_nG.copy(), U_nn, 0.0, kpt.vt_nG)
        for v_ni in kpt.vxx_ani.values():
            gemm(1.0, v_ni.copy(), U_nn, 0.0, v_ni)
        for v_nii in kpt.vxx_anii.values():
            gemm(1.0, v_nii.copy(), U_nn, 0.0, v_nii)
Example #11
0
class DensityFourierTransform(Observer):
    def __init__(self, timestep, frequencies, width=None, interval=1):
        """
        Parameters
        ----------
        timestep: float
            Time step in attoseconds (10^-18 s), e.g., 4.0 or 8.0
        frequencies: NumPy array or list of floats
            Frequencies in eV for Fourier transforms
        width: float or None
            Width of Gaussian envelope in eV, otherwise no envelope
        interval: int
            Number of timesteps between calls (used when attaching)
        """

        Observer.__init__(self, interval)
        self.timestep = interval * timestep * attosec_to_autime # autime
        self.omega_w = np.asarray(frequencies) * eV_to_aufrequency # autime^(-1)

        if width is None:
            self.sigma = None
        else:
            self.sigma = width * eV_to_aufrequency # autime^(-1)

        self.nw = len(self.omega_w)
        self.dtype = complex # np.complex128 really, but hey...
        self.Fnt_wsG = None
        self.Fnt_wsg = None

        self.Ant_sG = None
        self.Ant_sg = None

    def initialize(self, paw, allocate=True):
        self.allocated = False

        assert hasattr(paw, 'time') and hasattr(paw, 'niter'), 'Use TDDFT!'
        self.time = paw.time
        self.niter = paw.niter

        self.world = paw.wfs.world
        self.gd = paw.density.gd
        self.finegd = paw.density.finegd
        self.nspins = paw.density.nspins
        self.stencil = paw.input_parameters.stencils[1] # i.e. tar['InterpolationStencil']
        self.interpolator = paw.density.interpolator
        self.cinterpolator = Transformer(self.gd, self.finegd, self.stencil, \
                                        dtype=self.dtype, allocate=False)
        self.phase_cd = np.ones((3, 2), dtype=complex)

        self.Ant_sG = paw.density.nt_sG.copy() # TODO in allocate instead?

        # Attach to PAW-type object
        paw.attach(self, self.interval, density=paw.density)

        if allocate:
            self.allocate()

    def allocate(self):
        if not self.allocated:
            self.Fnt_wsG = self.gd.zeros((self.nw, self.nspins), \
                                        dtype=self.dtype)
            self.Fnt_wsg = None
            #self.Ant_sG = ...
            self.Ant_sg = None
            self.gamma_w = np.ones(self.nw, dtype=complex) * self.timestep
            self.cinterpolator.allocate()
            self.allocated = True

        if debug:
            assert is_contiguous(self.Fnt_wsG, self.dtype)

    def interpolate_fourier_transform(self):
        if self.Fnt_wsg is None:
            self.Fnt_wsg = self.finegd.empty((self.nw, self.nspins), \
                                            dtype=self.dtype)

        if self.dtype == float:
            intapply = self.interpolator.apply
        else:
            intapply = lambda Fnt_G, Fnt_g: self.cinterpolator.apply(Fnt_G, \
                Fnt_g, self.phase_cd)

        for w in range(self.nw):
            for s in range(self.nspins):
                intapply(self.Fnt_wsG[w,s], self.Fnt_wsg[w,s])

    def interpolate_average(self):
        if self.Ant_sg is None:
            self.Ant_sg = self.finegd.empty(self.nspins, dtype=float)

        for s in range(self.nspins):
            self.interpolator.apply(self.Ant_sG[s], self.Ant_sg[s])
            
    def update(self, density):

        # Update time
        # t[N] = t[N-1] + dt[N-1] #TODO better time-convention?
        self.time += self.timestep

        # Complex exponential with/without finite-width envelope
        f_w = np.exp(1.0j*self.omega_w*self.time)
        if self.sigma is not None:
            f_w *= np.exp(-self.time**2*self.sigma**2/2.0)

        # Update Fourier transformed density components
        # Fnt_wG[N] = Fnt_wG[N-1] + 1/sqrt(pi) * (nt_G[N]-avg_nt_G[N-1]) \
        #     * (f[N]*t[N] - gamma[N-1]) * dt[N]/(t[N]+dt[N])
        for w in range(self.nw):
            self.Fnt_wsG[w] += 1/np.pi**0.5 * (density.nt_sG - self.Ant_sG) \
                * (f_w[w]*self.time - self.gamma_w[w]) * self.timestep \
                / (self.time + self.timestep)

        # Update the cumulative phase factors
        # gamma[N] = gamma[N-1] + f[N]*dt[N]
        self.gamma_w += f_w * self.timestep

        # If dt[N] = dt for all N and sigma = 0, then this simplifies to:
        # gamma[N] = Sum_{n=0}^N exp(i*omega*n*dt) * dt
        # = (1 - exp(i*omega*(N+1)*dt)) / (1 - exp(i*omega*dt)) * dt

        # Update average density
        # Ant_G[N] = (t[N]*Ant_G[N-1] + nt_G[N]*dt[N])/(t[N]+dt[N])
        self.Ant_sG = (self.time*self.Ant_sG + density.nt_sG*self.timestep) \
            / (self.time + self.timestep)

    def get_fourier_transform(self, frequency=0, spin=0, gridrefinement=1):
        if gridrefinement == 1:
            return self.Fnt_wsG[frequency, spin]
        elif gridrefinement == 2:
            if self.Fnt_wsg is None:
                self.interpolate_fourier_transform()
            return self.Fnt_wsg[frequency, spin]
        else:
            raise NotImplementedError('Arbitrary refinement not implemented')

    def get_average(self, spin=0, gridrefinement=1):
        if gridrefinement == 1:
            return self.Ant_sG[spin]
        elif gridrefinement == 2:
            if self.Ant_sg is None:
                self.interpolate_average()
            return self.Ant_sg[spin]
        else:
            raise NotImplementedError('Arbitrary refinement not implemented')

    def read(self, filename, idiotproof=True):
        if idiotproof and not filename.endswith('.ftd'):
            raise IOError('Filename must end with `.ftd`.')

        tar = Reader(filename)

        # Test data type
        dtype = {'Float':float, 'Complex':complex}[tar['DataType']]
        if dtype != self.dtype:
            raise IOError('Data is an incompatible type.')

        # Test time
        time = tar['Time']
        if idiotproof and abs(time-self.time) >= 1e-9:
            raise IOError('Timestamp is incompatible with calculator.')

        # Test timestep (non-critical)
        timestep = tar['TimeStep']
        if abs(timestep - self.timestep) > 1e-12:
            print 'Warning: Time-step has been altered. (%lf -> %lf)' \
                % (self.timestep, timestep)
        self.timestep = timestep

        # Test dimensions
        nw = tar.dimension('nw')
        nspins = tar.dimension('nspins')
        ng = (tar.dimension('ngptsx'), tar.dimension('ngptsy'), \
              tar.dimension('ngptsz'),)

        if (nw != self.nw or nspins != self.nspins or
            (ng != self.gd.get_size_of_global_array()).any()):
            raise IOError('Data has incompatible shapes.')

        # Test width (non-critical)
        sigma = tar['Width']
        if ((sigma is None)!=(self.sigma is None) or # float <-> None
            (sigma is not None and self.sigma is not None and \
             abs(sigma - self.sigma) > 1e-12)): # float -> float
            print 'Warning: Width has been altered. (%s -> %s)' \
                % (self.sigma, sigma)
        self.sigma = sigma

        # Read frequencies
        self.omega_w[:] = tar.get('Frequency')

        # Read cumulative phase factors
        self.gamma_w[:] = tar.get('PhaseFactor')

        # Read average densities on master and distribute
        for s in range(self.nspins):
            all_Ant_G = tar.get('Average', s)
            self.gd.distribute(all_Ant_G, self.Ant_sG[s])

        # Read fourier transforms on master and distribute
        for w in range(self.nw):
            for s in range(self.nspins):
                all_Fnt_G = tar.get('FourierTransform', w, s)
                self.gd.distribute(all_Fnt_G, self.Fnt_wsG[w,s])

        # Close for good measure
        tar.close()

    def write(self, filename, idiotproof=True):
        if idiotproof and not filename.endswith('.ftd'):
            raise IOError('Filename must end with `.ftd`.')

        master = self.world.rank == 0

        # Open writer on master and set parameters/dimensions
        if master:
            tar = Writer(filename)
            tar['DataType'] = {float:'Float', complex:'Complex'}[self.dtype]
            tar['Time'] = self.time
            tar['TimeStep'] = self.timestep #non-essential
            tar['Width'] = self.sigma

            tar.dimension('nw', self.nw)
            tar.dimension('nspins', self.nspins)

            # Create dimensions for varioius netCDF variables:
            ng = self.gd.get_size_of_global_array()
            tar.dimension('ngptsx', ng[0])
            tar.dimension('ngptsy', ng[1])
            tar.dimension('ngptsz', ng[2])

            # Write frequencies
            tar.add('Frequency', ('nw',), self.omega_w, dtype=float)

            # Write cumulative phase factors
            tar.add('PhaseFactor', ('nw',), self.gamma_w, dtype=self.dtype)

        # Collect average densities on master and write
        if master:
            tar.add('Average', ('nspins', 'ngptsx', 'ngptsy', 
                'ngptsz', ), dtype=float)
        for s in range(self.nspins):
            big_Ant_G = self.gd.collect(self.Ant_sG[s])
            if master:
                tar.fill(big_Ant_G)

        # Collect fourier transforms on master and write
        if master:
            tar.add('FourierTransform', ('nw', 'nspins', 'ngptsx', 'ngptsy', \
                'ngptsz', ), dtype=self.dtype)
        for w in range(self.nw):
            for s in range(self.nspins):
                big_Fnt_G = self.gd.collect(self.Fnt_wsG[w,s])
                if master:
                    tar.fill(big_Fnt_G)

        # Close to flush changes
        if master:
            tar.close()

        # Make sure slaves don't return before master is done
        self.world.barrier()

    def dump(self, filename):
        if debug:
            assert is_contiguous(self.Fnt_wsG, self.dtype)
            assert is_contiguous(self.Ant_sG, float)

        all_Fnt_wsG = self.gd.collect(self.Fnt_wsG)
        all_Ant_sG = self.gd.collect(self.Ant_sG)

        if self.world.rank == 0:
            all_Fnt_wsG.dump(filename)
            all_Ant_sG.dump(filename+'_avg') # crude but easy
            self.omega_w.dump(filename+'_omega') # crude but easy
            self.gamma_w.dump(filename+'_gamma') # crude but easy

    def load(self, filename):
        if self.world.rank == 0:
            all_Fnt_wsG = np.load(filename)
            all_Ant_sG = np.load(filename+'_avg') # crude but easy
        else:
            all_Fnt_wsG = None
            all_Ant_sG = None

        if debug:
            assert all_Fnt_wsG is None or is_contiguous(all_Fnt_wsG, self.dtype)
            assert all_Ant_sG is None or is_contiguous(all_Ant_sG, float)

        if not self.allocated:
            self.allocate()

        self.gd.distribute(all_Fnt_wsG, self.Fnt_wsG)
        self.gd.distribute(all_Ant_sG, self.Ant_sG)

        self.omega_w = np.load(filename+'_omega') # crude but easy
        self.gamma_w = np.load(filename+'_gamma') # crude but easy
Example #12
0
class RealSpaceDensity(Density):
    def __init__(self, gd, finegd, nspins, charge, collinear=True,
                 stencil=3):
        Density.__init__(self, gd, finegd, nspins, charge, collinear)
        self.stencil = stencil

    def initialize(self, setups, timer, magmom_av, hund):
        Density.initialize(self, setups, timer, magmom_av, hund)

        # Interpolation function for the density:
        self.interpolator = Transformer(self.gd, self.finegd, self.stencil)
        
        spline_aj = []
        for setup in setups:
            if setup.nct is None:
                spline_aj.append([])
            else:
                spline_aj.append([setup.nct])
        self.nct = LFC(self.gd, spline_aj,
                       integral=[setup.Nct for setup in setups],
                       forces=True, cut=True)
        self.ghat = LFC(self.finegd, [setup.ghat_l for setup in setups],
                        integral=sqrt(4 * pi), forces=True)

    def set_positions(self, spos_ac, rank_a=None):
        Density.set_positions(self, spos_ac, rank_a)
        self.nct_G = self.gd.zeros()
        self.nct.add(self.nct_G, 1.0 / self.nspins)

    def interpolate_pseudo_density(self, comp_charge=None):
        """Interpolate pseudo density to fine grid."""
        if comp_charge is None:
            comp_charge = self.calculate_multipole_moments()

        self.nt_sg = self.interpolate(self.nt_sG, self.nt_sg)

        # With periodic boundary conditions, the interpolation will
        # conserve the number of electrons.
        if not self.gd.pbc_c.all():
            # With zero-boundary conditions in one or more directions,
            # this is not the case.
            pseudo_charge = -(self.charge + comp_charge)
            if abs(pseudo_charge) > 1.0e-14:
                x = (pseudo_charge /
                     self.finegd.integrate(self.nt_sg[:self.nspins]).sum())
                self.nt_sg *= x

    def interpolate(self, in_xR, out_xR=None):
        """Interpolate array(s)."""

        # ndim will be 3 in finite-difference mode and 1 when working
        # with the AtomPAW class (spherical atoms and 1d grids)
        ndim = self.gd.ndim

        if out_xR is None:
            out_xR = self.finegd.empty(in_xR.shape[:-ndim])

        a_xR = in_xR.reshape((-1,) + in_xR.shape[-ndim:])
        b_xR = out_xR.reshape((-1,) + out_xR.shape[-ndim:])
        
        for in_R, out_R in zip(a_xR, b_xR):
            self.interpolator.apply(in_R, out_R)

        return out_xR

    def calculate_pseudo_charge(self):
        self.nt_g = self.nt_sg[:self.nspins].sum(axis=0)
        self.rhot_g = self.nt_g.copy()
        self.ghat.add(self.rhot_g, self.Q_aL)

        if debug:
            charge = self.finegd.integrate(self.rhot_g) + self.charge
            if abs(charge) > self.charge_eps:
                raise RuntimeError('Charge not conserved: excess=%.9f' %
                                   charge)

    def get_pseudo_core_kinetic_energy_density_lfc(self):
        return LFC(self.gd,
                   [[setup.tauct] for setup in self.setups],
                   forces=True, cut=True)

    def calculate_dipole_moment(self):
        return self.finegd.calculate_dipole_moment(self.rhot_g)
Example #13
0
class SICSpin:
    def __init__(self,
                 kpt,
                 xc,
                 density,
                 hamiltonian,
                 wfs,
                 poissonsolver,
                 ghat,
                 finegd,
                 dtype=float,
                 coulomb_factor=0.5,
                 xc_factor=0.5,
                 uominres=1E-1,
                 uomaxres=1E-10,
                 uorelres=1E-4,
                 uonscres=1E-10,
                 rattle=-0.1,
                 stabpot=0.0,
                 maxuoiter=10,
                 logging=2):
        """Single spin SIC object.


        coulomb_factor:
            Scaling factor for Hartree-functional

        xc_factor:
            Scaling factor for xc-functional

        uominres:
            Minimum residual before unitary optimization starts

        uomaxres:
            Target accuracy for unitary optimization
            (absolute variance)

        uorelres:
            Target accuracy for unitary optimization
            (rel. to basis residual)

        maxuoiter:
            Maximum number of unitary optimization steps

        rattle:
            perturbation to the initial states
        """

        self.wfs = wfs
        self.kpt = kpt
        self.xc = xc
        self.poissonsolver = poissonsolver
        self.ghat = ghat
        self.pt = wfs.pt

        self.gd = wfs.gd
        self.finegd = finegd

        if self.finegd is self.gd:
            self.interpolator = None
            self.restrictor = None
        else:
            self.interpolator = Transformer(self.gd, self.finegd, 3)
            self.restrictor = Transformer(self.finegd, self.gd, 3)

        self.nspins = wfs.nspins
        self.spin = kpt.s
        self.timer = wfs.timer
        self.setups = wfs.setups

        self.dtype = dtype
        self.coulomb_factor = coulomb_factor
        self.xc_factor = xc_factor

        self.nocc = None  # number of occupied states
        self.W_mn = None  # unit. transf. to energy optimal states
        self.initial_W_mn = None  # initial unitary transformation
        self.vt_mG = None  # orbital dependent potentials
        self.exc_m = None  # PZ-SIC contributions (from E_xc)
        self.ecoulomb_m = None  # PZ-SIC contributions (from E_H)

        self.rattle = rattle  # perturb the initial unitary transformation
        self.stabpot = stabpot  # stabilizing constant potential to avoid
        # occupation of unoccupied states

        self.uominres = uominres
        self.uomaxres = uomaxres
        self.uorelres = uorelres
        self.uonscres = uonscres
        self.maxuoiter = maxuoiter
        self.maxlsiter = 1  # maximum number of line-search steps
        self.maxcgiter = 2  # maximum number of CG-iterations
        self.lsinterp = True  # interpolate for minimum during line search
        self.basiserror = 1E+20
        self.logging = logging

    def initialize_orbitals(self, rattle=-0.1, localize=True):
        if self.initial_W_mn is not None:
            self.nocc = self.initial_W_mn.shape[0]
        else:
            self.nocc, x = divmod(int(self.kpt.f_n.sum()), 3 - self.nspins)
            assert x == 0

        if self.nocc == 0:
            return

        Z_mmv = np.empty((self.nocc, self.nocc, 3), complex)
        for v in range(3):
            G_v = np.zeros(3)
            G_v[v] = 1
            Z_mmv[:, :,
                  v] = self.gd.wannier_matrix(self.kpt.psit_nG,
                                              self.kpt.psit_nG, G_v, self.nocc)
        self.finegd.comm.sum(Z_mmv)

        if self.initial_W_mn is not None:
            self.W_mn = self.initial_W_mn

        elif localize:
            W_nm = np.identity(self.nocc)
            localization = 0.0
            for iter in range(30):
                loc = _gpaw.localize(Z_mmv, W_nm)
                if loc - localization < 1e-6:
                    break
                localization = loc

            self.W_mn = W_nm.T.copy()
        else:
            self.W_mn = np.identity(self.nocc)

        if (rattle != 0.0 and self.W_mn is not None
                and self.initial_W_mn is None):
            U_mm = random_unitary_matrix(rattle, self.nocc)
            self.W_mn = np.dot(U_mm, self.W_mn)

        if self.W_mn is not None:
            self.finegd.comm.broadcast(self.W_mn, 0)

        spos_mc = -np.angle(Z_mmv.diagonal()).T / (2 * pi)
        self.initial_pos_mv = np.dot(spos_mc % 1.0, self.gd.cell_cv)

    def localize_orbitals(self):

        assert self.gd.orthogonal

        # calculate wannier matrixelements
        Z_mmv = np.empty((self.nocc, self.nocc, 3), complex)
        for v in range(3):
            G_v = np.zeros(3)
            G_v[v] = 1
            Z_mmv[:, :,
                  v] = self.gd.wannier_matrix(self.kpt.psit_nG,
                                              self.kpt.psit_nG, G_v, self.nocc)
        self.finegd.comm.sum(Z_mmv)

        # setup the initial configuration (identity)
        W_nm = np.identity(self.nocc)

        # localize the orbitals
        localization = 0.0
        for iter in range(30):
            loc = _gpaw.localize(Z_mmv, W_nm)
            if loc - localization < 1e-6:
                break
            localization = loc

        # apply localizing transformation
        self.W_mn = W_nm.T.copy()

    def rattle_orbitals(self, rattle=-0.1):

        # check for the trivial cases
        if rattle == 0.0:
            return

        if self.W_mn is None:
            return

        # setup a "random" unitary matrix
        nocc = self.W_mn.shape[0]
        U_mm = random_unitary_matrix(rattle, nocc)

        # apply unitary transformation
        self.W_mn = np.dot(U_mm, self.W_mn)

    def get_centers(self):
        assert self.gd.orthogonal

        # calculate energy optimal states (if necessary)
        if self.phit_mG is None:
            self.update_optimal_states()

        # calculate wannier matrixelements
        Z_mmv = np.empty((self.nocc, self.nocc, 3), complex)
        for v in range(3):
            G_v = np.zeros(3)
            G_v[v] = 1
            Z_mmv[:, :, v] = self.gd.wannier_matrix(self.phit_mG, self.phit_mG,
                                                    G_v, self.nocc)
        self.finegd.comm.sum(Z_mmv)

        # calculate positions of localized orbitals
        spos_mc = -np.angle(Z_mmv.diagonal()).T / (2 * pi)

        return np.dot(spos_mc % 1.0, self.gd.cell_cv) * Bohr

    def calculate_sic_matrixelements(self):
        # overlap of pseudo wavefunctions
        Htphit_mG = self.vt_mG * self.phit_mG
        V_mm = np.zeros((self.nocc, self.nocc), dtype=self.dtype)
        gemm(self.gd.dv, self.phit_mG, Htphit_mG, 0.0, V_mm, 't')

        # PAW
        for a, P_mi in self.P_ami.items():
            for m, dH_p in enumerate(self.dH_amp[a]):
                dH_ii = unpack(dH_p)
                V_mm[m, :] += np.dot(P_mi[m], np.dot(dH_ii, P_mi.T))

        # accumulate over grid-domains
        self.finegd.comm.sum(V_mm)
        self.V_mm = V_mm

        # Symmetrization of V and kappa-matrix:
        K_mm = 0.5 * (V_mm - V_mm.T.conj())
        V_mm = 0.5 * (V_mm + V_mm.T.conj())

        # evaluate the kinetic correction
        self.ekin = -np.trace(V_mm) * (3 - self.nspins)

        return V_mm, K_mm, np.vdot(K_mm, K_mm).real

    def update_optimal_states(self):
        # pseudo wavefunctions
        self.phit_mG = self.gd.zeros(self.nocc)
        if self.nocc > 0:
            gemm(1.0, self.kpt.psit_nG[:self.nocc], self.W_mn, 0.0,
                 self.phit_mG)

        # PAW
        self.P_ami = {}
        for a, P_ni in self.kpt.P_ani.items():
            self.P_ami[a] = np.dot(self.W_mn, P_ni[:self.nocc])

    def calculate_density(self, m):
        # pseudo density
        nt_G = self.phit_mG[m]**2

        if self.finegd is self.gd:
            nt_g = nt_G
        else:
            nt_g = self.finegd.empty()
            self.interpolator.apply(nt_G, nt_g)
            # normalize single-particle density
            nt_g *= self.gd.integrate(nt_G) / self.finegd.integrate(nt_g)

        # PAW corrections
        Q_aL = {}
        D_ap = {}
        for a, P_mi in self.P_ami.items():
            P_i = P_mi[m]
            D_ii = np.outer(P_i, P_i.conj()).real
            D_ap[a] = D_p = pack(D_ii)
            Q_aL[a] = np.dot(D_p, self.setups[a].Delta_pL)

        return nt_g, Q_aL, D_ap

    def update_potentials(self, save=False, restore=False):
        if restore:
            self.exc_m = self.exc_save_m
            self.ecoulomb_m = self.eco_save_m
            self.esic = self.esic_save
            self.vt_mG = self.vt_save_mG.copy()
            self.dH_amp = self.dH_save_amp.copy()
            return

        self.timer.start('ODD-potentials')

        nt_sg = self.finegd.empty(2)
        nt_sg[1] = 0.0
        vt_sg = self.finegd.empty(2)

        # PAW
        W_aL = self.ghat.dict()
        zero_initial_phi = False

        # initialize some bigger fields
        if self.vt_mG is None or self.nocc != self.phit_mG.shape[0]:
            self.vt_mG = self.gd.empty(self.nocc)
            self.exc_m = np.zeros(self.nocc)
            self.ecoulomb_m = np.zeros(self.nocc)
            self.vHt_mg = self.finegd.zeros(self.nocc)
            zero_initial_phi = True
        #
        # PAW
        self.dH_amp = {}
        for a, P_mi in self.P_ami.items():
            ni = P_mi.shape[1]
            self.dH_amp[a] = np.empty((self.nocc, ni * (ni + 1) // 2))
        #
        self.Q_maL = {}
        # loop all occupied orbitals
        for m, phit_G in enumerate(self.phit_mG):
            #
            # setup the single-particle density and PAW density-matrix
            nt_sg[0], Q_aL, D_ap = self.calculate_density(m)
            vt_sg[:] = 0.0
            #
            # xc-SIC
            self.timer.start('XC')
            if self.xc_factor != 0.0:
                exc = self.xc.calculate(self.finegd, nt_sg, vt_sg)
                exc /= self.finegd.comm.size
                vt_sg[0] *= -self.xc_factor
                #
                # PAW
                for a, D_p in D_ap.items():
                    setup = self.setups[a]
                    dH_p = self.dH_amp[a][m]
                    dH_sp = np.zeros((2, len(dH_p)))
                    #
                    D_sp = np.array([D_p, np.zeros_like(D_p)])
                    exc += self.xc.calculate_paw_correction(
                        setup, D_sp, dH_sp, addcoredensity=False)
                    dH_p[:] = -dH_sp[0] * self.xc_factor

                self.exc_m[m] = -self.xc_factor * exc
            self.timer.stop('XC')
            #
            # Hartree-SIC
            self.timer.start('Hartree')
            if self.coulomb_factor != 0.0:
                #
                # add compensation charges to pseudo density
                self.ghat.add(nt_sg[0], Q_aL)
                #
                # solve the coulomb problem
                self.poissonsolver.solve(self.vHt_mg[m],
                                         nt_sg[0],
                                         zero_initial_phi=zero_initial_phi)
                ecoulomb = 0.5 * self.finegd.integrate(
                    nt_sg[0] * self.vHt_mg[m])
                ecoulomb /= self.finegd.comm.size
                vt_sg[0] -= self.coulomb_factor * self.vHt_mg[m]
                #
                # PAW
                self.ghat.integrate(self.vHt_mg[m], W_aL)
                for a, D_p in D_ap.items():
                    setup = self.setups[a]
                    dH_p = self.dH_amp[a][m]
                    M_p = np.dot(setup.M_pp, D_p)
                    ecoulomb += np.dot(D_p, M_p)
                    #
                    dH_p -= self.coulomb_factor * (
                        2.0 * M_p + np.dot(setup.Delta_pL, W_aL[a]))
                #
                self.ecoulomb_m[m] = -self.coulomb_factor * ecoulomb
            self.timer.stop('Hartree')

            # restrict to course grid
            if self.finegd is self.gd:
                self.vt_mG[m] = vt_sg[0]
            else:
                self.timer.start('Restrictor')
                self.restrictor.apply(vt_sg[0], self.vt_mG[m])
                self.timer.stop('Restrictor')
            self.Q_maL[m] = Q_aL

        self.timer.stop('ODD-potentials')

        # accumulate total xc-SIC and coulomb-SIC
        self.finegd.comm.sum(self.exc_m)
        self.finegd.comm.sum(self.ecoulomb_m)

        # total sic (including spin-degeneracy)
        self.esic = self.exc_m.sum()
        self.esic += self.ecoulomb_m.sum()
        self.esic *= (3 - self.nspins)

        if save:
            self.exc_save_m = self.exc_m
            self.eco_save_m = self.ecoulomb_m
            self.esic_save = self.esic
            self.vt_save_mG = self.vt_mG.copy()
            self.dH_save_amp = self.dH_amp.copy()

    def apply_orbital_dependent_hamiltonian(self, psit_nG):
        """...

        Setup::

            |V phi_m> and <l|Vphi_m>,

        for occupied states m and unoccupied states l."""

        # nocc = self.nocc
        # nvirt = psit_nG.shape[0] - nocc

        self.Htphit_mG = self.vt_mG * self.phit_mG

    def correct_hamiltonian_matrix(self, H_nn):
        """ Add contributions of the non-local part of the
            interaction potential to the Hamiltonian matrix.

            on entry:
                          H_nn[n,m] = <n|H_{KS}|m>
            on exit:
                          H_nn[n,m] = <n|H_{KS} + V_{u}|m>

            where V_u is the unified Hamiltonian

                V_u = ...

        """

        nocc = self.nocc
        nvirt = H_nn.shape[0] - nocc

        W_mn = self.W_mn
        # R_mk = self.R_mk

        if self.gd.comm.rank == 0:
            V_mm = 0.5 * (self.V_mm + self.V_mm.T)
            H_nn[:nocc, :nocc] += np.dot(W_mn.T, np.dot(V_mm, W_mn))
            if self.stabpot != 0.0:
                H_nn[nocc:, nocc:] += np.eye(nvirt) * self.stabpot

        if nvirt != 0:
            H_nn[:nocc, nocc:] = 0.0  # R_nk
            H_nn[nocc:, :nocc] = 0.0  # R_nk.T
            # R_nk = np.dot(W_mn.T, R_mk) # CHECK THIS
            # H_nn[:nocc, nocc:] += R_nk
            # H_nn[nocc:, :nocc] += R_nk.T

    def calculate_residual(self, psit_nG, Htpsit_nG, P_ani, c_ani):
        """ Calculate the action of the unified Hamiltonian on an
            arbitrary state:

                H_u|Psi> =
        """

        nocc = self.nocc
        nvirt = psit_nG.shape[0] - nocc

        # constraint for unoccupied states
        R_mk = np.zeros((nocc, nvirt), dtype=self.dtype)
        if nvirt > 0:
            gemm(self.gd.dv, psit_nG[nocc:], self.Htphit_mG, 0.0, R_mk, 't')
            # PAW
            for a, P_mi in self.P_ami.items():
                P_ni = P_ani[a]

                for m, dH_p in enumerate(self.dH_amp[a]):
                    dH_ii = unpack(dH_p)
                    R_mk[m] += np.dot(P_mi[m], np.dot(dH_ii, P_ni[nocc:].T))

            self.finegd.comm.sum(R_mk)

        # self.R_mk = R_mk
        # R_mk  = self.R_mk
        W_mn = self.W_mn
        Htphit_mG = self.Htphit_mG
        phit_mG = self.phit_mG
        K_mm = 0.5 * (self.V_mm - self.V_mm.T)
        Q_mn = np.dot(K_mm, W_mn)

        # Action of unified Hamiltonian on occupied states:
        if nocc > 0:
            gemm(1.0, Htphit_mG, W_mn.T.copy(), 1.0, Htpsit_nG[:nocc])
            gemm(1.0, phit_mG, Q_mn.T.copy(), 1.0, Htpsit_nG[:nocc])
        if nvirt > 0:
            gemm(1.0, phit_mG, R_mk.T.copy(), 1.0, Htpsit_nG[nocc:])
            if self.stabpot != 0.0:
                Htpsit_nG[nocc:] += self.stabpot * psit_nG[nocc:]

        # PAW
        for a, P_mi in self.P_ami.items():
            #
            c_ni = c_ani[a]
            ct_mi = P_mi.copy()
            #
            dO_ii = self.setups[a].dO_ii
            c_ni[:nocc] += np.dot(Q_mn.T, np.dot(P_mi, dO_ii))
            c_ni[nocc:] += np.dot(R_mk.T, np.dot(P_mi, dO_ii))
            #
            for m, dH_p in enumerate(self.dH_amp[a]):
                dH_ii = unpack(dH_p)
                ct_mi[m] = np.dot(P_mi[m], dH_ii)
            c_ni[:nocc] += np.dot(W_mn.T, ct_mi)
            c_ni[nocc:] += self.stabpot * np.dot(P_ani[a][nocc:], dO_ii)

    def calculate_residual_change(self, psit_xG, Htpsit_xG, P_axi, c_axi, n_x):
        """

        """
        assert len(n_x) == 1

        Htphit_mG = self.Htphit_mG
        phit_mG = self.phit_mG

        w_mx = np.zeros((self.nocc, 1), dtype=self.dtype)
        v_mx = np.zeros((self.nocc, 1), dtype=self.dtype)

        gemm(self.gd.dv, psit_xG, phit_mG, 0.0, w_mx, 't')
        gemm(self.gd.dv, psit_xG, Htphit_mG, 0.0, v_mx, 't')

        # PAW
        for a, P_mi in self.P_ami.items():
            P_xi = P_axi[a]
            dO_ii = self.setups[a].dO_ii
            #
            w_mx += np.dot(P_mi, np.dot(dO_ii, P_xi.T))

            for m, dH_p in enumerate(self.dH_amp[a]):
                dH_ii = unpack(dH_p)
                v_mx[m] += np.dot(P_mi[m], np.dot(dH_ii, P_xi.T))

        # sum over grid-domains
        self.finegd.comm.sum(w_mx)
        self.finegd.comm.sum(v_mx)

        V_mm = 0.5 * (self.V_mm + self.V_mm.T)
        q_mx = v_mx - np.dot(V_mm, w_mx)

        if self.stabpot != 0.0:
            q_mx -= self.stabpot * w_mx

        gemm(1.0, Htphit_mG, w_mx.T.copy(), 1.0, Htpsit_xG)
        gemm(1.0, phit_mG, q_mx.T.copy(), 1.0, Htpsit_xG)

        # PAW
        for a, P_mi in self.P_ami.items():
            c_xi = c_axi[a]
            ct_mi = P_mi.copy()

            dO_ii = self.setups[a].dO_ii
            c_xi += np.dot(q_mx.T, np.dot(P_mi, dO_ii))

            for m, dH_p in enumerate(self.dH_amp[a]):
                dH_ii = unpack(dH_p)
                ct_mi[m] = np.dot(P_mi[m], dH_ii)
            c_xi += np.dot(w_mx.T, ct_mi)

        if self.stabpot != 0.0:
            Htphit_mG += self.stabpot * psit_xG
            for a, P_xi in P_axi.items():
                dO_ii = self.setups[a].dO_ii
                c_axi[a] += self.stabpot * np.dot(P_xi, dO_ii)

    def rotate(self, U_nn):
        """ Unitary transformations amongst the canonic states
            require to apply a counter-acting transformation to
            the energy optimal states. This subroutine takes
            care of it.

            Reorthogonalization is required whenever unoccupied
            states are mixed in.
        """
        # skip if no transformation to optimal states is set-up
        if self.W_mn is None:
            return

        # compensate the transformation amongst the occupied states
        self.W_mn = np.dot(self.W_mn, U_nn[:self.nocc, :self.nocc])

        # reorthogonalize if unoccupied states may have been mixed in
        if self.nocc != U_nn.shape[0]:
            self.W_mn = ortho(self.W_mn)
            # self.R_mk = np.dot(self.R_mk, U_nn[self.nocc:, self.nocc:].T)

    def add_forces(self, F_av):
        # Calculate changes in projections
        deg = 3 - self.nspins
        F_amiv = self.pt.dict(self.nocc, derivative=True)
        self.pt.derivative(self.phit_mG, F_amiv)
        for m in range(self.nocc):
            # Force from compensation charges:
            dF_aLv = self.ghat.dict(derivative=True)
            self.ghat.derivative(self.vHt_mg[m], dF_aLv)
            for a, dF_Lv in dF_aLv.items():
                F_av[a] -= deg * self.coulomb_factor * \
                    np.dot(self.Q_maL[m][a], dF_Lv)

            # Force from projectors
            for a, F_miv in F_amiv.items():
                F_vi = F_miv[m].T.conj()
                dH_ii = unpack(self.dH_amp[a][m])
                P_i = self.P_ami[a][m]
                F_v = np.dot(np.dot(F_vi, dH_ii), P_i)
                F_av[a] += deg * 2 * F_v.real

    def calculate(self):
        """ Evaluate the non-unitary invariant part of the
            energy functional and returns

            esic: float
                self-interaction energy

            ekin: float
                correction to the kinetic energy
        """
        # initialize transformation from canonic to energy
        # optimal states (if necessary)
        if self.W_mn is None:
            self.initialize_orbitals(rattle=self.rattle)

        # optimize the non-unitary invariant part of the
        # functional
        self.unitary_optimization()

        return self.esic, self.ekin

    def unitary_optimization(self):
        """ Optimization of the non-unitary invariant part of the
            energy functional.
        """

        optstep = 0.0
        Gold = 0.0
        cgiter = 0
        #
        epsstep = 0.005  # 0.005
        dltstep = 0.1  # 0.1
        prec = 1E-7

        # get the initial ODD potentials/energy/matrixelements
        self.update_optimal_states()
        self.update_potentials(save=True)
        ESI = self.esic
        V_mm, K_mm, norm = self.calculate_sic_matrixelements()

        if norm < self.uonscres and self.maxuoiter > 0:
            return

        if self.nocc <= 1:
            return

        # optimize the unitary transformation
        #
        # allocate arrays for the search direction,
        # i.e., the (conjugate) gradient
        D_old_mm = np.zeros_like(self.W_mn)

        for iter in range(abs(self.maxuoiter)):
            # copy the initial unitary transformation and orbital
            # dependent energies
            W_old_mn = self.W_mn.copy()

            # setup the steepest-descent/conjugate gradient
            # D_nn:  search direction
            # K_nn:  inverse gradient
            # G0  :  <K,D> (projected length of D along K)
            if Gold != 0.0:
                # conjugate gradient
                G0 = np.sum(K_mm * K_mm.conj()).real
                beta = G0 / Gold
                Gold = G0
                D_mm = K_mm + beta * D_old_mm

                G0 = np.sum(K_mm * D_mm.conj()).real
            else:
                # steepest-descent
                G0 = np.sum(K_mm * K_mm.conj()).real
                Gold = G0
                D_mm = K_mm

            updated = False
            minimum = False
            failed = True
            noise = False
            E0 = ESI

            # try to estimate optimal step-length from change in length
            # of the gradient (force-only)
            if (epsstep != 0.0 and norm > self.uomaxres):
                step = epsstep
                while (True):
                    U_mm = matrix_exponential(D_mm, step)
                    self.W_mn = np.dot(U_mm, W_old_mn)
                    self.update_optimal_states()
                    self.update_potentials()
                    E1 = self.esic
                    K0_mm = K_mm.copy()
                    V_mm, K_mm, norm = self.calculate_sic_matrixelements()

                    # projected length of the gradient at the new position
                    G1 = np.sum(((K_mm - K0_mm) / step) * D_mm.conj()).real
                    #
                    if (abs(E1 - E0) < prec and E1 >= E0):
                        #
                        eps_works = True
                        Eeps = E1
                        noise = True
                    elif (E1 < E0):
                        #
                        # trial step reduced energy
                        eps_works = True
                        Eeps = E1
                    else:
                        #
                        # epsilon step did not work
                        eps_works = False
                        optstep = 0.0
                        break

                    # compute the optimal step size
                    # optstep = step*G0/(G0-G1)
                    # print -G0/G1
                    optstep = -G0 / G1

                    if (eps_works):
                        break

                    # print eps_works, optstep, G0/((G0-G1)/step)

                    # decide on the method for stepping
                if (optstep > 0.0):

                    # convex region -> force only estimate for minimum
                    U_mm = matrix_exponential(D_mm, optstep)
                    self.W_mn = np.dot(U_mm, W_old_mn)
                    self.update_optimal_states()
                    self.update_potentials()
                    E1 = self.esic
                    if (abs(E1 - E0) < prec and E1 >= E0):
                        V_mm, K_mm, norm = self.calculate_sic_matrixelements()
                        ESI = E1
                        optstep = optstep
                        lsiter = 0
                        maxlsiter = -1
                        updated = True
                        minimum = True
                        failed = False
                        lsmethod = 'CV-N'
                        noise = True
                    elif (E1 < E0):
                        V_mm, K_mm, norm = self.calculate_sic_matrixelements()
                        ESI = E1
                        optstep = optstep
                        lsiter = 0
                        maxlsiter = -1
                        updated = True
                        minimum = True
                        failed = False
                        lsmethod = 'CV-S'
                    else:
                        # self.K_unn[q] = K_nn
                        ESI = E0
                        step = optstep
                        optstep = 0.0
                        lsiter = 0
                        maxlsiter = self.maxlsiter
                        updated = False
                        minimum = False
                        failed = True
                        lsmethod = 'CV-F-CC'
                else:
                    # self.K_unn[q] = K_nn
                    ESI = E0
                    step = optstep
                    optstep = 0.0
                    lsiter = 0
                    maxlsiter = self.maxlsiter
                    updated = False
                    minimum = False
                    failed = True
                    lsmethod = 'CC'
            else:
                maxlsiter = 0
                lsiter = -1
                optstep = epsstep
                updated = False
                minimum = True
                failed = False
                lsmethod = 'CC-EPS'

            if (optstep == 0.0):
                #
                # we are in the concave region or force-only estimate failed,
                # just follow the (conjugate) gradient
                step = dltstep * abs(step)
                U_mm = matrix_exponential(D_mm, step)
                self.W_mn = np.dot(U_mm, W_old_mn)
                self.update_optimal_states()
                self.update_potentials()
                E1 = self.esic
                #
                #
                if (abs(E1 - E0) < prec and E1 >= E0):
                    ESI = E1
                    optstep = 0.0
                    updated = False
                    minimum = True
                    failed = True
                    lsmethod = lsmethod + '-DLT-N'
                    noise = True
                    maxlsiter = -1
                elif (E1 < E0):
                    ESI = E1
                    optstep = step
                    updated = True
                    minimum = False
                    failed = False
                    lsmethod = lsmethod + '-DLT'
                    maxlsiter = self.maxlsiter
                elif (eps_works):
                    ESI = Eeps
                    E1 = Eeps
                    step = epsstep
                    updated = False
                    minimum = False
                    failed = False
                    lsmethod = lsmethod + '-EPS'
                    maxlsiter = self.maxlsiter
                else:
                    optstep = 0.0
                    updated = False
                    minimum = False
                    failed = True
                    lsmethod = lsmethod + '-EPS-F'
                    maxlsiter = -1
                #
                G = (E1 - E0) / step
                step0 = 0.0
                step1 = step
                step2 = 2 * step
                #
                for lsiter in range(maxlsiter):
                    #
                    # energy at the new position
                    U_mm = matrix_exponential(D_mm, step2)
                    self.W_mn = np.dot(U_mm, W_old_mn)
                    self.update_optimal_states()
                    self.update_potentials()
                    E2 = self.esic
                    G = (E2 - E1) / (step2 - step1)
                    #
                    #
                    if (G > 0.0):
                        if self.lsinterp:
                            a = (E0 / ((step2 - step0) * (step1 - step0)) +
                                 E2 / ((step2 - step1) * (step2 - step0)) -
                                 E1 / ((step2 - step1) * (step1 - step0)))
                            b = (E2 - E0) / (step2 - step0) - a * (step2 +
                                                                   step0)
                            if (a < step**2):
                                optstep = 0.5 * (step0 + step2)
                            else:
                                optstep = -0.5 * b / a
                            updated = False
                            minimum = True
                            break
                        else:
                            optstep = step1
                            updated = False
                            minimum = True
                            break

                    elif (G < 0.0):
                        optstep = step2
                        step0 = step1
                        step1 = step2
                        step2 = step2 + step
                        E0 = E1
                        E1 = E2
                        ESI = E2
                        updated = True
                        minimum = False

            if (cgiter != 0):
                lsmethod = lsmethod + ' CG'

            if (cgiter >= self.maxcgiter or not minimum):
                Gold = 0.0
                cgiter = 0
            else:
                cgiter = cgiter + 1
                D_old_mm[:, :] = D_mm[:, :]

            # update the energy and matrixelements of V and Kappa
            # and accumulate total residual of unitary optimization
            if (not updated):
                if optstep > 0.0:
                    U_mm = matrix_exponential(D_mm, optstep)
                    self.W_mn = np.dot(U_mm, W_old_mn)
                    self.update_optimal_states()
                    self.update_potentials()
                    ESI = self.esic
                    V_mm, K_mm, norm = self.calculate_sic_matrixelements()
                else:
                    self.W_mn = W_old_mn
                    self.update_optimal_states()
                    self.update_potentials(restore=True)
                    ESI = self.esic
                    V_mm, K_mm, norm = self.calculate_sic_matrixelements()

            if (lsiter == maxlsiter - 1):
                V_mm, K_mm, norm = self.calculate_sic_matrixelements()

            E0 = ESI

            # orthonormalize the energy optimal orbitals
            self.W_mn = ortho(self.W_mn)
            K = max(norm, 1.0e-16)

            if self.finegd.comm.rank == 0:
                if self.logging == 1:
                    print("           UO-%d: %2d %5.1f  %20.5f  " %
                          (self.spin, iter, np.log10(K), ESI * Hartree))
                elif self.logging == 2:
                    print("           UO-%d: %2d %5.1f  %20.5f  " %
                          (self.spin, iter, np.log10(K), ESI * Hartree) +
                          lsmethod)

            if ((K < self.uomaxres
                 or K < self.wfs.eigensolver.error * self.uorelres) or noise
                    or failed) and not self.maxuoiter < 0:
                break
Example #14
0
class DensityFourierTransform(Observer):
    def __init__(self, timestep, frequencies, width=None, interval=1):
        """
        Parameters

        timestep: float
            Time step in attoseconds (10^-18 s), e.g., 4.0 or 8.0
        frequencies: NumPy array or list of floats
            Frequencies in eV for Fourier transforms
        width: float or None
            Width of Gaussian envelope in eV, otherwise no envelope
        interval: int
            Number of timesteps between calls (used when attaching)
        """

        Observer.__init__(self, interval)
        self.timestep = interval * timestep * attosec_to_autime  # autime
        self.omega_w = np.asarray(
            frequencies) * eV_to_aufrequency  # autime^(-1)

        if width is None:
            self.sigma = None
        else:
            self.sigma = width * eV_to_aufrequency  # autime^(-1)

        self.nw = len(self.omega_w)
        self.dtype = complex  # np.complex128 really, but hey...
        self.Fnt_wsG = None
        self.Fnt_wsg = None

        self.Ant_sG = None
        self.Ant_sg = None

    def initialize(self, paw, allocate=True):
        self.allocated = False

        assert hasattr(paw, 'time') and hasattr(paw, 'niter'), 'Use TDDFT!'
        self.time = paw.time
        self.niter = paw.niter

        self.world = paw.wfs.world
        self.gd = paw.density.gd
        self.finegd = paw.density.finegd
        self.nspins = paw.density.nspins
        self.stencil = paw.input_parameters.stencils[
            1]  # i.e. tar['InterpolationStencil']
        self.interpolator = paw.density.interpolator
        self.cinterpolator = Transformer(self.gd, self.finegd, self.stencil, \
                                        dtype=self.dtype)
        self.phase_cd = np.ones((3, 2), dtype=complex)

        self.Ant_sG = paw.density.nt_sG.copy()  # TODO in allocate instead?

        # Attach to PAW-type object
        paw.attach(self, self.interval, density=paw.density)

        if allocate:
            self.allocate()

    def allocate(self):
        if not self.allocated:
            self.Fnt_wsG = self.gd.zeros((self.nw, self.nspins), \
                                        dtype=self.dtype)
            self.Fnt_wsg = None
            #self.Ant_sG = ...
            self.Ant_sg = None
            self.gamma_w = np.ones(self.nw, dtype=complex) * self.timestep
            self.allocated = True

        if debug:
            assert is_contiguous(self.Fnt_wsG, self.dtype)

    def interpolate_fourier_transform(self):
        if self.Fnt_wsg is None:
            self.Fnt_wsg = self.finegd.empty((self.nw, self.nspins), \
                                            dtype=self.dtype)

        if self.dtype == float:
            intapply = self.interpolator.apply
        else:
            intapply = lambda Fnt_G, Fnt_g: self.cinterpolator.apply(Fnt_G, \
                Fnt_g, self.phase_cd)

        for w in range(self.nw):
            for s in range(self.nspins):
                intapply(self.Fnt_wsG[w, s], self.Fnt_wsg[w, s])

    def interpolate_average(self):
        if self.Ant_sg is None:
            self.Ant_sg = self.finegd.empty(self.nspins, dtype=float)

        for s in range(self.nspins):
            self.interpolator.apply(self.Ant_sG[s], self.Ant_sg[s])

    def update(self, density):

        # Update time
        # t[N] = t[N-1] + dt[N-1] #TODO better time-convention?
        self.time += self.timestep

        # Complex exponential with/without finite-width envelope
        f_w = np.exp(1.0j * self.omega_w * self.time)
        if self.sigma is not None:
            f_w *= np.exp(-self.time**2 * self.sigma**2 / 2.0)

        # Update Fourier transformed density components
        # Fnt_wG[N] = Fnt_wG[N-1] + 1/sqrt(pi) * (nt_G[N]-avg_nt_G[N-1]) \
        #     * (f[N]*t[N] - gamma[N-1]) * dt[N]/(t[N]+dt[N])
        for w in range(self.nw):
            self.Fnt_wsG[w] += 1/np.pi**0.5 * (density.nt_sG - self.Ant_sG) \
                * (f_w[w]*self.time - self.gamma_w[w]) * self.timestep \
                / (self.time + self.timestep)

        # Update the cumulative phase factors
        # gamma[N] = gamma[N-1] + f[N]*dt[N]
        self.gamma_w += f_w * self.timestep

        # If dt[N] = dt for all N and sigma = 0, then this simplifies to:
        # gamma[N] = Sum_{n=0}^N exp(i*omega*n*dt) * dt
        # = (1 - exp(i*omega*(N+1)*dt)) / (1 - exp(i*omega*dt)) * dt

        # Update average density
        # Ant_G[N] = (t[N]*Ant_G[N-1] + nt_G[N]*dt[N])/(t[N]+dt[N])
        self.Ant_sG = (self.time*self.Ant_sG + density.nt_sG*self.timestep) \
            / (self.time + self.timestep)

    def get_fourier_transform(self, frequency=0, spin=0, gridrefinement=1):
        if gridrefinement == 1:
            return self.Fnt_wsG[frequency, spin]
        elif gridrefinement == 2:
            if self.Fnt_wsg is None:
                self.interpolate_fourier_transform()
            return self.Fnt_wsg[frequency, spin]
        else:
            raise NotImplementedError('Arbitrary refinement not implemented')

    def get_average(self, spin=0, gridrefinement=1):
        if gridrefinement == 1:
            return self.Ant_sG[spin]
        elif gridrefinement == 2:
            if self.Ant_sg is None:
                self.interpolate_average()
            return self.Ant_sg[spin]
        else:
            raise NotImplementedError('Arbitrary refinement not implemented')

    def read(self, filename, idiotproof=True):
        if idiotproof and not filename.endswith('.ftd'):
            raise IOError('Filename must end with `.ftd`.')

        tar = Reader(filename)

        # Test data type
        dtype = {'Float': float, 'Complex': complex}[tar['DataType']]
        if dtype != self.dtype:
            raise IOError('Data is an incompatible type.')

        # Test time
        time = tar['Time']
        if idiotproof and abs(time - self.time) >= 1e-9:
            raise IOError('Timestamp is incompatible with calculator.')

        # Test timestep (non-critical)
        timestep = tar['TimeStep']
        if abs(timestep - self.timestep) > 1e-12:
            print('Warning: Time-step has been altered. (%lf -> %lf)' \
                % (self.timestep, timestep))
        self.timestep = timestep

        # Test dimensions
        nw = tar.dimension('nw')
        nspins = tar.dimension('nspins')
        ng = (tar.dimension('ngptsx'), tar.dimension('ngptsy'), \
              tar.dimension('ngptsz'),)

        if (nw != self.nw or nspins != self.nspins
                or (ng != self.gd.get_size_of_global_array()).any()):
            raise IOError('Data has incompatible shapes.')

        # Test width (non-critical)
        sigma = tar['Width']
        if ((sigma is None)!=(self.sigma is None) or # float <-> None
            (sigma is not None and self.sigma is not None and \
             abs(sigma - self.sigma) > 1e-12)): # float -> float
            print('Warning: Width has been altered. (%s -> %s)' \
                % (self.sigma, sigma))
        self.sigma = sigma

        # Read frequencies
        self.omega_w[:] = tar.get('Frequency')

        # Read cumulative phase factors
        self.gamma_w[:] = tar.get('PhaseFactor')

        # Read average densities on master and distribute
        for s in range(self.nspins):
            all_Ant_G = tar.get('Average', s)
            self.gd.distribute(all_Ant_G, self.Ant_sG[s])

        # Read fourier transforms on master and distribute
        for w in range(self.nw):
            for s in range(self.nspins):
                all_Fnt_G = tar.get('FourierTransform', w, s)
                self.gd.distribute(all_Fnt_G, self.Fnt_wsG[w, s])

        # Close for good measure
        tar.close()

    def write(self, filename, idiotproof=True):
        if idiotproof and not filename.endswith('.ftd'):
            raise IOError('Filename must end with `.ftd`.')

        master = self.world.rank == 0

        # Open writer on master and set parameters/dimensions
        if master:
            tar = Writer(filename)
            tar['DataType'] = {float: 'Float', complex: 'Complex'}[self.dtype]
            tar['Time'] = self.time
            tar['TimeStep'] = self.timestep  #non-essential
            tar['Width'] = self.sigma

            tar.dimension('nw', self.nw)
            tar.dimension('nspins', self.nspins)

            # Create dimensions for varioius netCDF variables:
            ng = self.gd.get_size_of_global_array()
            tar.dimension('ngptsx', ng[0])
            tar.dimension('ngptsy', ng[1])
            tar.dimension('ngptsz', ng[2])

            # Write frequencies
            tar.add('Frequency', ('nw', ), self.omega_w, dtype=float)

            # Write cumulative phase factors
            tar.add('PhaseFactor', ('nw', ), self.gamma_w, dtype=self.dtype)

        # Collect average densities on master and write
        if master:
            tar.add('Average', (
                'nspins',
                'ngptsx',
                'ngptsy',
                'ngptsz',
            ),
                    dtype=float)
        for s in range(self.nspins):
            big_Ant_G = self.gd.collect(self.Ant_sG[s])
            if master:
                tar.fill(big_Ant_G)

        # Collect fourier transforms on master and write
        if master:
            tar.add('FourierTransform', ('nw', 'nspins', 'ngptsx', 'ngptsy', \
                'ngptsz', ), dtype=self.dtype)
        for w in range(self.nw):
            for s in range(self.nspins):
                big_Fnt_G = self.gd.collect(self.Fnt_wsG[w, s])
                if master:
                    tar.fill(big_Fnt_G)

        # Close to flush changes
        if master:
            tar.close()

        # Make sure slaves don't return before master is done
        self.world.barrier()

    def dump(self, filename):
        if debug:
            assert is_contiguous(self.Fnt_wsG, self.dtype)
            assert is_contiguous(self.Ant_sG, float)

        all_Fnt_wsG = self.gd.collect(self.Fnt_wsG)
        all_Ant_sG = self.gd.collect(self.Ant_sG)

        if self.world.rank == 0:
            all_Fnt_wsG.dump(filename)
            all_Ant_sG.dump(filename + '_avg')  # crude but easy
            self.omega_w.dump(filename + '_omega')  # crude but easy
            self.gamma_w.dump(filename + '_gamma')  # crude but easy

    def load(self, filename):
        if self.world.rank == 0:
            all_Fnt_wsG = np.load(filename)
            all_Ant_sG = np.load(filename + '_avg')  # crude but easy
        else:
            all_Fnt_wsG = None
            all_Ant_sG = None

        if debug:
            assert all_Fnt_wsG is None or is_contiguous(
                all_Fnt_wsG, self.dtype)
            assert all_Ant_sG is None or is_contiguous(all_Ant_sG, float)

        if not self.allocated:
            self.allocate()

        self.gd.distribute(all_Fnt_wsG, self.Fnt_wsG)
        self.gd.distribute(all_Ant_sG, self.Ant_sG)

        self.omega_w = np.load(filename + '_omega')  # crude but easy
        self.gamma_w = np.load(filename + '_gamma')  # crude but easy
Example #15
0
    def get_all_electron_density(self,
                                 atoms=None,
                                 gridrefinement=2,
                                 spos_ac=None,
                                 skip_core=False):
        """Return real all-electron density array.

           Usage: Either get_all_electron_density(atoms) or
                         get_all_electron_density(spos_ac=spos_ac)

           skip_core=True theoretically returns the
                          all-electron valence density (use with
                          care; will not in general integrate
                          to valence)
        """
        if spos_ac is None:
            spos_ac = atoms.get_scaled_positions() % 1.0

        # Refinement of coarse grid, for representation of the AE-density
        # XXXXXXXXXXXX think about distribution depending on gridrefinement!
        if gridrefinement == 1:
            gd = self.redistributor.aux_gd
            n_sg = self.nt_sG.copy()
            # This will get the density with the same distribution
            # as finegd:
            n_sg = self.redistributor.distribute(n_sg)
        elif gridrefinement == 2:
            gd = self.finegd
            if self.nt_sg is None:
                self.interpolate_pseudo_density()
            n_sg = self.nt_sg.copy()
        elif gridrefinement == 4:
            # Extra fine grid
            gd = self.finegd.refine()

            # Interpolation function for the density:
            interpolator = Transformer(self.finegd, gd, 3)  # XXX grids!

            # Transfer the pseudo-density to the fine grid:
            n_sg = gd.empty(self.nspins)
            if self.nt_sg is None:
                self.interpolate_pseudo_density()
            for s in range(self.nspins):
                interpolator.apply(self.nt_sg[s], n_sg[s])
        else:
            raise NotImplementedError

        # Add corrections to pseudo-density to get the AE-density
        splines = {}
        phi_aj = []
        phit_aj = []
        nc_a = []
        nct_a = []
        for a, id in enumerate(self.setups.id_a):
            if id in splines:
                phi_j, phit_j, nc, nct = splines[id]
            else:
                # Load splines:
                phi_j, phit_j, nc, nct = self.setups[a].get_partial_waves()[:4]
                splines[id] = (phi_j, phit_j, nc, nct)
            phi_aj.append(phi_j)
            phit_aj.append(phit_j)
            nc_a.append([nc])
            nct_a.append([nct])

        # Create localized functions from splines
        phi = BasisFunctions(gd, phi_aj)
        phit = BasisFunctions(gd, phit_aj)
        nc = LFC(gd, nc_a)
        nct = LFC(gd, nct_a)
        phi.set_positions(spos_ac)
        phit.set_positions(spos_ac)
        nc.set_positions(spos_ac)
        nct.set_positions(spos_ac)

        I_sa = np.zeros((self.nspins, len(spos_ac)))
        a_W = np.empty(len(phi.M_W), np.intc)
        W = 0
        for a in phi.atom_indices:
            nw = len(phi.sphere_a[a].M_w)
            a_W[W:W + nw] = a
            W += nw

        x_W = phi.create_displacement_arrays()[0]
        D_asp = self.D_asp  # XXX really?

        rho_MM = np.zeros((phi.Mmax, phi.Mmax))
        for s, I_a in enumerate(I_sa):
            M1 = 0
            for a, setup in enumerate(self.setups):
                ni = setup.ni
                D_sp = D_asp.get(a)
                if D_sp is None:
                    D_sp = np.empty((self.nspins, ni * (ni + 1) // 2))
                else:
                    I_a[a] = (
                        (setup.Nct) / self.nspins -
                        sqrt(4 * pi) * np.dot(D_sp[s], setup.Delta_pL[:, 0]))

                    if not skip_core:
                        I_a[a] -= setup.Nc / self.nspins

                if gd.comm.size > 1:
                    gd.comm.broadcast(D_sp, D_asp.partition.rank_a[a])
                M2 = M1 + ni
                rho_MM[M1:M2, M1:M2] = unpack2(D_sp[s])
                M1 = M2

            assert np.all(n_sg[s].shape == phi.gd.n_c)
            phi.lfc.ae_valence_density_correction(rho_MM, n_sg[s], a_W, I_a,
                                                  x_W)
            phit.lfc.ae_valence_density_correction(-rho_MM, n_sg[s], a_W, I_a,
                                                   x_W)

        a_W = np.empty(len(nc.M_W), np.intc)
        W = 0
        for a in nc.atom_indices:
            nw = len(nc.sphere_a[a].M_w)
            a_W[W:W + nw] = a
            W += nw
        scale = 1.0 / self.nspins

        for s, I_a in enumerate(I_sa):

            if not skip_core:
                nc.lfc.ae_core_density_correction(scale, n_sg[s], a_W, I_a)

            nct.lfc.ae_core_density_correction(-scale, n_sg[s], a_W, I_a)
            gd.comm.sum(I_a)
            N_c = gd.N_c
            g_ac = np.around(N_c * spos_ac).astype(int) % N_c - gd.beg_c

            if not skip_core:

                for I, g_c in zip(I_a, g_ac):
                    if (g_c >= 0).all() and (g_c < gd.n_c).all():
                        n_sg[s][tuple(g_c)] -= I / gd.dv

        return n_sg, gd
Example #16
0
class PhononPerturbation(Perturbation):
    """Implementation of a phonon perturbation.

    This class implements the change in the effective potential due to a
    displacement of an atom ``a`` in direction ``v`` with wave-vector ``q``.
    The action of the perturbing potential on a state vector is implemented in
    the ``apply`` member function.
    
    """
    
    def __init__(self, calc, kd, poisson_solver, dtype=float, **kwargs):
        """Store useful objects, e.g. lfc's for the various atomic functions.
            
        Depending on whether the system is periodic or finite, Poisson's equation
        is solved with FFT or multigrid techniques, respectively.

        Parameters
        ----------
        calc: Calculator
            Ground-state calculation.
        kd: KPointDescriptor
            Descriptor for the q-vectors of the dynamical matrix.
     
        """

        self.kd = kd
        self.dtype = dtype
        self.poisson = poisson_solver

        # Gamma wrt q-vector
        if self.kd.gamma:
            self.phase_cd = None
        else:
            assert self.kd.mynks == len(self.kd.ibzk_qc)

            self.phase_qcd = []
            sdisp_cd = calc.wfs.gd.sdisp_cd

            for q in range(self.kd.mynks):
                phase_cd = np.exp(2j * np.pi * \
                                  sdisp_cd * self.kd.ibzk_qc[q, :, np.newaxis])
                self.phase_qcd.append(phase_cd)
            
        # Store grid-descriptors
        self.gd = calc.density.gd
        self.finegd = calc.density.finegd

        # Steal setups for the lfc's
        setups = calc.wfs.setups

        # Store projector coefficients
        self.dH_asp = calc.hamiltonian.dH_asp.copy()
        
        # Localized functions:
        # core corections
        self.nct = LFC(self.gd, [[setup.nct] for setup in setups],
                       integral=[setup.Nct for setup in setups], dtype=self.dtype)
        # compensation charges
        #XXX what is the consequence of numerical errors in the integral ??
        self.ghat = LFC(self.finegd, [setup.ghat_l for setup in setups],
                        dtype=self.dtype)
        ## self.ghat = LFC(self.finegd, [setup.ghat_l for setup in setups],
        ##                 integral=sqrt(4 * pi), dtype=self.dtype)
        # vbar potential
        self.vbar = LFC(self.finegd, [[setup.vbar] for setup in setups],
                        dtype=self.dtype)

        # Expansion coefficients for the compensation charges
        self.Q_aL = calc.density.Q_aL.copy()
        
        # Grid transformer -- convert array from fine to coarse grid
        self.restrictor = Transformer(self.finegd, self.gd, nn=3,
                                      dtype=self.dtype, allocate=False)

        # Atom, cartesian coordinate and q-vector of the perturbation
        self.a = None
        self.v = None
        
        # Local q-vector index of the perturbation
        if self.kd.gamma:
            self.q = -1
        else:
            self.q = None

    def initialize(self, spos_ac):
        """Prepare the various attributes for a calculation."""

        # Set positions on LFC's
        self.nct.set_positions(spos_ac)
        self.ghat.set_positions(spos_ac)
        self.vbar.set_positions(spos_ac)

        if not self.kd.gamma:
            
            # Set q-vectors and update
            self.ghat.set_k_points(self.kd.ibzk_qc)
            self.ghat._update(spos_ac)
            # Set q-vectors and update
            self.vbar.set_k_points(self.kd.ibzk_qc)
            self.vbar._update(spos_ac)

            # Phase factor exp(iq.r) needed to obtian the periodic part of lfcs
            coor_vg = self.finegd.get_grid_point_coordinates()
            cell_cv = self.finegd.cell_cv
            # Convert to scaled coordinates
            scoor_cg = np.dot(la.inv(cell_cv), coor_vg.swapaxes(0, -2))
            scoor_cg = scoor_cg.swapaxes(1,-2)
            # Phase factor
            phase_qg = np.exp(2j * pi *
                              np.dot(self.kd.ibzk_qc, scoor_cg.swapaxes(0,-2)))
            self.phase_qg = phase_qg.swapaxes(1, -2)

        #XXX To be removed from this class !!
        # Setup the Poisson solver -- to be used on the fine grid
        self.poisson.set_grid_descriptor(self.finegd)
        self.poisson.initialize()

        # Grid transformer
        self.restrictor.allocate()

    def set_q(self, q):
        """Set the index of the q-vector of the perturbation."""

        assert not self.kd.gamma, "Gamma-point calculation"
        
        self.q = q

        # Update phases and Poisson solver
        self.phase_cd = self.phase_qcd[q]
        self.poisson.set_q(self.kd.ibzk_qc[q])

        # Invalidate calculated quantities
        # - local part of perturbing potential
        self.v1_G = None

    def set_av(self, a, v):
        """Set atom and cartesian component of the perturbation.

        Parameters
        ----------
        a: int
            Index of the atom.
        v: int 
            Cartesian component (0, 1 or 2) of the atomic displacement.
            
        """

        assert self.q is not None
        
        self.a = a
        self.v = v
        
        # Update derivative of local potential
        self.calculate_local_potential()
        
    def get_phase_cd(self):
        """Overwrite base class member function."""

        return self.phase_cd
    
    def has_q(self):
        """Overwrite base class member function."""

        return (not self.kd.gamma)

    def get_q(self):
        """Return q-vector."""

        assert not self.kd.gamma, "Gamma-point calculation."
        
        return self.kd.ibzk_qc[self.q]
    
    def solve_poisson(self, phi_g, rho_g):
        """Solve Poisson's equation for a Bloch-type charge distribution.

        More to come here ...
        
        Parameters
        ----------
        phi_g: GridDescriptor
            Grid for the solution of Poissons's equation.
        rho_g: GridDescriptor
            Grid with the charge distribution.

        """

        #assert phi_g.shape == rho_g.shape == self.phase_qg.shape[-3:], \
        #       ("Arrays have incompatible shapes.")
        assert self.q is not None, ("q-vector not set")
        
        # Gamma point calculation wrt the q-vector -> rho_g periodic
        if self.kd.gamma: 
            #XXX NOTICE: solve_neutral
            self.poisson.solve_neutral(phi_g, rho_g)
        else:
            # Divide out the phase factor to get the periodic part
            rhot_g = rho_g/self.phase_qg[self.q]

            # Solve Poisson's equation for the periodic part of the potential
            #XXX NOTICE: solve_neutral
            self.poisson.solve_neutral(phi_g, rhot_g)

            # Return to Bloch form
            phi_g *= self.phase_qg[self.q]

    def calculate_local_potential(self):
        """Derivate of the local potential wrt an atomic displacements.

        The local part of the PAW potential has contributions from the
        compensation charges (``ghat``) and a spherical symmetric atomic
        potential (``vbar``).
        
        """

        assert self.a is not None
        assert self.v is not None
        assert self.q is not None
        
        a = self.a
        v = self.v
        
        # Expansion coefficients for the ghat functions
        Q_aL = self.ghat.dict(zero=True)
        # Remember sign convention for add_derivative method
        # And be sure not to change the dtype of the arrays by assigning values
        # to array elements.
        Q_aL[a][:] = -1 * self.Q_aL[a]

        # Grid for derivative of compensation charges
        ghat1_g = self.finegd.zeros(dtype=self.dtype)
        self.ghat.add_derivative(a, v, ghat1_g, c_axi=Q_aL, q=self.q)
        
        # Solve Poisson's eq. for the potential from the periodic part of the
        # compensation charge derivative
        v1_g = self.finegd.zeros(dtype=self.dtype)
        self.solve_poisson(v1_g, ghat1_g)
        
        # Store potential from the compensation charge
        self.vghat1_g = v1_g.copy()
        
        # Add derivative of vbar - sign convention in add_derivative method
        c_ai = self.vbar.dict(zero=True)
        c_ai[a][0] = -1.
        self.vbar.add_derivative(a, v, v1_g, c_axi=c_ai, q=self.q)

        # Store potential for the evaluation of the energy derivative
        self.v1_g = v1_g.copy()
        
        # Transfer to coarse grid
        v1_G = self.gd.zeros(dtype=self.dtype)
        self.restrictor.apply(v1_g, v1_G, phases=self.phase_cd)

        self.v1_G = v1_G
        
    def apply(self, psi_nG, y_nG, wfs, k, kplusq):
        """Apply perturbation to unperturbed wave-functions.

        Parameters
        ----------
        psi_nG: ndarray
            Set of grid vectors to which the perturbation is applied.
        y_nG: ndarray
            Output vectors.
        wfs: WaveFunctions
            Instance of class ``WaveFunctions``.
        k: int
            Index of the k-point for the vectors.
        kplusq: int
            Index of the k+q vector.
            
        """

        assert self.a is not None
        assert self.v is not None
        assert self.q is not None
        assert psi_nG.ndim in (3, 4)
        assert tuple(self.gd.n_c) == psi_nG.shape[-3:]

        if psi_nG.ndim == 3:
            y_nG += self.v1_G * psi_nG
        else:
            y_nG += self.v1_G[np.newaxis, :] * psi_nG

        self.apply_nonlocal_potential(psi_nG, y_nG, wfs, k, kplusq)

    def apply_nonlocal_potential(self, psi_nG, y_nG, wfs, k, kplusq):
        """Derivate of the non-local PAW potential wrt an atomic displacement.

        Parameters
        ----------
        k: int
            Index of the k-point being operated on.
        kplusq: int
            Index of the k+q vector.
            
        """

        assert self.a is not None
        assert self.v is not None
        assert psi_nG.ndim in (3, 4)
        assert tuple(self.gd.n_c) == psi_nG.shape[-3:]
        
        if psi_nG.ndim == 3:
            n = 1
        else:
            n = psi_nG.shape[0] 
            
        a = self.a
        v = self.v
        
        P_ani = wfs.kpt_u[k].P_ani
        dP_aniv = wfs.kpt_u[k].dP_aniv
        pt = wfs.pt
        
        # < p_a^i | Psi_nk >
        P_ni = P_ani[a]
        # < dp_av^i | Psi_nk > - remember the sign convention of the derivative
        dP_ni = -1 * dP_aniv[a][...,v]
        
        # Expansion coefficients for the projectors on atom a
        dH_ii = unpack(self.dH_asp[a][0])
       
        # The derivative of the non-local PAW potential has two contributions
        # 1) Sum over projectors
        c_ni = np.dot(dP_ni, dH_ii)
        c_ani = pt.dict(shape=n, zero=True)
        c_ani[a] = c_ni
        # k+q !!
        pt.add(y_nG, c_ani, q=kplusq)

        # 2) Sum over derivatives of the projectors
        dc_ni = np.dot(P_ni, dH_ii)
        dc_ani = pt.dict(shape=n, zero=True)
        # Take care of sign of derivative in the coefficients
        dc_ani[a] = -1 * dc_ni
        # k+q !!
        pt.add_derivative(a, v, y_nG, dc_ani, q=kplusq)
Example #17
0
class RealSpaceDensity(Density):
    def __init__(self,
                 gd,
                 finegd,
                 nspins,
                 charge,
                 redistributor,
                 stencil=3,
                 background_charge=None):
        Density.__init__(self,
                         gd,
                         finegd,
                         nspins,
                         charge,
                         redistributor,
                         background_charge=background_charge)
        self.stencil = stencil

    def initialize(self, setups, timer, magmom_a, hund):
        Density.initialize(self, setups, timer, magmom_a, hund)

        # Interpolation function for the density:
        self.interpolator = Transformer(self.redistributor.aux_gd, self.finegd,
                                        self.stencil)

        spline_aj = []
        for setup in setups:
            if setup.nct is None:
                spline_aj.append([])
            else:
                spline_aj.append([setup.nct])
        self.nct = LFC(self.gd,
                       spline_aj,
                       integral=[setup.Nct for setup in setups],
                       forces=True,
                       cut=True)
        self.ghat = LFC(self.finegd, [setup.ghat_l for setup in setups],
                        integral=sqrt(4 * pi),
                        forces=True)

    def set_positions(self, spos_ac, rank_a=None):
        Density.set_positions(self, spos_ac, rank_a)
        self.nct_G = self.gd.zeros()
        self.nct.add(self.nct_G, 1.0 / self.nspins)

    def interpolate_pseudo_density(self, comp_charge=None):
        """Interpolate pseudo density to fine grid."""
        if comp_charge is None:
            comp_charge = self.calculate_multipole_moments()

        self.nt_sg = self.distribute_and_interpolate(self.nt_sG, self.nt_sg)

        # With periodic boundary conditions, the interpolation will
        # conserve the number of electrons.
        if not self.gd.pbc_c.all():
            # With zero-boundary conditions in one or more directions,
            # this is not the case.
            pseudo_charge = (self.background_charge.charge - self.charge -
                             comp_charge)
            if abs(pseudo_charge) > 1.0e-14:
                x = (pseudo_charge / self.finegd.integrate(self.nt_sg).sum())
                self.nt_sg *= x

    def interpolate(self, in_xR, out_xR=None):
        """Interpolate array(s)."""

        # ndim will be 3 in finite-difference mode and 1 when working
        # with the AtomPAW class (spherical atoms and 1d grids)
        ndim = self.gd.ndim

        if out_xR is None:
            out_xR = self.finegd.empty(in_xR.shape[:-ndim])

        a_xR = in_xR.reshape((-1, ) + in_xR.shape[-ndim:])
        b_xR = out_xR.reshape((-1, ) + out_xR.shape[-ndim:])

        for in_R, out_R in zip(a_xR, b_xR):
            self.interpolator.apply(in_R, out_R)

        return out_xR

    def distribute_and_interpolate(self, in_xR, out_xR=None):
        in_xR = self.redistributor.distribute(in_xR)
        return self.interpolate(in_xR, out_xR)

    def calculate_pseudo_charge(self):
        self.nt_g = self.nt_sg.sum(axis=0)
        self.rhot_g = self.nt_g.copy()
        self.ghat.add(self.rhot_g, self.Q_aL)
        self.background_charge.add_charge_to(self.rhot_g)

        if debug:
            charge = self.finegd.integrate(self.rhot_g) + self.charge
            if abs(charge) > self.charge_eps:
                raise RuntimeError('Charge not conserved: excess=%.9f' %
                                   charge)

    def get_pseudo_core_kinetic_energy_density_lfc(self):
        return LFC(self.gd, [[setup.tauct] for setup in self.setups],
                   forces=True,
                   cut=True)

    def calculate_dipole_moment(self):
        return self.finegd.calculate_dipole_moment(self.rhot_g)
Example #18
0
class RealSpaceHamiltonian(Hamiltonian):
    def __init__(self,
                 gd,
                 finegd,
                 nspins,
                 setups,
                 timer,
                 xc,
                 world,
                 vext=None,
                 psolver=None,
                 stencil=3,
                 redistributor=None):
        Hamiltonian.__init__(self,
                             gd,
                             finegd,
                             nspins,
                             setups,
                             timer,
                             xc,
                             world,
                             vext=vext,
                             redistributor=redistributor)

        # Solver for the Poisson equation:
        if psolver is None:
            psolver = {}
        if isinstance(psolver, dict):
            psolver = create_poisson_solver(**psolver)
        self.poisson = psolver
        self.poisson.set_grid_descriptor(self.finegd)

        # Restrictor function for the potential:
        self.restrictor = Transformer(self.finegd, self.redistributor.aux_gd,
                                      stencil)
        self.restrict = self.restrictor.apply

        self.vbar = LFC(self.finegd, [[setup.vbar] for setup in setups],
                        forces=True)
        self.vbar_g = None

    def restrict_and_collect(self, a_xg, b_xg=None, phases=None):
        if self.redistributor.enabled:
            tmp_xg = self.restrictor.apply(a_xg, output=None, phases=phases)
            b_xg = self.redistributor.collect(tmp_xg, b_xg)
        else:
            b_xg = self.restrictor.apply(a_xg, output=b_xg, phases=phases)
        return b_xg

    def __str__(self):
        s = Hamiltonian.__str__(self)

        degree = self.restrictor.nn * 2 - 1
        name = ['linear', 'cubic', 'quintic', 'heptic'][degree // 2]
        s += ('  Interpolation: tri-%s ' % name +
              '(%d. degree polynomial)\n' % degree)
        s += '  Poisson solver: %s' % self.poisson.get_description()
        return s

    def set_positions(self, spos_ac, rank_a):
        Hamiltonian.set_positions(self, spos_ac, rank_a)
        if self.vbar_g is None:
            self.vbar_g = self.finegd.empty()
        self.vbar_g[:] = 0.0
        self.vbar.add(self.vbar_g)

    def update_pseudo_potential(self, dens):
        self.timer.start('vbar')
        e_zero = self.finegd.integrate(self.vbar_g,
                                       dens.nt_g,
                                       global_integral=False)

        vt_g = self.vt_sg[0]
        vt_g[:] = self.vbar_g
        self.timer.stop('vbar')

        e_external = 0.0
        if self.vext is not None:
            vext_g = self.vext.get_potential(self.finegd)
            vt_g += vext_g
            e_external = self.finegd.integrate(vext_g,
                                               dens.rhot_g,
                                               global_integral=False)

        if self.nspins == 2:
            self.vt_sg[1] = vt_g

        self.timer.start('XC 3D grid')
        e_xc = self.xc.calculate(self.finegd, dens.nt_sg, self.vt_sg)
        e_xc /= self.finegd.comm.size
        self.timer.stop('XC 3D grid')

        self.timer.start('Poisson')
        # npoisson is the number of iterations:
        self.npoisson = self.poisson.solve(self.vHt_g,
                                           dens.rhot_g,
                                           charge=-dens.charge)
        self.timer.stop('Poisson')

        self.timer.start('Hartree integrate/restrict')
        e_coulomb = 0.5 * self.finegd.integrate(
            self.vHt_g, dens.rhot_g, global_integral=False)

        for vt_g in self.vt_sg:
            vt_g += self.vHt_g

        self.timer.stop('Hartree integrate/restrict')
        return np.array([e_coulomb, e_zero, e_external, e_xc])

    def calculate_kinetic_energy(self, density):
        # XXX new timer item for kinetic energy?
        self.timer.start('Hartree integrate/restrict')
        self.restrict_and_collect(self.vt_sg, self.vt_sG)

        e_kinetic = 0.0
        s = 0
        for vt_G, nt_G in zip(self.vt_sG, density.nt_sG):
            if self.ref_vt_sG is not None:
                vt_G += self.ref_vt_sG[s]

            if s < self.nspins:
                e_kinetic -= self.gd.integrate(vt_G,
                                               nt_G - density.nct_G,
                                               global_integral=False)
            else:
                e_kinetic -= self.gd.integrate(vt_G,
                                               nt_G,
                                               global_integral=False)
            s += 1
        self.timer.stop('Hartree integrate/restrict')
        return e_kinetic

    def calculate_atomic_hamiltonians(self, dens):
        def getshape(a):
            return sum(2 * l + 1 for l, _ in enumerate(self.setups[a].ghat_l)),

        W_aL = ArrayDict(self.atomdist.aux_partition, getshape, float)
        if self.vext:
            vext_g = self.vext.get_potential(self.finegd)
            dens.ghat.integrate(self.vHt_g + vext_g, W_aL)
        else:
            dens.ghat.integrate(self.vHt_g, W_aL)

        return self.atomdist.to_work(self.atomdist.from_aux(W_aL))

    def calculate_forces2(self, dens, ghat_aLv, nct_av, vbar_av):
        if self.nspins == 2:
            vt_G = self.vt_sG.mean(0)
        else:
            vt_G = self.vt_sG[0]

        self.vbar.derivative(dens.nt_g, vbar_av)
        if self.vext:
            vext_g = self.vext.get_potential(self.finegd)
            dens.ghat.derivative(self.vHt_g + vext_g, ghat_aLv)
        else:
            dens.ghat.derivative(self.vHt_g, ghat_aLv)
        dens.nct.derivative(vt_G, nct_av)

    def get_electrostatic_potential(self, dens):
        self.update(dens)

        v_g = self.finegd.collect(self.vHt_g, broadcast=True)
        v_g = self.finegd.zero_pad(v_g)
        if hasattr(self.poisson, 'correction'):
            assert self.poisson.c == 2
            v_g[:, :, 0] = self.poisson.correction
        return v_g