class KineticEnergyPreconditioner:
def __init__(self, gd, kin, dtype):
self.preconditioner = Preconditioner(gd, kin, dtype)
self.preconditioner.allocate()
def apply(self, kpt, psi, psin):
for i in range(len(psi)):
psin[i][:] = self.preconditioner(psi[i], kpt.phase_cd, None, None)
class KineticEnergyPreconditioner:
def __init__(self, gd, kin, dtype):
self.preconditioner = Preconditioner(gd, kin, dtype)
self.preconditioner.allocate()
def apply(self, kpt, psi, psin):
for i in range(len(psi)):
psin[i][:] = self.preconditioner(psi[i], kpt.phase_cd, None, None)
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()
def __init__(self, gd, project, dtype=float):
"""Init the gpaw preconditioner.
Parameters
----------
gd: GridDescriptor
Coarse grid
"""
self.project = project
self.gd = gd
# K-point for the preconditioner
self.kpt = None
kin = Laplace(gd, scale=-0.5, n=3, dtype=dtype, allocate=True)
self.pc = Preconditioner(gd, kin, dtype=dtype)
self.pc.allocate()
# For scipy's linear solver
N = np.prod(gd.n_c)
self.shape = (N,N)
self.dtype = dtype
def __init__(self, gd, kin, dtype):
self.preconditioner = Preconditioner(gd, kin, dtype)
self.preconditioner.allocate()
def __init__(self, gd, kin, dtype):
self.preconditioner = Preconditioner(gd, kin, dtype)
self.preconditioner.allocate()
def make_preconditioner(self, block=1):
return Preconditioner(self.gd, self.kin, self.dtype, block)
class ScipyPreconditioner:
def __init__(self, gd, project, dtype=float):
"""Init the gpaw preconditioner.
Parameters
----------
gd: GridDescriptor
Coarse grid
"""
self.project = project
self.gd = gd
# K-point for the preconditioner
self.kpt = None
kin = Laplace(gd, scale=-0.5, n=3, dtype=dtype, allocate=True)
self.pc = Preconditioner(gd, kin, dtype=dtype)
self.pc.allocate()
# For scipy's linear solver
N = np.prod(gd.n_c)
self.shape = (N,N)
self.dtype = dtype
def set_kpt(self, kpt):
"""Set k-point for ``Transformer`` objects inside the preconditioner.
Parameters
----------
kpt: KPoint or KPointContainer
Only requirement is that it must have an ``phase_cd`` attribute.
"""
self.kpt = kpt
def matvec(self, x):
"""Matrix vector multiplication for ``scipy.sparse.linalg`` solvers.
Parameters
----------
x: ndarry
1-dimensional array holding the grid representation of the vector
"""
# Output array
y_G = self.gd.zeros(dtype=self.dtype)
shape = y_G.shape
size = x.size
assert size == np.prod(shape)
x_G = x.reshape(shape)
# Call gpaw preconditioner
y_G = self.pc(x_G, kpt=self.kpt)
# Project out undesired (numerical) components
self.project(y_G)
y = y_G.ravel()
return y
