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