def saveVec(self, vecs, step): name = self.h5name # self.logger.debug("saveVec %s" % name) ViewHDF5 = PETSc.ViewerHDF5() # Init. Viewer ViewHDF5.create(f"./{self.saveDir}/{name}-{step:05d}.h5", mode=PETSc.Viewer.Mode.WRITE, comm=self.comm) ViewHDF5.pushGroup('/fields') for vec in vecs: ViewHDF5.view(obj=vec) # Put PETSc object into the viewer ViewHDF5.destroy() # Destroy Viewer
def saveVec(self, vec, timeStep=None): """Save the vector.""" name = vec.getName() # self.logger.debug("saveVec %s" % name) ViewHDF5 = PETSc.ViewerHDF5() # Init. Viewer if timeStep is None: ViewHDF5.create(name + '.h5', mode=PETSc.Viewer.Mode.WRITE, comm=self.comm) else: ViewHDF5.create(name + '-%04d.h5' % timeStep, mode=PETSc.Viewer.Mode.WRITE, comm=self.comm) ViewHDF5.pushGroup('/fields') ViewHDF5.view(obj=vec) # Put PETSc object into the viewer ViewHDF5.destroy() # Destroy Viewer
def read_initial_data_from_hdf5(self): 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(hdf5_filename, mode=PETSc.Viewer.Mode.READ, comm=PETSc.COMM_WORLD) hdf5_viewer.setTimestep(timestep) self.A.load(hdf5_viewer) self.J.load(hdf5_viewer) self.P.load(hdf5_viewer) self.O.load(hdf5_viewer) hdf5_viewer.destroy()
def testViewLoadCycle(self): grank = PETSc.COMM_WORLD.rank for i in range(self.NTIMES): if i == 0: infname = self.infile() informt = self.informat() else: infname = self.outfile() informt = self.outformat() if self.HETEROGENEOUS: mycolor = (grank > self.NTIMES - i) else: mycolor = 0 try: import mpi4py except ImportError: self.skipTest( 'mpi4py') # throws special exception to signal test skip mpicomm = PETSc.COMM_WORLD.tompi4py() comm = PETSc.Comm(comm=mpicomm.Split(color=mycolor, key=grank)) if mycolor == 0: self.outputText("Begin cycle %d\n" % i, comm) plex = PETSc.DMPlex() vwr = PETSc.ViewerHDF5() # Create plex plex.create(comm=comm) plex.setName("DMPlex Object") # Load data from XDMF into dm in parallel vwr.create(infname, mode='r', comm=comm) vwr.pushFormat(format=informt) plex.load(viewer=vwr) plex.setOptionsPrefix("loaded_") plex.setFromOptions() vwr.popFormat() vwr.destroy() self.outputPlex(plex) # Test DM is indeed distributed flg = plex.isDistributed() self.outputText( "Loaded mesh distributed? %s\n" % str(flg).upper(), comm) # Interpolate plex.interpolate() plex.setOptionsPrefix("interpolated_") plex.setFromOptions() self.outputPlex(plex) # Redistribute part = plex.getPartitioner() part.setType(self.partitionerType()) _ = plex.distribute(overlap=0) plex.setOptionsPrefix("redistributed_") plex.setFromOptions() self.outputPlex(plex) # Save redistributed dm to XDMF in parallel vwr.create(self.outfile(), mode='w', comm=comm) vwr.pushFormat(format=self.outformat()) plex.setName("DMPlex Object") plex.view(viewer=vwr) vwr.popFormat() vwr.destroy() # Destroy plex plex.destroy() self.outputText("End cycle %d\n--------\n" % i, comm) PETSc.COMM_WORLD.Barrier() # Check that the output is identical to that of plex/tutorial/ex5.c. self.assertTrue( filecmp.cmp(self.tmp_output_file(), self.ref_output_file(), shallow=False), 'Contents of the files not the same.') PETSc.COMM_WORLD.Barrier()
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, 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 ''' # 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 __init__(self, cfgfile, runid=None, cfg=None): ''' Constructor ''' assert cfgfile is not None assert cfgfile is not "" # if runid is empty use timestamp if runid == None or runid == "": runid = datetime.datetime.fromtimestamp( time.time()).strftime("%y%m%d%H%M%S") # stencil width stencil = 2 # load run config file if cfg != None: self.cfg = cfg elif cfgfile != None and len(cfgfile) > 0: self.cfg = Config(cfgfile) else: if PETSc.COMM_WORLD.getRank() == 0: print("ERROR: No valid config file or object passed.") sys.exit() # determine solver modules if cfg['solver']['method'] == 'explicit': self.vlasov_module = None else: self.vlasov_module = "vlasov.solvers." if cfg['solver']['mode'] == 'split': self.vlasov_module += 'vlasov' else: self.vlasov_module += self.cfg['solver']['mode'] self.vlasov_module += '.' + "PETSc" if cfg['solver']['type'] == 'newton' or cfg['solver'][ 'type'] == 'nonlinear': self.vlasov_module += "NL" if cfg['solver']['mode'] == 'split': self.vlasov_module += "Vlasov" # self.vlasov_module += self.cfg['solver']['poisson_bracket'] # if cfg['solver']['preconditioner_type'] != None and cfg['solver']['preconditioner_scheme'] != None: # self.vlasov_module += self.cfg['solver']['preconditioner_scheme'] self.vlasov_module += self.cfg['solver']['timestepping'].upper() if not cfg.is_dissipation_none: if cfg['solver']['dissipation'] == 'double_bracket': self.vlasov_module += "DB" self.poisson_module = "vlasov.solvers.poisson.PETScPoisson" self.poisson_module += self.cfg['solver']['laplace_operator'] # importing solver modules if PETSc.COMM_WORLD.getRank() == 0: print("Loading Vlasov solver %s" % (self.vlasov_module)) print("Loading Poisson solver %s" % (self.poisson_module)) print("") if not self.cfg.is_preconditioner_none(): print("Using Preconditioner %s" % (self.cfg.get_preconditioner())) if not self.cfg.is_dissipation_none(): if self.cfg.is_dissipation_collisions(): print("Using Collision Operator %s" % (self.cfg.get_collision_operator())) if self.cfg.is_dissipation_double_bracket(): print("Using Double Bracket %s" % (self.cfg.get_double_bracket())) print("") self.vlasov_object = __import__(self.vlasov_module, globals(), locals(), ['PETScVlasovSolver'], 0) self.poisson_object = __import__(self.poisson_module, globals(), locals(), ['PETScPoissonSolver'], 0) # timestep setup ht = self.cfg['grid']['ht'] # timestep size nt = self.cfg['grid']['nt'] # number of timesteps self.nsave = self.cfg['io'][ 'nsave'] # save only every nsave'th timestep # grid setup nx = self.cfg['grid']['nx'] # number of points in x nv = self.cfg['grid']['nv'] # number of points in v L = self.cfg['grid']['L'] vMax = self.cfg['grid']['vmax'] vMin = self.cfg['grid']['vmin'] if vMin == None: vMin = -vMax Lx = L Lv = vMax - vMin hx = Lx / nx # gridstep size in x hv = Lv / (nv - 1) # gridstep size in v self.time = PETSc.Vec().createMPI(1, 1, comm=PETSc.COMM_WORLD) self.time.setName('t') if PETSc.COMM_WORLD.getRank() == 0: self.time.setValue(0, 0.0) self.solver_package = self.cfg['solver']['lu_package'] self.nInitial = self.cfg['solver'][ 'initial_iter'] # number of iterations for initial guess self.coll_freq = self.cfg['solver']['coll_freq'] # collision frequency self.coll_drag = self.cfg['solver']['coll_drag'] # drag factor self.coll_diff = self.cfg['solver']['coll_diff'] # diff factor self.charge = self.cfg['initial_data']['charge'] # particle charge self.mass = self.cfg['initial_data']['mass'] # particle mass output_directory = self.cfg['io']['output_dir'] if output_directory == None or output_directory == "": output_directory = "." tindex = cfgfile.rfind('/') run_filename = cfgfile[tindex:].replace('.cfg', '.') + runid hdf_out_filename = output_directory + '/' + run_filename + ".hdf5" cfg_out_filename = output_directory + '/' + run_filename + ".cfg" # hdf_in_filename = self.cfg['io']['hdf5_input'] # hdf_out_filename = self.cfg['io']['hdf5_output'] self.cfg.write_current_config(cfg_out_filename) # set initial guess method initial_guess_options = { None: self.initial_guess_none, "None": self.initial_guess_none, "": self.initial_guess_none, "rk4": self.initial_guess_rk4, "gear": self.initial_guess_gear, "symplectic2": self.initial_guess_symplectic2, "symplectic4": self.initial_guess_symplectic4, } self.initial_guess_method = initial_guess_options[self.cfg['solver'] ['initial_guess']] # set some PETSc 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('snes_rtol', self.cfg['solver']['petsc_snes_rtol']) OptDB.setValue('snes_atol', self.cfg['solver']['petsc_snes_atol']) OptDB.setValue('snes_stol', self.cfg['solver']['petsc_snes_stol']) OptDB.setValue('snes_max_it', self.cfg['solver']['petsc_snes_max_iter']) if PETSc.COMM_WORLD.getRank() == 0: print("Initialising Distributed Arrays.") # create DA for 2d grid (f only) # self.da1 = PETSc.DMDA().create(dim=2, dof=1, # sizes=[nx, nv], # proc_sizes=[1, PETSc.COMM_WORLD.getSize()], # boundary_type=['periodic', 'periodic'], # stencil_width=stencil, # stencil_type='box') # self.da1 = PETSc.DMDA().create(dim=2, dof=1, # sizes=[nx, nv], # proc_sizes=[PETSc.COMM_WORLD.getSize(), 1], # boundary_type=['periodic', 'ghosted'], # stencil_width=stencil, # stencil_type='box') self.da1 = PETSc.DMDA().create( dim=2, dof=1, sizes=[nx, nv], proc_sizes=[1, PETSc.COMM_WORLD.getSize()], boundary_type=['periodic', 'ghosted'], stencil_width=stencil, stencil_type='box') # self.da1 = PETSc.DMDA().create(dim=2, dof=1, # sizes=[nx, nv], # proc_sizes=[PETSc.DECIDE, 2], # boundary_type=['periodic', 'ghosted'], # stencil_width=stencil, # stencil_type='box') # self.da1 = PETSc.DMDA().create(dim=2, dof=1, # sizes=[nx, nv], # proc_sizes=[PETSc.DECIDE, PETSc.DECIDE], # boundary_type=['periodic', 'ghosted'], # stencil_width=stencil, # stencil_type='box') # create VIDA for x grid self.dax = PETSc.DMDA().create(dim=1, dof=1, sizes=[nx], proc_sizes=[PETSc.COMM_WORLD.getSize()], boundary_type=('periodic'), stencil_width=stencil, stencil_type='box') # create VIDA for y grid self.day = PETSc.DMDA().create(dim=1, dof=1, sizes=[nv], proc_sizes=[PETSc.COMM_WORLD.getSize()], boundary_type=('ghosted'), stencil_width=stencil, stencil_type='box') # initialise grid self.da1.setUniformCoordinates(xmin=0.0, xmax=Lx, ymin=vMin, ymax=vMax) self.dax.setUniformCoordinates(xmin=0.0, xmax=Lx) self.day.setUniformCoordinates(xmin=vMin, xmax=vMax) # get local index ranges (xs, xe), (ys, ye) = self.da1.getRanges() (xsx, xex), = self.dax.getRanges() # get coordinate vectors coords_x = self.dax.getCoordinates() coords_v = 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(nx)).copy() scatter.destroy() xVec.destroy() # save v coordinate arrays scatter, vVec = PETSc.Scatter.toAll(coords_v) scatter.begin(coords_v, vVec, PETSc.InsertMode.INSERT, PETSc.ScatterMode.FORWARD) scatter.end(coords_v, vVec, PETSc.InsertMode.INSERT, PETSc.ScatterMode.FORWARD) vGrid = vVec.getValues(range(nv)).copy() scatter.destroy() vVec.destroy() # create grid object self.grid = Grid().create(xGrid, vGrid, nt, nx, nv, ht, hx, hv, stencil) # create vectors for Hamiltonians self.h0 = self.da1.createGlobalVec() # kinetic Hamiltonian self.h1c = self.da1.createGlobalVec() # current potential Hamiltonian self.h2c = self.da1.createGlobalVec() # current external Hamiltonian self.h1h = self.da1.createGlobalVec() # previous potential Hamiltonian self.h2h = self.da1.createGlobalVec() # previous external Hamiltonian # distribution functions self.fl = self.da1.createGlobalVec() # last (k+1, n ) self.fc = self.da1.createGlobalVec() # current (k+1, n+1) self.fh = self.da1.createGlobalVec() # history (k) # distribution function solver vectors self.fb = self.da1.createGlobalVec() # right hand side # moments self.N = self.dax.createGlobalVec() # density self.U = self.dax.createGlobalVec() # velocity density self.E = self.dax.createGlobalVec() # energy density self.A = self.dax.createGlobalVec() # collision factor # local moments self.nc = PETSc.Vec().createSeq(nx) # current density self.uc = PETSc.Vec().createSeq(nx) # current velocity density self.ec = PETSc.Vec().createSeq(nx) # current energy density self.ac = PETSc.Vec().createSeq(nx) # current collision factor self.nh = PETSc.Vec().createSeq(nx) # history density self.uh = PETSc.Vec().createSeq(nx) # history velocity density self.eh = PETSc.Vec().createSeq(nx) # history energy density self.ah = PETSc.Vec().createSeq(nx) # history collision factor # internal potential self.pc_int = self.dax.createGlobalVec() # current self.ph_int = self.dax.createGlobalVec() # history # external potential self.pc_ext = self.dax.createGlobalVec() # current self.ph_ext = self.dax.createGlobalVec() # history # potential solver vectors self.pb = self.dax.createGlobalVec() # right hand side self.pn = self.dax.createGlobalVec() # null vector # set variable names self.h0.setName('h0') self.h1c.setName('h1') self.h2c.setName('h2') self.fc.setName('f') self.pc_int.setName('phi_int') self.pc_ext.setName('phi_ext') self.N.setName('n') self.U.setName('u') self.E.setName('e') # initialise nullspace basis vectors # the Poisson equation has a null space of all constant vectors # that needs to be removed to avoid jumpy potentials self.pn.set(1.) self.pn.normalize() self.p_nullspace = PETSc.NullSpace().create(constant=False, vectors=(self.pn, )) # create Toolbox self.toolbox = Toolbox(self.da1, self.dax, self.grid) # initialise kinetic hamiltonian if PETSc.COMM_WORLD.getRank() == 0: print("Initialising kinetic Hamiltonian.") self.toolbox.initialise_kinetic_hamiltonian(self.h0, self.mass) # create Arakawa initial guess solver object if PETSc.COMM_WORLD.getRank() == 0: print("Instantiating Initial Guess Objects.") self.arakawa_rk4 = PETScArakawaRungeKutta(self.cfg, self.da1, self.grid, self.h0, self.h1h, self.h2h, self.nInitial) self.arakawa_gear = PETScArakawaGear(self.cfg, self.da1, self.grid, self.h0, self.h1h, self.h2h, self.nInitial) self.arakawa_symplectic = PETScArakawaSymplectic( self.cfg, self.da1, self.grid, self.h0, self.h1h, self.h2h, self.nInitial) # create solver dummies self.vlasov_solver = None self.poisson_solver = None self.poisson_ksp = None if PETSc.COMM_WORLD.getRank() == 0: print() print("Run ID: %s" % runid) print() print("Config File: %s" % cfgfile) print("Output File: %s" % hdf_out_filename) print() print("nt = %i" % (self.grid.nt)) print("nx = %i" % (self.grid.nx)) print("nv = %i" % (self.grid.nv)) print() print("ht = %e" % (self.grid.ht)) print("hx = %e" % (self.grid.hx)) print("hv = %e" % (self.grid.hv)) print() print("xMin = %+12.6e" % (self.grid.xMin())) print("xMax = %+12.6e" % (self.grid.vMax())) print("vMin = %+12.6e" % (self.grid.vMin())) print("vMax = %+12.6e" % (self.grid.vMax())) print() print("nu = %7.1e" % (self.coll_freq)) print() print("CFL = %e" % (self.grid.hx / vMax)) print() print() # set initial data N0 = self.dax.createGlobalVec() T0 = self.dax.createGlobalVec() N0.setName('n0') T0.setName('T0') n0 = PETSc.Vec().createSeq(nx) t0 = PETSc.Vec().createSeq(nx) if self.cfg['initial_data']['distribution_python'] != None: init_data = __import__( "runs." + self.cfg['initial_data']['distribution_python'], globals(), locals(), ['distribution'], 0) if PETSc.COMM_WORLD.getRank() == 0: print( "Initialising distribution function with Python function.") self.toolbox.initialise_distribution_function( self.fc, init_data.distribution) N0.set(0.) T0.set(0.) else: N0_arr = self.dax.getVecArray(N0) T0_arr = self.dax.getVecArray(T0) if self.cfg['initial_data']['density_python'] != None: init_data = __import__( "runs." + self.cfg['initial_data']['density_python'], globals(), locals(), ['density'], 0) if PETSc.COMM_WORLD.getRank() == 0: print("Initialising density with Python function.") for i in range(xsx, xex): N0_arr[i] = init_data.density(self.grid.x[i], self.grid.xLength()) else: N0_arr[xsx:xex] = self.cfg['initial_data']['density'] if self.cfg['initial_data']['temperature_python'] != None: init_data = __import__( "runs." + self.cfg['initial_data']['temperature_python'], globals(), locals(), ['temperature'], 0) if PETSc.COMM_WORLD.getRank() == 0: print("Initialising temperature with Python function.") for i in range(xsx, xex): T0_arr[i] = init_data.temperature(self.grid.x[i]) else: T0_arr[xsx:xex] = self.cfg['initial_data']['temperature'] if PETSc.COMM_WORLD.getRank() == 0: print("Initialising distribution function with Maxwellian.") self.copy_xvec_to_seq(N0, n0) self.copy_xvec_to_seq(T0, t0) self.toolbox.initialise_distribution_nT(self.fc, n0, t0) # Fourier Filtering nfourier = cfg['initial_data']['nfourier'] if nfourier > 0: from scipy.fftpack import rfft, irfft if PETSc.COMM_WORLD.getRank() == 0: print("Fourier Filtering") # obtain whole f vector everywhere scatter, Xglobal = PETSc.Scatter.toAll(self.fc) scatter.begin(self.fc, Xglobal, PETSc.InsertMode.INSERT, PETSc.ScatterMode.FORWARD) scatter.end(self.fc, Xglobal, PETSc.InsertMode.INSERT, PETSc.ScatterMode.FORWARD) petsc_indices = self.da1.getAO().app2petsc( np.arange(nx * nv, dtype=np.int32)) fTmp = Xglobal.getValues(petsc_indices).copy().reshape((nv, nx)).T scatter.destroy() Xglobal.destroy() # filter distribution function fFft = rfft(fTmp, axis=0) fFft[nfourier:, :] = 0. fTmp = irfft(fFft, axis=0) # store filtered distribution function f_arr = self.da1.getVecArray(self.fc) f_arr[:, :] = fTmp[xs:xe, ys:ye] # normalise f self.normalise_distribution_function() # calculate potential and moments if PETSc.COMM_WORLD.getRank() == 0: print("Calculate initial potential and moments.") self.calculate_moments(output=False) # initialise Gear History if self.cfg['solver']['initial_guess'] == "gear": self.arakawa_gear.initialise_history(self.fc) # check for external potential if self.cfg['initial_data']['external_python'] != None: if PETSc.COMM_WORLD.getRank() == 0: print("Calculate external potential.") external_data = __import__( "runs." + self.cfg['initial_data']['external_python'], globals(), locals(), ['external'], 0) self.external = external_data.external else: self.external = None # calculate external potential self.calculate_external(0.) # create HDF5 output file if PETSc.COMM_WORLD.getRank() == 0: print("Create HDF5 output file.") # use h5py to store attributes hdf5out = h5py.File(hdf_out_filename, 'w', driver='mpio', comm=PETSc.COMM_WORLD.tompi4py()) hdf5out.attrs['charge'] = self.charge hdf5out.close() # create PETSc HDF5 viewer self.hdf5_viewer = PETSc.ViewerHDF5().create( hdf_out_filename, mode=PETSc.Viewer.Mode.APPEND, comm=PETSc.COMM_WORLD) self.hdf5_viewer.pushGroup("/") if PETSc.COMM_WORLD.getRank() == 0: print("Saving initial data to HDF5.") # write grid data to hdf5 file coords_x.setName('x') coords_v.setName('v') self.hdf5_viewer(coords_x) self.hdf5_viewer(coords_v) # write initial data to hdf5 file # self.hdf5_viewer(N0) # self.hdf5_viewer(T0) # save to hdf5 self.hdf5_viewer.setTimestep(0) self.save_hdf5_vectors() if PETSc.COMM_WORLD.getRank() == 0: print("run_base.py: initialisation done.") print("")
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)
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()