def fci_to_vtk(infile, outfile, scale=5):
if not have_evtk:
return
with bdata.DataFile(infile, write=False, create=False) as f:
dx = f.read('dx')
dy = f.read('dy')
bx = f.read('bx')
by = np.ones(bx.shape)
bz = f.read('bz')
if bx is None:
xt_prime = f.read('forward_xt_prime')
zt_prime = f.read('forward_zt_prime')
array_indices = indices(xt_prime.shape)
bx = xt_prime - array_indices[0, ...]
by = by * dy
bz = zt_prime - array_indices[2, ...]
nx, ny, nz = bx.shape
dz = nx * dx / nz
x = np.linspace(0, nx * dx, nx)
y = np.linspace(0, ny * dy, ny, endpoint=False)
z = np.linspace(0, nz * dz, nz, endpoint=False)
gridToVTK(outfile,
x * scale,
y,
z * scale,
pointData={'B': (bx * scale, by, bz * scale)})
def slab(nx,
ny,
nz,
filename="fci.grid.nc",
Lx=0.1,
Ly=10.,
Lz=1.,
Bt=1.0,
Bp=0.1,
Bpprime=1.0):
"""
nx - Number of radial points
ny - Number of toroidal points (NOTE: Different to BOUT++ standard)
nz - Number of poloidal points
Lx - Radial domain size [m]
Ly - Toroidal domain size [m]
Lz - Poloidal domain size [m]
Bt - Toroidal magnetic field [T]
Bp - Poloidal magnetic field [T]
Bpprime - Gradient of Bp [T/m] Bp(x) = Bp + Bpprime * x
"""
MXG = 2
# Make sure input types are sane
nx = int(nx)
ny = int(ny)
nz = int(nz)
Lx = float(Lx)
Ly = float(Ly)
Lz = float(Lz)
delta_x = old_div(Lx, (nx - 2. * MXG))
delta_pol = old_div(Lz, (nz))
delta_tor = old_div(Ly, (ny))
# Coord arrays
x = Lx * (np.arange(nx) - MXG + 0.5) / (nx - 2. * MXG
) # 0 and 1 half-way between cells
y = np.linspace(0, Ly, ny)
z = np.linspace(0, Lz, nz, endpoint=False)
############################################################
# Effective major radius
R = old_div(Ly, (2. * pi))
# Set poloidal magnetic field
Bpx = Bp + (x - old_div(Lx, 2)) * Bpprime
Bpxy = np.transpose(np.resize(Bpx, (nz, ny, nx)), (2, 1, 0))
Bxy = np.sqrt(Bpxy**2 + Bt**2)[:, :, 0]
class Mappoint(object):
def __init__(self, xt, zt):
self.xt = xt
self.zt = zt
self.xt_prime = old_div(xt, delta_x) + MXG - 0.5
self.zt_prime = old_div(zt, delta_pol)
def unroll_map_coeff(map_list, coeff):
coeff_array = np.transpose(
np.resize(
np.array([getattr(f, coeff) for f in map_list]).reshape(
(nx, nz)), (ny, nx, nz)), (1, 0, 2))
return coeff_array
def b_field(vector, y):
x0 = old_div(Lx, 2.) # Centre of box, where bz = 0.
x, z = vector
bx = 0.
bz = Bp + (x - x0) * Bpprime
return [bx, bz]
def field_line_tracer(direction, map_list):
result = np.zeros((nx, nz, 2))
for i in np.arange(0, nx):
for k in np.arange(0, nz):
result[i, k, :] = odeint(b_field, [x[i], z[k]],
[0, delta_tor * direction])[1, :]
map_list.append(Mappoint(result[i, k, 0], result[i, k, 1]))
return result
forward_map = []
forward_coords = field_line_tracer(+1, forward_map)
backward_map = []
backward_coords = field_line_tracer(-1, backward_map)
X, Y = np.meshgrid(x, y, indexing='ij')
x0 = 0.5
g_22 = old_div(((Bp + (X - x0) * Lx * Bpprime)**2 + Bt**2), Bt**2)
with bdata.DataFile(filename, write=True, create=True) as f:
f.write('nx', nx)
f.write('ny', ny)
f.write('nz', nz)
f.write("dx", delta_x)
f.write("dy", delta_tor)
f.write("g_22", g_22)
f.write("Bxy", (Bxy))
xt_prime = unroll_map_coeff(forward_map, 'xt_prime')
f.write('forward_xt_prime', (xt_prime))
zt_prime = unroll_map_coeff(forward_map, 'zt_prime')
f.write('forward_zt_prime', (zt_prime))
xt_prime = unroll_map_coeff(backward_map, 'xt_prime')
f.write('backward_xt_prime', (xt_prime))
zt_prime = unroll_map_coeff(backward_map, 'zt_prime')
f.write('backward_zt_prime', (zt_prime))
map_list.append(Mappoint(result[i,k,0],result[i,k,1]))
return result
if __name__ == "__main__":
forward_map = []
forward_coords = field_line_tracer(+1, forward_map)
backward_map = []
backward_coords = field_line_tracer(-1, backward_map)
X,Y = np.meshgrid(x,y,indexing='ij')
x0 = 0.5
g_22 = np.sqrt(((Bp + (X-x0) * Lx * Bpprime)**2 + 1))
with bdata.DataFile('fci.grid.nc', write=True, create=True) as f:
f.write('nx', nx)
f.write('ny', ny)
f.write('nz', nz)
f.write("dx", delta_x)
f.write("dy", delta_tor)
f.write("g_22", g_22)
f.write("Bxy", transform3D(Bxy))
xt_prime = unroll_map_coeff(forward_map, 'xt_prime')
f.write('forward_xt_prime', transform3D(xt_prime))
zt_prime = unroll_map_coeff(forward_map, 'zt_prime')
f.write('forward_zt_prime', transform3D(zt_prime))
xt_prime = unroll_map_coeff(backward_map, 'xt_prime')
f.write('backward_xt_prime', transform3D(xt_prime))
import boututils.datafile as bdata
from boutdata.input import transform3D
# Parameters
nx = 10
ny = 20
nz = 8
shape = [nx, ny, nz]
xt_prime = zeros(shape)
zt_prime = zeros(shape)
for x in range(nx):
# No interpolation in x
xt_prime[x,:,:] = x
# Each y slice scans between neighbouring z points
for z in range(nz):
zt_prime[x,:,z] = z + concatenate([linspace(-1, 1, ny-1), [0]])
with bdata.DataFile('simple_test.nc', write=True, create=True) as f:
f.write('nx',nx)
f.write('ny',ny)
for direction_name in ['forward', 'backward']:
f.write(direction_name + '_xt_prime', transform3D(xt_prime))
f.write(direction_name + '_zt_prime', transform3D(zt_prime))
def write_maps(grid,
magnetic_field,
maps,
gridfile='fci.grid.nc',
new_names=False,
metric2d=True,
format="NETCDF3_64BIT",
quiet=False):
"""Write FCI maps to BOUT++ grid file
Parameters
----------
grid : :py:obj:`zoidberg.grid.Grid`
Grid generated by Zoidberg
magnetic_field : :py:obj:`zoidberg.field.MagneticField`
Zoidberg magnetic field object
maps : dict
Dictionary of FCI maps
gridfile : str, optional
Output filename
new_names : bool, optional
Write "g_yy" rather than "g_22"
metric2d : bool, optional
Output only 2D metrics
format : str, optional
Specifies file format to use, passed to boutdata.DataFile
quiet : bool, optional
Don't warn about 2D metrics
Returns
-------
Writes the following variables to the grid file
"""
nx, ny, nz = grid.shape
# Get metric tensor
metric = grid.metric()
# Check if the magnetic field is in cylindrical coordinates
# If so, we need to change the gyy and g_yy metrics
pol_grid, ypos = grid.getPoloidalGrid(0)
Rmaj = magnetic_field.Rfunc(pol_grid.R, pol_grid.Z, ypos)
if Rmaj is not None:
# In cylindrical coordinates
Rmaj = np.zeros(grid.shape)
for yindex in range(grid.numberOfPoloidalGrids()):
pol_grid, ypos = grid.getPoloidalGrid(yindex)
Rmaj[:, yindex, :] = magnetic_field.Rfunc(pol_grid.R, pol_grid.Z,
ypos)
metric["gyy"] = 1. / Rmaj**2
metric["g_yy"] = Rmaj**2
# Get magnetic field and pressure
Bmag = np.zeros(grid.shape)
pressure = np.zeros(grid.shape)
for yindex in range(grid.numberOfPoloidalGrids()):
pol_grid, ypos = grid.getPoloidalGrid(yindex)
Bmag[:, yindex, :] = magnetic_field.Bmag(pol_grid.R, pol_grid.Z, ypos)
pressure[:,
yindex, :] = magnetic_field.pressure(pol_grid.R, pol_grid.Z,
ypos)
metric["g_yy"][:, yindex, :] = (
metric["g_yy"][:, yindex, :] *
(Bmag[:, yindex, :] /
magnetic_field.Byfunc(pol_grid.R, pol_grid.Z, ypos))**2)
metric["gyy"][:, yindex, :] = (
metric["gyy"][:, yindex, :] *
(magnetic_field.Byfunc(pol_grid.R, pol_grid.Z, ypos) /
Bmag[:, yindex, :])**2)
# Get attributes from magnetic field (e.g. psi)
attributes = {}
for name in magnetic_field.attributes:
attribute = np.zeros(grid.shape)
for yindex in range(grid.numberOfPoloidalGrids()):
pol_grid, ypos = grid.getPoloidalGrid(yindex)
attribute[:,
yindex, :] = magnetic_field.attributes[name](pol_grid.R,
pol_grid.Z,
ypos)
attributes[name] = attribute
# Metric is now 3D
if metric2d:
# Remove the Z dimension from metric components
if not quiet:
print(
"WARNING: Outputting 2D metrics, discarding metric information."
)
for key in metric:
try:
metric[key] = metric[key][:, :, 0]
except TypeError:
pass
# Make dz a constant
metric["dz"] = metric["dz"][0, 0]
# Add Rxy, Bxy
metric["Rxy"] = maps["R"][:, :, 0]
metric["Bxy"] = Bmag[:, :, 0]
with bdata.DataFile(gridfile, write=True, create=True, format=format) as f:
ixseps = nx + 1
f.write('nx', nx)
f.write('ny', ny)
f.write('nz', nz)
f.write("dx", metric["dx"])
f.write("dy", metric["dy"])
f.write("dz", metric["dz"])
f.write("ixseps1", ixseps)
f.write("ixseps2", ixseps)
# Metric tensor
if new_names:
for key, val in metric.items():
f.write(key, val)
else:
# Translate between output variable names and metric names
# Map from new to old names. Anything not in this dict
# is output unchanged
name_changes = {
"g_yy": "g_22",
"gyy": "g22",
"gxx": "g11",
"gxz": "g13",
"gzz": "g33",
"g_xx": "g_11",
"g_xz": "g_13",
"g_zz": "g_33"
}
for key in metric:
name = key
if name in name_changes:
name = name_changes[name]
f.write(name, metric[key])
# Magnetic field
f.write("B", Bmag)
# Pressure
f.write("pressure", pressure)
# Attributes
for name in attributes:
f.write(name, attributes[name])
# Maps - write everything to file
for key in maps:
f.write(key, maps[key])
def write_maps(grid,
magnetic_field,
maps,
gridfile='fci.grid.nc',
legacy=False):
"""Write FCI maps to BOUT++ grid file
Inputs
------
grid - Grid generated by Zoidberg
magnetic_field - Zoidberg magnetic field object
maps - Dictionary of FCI maps
gridfile - Output filename
legacy - If true, write FCI maps using FFTs
"""
nx, ny, nz = (grid.nx, grid.ny, grid.nz)
xarray, yarray, zarray = (grid.xarray, grid.yarray, grid.zarray)
g_22 = np.zeros((nx, ny)) + 1. / grid.Rmaj**2
totalbx = np.zeros((nx, ny, nz))
totalbz = np.zeros((nx, ny, nz))
Bxy = np.zeros((nx, ny, nz))
for i in np.arange(0, nx):
for j in np.arange(0, ny):
for k in np.arange(0, nz):
Bxy[i, j, k] = np.sqrt((
magnetic_field.Bxfunc(xarray[i], zarray[k], yarray[j])**2 +
magnetic_field.Bzfunc(xarray[i], zarray[k], yarray[j])**2))
totalbx[i, j,
k] = magnetic_field.Bxfunc(xarray[i], zarray[k],
yarray[j])
totalbz[i, j,
k] = magnetic_field.Bzfunc(xarray[i], zarray[k],
yarray[j])
with bdata.DataFile(gridfile, write=True, create=True) as f:
ixseps = nx + 1
f.write('nx', grid.nx)
f.write('ny', grid.ny)
if not legacy:
# Legacy files don't need nz
f.write('nz', grid.nz)
f.write("dx", grid.delta_x)
f.write("dy", grid.delta_y)
f.write("ixseps1", ixseps)
f.write("ixseps2", ixseps)
f.write("g_22", g_22)
f.write("Bxy", Bxy[:, :, 0])
f.write("bx", totalbx)
f.write("bz", totalbz)
# Legacy grid files need to FFT 3D arrays
if legacy:
from boutdata.input import transform3D
f.write('forward_xt_prime', transform3D(maps['forward_xt_prime']))
f.write('forward_zt_prime', transform3D(maps['forward_zt_prime']))
f.write('backward_xt_prime',
transform3D(maps['backward_xt_prime']))
f.write('backward_zt_prime',
transform3D(maps['backward_zt_prime']))
else:
f.write('forward_xt_prime', maps['forward_xt_prime'])
f.write('forward_zt_prime', maps['forward_zt_prime'])
f.write('backward_xt_prime', maps['backward_xt_prime'])
f.write('backward_zt_prime', maps['backward_zt_prime'])
def write_maps(grid,
magnetic_field,
maps,
gridfile='fci.grid.nc',
new_names=False,
metric2d=True):
"""Write FCI maps to BOUT++ grid file
Inputs
------
grid - Grid generated by Zoidberg
magnetic_field - Zoidberg magnetic field object
maps - Dictionary of FCI maps
gridfile - Output filename
new_names - Write "g_yy" rather than "g_22"
metric2d - Output only 2D metrics.
"""
nx, ny, nz = grid.shape
# Get metric tensor
metric = grid.metric()
# Check if the magnetic field is in cylindrical coordinates
# If so, we need to change the gyy and g_yy metrics
pol_grid, ypos = grid.getPoloidalGrid(0)
Rmaj = magnetic_field.Rfunc(pol_grid.R, pol_grid.Z, ypos)
if Rmaj is not None:
# In cylindrical coordinates
Rmaj = np.zeros(grid.shape)
for yindex in range(grid.numberOfPoloidalGrids()):
pol_grid, ypos = grid.getPoloidalGrid(yindex)
Rmaj[:, yindex, :] = magnetic_field.Rfunc(pol_grid.R, pol_grid.Z,
ypos)
metric["gyy"] = 1. / Rmaj**2
metric["g_yy"] = Rmaj**2
# Get magnetic field
Bmag = np.zeros(grid.shape)
for yindex in range(grid.numberOfPoloidalGrids()):
pol_grid, ypos = grid.getPoloidalGrid(yindex)
Bmag[:, yindex, :] = magnetic_field.Bmag(pol_grid.R, pol_grid.Z, ypos)
# Metric is now 3D
if metric2d:
# Remove the Z dimension from metric components
print("WARNING: Outputting 2D metrics, discarding metric information.")
for key in metric:
try:
metric[key] = metric[key][:, :, 0]
except:
pass
# Make dz a constant
metric["dz"] = metric["dz"][0, 0]
# Add Rxy, Bxy
metric["Rxy"] = maps["R"][:, :, 0]
metric["Bxy"] = Bmag[:, :, 0]
with bdata.DataFile(gridfile, write=True, create=True) as f:
ixseps = nx + 1
f.write('nx', nx)
f.write('ny', ny)
f.write('nz', nz)
f.write("dx", metric["dx"])
f.write("dy", metric["dy"])
f.write("dz", metric["dz"])
f.write("ixseps1", ixseps)
f.write("ixseps2", ixseps)
# Metric tensor
if new_names:
for key, val in metric.items():
f.write(key, val)
else:
# Translate between output variable names and metric names
# Map from new to old names. Anything not in this dict
# is output unchanged
name_changes = {
"g_yy": "g_22",
"gyy": "g22",
"gxx": "g11",
"gxz": "g13",
"gzz": "g33",
"g_xx": "g_11",
"g_xz": "g_13",
"g_zz": "g_33"
}
for key in metric:
name = key
if name in name_changes:
name = name_changes[name]
f.write(name, metric[key])
# Magnetic field
f.write("B", Bmag)
# Maps - write everything to file
for key in maps:
f.write(key, maps[key])