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()
