Python grid_descriptor.GridDescriptor类代码示例

def go(comm, ngpts, repeat, narrays, out, prec):
    N_c = np.array((ngpts, ngpts, ngpts))
    a = 10.0
    gd = GridDescriptor(N_c, (a, a, a), comm=comm))
    gdcoarse = gd.coarsen()
    gdfine = gd.refine()
    kin1 = Laplace(gd, -0.5, 1).apply
    laplace = Laplace(gd, -0.5, 2)
    kin2 = laplace.apply
    restrict = Transformer(gd, gdcoarse, 1).apply
    interpolate = Transformer(gd, gdfine, 1).apply
    precondition = Preconditioner(gd, laplace, np_float)
    a1 = gd.empty(narrays)
    a1[:] = 1.0
    a2 = gd.empty(narrays)
    c = gdcoarse.empty(narrays)
    f = gdfine.empty(narrays)

    T = [0, 0, 0, 0, 0]
    for i in range(repeat):
        comm.barrier()
        kin1(a1, a2)
        comm.barrier()
        t0a = time()
        kin1(a1, a2)
        t0b = time()
        comm.barrier()
        t1a = time()
        kin2(a1, a2)
        t1b = time()
        comm.barrier()
        t2a = time()
        for A, C in zip(a1,c):
            restrict(A, C)
        t2b = time()
        comm.barrier()
        t3a = time()
        for A, F in zip(a1,f):
            interpolate(A, F)
        t3b = time()
        comm.barrier()
        if prec:
            t4a = time()
            for A in a1:
                precondition(A, None, None, None)
            t4b = time()
            comm.barrier()

        T[0] += t0b - t0a
        T[1] += t1b - t1a
        T[2] += t2b - t2a
        T[3] += t3b - t3a
        if prec:
            T[4] += t4b - t4a

    if mpi.rank == 0:
        out.write(' %2d %2d %2d' % tuple(gd.parsize_c))
        out.write(' %12.6f %12.6f %12.6f %12.6f %12.6f\n' %
                  tuple([t / repeat / narrays for t in T]))
        out.flush()

class WignerSeitzTruncatedCoulomb:
    def __init__(self, cell_cv, nk_c, txt=sys.stdout):
        self.nk_c = nk_c
        bigcell_cv = cell_cv * nk_c[:, np.newaxis]
        L_c = (np.linalg.inv(bigcell_cv)**2).sum(0)**-0.5
        
        rc = 0.5 * L_c.min()
        prnt('Inner radius for %dx%dx%d Wigner-Seitz cell: %.3f Ang' %
             (tuple(nk_c) + (rc * Bohr,)), file=txt)
        
        self.a = 5 / rc
        prnt('Range-separation parameter: %.3f Ang^-1' % (self.a / Bohr),
             file=txt)
        
