def current(src, dst, par, gd, l1, l2, m1, m2, n1, n2, nghost):
n1shift, n2shift, m1shift, m2shift, l1shift, l2shift = der_limits(
n1, n2, m1, m2, l1, l2, nghost)
aa = np.array([
src['ax'][n1shift:n2shift, m1shift:m2shift, l1shift:l2shift],
src['ay'][n1shift:n2shift, m1shift:m2shift, l1shift:l2shift],
src['az'][n1shift:n2shift, m1shift:m2shift, l1shift:l2shift]
])
var = curl2(aa, gd.dx, gd.dy, gd.dz)
n1r, m1r, l1r = under_limits(n1, m1, l1, n1shift, m1shift, l1shift,
nghost)
return var[:, n1r:n2 - n1 + n1r, m1r:m2 - m1 + m1r, l1r:l2 - l1 + l1r]
def current(src, dst, key, par, gd, l1, l2, m1, m2, n1, n2, nghost):
if key == "jj":
n1shift, n2shift, m1shift, m2shift, l1shift, l2shift = der_limits(
n1, n2, m1, m2, l1, l2, nghost)
aa = np.array([
src["ax"][n1shift:n2shift, m1shift:m2shift, l1shift:l2shift],
src["ay"][n1shift:n2shift, m1shift:m2shift, l1shift:l2shift],
src["az"][n1shift:n2shift, m1shift:m2shift, l1shift:l2shift],
])
var = curl2(aa, gd.dx, gd.dy, gd.dz)
n1r, m1r, l1r = under_limits(n1, m1, l1, n1shift, m1shift, l1shift,
nghost)
return var[:, n1r:n2 - n1 + n1r, m1r:m2 - m1 + m1r,
l1r:l2 - l1 + l1r]
def __init__(self,
varfile='',
datadir='data/',
proc=-1,
ivar=-1,
quiet=False,
trimall=False,
format='native',
param=None,
dim=None,
index=None,
run2D=False,
magic=None,
setup=None):
"""
Description:
-----------
Read VAR files from pencil code. if proc < 0, then load all data
and assemble. otherwise, load VAR file from specified processor.
format -- one of (['native', 'n'], ['ieee-le', 'l'],
['ieee-be', 'B']) for byte-ordering
Params:
------
varfile=''
datadir='data/'
proc=-1
ivar=-1
quiet=False
trimall=False
format='native'
param=None
dim=None
index=None
run2D=False
Example of usage
------
ff=pc.read_var(trimall=True,ivar=100,magic=['tt','vort'])
"""
if (setup is not None):
datadir = os.path.expanduser(setup.datadir)
dim = setup.dim
param = setup.param
index = setup.index
run2D = setup.run2D
else:
datadir = os.path.expanduser(datadir)
if dim is None:
ldownsampled = 'VARd' in varfile
dim = read_dim(datadir, proc, down=ldownsampled)
if param is None:
param = read_param(datadir=datadir, quiet=quiet)
if index is None:
index = read_index(datadir=datadir, down=ldownsampled)
if dim.precision == 'D':
precision = 'd'
else:
precision = 'f'
if param.lwrite_aux:
totalvars = dim.mvar + dim.maux
else:
totalvars = dim.mvar
if 'VARd' in varfile:
if param.mvar_down > 0:
totalvars = param.mvar_down
# Read index.pro to get positions and "names"
# of variables in f(mx,my,mz,nvar).
# Thomas: seems useless now ?
#exec(index) # this loads the indicies.
if (not varfile):
if ivar < 0:
varfile = 'var.dat'
else:
varfile = 'VAR' + str(ivar)
if proc < 0:
procdirs = natural_sort(
filter(lambda s: s.startswith('proc'), os.listdir(datadir)))
else:
procdirs = ['proc' + str(proc)]
#global array
if (not run2D):
f = np.zeros((totalvars, dim.mz, dim.my, dim.mx), dtype=precision)
else:
if dim.ny == 1:
f = np.zeros((totalvars, dim.mz, dim.mx), dtype=precision)
else:
f = np.zeros((totalvars, dim.my, dim.mx), dtype=precision)
x = np.zeros(dim.mx, dtype=precision)
y = np.zeros(dim.my, dtype=precision)
z = np.zeros(dim.mz, dtype=precision)
for directory in procdirs:
proc = int(directory[4:])
procdim = read_dim(datadir, proc, down=ldownsampled)
if (not quiet):
#print "reading data from processor %i of %i ..." \ # Python 2
#% (proc, len(procdirs)) # Python 2
print("reading data from processor {0} of {1} ...".format(
proc, len(procdirs)))
mxloc = procdim.mx
myloc = procdim.my
mzloc = procdim.mz
#read data
filename = os.path.join(datadir, directory, varfile)
infile = npfile(filename, endian=format)
if (not run2D):
f_loc = infile.fort_read(precision,
shape=(-1, mzloc, myloc, mxloc))
else:
if dim.ny == 1:
f_loc = infile.fort_read(precision,
shape=(-1, mzloc, mxloc))
else:
f_loc = infile.fort_read(precision,
shape=(-1, myloc, mxloc))
raw_etc = infile.fort_read(precision)
infile.close()
t = raw_etc[0]
x_loc = raw_etc[1:mxloc + 1]
y_loc = raw_etc[mxloc + 1:mxloc + myloc + 1]
z_loc = raw_etc[mxloc + myloc + 1:mxloc + myloc + mzloc + 1]
if (param.lshear):
shear_offset = 1
deltay = raw_etc[-1]
else:
shear_offset = 0
dx = raw_etc[-3 - shear_offset]
dy = raw_etc[-2 - shear_offset]
dz = raw_etc[-1 - shear_offset]
if len(procdirs) > 1:
# Calculate where the local processor will go in
# the global array.
# Don't overwrite ghost zones of processor to the left (and
# accordingly in y and z direction--makes a difference on the
# diagonals).
#
# Recall that in NumPy, slicing is NON-INCLUSIVE on the right end
# ie, x[0:4] will slice all of a 4-digit array, not produce
# an error like in idl.
if procdim.ipx == 0:
i0x = 0
i1x = i0x + procdim.mx
i0xloc = 0
i1xloc = procdim.mx
else:
i0x = procdim.ipx * procdim.nx + procdim.nghostx
i1x = i0x + procdim.mx - procdim.nghostx
i0xloc = procdim.nghostx
i1xloc = procdim.mx
if procdim.ipy == 0:
i0y = 0
i1y = i0y + procdim.my
i0yloc = 0
i1yloc = procdim.my
else:
i0y = procdim.ipy * procdim.ny + procdim.nghosty
i1y = i0y + procdim.my - procdim.nghosty
i0yloc = procdim.nghosty
i1yloc = procdim.my
if procdim.ipz == 0:
i0z = 0
i1z = i0z + procdim.mz
i0zloc = 0
i1zloc = procdim.mz
else:
i0z = procdim.ipz * procdim.nz + procdim.nghostz
i1z = i0z + procdim.mz - procdim.nghostz
i0zloc = procdim.nghostz
i1zloc = procdim.mz
x[i0x:i1x] = x_loc[i0xloc:i1xloc]
y[i0y:i1y] = y_loc[i0yloc:i1yloc]
z[i0z:i1z] = z_loc[i0zloc:i1zloc]
if (not run2D):
f[:, i0z:i1z, i0y:i1y, i0x:i1x] = \
f_loc[:, i0zloc:i1zloc, i0yloc:i1yloc, i0xloc:i1xloc]
else:
if dim.ny == 1:
f[:, i0z:i1z, i0x:i1x] = \
f_loc[:, i0zloc:i1zloc, i0xloc:i1xloc]
else:
f[:, i0y:i1y, i0x:i1x] = \
f_loc[:, i0yloc:i1yloc, i0xloc:i1xloc]
else:
f = f_loc
x = x_loc
y = y_loc
z = z_loc
#endif MPI run
#endfor directories loop
if (magic is not None):
if ('bb' in magic):
# Compute the magnetic field before doing trimall.
aa = f[index['ax'] - 1:index['az'], ...]
self.bb = curl(aa, dx, dy, dz, x, y, z, run2D=param.lwrite_2d)
if (trimall):
self.bb = self.bb[:, dim.n1:dim.n2 + 1, dim.m1:dim.m2 + 1,
dim.l1:dim.l2 + 1]
if ('jj' in magic):
# Compute the electric current field before doing trimall.
aa = f[index['ax'] - 1:index['az'], ...]
self.jj = curl2(aa, dx, dy, dz, x, y, z)
if (trimall):
self.jj = self.jj[:, dim.n1:dim.n2 + 1, dim.m1:dim.m2 + 1,
dim.l1:dim.l2 + 1]
if ('vort' in magic):
# Compute the vorticity field before doing trimall.
uu = f[index['ux'] - 1:index['uz'], ...]
self.vort = curl(uu,
dx,
dy,
dz,
x,
y,
z,
run2D=param.lwrite_2d)
if (trimall):
if (param.lwrite_2d):
if (dim.nz == 1):
self.vort = self.vort[:, dim.m1:dim.m2 + 1,
dim.l1:dim.l2 + 1]
else:
self.vort = self.vort[:, dim.n1:dim.n2 + 1,
dim.l1:dim.l2 + 1]
else:
self.vort = self.vort[:, dim.n1:dim.n2 + 1,
dim.m1:dim.m2 + 1,
dim.l1:dim.l2 + 1]
# Trim the ghost zones of the global f-array if asked.
if trimall:
self.x = x[dim.l1:dim.l2 + 1]
self.y = y[dim.m1:dim.m2 + 1]
self.z = z[dim.n1:dim.n2 + 1]
if (not run2D):
self.f = f[:, dim.n1:dim.n2 + 1, dim.m1:dim.m2 + 1,
dim.l1:dim.l2 + 1]
else:
if dim.ny == 1:
self.f = f[:, dim.n1:dim.n2 + 1, dim.l1:dim.l2 + 1]
else:
self.f = f[:, dim.m1:dim.m2 + 1, dim.l1:dim.l2 + 1]
else:
self.x = x
self.y = y
self.z = z
self.f = f
self.l1 = dim.l1
self.l2 = dim.l2 + 1
self.m1 = dim.m1
self.m2 = dim.m2 + 1
self.n1 = dim.n1
self.n2 = dim.n2 + 1
# Assign an attribute to self for each variable defined in
# 'data/index.pro' so that e.g. self.ux is the x-velocity.
for key, value in index.items():
# print key,value.
if key != 'global_gg':
setattr(self, key, self.f[value - 1, ...])
# special treatment for vector quantities
if 'uu' in index.keys():
self.uu = self.f[index['ux'] - 1:index['uz'], ...]
if 'aa' in index.keys():
self.aa = self.f[index['ax'] - 1:index['az'], ...]
# Also treat Fcr (from cosmicrayflux) as a vector.
if 'fcr' in index.keys():
self.fcr = self.f[index['fcr'] - 1:index['fcr'] + 2, ...]
self.fcrx = self.fcr[0]
self.fcry = self.fcr[1]
self.fcrz = self.fcr[2]
self.t = t
self.dx = dx
self.dy = dy
self.dz = dz
if param.lshear:
self.deltay = deltay
# Do the rest of magic after the trimall (i.e. no additional curl...).
self.magic = magic
if self.magic is not None:
self.__magicAttributes(param)
def __init__(self, varfile='', datadir='data/', proc=-1, ivar=-1,
quiet=False, trimall=False, format='native',
param=None, dim=None, index=None, run2D=False,
magic=None, setup=None):
"""
Description:
-----------
Read VAR files from pencil code. if proc < 0, then load all data
and assemble. otherwise, load VAR file from specified processor.
format -- one of (['native', 'n'], ['ieee-le', 'l'],
['ieee-be', 'B']) for byte-ordering
Params:
------
varfile=''
datadir='data/'
proc=-1
ivar=-1
quiet=False
trimall=False
format='native'
param=None
dim=None
index=None
run2D=False
"""
if (setup is not None):
datadir = os.path.expanduser(setup.datadir)
dim = setup.dim
param = setup.param
index = setup.index
run2D = setup.run2D
else:
datadir = os.path.expanduser(datadir)
if dim is None:
dim = read_dim(datadir,proc)
if param is None:
param = read_param(datadir=datadir, quiet=quiet)
if index is None:
index = read_index(datadir=datadir)
if dim.precision == 'D':
precision = 'd'
else:
precision = 'f'
if param.lwrite_aux:
totalvars = dim.mvar+dim.maux
else:
totalvars = dim.mvar
# Read index.pro to get positions and "names"
# of variables in f(mx,my,mz,nvar).
# Thomas: seems useless now ?
#exec(index) # this loads the indicies.
if (not varfile):
if ivar < 0:
varfile = 'var.dat'
else:
varfile = 'VAR'+str(ivar)
if proc < 0:
procdirs = natural_sort(filter(lambda s:s.startswith('proc'),
os.listdir(datadir)))
else:
procdirs = ['proc'+str(proc)]
#global array
if (not run2D):
f = np.zeros((totalvars, dim.mz, dim.my, dim.mx),
dtype=precision)
else:
if dim.ny == 1:
f = np.zeros((totalvars, dim.mz, dim.mx), dtype=precision)
else:
f = np.zeros((totalvars, dim.my, dim.mx), dtype=precision)
x = np.zeros(dim.mx, dtype=precision)
y = np.zeros(dim.my, dtype=precision)
z = np.zeros(dim.mz, dtype=precision)
for directory in procdirs:
proc = int(directory[4:])
procdim = read_dim(datadir, proc)
if (not quiet):
print "reading data from processor %i of %i ..." \
% (proc, len(procdirs))
mxloc = procdim.mx
myloc = procdim.my
mzloc = procdim.mz
#read data
filename = os.path.join(datadir,directory,varfile)
infile = npfile(filename, endian=format)
if (not run2D):
f_loc = infile.fort_read(precision,
shape=(-1, mzloc, myloc, mxloc))
else:
if dim.ny == 1:
f_loc = infile.fort_read(precision, shape=(-1, mzloc, mxloc))
else:
f_loc = infile.fort_read(precision, shape=(-1, myloc, mxloc))
raw_etc = infile.fort_read(precision)
infile.close()
t = raw_etc[0]
x_loc = raw_etc[1:mxloc+1]
y_loc = raw_etc[mxloc+1:mxloc+myloc+1]
z_loc = raw_etc[mxloc+myloc+1:mxloc+myloc+mzloc+1]
if (param.lshear):
shear_offset = 1
deltay = raw_etc[-1]
else:
shear_offset = 0
dx = raw_etc[-3-shear_offset]
dy = raw_etc[-2-shear_offset]
dz = raw_etc[-1-shear_offset]
if len(procdirs) > 1:
# Calculate where the local processor will go in
# the global array.
# Don't overwrite ghost zones of processor to the left (and
# accordingly in y and z direction--makes a difference on the
# diagonals).
#
# Recall that in NumPy, slicing is NON-INCLUSIVE on the right end
# ie, x[0:4] will slice all of a 4-digit array, not produce
# an error like in idl.
if procdim.ipx == 0:
i0x = 0
i1x = i0x+procdim.mx
i0xloc = 0
i1xloc = procdim.mx
else:
i0x = procdim.ipx*procdim.nx+procdim.nghostx
i1x = i0x+procdim.mx-procdim.nghostx
i0xloc = procdim.nghostx
i1xloc = procdim.mx
if procdim.ipy == 0:
i0y = 0
i1y = i0y+procdim.my
i0yloc = 0
i1yloc = procdim.my
else:
i0y = procdim.ipy*procdim.ny+procdim.nghosty
i1y = i0y+procdim.my-procdim.nghosty
i0yloc = procdim.nghosty
i1yloc = procdim.my
if procdim.ipz == 0:
i0z = 0
i1z = i0z+procdim.mz
i0zloc = 0
i1zloc = procdim.mz
else:
i0z = procdim.ipz*procdim.nz+procdim.nghostz
i1z = i0z+procdim.mz-procdim.nghostz
i0zloc = procdim.nghostz
i1zloc = procdim.mz
x[i0x:i1x] = x_loc[i0xloc:i1xloc]
y[i0y:i1y] = y_loc[i0yloc:i1yloc]
z[i0z:i1z] = z_loc[i0zloc:i1zloc]
if (not run2D):
f[:, i0z:i1z, i0y:i1y, i0x:i1x] = \
f_loc[:, i0zloc:i1zloc, i0yloc:i1yloc, i0xloc:i1xloc]
else:
if dim.ny == 1:
f[:, i0z:i1z, i0x:i1x] = \
f_loc[:, i0zloc:i1zloc, i0xloc:i1xloc]
else:
f[:, i0y:i1y, i0x:i1x] = \
f_loc[:, i0yloc:i1yloc, i0xloc:i1xloc]
else:
f = f_loc
x = x_loc
y = y_loc
z = z_loc
#endif MPI run
#endfor directories loop
if (magic is not None):
if ('bb' in magic):
# Compute the magnetic field before doing trimall.
aa = f[index['ax']-1:index['az'],...]
self.bb = curl(aa,dx,dy,dz,run2D=param.lwrite_2d)
if (trimall): self.bb=self.bb[:, dim.n1:dim.n2+1,
dim.m1:dim.m2+1, dim.l1:dim.l2+1]
if ('jj' in magic):
# Compute the electric current field before doing trimall.
aa = f[index['ax']-1:index['az'],...]
self.jj = curl2(aa,dx,dy,dz)
if (trimall): self.jj=self.jj[:, dim.n1:dim.n2+1,
dim.m1:dim.m2+1, dim.l1:dim.l2+1]
if ('vort' in magic):
# Compute the vorticity field before doing trimall.
uu = f[index['ux']-1:index['uz'],...]
self.vort = curl(uu,dx,dy,dz,run2D=param.lwrite_2d)
if (trimall):
if (param.lwrite_2d):
if (dim.nz == 1):
self.vort=self.vort[:, dim.m1:dim.m2+1,
dim.l1:dim.l2+1]
else:
self.vort=self.vort[:, dim.n1:dim.n2+1,
dim.l1:dim.l2+1]
else:
self.vort=self.vort[:, dim.n1:dim.n2+1,
dim.m1:dim.m2+1, dim.l1:dim.l2+1]
# Trim the ghost zones of the global f-array if asked.
if trimall:
self.x = x[dim.l1:dim.l2+1]
self.y = y[dim.m1:dim.m2+1]
self.z = z[dim.n1:dim.n2+1]
if (not run2D):
self.f = f[:, dim.n1:dim.n2+1, dim.m1:dim.m2+1, dim.l1:dim.l2+1]
else:
if dim.ny == 1:
self.f = f[:, dim.n1:dim.n2+1, dim.l1:dim.l2+1]
else:
self.f = f[:, dim.m1:dim.m2+1, dim.l1:dim.l2+1]
else:
self.x = x
self.y = y
self.z = z
self.f = f
self.l1 = dim.l1
self.l2 = dim.l2+1
self.m1 = dim.m1
self.m2 = dim.m2+1
self.n1 = dim.n1
self.n2 = dim.n2+1
# Assign an attribute to self for each variable defined in
# 'data/index.pro' so that e.g. self.ux is the x-velocity.
for key,value in index.items():
# print key,value.
if key != 'global_gg':
setattr(self,key,self.f[value-1,...])
# special treatment for vector quantities
if index.has_key('uu'):
self.uu = self.f[index['ux']-1:index['uz'],...]
if index.has_key('aa'):
self.aa = self.f[index['ax']-1:index['az'],...]
# Also treat Fcr (from cosmicrayflux) as a vector.
if index.has_key('fcr'):
self.fcr = self.f[index['fcr']-1:index['fcr']+2,...]
self.fcrx = self.fcr[0]
self.fcry = self.fcr[1]
self.fcrz = self.fcr[2]
self.t = t
self.dx = dx
self.dy = dy
self.dz = dz
if param.lshear:
self.deltay = deltay
# Do the rest of magic after the trimall (i.e. no additional curl...).
self.magic = magic
if self.magic is not None:
self.__magicAttributes(param)
def magnetic_reynolds(uu,
param,
grid,
aa=list(),
bb=list(),
jj=list(),
nghost=3,
lmix=True,
quiet=True):
"""
Computes the magnetic Reynolds number from the advective and effective
resistive expressions in the induction equation.
call signature:
magnetic_reynolds(uu, param, grid, aa=None, bb=None, jj=None, nghost=3):
Keyword arguments:
*uu*:
The velocity field [3,mz,my,mx] from the simulation data
*param*:
The Param simulation object with resistivity data information
*grid*:
The Grid simulation object
*aa*:
The vector potential if bb is not present or hyper diffusion
*bb*:
The magnetic field
*jj*:
The current density field
*nghost*:
The number of ghost zones appropriate to the order of accuracy
*lmix*:
Option not to include hyper values when Laplacian values present
"""
if len(bb) == 0 and len(aa) == 0 and len(jj) == 0:
print('magnetic_reynolds WARNING: no aa, bb nor jj provided\n' +
'aa or bb must be provided or aa for only hyper resistivity')
#resistive force
lres, lhyper3 = False, False
for iresi in param.iresistivity:
iresi = str.strip(iresi, '\n')
if 'hyper' not in iresi and len(iresi) > 0:
lres = True
if 'hyper3' in iresi:
lhyper3 = True
fresi = np.zeros_like(uu)
if lres:
if lhyper3:
lhyper3 = lhyper3 == lmix
if len(jj) == 0:
if len(aa) == 0:
print(
'magnetic_reynolds WARNING: calculating jj without aa\n',
'provide aa or jj directly for accurate boundary values')
jj = curl(bb,
grid.dx,
grid.dy,
grid.dz,
x=grid.x,
y=grid.y,
coordinate_system=param.coord_system)
else:
jj = curl2(aa,
grid.dx,
grid.dy,
grid.dz,
x=grid.x,
y=grid.y,
coordinate_system=param.coord_system)
for j in range(0, 3):
jj[j, :nghost, :, :] = jj[j, -2 * nghost:-nghost, :, :]
jj[j, -nghost:, :, :] = jj[j, nghost:2 * nghost, :, :]
jj[j, :, :nghost, :] = jj[j, :, -2 * nghost:-nghost, :]
jj[j, :, -nghost:, :] = jj[j, :, nghost:2 * nghost, :]
jj[j, :, :, :nghost] = jj[j, :, :, -2 * nghost:-nghost]
jj[j, :, :, -nghost:] = jj[j, :, :, nghost:2 * nghost]
fresi = fresi + param.eta * param.mu0 * jj
for iresi in param.iresistivity:
iresi = str.strip(iresi, '\n')
if 'eta-const' not in iresi and 'hyper' not in iresi\
and len(iresi) > 0:
print(
'magnetic_reynolds WARNING: ' + iresi +
' not implemented\n' +
'terms may be missing from the standard resistive forces')
if lhyper3:
if len(aa) == 0:
print('magnetic_reynolds WARNING: no aa provided\n' +
'aa must be provided for hyper resistivity')
return 1
else:
del6a = np.zeros_like(aa)
for j in range(0, 3):
del6a[j] = del6(aa[j], grid.dx, grid.dy, grid.dz)
del6a[j, :nghost, :, :] = del6a[j, -2 * nghost:-nghost, :, :]
del6a[j, -nghost:, :, :] = del6a[j, nghost:2 * nghost, :, :]
del6a[j, :, :nghost, :] = del6a[j, :, -2 * nghost:-nghost, :]
del6a[j, :, -nghost:, :] = del6a[j, :, nghost:2 * nghost, :]
del6a[j, :, :, :nghost] = del6a[j, :, :, -2 * nghost:-nghost]
del6a[j, :, :, -nghost:] = del6a[j, :, :, nghost:2 * nghost]
#del6 for non-cartesian tba
#del6a[j] = del6(aa[j],grid.dx,grid.dy,grid.dz,x=grid.x,y=grid.y,
# coordinate_system=param.coord_system)
#effective at l > 5 grid.dx?
fresi = fresi + param.eta_hyper3 * del6a
del (del6a)
fresi2 = np.sqrt(dot2(fresi))
del (fresi)
#advective force
if len(bb) == 0:
if len(aa) == 0:
print(
'magnetic_reynolds WARNING: calculating uu x bb without bb\n',
'provide aa or bb directly to proceed')
return 1
else:
bb = curl(aa,
grid.dx,
grid.dy,
grid.dz,
x=grid.x,
y=grid.y,
coordinate_system=param.coord_system)
for j in range(0, 3):
bb[j, :nghost, :, :] = bb[j, -2 * nghost:-nghost, :, :]
bb[j, -nghost:, :, :] = bb[j, nghost:2 * nghost, :, :]
bb[j, :, :nghost, :] = bb[j, :, -2 * nghost:-nghost, :]
bb[j, :, -nghost:, :] = bb[j, :, nghost:2 * nghost, :]
bb[j, :, :, :nghost] = bb[j, :, :, -2 * nghost:-nghost]
bb[j, :, :, -nghost:] = bb[j, :, :, nghost:2 * nghost]
advec = cross(uu, bb)
advec2 = np.sqrt(dot2(advec))
del (advec)
#avoid division by zero
if fresi2.max() > 0:
fresi2[np.where(fresi2 == 0)] = fresi2[np.where(fresi2 > 0)].min()
Rm = advec2 / fresi2
#set minimum floor to exclude zero-valued Rm
if Rm.max() > 0:
Rm[np.where(Rm == 0)] = Rm[np.where(Rm > 0)].min()
else:
print('Rm undefined')
else:
Rm = advec2
print('Rm undefined')
return Rm
def read(self, var_file='', datadir='data', proc=-1, ivar=-1, quiet=True,
trimall=False, magic=None, sim=None, precision='d',
lpersist=False, dtype=np.float64):
"""
Read VAR files from Pencil Code. If proc < 0, then load all data
and assemble, otherwise load VAR file from specified processor.
The file format written by output() (and used, e.g. in var.dat)
consists of the followinig Fortran records:
1. data(mx, my, mz, nvar)
2. t(1), x(mx), y(my), z(mz), dx(1), dy(1), dz(1), deltay(1)
Here nvar denotes the number of slots, i.e. 1 for one scalar field, 3
for one vector field, 8 for var.dat in the case of MHD with entropy.
but, deltay(1) is only there if lshear is on! need to know parameters.
call signature:
var(var_file='', datadir='data', proc=-1, ivar=-1, quiet=True,
trimall=False, magic=None, sim=None, precision='d')
Keyword arguments:
var_file: Name of the VAR file.
datadir: Directory where the data is stored.
proc: Processor to be read. If -1 read all and assemble to one array.
ivar: Index of the VAR file, if var_file is not specified.
quiet: Flag for switching off output.
trimall: Trim the data cube to exclude ghost zones.
magic: Values to be computed from the data, e.g. B = curl(A).
sim: Simulation sim object.
precision: Float (f), double (d) or half (half).
dtype: precision for var.obj, default double
"""
import os
from scipy.io import FortranFile
from pencil.math.derivatives import curl, curl2
from pencil import read
from pencil.sim import __Simulation__
def persist(self, infile=None, precision='d', quiet=quiet):
"""An open Fortran file potentially containing persistent variables appended
to the f array and grid data are read from the first proc data
Record types provide the labels and id record for the peristent
variables in the depricated fortran binary format
"""
record_types = {}
for key in read.record_types.keys():
if read.record_types[key][1] == 'd':
record_types[key]=(read.record_types[key][0],
precision)
else:
record_types[key] = read.record_types[key]
try:
tmp_id = infile.read_record('h')
except:
return -1
block_id = 0
for i in range(2000):
i += 1
tmp_id = infile.read_record('h')
block_id = tmp_id[0]
if block_id == 2000:
break
for key in record_types.keys():
if record_types[key][0] == block_id:
tmp_val = infile.read_record(record_types[key][1])
self.__setattr__(key, tmp_val[0])
if not quiet:
print(key, record_types[key][0],
record_types[key][1],tmp_val)
return self
dim = None
param = None
index = None
if isinstance(sim, __Simulation__):
datadir = os.path.expanduser(sim.datadir)
dim = sim.dim
param = read.param(datadir=sim.datadir, quiet=True,
conflicts_quiet=True)
index = read.index(datadir=sim.datadir)
else:
datadir = os.path.expanduser(datadir)
if dim is None:
if var_file[0:2].lower() == 'og':
dim = read.ogdim(datadir, proc)
else:
if var_file[0:4] == 'VARd':
dim = read.dim(datadir, proc, down=True)
else:
dim = read.dim(datadir, proc)
if param is None:
param = read.param(datadir=datadir, quiet=quiet,
conflicts_quiet=True)
if index is None:
index = read.index(datadir=datadir)
if param.lwrite_aux:
total_vars = dim.mvar + dim.maux
else:
total_vars = dim.mvar
if os.path.exists(os.path.join(datadir, 'grid.h5')):
#
# Read HDF5 files.
#
import h5py
run2D = param.lwrite_2d
# Set up the global array.
if not run2D:
self.f = np.zeros((total_vars, dim.mz, dim.my, dim.mx),
dtype=dtype)
else:
if dim.ny == 1:
self.f = np.zeros((total_vars, dim.mz, dim.mx),
dtype=dtype)
else:
self.f = np.zeros((total_vars, dim.my, dim.mx),
dtype=dtype)
if not var_file:
if ivar < 0:
var_file = 'var.h5'
else:
var_file = 'VAR' + str(ivar) + '.h5'
file_name = os.path.join(datadir, 'allprocs', var_file)
with h5py.File(file_name, 'r') as tmp:
for key in tmp['data'].keys():
self.f[index.__getattribute__(key)-1, :] = dtype(
tmp['data/'+key][:])
t = (tmp['time'][()]).astype(precision)
x = (tmp['grid/x'][()]).astype(precision)
y = (tmp['grid/y'][()]).astype(precision)
z = (tmp['grid/z'][()]).astype(precision)
dx = (tmp['grid/dx'][()]).astype(precision)
dy = (tmp['grid/dy'][()]).astype(precision)
dz = (tmp['grid/dz'][()]).astype(precision)
if param.lshear:
deltay = (tmp['persist/shear_delta_y'][(0)]).astype(precision)
if lpersist:
for key in tmp['persist'].keys():
self.__setattr__(key, (tmp['persist'][key][0]).astype(precision))
else:
#
# Read scattered Fortran binary files.
#
run2D = param.lwrite_2d
if dim.precision == 'D':
read_precision = 'd'
else:
read_precision = 'f'
if not var_file:
if ivar < 0:
var_file = 'var.dat'
else:
var_file = 'VAR' + str(ivar)
if proc < 0:
proc_dirs = self.__natural_sort(
filter(lambda s: s.startswith('proc'),
os.listdir(datadir)))
if (proc_dirs.count("proc_bounds.dat") > 0):
proc_dirs.remove("proc_bounds.dat")
if param.lcollective_io:
# A collective IO strategy is being used
proc_dirs = ['allprocs']
# else:
# proc_dirs = proc_dirs[::dim.nprocx*dim.nprocy]
else:
proc_dirs = ['proc' + str(proc)]
# Set up the global array.
if not run2D:
self.f = np.zeros((total_vars, dim.mz, dim.my, dim.mx),
dtype=dtype)
else:
if dim.ny == 1:
self.f = np.zeros((total_vars, dim.mz, dim.mx),
dtype=dtype)
else:
self.f = np.zeros((total_vars, dim.my, dim.mx),
dtype=dtype)
x = np.zeros(dim.mx, dtype=precision)
y = np.zeros(dim.my, dtype=precision)
z = np.zeros(dim.mz, dtype=precision)
for directory in proc_dirs:
if not param.lcollective_io:
proc = int(directory[4:])
if var_file[0:2].lower() == 'og':
procdim = read.ogdim(datadir, proc)
else:
if var_file[0:4] == 'VARd':
procdim = read.dim(datadir, proc, down=True)
else:
procdim = read.dim(datadir, proc)
if not quiet:
print("Reading data from processor"+
" {0} of {1} ...".format(proc, len(proc_dirs)))
else:
# A collective IO strategy is being used
procdim = dim
# else:
# procdim.mx = dim.mx
# procdim.my = dim.my
# procdim.nx = dim.nx
# procdim.ny = dim.ny
# procdim.ipx = dim.ipx
# procdim.ipy = dim.ipy
mxloc = procdim.mx
myloc = procdim.my
mzloc = procdim.mz
# Read the data.
file_name = os.path.join(datadir, directory, var_file)
infile = FortranFile(file_name)
if not run2D:
f_loc = dtype(infile.read_record(dtype=read_precision))
f_loc = f_loc.reshape((-1, mzloc, myloc, mxloc))
else:
if dim.ny == 1:
f_loc = dtype(infile.read_record(dtype=read_precision))
f_loc = f_loc.reshape((-1, mzloc, mxloc))
else:
f_loc = dtype(infile.read_record(dtype=read_precision))
f_loc = f_loc.reshape((-1, myloc, mxloc))
raw_etc = infile.read_record(dtype=read_precision)
if lpersist:
persist(self, infile=infile, precision=read_precision, quiet=quiet)
infile.close()
t = raw_etc[0]
x_loc = raw_etc[1:mxloc+1]
y_loc = raw_etc[mxloc+1:mxloc+myloc+1]
z_loc = raw_etc[mxloc+myloc+1:mxloc+myloc+mzloc+1]
if param.lshear:
shear_offset = 1
deltay = raw_etc[-1]
else:
shear_offset = 0
dx = raw_etc[-3-shear_offset]
dy = raw_etc[-2-shear_offset]
dz = raw_etc[-1-shear_offset]
if len(proc_dirs) > 1:
# Calculate where the local processor will go in
# the global array.
#
# Don't overwrite ghost zones of processor to the
# left (and accordingly in y and z direction -- makes
# a difference on the diagonals)
#
# Recall that in NumPy, slicing is NON-INCLUSIVE on
# the right end, ie, x[0:4] will slice all of a
# 4-digit array, not produce an error like in idl.
if procdim.ipx == 0:
i0x = 0
i1x = i0x + procdim.mx
i0xloc = 0
i1xloc = procdim.mx
else:
i0x = procdim.ipx*procdim.nx + procdim.nghostx
i1x = i0x + procdim.mx - procdim.nghostx
i0xloc = procdim.nghostx
i1xloc = procdim.mx
if procdim.ipy == 0:
i0y = 0
i1y = i0y + procdim.my
i0yloc = 0
i1yloc = procdim.my
else:
i0y = procdim.ipy*procdim.ny + procdim.nghosty
i1y = i0y + procdim.my - procdim.nghosty
i0yloc = procdim.nghosty
i1yloc = procdim.my
if procdim.ipz == 0:
i0z = 0
i1z = i0z+procdim.mz
i0zloc = 0
i1zloc = procdim.mz
else:
i0z = procdim.ipz*procdim.nz + procdim.nghostz
i1z = i0z + procdim.mz - procdim.nghostz
i0zloc = procdim.nghostz
i1zloc = procdim.mz
x[i0x:i1x] = x_loc[i0xloc:i1xloc]
y[i0y:i1y] = y_loc[i0yloc:i1yloc]
z[i0z:i1z] = z_loc[i0zloc:i1zloc]
if not run2D:
self.f[:, i0z:i1z, i0y:i1y, i0x:i1x] = f_loc[:, i0zloc:i1zloc,
i0yloc:i1yloc, i0xloc:i1xloc]
else:
if dim.ny == 1:
self.f[:, i0z:i1z, i0x:i1x] = f_loc[:, i0zloc:i1zloc, i0xloc:i1xloc]
else:
self.f[i0z:i1z, i0y:i1y, i0x:i1x] = f_loc[i0zloc:i1zloc,
i0yloc:i1yloc, i0xloc:i1xloc]
else:
self.f = f_loc
x = x_loc
y = y_loc
z = z_loc
if magic is not None:
if 'bb' in magic:
# Compute the magnetic field before doing trimall.
aa = self.f[index.ax-1:index.az, ...]
self.bb = dtype(curl(aa, dx, dy, dz, x=x, y=y, run2D=run2D,
coordinate_system=param.coord_system))
if trimall:
self.bb = self.bb[:, dim.n1:dim.n2+1,
dim.m1:dim.m2+1, dim.l1:dim.l2+1]
if 'jj' in magic:
# Compute the electric current field before doing trimall.
aa = self.f[index.ax-1:index.az, ...]
self.jj = dtype(curl2(aa, dx, dy, dz, x=x, y=y,
coordinate_system=param.coord_system))
if trimall:
self.jj = self.jj[:, dim.n1:dim.n2+1,
dim.m1:dim.m2+1, dim.l1:dim.l2+1]
if 'vort' in magic:
# Compute the vorticity field before doing trimall.
uu = self.f[index.ux-1:index.uz, ...]
self.vort = dtype(curl(uu, dx, dy, dz, x=x, y=y, run2D=run2D,
coordinate_system=param.coord_system))
if trimall:
if run2D:
if dim.nz == 1:
self.vort = self.vort[:, dim.m1:dim.m2+1,
dim.l1:dim.l2+1]
else:
self.vort = self.vort[:, dim.n1:dim.n2+1,
dim.l1:dim.l2+1]
else:
self.vort = self.vort[:, dim.n1:dim.n2+1,
dim.m1:dim.m2+1,
dim.l1:dim.l2+1]
# Trim the ghost zones of the global f-array if asked.
if trimall:
self.x = x[dim.l1:dim.l2+1]
self.y = y[dim.m1:dim.m2+1]
self.z = z[dim.n1:dim.n2+1]
if not run2D:
self.f = self.f[:, dim.n1:dim.n2+1,
dim.m1:dim.m2+1, dim.l1:dim.l2+1]
else:
if dim.ny == 1:
self.f = self.f[:, dim.n1:dim.n2+1, dim.l1:dim.l2+1]
else:
self.f = self.f[:, dim.m1:dim.m2+1, dim.l1:dim.l2+1]
else:
self.x = x
self.y = y
self.z = z
self.l1 = dim.l1
self.l2 = dim.l2 + 1
self.m1 = dim.m1
self.m2 = dim.m2 + 1
self.n1 = dim.n1
self.n2 = dim.n2 + 1
# Assign an attribute to self for each variable defined in
# 'data/index.pro' so that e.g. self.ux is the x-velocity
aatest = []
uutest = []
for key in index.__dict__.keys():
if 'aatest' in key:
aatest.append(key)
if 'uutest' in key:
uutest.append(key)
if key != 'global_gg' and key != 'keys' and 'aatest' not in key\
and 'uutest' not in key:
value = index.__dict__[key]
setattr(self, key, self.f[value-1, ...])
# Special treatment for vector quantities.
if hasattr(index, 'uu'):
self.uu = self.f[index.ux-1:index.uz, ...]
if hasattr(index, 'aa'):
self.aa = self.f[index.ax-1:index.az, ...]
if hasattr(index, 'uu_sph'):
self.uu_sph = self.f[index.uu_sphx-1:index.uu_sphz, ...]
if hasattr(index, 'bb_sph'):
self.bb_sph = self.f[index.bb_sphx-1:index.bb_sphz, ...]
# Special treatment for test method vector quantities.
#Note index 1,2,3,...,0 last vector may be the zero field/flow
if hasattr(index, 'aatest1'):
naatest = int(len(aatest)/3)
for j in range(0,naatest):
key = 'aatest'+str(np.mod(j+1,naatest))
value = index.__dict__['aatest1'] + 3*j
setattr(self, key, self.f[value-1:value+2, ...])
if hasattr(index, 'uutest1'):
nuutest = int(len(uutest)/3)
for j in range(0,nuutest):
key = 'uutest'+str(np.mod(j+1,nuutest))
value = index.__dict__['uutest'] + 3*j
setattr(self, key, self.f[value-1:value+2, ...])
self.t = t
self.dx = dx
self.dy = dy
self.dz = dz
if param.lshear:
self.deltay = deltay
# Do the rest of magic after the trimall (i.e. no additional curl.)
self.magic = magic
if self.magic is not None:
self.magic_attributes(param, dtype=dtype)
def read(
self,
var_file="",
datadir="data",
proc=-1,
ivar=-1,
quiet=True,
trimall=False,
magic=None,
sim=None,
precision="d",
lpersist=False,
dtype=np.float64,
):
"""
read(var_file='', datadir='data', proc=-1, ivar=-1, quiet=True,
trimall=False, magic=None, sim=None, precision='f')
Read VAR files from Pencil Code. If proc < 0, then load all data
and assemble, otherwise load VAR file from specified processor.
The file format written by output() (and used, e.g. in var.dat)
consists of the followinig Fortran records:
1. data(mx, my, mz, nvar)
2. t(1), x(mx), y(my), z(mz), dx(1), dy(1), dz(1), deltay(1)
Here nvar denotes the number of slots, i.e. 1 for one scalar field, 3
for one vector field, 8 for var.dat in the case of MHD with entropy.
but, deltay(1) is only there if lshear is on! need to know parameters.
Parameters
----------
var_file : string
Name of the VAR file.
If not specified, use var.dat (which is the latest snapshot of the fields)
datadir : string
Directory where the data is stored.
proc : int
Processor to be read. If -1 read all and assemble to one array.
ivar : int
Index of the VAR file, if var_file is not specified.
quiet : bool
Flag for switching off output.
trimall : bool
Trim the data cube to exclude ghost zones.
magic : bool
Values to be computed from the data, e.g. B = curl(A).
sim : pencil code simulation object
Contains information about the local simulation.
precision : string
Float 'f', double 'd' or half 'half'.
lpersist : bool
Read the persistent variables if they exist
Returns
-------
DataCube
Instance of the pencil.read.var.DataCube class.
All of the computed fields are imported as class members.
Examples
--------
Read the latest var.dat file and print the shape of the uu array:
>>> var = pc.read.var()
>>> print(var.uu.shape)
Read the VAR2 file, compute the magnetic field B = curl(A),
the vorticity omega = curl(u) and remove the ghost zones:
>>> var = pc.read.var(var_file='VAR2', magic=['bb', 'vort'], trimall=True)
>>> print(var.bb.shape)
"""
import os
from scipy.io import FortranFile
from pencil.math.derivatives import curl, curl2
from pencil import read
from pencil.sim import __Simulation__
def persist(self, infile=None, precision="d", quiet=quiet):
"""An open Fortran file potentially containing persistent variables appended
to the f array and grid data are read from the first proc data
Record types provide the labels and id record for the peristent
variables in the depricated fortran binary format
"""
record_types = {}
for key in read.record_types.keys():
if read.record_types[key][1] == "d":
record_types[key] = (read.record_types[key][0], precision)
else:
record_types[key] = read.record_types[key]
try:
tmp_id = infile.read_record("h")
except:
return -1
block_id = 0
for i in range(2000):
i += 1
tmp_id = infile.read_record("h")
block_id = tmp_id[0]
if block_id == 2000:
break
for key in record_types.keys():
if record_types[key][0] == block_id:
tmp_val = infile.read_record(record_types[key][1])
self.__setattr__(key, tmp_val[0])
if not quiet:
print(key, record_types[key][0],
record_types[key][1], tmp_val)
return self
dim = None
param = None
index = None
if isinstance(sim, __Simulation__):
datadir = os.path.expanduser(sim.datadir)
dim = sim.dim
param = read.param(datadir=sim.datadir,
quiet=True,
conflicts_quiet=True)
index = read.index(datadir=sim.datadir)
else:
datadir = os.path.expanduser(datadir)
if dim is None:
if var_file[0:2].lower() == "og":
dim = read.ogdim(datadir, proc)
else:
if var_file[0:4] == "VARd":
dim = read.dim(datadir, proc, down=True)
else:
dim = read.dim(datadir, proc)
if param is None:
param = read.param(datadir=datadir,
quiet=quiet,
conflicts_quiet=True)
if index is None:
index = read.index(datadir=datadir)
if param.lwrite_aux:
total_vars = dim.mvar + dim.maux
else:
total_vars = dim.mvar
if os.path.exists(os.path.join(datadir, "grid.h5")):
#
# Read HDF5 files.
#
import h5py
run2D = param.lwrite_2d
# Set up the global array.
if not run2D:
self.f = np.zeros((total_vars, dim.mz, dim.my, dim.mx),
dtype=dtype)
else:
if dim.ny == 1:
self.f = np.zeros((total_vars, dim.mz, dim.mx),
dtype=dtype)
else:
self.f = np.zeros((total_vars, dim.my, dim.mx),
dtype=dtype)
if not var_file:
if ivar < 0:
var_file = "var.h5"
else:
var_file = "VAR" + str(ivar) + ".h5"
file_name = os.path.join(datadir, "allprocs", var_file)
with h5py.File(file_name, "r") as tmp:
for key in tmp["data"].keys():
self.f[index.__getattribute__(key) - 1, :] = dtype(
tmp["data/" + key][:])
t = (tmp["time"][()]).astype(precision)
x = (tmp["grid/x"][()]).astype(precision)
y = (tmp["grid/y"][()]).astype(precision)
z = (tmp["grid/z"][()]).astype(precision)
dx = (tmp["grid/dx"][()]).astype(precision)
dy = (tmp["grid/dy"][()]).astype(precision)
dz = (tmp["grid/dz"][()]).astype(precision)
if param.lshear:
deltay = (tmp["persist/shear_delta_y"][(
0)]).astype(precision)
if lpersist:
for key in tmp["persist"].keys():
self.__setattr__(
key, (tmp["persist"][key][0]).astype(precision))
else:
#
# Read scattered Fortran binary files.
#
run2D = param.lwrite_2d
if dim.precision == "D":
read_precision = "d"
else:
read_precision = "f"
if not var_file:
if ivar < 0:
var_file = "var.dat"
else:
var_file = "VAR" + str(ivar)
if proc < 0:
proc_dirs = self.__natural_sort(
filter(lambda s: s.startswith("proc"),
os.listdir(datadir)))
if proc_dirs.count("proc_bounds.dat") > 0:
proc_dirs.remove("proc_bounds.dat")
if param.lcollective_io:
# A collective IO strategy is being used
proc_dirs = ["allprocs"]
# else:
# proc_dirs = proc_dirs[::dim.nprocx*dim.nprocy]
else:
proc_dirs = ["proc" + str(proc)]
# Set up the global array.
if not run2D:
self.f = np.zeros((total_vars, dim.mz, dim.my, dim.mx),
dtype=dtype)
else:
if dim.ny == 1:
self.f = np.zeros((total_vars, dim.mz, dim.mx),
dtype=dtype)
else:
self.f = np.zeros((total_vars, dim.my, dim.mx),
dtype=dtype)
x = np.zeros(dim.mx, dtype=precision)
y = np.zeros(dim.my, dtype=precision)
z = np.zeros(dim.mz, dtype=precision)
for directory in proc_dirs:
if not param.lcollective_io:
proc = int(directory[4:])
if var_file[0:2].lower() == "og":
procdim = read.ogdim(datadir, proc)
else:
if var_file[0:4] == "VARd":
procdim = read.dim(datadir, proc, down=True)
else:
procdim = read.dim(datadir, proc)
if not quiet:
print("Reading data from processor" +
" {0} of {1} ...".format(proc, len(proc_dirs)))
else:
# A collective IO strategy is being used
procdim = dim
# else:
# procdim.mx = dim.mx
# procdim.my = dim.my
# procdim.nx = dim.nx
# procdim.ny = dim.ny
# procdim.ipx = dim.ipx
# procdim.ipy = dim.ipy
mxloc = procdim.mx
myloc = procdim.my
mzloc = procdim.mz
# Read the data.
file_name = os.path.join(datadir, directory, var_file)
infile = FortranFile(file_name)
if not run2D:
f_loc = dtype(infile.read_record(dtype=read_precision))
f_loc = f_loc.reshape((-1, mzloc, myloc, mxloc))
else:
if dim.ny == 1:
f_loc = dtype(infile.read_record(dtype=read_precision))
f_loc = f_loc.reshape((-1, mzloc, mxloc))
else:
f_loc = dtype(infile.read_record(dtype=read_precision))
f_loc = f_loc.reshape((-1, myloc, mxloc))
raw_etc = infile.read_record(dtype=read_precision)
if lpersist:
persist(self,
infile=infile,
precision=read_precision,
quiet=quiet)
infile.close()
t = raw_etc[0]
x_loc = raw_etc[1:mxloc + 1]
y_loc = raw_etc[mxloc + 1:mxloc + myloc + 1]
z_loc = raw_etc[mxloc + myloc + 1:mxloc + myloc + mzloc + 1]
if param.lshear:
shear_offset = 1
deltay = raw_etc[-1]
else:
shear_offset = 0
dx = raw_etc[-3 - shear_offset]
dy = raw_etc[-2 - shear_offset]
dz = raw_etc[-1 - shear_offset]
if len(proc_dirs) > 1:
# Calculate where the local processor will go in
# the global array.
#
# Don't overwrite ghost zones of processor to the
# left (and accordingly in y and z direction -- makes
# a difference on the diagonals)
#
# Recall that in NumPy, slicing is NON-INCLUSIVE on
# the right end, ie, x[0:4] will slice all of a
# 4-digit array, not produce an error like in idl.
if procdim.ipx == 0:
i0x = 0
i1x = i0x + procdim.mx
i0xloc = 0
i1xloc = procdim.mx
else:
i0x = procdim.ipx * procdim.nx + procdim.nghostx
i1x = i0x + procdim.mx - procdim.nghostx
i0xloc = procdim.nghostx
i1xloc = procdim.mx
if procdim.ipy == 0:
i0y = 0
i1y = i0y + procdim.my
i0yloc = 0
i1yloc = procdim.my
else:
i0y = procdim.ipy * procdim.ny + procdim.nghosty
i1y = i0y + procdim.my - procdim.nghosty
i0yloc = procdim.nghosty
i1yloc = procdim.my
if procdim.ipz == 0:
i0z = 0
i1z = i0z + procdim.mz
i0zloc = 0
i1zloc = procdim.mz
else:
i0z = procdim.ipz * procdim.nz + procdim.nghostz
i1z = i0z + procdim.mz - procdim.nghostz
i0zloc = procdim.nghostz
i1zloc = procdim.mz
x[i0x:i1x] = x_loc[i0xloc:i1xloc]
y[i0y:i1y] = y_loc[i0yloc:i1yloc]
z[i0z:i1z] = z_loc[i0zloc:i1zloc]
if not run2D:
self.f[:, i0z:i1z, i0y:i1y,
i0x:i1x] = f_loc[:, i0zloc:i1zloc,
i0yloc:i1yloc, i0xloc:i1xloc]
else:
if dim.ny == 1:
self.f[:, i0z:i1z,
i0x:i1x] = f_loc[:, i0zloc:i1zloc,
i0xloc:i1xloc]
else:
self.f[i0z:i1z, i0y:i1y,
i0x:i1x] = f_loc[i0zloc:i1zloc,
i0yloc:i1yloc,
i0xloc:i1xloc]
else:
self.f = f_loc
x = x_loc
y = y_loc
z = z_loc
if magic is not None:
if not np.all(param.lequidist):
raise NotImplementedError(
"Magic functions are only implemented for equidistant grids."
)
if "bb" in magic:
# Compute the magnetic field before doing trimall.
aa = self.f[index.ax - 1:index.az, ...]
self.bb = dtype(
curl(
aa,
dx,
dy,
dz,
x=x,
y=y,
run2D=run2D,
coordinate_system=param.coord_system,
))
if trimall:
self.bb = self.bb[:, dim.n1:dim.n2 + 1, dim.m1:dim.m2 + 1,
dim.l1:dim.l2 + 1]
if "jj" in magic:
# Compute the electric current field before doing trimall.
aa = self.f[index.ax - 1:index.az, ...]
self.jj = dtype(
curl2(aa,
dx,
dy,
dz,
x=x,
y=y,
coordinate_system=param.coord_system))
if trimall:
self.jj = self.jj[:, dim.n1:dim.n2 + 1, dim.m1:dim.m2 + 1,
dim.l1:dim.l2 + 1]
if "vort" in magic:
# Compute the vorticity field before doing trimall.
uu = self.f[index.ux - 1:index.uz, ...]
self.vort = dtype(
curl(
uu,
dx,
dy,
dz,
x=x,
y=y,
run2D=run2D,
coordinate_system=param.coord_system,
))
if trimall:
if run2D:
if dim.nz == 1:
self.vort = self.vort[:, dim.m1:dim.m2 + 1,
dim.l1:dim.l2 + 1]
else:
self.vort = self.vort[:, dim.n1:dim.n2 + 1,
dim.l1:dim.l2 + 1]
else:
self.vort = self.vort[:, dim.n1:dim.n2 + 1,
dim.m1:dim.m2 + 1,
dim.l1:dim.l2 + 1, ]
# Trim the ghost zones of the global f-array if asked.
if trimall:
self.x = x[dim.l1:dim.l2 + 1]
self.y = y[dim.m1:dim.m2 + 1]
self.z = z[dim.n1:dim.n2 + 1]
if not run2D:
self.f = self.f[:, dim.n1:dim.n2 + 1, dim.m1:dim.m2 + 1,
dim.l1:dim.l2 + 1]
else:
if dim.ny == 1:
self.f = self.f[:, dim.n1:dim.n2 + 1, dim.l1:dim.l2 + 1]
else:
self.f = self.f[:, dim.m1:dim.m2 + 1, dim.l1:dim.l2 + 1]
else:
self.x = x
self.y = y
self.z = z
self.l1 = dim.l1
self.l2 = dim.l2 + 1
self.m1 = dim.m1
self.m2 = dim.m2 + 1
self.n1 = dim.n1
self.n2 = dim.n2 + 1
# Assign an attribute to self for each variable defined in
# 'data/index.pro' so that e.g. self.ux is the x-velocity
aatest = []
uutest = []
for key in index.__dict__.keys():
if "aatest" in key:
aatest.append(key)
if "uutest" in key:
uutest.append(key)
if (key != "global_gg" and key != "keys" and "aatest" not in key
and "uutest" not in key):
value = index.__dict__[key]
setattr(self, key, self.f[value - 1, ...])
# Special treatment for vector quantities.
if hasattr(index, "uu"):
self.uu = self.f[index.ux - 1:index.uz, ...]
if hasattr(index, "aa"):
self.aa = self.f[index.ax - 1:index.az, ...]
if hasattr(index, "uu_sph"):
self.uu_sph = self.f[index.uu_sphx - 1:index.uu_sphz, ...]
if hasattr(index, "bb_sph"):
self.bb_sph = self.f[index.bb_sphx - 1:index.bb_sphz, ...]
# Special treatment for test method vector quantities.
# Note index 1,2,3,...,0 last vector may be the zero field/flow
if hasattr(index, "aatest1"):
naatest = int(len(aatest) / 3)
for j in range(0, naatest):
key = "aatest" + str(np.mod(j + 1, naatest))
value = index.__dict__["aatest1"] + 3 * j
setattr(self, key, self.f[value - 1:value + 2, ...])
if hasattr(index, "uutest1"):
nuutest = int(len(uutest) / 3)
for j in range(0, nuutest):
key = "uutest" + str(np.mod(j + 1, nuutest))
value = index.__dict__["uutest"] + 3 * j
setattr(self, key, self.f[value - 1:value + 2, ...])
self.t = t
self.dx = dx
self.dy = dy
self.dz = dz
if param.lshear:
self.deltay = deltay
# Do the rest of magic after the trimall (i.e. no additional curl.)
self.magic = magic
if self.magic is not None:
self.magic_attributes(param, dtype=dtype)
