def __init__(self, out_dir):
self.out_dir = out_dir
self.pickle_dir = '{}/pickles'.format(out_dir)
self.fig_dir = '{}/figures'.format(out_dir)
os.chdir(out_dir)
os.system('mkdir -p {}'.format(self.pickle_dir))
os.system('mkdir -p {}'.format(self.fig_dir))
self.dat = DataFile('BOUT.dmp.0.nc')
def listKeys(self, simIndex=0, simType='1-base'):
if simType == '1-base':
os.chdir('{}/{}'.format(self.outDir, simIndex))
else:
os.chdir('{}/{}/{}'.format(self.outDir, simIndex, simType))
datFile = DataFile('BOUT.dmp.0.nc')
self.datFile = datFile
return datFile.keys()
def file_import(name):
f = DataFile(name) # Open file
varlist = f.list() # Get list of all variables in file
data = {} # Create empty dictionary
for v in varlist:
data[v] = f.read(v)
f.close()
return data
def check_test():
print("Checking output")
numFailures = 0
numTests = 0
try:
run = BoutOutputs(runOutputDir, info=False, yguards=True)
except TypeError:
# Option not implemented in boutdata.data
run = BoutOutputs(runOutputDir, yguards=True)
runExpected = DataFile(runExpectedOutput)
# Get number of guard cells
m_guards = (run["MXG"], run["MYG"])
m_guards_expected = (runExpected["MXG"], runExpected["MYG"])
# Get names of evolving variables, which we will test
try:
evolvingVariables = run.evolvingVariables()
except AttributeError:
# This part should be deleted once BOUT++ repo is updated so that the above works everywhere
print("Warning: Updated boutdata.data not found")
from get_evolving_fields import get_evolving_fields
evolvingVariables = get_evolving_fields(run)
# Test output
for name in evolvingVariables:
if len(run[name].shape) == 1:
# scalars are diagnostic outputs like wall-time, so not useful to test
continue
# exclude second x guard cells as they are not used and may not always be set consistently
data = testfield_slice(run[name], m_guards)
expectedData = testfield_slice(runExpected.read(name),
m_guards_expected)
diff_max, norm_max = testfield_max(data, expectedData)
numTests = numTests + 1
if diff_max / norm_max > tolerance and diff_max > abs_tolerance:
numFailures = numFailures + 1
print("FAILURE: Test of max error " + str(numTests) + " (" + name +
") failed, with diff_max/norm_max=" +
str(diff_max / norm_max) + " and diff_max=" + str(diff_max))
else:
print("Test of max error " + str(numTests) + " (" + name +
") passed, with diff_max/norm_max=" +
str(diff_max / norm_max) + " and diff_max=" + str(diff_max))
diff_mean, norm_mean = testfield_mean(data, expectedData)
print("Test of mean error " + str(numTests) + " (" + name +
") diff_mean/norm_mean=" + str(diff_mean / norm_mean) +
" and diff_mean=" + str(diff_mean))
print(
str(numTests - numFailures) + "/" + str(numTests) +
" tests passed in " + testname)
return numFailures, numTests
def calc_curvilinear_curvature(fname, field, grid):
from scipy.signal import savgol_filter
f = DataFile(str(fname), write=True)
B = f.read("B")
dBydz = np.zeros(np.shape(B))
dBydx = np.zeros(np.shape(B))
dBxdz = np.zeros(np.shape(B))
dBzdx = np.zeros(np.shape(B))
dx = grid.metric()["dx"]
dz = grid.metric()["dz"]
g_11 = grid.metric()["g_xx"]
g_22 = grid.metric()["g_yy"]
g_33 = grid.metric()["g_zz"]
g_12 = 0.0
g_13 = grid.metric()["g_xz"]
g_23 = 0.0
J = np.sqrt(g_11 * (g_22 * g_33 - g_23 * g_23) + g_12 *
(g_13 * g_23 - g_12 * g_33) + g_13 *
(g_12 * g_23 - g_22 * g_23))
Bx_smooth = np.zeros(B.shape)
By_smooth = np.zeros(B.shape)
Bz_smooth = np.zeros(B.shape)
for y in np.arange(0, B.shape[1]):
pol, _ = grid.getPoloidalGrid(y)
R = pol.R
Z = pol.Z
for x in np.arange(0, B.shape[0]):
Bx_smooth[x, y, :] = savgol_filter(
field.Bxfunc(R[x, :], y, Z[x, :]),
np.int(np.ceil(B.shape[-1] / 21) // 2 * 2 + 1), 5)
By_smooth[x, y, :] = savgol_filter(
field.Byfunc(R[x, :], y, Z[x, :]),
np.int(np.ceil(B.shape[-1] / 21) // 2 * 2 + 1), 5)
dBydz[x, y, :] = calc.deriv(By_smooth[x, y, :]) / dz[x, y, :]
dBxdz[x, y, :] = calc.deriv(Bx_smooth[x, y, :]) / dz[x, y, :]
for z in np.arange(0, B.shape[-1]):
Bz_smooth[:, y, z] = savgol_filter(
field.Bzfunc(R[:, z], y, Z[:, z]),
np.int(np.ceil(B.shape[0] / 7) // 2 * 2 + 1), 5)
dBzdx[:, y, z] = calc.deriv(Bz_smooth[:, y, z]) / dx[:, y, z]
dBydx[:, y, z] = calc.deriv(By_smooth[:, y, z]) / dx[:, y, z]
bxcvx = (-1 / J) * (dBydz / B**2.)
bxcvy = (1 / J) * ((dBxdz - dBzdx) / B**2.)
bxcvz = (1 / J) * (dBydx / B**2.)
f.write('bxcvz', bxcvz)
f.write('bxcvx', bxcvx)
f.write('bxcvy', bxcvy)
f.close()
def scalevar(var, factor, path="."):
"""
Scales a variable by a given factor, modifying
restart files in place
Inputs
------
var Name of the variable (string)
factor Factor to multiply (float)
path Path to the restart files
Returns
-------
None
"""
file_list = glob.glob(os.path.join(path, "BOUT.restart.*"))
nfiles = len(file_list)
print("Number of restart files: %d" % (nfiles, ))
for file in file_list:
print(file)
with DataFile(file, write=True) as d:
d[var] = d[var] * factor
def scalevar(var, factor, path="."):
"""Scales a variable by a given factor, modifying restart files in
place
.. warning:: Modifies restart files in place! This is in contrast
to most of the functions in this module!
Parameters
----------
var : str
Name of the variable
factor : float
Factor to multiply
path : str, optional
Path to the restart files (default: current directory)
"""
file_list = glob.glob(os.path.join(path, "BOUT.restart.*"))
nfiles = len(file_list)
print("Number of restart files: %d" % (nfiles, ))
for file in file_list:
print(file)
with DataFile(file, write=True) as d:
d[var] = d[var] * factor
def create_cache(path, prefix):
"""Create a list of DataFile objects to be passed repeatedly to
collect.
Parameters
----------
path : str
Path to data files
prefix : str
File prefix
Returns
-------
namedtuple : (list of str, bool, str, list of :py:obj:`~boututils.datafile.DataFile`)
The cache of DataFiles in a namedtuple along with the file_list,
and parallel and suffix attributes
"""
# define namedtuple to return as the result
from collections import namedtuple
datafile_cache_tuple = namedtuple(
"datafile_cache", ["file_list", "parallel", "suffix", "datafile_list"])
file_list, parallel, suffix = findFiles(path, prefix)
cache = []
for f in file_list:
cache.append(DataFile(f))
return datafile_cache_tuple(file_list=file_list, parallel=parallel, suffix=suffix, datafile_list=cache)
def checkForDifferentLengths(dmpFiles):
#{{{docstring
"""
Checks that the length of the variables are the same.
Paramters
---------
dmpFiles : iterable
Iterable containing the paths to the restart files.
Returns
-------
maxDiff : int
Maximum difference between the time indices
shortestCommonLen : int
Shortest common time indices
"""
#}}}
print("\nChecking if the output has the same number of outputs")
curMax = 0
curMin = float("inf")
for d in dmpFiles:
print("\nChecking {}".format(d))
with DataFile(d) as dmp:
tLen = len(dmp.read("t_array"))
curMax = curMax if curMax > tLen else tLen
curMin = curMin if curMin < tLen else tLen
maxDiff = curMax - curMin
shortestCommonLen = curMin
return maxDiff, shortestCommonLen
def checkForCorruption(restartFiles):
#{{{docstring
"""
Check for corruption by checking the field mean.
WARNING: This is not very water-proof.
Corruption could in theory still have occured, so use with care.
Paramters
---------
restartFiles : iterable
Iterable containing the paths to the restart files.
"""
#}}}
print("\nChecking for corrupted restart files by checking the field mean")
for r in restartFiles:
print("\nChecking {}".format(r))
with DataFile(r) as restart:
for var in restart.list():
if restart.ndims(var) == 3:
mean = restart.read(var).mean()
if np.isclose(mean, 0):
message="{} has a zero mean. File could be corrupted.".\
format(var)
raise RuntimeError(message)
else:
print("{} PASSED with a mean of {}.".format(var, mean))
def getUniformSpacing(path, coordinate, xguards=False, yguards=False):
#{{{docstring
"""
Fastest way to obtain the grid spacing assuming equidistant grid
Parameters
----------
path : str
Path to read from
coordinate : str
Coordinate to return size of
xguards : bool
If the ghost points in x should be included
yguards : bool
If the ghost points in y should be included
Returns
-------
spacing : array-like
The grid spacing
"""
#}}}
with DataFile(os.path.join(path, "BOUT.dmp.0.nc")) as f:
if coordinate == "x" or coordinate == "y":
if coordinate == "x":
# dx
spacing = f.read("dx")
elif coordinate == "y":
# dy
spacing = f.read("dy")
shape = spacing.shape
xSize = shape[0] - 2 * int(f.read("MXG"))
ySize = shape[1] - 2 * int(f.read("MYG"))
if not (xguards) and not (yguards):
spacingEmpty = np.empty((xSize*int(f.read("NXPE")),\
ySize*int(f.read("NYPE"))))
elif xguards and not (yguards):
spacingEmpty = np.empty((xSize*int(f.read("NXPE"))+\
2*int(f.read("MXG")),\
ySize*int(f.read("NYPE"))))
elif not (xguards) and yguards:
spacingEmpty = np.empty((xSize*int(f.read("NXPE")),\
ySize*int(f.read("NYPE"))+\
2*int(f.read("MYG"))))
elif xguards and yguards:
spacingEmpty = np.empty((xSize*int(f.read("NXPE"))+\
2*int(f.read("MXG")),\
ySize*int(f.read("NYPE"))+\
2*int(f.read("MYG"))))
spacingEmpty.fill(spacing[0, 0])
spacing = spacingEmpty
elif coordinate == "z":
# dz
spacing = f.read("dz")
else:
raise ValueError("Unknown coordinate {}".format(coordinate))
return spacing
def rmSpuriousTime(newFiles, dmpFiles, shortestCommonLen):
#{{{docstring
"""
Will remove the suprious times in the dump files.
Parameters
----------
newFiles : iterable
Iterable containing the paths to the new files to be created.
dmpFiles : iterable
Iterable containing the paths to the dump files.
shortestCommonLen : int
Shortest common time indices
"""
#}}}
for n, d in zip(newFiles, dmpFiles):
print("\nRemoving spurious time {}".format(d))
# Open the restart file in read mode and create the restart file
with DataFile(n, write=True, create=True) as newF,\
DataFile(d) as dmp:
# Loop over the variables in the dmp file
# Put a 4d variable in the front
theList = dmp.list()
ind = theList.index("lnN")
theList[0], theList[ind] = theList[ind], theList[0]
for var in theList:
# Read the data
dData = dmp.read(var)
# Find 4D variables and time traces
if dmp.ndims(var) == 4 or dmp.ndims(var) == 1:
print(" Using first {} timepoints in {}".\
format(shortestCommonLen, var))
# Read from restart
newData = dData[:shortestCommonLen, ...]
elif var == "iteration":
print(" Fixing 'iteration'")
newData = shortestCommonLen - 2
else:
print(" Nothing to be done for {}".format(var))
newData = dData.copy()
newF.write(var, newData, info=True)
print("{} written".format(n))
def getDataFile(i):
"""Get the DataFile from the cache, if present, otherwise open the
DataFile
"""
if datafile_cache is not None:
return datafile_cache.datafile_list[i]
else:
return DataFile(file_list[i])
def rmLastTime(newFiles, dmpFiles):
#{{{docstring
"""
Will remove the last time index in the dump files.
Parameters
----------
newFiles : iterable
Iterable containing the paths to the new files to be created.
dmpFiles : iterable
Iterable containing the paths to the dump files.
maxDiff : int
Maximum difference between the time indices
"""
#}}}
for n, d in zip(newFiles, dmpFiles):
print("\nRemoving last time point in {}".format(d))
# Open the restart file in read mode and create the restart file
with DataFile(n, write=True, create=True) as newF,\
DataFile(d) as dmp:
# Loop over the variables in the dmp file
# Put a 4d variable in the front
theList = dmp.list()
ind = theList.index("lnN")
theList[0], theList[ind] = theList[ind], theList[0]
for var in theList:
# Read the data
dData = dmp.read(var)
# Find 4D variables and time traces
if dmp.ndims(var) == 4 or dmp.ndims(var) == 1:
print(" Removing last time in " + var)
# Read from restart
newData = dData[:-1, ...]
elif var == "iteration":
print(" Fixing 'iteration'")
newData = dData - 1
else:
print(" Nothing to be done for {}".format(var))
newData = dData.copy()
newF.write(var, newData, info=True)
print("{} written".format(n))
def createNewProfile(baseGrid, newProfile, offset, pedestal):
os.system("cp {} {}".format(baseGrid, newProfile))
profile(newProfile, "Te0", 5, 100, hwid=0.1, alpha=0.1)
profile(newProfile, "Ti0", 5, 100, hwid=0.1, alpha=0.1)
profile(newProfile, "Ne0", offset, pedestal, hwid=0.1, alpha=0.1)
with DataFile(newProfile, write=True) as d:
d["Ni0"] = d["Ne0"]
print("generated {}".format(newProfile))
def file_import(name):
"""Read all variables from file into a dictionary
Parameters
----------
name : str
Name of file to read
Returns
-------
dict
Dictionary containing all the variables in the file
"""
f = DataFile(name) # Open file
varlist = f.list() # Get list of all variables in file
data = {} # Create empty dictionary
for v in varlist:
data[v] = f.read(v)
f.close()
return data
def file_import(name):
f = DataFile(name) # Open file
varlist = f.list() # Get list of all variables in file
data = {} # Create empty dictionary
for v in varlist:
data[v] = f.read(v)
f.close()
return data
def attributes(varname, path=".", prefix="BOUT.dmp"):
"""Return a dictionary of variable attributes in an output file
Parameters
----------
varname : str
Name of the variable
path : str, optional
Path to data files (default: ".")
prefix : str, optional
File prefix (default: "BOUT.dmp")
Returns
-------
dict
A dictionary of attributes of varname
"""
# Search for BOUT++ dump files in NetCDF format
file_list, _, _ = findFiles(path, prefix)
# Read data from the first file
f = DataFile(file_list[0])
return f.attributes(varname)
def attributes(varname, path=".", prefix="BOUT.dmp"):
"""Return a dictionary of variable attributes in an output file
Parameters
----------
varname : str
Name of the variable
path : str, optional
Path to data files (default: ".")
prefix : str, optional
File prefix (default: "BOUT.dmp")
Returns
-------
dict
A dictionary of attributes of varname
"""
# Search for BOUT++ dump files in NetCDF format
file_list, _, _ = findFiles(path, prefix)
# Read data from the first file
f = DataFile(file_list[0])
return f.attributes(varname)
def profile(filename, name, offset, pedestal, hwid=0.1, alpha=0.1):
"""
Calculate a radial profile, and add to file
"""
with DataFile(filename, write=True) as d:
nx = d["nx"]
ny = d["ny"]
x = np.arange(nx)
ix = d["ixseps1"]
prof = mtanh_profile(x, ix, hwid * nx, offset, pedestal, alpha)
prof2d = np.zeros([nx, ny])
for y in range(ny):
prof2d[:, y] = prof
# Handle X-points
# Lower inner PF region
j11 = d["jyseps1_1"]
if j11 >= 0:
# Reflect around separatrix
ix = d["ixseps1"]
for x in range(0, ix):
prof2d[x, 0:(j11 + 1)] = prof2d[np.clip(2 * ix - x, 0, nx - 1),
0:(j11 + 1)]
# Lower outer PF region
j22 = d["jyseps2_2"]
if j22 < ny - 1:
ix = d["ixseps1"]
for x in range(0, ix):
prof2d[x, (j22 + 1):] = prof2d[np.clip(2 * ix - x, 0, nx - 1),
(j22 + 1):]
# Upper PF region
j21 = d["jyseps2_1"]
j12 = d["jyseps1_2"]
if j21 != j12:
ix = d["ixseps2"]
for x in range(0, ix):
prof2d[x,
(j21 + 1):(j12 +
1)] = prof2d[np.clip(2 * ix - x, 0, nx - 1),
(j21 + 1):(j12 + 1)]
d.write(name, prof2d)
def addvar(var, value, path="."):
"""Adds a variable with constant value to all restart files.
.. warning:: Modifies restart files in place! This is in contrast
to most of the functions in this module!
This is useful for restarting simulations whilst turning on new
equations. By default BOUT++ throws an error if an evolving
variable is not in the restart file. By setting an option the
variable can be set to zero. This allows it to start with a
non-zero value.
Parameters
----------
var : str
The name of the variable to add
value : float
Constant value for the variable
path : str, optional
Input path to data files (default: current directory)
"""
file_list = glob.glob(os.path.join(path, "BOUT.restart.*"))
nfiles = len(file_list)
print("Number of restart files: %d" % (nfiles, ))
# Loop through all the restart files
for filename in file_list:
print(filename)
# Open the restart file for writing (modification)
with DataFile(filename, write=True) as df:
size = None
# Find a 3D variable and get its size
for varname in df.list():
size = df.size(varname)
if len(size) == 3:
break
if size is None:
raise Exception("no 3D variables found")
# Create a new 3D array with input value
data = np.zeros(size) + value
# Set the variable in the NetCDF file
df.write(var, data)
def collectData(self, quant, simIndex=0, simType='1-base'):
try:
quant2 = np.squeeze(self.unPickle(quant, simIndex, simType))
except(FileNotFoundError):
print('{} has not been pickled'.format(quant))
if simType == '1-base':
os.chdir('{}/{}'.format(self.outDir, simIndex))
else:
os.chdir('{}/{}/{}'.format(self.outDir, simIndex, simType))
try:
quant2 = np.squeeze(collect(quant))
except(ValueError):
print('quant in gridFile')
self.gridFile = fnmatch.filter(next(os.walk('./'))[2],
'*profile*')[0]
grid_dat = DataFile(self.gridFile)
quant2 = grid_dat[quant]
return quant2
def getMYG(path):
#{{{docstring
"""
Fastest way to obtain MYG
Parameters
----------
path : str
Path to read from
Returns
-------
MYG : int
Number of ghost points in y
"""
#}}}
with DataFile(os.path.join(path, "BOUT.dmp.0.nc")) as f:
return f.read("MYG")
def getGridSizes(path, coordinate, varName="lnN", includeGhost=False):
#{{{docstring
"""
Fastest way to obtain coordinate sizes.
Parameters
----------
path : str
Path to read from
coordinate : ["x"|"y"|"z"|"t"]
Coordinate to return size of
varName : str
Field to get the size of
includeGhost : bool
If the ghost points should be included
Returns
-------
coordinateSize : int
Size of the desired coordinate
"""
#}}}
with DataFile(os.path.join(path, "BOUT.dmp.0.nc")) as f:
if coordinate == "x":
# nx
coordinateSize =\
(f.size(varName)[1] - 2*int(f.read("MXG")))*f.read("NXPE")
if includeGhost:
coordinateSize += 2 * int(f.read("MXG"))
elif coordinate == "y":
# ny
coordinateSize =\
(f.size(varName)[2] - 2*int(f.read("MYG")))*f.read("NYPE")
if includeGhost:
coordinateSize += 2 * int(f.read("MYG"))
elif coordinate == "z":
# nz
coordinateSize = (f.size(varName)[3])
elif coordinate == "z":
coordinateSize = (f.size(varName)[0])
else:
raise ValueError("Unknown coordinate {}".format(coordinate))
return coordinateSize
def calc_qPar(self, simIndex=0, simType='3-addC'):
if simType == '1-base':
os.chdir('{}/{}'.format(self.outDir, simIndex))
else:
os.chdir('{}/{}/{}'.format(self.outDir, simIndex, simType))
datFile = DataFile('BOUT.dmp.0.nc')
Tnorm = float(datFile['Tnorm'])
Nnorm = float(datFile['Nnorm'])
gamma_e = 4
gamma_i = 2.5
mi = 3.34524e-27 # 2*Mp - deuterium
e = 1.6e-19
Te = self.collectData('Telim', simIndex, simType)[-1, :, -1]*Tnorm
Ti = self.collectData('Tilim', simIndex, simType)[-1, :, -1]*Tnorm
n = self.collectData('Ne', simIndex, simType)[-1, :, -1]*Nnorm
Cs = np.sqrt(Te + (5/3)*Ti)*np.sqrt(e/mi)
q_e = gamma_e * n * e * Te * Cs
q_i = gamma_i * n * e * Ti * Cs
return q_e + q_i
def dimensions(varname, path=".", prefix="BOUT.dmp"):
"""Return the names of dimensions of a variable in an output file
Parameters
----------
varname : str
Name of the variable
path : str, optional
Path to data files (default: ".")
prefix : str, optional
File prefix (default: "BOUT.dmp")
Returns
-------
tuple of strs
The elements of the tuple give the names of corresponding variable
dimensions
"""
file_list, _, _ = findFiles(path, prefix)
return DataFile(file_list[0]).dimensions(varname)
def __init__(self, path=".", prefix="BOUT.dmp"):
"""
Initialise BoutOutputs object
"""
self._path = path
self._prefix = prefix
# Label for this data
self.label = path
# Check that the path contains some data
file_list = glob.glob(os.path.join(path, prefix + "*.nc"))
if len(file_list) == 0:
raise ValueError("ERROR: No data files found")
# Available variables
self.varNames = []
with DataFile(file_list[0]) as f:
# Get variable names
self.varNames = f.keys()
def collectTime(paths, tInd=None):
#{{{docstring
"""
Collects the time
Parameters
-----------
paths : iterable of strings
What path to use when collecting the variable. Must be in
ascending temporal order as the variable will be
concatenated.
tInd : [None|tuple]
Start and end of the time if not None
Returns
-------
time : 1d-array
Array of the time at the difference time indices
"""
#}}}
# Initialize
time = None
for path in paths:
with DataFile(os.path.join(path, "BOUT.dmp.0.nc")) as f:
if time is None:
time = f.read("t_array")
else:
# Remove first point in time in the current time as this
# is the same as the last of the previous
time = np.concatenate((time, f.read("t_array")[1:]), axis=0)
if tInd is not None:
# NOTE: +1 since the collect ranges is INCLUSIVE, i.e. not working
# like a python slice
lastT = tInd[1] + 1 if tInd[1] is not None else None
time = time[tInd[0]:lastT]
return time
def gridData(self, simIndex=0):
os.chdir('{}/{}'.format(self.outDir, simIndex))
self.gridFile = fnmatch.filter(next(os.walk('./'))[2],
'*profile*')[0]
grid_dat = DataFile(self.gridFile)
self.grid_dat = grid_dat
self.j11 = int(grid_dat["jyseps1_1"])
self.j12 = int(grid_dat["jyseps1_2"])
self.j21 = int(grid_dat["jyseps2_1"])
self.j22 = int(grid_dat["jyseps2_2"])
self.ix1 = int(grid_dat["ixseps1"])
self.ix2 = int(grid_dat["ixseps2"])
try:
self.nin = int(grid_dat["ny_inner"])
except(KeyError):
self.nin = self.j12
self.nx = int(grid_dat["nx"])
self.ny = int(grid_dat["ny"])
self.R = grid_dat['Rxy']
self.Z = grid_dat['Zxy']
R2 = self.R[:, self.j12:self.j22]
self.outMid_idx = self.j12 + np.where(R2 == np.amax(R2))[1][0]
def checkProfiles(gridFiles=[], densities=[]):
if len(densities) < 1:
densities = np.zeros(len(gridFiles))
grids = []
ne = []
te = []
for i in gridFiles:
grd = DataFile(i)
grids.append(grd)
ne.append(grd["Ne0"] * 1e20)
te.append(grd["Te0"])
ix1 = grids[0]["ixseps1"]
j11 = grids[0]["jyseps1_1"]
j12 = grids[0]["jyseps1_2"]
j22 = grids[0]["jyseps2_2"]
j21 = grids[0]["jyseps2_1"]
mid = int((j12 + j22) / 2)
colors = getDistinctColors(len(gridFiles))
plt.figure(1)
plt.axvline(x=ix1, color="black", linestyle="--")
for i in range(len(ne)):
plt.plot(ne[i][:, mid], color=colors[i], label=gridFiles[i])
print(densities[i])
plt.axhline(y=eval("{}e19".format(densities[i])),
color=colors[i],
linestyle="--")
plt.figure(2)
plt.axvline(x=ix1, color="black", linestyle="--")
for i in range(len(te)):
plt.plot(te[i][:, mid], color=colors[i], label=gridFiles[i])
plt.axhline(y=te[i][ix1, mid], color=colors[i], linestyle="--")
plt.legend()
plt.show()
def addnoise(path=".", var=None, scale=1e-5):
"""Add random noise to restart files
.. warning:: Modifies restart files in place! This is in contrast
to most of the functions in this module!
Parameters
----------
path : str, optional
Path to restart files (default: current directory)
var : str, optional
The variable to modify. By default all 3D variables are modified
scale : float
Amplitude of the noise. Gaussian noise is used, with zero mean
and this parameter as the standard deviation
"""
file_list = glob.glob(os.path.join(path, "BOUT.restart.*"))
nfiles = len(file_list)
print("Number of restart files: %d" % (nfiles, ))
for file in file_list:
print(file)
with DataFile(file, write=True) as d:
if var is None:
for v in d.list():
if d.ndims(v) == 3:
print(" -> " + v)
data = d.read(v, asBoutArray=True)
data += normal(scale=scale, size=data.shape)
d.write(v, data)
else:
# Modify a single variable
print(" -> " + var)
data = d.read(var)
data += normal(scale=scale, size=data.shape)
d.write(var, data)
def collectSteadyN(steadyStatePath, yInd):
#{{{docstring
"""
Collects n in the steady state.
Parameters
----------
steadyStatePath : str
Path to collect from
yInd : int
Index for the parallel coordinate.
Returns
-------
n : array-4d
The density at the steady state
"""
#}}}
# Check last t index
with DataFile(os.path.join(steadyStatePath, "BOUT.dmp.0.nc")) as f:
tLast = len(f.read("t_array")) - 1
# In the steady state, the max gradient in "n" is the same
# throughout in the domain, so we use yInd=0, zInd=0 in the
# collect
lnN = collect("lnN",\
path=steadyStatePath ,\
xguards=False ,\
yguards=False ,\
yind = [yInd, yInd] ,\
zind = [0 , 0] ,\
tind = [tLast, tLast],\
info=False)
n = calcN(lnN, normalized=True)
return n
def addnoise(path=".", var=None, scale=1e-5):
"""
Add random noise to restart files
Inputs
------
path Path to the restart files
var The variable to modify. By default all 3D variables are modified
scale Amplitude of the noise. Gaussian noise is used, with zero mean
and this parameter as the standard deviation
Returns
-------
None
"""
file_list = glob.glob(os.path.join(path, "BOUT.restart.*"))
nfiles = len(file_list)
print("Number of restart files: %d" % (nfiles, ))
for file in file_list:
print(file)
with DataFile(file, write=True) as d:
if var is None:
for v in d.list():
if d.ndims(v) == 3:
print(" -> " + v)
data = d.read(v)
data += normal(scale=scale, size=data.shape)
d.write(v, data)
else:
# Modify a single variable
print(" -> " + var)
data = d.read(var)
data += normal(scale=scale, size=data.shape)
d.write(var, data)
# Only one region
yup_xsplit = [nx]
ydown_xsplit = [nx]
yup_xin = [0]
yup_xout = [-1]
ydown_xin = [0]
ydown_xout = [-1]
nrad = [nx]
npol = [ny]
######################################################
print("Writing grid to file "+output)
of = DataFile()
of.open(output, create=True)
of.write("nx", nx)
of.write("ny", ny)
# Topology for original scheme
of.write("ixseps1", ixseps1)
of.write("ixseps2", ixseps2)
of.write("jyseps1_1", jyseps1_1)
of.write("jyseps1_2", jyseps1_2)
of.write("jyseps2_1", jyseps2_1)
of.write("jyseps2_2", jyseps2_2)
of.write("ny_inner", ny_inner)
# Grid spacing
def generate(
nx,
ny,
R=2.0,
r=0.2, # Major & minor radius
dr=0.05, # Radial width of domain
Bt=1.0, # Toroidal magnetic field
q=5.0, # Safety factor
mxg=2,
file="circle.nc",
):
# q = rBt / RBp
Bp = r * Bt / (R * q)
# Minor radius as function of x. Choose so boundary
# is half-way between grid points
h = dr / (nx - 2.0 * mxg) # Grid spacing in r
rminor = linspace(r - 0.5 * dr - (mxg - 0.5) * h, r + 0.5 * dr + (mxg - 0.5) * h, nx)
# mesh spacing in x and y
dx = ndarray([nx, ny])
dx[:, :] = r * Bt * h # NOTE: dx is toroidal flux
dy = ndarray([nx, ny])
dy[:, :] = 2.0 * pi / ny
# LogB = log(1/(1+r/R cos(theta))) =(approx) -(r/R)*cos(theta)
logB = zeros([nx, ny, 3]) # (constant, n=1 real, n=1 imag)
# At y = 0, Rmaj = R + r*cos(theta)
logB[:, 0, 1] = -(rminor / R)
# Moving in y, phase shift by (toroidal angle) / q
for y in range(1, ny):
dtheta = y * 2.0 * pi / ny / q # Change in poloidal angle
logB[:, y, 1] = -(rminor / R) * cos(dtheta)
logB[:, y, 2] = -(rminor / R) * sin(dtheta)
# Shift angle from one end of y to the other
ShiftAngle = ndarray([nx])
ShiftAngle[:] = 2.0 * pi / q
Rxy = ndarray([nx, ny])
Rxy[:, :] = r # NOTE : opposite to standard BOUT convention
Btxy = ndarray([nx, ny])
Btxy[:, :] = Bp
Bpxy = ndarray([nx, ny])
Bpxy[:, :] = Bt
Bxy = ndarray([nx, ny])
Bxy[:, :] = sqrt(Bt ** 2 + Bp ** 2)
hthe = ndarray([nx, ny])
hthe[:, :] = R
print("Writing to file '" + file + "'")
f = DataFile()
f.open(file, create=True)
# Mesh size
f.write("nx", nx)
f.write("ny", ny)
# Mesh spacing
f.write("dx", dx)
f.write("dy", dy)
# Metric components
f.write("Rxy", Rxy)
f.write("Btxy", Btxy)
f.write("Bpxy", Bpxy)
f.write("Bxy", Bxy)
f.write("hthe", hthe)
# Shift
f.write("ShiftAngle", ShiftAngle)
# Curvature
f.write("logB", logB)
# Input parameters
f.write("R", R)
f.write("r", r)
f.write("dr", dr)
f.write("Bt", Bt)
f.write("q", q)
f.write("mxg", mxg)
f.close()
def slice(infile, outfile, region=None, xind=None, yind=None):
"""Copy an X-Y slice from one DataFile to another
Parameters
----------
infile : str
Name of DataFile to read slice from
outfile : str
Name of DataFile to write slice to. File will be created, and
will be overwritten if it already exists
region : {0, 1, 2, 3, 4, 5, None}, optional
Copy a whole region. The available regions are:
- 0: Lower inner leg
- 1: Inner core
- 2: Upper inner leg
- 3: Upper outer leg
- 4: Outer core
- 5: Lower outer leg
xind, yind : (int, int), optional
Index ranges for x and y. Range includes first point, but not
last point
TODO
----
- Rename to not clobber builtin `slice`
- Better regions?
"""
# Open input and output files
indf = DataFile(infile)
outdf = DataFile(outfile, create=True)
nx = indf["nx"][0]
ny = indf["ny"][0]
if region:
# Select a region of the mesh
xind = [0, nx]
if region == 0:
# Lower inner leg
yind = [0, indf["jyseps1_1"][0]+1]
elif region == 1:
# Inner core
yind = [indf["jyseps1_1"][0]+1, indf["jyseps2_1"][0]+1]
elif region == 2:
# Upper inner leg
yind = [indf["jyseps2_1"][0]+1, indf["ny_inner"][0]]
elif region == 3:
# Upper outer leg
yind = [indf["ny_inner"][0], indf["jyseps1_2"][0]+1]
elif region == 4:
# Outer core
yind = [indf["jyseps1_2"][0]+1, indf["jyseps2_2"][0]+1]
else:
# Lower outer leg
yind = [indf["jyseps2_2"][0]+1, ny]
else:
# Use indices
if not xind:
xind = [0, nx]
if not yind:
yind = [0, ny]
print("Indices: [%d:%d, %d:%d]" % (xind[0], xind[1], yind[0], yind[1]))
# List of variables requiring special handling
special = ["nx", "ny", "ny_inner",
"ixseps1", "ixseps2",
"jyseps1_1", "jyseps1_2", "jyseps2_1", "jyseps2_2",
"ShiftAngle"]
outdf["nx"] = xind[1] - xind[0]
outdf["ny"] = yind[1] - yind[0]
outdf["ny_inner"] = indf["ny_inner"][0] - yind[0]
outdf["ixseps1"] = indf["ixseps1"][0]
outdf["ixseps2"] = indf["ixseps2"][0]
outdf["jyseps1_1"] = indf["jyseps1_1"][0] - yind[0]
outdf["jyseps2_1"] = indf["jyseps2_1"][0] - yind[0]
outdf["jyseps1_2"] = indf["jyseps1_2"][0] - yind[0]
outdf["jyseps2_2"] = indf["jyseps2_2"][0] - yind[0]
outdf["ShiftAngle"] = indf["ShiftAngle"][xind[0]:xind[1]]
# Loop over all variables
for v in list(indf.keys()):
if v in special:
continue # Skip these variables
ndims = indf.ndims(v)
if ndims == 0:
# Copy scalars
print("Copying variable: " + v)
outdf[v] = indf[v][0]
elif ndims == 2:
# Assume [x,y]
print("Slicing variable: " + v);
outdf[v] = indf[v][xind[0]:xind[1], yind[0]:yind[1]]
else:
# Skip
print("Skipping variable: " + v)
indf.close()
outdf.close()
#!/usr/bin/env python
#
# Generate an input mesh
#
from boututils.datafile import DataFile # Wrapper around NetCDF4 libraries
nx = 5 # Minimum is 5: 2 boundary, one evolved
ny = 64 # Minimum 5. Should be divisible by number of processors (so powers of 2 nice)
f = DataFile()
f.open("conduct_grid.nc", create=True)
f.write("nx", nx)
f.write("ny", ny)
f.close()
def collect(varname, xind=None, yind=None, zind=None, tind=None, path=".",yguards=False, xguards=True, info=True,prefix="BOUT.dmp",strict=False):
"""Collect a variable from a set of BOUT++ outputs.
data = collect(name)
name Name of the variable (string)
Optional arguments:
xind = [min,max] Range of X indices to collect
yind = [min,max] Range of Y indices to collect
zind = [min,max] Range of Z indices to collect
tind = [min,max] Range of T indices to collect
path = "." Path to data files
prefix = "BOUT.dmp" File prefix
yguards = False Collect Y boundary guard cells?
xguards = True Collect X boundary guard cells?
(Set to True to be consistent with the
definition of nx)
info = True Print information about collect?
strict = False Fail if the exact variable name is not found?
"""
# Search for BOUT++ dump files in NetCDF format
file_list_nc = glob.glob(os.path.join(path, prefix+".nc"))
file_list_h5 = glob.glob(os.path.join(path, prefix+".hdf5"))
if file_list_nc != [] and file_list_h5 != []:
raise IOError("Error: Both NetCDF and HDF5 files are present: do not know which to read.")
elif file_list_h5 != []:
suffix = ".hdf5"
file_list = file_list_h5
else:
suffix = ".nc"
file_list = file_list_nc
if file_list != []:
print("Single (parallel) data file")
f = DataFile(file_list[0]) # Open the file
data = f.read(varname)
return data
file_list_nc = glob.glob(os.path.join(path, prefix+".*nc"))
file_list_h5 = glob.glob(os.path.join(path, prefix+".*hdf5"))
if file_list_nc != [] and file_list_h5 != []:
raise IOError("Error: Both NetCDF and HDF5 files are present: do not know which to read.")
elif file_list_h5 != []:
suffix = ".hdf5"
file_list = file_list_h5
else:
suffix = ".nc"
file_list = file_list_nc
file_list.sort()
if file_list == []:
raise IOError("ERROR: No data files found")
nfiles = len(file_list)
# Read data from the first file
f = DataFile(file_list[0])
try:
dimens = f.dimensions(varname)
#ndims = len(dimens)
ndims = f.ndims(varname)
except:
if strict:
raise
else:
# Find the variable
varname = findVar(varname, f.list())
dimens = f.dimensions(varname)
#ndims = len(dimens)
ndims = f.ndims(varname)
if ndims < 2:
# Just read from file
data = f.read(varname)
f.close()
return data
if ndims > 4:
raise ValueError("ERROR: Too many dimensions")
mxsub = f.read("MXSUB")
if mxsub is None:
raise ValueError("Missing MXSUB variable")
mysub = f.read("MYSUB")
mz = f.read("MZ")
myg = f.read("MYG")
t_array = f.read("t_array")
if t_array is None:
nt = 1
t_array = np.zeros(1)
else:
nt = len(t_array)
if info:
print("mxsub = %d mysub = %d mz = %d\n" % (mxsub, mysub, mz))
# Get the version of BOUT++ (should be > 0.6 for NetCDF anyway)
try:
v = f.read("BOUT_VERSION")
# 2D decomposition
nxpe = f.read("NXPE")
mxg = f.read("MXG")
nype = f.read("NYPE")
npe = nxpe * nype
if info:
print("nxpe = %d, nype = %d, npe = %d\n" % (nxpe, nype, npe))
if npe < nfiles:
print("WARNING: More files than expected (" + str(npe) + ")")
elif npe > nfiles:
print("WARNING: Some files missing. Expected " + str(npe))
if xguards:
nx = nxpe * mxsub + 2*mxg
else:
nx = nxpe * mxsub
except KeyError:
print("BOUT++ version : Pre-0.2")
# Assume number of files is correct
# No decomposition in X
nx = mxsub
mxg = 0
nxpe = 1
nype = nfiles
if yguards:
ny = mysub * nype + 2*myg
else:
ny = mysub * nype
f.close();
# Check ranges
def check_range(r, low, up, name="range"):
r2 = r
if r != None:
try:
n = len(r2)
except:
# No len attribute, so probably a single number
r2 = [r2,r2]
if (len(r2) < 1) or (len(r2) > 2):
print("WARNING: "+name+" must be [min, max]")
r2 = None
else:
if len(r2) == 1:
r2 = [r2,r2]
if r2[0] < low:
r2[0] = low
if r2[0] > up:
r2[0] = up
if r2[1] < 0:
r2[1] = 0
if r2[1] > up:
r2[1] = up
if r2[0] > r2[1]:
tmp = r2[0]
r2[0] = r2[1]
r2[1] = tmp
else:
r2 = [low, up]
return r2
xind = check_range(xind, 0, nx-1, "xind")
yind = check_range(yind, 0, ny-1, "yind")
zind = check_range(zind, 0, mz-2, "zind")
tind = check_range(tind, 0, nt-1, "tind")
xsize = xind[1] - xind[0] + 1
ysize = yind[1] - yind[0] + 1
zsize = zind[1] - zind[0] + 1
tsize = tind[1] - tind[0] + 1
# Map between dimension names and output size
sizes = {'x':xsize, 'y':ysize, 'z':zsize, 't':tsize}
# Create a list with size of each dimension
ddims = [sizes[d] for d in dimens]
# Create the data array
data = np.zeros(ddims)
for i in range(npe):
# Get X and Y processor indices
pe_yind = int(i/nxpe)
pe_xind = i % nxpe
inrange = True
if yguards:
# Get local ranges
ymin = yind[0] - pe_yind*mysub
ymax = yind[1] - pe_yind*mysub
# Check lower y boundary
if pe_yind == 0:
# Keeping inner boundary
if ymax < 0: inrange = False
if ymin < 0: ymin = 0
else:
if ymax < myg: inrange = False
if ymin < myg: ymin = myg
# Upper y boundary
if pe_yind == (nype - 1):
# Keeping outer boundary
if ymin >= (mysub + 2*myg): inrange = False
if ymax > (mysub + 2*myg - 1): ymax = (mysub + 2*myg - 1)
else:
if ymin >= (mysub + myg): inrange = False
if ymax >= (mysub + myg): ymax = (mysub+myg-1)
# Calculate global indices
ygmin = ymin + pe_yind * mysub
ygmax = ymax + pe_yind * mysub
else:
# Get local ranges
ymin = yind[0] - pe_yind*mysub + myg
ymax = yind[1] - pe_yind*mysub + myg
if (ymin >= (mysub + myg)) or (ymax < myg):
inrange = False # Y out of range
if ymin < myg:
ymin = myg
if ymax >= mysub+myg:
ymax = myg + mysub - 1
# Calculate global indices
ygmin = ymin + pe_yind * mysub - myg
ygmax = ymax + pe_yind * mysub - myg
if xguards:
# Get local ranges
xmin = xind[0] - pe_xind*mxsub
xmax = xind[1] - pe_xind*mxsub
# Check lower x boundary
if pe_xind == 0:
# Keeping inner boundary
if xmax < 0: inrange = False
if xmin < 0: xmin = 0
else:
if xmax < mxg: inrange = False
if xmin < mxg: xmin = mxg
# Upper x boundary
if pe_xind == (nxpe - 1):
# Keeping outer boundary
if xmin >= (mxsub + 2*mxg): inrange = False
if xmax > (mxsub + 2*mxg - 1): xmax = (mxsub + 2*mxg - 1)
else:
if xmin >= (mxsub + mxg): inrange = False
if xmax >= (mxsub + mxg): xmax = (mxsub+mxg-1)
# Calculate global indices
xgmin = xmin + pe_xind * mxsub
xgmax = xmax + pe_xind * mxsub
else:
# Get local ranges
xmin = xind[0] - pe_xind*mxsub + mxg
xmax = xind[1] - pe_xind*mxsub + mxg
if (xmin >= (mxsub + mxg)) or (xmax < mxg):
inrange = False # X out of range
if xmin < mxg:
xmin = mxg
if xmax >= mxsub+mxg:
xmax = mxg + mxsub - 1
# Calculate global indices
xgmin = xmin + pe_xind * mxsub - mxg
xgmax = xmax + pe_xind * mxsub - mxg
# Number of local values
nx_loc = xmax - xmin + 1
ny_loc = ymax - ymin + 1
if not inrange:
continue # Don't need this file
filename = os.path.join(path, prefix+"." + str(i) + suffix)
if info:
sys.stdout.write("\rReading from " + filename + ": [" + \
str(xmin) + "-" + str(xmax) + "][" + \
str(ymin) + "-" + str(ymax) + "] -> [" + \
str(xgmin) + "-" + str(xgmax) + "][" + \
str(ygmin) + "-" + str(ygmax) + "]")
f = DataFile(filename)
if ndims == 4:
d = f.read(varname, ranges=[tind[0],tind[1]+1,
xmin, xmax+1,
ymin, ymax+1,
zind[0],zind[1]+1])
data[:, (xgmin-xind[0]):(xgmin-xind[0]+nx_loc), (ygmin-yind[0]):(ygmin-yind[0]+ny_loc), :] = d
elif ndims == 3:
# Could be xyz or txy
if dimens[2] == 'z': # xyz
d = f.read(varname, ranges=[xmin, xmax+1,
ymin, ymax+1,
zind[0],zind[1]+1])
data[(xgmin-xind[0]):(xgmin-xind[0]+nx_loc), (ygmin-yind[0]):(ygmin-yind[0]+ny_loc), :] = d
else: # txy
d = f.read(varname, ranges=[tind[0],tind[1]+1,
xmin, xmax+1,
ymin, ymax+1])
data[:, (xgmin-xind[0]):(xgmin-xind[0]+nx_loc), (ygmin-yind[0]):(ygmin-yind[0]+ny_loc)] = d
elif ndims == 2:
# xy
d = f.read(varname, ranges=[xmin, xmax+1,
ymin, ymax+1])
data[(xgmin-xind[0]):(xgmin-xind[0]+nx_loc), (ygmin-yind[0]):(ygmin-yind[0]+ny_loc)] = d
elif ndims == 1:
if dimens[0] == 't':
# t
d = f.read(varname, ranges=[tind[0],tind[1]+1])
data[:] = d
f.close()
# Force the precision of arrays of dimension>1
if ndims>1:
try:
data = data.astype(t_array.dtype, copy=False)
except TypeError:
data = data.astype(t_array.dtype)
# Finished looping over all files
if info:
sys.stdout.write("\n")
return data
def split(nxpe, nype, path="data", output="./", informat="nc", outformat=None, mxg=2, myg=2):
"""Split restart files across NXPE x NYPE processors.
Returns True on success
Parameters
----------
nxpe, nype : int
The number of processors in x and y
path : str, optional
Path to original restart files (default: "data")
output : str, optional
Path to write new restart files (default: current directory)
informat : str, optional
File extension of original files (default: "nc")
outformat : str, optional
File extension of new files (default: use the same as `informat`)
mxg, myg : int, optional
The number of guard cells in x and y
TODO
----
- Replace printing errors with raising `ValueError`
- Fix undefined variables!
- Make informat work like `redistribute`
"""
if outformat is None:
outformat = informat
npes = nxpe * nype
if npes <= 0:
print("ERROR: Negative or zero number of processors")
return False
if path == output:
print("ERROR: Can't overwrite restart files")
return False
file_list = glob.glob(os.path.join(path, "BOUT.restart.*."+informat))
nfiles = len(file_list)
if nfiles == 0:
print("ERROR: No restart files found")
return False
# Read old processor layout
f = DataFile(os.path.join(path, file_list[0]))
old_layout = get_processor_layout(f, False)
f.close()
if nfiles != old_layout.npes:
print("WARNING: Number of restart files inconsistent with NPES")
print("Setting nfiles = " + str(old_npes))
nfiles = old_layout.npes
if old_layout.npes % old_layout.nxpe != 0:
print("ERROR: Old NPES is not a multiple of old NXPE")
return False
if nype % old_layout.nype != 0:
print("SORRY: New nype must be a multiple of old nype")
return False
if nxpe % old_layout.nxpe != 0:
print("SORRY: New nxpe must be a multiple of old nxpe")
return False
# Calculate total size of the grid
nx = old_layout.mxsub * old_layout.nxpe
ny = old_layout.mysub * old_layout.nype
print(("Grid sizes: ", nx, ny, mz))
# Create the new restart files
for mype in range(npes):
# Calculate X and Y processor numbers
pex = mype % nxpe
pey = int(mype / nxpe)
old_pex = int(pex / xs)
old_pey = int(pey / ys)
old_x = pex % xs
old_y = pey % ys
# Old restart file number
old_mype = old_layout.nxpe * old_pey + old_pex
# Calculate indices in old restart file
xmin = old_x*mxsub
xmax = xmin + mxsub - 1 + 2*mxg
ymin = old_y*mysub
ymax = ymin + mysub - 1 + 2*myg
print("New: "+str(mype)+" ("+str(pex)+", "+str(pey)+")")
print(" => "+str(old_layout.mype)+" ("+str(old_pex)+", " +
str(old_pey)+") : ("+str(old_x)+", "+str(old_y)+")")
def pol_slice(var3d, gridfile, n=1, zangle=0.0):
""" data2d = pol_slice(data3d, 'gridfile', n=1, zangle=0.0) """
n = int(n)
zangle = float(zangle)
s = np.shape(var3d)
if len(s) != 3:
print("ERROR: pol_slice expects a 3D variable")
return None
nx, ny, nz = s
dz = 2.*np.pi / float(n * (nz-1))
try:
# Open the grid file
gf = DataFile(gridfile)
# Check the grid size is correct
if gf.read("nx") != nx:
print("ERROR: Grid X size is different to the variable")
return None
if gf.read("ny") != ny:
print("ERROR: Grid Y size is different to the variable")
return None
# Get the toroidal shift
zShift = gf.read("qinty")
if zShift != None:
print("Using qinty as toroidal shift angle")
else:
zShift = gf.read("zShift")
if zShift != None:
print("Using zShift as toroidal shift angle")
else:
print("ERROR: Neither qinty nor zShift found")
return None
gf.close()
except:
print("ERROR: pol_slice couldn't read grid file")
return None
var2d = np.zeros([nx, ny])
######################################
# Perform 2D slice
zind = (zangle - zShift) / dz
z0f = np.floor(zind)
z0 = z0f.astype(int)
p = zind - z0f
# Make z0 between 0 and (nz-2)
z0 = ((z0 % (nz-1)) + (nz-1)) % (nz-1)
# Get z+ and z-
zp = (z0 + 1) % (nz-1)
zm = (z0 - 1 + (nz-1)) % (nz-1)
# There may be some more cunning way to do this indexing
for x in np.arange(nx):
for y in np.arange(ny):
var2d[x,y] = 0.5*p[x,y]*(p[x,y]-1.0) * var3d[x,y,zm[x,y]] + \
(1.0 - p[x,y]*p[x,y]) * var3d[x,y,z0[x,y]] + \
0.5*p[x,y]*(p[x,y]+1.0) * var3d[x,y,zp[x,y]]
return var2d
def redistribute(npes, path="data", nxpe=None, output=".", informat=None, outformat=None, mxg=2, myg=2):
"""Resize restart files across NPES processors.
Does not check if new processor arrangement is compatible with the
branch cuts. In this respect :py:func:`restart.split` is
safer. However, BOUT++ checks the topology during initialisation
anyway so this is not too serious.
Parameters
----------
npes : int
Number of processors for the new restart files
path : str, optional
Path to original restart files (default: "data")
nxpe : int, optional
Number of processors to use in the x-direction (determines
split: npes = nxpe * nype). Default is None which uses the
same algorithm as BoutMesh (but without topology information)
to determine a suitable value for nxpe.
output : str, optional
Location to save new restart files (default: current directory)
informat : str, optional
Specify file format of old restart files (must be a suffix
understood by DataFile, e.g. 'nc'). Default uses the format of
the first 'BOUT.restart.*' file listed by glob.glob.
outformat : str, optional
Specify file format of new restart files (must be a suffix
understood by DataFile, e.g. 'nc'). Default is to use the same
as informat.
Returns
-------
True on success
TODO
----
- Replace printing errors with raising `ValueError`
"""
if npes <= 0:
print("ERROR: Negative or zero number of processors")
return False
if path == output:
print("ERROR: Can't overwrite restart files")
return False
if informat is None:
file_list = glob.glob(os.path.join(path, "BOUT.restart.*"))
else:
file_list = glob.glob(os.path.join(path, "BOUT.restart.*."+informat))
nfiles = len(file_list)
# Read old processor layout
f = DataFile(file_list[0])
# Get list of variables
var_list = f.list()
if len(var_list) == 0:
print("ERROR: No data found")
return False
old_processor_layout = get_processor_layout(f, has_t_dimension=False)
print("Grid sizes: ", old_processor_layout.nx,
old_processor_layout.ny, old_processor_layout.mz)
if nfiles != old_processor_layout.npes:
print("WARNING: Number of restart files inconsistent with NPES")
print("Setting nfiles = " + str(old_processor_layout.npes))
nfiles = old_processor_layout.npes
if nfiles == 0:
print("ERROR: No restart files found")
return False
informat = file_list[0].split(".")[-1]
if outformat is None:
outformat = informat
try:
new_processor_layout = create_processor_layout(
old_processor_layout, npes, nxpe=nxpe)
except ValueError as e:
print("Could not find valid processor split. " + e.what())
nx = old_processor_layout.nx
ny = old_processor_layout.ny
mz = old_processor_layout.mz
mxg = old_processor_layout.mxg
myg = old_processor_layout.myg
old_npes = old_processor_layout.npes
old_nxpe = old_processor_layout.nxpe
old_nype = old_processor_layout.nype
old_mxsub = old_processor_layout.mxsub
old_mysub = old_processor_layout.mysub
nxpe = new_processor_layout.nxpe
nype = new_processor_layout.nype
mxsub = new_processor_layout.mxsub
mysub = new_processor_layout.mysub
outfile_list = []
for i in range(npes):
outpath = os.path.join(output, "BOUT.restart."+str(i)+"."+outformat)
outfile_list.append(DataFile(outpath, write=True, create=True))
infile_list = []
for i in range(old_npes):
inpath = os.path.join(path, "BOUT.restart."+str(i)+"."+outformat)
infile_list.append(DataFile(inpath))
for v in var_list:
ndims = f.ndims(v)
# collect data
if ndims == 0:
# scalar
data = f.read(v)
elif ndims == 2:
data = np.zeros((nx+2*mxg, ny+2*myg))
for i in range(old_npes):
ix = i % old_nxpe
iy = int(i/old_nxpe)
ixstart = mxg
if ix == 0:
ixstart = 0
ixend = -mxg
if ix == old_nxpe-1:
ixend = 0
iystart = myg
if iy == 0:
iystart = 0
iyend = -myg
if iy == old_nype-1:
iyend = 0
data[ix*old_mxsub+ixstart:(ix+1)*old_mxsub+2*mxg+ixend,
iy*old_mysub+iystart:(iy+1)*old_mysub+2*myg+iyend] = infile_list[i].read(v)[ixstart:old_mxsub+2*mxg+ixend, iystart:old_mysub+2*myg+iyend]
data = BoutArray(data, attributes=infile_list[0].attributes(v))
elif ndims == 3:
data = np.zeros((nx+2*mxg, ny+2*myg, mz))
for i in range(old_npes):
ix = i % old_nxpe
iy = int(i/old_nxpe)
ixstart = mxg
if ix == 0:
ixstart = 0
ixend = -mxg
if ix == old_nxpe-1:
ixend = 0
iystart = myg
if iy == 0:
iystart = 0
iyend = -myg
if iy == old_nype-1:
iyend = 0
data[ix*old_mxsub+ixstart:(ix+1)*old_mxsub+2*mxg+ixend, iy*old_mysub+iystart:(iy+1)*old_mysub+2*myg+iyend,
:] = infile_list[i].read(v)[ixstart:old_mxsub+2*mxg+ixend, iystart:old_mysub+2*myg+iyend, :]
data = BoutArray(data, attributes=infile_list[0].attributes(v))
else:
print("ERROR: variable found with unexpected number of dimensions,", ndims, v)
return False
# write data
for i in range(npes):
ix = i % nxpe
iy = int(i/nxpe)
outfile = outfile_list[i]
if v == "NPES":
outfile.write(v, npes)
elif v == "NXPE":
outfile.write(v, nxpe)
elif v == "NYPE":
outfile.write(v, nype)
elif ndims == 0:
# scalar
outfile.write(v, data)
elif ndims == 2:
# Field2D
outfile.write(
v, data[ix*mxsub:(ix+1)*mxsub+2*mxg, iy*mysub:(iy+1)*mysub+2*myg])
elif ndims == 3:
# Field3D
outfile.write(
v, data[ix*mxsub:(ix+1)*mxsub+2*mxg, iy*mysub:(iy+1)*mysub+2*myg, :])
else:
print(
"ERROR: variable found with unexpected number of dimensions,", f.ndims(v))
f.close()
for infile in infile_list:
infile.close()
for outfile in outfile_list:
outfile.close()
return True
def create(averagelast=1, final=-1, path="data", output="./", informat="nc", outformat=None):
"""Create restart files from data (dmp) files.
Parameters
----------
averagelast : int, optional
Number of time points (counting from `final`, inclusive) to
average over (default is 1 i.e. just take last time-point)
final : int, optional
The last time point to use (default is last, -1)
path : str, optional
Path to original restart files (default: "data")
output : str, optional
Path to write new restart files (default: current directory)
informat : str, optional
File extension of original files (default: "nc")
outformat : str, optional
File extension of new files (default: use the same as `informat`)
"""
if outformat is None:
outformat = informat
file_list = glob.glob(os.path.join(path, "BOUT.dmp.*."+informat))
nfiles = len(file_list)
print(("Number of data files: ", nfiles))
for i in range(nfiles):
# Open each data file
infname = os.path.join(path, "BOUT.dmp."+str(i)+"."+informat)
outfname = os.path.join(output, "BOUT.restart."+str(i)+"."+outformat)
print((infname, " -> ", outfname))
infile = DataFile(infname)
outfile = DataFile(outfname, create=True)
# Get the data always needed in restart files
hist_hi = infile.read("iteration")
print(("hist_hi = ", hist_hi))
outfile.write("hist_hi", hist_hi)
t_array = infile.read("t_array")
tt = t_array[final]
print(("tt = ", tt))
outfile.write("tt", tt)
tind = final
if tind < 0.0:
tind = len(t_array) + final
NXPE = infile.read("NXPE")
NYPE = infile.read("NYPE")
print(("NXPE = ", NXPE, " NYPE = ", NYPE))
outfile.write("NXPE", NXPE)
outfile.write("NYPE", NYPE)
# Get a list of variables
varnames = infile.list()
for var in varnames:
if infile.ndims(var) == 4:
# Could be an evolving variable
print((" -> ", var))
data = infile.read(var)
if averagelast == 1:
slice = data[final, :, :, :]
else:
slice = mean(data[(final - averagelast)
:final, :, :, :], axis=0)
print(slice.shape)
outfile.write(var, slice)
infile.close()
outfile.close()
# Normalisation
Lbar = 1.
Bbar = 1.
J0 = - J0 * shape.Bxy / (MU0 * Lbar) # Turn into A/m^2
P0 = P0 * Bbar**2 / (2.0*MU0) # Pascals
shape.add(P0, "pressure")
shape.add(J0, "Jpar0")
shape.add(bxcvz, "bxcvz")
for nx in nxlist:
# Generate a new mesh file
filename = "grid%d.nc" % nx
if isfile(filename):
print("Grid file '%s' already exists" % filename)
else:
print("Creating grid file '%s'" % filename)
f = DataFile(filename, create=True)
shape.write(nx,nx, f)
f.close()
# Generate BOUT.inp file
directory = "grid%d" % nx
shell("mkdir " + directory)
shell("cp data/BOUT.inp "+directory)
shell("sed -i 's/MZ = 17/MZ = %d/g' %s/BOUT.inp" % (nx, directory))
shell("sed -i 's/grid = \"grid16.nc\"/grid = \"%s\"/g' %s/BOUT.inp" % (filename, directory))
def resizeY(newy, path="data", output=".", informat="nc", outformat=None, myg=2):
"""Increase the number of Y points in restart files
NOTE:
* Can't overwrite
Parameters
----------
newy : int
ny for the new file
path : str, optional
Path to original restart files (default: "data")
output : str, optional
Path to write new restart files (default: current directory)
informat : str, optional
File extension of original files (default: "nc")
outformat : str, optional
File extension of new files (default: use the same as `informat`)
myg : int, optional
Number of ghost points in y (default: 2)
Returns
-------
True on success, else False
TODO
----
- Replace printing errors with raising `ValueError`
- Make informat work like `redistribute`
"""
if outformat is None:
outformat = informat
file_list = glob.glob(os.path.join(path, "BOUT.restart.*."+informat))
nfiles = len(file_list)
if nfiles == 0:
print("ERROR: No restart files found")
return False
for i in range(nfiles):
# Open each data file
infname = os.path.join(path, "BOUT.restart."+str(i)+"."+informat)
outfname = os.path.join(output, "BOUT.restart."+str(i)+"."+outformat)
print("Processing %s -> %s" % (infname, outfname))
infile = DataFile(infname)
outfile = DataFile(outfname, create=True)
# Copy basic information
for var in ["hist_hi", "NXPE", "NYPE", "tt"]:
data = infile.read(var)
try:
# Convert to scalar if necessary
data = data[0]
except:
pass
outfile.write(var, data)
# Get a list of variables
varnames = infile.list()
for var in varnames:
if infile.ndims(var) == 3:
# Could be an evolving variable [x,y,z]
print(" -> Resizing " + var)
# Read variable from input
indata = infile.read(var)
nx, ny, nz = indata.shape
# y coordinate in input and output data
iny = (arange(ny) - myg + 0.5) / (ny - 2*myg)
outy = (arange(newy) - myg + 0.5) / (newy - 2*myg)
outdata = zeros([nx, newy, nz])
for x in range(nx):
for z in range(nz):
f = interp1d(
iny, indata[x, :, z], bounds_error=False, fill_value=0.0)
outdata[x, :, z] = f(outy)
outfile.write(var, outdata)
elif infile.ndims(var) == 2:
# Assume evolving variable [x,y]
print(" -> Resizing " + var)
# Read variable from input
indata = infile.read(var)
nx, ny = indata.shape
# y coordinate in input and output data
iny = (arange(ny) - myg + 0.5) / (ny - 2*myg)
outy = (arange(newy) - myg + 0.5) / (newy - 2*myg)
outdata = zeros([nx, newy])
for x in range(nx):
f = interp1d(iny, indata[x, :],
bounds_error=False, fill_value=0.0)
outdata[x, :] = f(outy)
outfile.write(var, outdata)
else:
# Copy variable
print(" -> Copying " + var)
# Read variable from input
data = infile.read(var)
try:
# Convert to scalar if necessary
data = data[0]
except:
pass
outfile.write(var, data)
infile.close()
outfile.close()
from past.utils import old_div
from boututils.datafile import DataFile # Wrapper around NetCDF4 libraries
from math import pow
from sys import argv
length = 80. # Length of the domain in m
nx = 5 # Minimum is 5: 2 boundary, one evolved
if len(argv)>1:
ny = int(argv[1]) # Minimum 5. Should be divisible by number of processors (so powers of 2 nice)
else:
ny = 256 # Minimum 5. Should be divisible by number of processors (so powers of 2 nice)
#dy = [[1.]*ny]*nx # distance between points in y, in m/g22/lengthunit
g22 = [[pow(old_div(float(ny-1),length),2)]*ny]*nx
g_22 = [[pow(old_div(length,float(ny-1)),2)]*ny]*nx
ixseps1 = -1
ixseps2 = 0
f = DataFile()
f.open("conduct_grid.nc", create=True)
f.write("nx", nx)
f.write("ny", ny)
#f.write("dy", dy)
f.write("g22",g22)
f.write("g_22", g_22)
f.write("ixseps1", ixseps1)
f.write("ixseps2", ixseps2)
f.close()
def split(nxpe, nype, path="data", output="./", informat="nc", outformat=None):
"""Split restart files across NXPE x NYPE processors.
Returns True on success
"""
if outformat is None:
outformat = informat
mxg = 2
myg = 2
npes = nxpe * nype
if npes <= 0:
print("ERROR: Negative or zero number of processors")
return False
if path == output:
print("ERROR: Can't overwrite restart files")
return False
file_list = glob.glob(os.path.join(path, "BOUT.restart.*." + informat))
nfiles = len(file_list)
if nfiles == 0:
print("ERROR: No restart files found")
return False
# Read old processor layout
f = DataFile(os.path.join(path, file_list[0]))
# Get list of variables
var_list = f.list()
if len(var_list) == 0:
print("ERROR: No data found")
return False
old_npes = f.read("NPES")
old_nxpe = f.read("NXPE")
if nfiles != old_npes:
print("WARNING: Number of restart files inconsistent with NPES")
print("Setting nfiles = " + str(old_npes))
nfiles = old_npes
if old_npes % old_nxpe != 0:
print("ERROR: Old NPES is not a multiple of old NXPE")
return False
old_nype = int(old_npes / old_nxpe)
if nype % old_nype != 0:
print("SORRY: New nype must be a multiple of old nype")
return False
if nxpe % old_nxpe != 0:
print("SORRY: New nxpe must be a multiple of old nxpe")
return False
# Get dimension sizes
old_mxsub = 0
old_mysub = 0
mz = 0
for v in var_list:
if f.ndims(v) == 3:
s = f.size(v)
old_mxsub = s[0] - 2 * mxg
old_mysub = s[1] - 2 * myg
mz = s[2]
break
f.close()
# Calculate total size of the grid
nx = old_mxsub * old_nxpe
ny = old_mysub * old_nype
print(("Grid sizes: ", nx, ny, mz))
# Create the new restart files
for mype in range(npes):
# Calculate X and Y processor numbers
pex = mype % nxpe
pey = int(mype / nxpe)
old_pex = int(pex / xs)
old_pey = int(pey / ys)
old_x = pex % xs
old_y = pey % ys
# Old restart file number
old_mype = old_nxpe * old_pey + old_pex
# Calculate indices in old restart file
xmin = old_x * mxsub
xmax = xmin + mxsub - 1 + 2 * mxg
ymin = old_y * mysub
ymax = ymin + mysub - 1 + 2 * myg
print("New: " + str(mype) + " (" + str(pex) + ", " + str(pey) + ")")
print(
" => "
+ str(old_mype)
+ " ("
+ str(old_pex)
+ ", "
+ str(old_pey)
+ ") : ("
+ str(old_x)
+ ", "
+ str(old_y)
+ ")"
)
def create(averagelast=1, final=-1, path="data", output="./", informat="nc", outformat=None):
"""
Create restart files from data (dmp) files.
Inputs
======
averagelast Number of time points to average over.
Default is 1 i.e. just take last time-point
final The last time point to use. Default is last (-1)
path Path to the input data files
output Path where the output restart files should go
informat Format of the input data files
outformat Format of the output restart files
"""
if outformat is None:
outformat = informat
file_list = glob.glob(os.path.join(path, "BOUT.dmp.*." + informat))
nfiles = len(file_list)
print(("Number of data files: ", nfiles))
for i in range(nfiles):
# Open each data file
infname = os.path.join(path, "BOUT.dmp." + str(i) + "." + informat)
outfname = os.path.join(output, "BOUT.restart." + str(i) + "." + outformat)
print((infname, " -> ", outfname))
infile = DataFile(infname)
outfile = DataFile(outfname, create=True)
# Get the data always needed in restart files
hist_hi = infile.read("iteration")
print(("hist_hi = ", hist_hi))
outfile.write("hist_hi", hist_hi)
t_array = infile.read("t_array")
tt = t_array[final]
print(("tt = ", tt))
outfile.write("tt", tt)
tind = final
if tind < 0.0:
tind = len(t_array) + final
NXPE = infile.read("NXPE")
NYPE = infile.read("NYPE")
NPES = NXPE * NYPE
print(("NPES = ", NPES, " NXPE = ", NXPE))
outfile.write("NPES", NPES)
outfile.write("NXPE", NXPE)
# Get a list of variables
varnames = infile.list()
for var in varnames:
if infile.ndims(var) == 4:
# Could be an evolving variable
print((" -> ", var))
data = infile.read(var)
if averagelast == 1:
slice = data[final, :, :, :]
else:
slice = mean(data[(final - averagelast) : final, :, :, :], axis=0)
print(slice.shape)
outfile.write(var, slice)
infile.close()
outfile.close()
def redistribute(npes, path="data", nxpe=None, output=".", informat=None, outformat=None, mxg=2, myg=2):
"""Resize restart files across NPES processors.
Does not check if new processor arrangement is compatible with the branch cuts. In this respect restart.split is safer. However, BOUT++ checks the topology during initialisation anyway so this is not too serious.
Parameters
----------
npes : int
number of processors for the new restart files
path : string, optional
location of old restart files
nxpe : int, optional
number of processors to use in the x-direction (determines split: npes = nxpe * nype). Default is None which uses the same algorithm as BoutMesh (but without topology information) to determine a suitable value for nxpe.
output : string, optional
location to save new restart files
informat : string, optional
specify file format of old restart files (must be a suffix understood by DataFile, e.g. 'nc'). Default uses the format of the first 'BOUT.restart.*' file listed by glob.glob.
outformat : string, optional
specify file format of new restart files (must be a suffix understood by DataFile, e.g. 'nc'). Default is to use the same as informat.
Returns
-------
True on success
"""
if npes <= 0:
print("ERROR: Negative or zero number of processors")
return False
if path == output:
print("ERROR: Can't overwrite restart files")
return False
if informat is None:
file_list = glob.glob(os.path.join(path, "BOUT.restart.*"))
else:
file_list = glob.glob(os.path.join(path, "BOUT.restart.*." + informat))
nfiles = len(file_list)
# Read old processor layout
f = DataFile(file_list[0])
# Get list of variables
var_list = f.list()
if len(var_list) == 0:
print("ERROR: No data found")
return False
old_npes = f.read("NPES")
old_nxpe = f.read("NXPE")
old_nype = int(old_npes / old_nxpe)
if nfiles != old_npes:
print("WARNING: Number of restart files inconsistent with NPES")
print("Setting nfiles = " + str(old_npes))
nfiles = old_npes
if nfiles == 0:
print("ERROR: No restart files found")
return False
informat = file_list[0].split(".")[-1]
if outformat is None:
outformat = informat
old_mxsub = 0
old_mysub = 0
mz = 0
for v in var_list:
if f.ndims(v) == 3:
s = f.size(v)
old_mxsub = s[0] - 2 * mxg
if old_mxsub < 0:
if s[0] == 1:
old_mxsub = 1
mxg = 0
elif s[0] == 3:
old_mxsub = 1
mxg = 1
else:
print("Number of x points is wrong?")
return False
old_mysub = s[1] - 2 * myg
if old_mysub < 0:
if s[1] == 1:
old_mysub = 1
myg = 0
elif s[1] == 3:
old_mysub = 1
myg = 1
else:
print("Number of y points is wrong?")
return False
mz = s[2]
break
# Calculate total size of the grid
nx = old_mxsub * old_nxpe
ny = old_mysub * old_nype
print("Grid sizes: ", nx, ny, mz)
if nxpe is None: # Copy algorithm from BoutMesh for selecting nxpe
ideal = sqrt(float(nx) * float(npes) / float(ny)) # Results in square domain
for i in range(1, npes + 1):
if npes % i == 0 and nx % i == 0 and int(nx / i) >= mxg and ny % (npes / i) == 0:
# Found an acceptable value
# Warning: does not check branch cuts!
if nxpe is None or abs(ideal - i) < abs(ideal - nxpe):
nxpe = i # Keep value nearest to the ideal
if nxpe is None:
print("ERROR: could not find a valid value for nxpe")
return False
nype = int(npes / nxpe)
outfile_list = []
for i in range(npes):
outpath = os.path.join(output, "BOUT.restart." + str(i) + "." + outformat)
outfile_list.append(DataFile(outpath, write=True, create=True))
infile_list = []
for i in range(old_npes):
inpath = os.path.join(path, "BOUT.restart." + str(i) + "." + outformat)
infile_list.append(DataFile(inpath))
old_mxsub = int(nx / old_nxpe)
old_mysub = int(ny / old_nype)
mxsub = int(nx / nxpe)
mysub = int(ny / nype)
for v in var_list:
ndims = f.ndims(v)
# collect data
if ndims == 0:
# scalar
data = f.read(v)
elif ndims == 2:
data = np.zeros((nx + 2 * mxg, ny + 2 * myg))
for i in range(old_npes):
ix = i % old_nxpe
iy = int(i / old_nxpe)
ixstart = mxg
if ix == 0:
ixstart = 0
ixend = -mxg
if ix == old_nxpe - 1:
ixend = 0
iystart = myg
if iy == 0:
iystart = 0
iyend = -myg
if iy == old_nype - 1:
iyend = 0
data[
ix * old_mxsub + ixstart : (ix + 1) * old_mxsub + 2 * mxg + ixend,
iy * old_mysub + iystart : (iy + 1) * old_mysub + 2 * myg + iyend,
] = infile_list[i].read(v)[ixstart : old_mxsub + 2 * mxg + ixend, iystart : old_mysub + 2 * myg + iyend]
elif ndims == 3:
data = np.zeros((nx + 2 * mxg, ny + 2 * myg, mz))
for i in range(old_npes):
ix = i % old_nxpe
iy = int(i / old_nxpe)
ixstart = mxg
if ix == 0:
ixstart = 0
ixend = -mxg
if ix == old_nxpe - 1:
ixend = 0
iystart = myg
if iy == 0:
iystart = 0
iyend = -myg
if iy == old_nype - 1:
iyend = 0
data[
ix * old_mxsub + ixstart : (ix + 1) * old_mxsub + 2 * mxg + ixend,
iy * old_mysub + iystart : (iy + 1) * old_mysub + 2 * myg + iyend,
:,
] = infile_list[i].read(v)[
ixstart : old_mxsub + 2 * mxg + ixend, iystart : old_mysub + 2 * myg + iyend, :
]
else:
print("ERROR: variable found with unexpected number of dimensions,", ndims, v)
return False
# write data
for i in range(npes):
ix = i % nxpe
iy = int(i / nxpe)
outfile = outfile_list[i]
if v == "NPES":
outfile.write(v, npes)
elif v == "NXPE":
outfile.write(v, nxpe)
elif ndims == 0:
# scalar
outfile.write(v, data)
elif ndims == 2:
# Field2D
outfile.write(v, data[ix * mxsub : (ix + 1) * mxsub + 2 * mxg, iy * mysub : (iy + 1) * mysub + 2 * myg])
elif ndims == 3:
# Field3D
outfile.write(
v, data[ix * mxsub : (ix + 1) * mxsub + 2 * mxg, iy * mysub : (iy + 1) * mysub + 2 * myg, :]
)
else:
print("ERROR: variable found with unexpected number of dimensions,", f.ndims(v))
f.close()
for infile in infile_list:
infile.close()
for outfile in outfile_list:
outfile.close()
return True
def resizeY(newy, path="data", output=".", informat="nc", outformat=None, myg=2):
"""
Resize all the restart files in Y
"""
if outformat is None:
outformat = informat
file_list = glob.glob(os.path.join(path, "BOUT.restart.*." + informat))
nfiles = len(file_list)
if nfiles == 0:
print("ERROR: No restart files found")
return False
for i in range(nfiles):
# Open each data file
infname = os.path.join(path, "BOUT.restart." + str(i) + "." + informat)
outfname = os.path.join(output, "BOUT.restart." + str(i) + "." + outformat)
print("Processing %s -> %s", infname, outfname)
infile = DataFile(infname)
outfile = DataFile(outfname, create=True)
# Copy basic information
for var in ["hist_hi", "NPES", "NXPE", "tt"]:
data = infile.read(var)
try:
# Convert to scalar if necessary
data = data[0]
except:
pass
outfile.write(var, data)
# Get a list of variables
varnames = infile.list()
for var in varnames:
if infile.ndims(var) == 3:
# Could be an evolving variable [x,y,z]
print(" -> " + var)
# Read variable from input
indata = infile.read(var)
nx, ny, nz = indata.shape
# y coordinate in input and output data
iny = (arange(ny) - myg + 0.5) / (ny - 2 * myg)
outy = (arange(newy) - myg + 0.5) / (newy - 2 * myg)
outdata = zeros([nx, newy, nz])
for x in range(nx):
for z in range(nz):
f = interp1d(iny, indata[x, :, z], bounds_error=False, fill_value=0.0)
outdata[x, :, z] = f(outy)
outfile.write(var, outdata)
infile.close()
outfile.close()
def slice(infile, outfile, region = None, xind=None, yind=None):
"""
xind, yind - index ranges. Range includes first point, but not last point
"""
# Open input and output files
indf = DataFile(infile)
outdf = DataFile(outfile, create=True)
nx = indf["nx"][0]
ny = indf["ny"][0]
if region:
# Select a region of the mesh
xind = [0, nx]
if region == 0:
# Lower inner leg
yind = [0, indf["jyseps1_1"][0]+1]
elif region == 1:
# Inner core
yind = [indf["jyseps1_1"][0]+1, indf["jyseps2_1"][0]+1]
elif region == 2:
# Upper inner leg
yind = [indf["jyseps2_1"][0]+1, indf["ny_inner"][0]]
elif region == 3:
# Upper outer leg
yind = [indf["ny_inner"][0], indf["jyseps1_2"][0]+1]
elif region == 4:
# Outer core
yind = [indf["jyseps1_2"][0]+1, indf["jyseps2_2"][0]+1]
else:
# Lower outer leg
yind = [indf["jyseps2_2"][0]+1, ny]
else:
# Use indices
if not xind:
xind = [0, nx]
if not yind:
yind = [0, ny]
print("Indices: [%d:%d, %d:%d]" % (xind[0], xind[1], yind[0], yind[1]))
# List of variables requiring special handling
special = ["nx", "ny", "ny_inner",
"ixseps1", "ixseps2",
"jyseps1_1", "jyseps1_2", "jyseps2_1", "jyseps2_2",
"ShiftAngle"]
outdf["nx"] = xind[1] - xind[0]
outdf["ny"] = yind[1] - yind[0]
outdf["ny_inner"] = indf["ny_inner"][0] - yind[0]
outdf["ixseps1"] = indf["ixseps1"][0]
outdf["ixseps2"] = indf["ixseps2"][0]
outdf["jyseps1_1"] = indf["jyseps1_1"][0] - yind[0]
outdf["jyseps2_1"] = indf["jyseps2_1"][0] - yind[0]
outdf["jyseps1_2"] = indf["jyseps1_2"][0] - yind[0]
outdf["jyseps2_2"] = indf["jyseps2_2"][0] - yind[0]
outdf["ShiftAngle"] = indf["ShiftAngle"][xind[0]:xind[1]]
# Loop over all variables
for v in list(indf.keys()):
if v in special:
continue # Skip these variables
ndims = indf.ndims(v)
if ndims == 0:
# Copy scalars
print("Copying variable: " + v)
outdf[v] = indf[v][0]
elif ndims == 2:
# Assume [x,y]
print("Slicing variable: " + v);
outdf[v] = indf[v][xind[0]:xind[1], yind[0]:yind[1]]
else:
# Skip
print("Skipping variable: " + v)
indf.close()
outdf.close()
#!/usr/bin/env python
#
# Generate an input mesh
#
from boututils.datafile import DataFile # Wrapper around NetCDF4 libraries
nx = 5 # Minimum is 5: 2 boundary, one evolved
ny = 32 # Minimum 5. Should be divisible by number of processors (so powers of 2 nice)
dy = 1. # distance between points in y, in m/g22/lengthunit
ixseps1 = -1
ixseps2 = -1
f = DataFile()
f.open("test-staggered.nc", create=True)
f.write("nx", nx)
f.write("ny", ny)
f.write("dy", dy)
f.write("ixseps1", ixseps1)
f.write("ixseps2", ixseps2)
f.close()