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 set_grid_descriptor(self, gd):
# Should probably be renamed initialize
self.gd = gd
scale = -0.25 / pi
if self.nn == 'M':
if not gd.orthogonal:
raise RuntimeError('Cannot use Mehrstellen stencil with '
'non orthogonal cell.')
self.operators = [LaplaceA(gd, -scale)]
self.B = LaplaceB(gd)
else:
self.operators = [Laplace(gd, scale, self.nn)]
self.B = None
self.interpolators = []
self.restrictors = []
level = 0
self.presmooths = [2]
self.postsmooths = [1]
# Weights for the relaxation,
# only used if 'J' (Jacobi) is chosen as method
self.weights = [2.0 / 3.0]
while level < 8:
try:
gd2 = gd.coarsen()
except ValueError:
break
self.operators.append(Laplace(gd2, scale, 1))
self.interpolators.append(Transformer(gd2, gd))
self.restrictors.append(Transformer(gd, gd2))
self.presmooths.append(4)
self.postsmooths.append(4)
self.weights.append(1.0)
level += 1
gd = gd2
self.levels = level
if self.operators[-1].gd.N_c.max() > 36:
# Try to warn exactly once no matter how one uses the solver.
if gd.comm.parent is None:
warn = (gd.comm.rank == 0)
else:
warn = (gd.comm.parent.rank == 0)
if warn:
warntxt = '\n'.join([POISSON_GRID_WARNING, '',
self.get_description()])
else:
warntxt = ('Poisson warning from domain rank %d'
% self.gd.comm.rank)
# Warn from all ranks to avoid deadlocks.
warnings.warn(warntxt, stacklevel=2)
def __init__(self, hamiltonian, wfs, gd, dtype=float):
"""Save useful objects for the Sternheimer operator.
Parameters
----------
hamiltonian: Hamiltonian
Hamiltonian for a ground-state calculation.
wfs: WaveFunctions
Ground-state wave-functions.
gd: GridDescriptor
Grid on which the operator is defined.
dtype: dtype
dtype of the wave-function being operated on.
"""
self.hamiltonian = hamiltonian
self.kin = Laplace(gd, scale=-0.5, n=3, dtype=dtype)
self.kpt_u = wfs.kpt_u
self.pt = wfs.pt
self.gd = gd
# Variables for k-point and band index
self.k = None
self.n = None
self.kplusq = None
# For scipy's linear solver
N = np.prod(gd.n_c)
self.shape = (N, N)
self.dtype = dtype
def __init__(self,
stencil,
parallel,
initksl,
gd,
nvalence,
setups,
bd,
dtype,
world,
kd,
kptband_comm,
timer,
reuse_wfs_method=None,
collinear=True):
FDPWWaveFunctions.__init__(self,
parallel,
initksl,
reuse_wfs_method=reuse_wfs_method,
collinear=collinear,
gd=gd,
nvalence=nvalence,
setups=setups,
bd=bd,
dtype=dtype,
world=world,
kd=kd,
kptband_comm=kptband_comm,
timer=timer)
# Kinetic energy operator:
self.kin = Laplace(self.gd, -0.5, stencil, self.dtype)
self.taugrad_v = None # initialized by MGGA functional
def test_something(self):
laplace_uG = np.empty_like(self.laplace0_uG)
op = Laplace(self.gd, dtype=self.dtype)
for myu, laplace_G in enumerate(laplace_uG):
phase_cd = {float:None, complex:self.phase_ucd[myu]}[self.dtype]
op.apply(self.wf_uG[myu], laplace_G, phase_cd)
print 'myu:', myu, 'diff:', np.std(laplace_G-self.laplace0_uG[myu]), '/', np.abs(laplace_G-self.laplace0_uG[myu]).max()
def set_grid_descriptor(self, gd):
# Should probably be renamed initialize
self.gd = gd
self.dv = gd.dv
gd = self.gd
scale = -0.25 / pi
if self.nn == 'M':
if not gd.orthogonal:
raise RuntimeError('Cannot use Mehrstellen stencil with '
'non orthogonal cell.')
self.operators = [LaplaceA(gd, -scale, allocate=False)]
self.B = LaplaceB(gd, allocate=False)
else:
self.operators = [Laplace(gd, scale, self.nn, allocate=False)]
self.B = None
self.interpolators = []
self.restrictors = []
level = 0
self.presmooths = [2]
self.postsmooths = [1]
# Weights for the relaxation,
# only used if 'J' (Jacobi) is chosen as method
self.weights = [2.0 / 3.0]
while level < 4:
try:
gd2 = gd.coarsen()
except ValueError:
break
self.operators.append(Laplace(gd2, scale, 1, allocate=False))
self.interpolators.append(Transformer(gd2, gd, allocate=False))
self.restrictors.append(Transformer(gd, gd2, allocate=False))
self.presmooths.append(4)
self.postsmooths.append(4)
self.weights.append(1.0)
level += 1
gd = gd2
self.levels = level
def __init__(self, gd0, kin0, dtype=float, block=1):
gd1 = gd0.coarsen()
gd2 = gd1.coarsen()
self.kin0 = kin0
self.kin1 = Laplace(gd1, -0.5, 1, dtype)
self.kin2 = Laplace(gd2, -0.5, 1, dtype)
self.scratch0 = gd0.zeros((2, block), dtype, False)
self.scratch1 = gd1.zeros((3, block), dtype, False)
self.scratch2 = gd2.zeros((3, block), dtype, False)
self.step = 0.66666666 / kin0.get_diagonal_element()
self.restrictor_object0 = Transformer(gd0, gd1, 1, dtype)
self.restrictor_object1 = Transformer(gd1, gd2, 1, dtype)
self.interpolator_object2 = Transformer(gd2, gd1, 1, dtype)
self.interpolator_object1 = Transformer(gd1, gd0, 1, dtype)
self.restrictor0 = self.restrictor_object0.apply
self.restrictor1 = self.restrictor_object1.apply
self.interpolator2 = self.interpolator_object2.apply
self.interpolator1 = self.interpolator_object1.apply
def __init__(self, gd, bd, kd, setups, dtype): # override constructor
assert kd.comm.size == 1
WaveFunctions.__init__(self, gd, 1, setups, bd, dtype, world, kd, None)
self.kin = Laplace(gd, -0.5, dtype=dtype, allocate=False)
self.diagksl = None
self.orthoksl = BandLayouts(gd, bd, dtype)
self.initksl = None
self.overlap = None
self.rank_a = None
def __init__(self, stencil, diagksl, orthoksl, initksl, gd, nvalence,
setups, bd, dtype, world, kd, kptband_comm, timer):
FDPWWaveFunctions.__init__(self, diagksl, orthoksl, initksl, gd,
nvalence, setups, bd, dtype, world, kd,
kptband_comm, timer)
# Kinetic energy operator:
self.kin = Laplace(self.gd, -0.5, stencil, self.dtype)
self.matrixoperator = MatrixOperator(self.orthoksl)
self.taugrad_v = None # initialized by MGGA functional
def apply_t(self):
"""Apply kinetic energy operator and return new object."""
p = 2 # padding
newsize_c = self.size_c + 2 * p
gd = GridDescriptor(N_c=newsize_c + 1,
cell_cv=self.gd.h_c * (newsize_c + 1),
pbc_c=False,
comm=mpi.serial_comm)
T = Laplace(gd, scale =1/2., n=p)
f_ig = np.zeros((len(self.f_iG),) + tuple(newsize_c))
f_ig[:, p:-p, p:-p, p:-p] = self.f_iG
Tf_iG = np.empty_like(f_ig)
T.apply(f_ig, Tf_iG)
return LocalizedFunctions(self.gd, Tf_iG, self.corner_c - p,
self.index)
def set_absorbing_boundary(self, absorbing_boundary):
""" Sets up the absorbing boundary.
Parameters:
absorbing_boundary: absorbing boundary object of any kind.
"""
self.absorbing_boundary = absorbing_boundary
self.absorbing_boundary.set_up(self.hamiltonian.gd)
if self.absorbing_boundary.type == 'PML':
gd = self.hamiltonian.gd
self.laplace = Laplace(gd, n=2, dtype=complex)
self.gradient = np.array(
(Gradient(gd, 0, n=2,
dtype=complex), Gradient(gd, 1, n=2, dtype=complex),
Gradient(gd, 2, n=2, dtype=complex)))
self.lpsit = None
def initialize(self, load_gauss=False):
self.presmooths[self.levels] = 8
self.postsmooths[self.levels] = 8
self.phis = [None] + [gd.zeros() for gd in self.gds[1:]]
self.residuals = [gd.zeros() for gd in self.gds]
self.rhos = [gd.zeros() for gd in self.gds]
self.op_coarse_weights = [[g.empty() for g in (gd, ) * 4]
for gd in self.gds[1:]]
scale = -0.25 / np.pi
for i, gd in enumerate(self.gds):
if i == 0:
nn = self.nn
weights = self.dielectric.eps_gradeps
else:
nn = 1
weights = self.op_coarse_weights[i - 1]
operators = [Laplace(gd, scale, nn)] + \
[Gradient(gd, j, scale, nn) for j in (0, 1, 2)]
self.operators.append(WeightedFDOperator(weights, operators))
if load_gauss:
self.load_gauss()
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)
self.pc = Preconditioner(gd, kin, dtype=dtype)
# For scipy's linear solver
N = np.prod(gd.n_c)
self.shape = (N,N)
self.dtype = dtype
def __init__(self,
stencil,
diagksl,
orthoksl,
initksl,
gd,
nvalence,
setups,
bd,
dtype,
world,
kd,
timer=None):
FDPWWaveFunctions.__init__(self, diagksl, orthoksl, initksl, gd,
nvalence, setups, bd, dtype, world, kd,
timer)
self.wd = self.gd # wave function descriptor
# Kinetic energy operator:
self.kin = Laplace(self.gd, -0.5, stencil, self.dtype, allocate=False)
self.matrixoperator = MatrixOperator(orthoksl)
# Set up band and grid descriptors:
bd = BandDescriptor(N, band_comm, False)
gd = GridDescriptor((G, G, G), (a, a, a), True, domain_comm, parsize_c=D)
ksl = BandLayouts(gd, bd, block_comm, float)
# Random wave functions:
psit_mG = gd.empty(M)
for m in range(M):
np.random.seed(world.rank * M + m)
psit_mG[m] = np.random.uniform(-0.5, 0.5, tuple(gd.n_c))
if world.rank == 0:
print('Size of wave function array:', psit_mG.shape)
P_ani = {0: psit_mG[:, :2, 0, 0].copy(), 1: psit_mG[:, -1, -1, -3:].copy()}
kin = Laplace(gd, -0.5, 2).apply
vt_G = gd.empty()
vt_G.fill(0.567)
def run(psit_mG):
overlap = MatrixOperator(ksl, J)
def H(psit_xG):
Htpsit_xG = np.empty_like(psit_xG)
kin(psit_xG, Htpsit_xG)
for psit_G, y_G in zip(psit_xG, Htpsit_xG):
y_G += vt_G * psit_G
return Htpsit_xG
dH_aii = {0: np.ones((2, 2)) * 0.123, 1: np.ones((3, 3)) * 0.321}
""" Test the finite difference stencil """ from gpaw.grid_descriptor import GridDescriptor from gpaw.fd_operators import Laplace from gpaw.test import equal import numpy as np gpts = 64 nbands = 120 gd = GridDescriptor([gpts, gpts, gpts]) a = gd.empty(nbands) b = gd.empty(nbands) np.random.seed(10) a[:] = np.random.random(a.shape) op = Laplace(gd, 1.0, 3).apply op(a, b) equal(np.sum(b), -7.43198208966e-05, 1e-9)
def initialize(self, dens, ham, wfs, occ):
MGGA.initialize(self, dens, ham, wfs, occ)
self.kernel.world = wfs.world
self.kernel.gd = dens.finegd
self.kernel.lapl = Laplace(dens.finegd)
def create_laplace(self, gd, scale=1.0, n=1, dtype=float):
operators = [Laplace(gd, scale, n, dtype)]
operators += [Gradient(gd, j, scale, n, dtype) for j in (0, 1, 2)]
return WeightedFDOperator(operators)
def create_laplace(self, gd, scale=1.0, n=1, dtype=float):
"""Instantiate and return a Laplace operator
Allows subclasses to change the Laplace operator
"""
return Laplace(gd, scale, n, dtype)
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
def test(cellno, cellname, cell_cv, idiv, pbc, nn):
N_c = h2gpts(0.12, cell_cv, idiv=idiv)
if idiv == 1:
N_c += 1 - N_c % 2 # We want especially to test uneven grids
gd = GridDescriptor(N_c, cell_cv, pbc_c=pbc)
rho_g = gd.zeros()
phi_g = gd.zeros()
rho_g[:] = -0.3 + rng.rand(*rho_g.shape)
# Neutralize charge:
charge = gd.integrate(rho_g)
magic = gd.get_size_of_global_array().prod()
rho_g -= charge / gd.dv / magic
charge = gd.integrate(rho_g)
assert abs(charge) < 1e-12
# Check use_cholesky=True/False ?
from gpaw.poisson import FDPoissonSolver
ps = FastPoissonSolver(nn=nn)
#print('setgrid')
# Will raise BadAxesError for some pbc/cell combinations
ps.set_grid_descriptor(gd)
ps.solve(phi_g, rho_g)
laplace = Laplace(gd, scale=-1.0 / (4.0 * np.pi), n=nn)
def get_residual_err(phi_g):
rhotest_g = gd.zeros()
laplace.apply(phi_g, rhotest_g)
residual = np.abs(rhotest_g - rho_g)
# Residual is not accurate at end of non-periodic directions
# except for nn=1 (since effectively we use the right stencil
# only for nn=1 at the boundary).
#
# To do this check correctly, the Laplacian should have lower
# nn at the boundaries. Therefore we do not test the residual
# at these ends, only in between, by zeroing the bad ones:
if nn > 1:
exclude_points = nn - 1
for c in range(3):
if nn > 1 and not pbc[c]:
# get view ehere axis c refers becomes zeroth dimension:
X = residual.transpose(c, (c + 1) % 3, (c + 2) % 3)
if gd.beg_c[c] == 1:
X[:exclude_points] = 0.0
if gd.end_c[c] == gd.N_c[c]:
X[-exclude_points:] = 0.0
return residual.max()
maxerr = get_residual_err(phi_g)
pbcstring = '{}{}{}'.format(*pbc)
if 0:
ps2 = FDPoissonSolver(relax='J', nn=nn, eps=1e-18)
ps2.set_grid_descriptor(gd)
phi2_g = gd.zeros()
ps2.solve(phi2_g, rho_g)
phimaxerr = np.abs(phi2_g - phi_g).max()
maxerr2 = get_residual_err(phi2_g)
msg = ('{:2d} {:8s} pbc={} err={:8.5e} err[J]={:8.5e} '
'err[phi]={:8.5e} nn={:1d}'
.format(cellno, cellname, pbcstring, maxerr, maxerr2,
phimaxerr, nn))
state = 'ok' if maxerr < tolerance else 'FAIL'
msg = ('{:2d} {:8s} grid={} pbc={} err[fast]={:8.5e} nn={:1d} {}'
.format(cellno, cellname, N_c, pbcstring, maxerr, nn, state))
if world.rank == 0:
print(msg)
return maxerr
import numpy as np
from gpaw.grid_descriptor import GridDescriptor
from gpaw.fd_operators import Laplace
from gpaw.mpi import rank, size, world
G = 40
N = 2 * 3 * 2 * 5 * 7 * 8 * 3 * 11
B = 2
h = 0.2
a = h * R
M = N // B
assert M * B == N
D = size // B
assert D * B == size
r = rank // B * B
domain_comm = mpi.world.new_communicator(np.arange(r, r + D))
band_comm = mpi.world.new_communicator(np.arange(rank, size, D))
gd = GridDescriptor((G, G, G), (a, a, a), comm=domain_comm)
np.random.seed(rank)
psit_mG = np.random.uniform(size=(M, ) + tuple(gd.n_c))
send_mG = gd.empty(M)
recv_mG = gd.empty(M)
laplace = [Laplace(g, n=1).apply for g in gd]
for i in range(B // 2):
rrequest = band_comm.reveive(recv_mG, (rank + D) % size, 42, 0)
srequest = band_comm.send(send_mG, (rank - D) % size, 17, 0)
def set_grid_descriptor(self, gd):
self.gd = gd
axes = np.arange(3)
pbc_c = np.array(gd.pbc_c, dtype=bool)
periodic_axes = axes[pbc_c]
non_periodic_axes = axes[np.logical_not(pbc_c)]
# Find out which axes are orthogonal (0, 1 or 3)
# Note that one expects that the axes are always rotated in
# conventional form, thus for all axes to be
# classified as orthogonal, the cell_cv needs to be diagonal.
# This may always be achieved by rotating
# the unit-cell along with the atoms. The classification is
# inherited from grid_descriptor.orthogonal.
dotprods = np.dot(gd.cell_cv, gd.cell_cv.T)
# For each direction, check whether there is only one nonzero
# element in that row (necessarily being the diagonal element,
# since this is a cell vector length and must be > 0).
orthogonal_c = (np.abs(dotprods) > 1e-10).sum(axis=0) == 1
assert sum(orthogonal_c) in [0, 1, 3]
non_orthogonal_axes = axes[np.logical_not(orthogonal_c)]
if not all(pbc_c | orthogonal_c):
raise BadAxesError('Each axis must be periodic or orthogonal '
'to other axes. But we have pbc={} '
'and orthogonal={}'.format(
pbc_c.astype(int),
orthogonal_c.astype(int)))
# We sort them, and pick the longest non-periodic axes as the
# cholesky axis.
sorted_non_periodic_axes = sorted(non_periodic_axes,
key=lambda c: gd.N_c[c])
if self.use_cholesky:
if len(sorted_non_periodic_axes) > 0:
cholesky_axes = [sorted_non_periodic_axes[-1]]
if cholesky_axes[0] in non_orthogonal_axes:
msg = ('Cholesky axis cannot be non-orthogonal. '
'Do you really want a non-orthogonal non-periodic '
'axis? If so, run with use_cholesky=False.')
raise NotImplementedError(msg)
fst_axes = sorted_non_periodic_axes[0:-1]
else:
cholesky_axes = []
fst_axes = []
else:
cholesky_axes = []
fst_axes = sorted_non_periodic_axes
fft_axes = list(periodic_axes)
(self.cholesky_axes, self.fst_axes,
self.fft_axes) = cholesky_axes, fst_axes, fft_axes
fftfst_axes = self.fft_axes + self.fst_axes
axes = self.fft_axes + self.fst_axes + self.cholesky_axes
self.axes = axes
gd_x = [self.gd]
# Create xy flat decomposition (where x=axes[0] and y=axes[1])
domain = gd.N_c.copy()
domain[axes[0]] = 1
domain[axes[1]] = 1
parsize_c = decompose_domain(domain, gd.comm.size)
gd_x.append(gd.new_descriptor(parsize_c=parsize_c))
# Create z flat decomposition
domain = gd.N_c.copy()
domain[axes[2]] = 1
parsize_c = decompose_domain(domain, gd.comm.size)
gd_x.append(gd.new_descriptor(parsize_c=parsize_c))
self.gd_x = gd_x
# Calculate eigenvalues in fst/fft decomposition for
# non-cholesky axes in parallel
r_cx = np.indices(gd_x[-1].n_c)
r_cx += gd_x[-1].beg_c[:, np.newaxis, np.newaxis, np.newaxis]
r_cx = r_cx.astype(complex)
for c, axis in enumerate(fftfst_axes):
r_cx[axis] *= 2j * np.pi / gd_x[-1].N_c[axis]
if axis in fst_axes:
r_cx[axis] /= 2
for c, axis in enumerate(cholesky_axes):
r_cx[axis] = 0.0
np.exp(r_cx, out=r_cx)
fft_lambdas = np.zeros_like(r_cx[0], dtype=complex)
laplace = Laplace(gd_x[-1], -0.25 / pi, self.nn)
self.stencil_description = laplace.description
for coeff, offset_c in zip(laplace.coef_p, laplace.offset_pc):
offset_c = np.array(offset_c)
if not any(offset_c):
# The centerpoint is handled with (temp-1.0)
continue
non_zero_axes, = np.where(offset_c)
if set(non_zero_axes).issubset(fftfst_axes):
temp = np.ones_like(fft_lambdas)
for c, axis in enumerate(fftfst_axes):
temp *= r_cx[axis]**offset_c[axis]
fft_lambdas += coeff * (temp - 1.0)
assert np.linalg.norm(fft_lambdas.imag) < 1e-10
fft_lambdas = fft_lambdas.real.copy() # arr.real is not contiguous
# If there is no Cholesky decomposition, the system is already
# fully diagonal and we can directly invert the linear problem
# by dividing with the eigenvalues.
assert len(cholesky_axes) == 0
# if len(cholesky_axes) == 0:
with np.errstate(divide='ignore'):
self.inv_fft_lambdas = np.where(
np.abs(fft_lambdas) > 1e-10, 1.0 / fft_lambdas, 0)
offload_enabled = os.environ.get("GPAW_OFFLOAD");
device = mic.devices[0]
gpts = 64
gd = GridDescriptor([gpts, gpts, gpts])
a = gd.empty() # source
b = gd.empty() # result
np.random.seed(10)
a[:] = np.random.random(a.shape)
b[:] = np.zeros(b.shape)
op = Laplace(gd, 1.0, 3)
if offload_enabled:
offl_a = device.bind(a)
offl_b = device.bind(b)
print "--------------------------------------------------------------"
print "Starting relaxation step"
relax_start = time.time()
if offload_enabled:
op.relax(2, offl_b.array, offl_a.array, 4, 2.0 / 3.0)
else:
op.relax(2, b, a, 4, 2.0 / 3.0)
relax_end = time.time()
print "--------------------------------------------------------------"
# Strange one!
continue
print('------------------')
print(name, D)
print(cell[0])
print(cell[1])
print(cell[2])
for n in range(1, 6):
N = 2 * n + 2
gd = GridDescriptor((N, N, N), cell)
b_g = gd.zeros()
r_gv = gd.get_grid_point_coordinates().transpose((1, 2, 3, 0))
c_v = gd.cell_cv.sum(0) / 2
r_gv -= c_v
lap = Laplace(gd, n=n)
grad_v = [Gradient(gd, v, n=n) for v in range(3)]
assert lap.npoints == D * 2 * n + 1
for m in range(0, 2 * n + 1):
for ix in range(m + 1):
for iy in range(m - ix + 1):
iz = m - ix - iy
a_g = (r_gv**(ix, iy, iz)).prod(3)
if ix + iy + iz == 2 and max(ix, iy, iz) == 2:
r = 2.0
else:
r = 0.0
lap.apply(a_g, b_g)
e = b_g[n + 1, n + 1, n + 1] - r
assert abs(e) < 2e-12, e
for v in range(3):
def set_grid_descriptor(self, gd):
# If non-periodic boundary conditions is used,
# there is problems with auxiliary grid.
# Maybe with use_aux_grid=False it would work?
if gd.pbc_c.any():
raise NotImplementedError('Only non-periodic boundary '
'conditions are tested')
self.gd_small_fine = gd
assert np.all(self.gd_small_fine.N_c <= self.N_large_fine_c), \
'extended grid has to be larger than the original one'
if self.use_coarse:
# 1.1. Construct coarse chain on the small grid
self.coarser_i = []
gd = self.gd_small_fine
N_c = self.N_large_fine_c.copy()
for i in range(self.Ncoar):
gd2 = gd.coarsen()
self.coarser_i.append(Transformer(gd, gd2, self.nn_coarse))
N_c //= 2
gd = gd2
self.gd_small_coar = gd
else:
self.gd_small_coar = self.gd_small_fine
N_c = self.N_large_fine_c
# 1.2. Construct coarse extended grid
self.gd_large_coar = ext_gd(self.gd_small_coar, N_c=N_c)
# Initialize poissonsolvers
self.ps_large_coar.set_grid_descriptor(self.gd_large_coar)
if not self.use_coarse:
return
self.ps_small_fine.set_grid_descriptor(self.gd_small_fine)
if self.use_aux_grid:
# 2.1. Construct an auxiliary grid that is the small grid plus
# a buffer region allowing Laplace and refining
# with the used stencils
buf = self.nn_refine
for i in range(self.Ncoar):
buf = 2 * buf + self.nn_refine
buf += self.nn_laplace
div = 2**self.Ncoar
if buf % div != 0:
buf += div - buf % div
N_c = self.gd_small_fine.N_c + 2 * buf
if np.any(N_c > self.N_large_fine_c):
self.use_aux_grid = False
N_c = self.N_large_fine_c
self.gd_aux_fine = ext_gd(self.gd_small_fine, N_c=N_c)
else:
self.gd_aux_fine = ext_gd(self.gd_small_fine,
N_c=self.N_large_fine_c)
# 2.2 Construct Laplace on the aux grid
self.laplace_aux_fine = Laplace(self.gd_aux_fine, - 0.25 / np.pi,
self.nn_laplace)
# 2.3 Construct refine chain
self.refiner_i = []
gd = self.gd_aux_fine
N_c = gd.N_c.copy()
for i in range(self.Ncoar):
gd2 = gd.coarsen()
self.refiner_i.append(Transformer(gd2, gd, self.nn_refine))
N_c //= 2
gd = gd2
self.refiner_i = self.refiner_i[::-1]
self.gd_aux_coar = gd
if self.use_aux_grid:
# 2.4 Construct large coarse grid from aux grid
self.gd_large_coar_from_aux = ext_gd(self.gd_aux_coar,
N_c=self.gd_large_coar.N_c)
# Check the consistency of the grids
gd1 = self.gd_large_coar
gd2 = self.gd_large_coar_from_aux
assert np.all(gd1.N_c == gd2.N_c) and np.all(gd1.h_cv == gd2.h_cv)
self._initialized = False