#        nr_c = [get_efficient_fft_size(2 * int(L * self.a * 1.5))
        nr_c = [get_efficient_fft_size(2 * int(L * self.a * 3.0))
                for L in L_c]
        prnt('FFT size for calculating truncated Coulomb: %dx%dx%d' %
             tuple(nr_c), file=txt)
        
        self.gd = GridDescriptor(nr_c, bigcell_cv, comm=mpi.serial_comm)
        v_R = self.gd.empty()
        v_i = v_R.ravel()
        
        pos_iv = self.gd.get_grid_point_coordinates().reshape((3, -1)).T
        corner_jv = np.dot(np.indices((2, 2, 2)).reshape((3, 8)).T, bigcell_cv)
        for i, pos_v in enumerate(pos_iv):
            r = ((pos_v - corner_jv)**2).sum(axis=1).min()**0.5
            if r == 0:
                v_i[i] = 2 * self.a / pi**0.5
            else:
                v_i[i] = erf(self.a * r) / r
                
        self.K_Q = np.fft.fftn(v_R) * self.gd.dv
        
    def get_potential(self, pd):
        q_c = pd.kd.bzk_kc[0]
        shift_c = (q_c * self.nk_c).round().astype(int)
        max_c = self.gd.N_c // 2
        K_G = pd.zeros()
        N_c = pd.gd.N_c
        for G, Q in enumerate(pd.Q_qG[0]):
            Q_c = (np.unravel_index(Q, N_c) + N_c // 2) % N_c - N_c // 2
            Q_c = Q_c * self.nk_c + shift_c
            if (abs(Q_c) < max_c).all():
                K_G[G] = self.K_Q[tuple(Q_c)]

        G2_G = pd.G2_qG[0]
        a = self.a
        if pd.kd.gamma:
            K_G[0] += pi / a**2
        else:
            K_G[0] += 4 * pi * (1 - np.exp(-G2_G[0] / (4 * a**2))) / G2_G[0]
        K_G[1:] += 4 * pi * (1 - np.exp(-G2_G[1:] / (4 * a**2))) / G2_G[1:]
        assert pd.dtype == complex
        return K_G

 def __init__(self, gd, spline_j, spos_c, index=None):
     rcut = max([spline.get_cutoff() for spline in spline_j])
     corner_c = np.ceil(spos_c * gd.N_c - rcut / gd.h_c).astype(int)
     size_c = np.ceil(spos_c * gd.N_c + rcut / gd.h_c).astype(int) - corner_c
     smallgd = GridDescriptor(N_c=size_c + 1, cell_cv=gd.h_c * (size_c + 1), pbc_c=False, comm=mpi.serial_comm)
     lfc = LFC(smallgd, [spline_j])
     lfc.set_positions((spos_c[np.newaxis, :] * gd.N_c - corner_c + 1) / smallgd.N_c)
     ni = lfc.Mmax
     f_iG = smallgd.zeros(ni)
     lfc.add(f_iG, {0: np.eye(ni)})
     LocalizedFunctions.__init__(self, gd, f_iG, corner_c, index=index)

def extend_grid(gd, N_cd):
    N_cd = np.array(N_cd)
    
    N_c = gd.N_c + N_cd.sum(axis=1)
    cell_cv = gd.h_cv * N_c
    
    move_c = gd.get_grid_spacings() * N_cd[:,0]
    
    egd = GridDescriptor(N_c, cell_cv, gd.pbc_c, gd.comm)
    egd.extend_N_cd = N_cd
    
    return egd, cell_cv*Bohr, move_c*Bohr

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]

