def _create_DA(self): r"""Returns a PETSc DA and associated global Vec. Note that no local vector is returned. """ from petsc4py import PETSc if hasattr(PETSc.DA, 'PeriodicType'): if self.num_dim == 1: periodic_type = PETSc.DA.PeriodicType.X elif self.num_dim == 2: periodic_type = PETSc.DA.PeriodicType.XY elif self.num_dim == 3: periodic_type = PETSc.DA.PeriodicType.XYZ else: raise Exception("Invalid number of dimensions") DA = PETSc.DA().create(dim=self.num_dim, dof=1, sizes=self.num_cells, periodic_type=periodic_type, stencil_width=0, comm=PETSc.COMM_WORLD) else: DA = PETSc.DA().create( dim=self.num_dim, dof=1, sizes=self.num_cells, boundary_type=PETSc.DA.BoundaryType.PERIODIC, stencil_width=0, comm=PETSc.COMM_WORLD) return DA
def advection1D(self, M, cfl, T): '''Script to solve 1D advection equation: q_t + q_x = 0 Using first-order finite differences''' da = PETSc.DA().create([M]) da.setUniformCoordinates() # solves the problem from 0 to 1 da.view() xvec = da.getCoordinates() xvec.view() x = xvec.getArray() h = x[1] - x[0] k = cfl * h fg = da.createGlobalVector() fl = da.createLocalVector() N = int(round(T / k)) # Initial condition: q = np.exp(-10 * (x - 0.5)**2) fg.setArray(q) da.globalToLocal(fg, fl) fg.view() for n in range(N + 1): q = fl.getArray() q[1:] = q[1:] - cfl * (q[1:] - q[:-1]) fl.setArray(q) da.localToGlobal(fl, fg) fg.view() da.globalToLocal(fg, fl)
def solvePoisson(nx, ny): # Create the DA. da = PETSc.DA().create([nx, ny], stencil_width=1, boundary_type=('ghosted', 'ghosted')) pde = Poisson2D(da) b = pde.formRHS() pde.setupSolver() print "Solving using", pde.solver_name # Solve! tic = PETSc.Log().getTime() pde.solve(b, pde.x) toc = PETSc.Log().getTime() time = toc - tic timePerGrid = time / nx / ny # output performance/convergence information its = pde.ksp.getIterationNumber() rnorm = pde.ksp.getResidualNorm() print "Nx Ny its ||r||_2 ElapsedTime (s) Elapsed Time/N_tot" print nx, ny, its, rnorm, time, timePerGrid rh = pde.ksp.getConvergenceHistory() filename = '{0}_{1}x{2}.npz'.format(pde.solver_name, nx, ny) np.savez(filename, rhist=rh, npts=np.array([nx, ny]), name=pde.solver_name, time=time) return pde, time, timePerGrid
def advection1D(self, N): '''David: If you put in the linear algebra that you need (comments), I will move it to code''' da = PETSc.DA().create([N]) da.view() f = da.createGlobalVector() f.view() a = f.getArray() for i in range(f.getSize()): a[i] = i*i f.view() return
def init(f): da = PETSc.DA().create(comm=f.comm, dim=3, dof=d + 1, sizes=N) rhoVec = PETSc.Vec() rhoVec.createSeq(comm=f.comm, size=N[0] * N[1] * N[2] * d) rho = rhoVec.getArray().reshape((N[0], N[1], N[2], d)) from math import sin as sin, pi as pi for i in range(N[0]): for j in range(N[1]): for k in range(N[2]): for s in range(d): rho[i, j, k, s] = sin(2 * pi * i / N[0]) * sin( 2 * pi * j / N[1]) * sin(2 * pi * k / N[2]) f.compose("mesh", da) f.compose("rho", rhoVec)
def _create_DA(self, dof, num_ghost=0): r"""Returns a PETSc DA and associated global Vec. Note that no local vector is returned. """ from petsc4py import PETSc #Due to the way PETSc works, we just make the patch always periodic, #regardless of the boundary conditions actually selected. #This works because in solver.qbc() we first call globalToLocal() #and then impose the real boundary conditions (if non-periodic). if hasattr(PETSc.DA, 'PeriodicType'): if self.num_dim == 1: periodic_type = PETSc.DA.PeriodicType.X elif self.num_dim == 2: periodic_type = PETSc.DA.PeriodicType.XY elif self.num_dim == 3: periodic_type = PETSc.DA.PeriodicType.XYZ else: raise Exception("Invalid number of dimensions") DA = PETSc.DA().create(dim=self.num_dim, dof=dof, sizes=self.patch.num_cells_global, periodic_type=periodic_type, stencil_width=num_ghost, comm=PETSc.COMM_WORLD) else: DA = PETSc.DA().create( dim=self.num_dim, dof=dof, sizes=self.patch.num_cells_global, boundary_type=PETSc.DA.BoundaryType.PERIODIC, stencil_width=num_ghost, comm=PETSc.COMM_WORLD) return DA
def help(args=None): import sys, shlex # program name try: prog = sys.argv[0] except Exception: prog = getattr(sys, 'executable', 'python') if args is None: args = sys.argv[1:] elif isinstance(args, str): args = shlex.split(args) else: args = [str(a) for a in args] import petsc4py petsc4py.init([prog, '-help'] + args) from petsc4py import PETSc COMM = PETSc.COMM_SELF if 'vec' in args: vec = PETSc.Vec().create(comm=COMM) vec.setSizes(0) vec.setFromOptions() del vec if 'mat' in args: mat = PETSc.Mat().create(comm=COMM) mat.setSizes([0, 0]) mat.setFromOptions() del mat if 'ksp' in args: ksp = PETSc.KSP().create(comm=COMM) ksp.setFromOptions() del ksp if 'pc' in args: pc = PETSc.PC().create(comm=COMM) pc.setFromOptions() del pc if 'snes' in args: snes = PETSc.SNES().create(comm=COMM) snes.setFromOptions() del snes if 'ts' in args: ts = PETSc.TS().create(comm=COMM) ts.setFromOptions() del ts if 'da' in args: da = PETSc.DA().create(comm=COMM) da.setFromOptions() del da
def solveDiffusion(nx, ny, inputd): # Create the DA and setup the problem if inputd.kspType == 'sd': U = 0.0 self = 0.0 [M, R, row, col, val] = build_A(inputd, self) tic = PETSc.Log().getTime() U = dsolve.spsolve(M, R, use_umfpack=True) #Direct solver toc = PETSc.Log().getTime() its = 1.0 rnorm = np.linalg.norm(M * U - R) pde = 0.0 else: da = PETSc.DA().create([nx, ny], stencil_width=1, boundary_type=('ghosted', 'ghosted')) pde = diffusion2D(da) pde.setupSolver(inputd) print "Solving using", pde.solver_name #Now solve tic = PETSc.Log().getTime() pde.solve(pde.b, pde.x) toc = PETSc.Log().getTime() its = pde.ksp.getIterationNumber() rnorm = pde.ksp.getResidualNorm() M = 0.0 U = M # output performance/convergence information time = toc - tic timePerGrid = time / nx / ny #rnorm = np.linalg.norm(M*U-R) print "Nx Ny its ||r||_2 ElapsedTime (s) Elapsed Time/N_tot" print nx, ny, its, rnorm, time, timePerGrid return pde, time, timePerGrid, its, rnorm, U, M
import sys, petsc4py petsc4py.init(sys.argv) from petsc4py import PETSc import Bratu3D as Bratu3D OptDB = PETSc.Options() N = OptDB.getInt('N', 16) lambda_ = OptDB.getReal('lambda', 6.0) do_plot = OptDB.getBool('plot', False) da = PETSc.DA().create([N, N, N], stencil_width=1) #app = App(da, lambda_) snes = PETSc.SNES().create() F = da.createGlobalVec() snes.setFunction(Bratu3D.formFunction, F, args=(da, lambda_)) J = da.createMat() snes.setJacobian(Bratu3D.formJacobian, J, args=(da, lambda_)) snes.setFromOptions() X = da.createGlobalVec() Bratu3D.formInitGuess(X, da, lambda_) snes.solve(None, X) U = da.createNaturalVec() da.globalToNatural(X, U)
def __init__(self, cfgfile): ''' Constructor ''' # stencil = 1 stencil = 2 # load run config file cfg = Config(cfgfile) # set some PETSc options OptDB = PETSc.Options() # OptDB.setValue('snes_lag_preconditioner', 5) OptDB.setValue('snes_atol', cfg['solver']['petsc_residual']) OptDB.setValue('snes_rtol', 1E-16) OptDB.setValue('snes_stol', 1E-18) OptDB.setValue('snes_max_it', 10) OptDB.setValue('ksp_atol', cfg['solver']['petsc_residual'] * 1E-1) OptDB.setValue('ksp_rtol', 1E-10) OptDB.setValue('ksp_max_it', 10) # OptDB.setValue('ksp_convergence_test', 'skip') OptDB.setValue('snes_monitor', '') OptDB.setValue('ksp_monitor', '') # timestep setup self.ht = cfg['grid']['ht'] # timestep size self.nt = cfg['grid']['nt'] # number of timesteps self.nsave = cfg['io']['nsave'] # save only every nsave'th timestep # grid setup nx = cfg['grid']['nx'] # number of points in x ny = cfg['grid']['ny'] # number of points in y Lx = cfg['grid']['Lx'] # spatial domain in x x1 = cfg['grid']['x1'] # x2 = cfg['grid']['x2'] # Ly = cfg['grid']['Ly'] # spatial domain in y y1 = cfg['grid']['y1'] # y2 = cfg['grid']['y2'] # self.nx = nx self.ny = ny if x1 != x2: Lx = x2-x1 else: x1 = 0.0 x2 = Lx if y1 != y2: Ly = y2-y1 else: y1 = 0.0 y2 = Ly self.hx = Lx / nx # gridstep size in x self.hy = Ly / ny # gridstep size in y if PETSc.COMM_WORLD.getRank() == 0: print() print("nt = %i" % (self.nt)) print("nx = %i" % (self.nx)) print("ny = %i" % (self.ny)) print() print("ht = %e" % (self.ht)) print("hx = %e" % (self.hx)) print("hy = %e" % (self.hy)) print() print("Lx = %e" % (Lx)) print("Ly = %e" % (Ly)) print() self.time = PETSc.Vec().createMPI(1, PETSc.DECIDE, comm=PETSc.COMM_WORLD) self.time.setName('t') if PETSc.COMM_WORLD.getRank() == 0: self.time.setValue(0, 0.0) # create DA with single dof self.da1 = PETSc.DA().create(dim=2, dof=1, sizes=[nx, ny], proc_sizes=[PETSc.DECIDE, PETSc.DECIDE], boundary_type=('periodic', 'periodic'), stencil_width=stencil, stencil_type='box') # create DA (dof = 4 for Bx, By, Vx, Vy) self.da4 = PETSc.DA().create(dim=2, dof=4, sizes=[nx, ny], proc_sizes=[PETSc.DECIDE, PETSc.DECIDE], boundary_type=('periodic', 'periodic'), stencil_width=stencil, stencil_type='box') # create DA (dof = 5 for Bx, By, Vx, Vy, P) self.da5 = PETSc.DA().create(dim=2, dof=5, sizes=[nx, ny], proc_sizes=[PETSc.DECIDE, PETSc.DECIDE], boundary_type=('periodic', 'periodic'), stencil_width=stencil, stencil_type='box') # create DA for x grid self.dax = PETSc.DA().create(dim=1, dof=1, sizes=[nx], proc_sizes=[PETSc.DECIDE], boundary_type=('periodic')) # create DA for y grid self.day = PETSc.DA().create(dim=1, dof=1, sizes=[ny], proc_sizes=[PETSc.DECIDE], boundary_type=('periodic')) # initialise grid self.da1.setUniformCoordinates(xmin=x1, xmax=x2, ymin=y1, ymax=y2) self.da4.setUniformCoordinates(xmin=x1, xmax=x2, ymin=y1, ymax=y2) self.da5.setUniformCoordinates(xmin=x1, xmax=x2, ymin=y1, ymax=y2) self.dax.setUniformCoordinates(xmin=x1, xmax=x2) self.day.setUniformCoordinates(xmin=y1, xmax=y2) # create solution and RHS vector self.f = self.da5.createGlobalVec() self.x = self.da5.createGlobalVec() self.b = self.da5.createGlobalVec() # create global RK4 vectors self.Y = self.da5.createGlobalVec() self.X0 = self.da5.createGlobalVec() self.X1 = self.da5.createGlobalVec() self.X2 = self.da5.createGlobalVec() self.X3 = self.da5.createGlobalVec() self.X4 = self.da5.createGlobalVec() # create local RK4 vectors self.localX0 = self.da5.createLocalVec() self.localX1 = self.da5.createLocalVec() self.localX2 = self.da5.createLocalVec() self.localX3 = self.da5.createLocalVec() self.localX4 = self.da5.createLocalVec() # self.localP = self.da1.createLocalVec() # create vectors for magnetic and velocity field self.Bx = self.da1.createGlobalVec() self.By = self.da1.createGlobalVec() self.Vx = self.da1.createGlobalVec() self.Vy = self.da1.createGlobalVec() self.P = self.da1.createGlobalVec() # create local vectors for initialisation of pressure self.localBx = self.da1.createLocalVec() self.localBy = self.da1.createLocalVec() self.localVx = self.da1.createLocalVec() self.localVy = self.da1.createLocalVec() # set variable names self.Bx.setName('Bx') self.By.setName('By') self.Vx.setName('Vx') self.Vy.setName('Vy') self.P.setName('P') # create Matrix object self.petsc_matrix = PETScMatrix (self.da1, self.da5, nx, ny, self.ht, self.hx, self.hy) self.petsc_jacobian = PETScJacobian(self.da1, self.da5, nx, ny, self.ht, self.hx, self.hy) self.petsc_function = PETScFunction(self.da1, self.da5, nx, ny, self.ht, self.hx, self.hy) # self.petsc_jacobian_4d = PETSc_MHD_NL_Jacobian_Matrix.PETScJacobian(self.da1, self.da5, nx, ny, self.ht, self.hx, self.hy) # initialise matrix self.A = self.da5.createMat() self.A.setOption(self.A.Option.NEW_NONZERO_ALLOCATION_ERR, False) self.A.setUp() # create jacobian self.J = self.da5.createMat() self.J.setOption(self.J.Option.NEW_NONZERO_ALLOCATION_ERR, False) self.J.setUp() # create nonlinear solver self.snes = PETSc.SNES().create() self.snes.setFunction(self.petsc_function.snes_mult, self.f) self.snes.setJacobian(self.updateJacobian, self.J) self.snes.setFromOptions() self.snes.getKSP().setType('preonly') # self.snes.getKSP().setType('gmres') self.snes.getKSP().getPC().setType('lu') # self.snes.getKSP().getPC().setFactorSolverPackage('superlu_dist') self.snes.getKSP().getPC().setFactorSolverPackage('mumps') self.ksp = None # set initial data (xs, xe), (ys, ye) = self.da1.getRanges() coords = self.da1.getCoordinatesLocal() Bx_arr = self.da1.getVecArray(self.Bx) By_arr = self.da1.getVecArray(self.By) Vx_arr = self.da1.getVecArray(self.Vx) Vy_arr = self.da1.getVecArray(self.Vy) if cfg['initial_data']['magnetic_python'] != None: init_data = __import__("runs." + cfg['initial_data']['magnetic_python'], globals(), locals(), ['magnetic_x', 'magnetic_y'], 0) for i in range(xs, xe): for j in range(ys, ye): Bx_arr[i,j] = init_data.magnetic_x(coords[i,j][0], coords[i,j][1] + 0.5 * self.hy, Lx, Ly) By_arr[i,j] = init_data.magnetic_y(coords[i,j][0] + 0.5 * self.hx, coords[i,j][1], Lx, Ly) else: Bx_arr[xs:xe, ys:ye] = cfg['initial_data']['magnetic'] By_arr[xs:xe, ys:ye] = cfg['initial_data']['magnetic'] if cfg['initial_data']['velocity_python'] != None: init_data = __import__("runs." + cfg['initial_data']['velocity_python'], globals(), locals(), ['velocity_x', 'velocity_y'], 0) for i in range(xs, xe): for j in range(ys, ye): Vx_arr[i,j] = init_data.velocity_x(coords[i,j][0], coords[i,j][1] + 0.5 * self.hy, Lx, Ly) Vy_arr[i,j] = init_data.velocity_y(coords[i,j][0] + 0.5 * self.hx, coords[i,j][1], Lx, Ly) else: Vx_arr[xs:xe, ys:ye] = cfg['initial_data']['velocity'] Vy_arr[xs:xe, ys:ye] = cfg['initial_data']['velocity'] if cfg['initial_data']['pressure_python'] != None: init_data = __import__("runs." + cfg['initial_data']['pressure_python'], globals(), locals(), ['pressure', ''], 0) x_arr = self.da5.getVecArray(self.x) x_arr[xs:xe, ys:ye, 0] = Vx_arr[xs:xe, ys:ye] x_arr[xs:xe, ys:ye, 1] = Vy_arr[xs:xe, ys:ye] x_arr[xs:xe, ys:ye, 2] = Bx_arr[xs:xe, ys:ye] x_arr[xs:xe, ys:ye, 3] = By_arr[xs:xe, ys:ye] self.da1.globalToLocal(self.Bx, self.localBx) self.da1.globalToLocal(self.By, self.localBy) self.da1.globalToLocal(self.Vx, self.localVx) self.da1.globalToLocal(self.Vy, self.localVy) Bx_arr = self.da1.getVecArray(self.localBx) By_arr = self.da1.getVecArray(self.localBy) Vx_arr = self.da1.getVecArray(self.localVx) Vy_arr = self.da1.getVecArray(self.localVy) P_arr = self.da1.getVecArray(self.P) for i in range(xs, xe): for j in range(ys, ye): P_arr[i,j] = init_data.pressure(coords[i,j][0] + 0.5 * self.hx, coords[i,j][1] + 0.5 * self.hy, Lx, Ly) # P_arr[i,j] = init_data.pressure(coords[i,j][0] + 0.5 * self.hx, coords[i,j][1] + 0.5 * self.hy, Lx, Ly) \ # + 0.5 * (0.25 * (Bx_arr[i,j] + Bx_arr[i+1,j])**2 + 0.25 * (By_arr[i,j] + By_arr[i,j+1])**2) # P_arr[i,j] = init_data.pressure(coords[i,j][0] + 0.5 * self.hx, coords[i,j][1] + 0.5 * self.hy, Lx, Ly) \ # + 0.5 * (0.25 * (Vx_arr[i,j] + Vx_arr[i+1,j])**2 + 0.25 * (Vy_arr[i,j] + Vy_arr[i,j+1])**2) # P_arr[i,j] = init_data.pressure(coords[i,j][0] + 0.5 * self.hx, coords[i,j][1] + 0.5 * self.hy, Lx, Ly) \ # + 0.5 * (0.25 * (Vx_arr[i,j] + Vx_arr[i+1,j])**2 + 0.25 * (Vy_arr[i,j] + Vy_arr[i,j+1])**2) \ # - 1.0 * (0.25 * (Bx_arr[i,j] + Bx_arr[i+1,j])**2 + 0.25 * (By_arr[i,j] + By_arr[i,j+1])**2) # copy distribution function to solution vector x_arr = self.da5.getVecArray(self.x) x_arr[xs:xe, ys:ye, 4] = P_arr [xs:xe, ys:ye] # update solution history self.petsc_matrix.update_history(self.x) self.petsc_jacobian.update_history(self.x) self.petsc_function.update_history(self.x) # create HDF5 output file self.hdf5_viewer = PETSc.ViewerHDF5().create(cfg['io']['hdf5_output'], mode=PETSc.Viewer.Mode.WRITE, comm=PETSc.COMM_WORLD) self.hdf5_viewer.pushGroup("/") # write grid data to hdf5 file coords_x = self.dax.getCoordinates() coords_y = self.day.getCoordinates() coords_x.setName('x') coords_y.setName('y') self.hdf5_viewer(coords_x) self.hdf5_viewer(coords_y) # write initial data to hdf5 file self.hdf5_viewer.setTimestep(0) self.save_hdf5_vectors()
def solver_FE_MATfree_PETSc(I, a, L, Nx, C, T): """ This solves the system using a Forward Euler scheme, but without using a explisit matrix. This calls advance-function to do the calculations at each time step. This saves memory (no need to store a matrix in the memory) and time (dont use any time setting up a matrix). For now we can call to advancer-functions (could be extended using Cython and f2py for speed): - advance_FE_looper: a plain scalar looper. - advance_FE_fastnumpy: a vectorized looper using Numpy. Note that in this function there is no hardcopying in the conversion between PETSc arrays and Numpy arrays, as they share the same memory location. This saves time during the calculations. """ x = np.linspace(0, L, Nx + 1) dx = x[1] - x[0] dt = C * dx**2 / a Nt = int(round(T / float(dt))) t = np.linspace(0, T, Nt + 1) t0 = time.clock() # COMMENT da = PETSc.DA().create(sizes=[Nx + 1], boundary_type=2, stencil_type=0, stencil_width=1) da.setUniformCoordinates(0, L) # We never really use this in here u = da.createGlobalVector() u_1 = da.createGlobalVector() #A = da.createMatrix(); A.setType('aij') # type is seqaij as default, must change this # Initialize init time step [Istart, Iend] = u_1.getOwnershipRange() u_1.setValues(range(Istart, Iend), I[Istart:Iend]) local_vec = da.createLocalVector() local_vec_new = da.createLocalVector() # Assemble the vectors and the matrix. This distribute the objects out over the procs. #A.assemble(); u.assemble() u_1.assemble() # This can be cut out, this is for plotting in external # programs if PETSc.Options().getBool('toFile', default=False): if PETSc.COMM_WORLD.getSize() > 1: if PETSc.COMM_WORLD.getRank() == 0: print 'Warning: NOT writing to file (using parallel)' PETSc.Options().setValue('toFile', False) else: W = PETSc.Viewer().createASCII('test3.txt', format=1) u_1.view(W) for n in range(0, Nt): da.globalToLocal(u_1, local_vec) if PETSc.Options().getString('advance_method', default='fastnumpy') == 'looper': local_vec_new = advance_FE_looper(local_vec, C, da) else: local_vec_new = advance_FE_fastnumpy(local_vec, C, da) da.localToGlobal(local_vec_new, u_1) if PETSc.Options().getBool('draw', default=False): U = da.createNaturalVector() da.globalToNatural(u_1, U) petsc_viz(U, 0.001) return u_1, time.clock() - t0
def __init__(self, cfgfile): ''' Constructor ''' # stencil = 1 stencil = 2 # load run config file cfg = Config(cfgfile) cfg.set_petsc_options() # set some PETSc options # OptDB = PETSc.Options() # # OptDB.setValue('snes_lag_preconditioner', 5) # self.snes.getKSP().getPC().setReusePreconditioner(True) # timestep setup self.ht = cfg['grid']['ht'] # timestep size self.nt = cfg['grid']['nt'] # number of timesteps self.nsave = cfg['io']['nsave'] # save only every nsave'th timestep # grid setup nx = cfg['grid']['nx'] # number of points in x ny = cfg['grid']['ny'] # number of points in y Lx = cfg['grid']['Lx'] # spatial domain in x x1 = cfg['grid']['x1'] # x2 = cfg['grid']['x2'] # Ly = cfg['grid']['Ly'] # spatial domain in y y1 = cfg['grid']['y1'] # y2 = cfg['grid']['y2'] # self.nx = nx self.ny = ny if x1 != x2: Lx = x2-x1 else: x1 = 0.0 x2 = Lx if y1 != y2: Ly = y2-y1 else: y1 = 0.0 y2 = Ly self.hx = Lx / nx # gridstep size in x self.hy = Ly / ny # gridstep size in y # friction, viscosity and resistivity mu = cfg['initial_data']['mu'] # friction nu = cfg['initial_data']['nu'] # viscosity eta = cfg['initial_data']['eta'] # resistivity # self.update_jacobian = True if PETSc.COMM_WORLD.getRank() == 0: print() print("nt = %i" % (self.nt)) print("nx = %i" % (self.nx)) print("ny = %i" % (self.ny)) print() print("ht = %e" % (self.ht)) print("hx = %e" % (self.hx)) print("hy = %e" % (self.hy)) print() print("Lx = %e" % (Lx)) print("Ly = %e" % (Ly)) print() print("mu = %e" % (mu)) print("nu = %e" % (nu)) print("eta = %e" % (eta)) print() self.time = PETSc.Vec().createMPI(1, PETSc.DECIDE, comm=PETSc.COMM_WORLD) self.time.setName('t') if PETSc.COMM_WORLD.getRank() == 0: self.time.setValue(0, 0.0) # create DA with single dof self.da1 = PETSc.DA().create(dim=2, dof=1, sizes=[nx, ny], proc_sizes=[PETSc.DECIDE, PETSc.DECIDE], boundary_type=('periodic', 'periodic'), stencil_width=stencil, stencil_type='box') # create DA (dof = 4 for Bx, By, Vx, Vy) self.da4 = PETSc.DA().create(dim=2, dof=4, sizes=[nx, ny], proc_sizes=[PETSc.DECIDE, PETSc.DECIDE], boundary_type=('periodic', 'periodic'), stencil_width=stencil, stencil_type='box') # create DA (dof = 5 for Bx, By, Vx, Vy, P) self.da5 = PETSc.DA().create(dim=2, dof=5, sizes=[nx, ny], proc_sizes=[PETSc.DECIDE, PETSc.DECIDE], boundary_type=('periodic', 'periodic'), stencil_width=stencil, stencil_type='box') # create DA for x grid self.dax = PETSc.DA().create(dim=1, dof=1, sizes=[nx], proc_sizes=[PETSc.DECIDE], boundary_type=('periodic')) # create DA for y grid self.day = PETSc.DA().create(dim=1, dof=1, sizes=[ny], proc_sizes=[PETSc.DECIDE], boundary_type=('periodic')) # initialise grid self.da1.setUniformCoordinates(xmin=x1, xmax=x2, ymin=y1, ymax=y2) self.da4.setUniformCoordinates(xmin=x1, xmax=x2, ymin=y1, ymax=y2) self.da5.setUniformCoordinates(xmin=x1, xmax=x2, ymin=y1, ymax=y2) self.dax.setUniformCoordinates(xmin=x1, xmax=x2) self.day.setUniformCoordinates(xmin=y1, xmax=y2) # create solution and RHS vector self.f = self.da5.createGlobalVec() self.x = self.da5.createGlobalVec() self.b = self.da5.createGlobalVec() # create global RK4 vectors self.Y = self.da5.createGlobalVec() self.X0 = self.da5.createGlobalVec() self.X1 = self.da5.createGlobalVec() self.X2 = self.da5.createGlobalVec() self.X3 = self.da5.createGlobalVec() self.X4 = self.da5.createGlobalVec() # create local RK4 vectors self.localX0 = self.da5.createLocalVec() self.localX1 = self.da5.createLocalVec() self.localX2 = self.da5.createLocalVec() self.localX3 = self.da5.createLocalVec() self.localX4 = self.da5.createLocalVec() # self.localP = self.da1.createLocalVec() # create vectors for magnetic and velocity field self.Bx = self.da1.createGlobalVec() self.By = self.da1.createGlobalVec() self.Vx = self.da1.createGlobalVec() self.Vy = self.da1.createGlobalVec() self.P = self.da1.createGlobalVec() self.xcoords = self.da1.createGlobalVec() self.ycoords = self.da1.createGlobalVec() # create local vectors for initialisation of pressure self.localBx = self.da1.createLocalVec() self.localBy = self.da1.createLocalVec() self.localVx = self.da1.createLocalVec() self.localVy = self.da1.createLocalVec() # set variable names self.Bx.setName('Bx') self.By.setName('By') self.Vx.setName('Vx') self.Vy.setName('Vy') self.P.setName('P') # create Matrix object # self.petsc_matrix = PETScMatrix (self.da1, self.da5, nx, ny, self.ht, self.hx, self.hy, mu, nu, eta) self.petsc_function = PETScFunction(self.da1, self.da5, nx, ny, self.ht, self.hx, self.hy, mu, nu, eta) # initialise matrix self.A = self.da5.createMat() self.A.setOption(self.A.Option.NEW_NONZERO_ALLOCATION_ERR, False) self.A.setUp() # create jacobian self.J = self.da5.createMat() self.J.setOption(self.J.Option.NEW_NONZERO_ALLOCATION_ERR, False) self.J.setUp() # create nonlinear solver self.snes = PETSc.SNES().create() self.snes.setFunction(self.petsc_function.snes_mult, self.f) self.snes.setJacobian(self.updateJacobian, self.J) self.snes.setFromOptions() # self.snes.getKSP().setInitialGuessNonzero(True) # self.snes.getKSP().getPC().setReusePreconditioner(True) self.ksp = None # set initial data (xs, xe), (ys, ye) = self.da1.getRanges() # coords = self.da1.getCoordinatesLocal() xc_arr = self.da1.getVecArray(self.xcoords) yc_arr = self.da1.getVecArray(self.ycoords) for i in range(xs, xe): for j in range(ys, ye): xc_arr[i,j] = x1 + i*self.hx yc_arr[i,j] = y1 + j*self.hy Bx_arr = self.da1.getVecArray(self.Bx) By_arr = self.da1.getVecArray(self.By) Vx_arr = self.da1.getVecArray(self.Vx) Vy_arr = self.da1.getVecArray(self.Vy) xc_arr = self.da1.getVecArray(self.xcoords) yc_arr = self.da1.getVecArray(self.ycoords) if cfg['initial_data']['magnetic_python'] != None: init_data = __import__("examples." + cfg['initial_data']['magnetic_python'], globals(), locals(), ['magnetic_x', 'magnetic_y'], 0) for i in range(xs, xe): for j in range(ys, ye): Bx_arr[i,j] = init_data.magnetic_x(xc_arr[i,j], yc_arr[i,j] + 0.5 * self.hy, self.hx, self.hy) By_arr[i,j] = init_data.magnetic_y(xc_arr[i,j] + 0.5 * self.hx, yc_arr[i,j], self.hx, self.hy) else: Bx_arr[xs:xe, ys:ye] = cfg['initial_data']['magnetic'] By_arr[xs:xe, ys:ye] = cfg['initial_data']['magnetic'] if cfg['initial_data']['velocity_python'] != None: init_data = __import__("examples." + cfg['initial_data']['velocity_python'], globals(), locals(), ['velocity_x', 'velocity_y'], 0) for i in range(xs, xe): for j in range(ys, ye): Vx_arr[i,j] = init_data.velocity_x(xc_arr[i,j], yc_arr[i,j] + 0.5 * self.hy, self.hx, self.hy) Vy_arr[i,j] = init_data.velocity_y(xc_arr[i,j] + 0.5 * self.hx, yc_arr[i,j], self.hx, self.hy) else: Vx_arr[xs:xe, ys:ye] = cfg['initial_data']['velocity'] Vy_arr[xs:xe, ys:ye] = cfg['initial_data']['velocity'] if cfg['initial_data']['pressure_python'] != None: init_data = __import__("examples." + cfg['initial_data']['pressure_python'], globals(), locals(), ['pressure', ''], 0) # Fourier Filtering from scipy.fftpack import rfft, irfft nfourier_x = cfg['initial_data']['nfourier_Bx'] nfourier_y = cfg['initial_data']['nfourier_By'] if nfourier_x >= 0 or nfourier_y >= 0: print("Fourier Filtering B") # obtain whole Bx vector everywhere scatter, Xglobal = PETSc.Scatter.toAll(self.Bx) scatter.begin(self.Bx, Xglobal, PETSc.InsertMode.INSERT, PETSc.ScatterMode.FORWARD) scatter.end (self.Bx, Xglobal, PETSc.InsertMode.INSERT, PETSc.ScatterMode.FORWARD) petsc_indices = self.da1.getAO().app2petsc(np.arange(self.nx*self.ny, dtype=np.int32)) BxTmp = Xglobal.getValues(petsc_indices).copy().reshape((self.ny, self.nx)).T scatter.destroy() Xglobal.destroy() # obtain whole By vector everywhere scatter, Xglobal = PETSc.Scatter.toAll(self.By) scatter.begin(self.By, Xglobal, PETSc.InsertMode.INSERT, PETSc.ScatterMode.FORWARD) scatter.end (self.By, Xglobal, PETSc.InsertMode.INSERT, PETSc.ScatterMode.FORWARD) petsc_indices = self.da1.getAO().app2petsc(np.arange(self.nx*self.ny, dtype=np.int32)) ByTmp = Xglobal.getValues(petsc_indices).copy().reshape((self.ny, self.nx)).T scatter.destroy() Xglobal.destroy() if nfourier_x >= 0: # compute FFT, cut, compute inverse FFT BxFft = rfft(BxTmp, axis=1) ByFft = rfft(ByTmp, axis=1) BxFft[:,nfourier_x+1:] = 0. ByFft[:,nfourier_x+1:] = 0. BxTmp = irfft(BxFft, axis=1) ByTmp = irfft(ByFft, axis=1) if nfourier_y >= 0: BxFft = rfft(BxTmp, axis=0) ByFft = rfft(ByTmp, axis=0) BxFft[nfourier_y+1:,:] = 0. ByFft[nfourier_y+1:,:] = 0. BxTmp = irfft(BxFft, axis=0) ByTmp = irfft(ByFft, axis=0) Bx_arr = self.da1.getVecArray(self.Bx) By_arr = self.da1.getVecArray(self.By) Bx_arr[:,:] = BxTmp[xs:xe, ys:ye] By_arr[:,:] = ByTmp[xs:xe, ys:ye] x_arr = self.da5.getVecArray(self.x) x_arr[xs:xe, ys:ye, 0] = Vx_arr[xs:xe, ys:ye] x_arr[xs:xe, ys:ye, 1] = Vy_arr[xs:xe, ys:ye] x_arr[xs:xe, ys:ye, 2] = Bx_arr[xs:xe, ys:ye] x_arr[xs:xe, ys:ye, 3] = By_arr[xs:xe, ys:ye] self.da1.globalToLocal(self.Bx, self.localBx) self.da1.globalToLocal(self.By, self.localBy) self.da1.globalToLocal(self.Vx, self.localVx) self.da1.globalToLocal(self.Vy, self.localVy) Bx_arr = self.da1.getVecArray(self.localBx) By_arr = self.da1.getVecArray(self.localBy) Vx_arr = self.da1.getVecArray(self.localVx) Vy_arr = self.da1.getVecArray(self.localVy) P_arr = self.da1.getVecArray(self.P) for i in range(xs, xe): for j in range(ys, ye): P_arr[i,j] = init_data.pressure(xc_arr[i,j] + 0.5 * self.hx, yc_arr[i,j] + 0.5 * self.hy, self.hx, self.hy) # P_arr[i,j] = init_data.pressure(coords[i,j][0] + 0.5 * self.hx, coords[i,j][1] + 0.5 * self.hy, self.hx, self.hy) \ # + 0.5 * (0.25 * (Bx_arr[i,j] + Bx_arr[i+1,j])**2 + 0.25 * (By_arr[i,j] + By_arr[i,j+1])**2) # P_arr[i,j] = init_data.pressure(coords[i,j][0] + 0.5 * self.hx, coords[i,j][1] + 0.5 * self.hy, self.hx, self.hy) \ # + 0.5 * (0.25 * (Vx_arr[i,j] + Vx_arr[i+1,j])**2 + 0.25 * (Vy_arr[i,j] + Vy_arr[i,j+1])**2) # P_arr[i,j] = init_data.pressure(coords[i,j][0] + 0.5 * self.hx, coords[i,j][1] + 0.5 * self.hy, self.hx, self.hy) \ # + 0.5 * (0.25 * (Vx_arr[i,j] + Vx_arr[i+1,j])**2 + 0.25 * (Vy_arr[i,j] + Vy_arr[i,j+1])**2) \ # - 1.0 * (0.25 * (Bx_arr[i,j] + Bx_arr[i+1,j])**2 + 0.25 * (By_arr[i,j] + By_arr[i,j+1])**2) # copy distribution function to solution vector x_arr = self.da5.getVecArray(self.x) x_arr[xs:xe, ys:ye, 4] = P_arr [xs:xe, ys:ye] # update solution history # self.petsc_matrix.update_history(self.x) self.petsc_function.update_history(self.x) # create HDF5 output file self.hdf5_viewer = PETSc.ViewerHDF5().create(cfg['io']['hdf5_output'], mode=PETSc.Viewer.Mode.WRITE, comm=PETSc.COMM_WORLD) self.hdf5_viewer.pushGroup("/") # write grid data to hdf5 file coords_x = self.dax.getCoordinates() coords_y = self.day.getCoordinates() coords_x.setName('x') coords_y.setName('y') self.hdf5_viewer(coords_x) self.hdf5_viewer(coords_y) # write initial data to hdf5 file self.hdf5_viewer.setTimestep(0) self.save_hdf5_vectors()
P.setValues([0], [0], 1.0 / hx) lxs += 1 if lxe == mx: P.setValues([mx - 1], [mx - 1], 1.0 / hx) lxe -= 1 for i in range(lxs, lxe): P.setValues([i], [i - 1, i, i + 1], [ -1.0 / hx, 2.0 / hx - hx * math.exp(xx[i - gxs]) + shift, -1.0 / hx ]) P.assemble() return True # same_nz M = PETSc.Options().getInt('M', 9) da = PETSc.DA().create([M], comm=PETSc.COMM_WORLD) f = da.createGlobalVector() x = f.duplicate() J = da.getMatrix(PETSc.Mat.Type.AIJ) ts = PETSc.TS().create(PETSc.COMM_WORLD) ts.setProblemType(PETSc.TS.ProblemType.NONLINEAR) ts.setType(ts.Type.GL) ode = MyODE(da) ts.setIFunction(ode.function, f) ts.setIJacobian(ode.jacobian, J) ts.setTimeStep(0.1) ts.setDuration(10, 1.0) ts.setFromOptions()
def advection1D(self, M, cfl, T): '''Script to solve 1D advection equation: q_t + q_x = 0 Using first-order finite differences''' # PETSc DA object, which handles structured grids, here are # requesting M global points da = PETSc.DA().create([M]) print(da.getSizes()) ranges = da.getRanges() print(ranges[0][1] - ranges[0][0]) # this solves the problem on the domain [0 1] by default da.setUniformCoordinates() # view commands dump to the screen da.view() # this gets the coordinate vector, this is akin to the # linspace(0,1,M) command xvec = da.getCoordinates() xvec.view() # access the local data array within the globally distributed # vector x = xvec.getArray() h = x[1] - x[0] k = cfl * h # global vector represents coordinated data with other # processors fg = da.createGlobalVector() # local vector represents local data that is not shared fl = da.createLocalVector() # we will operate on the local vector, then coordinate with # other processors on the global vector N = int(round(T / k)) # Initial condition: q = np.exp(-10 * (x - 0.5)**2) # this is a little tricky. the local vector contains ghost # points, which represents data that comes in from the global # vector when we are coordinating. On the other hand, the # global vector only ever contains the data we own. So when # we set initial conditions, we do it on the global vector, # then scatter to the local fg.setArray(q) da.globalToLocal(fg, fl) # this should dump the properly set initial conditions fg.view() for n in range(N + 1): # grab the working array out of the local vector # (including ghost points) q = fl.getArray() # operate on the local data q[1:] = q[1:] - cfl * (q[1:] - q[:-1]) # restore the working array fl.setArray(q) # this is a local update, the local array in the local # vector is copied into the corresponding local array in # the global vector, ghost points are discarded da.localToGlobal(fl, fg) fg.view() # this is a global update coordinated over all processes. # The local array in the global vector is copied into the # corresponding local array in the local vector. # In addition, the ghost values needed to operate locally # are sent over MPI to the correct positions in the local # array in the local vector. da.globalToLocal(fg, fl)
finally: fh.close() import petsc4py, sys petsc4py.init(sys.argv) from petsc4py import PETSc execfile('petsc-mat.py') execfile('petsc-ksp.py') OptDB = PETSc.Options() if OptDB.getBool('plot', True): da = PETSc.DA().create([m, n]) u = da.createGlobalVec() x.copy(u) draw = PETSc.Viewer.DRAW() OptDB['draw_pause'] = 1 draw(u) if OptDB.getBool('plot_mpl', False): try: from matplotlib import pylab except ImportError: print("matplotlib not available") else: from numpy import mgrid X, Y = mgrid[0:1:1j * m, 0:1:1j * n] Z = x[...].reshape(m, n)
if i > 0: u_w = x[i - 1, j] # west if i < mx - 1: u_e = x[i + 1, j] # east if j > 0: u_s = x[i, j - 1] # south if j < ny - 1: u_n = x[i, j + 1] # north u_xx = (-u_e + 2 * u - u_w) * hy / hx u_yy = (-u_n + 2 * u - u_s) * hx / hy y[i, j] = u_xx + u_yy OptDB = PETSc.Options() n = OptDB.getInt('n', 16) nx = OptDB.getInt('nx', n) ny = OptDB.getInt('ny', n) da = PETSc.DA().create([nx, ny], stencil_width=1) pde = Poisson2D(da) x = da.createGlobalVec() b = da.createGlobalVec() # A = da.createMat('python') A = PETSc.Mat().createPython([x.getSizes(), b.getSizes()], comm=da.comm) A.setPythonContext(pde) A.setUp() ksp = PETSc.KSP().create() ksp.setOperators(A) ksp.setType('cg') pc = ksp.getPC() pc.setType('none') ksp.setFromOptions()
# Run this example with, say, 2 procs import sys sys.path.append('..') from petsc4py import PETSc as petsc from pycuda import autoinit import pycuda.gpuarray as gpuarray import pycuda.driver as cuda from pycuda.compiler import SourceModule import GPUArray da = petsc.DA().create([16, 1], stencil_width=1, comm=petsc.COMM_WORLD, stencil_type='star') da.setVecType('cusp') V_local = da.createLocalVec() V_global = da.createGlobalVec() V_global.set(1.0) # Get a handle to the global vec: varray = GPUArray.getGPUArray(V_global) # CUDA kernel: mod = SourceModule(""" __global__ void doublify(double *a) {
# -------- MAIN PROGRAM -------------- # create DA and allocate global and local variables from petsc4py import PETSc if mat_dispersion: from fdtd_da_pml import fdtddispersion2d as fdtd_2d from fdtd_da_pml import calcdispersion2d else: from fdtd_da_pml import fdtd2d as fdtd_2d stype = PETSc.DA.StencilType.STAR swidth = 1 da = PETSc.DA().create([nx, ny], dof=1, stencil_type=stype, stencil_width=swidth) (xi, xf), (yi, yf) = da.getRanges() (gxi, gxf), (gyi, gyf) = da.getGhostRanges() Q1 = da.createGlobalVec() Q2 = da.createGlobalVec() Q3 = da.createGlobalVec() Q1loc = da.createLocalVec() Q2loc = da.createLocalVec() Q3loc = da.createLocalVec() q1 = Q1loc.getArray().reshape([gxf - gxi, gyf - gyi], order='F') q2 = Q2loc.getArray().reshape([gxf - gxi, gyf - gyi], order='F') q3 = Q3loc.getArray().reshape([gxf - gxi, gyf - gyi], order='F')
from petsc4py import PETSc import numpy as np import DMPFOR global_nx = 3 global_ny = 2 dof = 4 da = PETSc.DA().create( dim=2, dof=dof, sizes=[global_nx, global_ny], #periodic_type = PETSc.DA.PeriodicType.GHOSTED_XYZ, #stencil_type=self.STENCIL, #stencil_width=2, comm=PETSc.COMM_WORLD) gVec = da.createGlobalVector() lVec = da.createLocalVector() ranges = da.getRanges() nx_start = ranges[0][0] nx_end = ranges[0][1] ny_start = ranges[1][0] ny_end = ranges[1][1] nx = nx_end - nx_start ny = ny_end - ny_start q = np.empty((dof, nx, ny), order='F')
f = PETSc.Vec().createSeq(nx * ny) snes.setFunction(appc.evalFunction, f) # configure the nonlinear solver # to use a matrix-free Jacobian snes.setUseMF(True) snes.getKSP().setType('cg') snes.setFromOptions() # solve the nonlinear problem b, x = None, f.duplicate() x.set(0) # zero inital guess snes.solve(b, x) if OptDB.getBool('plot', True): da = PETSc.DA().create([nx, ny]) u = da.createGlobalVec() x.copy(u) draw = PETSc.Viewer.DRAW() OptDB['draw_pause'] = 1 draw(u) if OptDB.getBool('plot_mpl', False): try: from matplotlib import pylab except ImportError: PETSc.Sys.Print("matplotlib not available") else: from numpy import mgrid X, Y = mgrid[0:1:1j * nx, 0:1:1j * ny] Z = x[...].reshape(nx, ny)
def read_petsc(solution, frame, path='./', file_prefix='claw', read_aux=False, options={}): r""" Read in pickles and PETSc data files representing the solution :Input: - *solution* - (:class:`~pyclaw.solution.Solution`) Solution object to read the data into. - *frame* - (int) Frame number to be read in - *path* - (string) Path to the current directory of the file - *file_prefix* - (string) Prefix of the files to be read in. ``default = 'fort'`` - *read_aux* (bool) Whether or not an auxillary file will try to be read in. ``default = False`` - *options* - (dict) Optional argument dictionary, see `PETScIO Option Table`_ .. _`PETScIO Option Table`: format : one of 'ascii' or 'binary' """ # Option parsing option_defaults = {'format': 'binary'} for (k, v) in option_defaults.iteritems(): if options.has_key(k): pass else: options[k] = option_defaults[k] pickle_filename = os.path.join( path, '%s.pkl' % file_prefix) + str(frame).zfill(4) viewer_filename = os.path.join( path, '%s.ptc' % file_prefix) + str(frame).zfill(4) aux_filename = os.path.join( path, '%s_aux.ptc' % file_prefix) + str(frame).zfill(4) if frame < 0: # Don't construct file names with negative frameno values. raise IOError("Frame " + str(frame) + " does not exist ***") pickle_file = open(pickle_filename, 'rb') # this dictionary is mostly holding debugging information, only ngrids is needed # most of this information is explicitly saved in the individual grids value_dict = pickle.load(pickle_file) ngrids = value_dict['ngrids'] read_aux = value_dict['write_aux'] ndim = value_dict['ndim'] # now set up the PETSc viewer if options['format'] == 'ascii': viewer = PETSc.Viewer().createASCII(viewer_filename, PETSc.Viewer.Mode.READ) if read_aux: aux_viewer = PETSc.Viewer().createASCII(aux_viewer_filename, PETSc.Viewer.Mode.READ) elif options['format'] == 'binary': viewer = PETSc.Viewer().createBinary(viewer_filename, PETSc.Viewer.Mode.READ) if read_aux: aux_viewer = PETSc.Viewer().createBinary(aux_viewer_filename, PETSc.Viewer.Mode.READ) else: raise IOError('format type %s not supported' % options['format']) for m in xrange(ngrids): grid_dict = pickle.load(pickle_file) gridno = grid_dict['gridno'] level = grid_dict['level'] names = grid_dict['names'] lower = grid_dict['lower'] n = grid_dict['n'] d = grid_dict['d'] dimensions = [] for i in xrange(ndim): dimensions.append( pyclaw.solution.Dimension(names[i], lower[i], lower[i] + n[i] * d[i], n[i])) grid = pyclaw.solution.Grid(dimensions) grid.t = value_dict['t'] grid.meqn = value_dict['meqn'] nbc = [x + (2 * grid.mbc) for x in grid.n] grid.q_da = PETSc.DA().create( dim=grid.ndim, dof=grid.meqn, # should be modified to reflect the update sizes=nbc, #periodic_type = PETSc.DA.PeriodicType.X, #periodic_type=grid.PERIODIC, #stencil_type=grid.STENCIL, stencil_width=grid.mbc, comm=PETSc.COMM_WORLD) grid.gqVec = PETSc.Vec().load(viewer) grid.q = grid.gqVec.getArray().copy() grid.q.shape = (grid.q.size / grid.meqn, grid.meqn) if read_aux: nbc = [x + (2 * grid.mbc) for x in grid.n] grid.aux_da = PETSc.DA().create( dim=grid.ndim, dof=maux, # should be modified to reflect the update sizes=nbc, #Amal: what about for 2D, 3D #periodic_type = PETSc.DA.PeriodicType.X, #periodic_type=grid.PERIODIC, #stencil_type=grid.STENCIL, stencil_width=grid.mbc, comm=PETSc.COMM_WORLD) grid.gauxVec = PETSc.Vec().load(aux_viewer) grid.aux = grid.gauxVec.getArray().copy() grid.aux.shape = (grid.aux.size / grid.meqn, grid.meqn) # Add AMR attributes: grid.gridno = gridno grid.level = level solution.grids.append(grid) pickle_file.close() viewer.destroy() if read_aux: aux_viewer.destroy()
def __init__(self, cfgfile): ''' Constructor ''' # load run config file cfg = Config(cfgfile) # timestep setup self.ht = cfg['grid']['ht'] # timestep size self.nt = cfg['grid']['nt'] # number of timesteps self.nsave = cfg['io']['nsave'] # save only every nsave'th timestep self.omega = 0.1 # relaxation parameter # grid setup nx = cfg['grid']['nx'] # number of points in x ny = cfg['grid']['ny'] # number of points in y Lx = cfg['grid']['Lx'] # spatial domain in x x1 = cfg['grid']['x1'] # x2 = cfg['grid']['x2'] # Ly = cfg['grid']['Ly'] # spatial domain in y y1 = cfg['grid']['y1'] # y2 = cfg['grid']['y2'] # if x1 != x2: Lx = x2 - x1 else: x1 = 0.0 x2 = Lx if y1 != y2: Ly = y2 - y1 else: y1 = 0.0 y2 = Ly self.hx = Lx / nx # gridstep size in x self.hy = Ly / ny # gridstep size in y self.time = PETSc.Vec().createMPI(1, PETSc.DECIDE, comm=PETSc.COMM_WORLD) self.time.setName('t') if PETSc.COMM_WORLD.getRank() == 0: self.time.setValue(0, 0.0) # set some PETSc options OptDB = PETSc.Options() OptDB.setValue('ksp_rtol', cfg['solver']['petsc_residual']) # OptDB.setValue('ksp_max_it', 100) OptDB.setValue('ksp_max_it', 200) # OptDB.setValue('ksp_max_it', 1000) # OptDB.setValue('ksp_max_it', 2000) # OptDB.setValue('ksp_monitor', '') # OptDB.setValue('log_info', '') # OptDB.setValue('log_summary', '') # create DA with single dof self.da1 = PETSc.DA().create(dim=2, dof=1, sizes=[nx, ny], proc_sizes=[PETSc.DECIDE, PETSc.DECIDE], boundary_type=('periodic', 'periodic'), stencil_width=1, stencil_type='box') # create DA (dof = 4 for Bx, By, Vx, Vy) self.da4 = PETSc.DA().create(dim=2, dof=4, sizes=[nx, ny], proc_sizes=[PETSc.DECIDE, PETSc.DECIDE], boundary_type=('periodic', 'periodic'), stencil_width=1, stencil_type='box') # create DA (dof = 5 for Bx, By, Vx, Vy, P) self.da5 = PETSc.DA().create(dim=2, dof=5, sizes=[nx, ny], proc_sizes=[PETSc.DECIDE, PETSc.DECIDE], boundary_type=('periodic', 'periodic'), stencil_width=1, stencil_type='box') # create DA for x grid self.dax = PETSc.DA().create(dim=1, dof=1, sizes=[nx], proc_sizes=[PETSc.DECIDE], boundary_type=('periodic')) # create DA for y grid self.day = PETSc.DA().create(dim=1, dof=1, sizes=[ny], proc_sizes=[PETSc.DECIDE], boundary_type=('periodic')) # initialise grid self.da1.setUniformCoordinates(xmin=x1, xmax=x2, ymin=y1, ymax=y2) self.da4.setUniformCoordinates(xmin=x1, xmax=x2, ymin=y1, ymax=y2) self.da5.setUniformCoordinates(xmin=x1, xmax=x2, ymin=y1, ymax=y2) self.dax.setUniformCoordinates(xmin=x1, xmax=x2) self.day.setUniformCoordinates(xmin=y1, xmax=y2) # create solution and RHS vector self.x = self.da4.createGlobalVec() self.b = self.da4.createGlobalVec() self.u = self.da5.createGlobalVec() # create residual vectors self.ru = self.da5.createGlobalVec() self.rx = self.da4.createGlobalVec() self.rp = self.da1.createGlobalVec() # create global RK4 vectors self.X1 = self.da4.createGlobalVec() self.X2 = self.da4.createGlobalVec() self.X3 = self.da4.createGlobalVec() self.X4 = self.da4.createGlobalVec() # create local RK4 vectors self.localX = self.da4.createLocalVec() self.localX1 = self.da4.createLocalVec() self.localX2 = self.da4.createLocalVec() self.localX3 = self.da4.createLocalVec() self.localX4 = self.da4.createLocalVec() self.localP = self.da1.createLocalVec() # create vectors for magnetic and velocity field self.Bx = self.da1.createGlobalVec() self.By = self.da1.createGlobalVec() self.Vx = self.da1.createGlobalVec() self.Vy = self.da1.createGlobalVec() self.P = self.da1.createGlobalVec() # set variable names self.x.setName('solver_x') self.b.setName('solver_b') self.Bx.setName('Bx') self.By.setName('By') self.Vx.setName('Vx') self.Vy.setName('Vy') self.P.setName('P') # create Matrix object self.petsc_matrix = PETScSolver(self.da1, self.da4, nx, ny, self.ht, self.hx, self.hy, self.omega) self.petsc_function = PETScFunction(self.da1, self.da5, nx, ny, self.ht, self.hx, self.hy) # create sparse matrix self.mat = PETSc.Mat().createPython( [self.x.getSizes(), self.b.getSizes()], comm=PETSc.COMM_WORLD) self.mat.setPythonContext(self.petsc_matrix) self.mat.setUp() # create linear solver and preconditioner self.ksp = PETSc.KSP().create() self.ksp.setFromOptions() self.ksp.setOperators(self.mat) self.ksp.setType(cfg['solver']['petsc_ksp_type']) self.ksp.setInitialGuessNonzero(True) self.pc = self.ksp.getPC() self.pc.setType(cfg['solver']['petsc_pc_type']) # # create Preconditioner matrix and solver # self.pc_mat = PETScPreconditioner(self.da1, self.da4, self.P, nx, ny, self.ht, self.hx, self.hy) # # # create sparse matrix # self.pc_A = PETSc.Mat().createPython([self.x.getSizes(), self.b.getSizes()], comm=PETSc.COMM_WORLD) # self.pc_A.setPythonContext(self.pc_mat) # self.pc_A.setUp() # # # create linear solver and preconditioner # self.pc_ksp = PETSc.KSP().create() # self.pc_ksp.setFromOptions() # self.pc_ksp.setOperators(self.pc_A) # self.pc_ksp.setType(cfg['solver']['petsc_ksp_type']) # self.pc_ksp.setInitialGuessNonzero(True) # # self.pc_pc = self.pc_ksp.getPC() # self.pc_pc.setType('none') # # # # create Arakawa solver object # self.mhd_rk4 = PETScRK4(self.da4, nx, ny, self.ht, self.hx, self.hy) # set initial data (xs, xe), (ys, ye) = self.da1.getRanges() coords = self.da1.getCoordinateDA().getVecArray( self.da1.getCoordinates()) # print # print(self.hx) # print(coords[1,0][0] - coords[0,0][0]) # print # print(self.hy) # print(coords[0,1][1] - coords[0,0][1]) # print # print(Lx) # print(coords[-1,0][0]+self.hx) # print # print(Ly) # print(coords[0,-1][1]+self.hy) # print x_arr = self.da4.getVecArray(self.x) Bx_arr = self.da1.getVecArray(self.Bx) By_arr = self.da1.getVecArray(self.By) Vx_arr = self.da1.getVecArray(self.Vx) Vy_arr = self.da1.getVecArray(self.Vy) P_arr = self.da1.getVecArray(self.P) if cfg['initial_data']['magnetic_python'] != None: init_data = __import__( "runs." + cfg['initial_data']['magnetic_python'], globals(), locals(), ['magnetic_x', 'magnetic_y'], 0) for i in range(xs, xe): for j in range(ys, ye): Bx_arr[i, j] = init_data.magnetic_x(coords[i, j][0], coords[i, j][1], Lx, Ly) By_arr[i, j] = init_data.magnetic_y(coords[i, j][0], coords[i, j][1], Lx, Ly) else: Bx_arr[xs:xe, ys:ye] = cfg['initial_data']['magnetic'] By_arr[xs:xe, ys:ye] = cfg['initial_data']['magnetic'] if cfg['initial_data']['velocity_python'] != None: init_data = __import__( "runs." + cfg['initial_data']['velocity_python'], globals(), locals(), ['velocity_x', 'velocity_y'], 0) for i in range(xs, xe): for j in range(ys, ye): Vx_arr[i, j] = init_data.velocity_x(coords[i, j][0], coords[i, j][1], Lx, Ly) Vy_arr[i, j] = init_data.velocity_y(coords[i, j][0], coords[i, j][1], Lx, Ly) else: Vx_arr[xs:xe, ys:ye] = cfg['initial_data']['velocity'] Vy_arr[xs:xe, ys:ye] = cfg['initial_data']['velocity'] if cfg['initial_data']['pressure_python'] != None: init_data = __import__( "runs." + cfg['initial_data']['pressure_python'], globals(), locals(), ['pressure', ''], 0) for i in range(xs, xe): for j in range(ys, ye): P_arr[i, j] = init_data.pressure(coords[i, j][0], coords[i, j][1], Lx, Ly) # copy distribution function to solution vector x_arr[xs:xe, ys:ye, 0] = Bx_arr[xs:xe, ys:ye] x_arr[xs:xe, ys:ye, 1] = By_arr[xs:xe, ys:ye] x_arr[xs:xe, ys:ye, 2] = Vx_arr[xs:xe, ys:ye] x_arr[xs:xe, ys:ye, 3] = Vy_arr[xs:xe, ys:ye] # copy 4D solution vector to 5D solution vector self.copy_solution_to_u() # update solution history self.petsc_matrix.update_history(self.x, self.P) self.petsc_function.update_history(self.u) # create HDF5 output file self.hdf5_viewer = PETSc.Viewer().createHDF5( cfg['io']['hdf5_output'], mode=PETSc.Viewer.Mode.WRITE, comm=PETSc.COMM_WORLD) self.hdf5_viewer.HDF5PushGroup("/") # write grid data to hdf5 file coords_x = self.dax.getCoordinates() coords_y = self.day.getCoordinates() coords_x.setName('x') coords_y.setName('y') self.hdf5_viewer(coords_x) self.hdf5_viewer(coords_y) # write initial data to hdf5 file self.hdf5_viewer.HDF5SetTimestep(0) self.hdf5_viewer(self.time) # self.hdf5_viewer(self.x) # self.hdf5_viewer(self.b) self.hdf5_viewer(self.Bx) self.hdf5_viewer(self.By) self.hdf5_viewer(self.Vx) self.hdf5_viewer(self.Vy) self.hdf5_viewer(self.P)
import petsc4py import sys petsc4py.init(sys.argv) from petsc4py import PETSc from matplotlib import pylab # Dimensions of the 2D grid. nx = 101 ny = 101 w = 2. / 10. # Angular frequency of wave (2*pi / period). # Create the DA. da = PETSc.DA().create([nx, ny], \ stencil_width=1, \ boundary_type=('ghosted', 'ghosted')) # Create the rhs vector based on the DA. b = da.createGlobalVec() b_val = da.getVecArray(b) # Obtain access to elements of b. b_val[50, 50] = 1 # Set central value to 1. # Create (a vector to store) the solution vector. x = da.createGlobalVec() # Create the matrix. A = da.getMatrix('aij') # Stencil objects make it easy to set the values of the matrix elements.
def __init__(self, cfgfile, mode="none"): ''' Constructor ''' petsc4py.init(sys.argv) if PETSc.COMM_WORLD.getRank() == 0: print("") print("Reduced MHD 2D") print("==============") print("") # solver mode self.mode = mode # set run id to timestamp self.run_id = datetime.datetime.fromtimestamp( time.time()).strftime("%y%m%d%H%M%S") if PETSc.COMM_WORLD.getRank() == 0: print(" Config: %s" % cfgfile) # load run config file self.cfg = Config(cfgfile) # timestep setup self.ht = self.cfg['grid']['ht'] # timestep size self.nt = self.cfg['grid']['nt'] # number of timesteps self.nsave = self.cfg['io'][ 'nsave'] # save only every nsave'th timestep # grid setup self.nx = self.cfg['grid']['nx'] # number of points in x self.ny = self.cfg['grid']['ny'] # number of points in y self.Lx = self.cfg['grid']['Lx'] # spatial domain in x x1 = self.cfg['grid']['x1'] # x2 = self.cfg['grid']['x2'] # self.Ly = self.cfg['grid']['Ly'] # spatial domain in y y1 = self.cfg['grid']['y1'] # y2 = self.cfg['grid']['y2'] # self.hx = self.cfg['grid']['hx'] # gridstep size in x self.hy = self.cfg['grid']['hy'] # gridstep size in y # create time vector self.time = PETSc.Vec().createMPI(1, comm=PETSc.COMM_WORLD) self.time.setName('t') # electron skin depth self.de = self.cfg['initial_data']['skin_depth'] # double bracket dissipation self.nu = self.cfg['initial_data']['dissipation'] # set global tolerance self.tolerance = self.cfg['solver'][ 'petsc_snes_atol'] * self.nx * self.ny # direct solver package self.solver_package = self.cfg['solver']['lu_solver_package'] # set some PETSc solver options OptDB = PETSc.Options() OptDB.setValue('ksp_rtol', self.cfg['solver']['petsc_ksp_rtol']) OptDB.setValue('ksp_atol', self.cfg['solver']['petsc_ksp_atol']) OptDB.setValue('ksp_max_it', self.cfg['solver']['petsc_ksp_max_iter']) OptDB.setValue('pc_type', 'hypre') OptDB.setValue('pc_hypre_type', 'boomeramg') # create DA with single dof self.da1 = PETSc.DA().create(dim=2, dof=1, sizes=[self.nx, self.ny], proc_sizes=[PETSc.DECIDE, PETSc.DECIDE], boundary_type=('periodic', 'periodic'), stencil_width=1, stencil_type='box') # create DA (dof = 4 for A, J, P, O) self.da4 = PETSc.DA().create(dim=2, dof=4, sizes=[self.nx, self.ny], proc_sizes=[PETSc.DECIDE, PETSc.DECIDE], boundary_type=('periodic', 'periodic'), stencil_width=1, stencil_type='box') # create DA for x grid self.dax = PETSc.DA().create(dim=1, dof=1, sizes=[self.nx], proc_sizes=[PETSc.DECIDE], boundary_type=('periodic')) # create DA for y grid self.day = PETSc.DA().create(dim=1, dof=1, sizes=[self.ny], proc_sizes=[PETSc.DECIDE], boundary_type=('periodic')) # initialise grid self.da1.setUniformCoordinates(xmin=x1, xmax=x2, ymin=y1, ymax=y2) self.da4.setUniformCoordinates(xmin=x1, xmax=x2, ymin=y1, ymax=y2) self.dax.setUniformCoordinates(xmin=x1, xmax=x2) self.day.setUniformCoordinates(xmin=y1, xmax=y2) # create solution and RHS vector self.dx = self.da4.createGlobalVec() self.dy = self.da4.createGlobalVec() self.x = self.da4.createGlobalVec() self.b = self.da4.createGlobalVec() self.f = self.da4.createGlobalVec() self.Pb = self.da1.createGlobalVec() self.FA = self.da1.createGlobalVec() self.FJ = self.da1.createGlobalVec() self.FP = self.da1.createGlobalVec() self.FO = self.da1.createGlobalVec() # create initial guess vectors self.igFA1 = self.da1.createGlobalVec() self.igFA2 = self.da1.createGlobalVec() self.igFO1 = self.da1.createGlobalVec() self.igFO2 = self.da1.createGlobalVec() # nullspace vectors self.x0 = self.da4.createGlobalVec() self.P0 = self.da1.createGlobalVec() # create vectors for magnetic and velocity field self.A = self.da1.createGlobalVec() # magnetic vector potential A self.J = self.da1.createGlobalVec() # current density J self.P = self.da1.createGlobalVec() # streaming function psi self.O = self.da1.createGlobalVec() # vorticity omega self.Bx = self.da1.createGlobalVec() self.By = self.da1.createGlobalVec() self.Vx = self.da1.createGlobalVec() self.Vy = self.da1.createGlobalVec() # set variable names self.A.setName('A') self.J.setName('J') self.P.setName('P') self.O.setName('O') self.Bx.setName('Bx') self.By.setName('By') self.Vx.setName('Vx') self.Vy.setName('Vy') # initialise nullspace self.x0.set(0.) x0_arr = self.da4.getVecArray(self.x0)[...] x0_arr[:, :, 2] = 1. self.x0.assemble() self.x0.normalize() self.solver_nullspace = PETSc.NullSpace().create(constant=False, vectors=(self.x0, )) self.poisson_nullspace = PETSc.NullSpace().create(constant=True) # initialise Poisson matrix self.Pm = self.da1.createMat() self.Pm.setOption(PETSc.Mat.Option.NEW_NONZERO_ALLOCATION_ERR, False) self.Pm.setUp() self.Pm.setNullSpace(self.poisson_nullspace) # create Poisson solver object self.petsc_poisson = PETScPoisson(self.da1, self.nx, self.ny, self.hx, self.hy) # setup linear Poisson solver self.poisson_ksp = PETSc.KSP().create() self.poisson_ksp.setFromOptions() self.poisson_ksp.setOperators(self.Pm) self.poisson_ksp.setTolerances( rtol=self.cfg['solver']['poisson_ksp_rtol'], atol=self.cfg['solver']['poisson_ksp_atol'], max_it=self.cfg['solver']['poisson_ksp_max_iter']) self.poisson_ksp.setType('cg') self.poisson_ksp.getPC().setType('hypre') self.poisson_ksp.setUp() self.petsc_poisson.formMat(self.Pm) # create derivatives object self.derivatives = PETScDerivatives(self.da1, self.nx, self.ny, self.ht, self.hx, self.hy) # read initial data if self.cfg["io"]["hdf5_input"] != None and self.cfg["io"][ "hdf5_input"] != "": if self.cfg["initial_data"]["python"] != None and self.cfg[ "initial_data"]["python"] != "": if PETSc.COMM_WORLD.getRank() == 0: print( "WARNING: Both io.hdf5_input and initial_data.python are set!" ) print(" Reading initial data from HDF5 file.") self.read_initial_data_from_hdf5() else: self.read_initial_data_from_python() # copy initial data vectors to x self.copy_x_from_da1_to_da4() # create HDF5 output file and write parameters hdf5_filename = self.cfg['io']['hdf5_output'] last_dot = hdf5_filename.rfind('.') hdf5_filename = hdf5_filename[:last_dot] + "." + str( self.run_id) + hdf5_filename[last_dot:] if PETSc.COMM_WORLD.getRank() == 0: print(" Output: %s" % hdf5_filename) hdf5out = h5py.File(hdf5_filename, "w", driver="mpio", comm=PETSc.COMM_WORLD.tompi4py()) hdf5out.attrs["run_id"] = self.run_id for cfg_group in self.cfg: for cfg_item in self.cfg[cfg_group]: if self.cfg[cfg_group][cfg_item] != None: value = self.cfg[cfg_group][cfg_item] else: value = "" hdf5out.attrs[cfg_group + "." + cfg_item] = value hdf5out.attrs["solver.solver_mode"] = self.mode if self.cfg["initial_data"]["python"] != None and self.cfg[ "initial_data"]["python"] != "": python_input = open( "examples/" + self.cfg['initial_data']['python'] + ".py", 'r') python_file = python_input.read() python_input.close() else: python_file = "" hdf5out.attrs["initial_data.python_file"] = python_file hdf5out.close() # create HDF5 viewer self.hdf5_viewer = PETSc.ViewerHDF5().create( hdf5_filename, mode=PETSc.Viewer.Mode.APPEND, comm=PETSc.COMM_WORLD) self.hdf5_viewer.pushGroup("/") # write grid data to hdf5 file coords_x = self.dax.getCoordinates() coords_y = self.day.getCoordinates() coords_x.setName('x') coords_y.setName('y') self.hdf5_viewer(coords_x) self.hdf5_viewer(coords_y) # write initial data to hdf5 file self.save_to_hdf5(0) # output some more information if PETSc.COMM_WORLD.getRank() == 0: print("") print(" nt = %i" % (self.nt)) print(" nx = %i" % (self.nx)) print(" ny = %i" % (self.ny)) print("") print(" ht = %f" % (self.ht)) print(" hx = %f" % (self.hx)) print(" hy = %f" % (self.hy)) print("") print(" PETSc SNES rtol = %e" % self.cfg['solver']['petsc_snes_rtol']) print(" atol = %e" % self.cfg['solver']['petsc_snes_atol']) print(" stol = %e" % self.cfg['solver']['petsc_snes_stol']) print(" max iter = %i" % self.cfg['solver']['petsc_snes_max_iter']) print("") print(" PETSc KSP rtol = %e" % self.cfg['solver']['petsc_ksp_rtol']) print(" atol = %e" % self.cfg['solver']['petsc_ksp_atol']) print(" max iter = %i" % self.cfg['solver']['petsc_ksp_max_iter']) print("")
def __init__(self, cfgfile): ''' Constructor ''' super().__init__(cfgfile, mode="split") OptDB = PETSc.Options() # OptDB.setValue('ksp_monitor', '') # OptDB.setValue('snes_monitor', '') # # OptDB.setValue('log_info', '') # OptDB.setValue('log_summary', '') OptDB.setValue('ksp_rtol', self.cfg['solver']['petsc_ksp_rtol']) OptDB.setValue('ksp_atol', self.cfg['solver']['petsc_ksp_atol']) OptDB.setValue('ksp_max_it', self.cfg['solver']['petsc_ksp_max_iter']) # OptDB.setValue('ksp_initial_guess_nonzero', 1) OptDB.setValue('pc_type', 'hypre') OptDB.setValue('pc_hypre_type', 'boomeramg') OptDB.setValue('pc_hypre_boomeramg_max_iter', 2) # OptDB.setValue('pc_hypre_boomeramg_max_levels', 6) # OptDB.setValue('pc_hypre_boomeramg_tol', 1e-7) # create DA (dof = 2 for A, P) self.da2 = PETSc.DA().create(dim=2, dof=2, sizes=[self.nx, self.ny], proc_sizes=[PETSc.DECIDE, PETSc.DECIDE], boundary_type=('periodic', 'periodic'), stencil_width=1, stencil_type='box') # create solution and RHS vector self.dx2 = self.da2.createGlobalVec() self.dy2 = self.da2.createGlobalVec() self.b = self.da2.createGlobalVec() self.Ad = self.da1.createGlobalVec() self.Jd = self.da1.createGlobalVec() self.Pd = self.da1.createGlobalVec() self.Od = self.da1.createGlobalVec() # create Jacobian, Function, and linear Matrix objects self.petsc_precon = PETScPreconditioner(self.da1, self.da2, self.nx, self.ny, self.ht, self.hx, self.hy) # self.petsc_solver2 = PETScSolverDOF2(self.da1, self.da2, self.nx, self.ny, self.ht, self.hx, self.hy) self.petsc_solver2 = PETScSolverDOF2(self.da1, self.da2, self.nx, self.ny, self.ht, self.hx, self.hy, self.petsc_precon) self.petsc_solver = PETScSolver(self.da1, self.da4, self.nx, self.ny, self.ht, self.hx, self.hy) self.petsc_precon.set_tolerances( poisson_rtol=self.cfg['solver']['pc_poisson_rtol'], poisson_atol=self.cfg['solver']['pc_poisson_atol'], poisson_max_it=self.cfg['solver']['pc_poisson_max_iter'], parabol_rtol=self.cfg['solver']['pc_parabol_rtol'], parabol_atol=self.cfg['solver']['pc_parabol_atol'], parabol_max_it=self.cfg['solver']['pc_parabol_max_iter'], jacobi_max_it=self.cfg['solver']['pc_jacobi_max_iter']) # initialise matrixfree Jacobian self.Jmf = PETSc.Mat().createPython( [self.b.getSizes(), self.b.getSizes()], context=self.petsc_solver2, comm=PETSc.COMM_WORLD) self.Jmf.setUp() # create linear solver self.ksp = PETSc.KSP().create() self.ksp.setFromOptions() self.ksp.setOperators(self.Jmf) self.ksp.setInitialGuessNonzero(True) self.ksp.setType('fgmres') self.ksp.getPC().setType('none') # update solution history self.petsc_solver.update_previous(self.x) self.petsc_solver2.update_previous(self.A, self.J, self.P, self.O)
def __init__(self, cfgfile): ''' Constructor ''' # stencil = 1 stencil = 2 # load run config file cfg = Config(cfgfile) self.cfg = cfg cfg.set_petsc_options() # timestep setup self.ht = cfg['grid']['ht'] # timestep size self.nt = cfg['grid']['nt'] # number of timesteps self.nsave = cfg['io']['nsave'] # save only every nsave'th timestep # grid setup nx = cfg['grid']['nx'] # number of points in x ny = cfg['grid']['ny'] # number of points in y Lx = cfg['grid']['Lx'] # spatial domain in x x1 = cfg['grid']['x1'] # x2 = cfg['grid']['x2'] # Ly = cfg['grid']['Ly'] # spatial domain in y y1 = cfg['grid']['y1'] # y2 = cfg['grid']['y2'] # self.nx = nx self.ny = ny if x1 != x2: Lx = x2 - x1 else: x1 = 0.0 x2 = Lx if y1 != y2: Ly = y2 - y1 else: y1 = 0.0 y2 = Ly self.hx = Lx / nx # gridstep size in x self.hy = Ly / ny # gridstep size in y # friction, viscosity and resistivity mu = cfg['initial_data']['mu'] # friction nu = cfg['initial_data']['nu'] # viscosity eta = cfg['initial_data']['eta'] # resistivity de = cfg['initial_data']['de'] # electron skin depth # self.update_jacobian = True if PETSc.COMM_WORLD.getRank() == 0: print() print("nt = %i" % (self.nt)) print("nx = %i" % (self.nx)) print("ny = %i" % (self.ny)) print() print("ht = %e" % (self.ht)) print("hx = %e" % (self.hx)) print("hy = %e" % (self.hy)) print() print("Lx = %e" % (Lx)) print("Ly = %e" % (Ly)) print() print("mu = %e" % (mu)) print("nu = %e" % (nu)) print("eta = %e" % (eta)) print() self.time = PETSc.Vec().createMPI(1, PETSc.DECIDE, comm=PETSc.COMM_WORLD) self.time.setName('t') if PETSc.COMM_WORLD.getRank() == 0: self.time.setValue(0, 0.0) # create derivatives object self.derivatives = MHD_Derivatives(nx, ny, self.ht, self.hx, self.hy) # create DA with single dof self.da1 = PETSc.DA().create(dim=2, dof=1, sizes=[nx, ny], proc_sizes=[PETSc.DECIDE, PETSc.DECIDE], boundary_type=('periodic', 'periodic'), stencil_width=stencil, stencil_type='box') # create DA (dof = 7 for Vx, Vy, Bx, By, Bix, Biy, P) self.da7 = PETSc.DA().create(dim=2, dof=7, sizes=[nx, ny], proc_sizes=[PETSc.DECIDE, PETSc.DECIDE], boundary_type=('periodic', 'periodic'), stencil_width=stencil, stencil_type='box') # create DA for x grid self.dax = PETSc.DA().create(dim=1, dof=1, sizes=[nx], proc_sizes=[PETSc.DECIDE], boundary_type=('periodic')) # create DA for y grid self.day = PETSc.DA().create(dim=1, dof=1, sizes=[ny], proc_sizes=[PETSc.DECIDE], boundary_type=('periodic')) # initialise grid self.da1.setUniformCoordinates(xmin=x1, xmax=x2, ymin=y1, ymax=y2) self.da7.setUniformCoordinates(xmin=x1, xmax=x2, ymin=y1, ymax=y2) self.dax.setUniformCoordinates(xmin=x1, xmax=x2) self.day.setUniformCoordinates(xmin=y1, xmax=y2) # create solution and RHS vector self.f = self.da7.createGlobalVec() self.x = self.da7.createGlobalVec() self.b = self.da7.createGlobalVec() # create global RK4 vectors self.Y = self.da7.createGlobalVec() self.X0 = self.da7.createGlobalVec() self.X1 = self.da7.createGlobalVec() self.X2 = self.da7.createGlobalVec() self.X3 = self.da7.createGlobalVec() self.X4 = self.da7.createGlobalVec() # create local RK4 vectors self.localX0 = self.da7.createLocalVec() self.localX1 = self.da7.createLocalVec() self.localX2 = self.da7.createLocalVec() self.localX3 = self.da7.createLocalVec() self.localX4 = self.da7.createLocalVec() # self.localP = self.da1.createLocalVec() # create vectors for magnetic and velocity field self.Bix = self.da1.createGlobalVec() self.Biy = self.da1.createGlobalVec() self.Bx = self.da1.createGlobalVec() self.By = self.da1.createGlobalVec() self.Vx = self.da1.createGlobalVec() self.Vy = self.da1.createGlobalVec() self.P = self.da1.createGlobalVec() self.xcoords = self.da1.createGlobalVec() self.ycoords = self.da1.createGlobalVec() # create local vectors for initialisation of pressure self.localBix = self.da1.createLocalVec() self.localBiy = self.da1.createLocalVec() self.localBx = self.da1.createLocalVec() self.localBy = self.da1.createLocalVec() self.localVx = self.da1.createLocalVec() self.localVy = self.da1.createLocalVec() # set variable names self.Bix.setName('Bix') self.Biy.setName('Biy') self.Bx.setName('Bx') self.By.setName('By') self.Vx.setName('Vx') self.Vy.setName('Vy') self.P.setName('P') # create Matrix object self.petsc_matrix = PETScMatrix(self.da1, self.da7, nx, ny, self.ht, self.hx, self.hy, mu, nu, eta, de) self.petsc_function = PETScFunction(self.da1, self.da7, nx, ny, self.ht, self.hx, self.hy, mu, nu, eta, de) # initialise matrix self.A = self.da7.createMat() self.A.setOption(self.A.Option.NEW_NONZERO_ALLOCATION_ERR, False) self.A.setUp() # create jacobian self.J = self.da7.createMat() self.J.setOption(self.J.Option.NEW_NONZERO_ALLOCATION_ERR, False) self.J.setUp() # create nonlinear solver self.snes = PETSc.SNES().create() self.snes.setFunction(self.petsc_function.snes_mult, self.f) self.snes.setJacobian(self.updateJacobian, self.J) self.snes.setFromOptions() # self.snes.getKSP().setInitialGuessNonzero(True) # self.snes.getKSP().getPC().setReusePreconditioner(True) self.ksp = None # set initial data (xs, xe), (ys, ye) = self.da1.getRanges() # coords = self.da1.getCoordinatesLocal() xc_arr = self.da1.getVecArray(self.xcoords) yc_arr = self.da1.getVecArray(self.ycoords) for i in range(xs, xe): for j in range(ys, ye): xc_arr[i, j] = x1 + i * self.hx yc_arr[i, j] = y1 + j * self.hy if cfg['io']['hdf5_input'] != None: hdf5_filename = self.cfg["io"]["hdf5_input"] if PETSc.COMM_WORLD.getRank() == 0: print(" Input: %s" % hdf5_filename) hdf5in = h5py.File(hdf5_filename, "r", driver="mpio", comm=PETSc.COMM_WORLD.tompi4py()) # assert self.nx == hdf5in.attrs["grid.nx"] # assert self.ny == hdf5in.attrs["grid.ny"] # assert self.hx == hdf5in.attrs["grid.hx"] # assert self.hy == hdf5in.attrs["grid.hy"] # assert self.Lx == hdf5in.attrs["grid.Lx"] # assert self.Ly == hdf5in.attrs["grid.Ly"] # # assert self.de == hdf5in.attrs["initial_data.skin_depth"] timestep = len(hdf5in["t"][...].flatten()) - 1 hdf5in.close() hdf5_viewer = PETSc.ViewerHDF5().create( cfg['io']['hdf5_input'], mode=PETSc.Viewer.Mode.READ, comm=PETSc.COMM_WORLD) hdf5_viewer.setTimestep(timestep) self.Bix.load(hdf5_viewer) self.Biy.load(hdf5_viewer) self.Bx.load(hdf5_viewer) self.By.load(hdf5_viewer) self.Vx.load(hdf5_viewer) self.Vy.load(hdf5_viewer) self.P.load(hdf5_viewer) hdf5_viewer.destroy() # copy modified magnetic induction to solution vector Bx_arr = self.da1.getVecArray(self.Bx) By_arr = self.da1.getVecArray(self.By) Vx_arr = self.da1.getVecArray(self.Vx) Vy_arr = self.da1.getVecArray(self.Vy) Bix_arr = self.da1.getVecArray(self.Bix) Biy_arr = self.da1.getVecArray(self.Biy) P_arr = self.da1.getVecArray(self.P) x_arr = self.da7.getVecArray(self.x) x_arr[xs:xe, ys:ye, 0] = Vx_arr[xs:xe, ys:ye] x_arr[xs:xe, ys:ye, 1] = Vy_arr[xs:xe, ys:ye] x_arr[xs:xe, ys:ye, 2] = Bx_arr[xs:xe, ys:ye] x_arr[xs:xe, ys:ye, 3] = By_arr[xs:xe, ys:ye] x_arr[xs:xe, ys:ye, 4] = Bix_arr[xs:xe, ys:ye] x_arr[xs:xe, ys:ye, 5] = Biy_arr[xs:xe, ys:ye] x_arr[xs:xe, ys:ye, 6] = P_arr[xs:xe, ys:ye] else: Bx_arr = self.da1.getVecArray(self.Bx) By_arr = self.da1.getVecArray(self.By) Vx_arr = self.da1.getVecArray(self.Vx) Vy_arr = self.da1.getVecArray(self.Vy) xc_arr = self.da1.getVecArray(self.xcoords) yc_arr = self.da1.getVecArray(self.ycoords) if cfg['initial_data']['magnetic_python'] != None: init_data = __import__( "examples." + cfg['initial_data']['magnetic_python'], globals(), locals(), ['magnetic_x', 'magnetic_y'], 0) for i in range(xs, xe): for j in range(ys, ye): Bx_arr[i, j] = init_data.magnetic_x( xc_arr[i, j], yc_arr[i, j] + 0.5 * self.hy, self.hx, self.hy) By_arr[i, j] = init_data.magnetic_y( xc_arr[i, j] + 0.5 * self.hx, yc_arr[i, j], self.hx, self.hy) else: Bx_arr[xs:xe, ys:ye] = cfg['initial_data']['magnetic'] By_arr[xs:xe, ys:ye] = cfg['initial_data']['magnetic'] if cfg['initial_data']['velocity_python'] != None: init_data = __import__( "examples." + cfg['initial_data']['velocity_python'], globals(), locals(), ['velocity_x', 'velocity_y'], 0) for i in range(xs, xe): for j in range(ys, ye): Vx_arr[i, j] = init_data.velocity_x( xc_arr[i, j], yc_arr[i, j] + 0.5 * self.hy, self.hx, self.hy) Vy_arr[i, j] = init_data.velocity_y( xc_arr[i, j] + 0.5 * self.hx, yc_arr[i, j], self.hx, self.hy) else: Vx_arr[xs:xe, ys:ye] = cfg['initial_data']['velocity'] Vy_arr[xs:xe, ys:ye] = cfg['initial_data']['velocity'] if cfg['initial_data']['pressure_python'] != None: init_data = __import__( "examples." + cfg['initial_data']['pressure_python'], globals(), locals(), ['pressure', ''], 0) # Fourier Filtering from scipy.fftpack import rfft, irfft nfourier_x = cfg['initial_data']['nfourier_Bx'] nfourier_y = cfg['initial_data']['nfourier_By'] if nfourier_x >= 0 or nfourier_y >= 0: print("Fourier Filtering B") # obtain whole Bx vector everywhere scatter, Xglobal = PETSc.Scatter.toAll(self.Bx) scatter.begin(self.Bx, Xglobal, PETSc.InsertMode.INSERT, PETSc.ScatterMode.FORWARD) scatter.end(self.Bx, Xglobal, PETSc.InsertMode.INSERT, PETSc.ScatterMode.FORWARD) petsc_indices = self.da1.getAO().app2petsc( np.arange(self.nx * self.ny, dtype=np.int32)) BxTmp = Xglobal.getValues(petsc_indices).copy().reshape( (self.ny, self.nx)).T scatter.destroy() Xglobal.destroy() # obtain whole By vector everywhere scatter, Xglobal = PETSc.Scatter.toAll(self.By) scatter.begin(self.By, Xglobal, PETSc.InsertMode.INSERT, PETSc.ScatterMode.FORWARD) scatter.end(self.By, Xglobal, PETSc.InsertMode.INSERT, PETSc.ScatterMode.FORWARD) petsc_indices = self.da1.getAO().app2petsc( np.arange(self.nx * self.ny, dtype=np.int32)) ByTmp = Xglobal.getValues(petsc_indices).copy().reshape( (self.ny, self.nx)).T scatter.destroy() Xglobal.destroy() if nfourier_x >= 0: # compute FFT, cut, compute inverse FFT BxFft = rfft(BxTmp, axis=1) ByFft = rfft(ByTmp, axis=1) BxFft[:, nfourier_x + 1:] = 0. ByFft[:, nfourier_x + 1:] = 0. BxTmp = irfft(BxFft, axis=1) ByTmp = irfft(ByFft, axis=1) if nfourier_y >= 0: BxFft = rfft(BxTmp, axis=0) ByFft = rfft(ByTmp, axis=0) BxFft[nfourier_y + 1:, :] = 0. ByFft[nfourier_y + 1:, :] = 0. BxTmp = irfft(BxFft, axis=0) ByTmp = irfft(ByFft, axis=0) Bx_arr = self.da1.getVecArray(self.Bx) By_arr = self.da1.getVecArray(self.By) Bx_arr[:, :] = BxTmp[xs:xe, ys:ye] By_arr[:, :] = ByTmp[xs:xe, ys:ye] Bx_arr = self.da1.getVecArray(self.Bx) By_arr = self.da1.getVecArray(self.By) Vx_arr = self.da1.getVecArray(self.Vx) Vy_arr = self.da1.getVecArray(self.Vy) x_arr = self.da7.getVecArray(self.x) x_arr[xs:xe, ys:ye, 0] = Vx_arr[xs:xe, ys:ye] x_arr[xs:xe, ys:ye, 1] = Vy_arr[xs:xe, ys:ye] x_arr[xs:xe, ys:ye, 2] = Bx_arr[xs:xe, ys:ye] x_arr[xs:xe, ys:ye, 3] = By_arr[xs:xe, ys:ye] # compure generalised magnetic induction self.da1.globalToLocal(self.Bx, self.localBx) self.da1.globalToLocal(self.By, self.localBy) self.da1.globalToLocal(self.Vx, self.localVx) self.da1.globalToLocal(self.Vy, self.localVy) Bx_arr = self.da1.getVecArray(self.localBx) By_arr = self.da1.getVecArray(self.localBy) Vx_arr = self.da1.getVecArray(self.localVx) Vy_arr = self.da1.getVecArray(self.localVy) Bix_arr = self.da1.getVecArray(self.Bix) Biy_arr = self.da1.getVecArray(self.Biy) for i in range(xs, xe): for j in range(ys, ye): Bix_arr[i, j] = self.derivatives.Bix(Bx_arr[...], By_arr[...], i - xs + 2, j - ys + 2, de) Biy_arr[i, j] = self.derivatives.Biy(Bx_arr[...], By_arr[...], i - xs + 2, j - ys + 2, de) # copy modified magnetic induction to solution vector x_arr = self.da7.getVecArray(self.x) x_arr[xs:xe, ys:ye, 4] = Bix_arr[xs:xe, ys:ye] x_arr[xs:xe, ys:ye, 5] = Biy_arr[xs:xe, ys:ye] # compute pressure self.da1.globalToLocal(self.Bix, self.localBix) self.da1.globalToLocal(self.Biy, self.localBiy) Bix_arr = self.da1.getVecArray(self.localBix) Biy_arr = self.da1.getVecArray(self.localBiy) P_arr = self.da1.getVecArray(self.P) for i in range(xs, xe): for j in range(ys, ye): P_arr[i,j] = init_data.pressure(xc_arr[i,j] + 0.5 * self.hx, yc_arr[i,j] + 0.5 * self.hy, self.hx, self.hy) \ - 0.5 * 0.25 * (Bix_arr[i,j] + Bix_arr[i+1,j]) * (Bx_arr[i,j] + Bx_arr[i+1,j]) \ - 0.5 * 0.25 * (Biy_arr[i,j] + Biy_arr[i,j+1]) * (By_arr[i,j] + By_arr[i,j+1]) \ # - 0.5 * (0.25 * (Vx_arr[i,j] + Vx_arr[i+1,j])**2 + 0.25 * (Vy_arr[i,j] + Vy_arr[i,j+1])**2) # copy pressure to solution vector x_arr = self.da7.getVecArray(self.x) x_arr[xs:xe, ys:ye, 6] = P_arr[xs:xe, ys:ye] # update solution history self.petsc_matrix.update_history(self.x) self.petsc_function.update_history(self.x) hdf5_filename = cfg['io']['hdf5_output'] if PETSc.COMM_WORLD.getRank() == 0: print(" Output: %s" % hdf5_filename) hdf5out = h5py.File(hdf5_filename, "w", driver="mpio", comm=PETSc.COMM_WORLD.tompi4py()) for cfg_group in self.cfg: for cfg_item in self.cfg[cfg_group]: if self.cfg[cfg_group][cfg_item] != None: value = self.cfg[cfg_group][cfg_item] else: value = "" hdf5out.attrs[cfg_group + "." + cfg_item] = value # if self.cfg["initial_data"]["python"] != None and self.cfg["initial_data"]["python"] != "": # python_input = open("runs/" + self.cfg['initial_data']['python'] + ".py", 'r') # python_file = python_input.read() # python_input.close() # else: # python_file = "" # hdf5out.attrs["initial_data.python_file"] = python_file hdf5out.close() # create HDF5 output file self.hdf5_viewer = PETSc.ViewerHDF5().create( hdf5_filename, mode=PETSc.Viewer.Mode.WRITE, comm=PETSc.COMM_WORLD) self.hdf5_viewer.pushGroup("/") # write grid data to hdf5 file coords_x = self.dax.getCoordinates() coords_y = self.day.getCoordinates() coords_x.setName('x') coords_y.setName('y') self.hdf5_viewer(coords_x) self.hdf5_viewer(coords_y) # write initial data to hdf5 file self.hdf5_viewer.setTimestep(0) self.save_hdf5_vectors()
def __init__(self, cfgfile): ''' Constructor ''' # load run config file cfg = Config(cfgfile) # timestep setup self.ht = cfg['grid']['ht'] # timestep size self.nt = cfg['grid']['nt'] # number of timesteps self.nsave = cfg['io']['nsave'] # save only every nsave'th timestep # grid setup self.nx = cfg['grid']['nx'] # number of points in x self.ny = cfg['grid']['ny'] # number of points in y Lx = cfg['grid']['Lx'] # spatial domain in x x1 = cfg['grid']['x1'] # x2 = cfg['grid']['x2'] # Ly = cfg['grid']['Ly'] # spatial domain in y y1 = cfg['grid']['y1'] # y2 = cfg['grid']['y2'] # if x1 != x2: Lx = x2 - x1 else: x1 = 0.0 x2 = Lx if y1 != y2: Ly = y2 - y1 else: y1 = 0.0 y2 = Ly self.hx = Lx / self.nx # gridstep size in x self.hy = Ly / self.ny # gridstep size in y self.time = PETSc.Vec().createMPI(1, PETSc.DECIDE, comm=PETSc.COMM_WORLD) self.time.setName('t') if PETSc.COMM_WORLD.getRank() == 0: self.time.setValue(0, 0.0) OptDB = PETSc.Options() OptDB.setValue('snes_rtol', cfg['solver']['petsc_snes_rtol']) OptDB.setValue('snes_atol', cfg['solver']['petsc_snes_atol']) OptDB.setValue('snes_stol', cfg['solver']['petsc_snes_stol']) OptDB.setValue('snes_max_it', cfg['solver']['petsc_snes_max_iter']) OptDB.setValue('ksp_rtol', cfg['solver']['petsc_ksp_rtol']) OptDB.setValue('ksp_atol', cfg['solver']['petsc_ksp_atol']) OptDB.setValue('ksp_max_it', cfg['solver']['petsc_ksp_max_iter']) # OptDB.setValue('ksp_monitor', '') # OptDB.setValue('snes_monitor', '') # # OptDB.setValue('log_info', '') # OptDB.setValue('log_summary', '') self.snes_rtol = cfg['solver']['petsc_snes_rtol'] self.snes_atol = cfg['solver']['petsc_snes_atol'] self.snes_max_iter = cfg['solver']['petsc_snes_max_iter'] # create DA with single dof self.da1 = PETSc.DA().create(dim=2, dof=1, sizes=[self.nx, self.ny], proc_sizes=[PETSc.DECIDE, PETSc.DECIDE], boundary_type=('periodic', 'periodic'), stencil_width=2, stencil_type='box') # create DA (dof = 4 for A, J, P, O) self.da4 = PETSc.DA().create(dim=2, dof=4, sizes=[self.nx, self.ny], proc_sizes=[PETSc.DECIDE, PETSc.DECIDE], boundary_type=('periodic', 'periodic'), stencil_width=2, stencil_type='box') # create DA for x grid self.dax = PETSc.DA().create(dim=1, dof=1, sizes=[self.nx], proc_sizes=[PETSc.DECIDE], boundary_type=('periodic')) # create DA for y grid self.day = PETSc.DA().create(dim=1, dof=1, sizes=[self.ny], proc_sizes=[PETSc.DECIDE], boundary_type=('periodic')) # initialise grid self.da1.setUniformCoordinates(xmin=x1, xmax=x2, ymin=y1, ymax=y2) self.da4.setUniformCoordinates(xmin=x1, xmax=x2, ymin=y1, ymax=y2) self.dax.setUniformCoordinates(xmin=x1, xmax=x2) self.day.setUniformCoordinates(xmin=y1, xmax=y2) # create RHS vector self.Ab = self.da1.createGlobalVec() self.Ob = self.da1.createGlobalVec() self.Pb = self.da1.createGlobalVec() # create vectors for magnetic and velocity field self.A = self.da1.createGlobalVec() # magnetic vector potential A self.J = self.da1.createGlobalVec() # current density J self.P = self.da1.createGlobalVec() # streaming function psi self.O = self.da1.createGlobalVec() # vorticity omega self.Bx = self.da1.createGlobalVec() self.By = self.da1.createGlobalVec() self.Vx = self.da1.createGlobalVec() self.Vy = self.da1.createGlobalVec() # set variable names self.A.setName('A') self.J.setName('J') self.P.setName('P') self.O.setName('O') self.Bx.setName('Bx') self.By.setName('By') self.Vx.setName('Vx') self.Vy.setName('Vy') # create nullspace self.poisson_nullspace = PETSc.NullSpace().create(constant=True) # create jacobian and matrix objects self.petsc_vorticity = PETScVorticity(self.da1, self.nx, self.ny, self.ht, self.hx, self.hy) self.petsc_ohmslaw = PETScOhmsLaw(self.da1, self.nx, self.ny, self.ht, self.hx, self.hy) self.petsc_poisson = PETScPoisson(self.da1, self.nx, self.ny, self.hx, self.hy) # initialise vorticity matrix self.vorticity_matrix = self.da1.createMat() self.vorticity_matrix.setOption( self.vorticity_matrix.Option.NEW_NONZERO_ALLOCATION_ERR, False) self.vorticity_matrix.setUp() # initialise Ohms's law matrix self.ohmslaw_matrix = self.da1.createMat() self.ohmslaw_matrix.setOption( self.ohmslaw_matrix.Option.NEW_NONZERO_ALLOCATION_ERR, False) self.ohmslaw_matrix.setUp() # initialise Poisson matrix self.poisson_matrix = self.da1.createMat() self.poisson_matrix.setOption( self.poisson_matrix.Option.NEW_NONZERO_ALLOCATION_ERR, False) self.poisson_matrix.setUp() self.poisson_matrix.setNullSpace(self.poisson_nullspace) # create nonlinear vorticity solver self.vorticity_snes = PETSc.SNES().create() self.vorticity_snes.setType('ksponly') self.vorticity_snes.setFunction(self.petsc_vorticity.snes_mult, self.Ob) self.vorticity_snes.setJacobian(self.update_vorticity_jacobian, self.vorticity_matrix) self.vorticity_snes.setFromOptions() # self.vorticity_snes.getKSP().setType('gmres') self.vorticity_snes.getKSP().setType('preonly') self.vorticity_snes.getKSP().getPC().setType('lu') self.vorticity_snes.getKSP().getPC().setFactorSolverPackage( solver_package) # create nonlinear Ohms's law solver self.ohmslaw_snes = PETSc.SNES().create() self.ohmslaw_snes.setType('ksponly') self.ohmslaw_snes.setFunction(self.petsc_ohmslaw.snes_mult, self.Ab) self.ohmslaw_snes.setJacobian(self.update_ohmslaw_jacobian, self.ohmslaw_matrix) self.ohmslaw_snes.setFromOptions() # self.ohmslaw_snes.getKSP().setType('gmres') self.ohmslaw_snes.getKSP().setType('preonly') self.ohmslaw_snes.getKSP().getPC().setType('lu') self.ohmslaw_snes.getKSP().getPC().setFactorSolverPackage( solver_package) # create linear Poisson solver self.poisson_ksp = PETSc.KSP().create() self.poisson_ksp.setFromOptions() self.poisson_ksp.setOperators(self.poisson_matrix) self.poisson_ksp.setType('preonly') self.poisson_ksp.getPC().setType('lu') self.poisson_ksp.getPC().setFactorSolverPackage(solver_package) # self.poisson_ksp.setNullSpace(self.poisson_nullspace) self.petsc_poisson.formMat(self.poisson_matrix) # create derivatives object self.derivatives = PETScDerivatives(self.da1, self.nx, self.ny, self.ht, self.hx, self.hy) # get coordinate vectors coords_x = self.dax.getCoordinates() coords_y = self.day.getCoordinates() # save x coordinate arrays scatter, xVec = PETSc.Scatter.toAll(coords_x) scatter.begin(coords_x, xVec, PETSc.InsertMode.INSERT, PETSc.ScatterMode.FORWARD) scatter.end(coords_x, xVec, PETSc.InsertMode.INSERT, PETSc.ScatterMode.FORWARD) xGrid = xVec.getValues(range(0, self.nx)).copy() scatter.destroy() xVec.destroy() # save y coordinate arrays scatter, yVec = PETSc.Scatter.toAll(coords_y) scatter.begin(coords_y, yVec, PETSc.InsertMode.INSERT, PETSc.ScatterMode.FORWARD) scatter.end(coords_y, yVec, PETSc.InsertMode.INSERT, PETSc.ScatterMode.FORWARD) yGrid = yVec.getValues(range(0, self.ny)).copy() scatter.destroy() yVec.destroy() # set initial data (xs, xe), (ys, ye) = self.da1.getRanges() A_arr = self.da1.getVecArray(self.A) P_arr = self.da1.getVecArray(self.P) init_data = __import__("runs." + cfg['initial_data']['python'], globals(), locals(), ['magnetic_A', 'velocity_P'], 0) for i in range(xs, xe): for j in range(ys, ye): A_arr[i, j] = init_data.magnetic_A(xGrid[i], yGrid[j], Lx, Ly) P_arr[i, j] = init_data.velocity_P(xGrid[i], yGrid[j], Lx, Ly) # Fourier Filtering self.nfourier = cfg['initial_data']['nfourier'] if self.nfourier >= 0: # obtain whole A vector everywhere scatter, Aglobal = PETSc.Scatter.toAll(self.A) scatter.begin(self.A, Aglobal, PETSc.InsertMode.INSERT, PETSc.ScatterMode.FORWARD) scatter.end(self.A, Aglobal, PETSc.InsertMode.INSERT, PETSc.ScatterMode.FORWARD) petsc_indices = self.da1.getAO().app2petsc( np.arange(self.nx * self.ny, dtype=np.int32)) Ainit = Aglobal.getValues(petsc_indices).copy().reshape( (self.ny, self.nx)) scatter.destroy() Aglobal.destroy() # compute FFT, cut, compute inverse FFT from scipy.fftpack import rfft, irfft Afft = rfft(Ainit, axis=1) Afft[:, 0] = 0. Afft[:, self.nfourier + 1:] = 0. A_arr = self.da1.getVecArray(self.A) A_arr[:, :] = irfft(Afft).T[xs:xe, ys:ye] # compute current and vorticity self.derivatives.laplace_vec(self.A, self.J, -1.) self.derivatives.laplace_vec(self.P, self.O, -1.) J_arr = self.da1.getVecArray(self.J) O_arr = self.da1.getVecArray(self.O) # add perturbations for i in range(xs, xe): for j in range(ys, ye): J_arr[i, j] += init_data.current_perturbation( xGrid[i], yGrid[j], Lx, Ly) O_arr[i, j] += init_data.vorticity_perturbation( xGrid[i], yGrid[j], Lx, Ly) # create HDF5 output file self.hdf5_viewer = PETSc.ViewerHDF5().create( cfg['io']['hdf5_output'], mode=PETSc.Viewer.Mode.WRITE, comm=PETSc.COMM_WORLD) self.hdf5_viewer.pushGroup("/") # write grid data to hdf5 file coords_x = self.dax.getCoordinates() coords_y = self.day.getCoordinates() coords_x.setName('x') coords_y.setName('y') self.hdf5_viewer(coords_x) self.hdf5_viewer(coords_y) # write initial data to hdf5 file self.save_to_hdf5(0)