Beispiel #1
0
    def __init__(self, comm, broadcast_comm, gd, aux_gd, spos_ac):
        self.comm = comm
        self.broadcast_comm = broadcast_comm
        self.gd = gd
        self.aux_gd = aux_gd

        rank_a = gd.get_ranks_from_positions(spos_ac)
        aux_rank_a = aux_gd.get_ranks_from_positions(spos_ac)
        self.partition = AtomPartition(gd.comm, rank_a, name='gd')

        if gd is aux_gd:
            name = 'aux-unextended'
        else:
            name = 'aux-extended'
        self.aux_partition = AtomPartition(aux_gd.comm, aux_rank_a, name=name)

        self.work_partition = AtomPartition(comm,
                                            np.zeros(len(spos_ac)),
                                            name='work').as_even_partition()

        if gd is aux_gd:
            aux_broadcast_comm = gd.comm.new_communicator([gd.comm.rank])
        else:
            aux_broadcast_comm = broadcast_comm

        self.aux_dist = AtomicMatrixDistributor(self.partition,
                                                aux_broadcast_comm,
                                                self.aux_partition)
        self.work_dist = AtomicMatrixDistributor(self.partition,
                                                 broadcast_comm,
                                                 self.work_partition)
Beispiel #2
0
    def set_positions(self, spos_ac, rank_a):
        atom_partition = AtomPartition(self.gd.comm, rank_a)

        self.spos_ac = spos_ac
        self.vbar.set_positions(spos_ac)
        self.xc.set_positions(spos_ac)

        # 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.
        # XXX what purpose does this serve?  In what case does it happen?
        # How would one even go about figuring it out?  Why does it all have
        # to be so unreadable? -Ask
        #
        if (self.rank_a is not None and self.dH_asp is None
                and (rank_a == self.gd.comm.rank).any()):
            self.dH_asp = {}

        if self.rank_a is not None and self.dH_asp is not None:
            self.timer.start('Redistribute')

            def get_empty(a):
                ni = self.setups[a].ni
                return np.empty((self.ns, ni * (ni + 1) // 2))

            self.atom_partition.redistribute(atom_partition, self.dH_asp,
                                             get_empty)
            self.timer.stop('Redistribute')

        self.rank_a = rank_a
        self.atom_partition = atom_partition
        self.dh_distributor = AtomicMatrixDistributor(atom_partition,
                                                      self.setups,
                                                      self.kptband_comm,
                                                      self.ns)
 def distribute_Dresp_asp(self, Dresp_asp):
     d("distribute_Dresp_asp")
     # okay, this is a bit hacky since we have to call this from
     # several different places.  Maybe Mikael can figure out something
     # smarter.  -Ask
     from gpaw.utilities.partition import AtomicMatrixDistributor
     amd = AtomicMatrixDistributor(self.density.atom_partition,
                                 self.density.setups, self.wfs.kptband_comm,
                                 self.density.ns)
     return amd.distribute(Dresp_asp)
 def distribute_Dresp_asp(self, Dresp_asp):
     d("distribute_Dresp_asp")
     # okay, this is a bit hacky since we have to call this from
     # several different places.  Maybe Mikael can figure out something
     # smarter.  -Ask
     from gpaw.utilities.partition import AtomicMatrixDistributor
     amd = AtomicMatrixDistributor(self.density.atom_partition,
                                   self.density.setups,
                                   self.wfs.kptband_comm, self.density.ns)
     return amd.distribute(Dresp_asp)
Beispiel #5
0
class AtomDistributions:
    def __init__(self, comm, broadcast_comm, gd, aux_gd, spos_ac):
        self.comm = comm
        self.broadcast_comm = broadcast_comm
        self.gd = gd
        self.aux_gd = aux_gd

        rank_a = gd.get_ranks_from_positions(spos_ac)
        aux_rank_a = aux_gd.get_ranks_from_positions(spos_ac)
        self.partition = AtomPartition(gd.comm, rank_a, name='gd')

        if gd is aux_gd:
            name = 'aux-unextended'
        else:
            name = 'aux-extended'
        self.aux_partition = AtomPartition(aux_gd.comm, aux_rank_a, name=name)

        self.work_partition = AtomPartition(comm,
                                            np.zeros(len(spos_ac)),
                                            name='work').as_even_partition()

        if gd is aux_gd:
            aux_broadcast_comm = gd.comm.new_communicator([gd.comm.rank])
        else:
            aux_broadcast_comm = broadcast_comm

        self.aux_dist = AtomicMatrixDistributor(self.partition,
                                                aux_broadcast_comm,
                                                self.aux_partition)
        self.work_dist = AtomicMatrixDistributor(self.partition,
                                                 broadcast_comm,
                                                 self.work_partition)

    def to_aux(self, arraydict):
        if self.gd is self.aux_gd:
            return arraydict.copy()
        return self.aux_dist.distribute(arraydict)

    def from_aux(self, arraydict):
        if self.gd is self.aux_gd:
            return arraydict.copy()
        return self.aux_dist.collect(arraydict)

    def to_work(self, arraydict):
        return self.work_dist.distribute(arraydict)

    def from_work(self, arraydict):
        return self.work_dist.collect(arraydict)
    def set_positions(self, spos_ac, rank_a):
        atom_partition = AtomPartition(self.gd.comm, rank_a)

        self.spos_ac = spos_ac
        self.vbar.set_positions(spos_ac)
        self.xc.set_positions(spos_ac)

        # 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.
        # XXX what purpose does this serve?  In what case does it happen?
        # How would one even go about figuring it out?  Why does it all have
        # to be so unreadable? -Ask
        #
        if self.rank_a is not None and self.dH_asp is None and (rank_a == self.gd.comm.rank).any():
            self.dH_asp = {}

        if self.rank_a is not None and self.dH_asp is not None:
            self.timer.start("Redistribute")

            def get_empty(a):
                ni = self.setups[a].ni
                return np.empty((self.ns, ni * (ni + 1) // 2))

            self.atom_partition.redistribute(atom_partition, self.dH_asp, get_empty)
            self.timer.stop("Redistribute")

        self.rank_a = rank_a
        self.atom_partition = atom_partition
        self.dh_distributor = AtomicMatrixDistributor(atom_partition, self.setups, self.kptband_comm, self.ns)
class Hamiltonian:
    """Hamiltonian object.

    Attributes:
     =============== =====================================================
     ``xc``          ``XC3DGrid`` object.
     ``poisson``     ``PoissonSolver``.
     ``gd``          Grid descriptor for coarse grids.
     ``finegd``      Grid descriptor for fine grids.
     ``restrict``    Function for restricting the effective potential.
     =============== =====================================================

    Soft and smooth pseudo functions on uniform 3D grids:
     ========== =========================================
     ``vHt_g``  Hartree potential on the fine grid.
     ``vt_sG``  Effective potential on the coarse grid.
     ``vt_sg``  Effective potential on the fine grid.
     ========== =========================================

    Energy contributions and forces:

    =========== ==========================================
                Description
    =========== ==========================================
    ``Ekin``    Kinetic energy.
    ``Epot``    Potential energy.
    ``Etot``    Total energy.
    ``Exc``     Exchange-Correlation energy.
    ``Eext``    Energy of external potential
    ``Eref``    Reference energy for all-electron atoms.
    ``S``       Entropy.
    ``Ebar``    Should be close to zero!
    =========== ==========================================

    """

    def __init__(self, gd, finegd, nspins, setups, timer, xc, world, kptband_comm, vext=None, collinear=True):
        """Create the Hamiltonian."""
        self.gd = gd
        self.finegd = finegd
        self.nspins = nspins
        self.setups = setups
        self.timer = timer
        self.xc = xc
        self.collinear = collinear
        self.ncomp = 2 - int(collinear)
        self.ns = self.nspins * self.ncomp ** 2
        self.world = world
        self.kptband_comm = kptband_comm

        self.dH_asp = None

        # The external potential
        self.vext = vext

        self.vt_sG = None
        self.vHt_g = None
        self.vt_sg = None

        self.rank_a = None
        self.atom_partition = None

        self.Ekin0 = None
        self.Ekin = None
        self.Epot = None
        self.Ebar = None
        self.Eext = None
        self.Exc = None
        self.Etot = None
        self.S = None

        self.ref_vt_sG = None
        self.ref_dH_asp = None

    def summary(self, fd):
        fd.write("XC and Coulomb potentials evaluated on a %d*%d*%d grid\n" % tuple(self.finegd.N_c))

    def set_positions(self, spos_ac, rank_a):
        atom_partition = AtomPartition(self.gd.comm, rank_a)

        self.spos_ac = spos_ac
        self.vbar.set_positions(spos_ac)
        self.xc.set_positions(spos_ac)

        # 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.
        # XXX what purpose does this serve?  In what case does it happen?
        # How would one even go about figuring it out?  Why does it all have
        # to be so unreadable? -Ask
        #
        if self.rank_a is not None and self.dH_asp is None and (rank_a == self.gd.comm.rank).any():
            self.dH_asp = {}

        if self.rank_a is not None and self.dH_asp is not None:
            self.timer.start("Redistribute")

            def get_empty(a):
                ni = self.setups[a].ni
                return np.empty((self.ns, ni * (ni + 1) // 2))

            self.atom_partition.redistribute(atom_partition, self.dH_asp, get_empty)
            self.timer.stop("Redistribute")

        self.rank_a = rank_a
        self.atom_partition = atom_partition
        self.dh_distributor = AtomicMatrixDistributor(atom_partition, self.setups, self.kptband_comm, self.ns)

    def aoom(self, DM, a, l, scale=1):
        """Atomic Orbital Occupation Matrix.
        
        Determine the Atomic Orbital Occupation Matrix (aoom) for a
        given l-quantum number.
        
        This operation, takes the density matrix (DM), which for
        example is given by unpack2(D_asq[i][spin]), and corrects for
        the overlap between the selected orbitals (l) upon which the
        the density is expanded (ex <p|p*>,<p|p>,<p*|p*> ).

        Returned is only the "corrected" part of the density matrix,
        which represents the orbital occupation matrix for l=2 this is
        a 5x5 matrix.
        """
        S = self.setups[a]
        l_j = S.l_j
        lq = S.lq
        nl = np.where(np.equal(l_j, l))[0]
        V = np.zeros(np.shape(DM))
        if len(nl) == 2:
            aa = (nl[0]) * len(l_j) - ((nl[0] - 1) * (nl[0]) / 2)
            bb = (nl[1]) * len(l_j) - ((nl[1] - 1) * (nl[1]) / 2)
            ab = aa + nl[1] - nl[0]

            if not scale:
                lq_a = lq[aa]
                lq_ab = lq[ab]
                lq_b = lq[bb]
            else:
                lq_a = 1
                lq_ab = lq[ab] / lq[aa]
                lq_b = lq[bb] / lq[aa]

            # and the correct entrances in the DM
            nn = (2 * np.array(l_j) + 1)[0 : nl[0]].sum()
            mm = (2 * np.array(l_j) + 1)[0 : nl[1]].sum()

            # finally correct and add the four submatrices of NC_DM
            A = DM[nn : nn + 2 * l + 1, nn : nn + 2 * l + 1] * (lq_a)
            B = DM[nn : nn + 2 * l + 1, mm : mm + 2 * l + 1] * (lq_ab)
            C = DM[mm : mm + 2 * l + 1, nn : nn + 2 * l + 1] * (lq_ab)
            D = DM[mm : mm + 2 * l + 1, mm : mm + 2 * l + 1] * (lq_b)

            V[nn : nn + 2 * l + 1, nn : nn + 2 * l + 1] = +(lq_a)
            V[nn : nn + 2 * l + 1, mm : mm + 2 * l + 1] = +(lq_ab)
            V[mm : mm + 2 * l + 1, nn : nn + 2 * l + 1] = +(lq_ab)
            V[mm : mm + 2 * l + 1, mm : mm + 2 * l + 1] = +(lq_b)

            return A + B + C + D, V
        else:
            nn = (2 * np.array(l_j) + 1)[0 : nl[0]].sum()
            A = DM[nn : nn + 2 * l + 1, nn : nn + 2 * l + 1] * lq[-1]
            V[nn : nn + 2 * l + 1, nn : nn + 2 * l + 1] = +lq[-1]
            return A, V

    def update(self, density):
        """Calculate effective potential.

        The XC-potential and the Hartree potential are evaluated on
        the fine grid, and the sum is then restricted to the coarse
        grid."""

        self.timer.start("Hamiltonian")

        if self.vt_sg is None:
            self.timer.start("Initialize Hamiltonian")
            self.vt_sg = self.finegd.empty(self.ns)
            self.vHt_g = self.finegd.zeros()
            self.vt_sG = self.gd.empty(self.ns)
            self.poisson.initialize()
            self.timer.stop("Initialize Hamiltonian")

        Ekin, Epot, Ebar, Eext, Exc, W_aL = self.update_pseudo_potential(density)

        self.timer.start("Atomic")
        self.dH_asp = None  # XXXX

        dH_asp = {}
        for a, D_sp in density.D_asp.items():
            W_L = W_aL[a]
            setup = self.setups[a]

            D_p = D_sp[: self.nspins].sum(0)
            dH_p = setup.K_p + setup.M_p + setup.MB_p + 2.0 * np.dot(setup.M_pp, D_p) + np.dot(setup.Delta_pL, W_L)
            Ekin += np.dot(setup.K_p, D_p) + setup.Kc
            Ebar += setup.MB + np.dot(setup.MB_p, D_p)
            Epot += setup.M + np.dot(D_p, (setup.M_p + np.dot(setup.M_pp, D_p)))

            if self.vext is not None:
                vext = self.vext.get_taylor(spos_c=self.spos_ac[a, :])
                # Tailor expansion to the zeroth order
                Eext += vext[0][0] * (sqrt(4 * pi) * density.Q_aL[a][0] + setup.Z)
                dH_p += vext[0][0] * sqrt(4 * pi) * setup.Delta_pL[:, 0]
                if len(vext) > 1:
                    # Tailor expansion to the first order
                    Eext += sqrt(4 * pi / 3) * np.dot(vext[1], density.Q_aL[a][1:4])
                    # there must be a better way XXXX
                    Delta_p1 = np.array([setup.Delta_pL[:, 1], setup.Delta_pL[:, 2], setup.Delta_pL[:, 3]])
                    dH_p += sqrt(4 * pi / 3) * np.dot(vext[1], Delta_p1)

            dH_asp[a] = dH_sp = np.zeros_like(D_sp)

            if setup.HubU is not None:
                assert self.collinear
                nspins = len(D_sp)

                l_j = setup.l_j
                l = setup.Hubl
                scale = setup.Hubs
                nl = np.where(np.equal(l_j, l))[0]
                nn = (2 * np.array(l_j) + 1)[0 : nl[0]].sum()

                for D_p, H_p in zip(D_sp, dH_asp[a]):
                    [N_mm, V] = self.aoom(unpack2(D_p), a, l, scale)
                    N_mm = N_mm / 2 * nspins

                    Eorb = setup.HubU / 2.0 * (N_mm - np.dot(N_mm, N_mm)).trace()
                    Vorb = setup.HubU * (0.5 * np.eye(2 * l + 1) - N_mm)
                    Exc += Eorb
                    if nspins == 1:
                        # add contribution of other spin manyfold
                        Exc += Eorb

                    if len(nl) == 2:
                        mm = (2 * np.array(l_j) + 1)[0 : nl[1]].sum()

                        V[nn : nn + 2 * l + 1, nn : nn + 2 * l + 1] *= Vorb
                        V[mm : mm + 2 * l + 1, nn : nn + 2 * l + 1] *= Vorb
                        V[nn : nn + 2 * l + 1, mm : mm + 2 * l + 1] *= Vorb
                        V[mm : mm + 2 * l + 1, mm : mm + 2 * l + 1] *= Vorb
                    else:
                        V[nn : nn + 2 * l + 1, nn : nn + 2 * l + 1] *= Vorb

                    Htemp = unpack(H_p)
                    Htemp += V
                    H_p[:] = pack2(Htemp)

            dH_sp[: self.nspins] += dH_p
            if self.ref_dH_asp:
                dH_sp += self.ref_dH_asp[a]
            # We are not yet done with dH_sp; still need XC correction below

        Ddist_asp = self.dh_distributor.distribute(density.D_asp)

        dHdist_asp = {}
        Exca = 0.0
        self.timer.start("XC Correction")
        for a, D_sp in Ddist_asp.items():
            setup = self.setups[a]
            dH_sp = np.zeros_like(D_sp)
            Exca += self.xc.calculate_paw_correction(setup, D_sp, dH_sp, a=a)
            # XXX Exc are added on the "wrong" distribution; sum only works
            # when gd.comm and distribution comm are the same
            dHdist_asp[a] = dH_sp
        self.timer.stop("XC Correction")

        dHdist_asp = self.dh_distributor.collect(dHdist_asp)

        # Exca has contributions from all cores so modify it so it is
        # parallel in the same way as the other energies.
        Exca = self.world.sum(Exca)
        if self.gd.comm.rank == 0:
            Exc += Exca

        assert len(dHdist_asp) == len(self.atom_partition.my_indices)

        for a, D_sp in density.D_asp.items():
            dH_sp = dH_asp[a]
            dH_sp += dHdist_asp[a]
            Ekin -= (D_sp * dH_sp).sum()  # NCXXX
        self.dH_asp = dH_asp
        self.timer.stop("Atomic")

        # Make corrections due to non-local xc:
        # xcfunc = self.xc.xcfunc
        self.Enlxc = 0.0  # XXXxcfunc.get_non_local_energy()
        Ekin += self.xc.get_kinetic_energy_correction() / self.gd.comm.size

        energies = np.array([Ekin, Epot, Ebar, Eext, Exc])
        self.timer.start("Communicate energies")
        self.gd.comm.sum(energies)
        # Make sure that all CPUs have the same energies
        self.world.broadcast(energies, 0)
        self.timer.stop("Communicate energies")
        (self.Ekin0, self.Epot, self.Ebar, self.Eext, self.Exc) = energies

        # self.Exc += self.Enlxc
        # self.Ekin0 += self.Enlkin

        self.timer.stop("Hamiltonian")

    def get_energy(self, occupations):
        self.Ekin = self.Ekin0 + occupations.e_band
        self.S = occupations.e_entropy

        # Total free energy:
        self.Etot = self.Ekin + self.Epot + self.Eext + self.Ebar + self.Exc - self.S

        return self.Etot

    def linearize_to_xc(self, new_xc, density):
        # Store old hamiltonian
        ref_vt_sG = self.vt_sG.copy()
        ref_dH_asp = {}
        for a, dH_sp in self.dH_asp.items():
            ref_dH_asp[a] = dH_sp.copy()
        self.xc = new_xc
        self.xc.set_positions(self.spos_ac)
        self.update(density)

        ref_vt_sG -= self.vt_sG
        for a, dH_sp in self.dH_asp.items():
            ref_dH_asp[a] -= dH_sp
        self.ref_vt_sG = ref_vt_sG
        self.ref_dH_asp = ref_dH_asp

    def calculate_forces(self, dens, F_av):
        ghat_aLv = dens.ghat.dict(derivative=True)
        nct_av = dens.nct.dict(derivative=True)
        vbar_av = self.vbar.dict(derivative=True)

        self.calculate_forces2(dens, ghat_aLv, nct_av, vbar_av)

        # Force from compensation charges:
        for a, dF_Lv in ghat_aLv.items():
            F_av[a] += np.dot(dens.Q_aL[a], dF_Lv)

        # Force from smooth core charge:
        for a, dF_v in nct_av.items():
            F_av[a] += dF_v[0]

        # Force from zero potential:
        for a, dF_v in vbar_av.items():
            F_av[a] += dF_v[0]

        self.xc.add_forces(F_av)
        self.gd.comm.sum(F_av, 0)

    def apply_local_potential(self, psit_nG, Htpsit_nG, s):
        """Apply the Hamiltonian operator to a set of vectors.

        XXX Parameter description is deprecated!
        
        Parameters:

        a_nG: ndarray
            Set of vectors to which the overlap operator is applied.
        b_nG: ndarray, output
            Resulting H times a_nG vectors.
        kpt: KPoint object
            k-point object defined in kpoint.py.
        calculate_projections: bool
            When True, the integrals of projector times vectors
            P_ni = <p_i | a_nG> are calculated.
            When False, existing P_uni are used
        local_part_only: bool
            When True, the non-local atomic parts of the Hamiltonian
            are not applied and calculate_projections is ignored.
        
        """
        vt_G = self.vt_sG[s]
        if psit_nG.ndim == 3:
            Htpsit_nG += psit_nG * vt_G
        else:
            for psit_G, Htpsit_G in zip(psit_nG, Htpsit_nG):
                Htpsit_G += psit_G * vt_G

    def apply(self, a_xG, b_xG, wfs, kpt, calculate_P_ani=True):
        """Apply the Hamiltonian operator to a set of vectors.

        Parameters:

        a_nG: ndarray
            Set of vectors to which the overlap operator is applied.
        b_nG: ndarray, output
            Resulting S times a_nG vectors.
        wfs: WaveFunctions
            Wave-function object defined in wavefunctions.py
        kpt: KPoint object
            k-point object defined in kpoint.py.
        calculate_P_ani: bool
            When True, the integrals of projector times vectors
            P_ni = <p_i | a_nG> are calculated.
            When False, existing P_ani are used
        
        """

        wfs.kin.apply(a_xG, b_xG, kpt.phase_cd)
        self.apply_local_potential(a_xG, b_xG, kpt.s)
        shape = a_xG.shape[:-3]
        P_axi = wfs.pt.dict(shape)

        if calculate_P_ani:  # TODO calculate_P_ani=False is experimental
            wfs.pt.integrate(a_xG, P_axi, kpt.q)
        else:
            for a, P_ni in kpt.P_ani.items():
                P_axi[a][:] = P_ni

        for a, P_xi in P_axi.items():
            dH_ii = unpack(self.dH_asp[a][kpt.s])
            P_axi[a] = np.dot(P_xi, dH_ii)
        wfs.pt.add(b_xG, P_axi, kpt.q)

    def get_xc_difference(self, xc, density):
        """Calculate non-selfconsistent XC-energy difference."""
        if density.nt_sg is None:
            density.interpolate_pseudo_density()
        nt_sg = density.nt_sg
        if hasattr(xc, "hybrid"):
            xc.calculate_exx()
        Exc = xc.calculate(density.finegd, nt_sg) / self.gd.comm.size
        for a, D_sp in density.D_asp.items():
            setup = self.setups[a]
            Exc += xc.calculate_paw_correction(setup, D_sp)
        Exc = self.gd.comm.sum(Exc)
        return Exc - self.Exc

    def estimate_memory(self, mem):
        nbytes = self.gd.bytecount()
        nfinebytes = self.finegd.bytecount()
        arrays = mem.subnode("Arrays", 0)
        arrays.subnode("vHt_g", nfinebytes)
        arrays.subnode("vt_sG", self.nspins * nbytes)
        arrays.subnode("vt_sg", self.nspins * nfinebytes)
        self.xc.estimate_memory(mem.subnode("XC"))
        self.poisson.estimate_memory(mem.subnode("Poisson"))
        self.vbar.estimate_memory(mem.subnode("vbar"))

    def read(self, reader, parallel):
        self.Ekin = reader["Ekin"]
        self.Epot = reader["Epot"]
        self.Ebar = reader["Ebar"]
        try:
            self.Eext = reader["Eext"]
        except (AttributeError, KeyError):
            self.Eext = 0.0
        self.Exc = reader["Exc"]
        self.S = reader["S"]
        self.Etot = reader.get("PotentialEnergy", broadcast=True) - 0.5 * self.S

        if not reader.has_array("PseudoPotential"):
            return

        hdf5 = hasattr(reader, "hdf5")
        version = reader["version"]

        # Read pseudo potential on the coarse grid
        # and broadcast on kpt/band comm:
        if version > 0.3:
            self.vt_sG = self.gd.empty(self.nspins)
            if hdf5:
                indices = [slice(0, self.nspins)] + self.gd.get_slice()
                do_read = self.kptband_comm.rank == 0
                reader.get("PseudoPotential", out=self.vt_sG, parallel=parallel, read=do_read, *indices)  # XXX read=?
                self.kptband_comm.broadcast(self.vt_sG, 0)
            else:
                for s in range(self.nspins):
                    self.gd.distribute(reader.get("PseudoPotential", s), self.vt_sG[s])

        # Read non-local part of hamiltonian
        self.dH_asp = {}
        natoms = len(self.setups)
        self.rank_a = np.zeros(natoms, int)
        if version > 0.3:
            all_H_sp = reader.get("NonLocalPartOfHamiltonian", broadcast=True)

        if self.gd.comm.rank == 0 and version > 0.3:
            self.dH_asp = read_atomic_matrices(all_H_sp, self.setups)
Beispiel #8
0
class Hamiltonian:
    """Hamiltonian object.

    Attributes:
     =============== =====================================================
     ``xc``          ``XC3DGrid`` object.
     ``poisson``     ``PoissonSolver``.
     ``gd``          Grid descriptor for coarse grids.
     ``finegd``      Grid descriptor for fine grids.
     ``restrict``    Function for restricting the effective potential.
     =============== =====================================================

    Soft and smooth pseudo functions on uniform 3D grids:
     ========== =========================================
     ``vHt_g``  Hartree potential on the fine grid.
     ``vt_sG``  Effective potential on the coarse grid.
     ``vt_sg``  Effective potential on the fine grid.
     ========== =========================================

    Energy contributions and forces:

    =========== ==========================================
                Description
    =========== ==========================================
    ``Ekin``    Kinetic energy.
    ``Epot``    Potential energy.
    ``Etot``    Total energy.
    ``Exc``     Exchange-Correlation energy.
    ``Eext``    Energy of external potential
    ``Eref``    Reference energy for all-electron atoms.
    ``S``       Entropy.
    ``Ebar``    Should be close to zero!
    =========== ==========================================

    """
    def __init__(self,
                 gd,
                 finegd,
                 nspins,
                 setups,
                 timer,
                 xc,
                 world,
                 kptband_comm,
                 vext=None,
                 collinear=True):
        """Create the Hamiltonian."""
        self.gd = gd
        self.finegd = finegd
        self.nspins = nspins
        self.setups = setups
        self.timer = timer
        self.xc = xc
        self.collinear = collinear
        self.ncomp = 2 - int(collinear)
        self.ns = self.nspins * self.ncomp**2
        self.world = world
        self.kptband_comm = kptband_comm

        self.dH_asp = None

        # The external potential
        self.vext = vext

        self.vt_sG = None
        self.vHt_g = None
        self.vt_sg = None

        self.rank_a = None
        self.atom_partition = None

        self.Ekin0 = None
        self.Ekin = None
        self.Epot = None
        self.Ebar = None
        self.Eext = None
        self.Exc = None
        self.Etot = None
        self.S = None

        self.ref_vt_sG = None
        self.ref_dH_asp = None

    def summary(self, fd):
        fd.write('XC and Coulomb potentials evaluated on a %d*%d*%d grid\n' %
                 tuple(self.finegd.N_c))

    def set_positions(self, spos_ac, rank_a):
        atom_partition = AtomPartition(self.gd.comm, rank_a)

        self.spos_ac = spos_ac
        self.vbar.set_positions(spos_ac)
        self.xc.set_positions(spos_ac)

        # 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.
        # XXX what purpose does this serve?  In what case does it happen?
        # How would one even go about figuring it out?  Why does it all have
        # to be so unreadable? -Ask
        #
        if (self.rank_a is not None and self.dH_asp is None
                and (rank_a == self.gd.comm.rank).any()):
            self.dH_asp = {}

        if self.rank_a is not None and self.dH_asp is not None:
            self.timer.start('Redistribute')

            def get_empty(a):
                ni = self.setups[a].ni
                return np.empty((self.ns, ni * (ni + 1) // 2))

            self.atom_partition.redistribute(atom_partition, self.dH_asp,
                                             get_empty)
            self.timer.stop('Redistribute')

        self.rank_a = rank_a
        self.atom_partition = atom_partition
        self.dh_distributor = AtomicMatrixDistributor(atom_partition,
                                                      self.setups,
                                                      self.kptband_comm,
                                                      self.ns)

    def aoom(self, DM, a, l, scale=1):
        """Atomic Orbital Occupation Matrix.
        
        Determine the Atomic Orbital Occupation Matrix (aoom) for a
        given l-quantum number.
        
        This operation, takes the density matrix (DM), which for
        example is given by unpack2(D_asq[i][spin]), and corrects for
        the overlap between the selected orbitals (l) upon which the
        the density is expanded (ex <p|p*>,<p|p>,<p*|p*> ).

        Returned is only the "corrected" part of the density matrix,
        which represents the orbital occupation matrix for l=2 this is
        a 5x5 matrix.
        """
        S = self.setups[a]
        l_j = S.l_j
        lq = S.lq
        nl = np.where(np.equal(l_j, l))[0]
        V = np.zeros(np.shape(DM))
        if len(nl) == 2:
            aa = (nl[0]) * len(l_j) - ((nl[0] - 1) * (nl[0]) / 2)
            bb = (nl[1]) * len(l_j) - ((nl[1] - 1) * (nl[1]) / 2)
            ab = aa + nl[1] - nl[0]

            if not scale:
                lq_a = lq[aa]
                lq_ab = lq[ab]
                lq_b = lq[bb]
            else:
                lq_a = 1
                lq_ab = lq[ab] / lq[aa]
                lq_b = lq[bb] / lq[aa]

            # and the correct entrances in the DM
            nn = (2 * np.array(l_j) + 1)[0:nl[0]].sum()
            mm = (2 * np.array(l_j) + 1)[0:nl[1]].sum()

            # finally correct and add the four submatrices of NC_DM
            A = DM[nn:nn + 2 * l + 1, nn:nn + 2 * l + 1] * (lq_a)
            B = DM[nn:nn + 2 * l + 1, mm:mm + 2 * l + 1] * (lq_ab)
            C = DM[mm:mm + 2 * l + 1, nn:nn + 2 * l + 1] * (lq_ab)
            D = DM[mm:mm + 2 * l + 1, mm:mm + 2 * l + 1] * (lq_b)

            V[nn:nn + 2 * l + 1, nn:nn + 2 * l + 1] = +(lq_a)
            V[nn:nn + 2 * l + 1, mm:mm + 2 * l + 1] = +(lq_ab)
            V[mm:mm + 2 * l + 1, nn:nn + 2 * l + 1] = +(lq_ab)
            V[mm:mm + 2 * l + 1, mm:mm + 2 * l + 1] = +(lq_b)

            return A + B + C + D, V
        else:
            nn = (2 * np.array(l_j) + 1)[0:nl[0]].sum()
            A = DM[nn:nn + 2 * l + 1, nn:nn + 2 * l + 1] * lq[-1]
            V[nn:nn + 2 * l + 1, nn:nn + 2 * l + 1] = +lq[-1]
            return A, V

    def update(self, density):
        """Calculate effective potential.

        The XC-potential and the Hartree potential are evaluated on
        the fine grid, and the sum is then restricted to the coarse
        grid."""

        self.timer.start('Hamiltonian')

        if self.vt_sg is None:
            self.timer.start('Initialize Hamiltonian')
            self.vt_sg = self.finegd.empty(self.ns)
            self.vHt_g = self.finegd.zeros()
            self.vt_sG = self.gd.empty(self.ns)
            self.poisson.initialize()
            self.timer.stop('Initialize Hamiltonian')

        Ekin, Epot, Ebar, Eext, Exc, W_aL = \
            self.update_pseudo_potential(density)

        self.timer.start('Atomic')
        self.dH_asp = None  # XXXX

        dH_asp = {}
        for a, D_sp in density.D_asp.items():
            W_L = W_aL[a]
            setup = self.setups[a]

            D_p = D_sp[:self.nspins].sum(0)
            dH_p = (setup.K_p + setup.M_p + setup.MB_p +
                    2.0 * np.dot(setup.M_pp, D_p) +
                    np.dot(setup.Delta_pL, W_L))
            Ekin += np.dot(setup.K_p, D_p) + setup.Kc
            Ebar += setup.MB + np.dot(setup.MB_p, D_p)
            Epot += setup.M + np.dot(D_p,
                                     (setup.M_p + np.dot(setup.M_pp, D_p)))

            if self.vext is not None:
                vext = self.vext.get_taylor(spos_c=self.spos_ac[a, :])
                # Tailor expansion to the zeroth order
                Eext += vext[0][0] * (sqrt(4 * pi) * density.Q_aL[a][0] +
                                      setup.Z)
                dH_p += vext[0][0] * sqrt(4 * pi) * setup.Delta_pL[:, 0]
                if len(vext) > 1:
                    # Tailor expansion to the first order
                    Eext += sqrt(4 * pi / 3) * np.dot(vext[1],
                                                      density.Q_aL[a][1:4])
                    # there must be a better way XXXX
                    Delta_p1 = np.array([
                        setup.Delta_pL[:, 1], setup.Delta_pL[:, 2],
                        setup.Delta_pL[:, 3]
                    ])
                    dH_p += sqrt(4 * pi / 3) * np.dot(vext[1], Delta_p1)

            dH_asp[a] = dH_sp = np.zeros_like(D_sp)

            if setup.HubU is not None:
                assert self.collinear
                nspins = len(D_sp)

                l_j = setup.l_j
                l = setup.Hubl
                scale = setup.Hubs
                nl = np.where(np.equal(l_j, l))[0]
                nn = (2 * np.array(l_j) + 1)[0:nl[0]].sum()

                for D_p, H_p in zip(D_sp, dH_asp[a]):
                    [N_mm, V] = self.aoom(unpack2(D_p), a, l, scale)
                    N_mm = N_mm / 2 * nspins

                    Eorb = setup.HubU / 2. * (N_mm -
                                              np.dot(N_mm, N_mm)).trace()
                    Vorb = setup.HubU * (0.5 * np.eye(2 * l + 1) - N_mm)
                    Exc += Eorb
                    if nspins == 1:
                        # add contribution of other spin manyfold
                        Exc += Eorb

                    if len(nl) == 2:
                        mm = (2 * np.array(l_j) + 1)[0:nl[1]].sum()

                        V[nn:nn + 2 * l + 1, nn:nn + 2 * l + 1] *= Vorb
                        V[mm:mm + 2 * l + 1, nn:nn + 2 * l + 1] *= Vorb
                        V[nn:nn + 2 * l + 1, mm:mm + 2 * l + 1] *= Vorb
                        V[mm:mm + 2 * l + 1, mm:mm + 2 * l + 1] *= Vorb
                    else:
                        V[nn:nn + 2 * l + 1, nn:nn + 2 * l + 1] *= Vorb

                    Htemp = unpack(H_p)
                    Htemp += V
                    H_p[:] = pack2(Htemp)

            dH_sp[:self.nspins] += dH_p
            if self.ref_dH_asp:
                dH_sp += self.ref_dH_asp[a]
            # We are not yet done with dH_sp; still need XC correction below

        Ddist_asp = self.dh_distributor.distribute(density.D_asp)

        dHdist_asp = {}
        Exca = 0.0
        self.timer.start('XC Correction')
        for a, D_sp in Ddist_asp.items():
            setup = self.setups[a]
            dH_sp = np.zeros_like(D_sp)
            Exca += self.xc.calculate_paw_correction(setup, D_sp, dH_sp, a=a)
            # XXX Exc are added on the "wrong" distribution; sum only works
            # when gd.comm and distribution comm are the same
            dHdist_asp[a] = dH_sp
        self.timer.stop('XC Correction')

        dHdist_asp = self.dh_distributor.collect(dHdist_asp)

        # Exca has contributions from all cores so modify it so it is
        # parallel in the same way as the other energies.
        Exca = self.world.sum(Exca)
        if self.gd.comm.rank == 0:
            Exc += Exca

        assert len(dHdist_asp) == len(self.atom_partition.my_indices)

        for a, D_sp in density.D_asp.items():
            dH_sp = dH_asp[a]
            dH_sp += dHdist_asp[a]
            Ekin -= (D_sp * dH_sp).sum()  # NCXXX
        self.dH_asp = dH_asp
        self.timer.stop('Atomic')

        # Make corrections due to non-local xc:
        #xcfunc = self.xc.xcfunc
        self.Enlxc = 0.0  # XXXxcfunc.get_non_local_energy()
        Ekin += self.xc.get_kinetic_energy_correction() / self.gd.comm.size

        energies = np.array([Ekin, Epot, Ebar, Eext, Exc])
        self.timer.start('Communicate energies')
        self.gd.comm.sum(energies)
        # Make sure that all CPUs have the same energies
        self.world.broadcast(energies, 0)
        self.timer.stop('Communicate energies')
        (self.Ekin0, self.Epot, self.Ebar, self.Eext, self.Exc) = energies

        #self.Exc += self.Enlxc
        #self.Ekin0 += self.Enlkin

        self.timer.stop('Hamiltonian')

    def get_energy(self, occupations):
        self.Ekin = self.Ekin0 + occupations.e_band
        self.S = occupations.e_entropy

        # Total free energy:
        self.Etot = (self.Ekin + self.Epot + self.Eext + self.Ebar + self.Exc -
                     self.S)

        return self.Etot

    def linearize_to_xc(self, new_xc, density):
        # Store old hamiltonian
        ref_vt_sG = self.vt_sG.copy()
        ref_dH_asp = {}
        for a, dH_sp in self.dH_asp.items():
            ref_dH_asp[a] = dH_sp.copy()
        self.xc = new_xc
        self.xc.set_positions(self.spos_ac)
        self.update(density)

        ref_vt_sG -= self.vt_sG
        for a, dH_sp in self.dH_asp.items():
            ref_dH_asp[a] -= dH_sp
        self.ref_vt_sG = ref_vt_sG
        self.ref_dH_asp = ref_dH_asp

    def calculate_forces(self, dens, F_av):
        ghat_aLv = dens.ghat.dict(derivative=True)
        nct_av = dens.nct.dict(derivative=True)
        vbar_av = self.vbar.dict(derivative=True)

        self.calculate_forces2(dens, ghat_aLv, nct_av, vbar_av)

        # Force from compensation charges:
        for a, dF_Lv in ghat_aLv.items():
            F_av[a] += np.dot(dens.Q_aL[a], dF_Lv)

        # Force from smooth core charge:
        for a, dF_v in nct_av.items():
            F_av[a] += dF_v[0]

        # Force from zero potential:
        for a, dF_v in vbar_av.items():
            F_av[a] += dF_v[0]

        self.xc.add_forces(F_av)
        self.gd.comm.sum(F_av, 0)

    def apply_local_potential(self, psit_nG, Htpsit_nG, s):
        """Apply the Hamiltonian operator to a set of vectors.

        XXX Parameter description is deprecated!
        
        Parameters:

        a_nG: ndarray
            Set of vectors to which the overlap operator is applied.
        b_nG: ndarray, output
            Resulting H times a_nG vectors.
        kpt: KPoint object
            k-point object defined in kpoint.py.
        calculate_projections: bool
            When True, the integrals of projector times vectors
            P_ni = <p_i | a_nG> are calculated.
            When False, existing P_uni are used
        local_part_only: bool
            When True, the non-local atomic parts of the Hamiltonian
            are not applied and calculate_projections is ignored.
        
        """
        vt_G = self.vt_sG[s]
        if psit_nG.ndim == 3:
            Htpsit_nG += psit_nG * vt_G
        else:
            for psit_G, Htpsit_G in zip(psit_nG, Htpsit_nG):
                Htpsit_G += psit_G * vt_G

    def apply(self, a_xG, b_xG, wfs, kpt, calculate_P_ani=True):
        """Apply the Hamiltonian operator to a set of vectors.

        Parameters:

        a_nG: ndarray
            Set of vectors to which the overlap operator is applied.
        b_nG: ndarray, output
            Resulting S times a_nG vectors.
        wfs: WaveFunctions
            Wave-function object defined in wavefunctions.py
        kpt: KPoint object
            k-point object defined in kpoint.py.
        calculate_P_ani: bool
            When True, the integrals of projector times vectors
            P_ni = <p_i | a_nG> are calculated.
            When False, existing P_ani are used
        
        """

        wfs.kin.apply(a_xG, b_xG, kpt.phase_cd)
        self.apply_local_potential(a_xG, b_xG, kpt.s)
        shape = a_xG.shape[:-3]
        P_axi = wfs.pt.dict(shape)

        if calculate_P_ani:  # TODO calculate_P_ani=False is experimental
            wfs.pt.integrate(a_xG, P_axi, kpt.q)
        else:
            for a, P_ni in kpt.P_ani.items():
                P_axi[a][:] = P_ni

        for a, P_xi in P_axi.items():
            dH_ii = unpack(self.dH_asp[a][kpt.s])
            P_axi[a] = np.dot(P_xi, dH_ii)
        wfs.pt.add(b_xG, P_axi, kpt.q)

    def get_xc_difference(self, xc, density):
        """Calculate non-selfconsistent XC-energy difference."""
        if density.nt_sg is None:
            density.interpolate_pseudo_density()
        nt_sg = density.nt_sg
        if hasattr(xc, 'hybrid'):
            xc.calculate_exx()
        Exc = xc.calculate(density.finegd, nt_sg) / self.gd.comm.size
        for a, D_sp in density.D_asp.items():
            setup = self.setups[a]
            Exc += xc.calculate_paw_correction(setup, D_sp)
        Exc = self.gd.comm.sum(Exc)
        return Exc - self.Exc

    def estimate_memory(self, mem):
        nbytes = self.gd.bytecount()
        nfinebytes = self.finegd.bytecount()
        arrays = mem.subnode('Arrays', 0)
        arrays.subnode('vHt_g', nfinebytes)
        arrays.subnode('vt_sG', self.nspins * nbytes)
        arrays.subnode('vt_sg', self.nspins * nfinebytes)
        self.xc.estimate_memory(mem.subnode('XC'))
        self.poisson.estimate_memory(mem.subnode('Poisson'))
        self.vbar.estimate_memory(mem.subnode('vbar'))

    def read(self, reader, parallel):
        self.Ekin = reader['Ekin']
        self.Epot = reader['Epot']
        self.Ebar = reader['Ebar']
        try:
            self.Eext = reader['Eext']
        except (AttributeError, KeyError):
            self.Eext = 0.0
        self.Exc = reader['Exc']
        self.S = reader['S']
        self.Etot = reader.get('PotentialEnergy',
                               broadcast=True) - 0.5 * self.S

        if not reader.has_array('PseudoPotential'):
            return

        hdf5 = hasattr(reader, 'hdf5')
        version = reader['version']

        # Read pseudo potential on the coarse grid
        # and broadcast on kpt/band comm:
        if version > 0.3:
            self.vt_sG = self.gd.empty(self.nspins)
            if hdf5:
                indices = [
                    slice(0, self.nspins),
                ] + self.gd.get_slice()
                do_read = (self.kptband_comm.rank == 0)
                reader.get('PseudoPotential',
                           out=self.vt_sG,
                           parallel=parallel,
                           read=do_read,
                           *indices)  # XXX read=?
                self.kptband_comm.broadcast(self.vt_sG, 0)
            else:
                for s in range(self.nspins):
                    self.gd.distribute(reader.get('PseudoPotential', s),
                                       self.vt_sG[s])

        # Read non-local part of hamiltonian
        self.dH_asp = {}
        natoms = len(self.setups)
        self.rank_a = np.zeros(natoms, int)
        if version > 0.3:
            all_H_sp = reader.get('NonLocalPartOfHamiltonian', broadcast=True)

        if self.gd.comm.rank == 0 and version > 0.3:
            self.dH_asp = read_atomic_matrices(all_H_sp, self.setups)