def __init__(self, cfgfile, runid):
super().__init__(cfgfile, runid)
self.arakawa_lf = PETScArakawaLeapfrog(self.da1, self.dax,
self.h0, self.vGrid,
self.nx, self.nv, self.ht, self.hx, self.hv)
self.arakawa_rk = PETScArakawaRungeKutta(self.da1, self.dax,
self.h0, self.vGrid,
self.nx, self.nv, -self.ht, self.hx, self.hv)
self.f.copy(self.fh)
self.arakawa_rk.rk18_J4(self.f, self.h1)
self.calculate_moments(output=False)
# self.calculate_external(-self.ht)
# self.arakawa_lf.update_external(self.p_ext)
self.arakawa_lf.update_history(self.f, self.h1)
self.fh.copy(self.f)
self.calculate_moments(output=False)
# self.calculate_external(0.)
# self.arakawa_lf.update_external(self.p_ext)
self.arakawa_lf.update_history(self.f, self.h1)
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, runid):
super().__init__(cfgfile, runid)
self.arakawa_rk = PETScArakawaRungeKutta(self.da1, self.dax, self.h0,
self.vGrid, self.nx, self.nv,
self.ht, self.hx, self.hv)