def make_dummy_kpt_reference(l, function, k_c,
                             rcut=6., a=10., n=60, dtype=complex):
    r = np.linspace(0., rcut, 300)
    mcount = 2*l + 1
    fcount = 1
    kcount = 1
    gd = GridDescriptor((n, n, n), (a, a, a), (True, True, True))
    kpt = KPoint([], gd, 1., 0, 0, 0, k_c, dtype)
    spline = Spline(l, r[-1], function(r))
    center = (.5, .5, .5)
    lf = create_localized_functions([spline], gd, center, dtype=dtype)
    lf.set_phase_factors([kpt.k_c])
    psit_nG = gd.zeros(mcount, dtype=dtype)
    coef_xi = np.identity(mcount * fcount, dtype=dtype)
    lf.add(psit_nG, coef_xi, k=0)
    kpt.psit_nG = psit_nG
    print 'Number of boxes', len(lf.box_b)
    print 'Phase kb factors shape', lf.phase_kb.shape
    return gd, kpt, center

    def setUp(self):
        for virtvar in ['boundaries', 'celltype']:
            assert getattr(self,virtvar) is not None, 'Virtual "%s"!' % virtvar

        # Basic unit cell information:
        pbc_c = {'zero'    : (False,False,False), \
                 'periodic': (True,True,True), \
                 'mixed'   : (True, False, True)}[self.boundaries]
        a, b = self.a, 2**0.5*self.a
        cell_cv = {'general'   : np.array([[0,a,a],[a/2,0,a/2],[a/2,a/2,0]]),
                   'rotated'   : np.array([[0,0,b],[b/2,0,0],[0,b/2,0]]),
                   'inverted'   : np.array([[0,0,b],[0,b/2,0],[b/2,0,0]]),
                   'orthogonal': np.diag([b, b/2, b/2])}[self.celltype]
        cell_cv = np.array([(4-3*pbc)*c_v for pbc,c_v in zip(pbc_c, cell_cv)])

        # Decide how many kpoints to sample from the 1st Brillouin Zone
        kpts_c = np.ceil((10/Bohr)/np.sum(cell_cv**2,axis=1)**0.5).astype(int)
        kpts_c = tuple(kpts_c*pbc_c + 1-pbc_c)
        bzk_kc = kpts2ndarray(kpts_c)
        self.gamma = len(bzk_kc) == 1 and not bzk_kc[0].any()

        #p = InputParameters()
        #Z_a = self.atoms.get_atomic_numbers()
        #xcfunc = XC(p.xc)
        #setups = Setups(Z_a, p.setups, p.basis, p.lmax, xcfunc)
        #symmetry, weight_k, self.ibzk_kc = reduce_kpoints(self.atoms, bzk_kc,
        #                                                  setups, p.usesymm)

        self.ibzk_kc = bzk_kc.copy() # don't use symmetry reduction of kpoints
        self.nibzkpts = len(self.ibzk_kc)
        self.ibzk_kv = kpoint_convert(cell_cv, skpts_kc=self.ibzk_kc)

        # Parse parallelization parameters and create suitable communicators.
        #parsize_domain, parsize_bands = create_parsize_minbands(self.nbands, world.size)
        parsize_domain, parsize_bands = world.size//gcd(world.size, self.nibzkpts), 1
        assert self.nbands % np.prod(parsize_bands) == 0
        domain_comm, kpt_comm, band_comm = distribute_cpus(parsize_domain,
            parsize_bands, self.nspins, self.nibzkpts)

        # Set up band descriptor:
        self.bd = BandDescriptor(self.nbands, band_comm)

        # Set up grid descriptor:
        N_c = np.round(np.sum(cell_cv**2, axis=1)**0.5 / self.h)
        N_c += 4-N_c % 4 # makes domain decomposition easier
        self.gd = GridDescriptor(N_c, cell_cv, pbc_c, domain_comm, parsize_domain)
        self.assertEqual(self.gamma, np.all(~self.gd.pbc_c))

        # What to do about kpoints?
        self.kpt_comm = kpt_comm

        if debug and world.rank == 0:
            comm_sizes = tuple([comm.size for comm in [world, self.bd.comm, \
                                                   self.gd.comm, self.kpt_comm]])
            print '%d world, %d band, %d domain, %d kpt' % comm_sizes

    def get_combined_data(self, qmdata=None, cldata=None, spacing=None):
        
        if qmdata is None:
            qmdata = self.density.rhot_g
        
        if cldata is None:
            cldata = self.classical_material.charge_density
        
        if spacing is None:
            spacing = self.cl.gd.h_cv[0, 0]
        
        spacing_au = spacing / Bohr  # from Angstroms to a.u.
        
        # Collect data from different processes
        cln = self.cl.gd.collect(cldata)
        qmn = self.qm.gd.collect(qmdata)

        clgd = GridDescriptor(self.cl.gd.N_c,
                              self.cl.cell,
                              False,
                              serial_comm,
                              None)

        if world.rank == 0:
            cln *= self.classical_material.sign
            # refine classical part
            while clgd.h_cv[0, 0] > spacing_au * 1.50:  # 45:
                cln = Transformer(clgd, clgd.refine()).apply(cln)
                clgd = clgd.refine()
                
            # refine quantum part
            qmgd = GridDescriptor(self.qm.gd.N_c,
                                  self.qm.cell,
                                  False,
                                  serial_comm,
                                  None)                           
            while qmgd.h_cv[0, 0] < clgd.h_cv[0, 0] * 0.95:
                qmn = Transformer(qmgd, qmgd.coarsen()).apply(qmn)
                qmgd = qmgd.coarsen()
            
            assert np.all(qmgd.h_cv == clgd.h_cv), " Spacings %.8f (qm) and %.8f (cl) Angstroms" % (qmgd.h_cv[0][0] * Bohr, clgd.h_cv[0][0] * Bohr)
            
            # now find the corners
            r_gv_cl = clgd.get_grid_point_coordinates().transpose((1, 2, 3, 0))
            cind = self.qm.corner1 / np.diag(clgd.h_cv) - 1
            
            n = qmn.shape

            # print 'Corner points:     ', self.qm.corner1*Bohr,      ' - ', self.qm.corner2*Bohr
            # print 'Calculated points: ', r_gv_cl[tuple(cind)]*Bohr, ' - ', r_gv_cl[tuple(cind+n+1)]*Bohr
                        
            cln[cind[0] + 1:cind[0] + n[0] + 1,
                cind[1] + 1:cind[1] + n[1] + 1,
                cind[2] + 1:cind[2] + n[2] + 1] += qmn
        
        world.barrier()
        return cln, clgd

