def __init__(self, problem_params, dtype_u=particles, dtype_f=acceleration): """ Initialization routine Args: problem_params (dict): custom parameters for the example dtype_u: particle data type (will be passed to parent class) dtype_f: acceleration data type (will be passed to parent class) """ # these parameters will be used later, so assert their existence essential_keys = ['k', 'phase', 'amp'] for key in essential_keys: if key not in problem_params: msg = 'need %s to instantiate problem, only got %s' % ( key, str(problem_params.keys())) raise ParameterError(msg) # invoke super init, passing nparts, dtype_u and dtype_f super(harmonic_oscillator, self).__init__(1, dtype_u, dtype_f, problem_params) if self.params.phase != 0.0: raise ProblemError('Phase != 0 not implemented yet') if self.params.amp != 1.0: raise ProblemError('amp != 1 not implemented yet')
def __init__(self, problem_params, dtype_u, dtype_f): """ Initialization routine Args: problem_params (dict): custom parameters for the example dtype_u: mesh data type (will be passed parent class) dtype_f: mesh data type (will be passed parent class) """ # these parameters will be used later, so assert their existence essential_keys = ['nvars', 'nu', 'freq'] for key in essential_keys: if key not in problem_params: msg = 'need %s to instantiate problem, only got %s' % (key, str(problem_params.keys())) raise ParameterError(msg) # make sure parameters have the correct form if problem_params['freq'] >= 0 and problem_params['freq'] % 2 != 0: raise ProblemError('need even number of frequencies due to periodic BCs') if problem_params['nvars'] % 2 != 0: raise ProblemError('the setup requires nvars = 2^p per dimension') # invoke super init, passing number of dofs, dtype_u and dtype_f super(heat1d_periodic, self).__init__(init=problem_params['nvars'], dtype_u=dtype_u, dtype_f=dtype_f, params=problem_params) # compute dx (equal in both dimensions) and get discretization matrix A self.dx = 1.0 / self.params.nvars self.A = self.__get_A(self.params.nvars, self.params.nu, self.dx)
def __init__(self, problem_params, dtype_u=mesh, dtype_f=rhs_imex_mesh): """ Initialization routine Args: problem_params (dict): custom parameters for the example dtype_u: mesh data type (will be passed to parent class) dtype_f: mesh data type wuth implicit and explicit parts (will be passed to parent class) """ if 'L' not in problem_params: problem_params['L'] = 1.0 if 'init_type' not in problem_params: problem_params['init_type'] = 'circle' # these parameters will be used later, so assert their existence essential_keys = ['nvars', 'nu', 'eps', 'L', 'radius'] for key in essential_keys: if key not in problem_params: msg = 'need %s to instantiate problem, only got %s' % ( key, str(problem_params.keys())) raise ParameterError(msg) # we assert that nvars looks very particular here.. this will be necessary for coarsening in space later on if len(problem_params['nvars']) != 2: raise ProblemError('this is a 2d example, got %s' % problem_params['nvars']) if problem_params['nvars'][0] != problem_params['nvars'][1]: raise ProblemError('need a square domain, got %s' % problem_params['nvars']) if problem_params['nvars'][0] % 2 != 0: raise ProblemError('the setup requires nvars = 2^p per dimension') # invoke super init, passing number of dofs, dtype_u and dtype_f super(allencahn2d_imex, self).__init__(init=problem_params['nvars'], dtype_u=dtype_u, dtype_f=dtype_f, params=problem_params) self.dx = self.params.L / self.params.nvars[ 0] # could be useful for hooks, too. self.xvalues = np.array([ i * self.dx - self.params.L / 2.0 for i in range(self.params.nvars[0]) ]) kx = np.zeros(self.init[0]) ky = np.zeros(self.init[1] // 2 + 1) kx[:int(self.init[0] / 2) + 1] = 2 * np.pi / self.params.L * np.arange( 0, int(self.init[0] / 2) + 1) kx[int(self.init[0] / 2) + 1:] = 2 * np.pi / self.params.L * \ np.arange(int(self.init[0] / 2) + 1 - self.init[0], 0) ky[:] = 2 * np.pi / self.params.L * np.arange(0, self.init[1] // 2 + 1) xv, yv = np.meshgrid(kx, ky, indexing='ij') self.lap = -xv**2 - yv**2
def __init__(self, problem_params, dtype_u, dtype_f): """ Initialization routine Args: problem_params (dict): custom parameters for the example dtype_u: mesh data type (will be passed parent class) dtype_f: mesh data type (will be passed parent class) """ # these parameters will be used later, so assert their existence essential_keys = [ 'nvars', 'nu', 'lambda0', 'newton_maxiter', 'newton_tol', 'interval' ] for key in essential_keys: if key not in problem_params: msg = 'need %s to instantiate problem, only got %s' % ( key, str(problem_params.keys())) raise ParameterError(msg) # we assert that nvars looks very particular here.. this will be necessary for coarsening in space later on if (problem_params['nvars'] + 1) % 2 != 0: raise ProblemError('setup requires nvars = 2^p - 1') # invoke super init, passing number of dofs, dtype_u and dtype_f super(generalized_fisher, self).__init__(problem_params['nvars'], dtype_u, dtype_f, problem_params) # compute dx and get discretization matrix A self.dx = (self.params.interval[1] - self.params.interval[0]) / (self.params.nvars + 1) self.A = self.__get_A(self.params.nvars, self.dx)
def __init__(self, problem_params, dtype_u=mesh, dtype_f=mesh): """ Initialization routine Args: problem_params (dict): custom parameters for the example dtype_u: mesh data type for solution dtype_f: mesh data type for RHS """ # these parameters will be used later, so assert their existence essential_keys = ['nvars', 'nu', 'freq'] for key in essential_keys: if key not in problem_params: msg = 'need %s to instantiate problem, only got %s' % ( key, str(problem_params.keys())) raise ParameterError(msg) # we assert that nvars looks very particular here.. this will be necessary for coarsening in space later on if (problem_params['nvars'] + 1) % 2 != 0: raise ProblemError('setup requires nvars = 2^p - 1') # invoke super init, passing number of dofs, dtype_u and dtype_f super(heat1d, self).__init__(init=problem_params['nvars'], dtype_u=dtype_u, dtype_f=dtype_f, params=problem_params) # compute dx and get discretization matrix A self.dx = 1.0 / (self.params.nvars + 1) self.A = self.__get_A(self.params.nvars, self.params.nu, self.dx)
def __init__(self, problem_params, dtype_u=mesh, dtype_f=rhs_comp2_mesh): """ Initialization routine Args: problem_params (dict): custom parameters for the example dtype_u: mesh data type (will be passed parent class) dtype_f: mesh data type (will be passed parent class) """ # these parameters will be used later, so assert their existence essential_keys = [ 'nvars', 'dw', 'eps', 'newton_maxiter', 'newton_tol', 'interval', 'radius' ] for key in essential_keys: if key not in problem_params: msg = 'need %s to instantiate problem, only got %s' % ( key, str(problem_params.keys())) raise ParameterError(msg) # we assert that nvars looks very particular here.. this will be necessary for coarsening in space later on if (problem_params['nvars']) % 2 != 0: raise ProblemError('setup requires nvars = 2^p') if 'stop_at_nan' not in problem_params: problem_params['stop_at_nan'] = True # invoke super init, passing number of dofs, dtype_u and dtype_f super(allencahn_periodic_multiimplicit, self).__init__(problem_params, dtype_u, dtype_f) self.A -= sp.eye(self.init) * 0.0 / self.params.eps**2
def build_f(self, f, part, t): """ Helper function to assemble the correct right-hand side out of B and E field Args: f (dtype_f): the field values part (dtype_u): the current particles data t (float): the current time Returns: acceleration: correct RHS of type acceleration """ if not isinstance(part, particles): raise ProblemError('something is wrong during build_f, got %s' % type(part)) N = self.params.nparts rhs = acceleration((3, self.params.nparts)) for n in range(N): rhs.values[:, n] = part.q[n] / part.m[n] * ( f.elec.values[:, n] + np.cross(part.vel.values[:, n], f.magn.values[:, n])) return rhs
def __init__(self, problem_params, dtype_u, dtype_f): """ Initialization routine Args: problem_params (dict): custom parameters for the example dtype_u: mesh data type (will be passed parent class) dtype_f: mesh data type (will be passed parent class) """ # these parameters will be used later, so assert their existence essential_keys = [ 'nvars', 'nu', 'eps', 'newton_maxiter', 'newton_tol', 'lin_tol', 'lin_maxiter', 'radius' ] for key in essential_keys: if key not in problem_params: msg = 'need %s to instantiate problem, only got %s' % ( key, str(problem_params.keys())) raise ParameterError(msg) # we assert that nvars looks very particular here.. this will be necessary for coarsening in space later on if len(problem_params['nvars']) != 2: raise ProblemError('this is a 2d example, got %s' % problem_params['nvars']) if problem_params['nvars'][0] != problem_params['nvars'][1]: raise ProblemError('need a square domain, got %s' % problem_params['nvars']) if problem_params['nvars'][0] % 2 != 0: raise ProblemError('the setup requires nvars = 2^p per dimension') # invoke super init, passing number of dofs, dtype_u and dtype_f super(allencahn_fullyimplicit, self).__init__(problem_params['nvars'], dtype_u, dtype_f, problem_params) # compute dx and get discretization matrix A self.dx = 1.0 / self.params.nvars[0] self.A = self.__get_A(self.params.nvars, self.dx) self.xvalues = np.array( [i * self.dx - 0.5 for i in range(self.params.nvars[0])]) self.newton_itercount = 0 self.lin_itercount = 0 self.newton_ncalls = 0 self.lin_ncalls = 0
def u_init(self): """ Routine to compute the starting values for the particles Returns: dtype_u: particle set filled with initial data """ u0 = self.params.u0 N = self.params.nparts u = self.dtype_u((3, N)) if u0[2][0] is not 1 or u0[3][0] is not 1: raise ProblemError('so far only q = m = 1 is implemented') # set first particle to u0 u.pos.values[0, 0] = u0[0][0] u.pos.values[1, 0] = u0[0][1] u.pos.values[2, 0] = u0[0][2] u.vel.values[0, 0] = u0[1][0] u.vel.values[1, 0] = u0[1][1] u.vel.values[2, 0] = u0[1][2] u.q[0] = u0[2][0] u.m[0] = u0[3][0] # initialize random seed np.random.seed(N) comx = u.pos.values[0, 0] comy = u.pos.values[1, 0] comz = u.pos.values[2, 0] for n in range(1, N): # draw 3 random variables in [-1,1] to shift positions r = np.random.random_sample(3) - 1 u.pos.values[0, n] = r[0] + u0[0][0] u.pos.values[1, n] = r[1] + u0[0][1] u.pos.values[2, n] = r[2] + u0[0][2] # draw 3 random variables in [-5,5] to shift velocities r = np.random.random_sample(3) - 5 u.vel.values[0, n] = r[0] + u0[1][0] u.vel.values[1, n] = r[1] + u0[1][1] u.vel.values[2, n] = r[2] + u0[1][2] u.q[n] = u0[2][0] u.m[n] = u0[3][0] # gather positions to check center comx += u.pos.values[0, n] comy += u.pos.values[1, n] comz += u.pos.values[2, n] # print('Center of positions:',comx/N,comy/N,comz/N) return u
def __get_A(N, c, dx, order, type): """ Helper function to assemble FD matrix A in sparse format Args: N (int): number of dofs c (float): diffusion coefficient dx (float): distance between two spatial nodes order (int): specifies order of discretization type (string): upwind or centered differences Returns: scipy.sparse.csc_matrix: matrix A in CSC format """ coeff = None stencil = None zero_pos = None if type == 'center': if order == 2: stencil = [-1.0, 0.0, 1.0] zero_pos = 2 coeff = 1.0 / 2.0 else: raise ProblemError("Order " + str(order) + " not implemented.") else: if order == 1: stencil = [-1.0, 1.0] coeff = 1.0 zero_pos = 2 else: raise ProblemError("Order " + str(order) + " not implemented.") offsets = [pos - zero_pos + 1 for pos in range(len(stencil))] A = sp.diags(stencil, offsets, shape=(N, N), format='csc') A *= c * coeff * (1.0 / dx) return A
def __init__(self, problem_params, dtype_u=dedalus_field, dtype_f=rhs_imex_dedalus_field): """ Initialization routine Args: problem_params (dict): custom parameters for the example dtype_u: mesh data type (will be passed parent class) dtype_f: mesh data type (will be passed parent class) """ if 'comm' not in problem_params: problem_params['comm'] = None # these parameters will be used later, so assert their existence essential_keys = ['nvars', 'nu', 'freq', 'comm'] for key in essential_keys: if key not in problem_params: msg = 'need %s to instantiate problem, only got %s' % ( key, str(problem_params.keys())) raise ParameterError(msg) # we assert that nvars looks very particular here.. this will be necessary for coarsening in space later on if problem_params['freq'] % 2 != 0: raise ProblemError('setup requires freq to be an equal number') xbasis = de.Fourier('x', problem_params['nvars'][0], interval=(0, 1), dealias=1) ybasis = de.Fourier('y', problem_params['nvars'][1], interval=(0, 1), dealias=1) domain = de.Domain([xbasis, ybasis], grid_dtype=np.float64, comm=problem_params['comm']) # invoke super init, passing number of dofs, dtype_u and dtype_f super(heat2d_dedalus_forced, self).__init__(init=domain, dtype_u=dtype_u, dtype_f=dtype_f, params=problem_params) self.x = self.init.grid(0, scales=1) self.y = self.init.grid(1, scales=1) self.rhs = self.dtype_u(self.init, val=0.0) self.problem = de.IVP(domain=self.init, variables=['u']) self.problem.parameters['nu'] = self.params.nu self.problem.add_equation( "dt(u) - nu * dx(dx(u)) - nu * dy(dy(u)) = 0") self.solver = self.problem.build_solver(de.timesteppers.SBDF1) self.u = self.solver.state['u']
def __init__(self, problem_params, dtype_u, dtype_f): """ Initialization routine Args: problem_params (dict): custom parameters for the example dtype_u: mesh data type (will be passed parent class) dtype_f: mesh data type (will be passed parent class) """ if 'L' not in problem_params: problem_params['L'] = 1.0 # these parameters will be used later, so assert their existence essential_keys = ['nvars', 'c', 'freq', 'nu', 'L'] for key in essential_keys: if key not in problem_params: msg = 'need %s to instantiate problem, only got %s' % ( key, str(problem_params.keys())) raise ParameterError(msg) # we assert that nvars looks very particular here.. this will be necessary for coarsening in space later on if (problem_params['nvars']) % 2 != 0: raise ProblemError('setup requires nvars = 2^p') # invoke super init, passing number of dofs, dtype_u and dtype_f super(advectiondiffusion1d_imex, self).__init__(init=problem_params['nvars'], dtype_u=dtype_u, dtype_f=dtype_f, params=problem_params) self.xvalues = np.array([ i * self.params.L / self.params.nvars - self.params.L / 2.0 for i in range(self.params.nvars) ]) kx = np.zeros(self.init // 2 + 1) for i in range(0, len(kx)): kx[i] = 2 * np.pi / self.params.L * i self.ddx = kx * 1j self.lap = -kx**2 rfft_in = pyfftw.empty_aligned(self.init, dtype='float64') fft_out = pyfftw.empty_aligned(self.init // 2 + 1, dtype='complex128') ifft_in = pyfftw.empty_aligned(self.init // 2 + 1, dtype='complex128') irfft_out = pyfftw.empty_aligned(self.init, dtype='float64') self.rfft_object = pyfftw.FFTW(rfft_in, fft_out, direction='FFTW_FORWARD') self.irfft_object = pyfftw.FFTW(ifft_in, irfft_out, direction='FFTW_BACKWARD')
def __init__(self, problem_params, dtype_u, dtype_f): """ Initialization routine Args: problem_params (dict): custom parameters for the example dtype_u: mesh data type (will be passed parent class) dtype_f: mesh data type (will be passed parent class) """ # these parameters will be used later, so assert their existence if 'comm' not in problem_params: problem_params['comm'] = PETSc.COMM_WORLD if 'sol_tol' not in problem_params: problem_params['sol_tol'] = 1E-10 if 'sol_maxiter' not in problem_params: problem_params['sol_maxiter'] = None essential_keys = ['nvars', 'nu', 'freq', 'comm'] for key in essential_keys: if key not in problem_params: msg = 'need %s to instantiate problem, only got %s' % (key, str(problem_params.keys())) raise ParameterError(msg) # make sure parameters have the correct form if len(problem_params['nvars']) != 2: raise ProblemError('this is a 2d example, got %s' % problem_params['nvars']) # create DMDA object which will be used for all grid operations da = PETSc.DMDA().create([problem_params['nvars'][0], problem_params['nvars'][1]], stencil_width=1, comm=problem_params['comm']) # invoke super init, passing number of dofs, dtype_u and dtype_f super(heat2d_petsc_forced, self).__init__(init=da, dtype_u=dtype_u, dtype_f=dtype_f, params=problem_params) # compute dx, dy and get local ranges self.dx = 1.0 / (self.params.nvars[0] - 1) self.dy = 1.0 / (self.params.nvars[1] - 1) (self.xs, self.xe), (self.ys, self.ye) = self.init.getRanges() # compute discretization matrix A and identity self.A = self.__get_A() self.Id = self.__get_Id() # setup solver self.ksp = PETSc.KSP() self.ksp.create(comm=self.params.comm) self.ksp.setType('cg') pc = self.ksp.getPC() pc.setType('none') self.ksp.setInitialGuessNonzero(True) self.ksp.setFromOptions() self.ksp.setTolerances(rtol=self.params.sol_tol, atol=self.params.sol_tol, max_it=self.params.sol_maxiter)
def __init__(self, problem_params, dtype_u, dtype_f): """ Initialization routine Args: problem_params (dict): custom parameters for the example dtype_u: mesh data type (will be passed parent class) dtype_f: mesh data type (will be passed parent class) """ # these parameters will be used later, so assert their existence essential_keys = ['nvars', 'c', 'freq'] for key in essential_keys: if key not in problem_params: msg = 'need %s to instantiate problem, only got %s' % ( key, str(problem_params.keys())) raise ParameterError(msg) # we assert that nvars looks very particular here.. this will be necessary for coarsening in space later on if (problem_params['nvars']) % 2 != 0: raise ProblemError('setup requires nvars = 2^p') if problem_params['freq'] >= 0 and problem_params['freq'] % 2 != 0: raise ProblemError( 'need even number of frequencies due to periodic BCs') if 'order' not in problem_params: problem_params['order'] = 1 if 'type' not in problem_params: problem_params['type'] = 'upwind' # invoke super init, passing number of dofs, dtype_u and dtype_f super(advection1d, self).__init__(init=problem_params['nvars'], dtype_u=dtype_u, dtype_f=dtype_f, params=problem_params) # compute dx and get discretization matrix A self.dx = 1.0 / self.params.nvars self.A = self.__get_A(self.params.nvars, self.params.c, self.dx, self.params.order, self.params.type)
def __get_A(N, nu, dx, ndim, order): """ Helper function to assemble FD matrix A in sparse format Args: N (list): number of dofs nu (float): diffusion coefficient dx (float): distance between two spatial nodes ndim (int): number of dimensions Returns: scipy.sparse.csc_matrix: matrix A in CSC format """ if order == 2: stencil = [1, -2, 1] zero_pos = 2 elif order == 4: stencil = [-1 / 12, 4 / 3, -5 / 2, 4 / 3, -1 / 12] zero_pos = 3 elif order == 6: stencil = [ 1 / 90, -3 / 20, 3 / 2, -49 / 18, 3 / 2, -3 / 20, 1 / 90 ] zero_pos = 4 elif order == 8: stencil = [ -1 / 560, 8 / 315, -1 / 5, 8 / 5, -205 / 72, 8 / 5, -1 / 5, 8 / 315, -1 / 560 ] zero_pos = 5 else: raise ProblemError( f'wrong order given, has to be 2, 4, 6, or 8, got {order}') dstencil = np.concatenate((stencil, np.delete(stencil, zero_pos - 1))) offsets = np.concatenate( ([N[0] - i - 1 for i in reversed(range(zero_pos - 1))], [i - zero_pos + 1 for i in range(zero_pos - 1, len(stencil))])) doffsets = np.concatenate( (offsets, np.delete(offsets, zero_pos - 1) - N[0])) A = sp.diags(dstencil, doffsets, shape=(N[0], N[0]), format='csc') # stencil = [1, -2, 1] # A = sp.diags(stencil, [-1, 0, 1], shape=(N[0], N[0]), format='csc') if ndim == 2: A = sp.kron(A, sp.eye(N[0])) + sp.kron(sp.eye(N[1]), A) elif ndim == 3: A = sp.kron(A, sp.eye(N[1] * N[0])) + sp.kron(sp.eye(N[2] * N[1]), A) + \ sp.kron(sp.kron(sp.eye(N[2]), A), sp.eye(N[0])) A *= nu / (dx**2) return A
def u_exact(self, t): """ Routine to compute the exact trajectory at time t (only for single-particle setup) Args: t (float): current time Returns: dtype_u: particle type containing the exact position and velocity """ # some abbreviations wE = self.params.omega_E wB = self.params.omega_B N = self.params.nparts u0 = self.params.u0 if N != 1: raise ProblemError('u_exact is only valid for a single particle') u = self.dtype_u((3, 1)) wbar = np.sqrt(2) * wE # position and velocity in z direction is easy to compute u.pos.values[2, 0] = u0[0][2] * np.cos( wbar * t) + u0[1][2] / wbar * np.sin(wbar * t) u.vel.values[2, 0] = -u0[0][2] * wbar * np.sin( wbar * t) + u0[1][2] * np.cos(wbar * t) # define temp. variables to compute complex position Op = 1 / 2 * (wB + np.sqrt(wB**2 - 4 * wE**2)) Om = 1 / 2 * (wB - np.sqrt(wB**2 - 4 * wE**2)) Rm = (Op * u0[0][0] + u0[1][1]) / (Op - Om) Rp = u0[0][0] - Rm Im = (Op * u0[0][1] - u0[1][0]) / (Op - Om) Ip = u0[0][1] - Im # compute position in complex notation w = np.complex(Rp, Ip) * np.exp(-np.complex(0, Op * t)) + np.complex( Rm, Im) * np.exp(-np.complex(0, Om * t)) # compute velocity as time derivative of the position dw = -1j * Op * np.complex(Rp, Ip) * \ np.exp(-np.complex(0, Op * t)) - 1j * Om * np.complex(Rm, Im) * np.exp(-np.complex(0, Om * t)) # get the appropriate real and imaginary parts u.pos.values[0, 0] = w.real u.vel.values[0, 0] = dw.real u.pos.values[1, 0] = w.imag u.vel.values[1, 0] = dw.imag return u
def __init__(self, problem_params, dtype_u=mesh, dtype_f=mesh): """ Initialization routine Args: problem_params (dict): custom parameters for the example dtype_u: mesh data type (will be passed parent class) dtype_f: mesh data type (will be passed parent class) """ # these parameters will be used later, so assert their existence essential_keys = [ 'nvars', 'dw', 'eps', 'newton_maxiter', 'newton_tol', 'interval', 'radius' ] for key in essential_keys: if key not in problem_params: msg = 'need %s to instantiate problem, only got %s' % ( key, str(problem_params.keys())) raise ParameterError(msg) # we assert that nvars looks very particular here.. this will be necessary for coarsening in space later on if (problem_params['nvars']) % 2 != 0: raise ProblemError('setup requires nvars = 2^p') if 'stop_at_nan' not in problem_params: problem_params['stop_at_nan'] = True # invoke super init, passing number of dofs, dtype_u and dtype_f super(allencahn_periodic_fullyimplicit, self).__init__(problem_params['nvars'], dtype_u, dtype_f, problem_params) # compute dx and get discretization matrix A self.dx = (self.params.interval[1] - self.params.interval[0]) / self.params.nvars self.xvalues = np.array([ self.params.interval[0] + i * self.dx for i in range(self.params.nvars) ]) self.A = self.__get_A(self.params.nvars, self.dx) self.newton_itercount = 0 self.lin_itercount = 0 self.newton_ncalls = 0 self.lin_ncalls = 0
def __init__(self, problem_params, dtype_u=pmesh_datatype, dtype_f=rhs_imex_pmesh): """ Initialization routine Args: problem_params (dict): custom parameters for the example dtype_u: pmesh data type (will be passed to parent class) dtype_f: pmesh data type wuth implicit and explicit parts (will be passed to parent class) """ if 'L' not in problem_params: problem_params['L'] = 1.0 if 'init_type' not in problem_params: problem_params['init_type'] = 'circle' if 'comm' not in problem_params: problem_params['comm'] = None if 'dw' not in problem_params: problem_params['dw'] = 0.0 # these parameters will be used later, so assert their existence essential_keys = ['nvars', 'eps', 'L', 'radius', 'dw'] for key in essential_keys: if key not in problem_params: msg = 'need %s to instantiate problem, only got %s' % (key, str(problem_params.keys())) raise ParameterError(msg) if not (isinstance(problem_params['nvars'], tuple) and len(problem_params['nvars']) > 1): raise ProblemError('Need at least two dimensions') # Creating ParticleMesh structure self.pm = ParticleMesh(BoxSize=problem_params['L'], Nmesh=list(problem_params['nvars']), dtype='f8', plan_method='measure', comm=problem_params['comm']) # create test RealField to get the local dimensions (there's probably a better way to do that) tmp = self.pm.create(type='real') # invoke super init, passing the communicator and the local dimensions as init super(allencahn_imex, self).__init__(init=(self.pm.comm, tmp.value.shape), dtype_u=dtype_u, dtype_f=dtype_f, params=problem_params) # Need this for diagnostics self.dx = self.params.L / problem_params['nvars'][0] self.dy = self.params.L / problem_params['nvars'][1] self.xvalues = [i * self.dx - problem_params['L'] / 2 for i in range(problem_params['nvars'][0])] self.yvalues = [i * self.dy - problem_params['L'] / 2 for i in range(problem_params['nvars'][1])]
def solve_system(self, rhs, dt, u0, t): """ Simple Newton solver for the nonlinear system Args: rhs (dtype_f): right-hand side for the nonlinear system dt (float): abbrev. for the node-to-node stepsize (or any other factor required) u0 (dtype_u): initial guess for the iterative solver t (float): current time (e.g. for time-dependent BCs) Returns: dtype_u: solution u """ mu = self.params.mu # create new mesh object from u0 and set initial values for iteration u = self.dtype_u(u0) x1 = u.values[0] x2 = u.values[1] # start newton iteration n = 0 res = 99 while n < self.params.newton_maxiter: # form the function g with g(u) = 0 g = np.array([ x1 - dt * x2 - rhs.values[0], x2 - dt * (mu * (1 - x1**2) * x2 - x1) - rhs.values[1] ]) # if g is close to 0, then we are done res = np.linalg.norm(g, np.inf) if res < self.params.newton_tol or np.isnan(res): break # prefactor for dg/du c = 1.0 / (-2 * dt**2 * mu * x1 * x2 - dt**2 - 1 + dt * mu * (1 - x1**2)) # assemble dg/du dg = c * np.array([[dt * mu * (1 - x1**2) - 1, -dt], [2 * dt * mu * x1 * x2 + dt, -1]]) # newton update: u1 = u0 - g/dg u.values -= np.dot(dg, g) # set new values and increase iteration count x1 = u.values[0] x2 = u.values[1] n += 1 if np.isnan(res) and self.params.stop_at_nan: raise ProblemError( 'Newton got nan after %i iterations, aborting...' % n) elif np.isnan(res): self.logger.warning('Newton got nan after %i iterations...' % n) if n == self.params.newton_maxiter: raise ProblemError( 'Newton did not converge after %i iterations, error is %s' % (n, res)) return u
def __init__(self, problem_params, dtype_u, dtype_f): """ Initialization routine Args: problem_params (dict): custom parameters for the example dtype_u: mesh data type (will be passed parent class) dtype_f: mesh data type (will be passed parent class) """ if 'L' not in problem_params: problem_params['L'] = 1.0 if 'init_type' not in problem_params: problem_params['init_type'] = 'circle' # these parameters will be used later, so assert their existence essential_keys = ['nvars', 'nu', 'eps', 'L', 'radius'] for key in essential_keys: if key not in problem_params: msg = 'need %s to instantiate problem, only got %s' % ( key, str(problem_params.keys())) raise ParameterError(msg) # we assert that nvars looks very particular here.. this will be necessary for coarsening in space later on if len(problem_params['nvars']) != 2: raise ProblemError('this is a 2d example, got %s' % problem_params['nvars']) if problem_params['nvars'][0] != problem_params['nvars'][1]: raise ProblemError('need a square domain, got %s' % problem_params['nvars']) if problem_params['nvars'][0] % 2 != 0: raise ProblemError('the setup requires nvars = 2^p per dimension') # invoke super init, passing number of dofs, dtype_u and dtype_f super(allencahn2d_imex, self).__init__(init=problem_params['nvars'], dtype_u=dtype_u, dtype_f=dtype_f, params=problem_params) dx = self.params.L / self.params.nvars[0] self.xvalues = np.array([ i * dx - self.params.L / 2.0 for i in range(self.params.nvars[0]) ]) kx = np.zeros(self.init[0]) ky = np.zeros(self.init[1]) for i in range(0, int(self.init[0] / 2) + 1): kx[i] = 2 * np.pi / self.params.L * i for i in range(0, int(self.init[1] / 2) + 1): ky[i] = 2 * np.pi / self.params.L * i for i in range(int(self.init[0] / 2) + 1, self.init[0]): kx[i] = 2 * np.pi / self.params.L * (-self.init[0] + i) for i in range(int(self.init[1] / 2) + 1, self.init[1]): ky[i] = 2 * np.pi / self.params.L * (-self.init[1] + i) self.lap = np.zeros(self.init) for i in range(self.init[0]): for j in range(self.init[1]): self.lap[i, j] = -kx[i]**2 - ky[j]**2 # TODO: cleanup and move to real-valued FFT fft_in = pyfftw.empty_aligned(self.init, dtype='complex128') fft_out = pyfftw.empty_aligned(self.init, dtype='complex128') ifft_in = pyfftw.empty_aligned(self.init, dtype='complex128') ifft_out = pyfftw.empty_aligned(self.init, dtype='complex128') self.fft_object = pyfftw.FFTW(fft_in, fft_out, direction='FFTW_FORWARD', axes=(0, 1)) self.ifft_object = pyfftw.FFTW(ifft_in, ifft_out, direction='FFTW_BACKWARD', axes=(0, 1))
def solve_system_2(self, rhs, factor, u0, t): """ Simple linear solver for (I-factor*A)u = rhs Args: rhs (dtype_f): right-hand side for the linear system factor (float): abbrev. for the local stepsize (or any other factor required) u0 (dtype_u): initial guess for the iterative solver t (float): current time (e.g. for time-dependent BCs) Returns: dtype_u: solution as mesh """ u = self.dtype_u(u0).values eps2 = self.params.eps**2 dw = self.params.dw Id = sp.eye(self.params.nvars) # start newton iteration n = 0 res = 99 while n < self.params.newton_maxiter: # print(n) # form the function g with g(u) = 0 g = u - rhs.values - factor * (-2.0 / eps2 * u * (1.0 - u) * (1.0 - 2.0 * u) - 6.0 * dw * u * (1.0 - u) + 0.0 / self.params.eps**2 * u) # if g is close to 0, then we are done res = np.linalg.norm(g, np.inf) if res < self.params.newton_tol: break # assemble dg dg = Id - factor * (-2.0 / eps2 * sp.diags( (1.0 - u) * (1.0 - 2.0 * u) - u * ((1.0 - 2.0 * u) + 2.0 * (1.0 - u)), offsets=0) - 6.0 * dw * sp.diags( (1.0 - u) - u, offsets=0) + 0.0 / self.params.eps**2 * Id) # newton update: u1 = u0 - g/dg u -= spsolve(dg, g) # u -= gmres(dg, g, x0=z, tol=self.params.lin_tol)[0] # increase iteration count n += 1 if np.isnan(res) and self.params.stop_at_nan: raise ProblemError( 'Newton got nan after %i iterations, aborting...' % n) elif np.isnan(res): self.logger.warning('Newton got nan after %i iterations...' % n) if n == self.params.newton_maxiter: self.logger.warning( 'Newton did not converge after %i iterations, error is %s' % (n, res)) self.newton_ncalls += 1 self.newton_itercount += n me = self.dtype_u(self.init) me.values = u return me
def solve_system(self, rhs, factor, u0, t): """ Simple Newton solver Args: rhs (dtype_f): right-hand side for the nonlinear system factor (float): abbrev. for the node-to-node stepsize (or any other factor required) u0 (dtype_u): initial guess for the iterative solver t (float): current time (required here for the BC) Returns: dtype_u: solution u """ u = self.dtype_u(u0) nu = self.params.nu lambda0 = self.params.lambda0 # set up boundary values to embed inner points lam1 = lambda0 / 2.0 * ((nu / 2.0 + 1)**0.5 + (nu / 2.0 + 1)**(-0.5)) sig1 = lam1 - np.sqrt(lam1**2 - lambda0**2) ul = (1 + (2**(nu / 2.0) - 1) * np.exp(-nu / 2.0 * sig1 * (self.params.interval[0] + 2 * lam1 * t)))**(-2.0 / nu) ur = (1 + (2**(nu / 2.0) - 1) * np.exp(-nu / 2.0 * sig1 * (self.params.interval[1] + 2 * lam1 * t)))**(-2.0 / nu) # start newton iteration n = 0 res = 99 while n < self.params.newton_maxiter: # form the function g with g(u) = 0 uext = np.concatenate(([ul], u.values, [ur])) g = u.values - \ factor * (self.A.dot(uext)[1:-1] + lambda0 ** 2 * u.values * (1 - u.values ** nu)) - rhs.values # if g is close to 0, then we are done res = np.linalg.norm(g, np.inf) if res < self.params.newton_tol: break # assemble dg dg = sp.eye(self.params.nvars) - factor * \ (self.A[1:-1, 1:-1] + sp.diags(lambda0 ** 2 - lambda0 ** 2 * (nu + 1) * u.values ** nu, offsets=0)) # newton update: u1 = u0 - g/dg u.values -= spsolve(dg, g) # increase iteration count n += 1 if n == self.params.newton_maxiter: raise ProblemError( 'Newton did not converge after %i iterations, error is %s' % (n, res)) return u
def solve_system(self, rhs, factor, u0, t): """ Simple Newton solver Args: rhs (dtype_f): right-hand side for the nonlinear system factor (float): abbrev. for the node-to-node stepsize (or any other factor required) u0 (dtype_u): initial guess for the iterative solver t (float): current time (required here for the BC) Returns: dtype_u: solution u """ u = self.dtype_u(u0).values eps2 = self.params.eps**2 dw = self.params.dw Id = sp.eye(self.params.nvars) v = 3.0 * np.sqrt(2) * self.params.eps * self.params.dw self.uext.values[0] = 0.5 * (1 + np.tanh( (self.params.interval[0] - v * t) / (np.sqrt(2) * self.params.eps))) self.uext.values[-1] = 0.5 * (1 + np.tanh( (self.params.interval[1] - v * t) / (np.sqrt(2) * self.params.eps))) A = self.A[1:-1, 1:-1] # start newton iteration n = 0 res = 99 while n < self.params.newton_maxiter: # print(n) # form the function g with g(u) = 0 self.uext.values[1:-1] = u[:] g = u - rhs.values \ - factor * (self.A.dot(self.uext.values)[1:-1] - 2.0 / eps2 * u * (1.0 - u) * (1.0 - 2.0 * u) - 6.0 * dw * u * (1.0 - u)) # if g is close to 0, then we are done res = np.linalg.norm(g, np.inf) if res < self.params.newton_tol: break # assemble dg dg = Id - factor * (A - 2.0 / eps2 * sp.diags( (1.0 - u) * (1.0 - 2.0 * u) - u * ((1.0 - 2.0 * u) + 2.0 * (1.0 - u)), offsets=0) - 6.0 * dw * sp.diags((1.0 - u) - u, offsets=0)) # newton update: u1 = u0 - g/dg u -= spsolve(dg, g) # u -= gmres(dg, g, x0=z, tol=self.params.lin_tol)[0] # increase iteration count n += 1 if np.isnan(res) and self.params.stop_at_nan: raise ProblemError( 'Newton got nan after %i iterations, aborting...' % n) elif np.isnan(res): self.logger.warning('Newton got nan after %i iterations...' % n) if n == self.params.newton_maxiter: self.logger.warning( 'Newton did not converge after %i iterations, error is %s' % (n, res)) self.newton_ncalls += 1 self.newton_itercount += n me = self.dtype_u(self.init) me.values = u return me
def __get_A(N, c, dx, order, type): """ Helper function to assemble FD matrix A in sparse format Args: N (int): number of dofs c (float): diffusion coefficient dx (float): distance between two spatial nodes order (int): specifies order of discretization type (string): upwind or centered differences Returns: scipy.sparse.csc_matrix: matrix A in CSC format """ coeff = None stencil = None zero_pos = None if type == 'center': if order == 2: stencil = [-1.0, 0.0, 1.0] zero_pos = 2 coeff = 1.0 / 2.0 elif order == 4: stencil = [1.0, -8.0, 0.0, 8.0, -1.0] zero_pos = 3 coeff = 1.0 / 12.0 elif order == 6: stencil = [-1.0, 9.0, -45.0, 0.0, 45.0, -9.0, 1.0] zero_pos = 4 coeff = 1.0 / 60.0 else: raise ProblemError("Order " + str(order) + " not implemented.") else: if order == 1: stencil = [-1.0, 1.0] coeff = 1.0 zero_pos = 2 elif order == 2: stencil = [1.0, -4.0, 3.0] coeff = 1.0 / 2.0 zero_pos = 3 elif order == 3: stencil = [1.0, -6.0, 3.0, 2.0] coeff = 1.0 / 6.0 zero_pos = 3 elif order == 4: stencil = [-5.0, 30.0, -90.0, 50.0, 15.0] coeff = 1.0 / 60.0 zero_pos = 4 elif order == 5: stencil = [3.0, -20.0, 60.0, -120.0, 65.0, 12.0] coeff = 1.0 / 60.0 zero_pos = 5 else: raise ProblemError("Order " + str(order) + " not implemented.") dstencil = np.concatenate((stencil, np.delete(stencil, zero_pos - 1))) offsets = np.concatenate( ([N - i - 1 for i in reversed(range(zero_pos - 1))], [i - zero_pos + 1 for i in range(zero_pos - 1, len(stencil))])) doffsets = np.concatenate( (offsets, np.delete(offsets, zero_pos - 1) - N)) A = sp.diags(dstencil, doffsets, shape=(N, N), format='csc') A *= -c * coeff * (1.0 / dx) return A
def __init__(self, problem_params, dtype_u=parallel_mesh, dtype_f=parallel_imex_mesh): """ Initialization routine Args: problem_params (dict): custom parameters for the example dtype_u: fft data type (will be passed to parent class) dtype_f: fft data type wuth implicit and explicit parts (will be passed to parent class) """ if 'L' not in problem_params: problem_params['L'] = 1.0 if 'init_type' not in problem_params: problem_params['init_type'] = 'circle' if 'comm' not in problem_params: problem_params['comm'] = None if 'dw' not in problem_params: problem_params['dw'] = 0.0 # these parameters will be used later, so assert their existence essential_keys = ['nvars', 'eps', 'L', 'radius', 'dw', 'spectral'] for key in essential_keys: if key not in problem_params: msg = 'need %s to instantiate problem, only got %s' % ( key, str(problem_params.keys())) raise ParameterError(msg) if not (isinstance(problem_params['nvars'], tuple) and len(problem_params['nvars']) > 1): raise ProblemError('Need at least two dimensions') # Creating FFT structure ndim = len(problem_params['nvars']) axes = tuple(range(ndim)) self.fft = PFFT(problem_params['comm'], list(problem_params['nvars']), axes=axes, dtype=np.float, collapse=True) # get test data to figure out type and dimensions tmp_u = newDistArray(self.fft, problem_params['spectral']) # invoke super init, passing the communicator and the local dimensions as init super(allencahn_imex, self).__init__(init=(tmp_u.shape, problem_params['comm'], tmp_u.dtype), dtype_u=dtype_u, dtype_f=dtype_f, params=problem_params) L = np.array([self.params.L] * ndim, dtype=float) # get local mesh X = np.ogrid[self.fft.local_slice(False)] N = self.fft.global_shape() for i in range(len(N)): X[i] = (X[i] * L[i] / N[i]) self.X = [np.broadcast_to(x, self.fft.shape(False)) for x in X] # get local wavenumbers and Laplace operator s = self.fft.local_slice() N = self.fft.global_shape() k = [np.fft.fftfreq(n, 1. / n).astype(int) for n in N[:-1]] k.append(np.fft.rfftfreq(N[-1], 1. / N[-1]).astype(int)) K = [ki[si] for ki, si in zip(k, s)] Ks = np.meshgrid(*K, indexing='ij', sparse=True) Lp = 2 * np.pi / L for i in range(ndim): Ks[i] = (Ks[i] * Lp[i]).astype(float) K = [np.broadcast_to(k, self.fft.shape(True)) for k in Ks] K = np.array(K).astype(float) self.K2 = np.sum(K * K, 0, dtype=float) # Need this for diagnostics self.dx = self.params.L / problem_params['nvars'][0] self.dy = self.params.L / problem_params['nvars'][1]
def __init__(self, problem_params, dtype_u=mesh, dtype_f=mesh): """ Initialization routine Args: problem_params (dict): custom parameters for the example dtype_u: mesh data type (will be passed parent class) dtype_f: mesh data type (will be passed parent class) """ # these parameters will be used later, so assert their existence if 'order' not in problem_params: problem_params['order'] = 2 if 'lintol' not in problem_params: problem_params['lintol'] = 1E-12 if 'liniter' not in problem_params: problem_params['liniter'] = 10000 if 'direct_solver' not in problem_params: problem_params['direct_solver'] = False essential_keys = [ 'nvars', 'c', 'freq', 'type', 'order', 'ndim', 'lintol', 'liniter', 'direct_solver' ] for key in essential_keys: if key not in problem_params: msg = 'need %s to instantiate problem, only got %s' % ( key, str(problem_params.keys())) raise ParameterError(msg) # make sure parameters have the correct form if problem_params['ndim'] > 3: raise ProblemError( f'can work with up to three dimensions, got {problem_params["ndim"]}' ) if type(problem_params['freq']) is not tuple or len( problem_params['freq']) != problem_params['ndim']: raise ProblemError( f'need {problem_params["ndim"]} frequencies, got {problem_params["freq"]}' ) for freq in problem_params['freq']: if freq % 2 != 0: raise ProblemError( 'need even number of frequencies due to periodic BCs') if type(problem_params['nvars']) is not tuple or len( problem_params['nvars']) != problem_params['ndim']: raise ProblemError( f'need {problem_params["ndim"]} nvars, got {problem_params["nvars"]}' ) for nvars in problem_params['nvars']: if nvars % 2 != 0: raise ProblemError( 'the setup requires nvars = 2^p per dimension') if problem_params['nvars'][1:] != problem_params['nvars'][:-1]: raise ProblemError('need a square domain, got %s' % problem_params['nvars']) # invoke super init, passing number of dofs, dtype_u and dtype_f super(advectionNd_periodic, self).__init__(init=problem_params['nvars'], dtype_u=dtype_u, dtype_f=dtype_f, params=problem_params) # compute dx (equal in both dimensions) and get discretization matrix A self.dx = 1.0 / self.params.nvars[0] self.A = self.__get_A(self.params.nvars, self.params.c, self.dx, self.params.ndim, self.params.type, self.params.order) xvalues = np.array([i * self.dx for i in range(self.params.nvars[0])]) self.xv = np.meshgrid(*[xvalues for _ in range(self.params.ndim)]) self.Id = sp.eye(np.prod(self.params.nvars), format='csc')
def solve_system_2(self, rhs, factor, u0, t): """ Newton-Solver for the second component Args: rhs (dtype_f): right-hand side for the linear system factor (float) : abbrev. for the node-to-node stepsize (or any other factor required) u0 (dtype_u): initial guess for the iterative solver (not used here so far) t (float): current time (e.g. for time-dependent BCs) Returns: dtype_u: solution as mesh """ u = self.dtype_u(u0) if self.params.spectral: tmpu = newDistArray(self.fft, False) tmpv = newDistArray(self.fft, False) tmpu[:] = self.fft.backward(u[..., 0], tmpu) tmpv[:] = self.fft.backward(u[..., 1], tmpv) tmprhsu = newDistArray(self.fft, False) tmprhsv = newDistArray(self.fft, False) tmprhsu[:] = self.fft.backward(rhs[..., 0], tmprhsu) tmprhsv[:] = self.fft.backward(rhs[..., 1], tmprhsv) else: tmpu = u[..., 0] tmpv = u[..., 1] tmprhsu = rhs[..., 0] tmprhsv = rhs[..., 1] # start newton iteration n = 0 res = 99 while n < self.params.newton_maxiter: # print(n, res) # form the function g with g(u) = 0 tmpgu = tmpu - tmprhsu - factor * (-tmpu * tmpv**2 + self.params.A) tmpgv = tmpv - tmprhsv - factor * (tmpu * tmpv**2) # if g is close to 0, then we are done res = max(np.linalg.norm(tmpgu, np.inf), np.linalg.norm(tmpgv, np.inf)) if res < self.params.newton_tol: break # assemble dg dg00 = 1 - factor * (-tmpv**2) dg01 = -factor * (-2 * tmpu * tmpv) dg10 = -factor * (tmpv**2) dg11 = 1 - factor * (2 * tmpu * tmpv) # interleave and unravel to put into sparse matrix dg00I = np.ravel(np.kron(dg00, np.array([1, 0]))) dg01I = np.ravel(np.kron(dg01, np.array([1, 0]))) dg10I = np.ravel(np.kron(dg10, np.array([1, 0]))) dg11I = np.ravel(np.kron(dg11, np.array([0, 1]))) # put into sparse matrix dg = sp.diags(dg00I, offsets=0) + sp.diags(dg11I, offsets=0) dg += sp.diags(dg01I, offsets=1, shape=dg.shape) + sp.diags( dg10I, offsets=-1, shape=dg.shape) # interleave g terms to apply inverse to it g = np.kron(tmpgu.flatten(), np.array([1, 0])) + np.kron( tmpgv.flatten(), np.array([0, 1])) # invert dg matrix b = sp.linalg.spsolve(dg, g) # update real-space vectors tmpu[:] -= b[::2].reshape(self.params.nvars) tmpv[:] -= b[1::2].reshape(self.params.nvars) # increase iteration count n += 1 if np.isnan(res) and self.params.stop_at_nan: raise ProblemError( 'Newton got nan after %i iterations, aborting...' % n) elif np.isnan(res): self.logger.warning('Newton got nan after %i iterations...' % n) if n == self.params.newton_maxiter: self.logger.warning( 'Newton did not converge after %i iterations, error is %s' % (n, res)) # self.newton_ncalls += 1 # self.newton_itercount += n me = self.dtype_u(self.init) if self.params.spectral: me[..., 0] = self.fft.forward(tmpu) me[..., 1] = self.fft.forward(tmpv) else: me[..., 0] = tmpu me[..., 1] = tmpv return me
def __get_A(N, c, dx, ndim, type, order): """ Helper function to assemble FD matrix A in sparse format Args: N (list): number of dofs nu (float): diffusion coefficient dx (float): distance between two spatial nodes type (str): disctretization type ndim (int): number of dimensions Returns: scipy.sparse.csc_matrix: matrix A in CSC format """ coeff = None stencil = None zero_pos = None if type == 'center': if order == 2: stencil = [-1.0, 0.0, 1.0] zero_pos = 2 coeff = 1.0 / 2.0 elif order == 4: stencil = [1.0, -8.0, 0.0, 8.0, -1.0] zero_pos = 3 coeff = 1.0 / 12.0 elif order == 6: stencil = [-1.0, 9.0, -45.0, 0.0, 45.0, -9.0, 1.0] zero_pos = 4 coeff = 1.0 / 60.0 else: raise ProblemError("Order " + str(order) + " not implemented.") else: if order == 1: stencil = [-1.0, 1.0] coeff = 1.0 zero_pos = 2 elif order == 2: stencil = [1.0, -4.0, 3.0] coeff = 1.0 / 2.0 zero_pos = 3 elif order == 3: stencil = [1.0, -6.0, 3.0, 2.0] coeff = 1.0 / 6.0 zero_pos = 3 elif order == 4: stencil = [-5.0, 30.0, -90.0, 50.0, 15.0] coeff = 1.0 / 60.0 zero_pos = 4 elif order == 5: stencil = [3.0, -20.0, 60.0, -120.0, 65.0, 12.0] coeff = 1.0 / 60.0 zero_pos = 5 else: raise ProblemError("Order " + str(order) + " not implemented.") dstencil = np.concatenate((stencil, np.delete(stencil, zero_pos - 1))) offsets = np.concatenate( ([N[0] - i - 1 for i in reversed(range(zero_pos - 1))], [i - zero_pos + 1 for i in range(zero_pos - 1, len(stencil))])) doffsets = np.concatenate( (offsets, np.delete(offsets, zero_pos - 1) - N[0])) A = coeff * sp.diags( dstencil, doffsets, shape=(N[0], N[0]), format='csc') if ndim == 2: A = sp.kron(A, sp.eye(N[0])) + sp.kron(sp.eye(N[1]), A) elif ndim == 3: A = sp.kron(A, sp.eye(N[1] * N[0])) + sp.kron(sp.eye(N[2] * N[1]), A) + \ sp.kron(sp.kron(sp.eye(N[2]), A), sp.eye(N[0])) A *= -c * (1.0 / dx) return A