def test():
    from gpaw.grid_descriptor import GridDescriptor

    ngpts = 40
    h = 1.0 / ngpts
    N_c = (ngpts, ngpts, ngpts)
    a = h * ngpts
    gd = GridDescriptor(N_c, (a, a, a))
    
    from gpaw.spline import Spline
    a = np.array([1, 0.9, 0.8, 0.0])
    s = Spline(0, 0.2, a)
    x = LocalizedFunctionsCollection(gd, [[s], [s]])
    x.set_positions([(0.5, 0.45, 0.5), (0.5, 0.55, 0.5)])
    n_G = gd.zeros()
    x.add(n_G)
    import pylab as plt
    plt.contourf(n_G[20, :, :])
    plt.axis('equal')
    plt.show()

def make_dummy_reference(l, function=None, rcut=6., a=12., n=60,
                         dtype=float):
    """Make a mock reference wave function using a made-up radial function
    as reference"""
    #print 'Dummy reference: l=%d, rcut=%.02f, alpha=%.02f' % (l, rcut, alpha)
    r = np.arange(0., rcut, .01)

    if function is None:
        function = QuasiGaussian(4., rcut)

    norm = get_norm(r, function(r), l)
    function.renormalize(norm)
    #g = QuasiGaussian(alpha, rcut)
    
    mcount = 2*l + 1
    fcount = 1
    gd = GridDescriptor((n, n, n), (a, a, a), (False, False, False))
    spline = Spline(l, r[-1], function(r), points=50)
    center = (.5, .5, .5)
    lf = create_localized_functions([spline], gd, center, dtype=dtype)
    psit_k = gd.zeros(mcount, dtype=dtype)
    coef_xi = np.identity(mcount * fcount, dtype=dtype)
    lf.add(psit_k, coef_xi)
    return gd, psit_k, center, function

def f(n, p):
    N = 2 * n
    gd = GridDescriptor((N, N, N), (L, L, L))
    a = gd.zeros()
    print(a.shape)
    #p = PoissonSolver(nn=1, relax=relax)
    p.set_grid_descriptor(gd)
    p.initialize()
    cut = N / 2.0 * 0.9
    s = Spline(l=0, rmax=cut, f_g=np.array([1, 0.5, 0.0]))
    c = LFC(gd, [[s], [s]])
    c.set_positions([(0, 0, 0), (0.5, 0.5, 0.5)])
    c.add(a)

    I0 = gd.integrate(a)
    a -= gd.integrate(a) / L**3

    I = gd.integrate(a)
    b = gd.zeros()
    p.solve(b, a, charge=0)#, eps=1e-20)
    return gd.collect(b, broadcast=1)

from gpaw.test import equal
from gpaw.grid_descriptor import GridDescriptor
from gpaw.spline import Spline
import gpaw.mpi as mpi
from gpaw.lfc import LocalizedFunctionsCollection as LFC

s = Spline(0, 1.0, [1.0, 0.5, 0.0])
n = 40
a = 8.0
gd = GridDescriptor((n, n, n), (a, a, a), comm=mpi.serial_comm)
c = LFC(gd, [[s], [s], [s]])
c.set_positions([(0.5, 0.5, 0.25 + 0.25 * i) for i in [0, 1, 2]])
b = gd.zeros()
c.add(b)
x = gd.integrate(b)

gd = GridDescriptor((n, n, n), (a, a, a), comm=mpi.serial_comm)
c = LFC(gd, [[s], [s], [s]])
c.set_positions([(0.5, 0.5, 0.25 + 0.25 * i) for i in [0, 1, 2]])
b = gd.zeros()
c.add(b)
y = gd.integrate(b)
equal(x, y, 1e-13)

from gpaw.fd_operators import Gradient
import numpy as np
from gpaw.grid_descriptor import GridDescriptor
from gpaw.mpi import world

if world.size > 4:
    # Grid is so small that domain decomposition cannot exceed 4 domains
    assert world.size % 4 == 0
    group, other = divmod(world.rank, 4)
    ranks = np.arange(4*group, 4*(group+1))
    domain_comm = world.new_communicator(ranks)
else:
    domain_comm = world

gd = GridDescriptor((8, 1, 1), (8.0, 1.0, 1.0), comm=domain_comm)
a = gd.zeros()
dadx = gd.zeros()
a[:, 0, 0] = np.arange(gd.beg_c[0], gd.end_c[0])
gradx = Gradient(gd, v=0)
print a.itemsize, a.dtype, a.shape
print dadx.itemsize, dadx.dtype, dadx.shape
gradx.apply(a, dadx)

#   a = [ 0.  1.  2.  3.  4.  5.  6.  7.]
#
#   da
#   -- = [-2.5  1.   1.   1.   1.   1.  1.  -2.5]
#   dx

dadx = gd.collect(dadx, broadcast=True)
assert dadx[3, 0, 0] == 1.0 and np.sum(dadx[:, 0, 0]) == 0.0

import numpy as np
from gpaw.lfc import LocalizedFunctionsCollection as LFC
from gpaw.grid_descriptor import GridDescriptor
from gpaw.spline import Spline
a = 4.0
gd = GridDescriptor(N_c=[16, 20, 20], cell_cv=[a, a + 1, a + 2],
                    pbc_c=(0, 1, 1))
spos_ac = np.array([[0.25, 0.15, 0.35], [0.5, 0.5, 0.5]])
kpts_kc = None
s = Spline(l=0, rmax=2.0, f_g=np.array([1, 0.9, 0.1, 0.0]))
p = Spline(l=1, rmax=2.0, f_g=np.array([1, 0.9, 0.1, 0.0]))
spline_aj = [[s], [s, p]]
c = LFC(gd, spline_aj, cut=True, forces=True)
c.set_positions(spos_ac)
C_ani = c.dict(3, zero=True)
if 1 in C_ani:
    C_ani[1][:, 1:] = np.eye(3)
psi = gd.zeros(3)
c.add(psi, C_ani)
c.integrate(psi, C_ani)
if 1 in C_ani:
    d = C_ani[1][:, 1:].diagonal()
    assert d.ptp() < 4e-6
    C_ani[1][:, 1:] -= np.diag(d)
    assert abs(C_ani[1]).max() < 5e-17
d_aniv = c.dict(3, derivative=True)
c.derivative(psi, d_aniv)
if 1 in d_aniv:
    for v in range(3):
        assert abs(d_aniv[1][v - 1, 0, v] + 0.2144) < 5e-5
        d_aniv[1][v - 1, 0, v] = 0

from time import time
from gpaw.grid_descriptor import GridDescriptor
from gpaw.fd_operators import Laplace
import gpaw.mpi as mpi

n = 96
h = 0.1
L = n * h
gd = GridDescriptor((n, n, n), [L, L, L])

# Allocate arrays:
a = gd.zeros(100) + 1.2
b = gd.empty(100)

o = Laplace(gd, 2).apply

t0 = time()
for r in range(10):
    o(a, b)

if mpi.rank == 0:
    print time() - t0, a.shape