def mag_pressure(src, dst, par, gd, l1, l2, m1, m2, n1, n2, nghost):
if 'bb' in dst.keys():
bb = dst['bb'][:, n1:n2, m1:m2, l1:l2]
else:
bb = bfield(src, dst, par, gd, l1, l2, m1, m2, n1, n2, nghost)
var = 0.5 * dot2(bb) / par.mu0
return var
def mag_pressure(src, dst, key, par, gd, l1, l2, m1, m2, n1, n2, nghost):
if key == "pb":
if "bb" in dst.keys():
bb = dst["bb"][:, n1:n2, m1:m2, l1:l2]
else:
bb = bfield(src, dst, "bb", par, gd, l1, l2, m1, m2, n1, n2,
nghost)
var = 0.5 * dot2(bb) / par.mu0
return var
def mean_mag_pressure(src, dst, key, par, gd, lindx, lindy, lindz, nghost,
sigma, typ, mode):
if key == "meaneb":
if "eb" in dst.keys():
eb = load_dataset(dst, 'eb', lindx, lindy, lindz, nghost)
elif "bb" in dst.keys():
var = load_dataset(dst, 'bb', lindx, lindy, lindz, nghost)
eb = 0.5 * dot2(var) / par.mu0
var = gauss_3Dsmooth(eb, sigma=sigma, typ=typ, mode=mode)
return var
def Mach_Av(src, dst, par, gd, l1, l2, m1, m2, n1, n2, nghost):
if 'bb' in dst.keys():
bb = dst['bb'][:, n1:n2, m1:m2, l1:l2]
else:
bb = bfield(src, dst, par, gd, l1, l2, m1, m2, n1, n2, nghost)
if 'rho' in src.keys():
rho = src['rho'][n1:n2, m1:m2, l1:l2]
elif 'lnrho' in src.keys():
rho = np.exp(src['lnrho'][n1:n2, m1:m2, l1:l2])
else:
print('no density used setting rho=1 in pressure calculation')
rho = 1
var = np.sqrt(dot2(bb) / (par.mu0 * rho))
return var
def Mach_Av(src, dst, key, par, gd, l1, l2, m1, m2, n1, n2, nghost):
if key == "Ma":
if "bb" in dst.keys():
bb = dst["bb"][:, n1:n2, m1:m2, l1:l2]
else:
bb = bfield(src, dst, "bb", par, gd, l1, l2, m1, m2, n1, n2,
nghost)
if "rho" in src.keys():
rho = src["rho"][n1:n2, m1:m2, l1:l2]
elif "lnrho" in src.keys():
rho = np.exp(src["lnrho"][n1:n2, m1:m2, l1:l2])
else:
print("no density used setting rho=1 in pressure calculation")
rho = 1
var = np.sqrt(dot2(bb) / (par.mu0 * rho))
return var
def Mach_cs(src, dst, par, gd, l1, l2, m1, m2, n1, n2, nghost):
uu = np.array([
src['ux'][n1:n2, m1:m2, l1:l2], src['uy'][n1:n2, m1:m2, l1:l2],
src['uz'][n1:n2, m1:m2, l1:l2]
])
if 'tt' in dst.keys():
tt = dst['tt'][n1:n2, m1:m2, l1:l2]
else:
tt = temperature(src, dst, par, gd, l1, l2, m1, m2, n1, n2, nghost)
if not par.gamma == 1:
cs2 = par.cp * (par.gamma - 1) * tt
else:
cs2 = par.cp * tt
print('tt min {} max {}'.format(tt.min(), tt.max()))
var = np.sqrt(dot2(uu) / cs2)
return var
def Mach_cs(src, dst, key, par, gd, l1, l2, m1, m2, n1, n2, nghost):
if key == "Ms":
uu = np.array([
src["ux"][n1:n2, m1:m2, l1:l2],
src["uy"][n1:n2, m1:m2, l1:l2],
src["uz"][n1:n2, m1:m2, l1:l2],
])
if "tt" in dst.keys():
tt = dst["tt"][n1:n2, m1:m2, l1:l2]
else:
tt = temperature(src, dst, "tt", par, gd, l1, l2, m1, m2, n1,
n2, nghost)
if not par.gamma == 1:
cs2 = par.cp * (par.gamma - 1) * tt
else:
cs2 = par.cp * tt
print("tt min {} max {}".format(tt.min(), tt.max()))
var = np.sqrt(dot2(uu) / cs2)
return var
def ohmic_heat(src, dst, key, par, gd, l1, l2, m1, m2, n1, n2, nghost):
if key == "ohmic":
if "jj" in dst["data"].keys():
jj = dst["data/jj"][:,n1:n2,m1:m2,l1:l2]
else:
calc_derived_data(src["data"], dst["data"], "jj", par,
gd, l1, l2, m1, m2, n1, n2, nghost)
jj = dst["data/jj"][:,n1:n2,m1:m2,l1:l2]
if "tt" in dst["data"].keys():
tt = dst["tt"][n1:n2,m1:m2,l1:l2]
else:
calc_derived_data(src["data"], dst["data"], "tt", par,
gd, l1, l2, m1, m2, n1, n2, nghost)
tt = dst["data/tt"][n1:n2,m1:m2,l1:l2]
if "rho" in dst["data"].keys():
rho = src["data/rho"][n1:n2,m1:m2,l1:l2]
elif "lnrho" in dst["data"].keys():
rho = np.exp(src["data/lnrho"][n1:n2,m1:m2,l1:l2])
else:
rho=1
var = par.eta*par.mu0*dot2(jj)/(rho*tt)
return var
def urand(src, dst, key, par, gd, l1, l2, m1, m2, n1, n2, nghost):
if key == "u2rand":
sigma = 0.02
if "sigma" in par.keys:
sigma = par.sigma
n1shift, n2shift, m1shift, m2shift, l1shift, l2shift = der_limits(
n1, n2, m1, m2, l1, l2, nghost)
uu = np.array([
src["ux"][n1shift:n2shift, m1shift:m2shift, l1shift:l2shift],
src["uy"][n1shift:n2shift, m1shift:m2shift, l1shift:l2shift],
src["uz"][n1shift:n2shift, m1shift:m2shift, l1shift:l2shift]
])
umean = uu.copy()
urand = uu.copy()
dx = min(gd.dx, gd.dy, gd.dz)
sigma_dx = sigma / dx
print("sigma_dx {:.2f} sigma {}".format(sigma_dx, sigma))
for j in range(3):
umean[j] = gf(uu[j], sigma_dx, mode="reflect")
urand[j] = uu[j] - umean[j]
var = dot2(urand)
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 derive_stats(
sim_path,
src,
dst,
stat_keys=["Rm", "uu", "Ms"],
par=[],
comm=None,
overwrite=False,
rank=0,
size=1,
nghost=3,
status="a",
chunksize=1000.0,
quiet=True,
nmin=32,
lmask=False,
mask_key="hot",
):
if comm:
overwrite = False
if isinstance(par, list):
os.chdir(sim_path)
par = read.param(quiet=True, conflicts_quiet=True)
# get data dimensions
nx, ny, nz = (
src["settings"]["nx"][0],
src["settings"]["ny"][0],
src["settings"]["nz"][0],
)
mx, my, mz = (
src["settings"]["mx"][0],
src["settings"]["my"][0],
src["settings"]["mz"][0],
)
# split data into manageable memory chunks
dstchunksize = 8 * nx * ny * nz / 1024 * 1024
if dstchunksize > chunksize:
nchunks = cpu_optimal(
nx,
ny,
nz,
quiet=quiet,
mvar=src["settings/mvar"][0],
maux=src["settings/maux"][0],
MBmin=chunksize,
nmin=nmin,
size=size,
)[1]
else:
nchunks = [1, 1, 1]
print("nchunks {}".format(nchunks))
# for mpi split chunks across processes
if size > 1:
locindx = np.array_split(np.arange(nx) + nghost, nchunks[0])
locindy = np.array_split(np.arange(ny) + nghost, nchunks[1])
locindz = np.array_split(np.arange(nz) + nghost, nchunks[2])
indx = [
locindx[np.mod(
rank + int(rank / nchunks[2]) + int(rank / nchunks[1]),
nchunks[0])]
]
indy = [locindy[np.mod(rank + int(rank / nchunks[2]), nchunks[1])]]
indz = [locindz[np.mod(rank, nchunks[2])]]
allchunks = 1
else:
locindx = np.array_split(np.arange(nx) + nghost, nchunks[0])
locindy = np.array_split(np.arange(ny) + nghost, nchunks[1])
locindz = np.array_split(np.arange(nz) + nghost, nchunks[2])
indx = np.array_split(np.arange(nx) + nghost, nchunks[0])
indy = np.array_split(np.arange(ny) + nghost, nchunks[1])
indz = np.array_split(np.arange(nz) + nghost, nchunks[2])
allchunks = nchunks[0] * nchunks[1] * nchunks[2]
# ensure derived variables are in a list
if isinstance(stat_keys, list):
stat_keys = stat_keys
else:
stat_keys = [stat_keys]
# initialise group
group = group_h5(dst,
"stats",
status="a",
overwrite=overwrite,
comm=comm,
rank=rank,
size=size)
for key in stat_keys:
mean_stat = list()
stdv_stat = list()
mean_mask = list()
stdv_mask = list()
nmask_msk = list()
mean_nmsk = list()
stdv_nmsk = list()
nmask_nmk = list()
for ichunk in range(allchunks):
for iz in [indz[np.mod(ichunk, nchunks[2])]]:
n1, n2 = iz[0], iz[-1] + 1
for iy in [
indy[np.mod(ichunk + int(ichunk / nchunks[2]),
nchunks[1])]
]:
m1, m2 = iy[0], iy[-1] + 1
for ix in [
indx[np.mod(
ichunk + int(ichunk / nchunks[2]) +
int(ichunk / nchunks[1]),
nchunks[0],
)]
]:
l1, l2 = ix[0], ix[-1] + 1
if key in src["data"].keys():
var = src["data"][key][n1:n2, m1:m2, l1:l2]
elif key == "uu" or key == "aa":
tmp = np.array([
src["data"][key[0] + "x"][n1:n2, m1:m2, l1:l2],
src["data"][key[0] + "y"][n1:n2, m1:m2, l1:l2],
src["data"][key[0] + "z"][n1:n2, m1:m2, l1:l2],
])
var = np.sqrt(dot2(tmp))
else:
if key in dst["data"].keys():
if is_vector(key):
var = np.sqrt(
dot2(dst["data"][key][:, n1:n2, m1:m2,
l1:l2]))
else:
var = dst["data"][key][n1:n2, m1:m2, l1:l2]
else:
print(
"stats: " + key + " does not exist in ",
src,
"or",
dst,
)
continue
if lmask:
mask = dst["masks"][mask_key][0, n1:n2, m1:m2,
l1:l2]
Nmask = mask[mask == False].size
if Nmask > 0:
mean_mask.append(var[mask == False].mean() *
Nmask)
stdv_mask.append(var[mask == False].std() *
Nmask)
else:
mean_mask.append(0)
stdv_mask.append(0)
nmask_msk.append(Nmask)
nmask = mask[mask == True].size
if nmask > 0:
mean_nmsk.append(var[mask == True].mean() *
nmask)
stdv_nmsk.append(var[mask == True].std() *
nmask)
else:
mean_nmsk.append(0)
stdv_nmsk.append(0)
nmask_nmk.append(nmask)
mean_stat.append(var.mean())
stdv_stat.append(var.std())
if comm:
if lmask:
mean_mask = comm.gather(mean_mask, root=0)
stdv_mask = comm.gather(stdv_mask, root=0)
mean_mask = comm.bcast(mean_mask, root=0)
stdv_mask = comm.bcast(stdv_mask, root=0)
mean_nmsk = comm.gather(mean_nmsk, root=0)
stdv_nmsk = comm.gather(stdv_nmsk, root=0)
mean_nmsk = comm.bcast(mean_nmsk, root=0)
stdv_nmsk = comm.bcast(stdv_nmsk, root=0)
nmask_msk = comm.gather(nmask_msk, root=0)
nmask_nmk = comm.gather(nmask_nmk, root=0)
nmask_msk = comm.bcast(nmask_msk, root=0)
nmask_nmk = comm.bcast(nmask_nmk, root=0)
mean_stat = comm.gather(mean_stat, root=0)
stdv_stat = comm.gather(stdv_stat, root=0)
mean_stat = comm.bcast(mean_stat, root=0)
stdv_stat = comm.bcast(stdv_stat, root=0)
if lmask:
summk = np.sum(nmask_msk)
if summk > 0:
meanm = np.sum(mean_mask) / summk
stdvm = np.sum(stdv_mask) / summk
else:
meanm = 0
stdvm = 0
sumnk = np.sum(nmask_nmk)
if sumnk > 0:
meann = np.sum(mean_nmsk) / sumnk
stdvn = np.sum(stdv_nmsk) / sumnk
else:
meann = 0
stdvn = 0
print(mask_key + "-" + key + "-mean = {}, ".format(meanm) +
mask_key + "-" + key + "-std = {}".format(stdvm))
print("not-" + mask_key + "-" + key +
"-mean = {}, ".format(meann) + "not-" + mask_key + "-" +
key + "-std = {}".format(stdvn))
dataset_h5(
group,
mask_key + "-" + key + "-mean",
status=status,
data=meanm,
comm=comm,
size=size,
rank=rank,
overwrite=True,
)
dataset_h5(
group,
mask_key + "-" + key + "-std",
status=status,
data=stdvm,
comm=comm,
size=size,
rank=rank,
overwrite=True,
)
dataset_h5(
group,
"not-" + mask_key + "-" + key + "-mean",
status=status,
data=meann,
comm=comm,
size=size,
rank=rank,
overwrite=True,
)
dataset_h5(
group,
"not-" + mask_key + "-" + key + "-std",
status=status,
data=stdvn,
comm=comm,
size=size,
rank=rank,
overwrite=True,
)
mstat = np.mean(mean_stat)
dstat = np.mean(stdv_stat)
print(key + "-mean = {}, ".format(mstat) + key +
"-std = {}".format(dstat))
dataset_h5(
group,
key + "-mean",
status=status,
data=mstat,
comm=comm,
size=size,
rank=rank,
overwrite=True,
)
dataset_h5(
group,
key + "-std",
status=status,
data=dstat,
comm=comm,
size=size,
rank=rank,
overwrite=True,
)
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 fluid_reynolds(uu,
param,
grid,
lnrho=list(),
shock=list(),
nghost=3,
lmix=True,
quiet=True):
"""
Computes the fluid Reynolds number from the advective and effective
viscous expressions in the momentum equation.
call signature:
fluid_reynolds(uu, ivisc, grid, rho=None, shock=None, nghost=3)
Keyword arguments:
*uu*:
The velocity field [3,mz,my,mx] from the simulation data
*param*:
The Param simulation object with viscosity data information
*grid*:
The Grid simulation object
*lnrho*:
The log density field if it is non-uniform
*shock*:
The shock variable if shock viscosity is applied
*nghost*:
The number of ghost zones appropriate to the order of accuracy
*lmix*:
Option not to include hyper values when Laplacian values present
"""
#viscous forces
th2 = 2. / 3
th1 = 1. / 3
fvisc = np.zeros_like(uu)
#molecular viscosity contribution
ldel2, lshock, lhyper3 = False, False, False
for ivisc in param.ivisc:
if not 'shock' in ivisc and not 'hyper' in ivisc\
and not '\n' in ivisc:
ldel2 = True
if 'shock' in ivisc:
lshock = True
if 'hyper3' in ivisc:
lhyper3 = True
if ldel2:
if lhyper3:
lhyper3 = lhyper3 == lmix
del2u = np.zeros_like(uu)
for j in range(0, 3):
del2u[j] = del2(uu[j],
grid.dx,
grid.dy,
grid.dz,
x=grid.x,
y=grid.y,
coordinate_system=param.coord_system)
del2u[j, :nghost, nghost:-nghost,
nghost:-nghost] = del2u[j, -2 * nghost:-nghost,
nghost:-nghost, nghost:-nghost]
del2u[j, -nghost:, nghost:-nghost,
nghost:-nghost] = del2u[j, nghost:2 * nghost, nghost:-nghost,
nghost:-nghost]
del2u[j, nghost:-nghost, :nghost,
nghost:-nghost] = del2u[j, nghost:-nghost,
-2 * nghost:-nghost, nghost:-nghost]
del2u[j, nghost:-nghost, -nghost:,
nghost:-nghost] = del2u[j, nghost:-nghost, nghost:2 * nghost,
nghost:-nghost]
del2u[j, nghost:-nghost,
nghost:-nghost, :nghost] = del2u[j, nghost:-nghost,
nghost:-nghost,
-2 * nghost:-nghost]
del2u[j, nghost:-nghost, nghost:-nghost,
-nghost:] = del2u[j, nghost:-nghost, nghost:-nghost,
nghost:2 * nghost]
for ivisc in param.ivisc:
ivisc = str.strip(ivisc, '\n')
if 'nu-const' not in ivisc and 'shock' not in ivisc\
and 'hyper' not in ivisc and len(ivisc) > 0:
print(
'fluid_reynolds WARNING: ' + ivisc + ' not implemented\n' +
'terms may be missing from the standard rate of strain tensor'
)
fvisc = fvisc + param.nu * del2u
del (del2u)
tmp0 = grad(uu[0],
grid.dx,
grid.dy,
grid.dz,
x=grid.x,
y=grid.y,
coordinate_system=param.coord_system)
for j in range(0, 3):
tmp0[j, :nghost, nghost:-nghost,
nghost:-nghost] = tmp0[j, -2 * nghost:-nghost, nghost:-nghost,
nghost:-nghost]
tmp0[j, -nghost:, nghost:-nghost,
nghost:-nghost] = tmp0[j, nghost:2 * nghost, nghost:-nghost,
nghost:-nghost]
tmp0[j, nghost:-nghost, :nghost,
nghost:-nghost] = tmp0[j, nghost:-nghost, -2 * nghost:-nghost,
nghost:-nghost]
tmp0[j, nghost:-nghost, -nghost:,
nghost:-nghost] = tmp0[j, nghost:-nghost, nghost:2 * nghost,
nghost:-nghost]
tmp0[j, nghost:-nghost,
nghost:-nghost, :nghost] = tmp0[j, nghost:-nghost, nghost:-nghost,
-2 * nghost:-nghost]
tmp0[j, nghost:-nghost, nghost:-nghost,
-nghost:] = tmp0[j, nghost:-nghost, nghost:-nghost,
nghost:2 * nghost]
tmp1 = grad(uu[1],
grid.dx,
grid.dy,
grid.dz,
x=grid.x,
y=grid.y,
coordinate_system=param.coord_system)
for j in range(0, 3):
tmp1[j, :nghost, nghost:-nghost,
nghost:-nghost] = tmp1[j, -2 * nghost:-nghost, nghost:-nghost,
nghost:-nghost]
tmp1[j, -nghost:, nghost:-nghost,
nghost:-nghost] = tmp1[j, nghost:2 * nghost, nghost:-nghost,
nghost:-nghost]
tmp1[j, nghost:-nghost, :nghost,
nghost:-nghost] = tmp1[j, nghost:-nghost, -2 * nghost:-nghost,
nghost:-nghost]
tmp1[j, nghost:-nghost, -nghost:,
nghost:-nghost] = tmp1[j, nghost:-nghost, nghost:2 * nghost,
nghost:-nghost]
tmp1[j, nghost:-nghost,
nghost:-nghost, :nghost] = tmp1[j, nghost:-nghost, nghost:-nghost,
-2 * nghost:-nghost]
tmp1[j, nghost:-nghost, nghost:-nghost,
-nghost:] = tmp1[j, nghost:-nghost, nghost:-nghost,
nghost:2 * nghost]
tmp2 = grad(uu[2],
grid.dx,
grid.dy,
grid.dz,
x=grid.x,
y=grid.y,
coordinate_system=param.coord_system)
for j in range(0, 3):
tmp2[j, :nghost, nghost:-nghost,
nghost:-nghost] = tmp2[j, -2 * nghost:-nghost, nghost:-nghost,
nghost:-nghost]
tmp2[j, -nghost:, nghost:-nghost,
nghost:-nghost] = tmp2[j, nghost:2 * nghost, nghost:-nghost,
nghost:-nghost]
tmp2[j, nghost:-nghost, :nghost,
nghost:-nghost] = tmp2[j, nghost:-nghost, -2 * nghost:-nghost,
nghost:-nghost]
tmp2[j, nghost:-nghost, -nghost:,
nghost:-nghost] = tmp2[j, nghost:-nghost, nghost:2 * nghost,
nghost:-nghost]
tmp2[j, nghost:-nghost,
nghost:-nghost, :nghost] = tmp2[j, nghost:-nghost, nghost:-nghost,
-2 * nghost:-nghost]
tmp2[j, nghost:-nghost, nghost:-nghost,
-nghost:] = tmp2[j, nghost:-nghost, nghost:-nghost,
nghost:2 * nghost]
#effect of compressibility
if len(lnrho) > 0:
divu = div(uu,
grid.dx,
grid.dy,
grid.dz,
x=grid.x,
y=grid.y,
coordinate_system=param.coord_system)
divu[:nghost, nghost:-nghost,
nghost:-nghost] = divu[-2 * nghost:-nghost, nghost:-nghost,
nghost:-nghost]
divu[-nghost:, nghost:-nghost,
nghost:-nghost] = divu[nghost:2 * nghost, nghost:-nghost,
nghost:-nghost]
divu[nghost:-nghost, :nghost,
nghost:-nghost] = divu[nghost:-nghost, -2 * nghost:-nghost,
nghost:-nghost]
divu[nghost:-nghost, -nghost:,
nghost:-nghost] = divu[nghost:-nghost, nghost:2 * nghost,
nghost:-nghost]
divu[nghost:-nghost,
nghost:-nghost, :nghost] = divu[nghost:-nghost, nghost:-nghost,
-2 * nghost:-nghost]
divu[nghost:-nghost, nghost:-nghost,
-nghost:] = divu[nghost:-nghost, nghost:-nghost,
nghost:2 * nghost]
gradlnrho = grad(lnrho,
grid.dx,
grid.dy,
grid.dz,
x=grid.x,
y=grid.y,
coordinate_system=param.coord_system)
for j in range(0, 3):
gradlnrho[j, :nghost, nghost:-nghost,
nghost:-nghost] = gradlnrho[j, -2 * nghost:-nghost,
nghost:-nghost,
nghost:-nghost]
gradlnrho[j, -nghost:, nghost:-nghost,
nghost:-nghost] = gradlnrho[j, nghost:2 * nghost,
nghost:-nghost,
nghost:-nghost]
gradlnrho[j, nghost:-nghost, :nghost,
nghost:-nghost] = gradlnrho[j, nghost:-nghost,
-2 * nghost:-nghost,
nghost:-nghost]
gradlnrho[j, nghost:-nghost, -nghost:,
nghost:-nghost] = gradlnrho[j, nghost:-nghost,
nghost:2 * nghost,
nghost:-nghost]
gradlnrho[j, nghost:-nghost,
nghost:-nghost, :nghost] = gradlnrho[j, nghost:-nghost,
nghost:-nghost,
-2 * nghost:-nghost]
gradlnrho[j, nghost:-nghost, nghost:-nghost,
-nghost:] = gradlnrho[j, nghost:-nghost, nghost:-nghost,
nghost:2 * nghost]
Sglnrho = np.zeros_like(uu)
Sglnrho[0] = dot(tmp0,gradlnrho) +\
(tmp0[0]+tmp1[0]+tmp2[0]-th2*divu)*gradlnrho[0]
Sglnrho[1] = dot(tmp1,gradlnrho) +\
(tmp0[1]+tmp1[1]+tmp2[1]-th2*divu)*gradlnrho[1]
Sglnrho[2] = dot(tmp2,gradlnrho) +\
(tmp0[2]+tmp1[2]+tmp2[2]-th2*divu)*gradlnrho[2]
graddivu = grad(divu,
grid.dx,
grid.dy,
grid.dz,
x=grid.x,
y=grid.y,
coordinate_system=param.coord_system)
for j in range(0, 3):
graddivu[j, :nghost, nghost:-nghost,
nghost:-nghost] = graddivu[j, -2 * nghost:-nghost,
nghost:-nghost, nghost:-nghost]
graddivu[j, -nghost:, nghost:-nghost,
nghost:-nghost] = graddivu[j, nghost:2 * nghost,
nghost:-nghost, nghost:-nghost]
graddivu[j, nghost:-nghost, :nghost,
nghost:-nghost] = graddivu[j, nghost:-nghost,
-2 * nghost:-nghost,
nghost:-nghost]
graddivu[j, nghost:-nghost, -nghost:,
nghost:-nghost] = graddivu[j, nghost:-nghost,
nghost:2 * nghost,
nghost:-nghost]
graddivu[j, nghost:-nghost,
nghost:-nghost, :nghost] = graddivu[j, nghost:-nghost,
nghost:-nghost,
-2 * nghost:-nghost]
graddivu[j, nghost:-nghost, nghost:-nghost,
-nghost:] = graddivu[j, nghost:-nghost, nghost:-nghost,
nghost:2 * nghost]
fvisc = fvisc + param.nu * (th1 * graddivu + Sglnrho)
del (Sglnrho)
elif param.ldensity:
print('fluid_reynolds WARNING: no lnrho provided\n' +
'rate of strain tensor likely incomplete')
#shock contribution
if lshock:
if len(shock) == 0:
print('fluid_reynolds WARNING: no shock provided\n' +
'rate of strain tensor likely incomplete')
else:
shock[:nghost, nghost:-nghost,
nghost:-nghost] = shock[-2 * nghost:-nghost, nghost:-nghost,
nghost:-nghost]
shock[-nghost:, nghost:-nghost,
nghost:-nghost] = shock[nghost:2 * nghost, nghost:-nghost,
nghost:-nghost]
shock[nghost:-nghost, :nghost,
nghost:-nghost] = shock[nghost:-nghost, -2 * nghost:-nghost,
nghost:-nghost]
shock[nghost:-nghost, -nghost:,
nghost:-nghost] = shock[nghost:-nghost, nghost:2 * nghost,
nghost:-nghost]
shock[nghost:-nghost,
nghost:-nghost, :nghost] = shock[nghost:-nghost,
nghost:-nghost,
-2 * nghost:-nghost]
shock[nghost:-nghost, nghost:-nghost,
-nghost:] = shock[nghost:-nghost, nghost:-nghost,
nghost:2 * nghost]
divugradlnrho = np.zeros_like(uu)
gradshock = grad(shock,
grid.dx,
grid.dy,
grid.dz,
x=grid.x,
y=grid.y,
coordinate_system=param.coord_system)
for j in range(0, 3):
gradshock[j, :nghost, nghost:-nghost,
nghost:-nghost] = gradshock[j, -2 * nghost:-nghost,
nghost:-nghost,
nghost:-nghost]
gradshock[j, -nghost:, nghost:-nghost,
nghost:-nghost] = gradshock[j, nghost:2 * nghost,
nghost:-nghost,
nghost:-nghost]
gradshock[j, nghost:-nghost, :nghost,
nghost:-nghost] = gradshock[j, nghost:-nghost,
-2 * nghost:-nghost,
nghost:-nghost]
gradshock[j, nghost:-nghost, -nghost:,
nghost:-nghost] = gradshock[j, nghost:-nghost,
nghost:2 * nghost,
nghost:-nghost]
gradshock[j, nghost:-nghost,
nghost:-nghost, :nghost] = gradshock[j,
nghost:-nghost,
nghost:-nghost,
-2 *
nghost:-nghost]
gradshock[j, nghost:-nghost, nghost:-nghost,
-nghost:] = gradshock[j, nghost:-nghost,
nghost:-nghost,
nghost:2 * nghost]
for j in range(0, 3):
divugradlnrho[j] = param.nu_shock*divu*gradshock[j] +\
param.nu_shock*shock*(divu*gradlnrho[j] + graddivu[j])
del (divu, gradshock, gradlnrho, graddivu)
fvisc = fvisc + divugradlnrho
del (divugradlnrho)
if lhyper3:
#deluij5 = np.zeros_like([uu,uu,uu])
#uij5glnrho to be included
del6u = np.zeros_like(uu)
for j in range(0, 3):
del6u[j] = del6(uu[j], grid.dx, grid.dy, grid.dz)
del6u[j, :nghost, nghost:-nghost,
nghost:-nghost] = del6u[j, -2 * nghost:-nghost,
nghost:-nghost, nghost:-nghost]
del6u[j, -nghost:, nghost:-nghost,
nghost:-nghost] = del6u[j, nghost:2 * nghost, nghost:-nghost,
nghost:-nghost]
del6u[j, nghost:-nghost, :nghost,
nghost:-nghost] = del6u[j, nghost:-nghost,
-2 * nghost:-nghost, nghost:-nghost]
del6u[j, nghost:-nghost, -nghost:,
nghost:-nghost] = del6u[j, nghost:-nghost, nghost:2 * nghost,
nghost:-nghost]
del6u[j, nghost:-nghost,
nghost:-nghost, :nghost] = del6u[j, nghost:-nghost,
nghost:-nghost,
-2 * nghost:-nghost]
del6u[j, nghost:-nghost, nghost:-nghost,
-nghost:] = del6u[j, nghost:-nghost, nghost:-nghost,
nghost:2 * nghost]
#del6 for non-cartesian tba
#del6u[j] = del6(uu[j],grid.dx,grid.dy,grid.dz,x=grid.x,y=grid.y,
# coordinate_system=param.coord_system)
fvisc = fvisc + param.nu_hyper3 * del6u
del (del6u)
fvisc2 = np.sqrt(dot2(fvisc))
#advective forces
advec = np.zeros_like(uu)
advec[0] = dot(uu, tmp0)
advec[1] = dot(uu, tmp1)
advec[0] = dot(uu, tmp2)
del (tmp0, tmp1, tmp2)
advec2 = np.sqrt(dot2(advec))
del (advec)
#avoid division by zero
if fvisc2.max() > 0:
fvisc2[np.where(fvisc2 == 0)] = fvisc2[np.where(fvisc2 > 0)].min()
Re = advec2 / fvisc2
#set minimum floor to exclude zero-valued Re
Re[np.where(Re == 0)] = Re[np.where(Re > 0)].min()
else:
Re = advec2
print('Re undefined')
return Re
def helmholtz_fft(
tot_field,
grid,
params,
nghost=3,
pot=True,
rot=True,
lno_mean=False,
nonperi_bc=None,
field_scalar=[],
s=None,
quiet=True,
):
"""
helmholz_fft(field, grid, params)
Creates the decomposition vector pair for the supplied vector field.
Parameters
----------
tot_field : ndarray
Vector field of dimension [3, mz, my, mx], which is decomposed.
grid: obj
Grid object with grid spacing dx, dy, dz.
params : obj
Simulation Params object with domain dimensions Lx, Ly and Lz.
nghost : int
Number of ghost zones to exclude from the fft.
lno_mean : float
Exclude any mean flow from the decomposition - should drop anyway.
nonperi_bc : string
String if not None with boundary condition label.
How to apply the bc needs to be implemented as required.
field_scalar : ndarray
Scalar field (density) as debug tool for energy comparison.
s : list of int
List of three integers if not None for fft dimension.
If none the dimension of the field [nz,ny,nx] is used.
Returns
-------
ndarray with decomposed field.
"""
if lno_mean:
# Exclude volume mean flows.
field = np.zeros_like(tot_field)
for j in range(0, 3):
field[j] = tot_field[j] - tot_field[j].mean()
else:
field = tot_field
# Use mean speed and grid spacing in normalization of div/curl check.
amp_field_1 = 1.0 / np.sqrt(dot2(field)).mean()
nz, ny, nx = (
field[:, nghost:-nghost, nghost:-nghost, nghost:-nghost].shape[-3],
field[:, nghost:-nghost, nghost:-nghost, nghost:-nghost].shape[-2],
field[:, nghost:-nghost, nghost:-nghost, nghost:-nghost].shape[-1],
)
invs = [nz, ny, nx]
if not s:
s = [nz, ny, nx]
knz, kny, knx = s[0], s[1], s[2]
nz2, ny2, nx2 = int(knz / 2), int(kny / 2), int(knx / 2)
# Derive wavenumbers k scaled to dimension of simulation domain.
kk = np.empty(shape=[3, s[0], s[1], s[2]])
k0 = np.arange(knx)
k0[nx2:] = -k0[nx2 - 1::-1] - 1
k1 = np.arange(kny)
k1[ny2:] = -k1[ny2 - 1::-1] - 1
k2 = np.arange(knz)
k2[nz2:] = -k2[nz2 - 1::-1] - 1
for j in range(0, k0.size):
kk[0, :, :, j] = k0[j] * 2 * np.pi / params.lxyz[0]
for j in range(0, k1.size):
kk[1, :, j, :] = k1[j] * 2 * np.pi / params.lxyz[1]
for j in range(0, k2.size):
kk[2, j, :, :] = k2[j] * 2 * np.pi / params.lxyz[2]
knorm = dot2(kk)
# Apply fast Fourier transform to the vector field.
kfield = np.empty(shape=[3, s[0], s[1], s[2]], dtype=complex)
for j in range(0, 3):
kfield[j] = np.fft.fftn(field[j, nghost:-nghost, nghost:-nghost,
nghost:-nghost],
s=s)
if pot:
# Reverse fft to obtain the scalar potential.
pfield = -1j * dot(kk, kfield)
pfield[np.where(knorm == 0)] = 0.0
if rot:
# Reverse fft to obtain the vector potential.
rfield = 1j * cross(kk, kfield)
for j in range(3):
rfield[j][np.where(knorm == 0)] = 0.0
# Avoid division by zero.
knorm[np.where(knorm == 0)] = 1.0
if pot:
pfield /= knorm
pot_field = np.zeros_like(field)
if rot:
for j in range(3):
rfield[j] /= knorm
rot_field = np.zeros_like(field)
if nonperi_bc:
print(
"Please implement new nonperi_bc not yet implemented.\n",
"Applying periodic boundary conditions for now.",
)
for j in range(0, 3):
if pot:
pot_field[j, nghost:-nghost, nghost:-nghost,
nghost:-nghost] = np.fft.ifftn(1j * pfield * kk[j],
s=invs).real
if rot:
rot_field[j, nghost:-nghost, nghost:-nghost,
nghost:-nghost] = np.fft.ifftn(cross(1j * kk, rfield)[j],
s=invs).real
# Apply the periodic boundary conditions for the ghost zones.
for j in range(0, 3):
if pot:
pot_field[j, :nghost, :, :] = pot_field[j,
-2 * nghost:-nghost, :, :]
pot_field[j, -nghost:, :, :] = pot_field[j,
nghost:2 * nghost, :, :]
pot_field[j, :, :nghost, :] = pot_field[j, :,
-2 * nghost:-nghost, :]
pot_field[j, :, -nghost:, :] = pot_field[j, :,
nghost:2 * nghost, :]
pot_field[j, :, :, :nghost] = pot_field[j, :, :,
-2 * nghost:-nghost]
pot_field[j, :, :, -nghost:] = pot_field[j, :, :,
nghost:2 * nghost]
if rot:
rot_field[j, :nghost, :, :] = rot_field[j,
-2 * nghost:-nghost, :, :]
rot_field[j, -nghost:, :, :] = rot_field[j,
nghost:2 * nghost, :, :]
rot_field[j, :, :nghost, :] = rot_field[j, :,
-2 * nghost:-nghost, :]
rot_field[j, :, -nghost:, :] = rot_field[j, :,
nghost:2 * nghost, :]
rot_field[j, :, :, :nghost] = rot_field[j, :, :,
-2 * nghost:-nghost]
rot_field[j, :, :, -nghost:] = rot_field[j, :, :,
nghost:2 * nghost]
# Compare internal energy of original and sum of decomposed vectors.
if pot:
pot2 = dot2(pot_field)[nghost:-nghost, nghost:-nghost, nghost:-nghost]
if rot:
rot2 = dot2(rot_field)[nghost:-nghost, nghost:-nghost, nghost:-nghost]
field2 = dot2(field)[nghost:-nghost, nghost:-nghost, nghost:-nghost]
if len(field_scalar) > 0:
# Compare kinetic energy of original and sum of decomposed vectors.
field2 *= field_scalar[nghost:-nghost, nghost:-nghost, nghost:-nghost]
if pot:
pot2 *= field_scalar[nghost:-nghost, nghost:-nghost,
nghost:-nghost]
if rot:
rot2 *= field_scalar[nghost:-nghost, nghost:-nghost,
nghost:-nghost]
if rot and not pot:
if not quiet:
print(
"mean total field energy {} mean rotational energy {}".format(
np.mean(field2), np.mean(rot2)))
elif pot and not rot:
if not quiet:
print("mean total field energy {} mean irrotational energy {}".
format(np.mean(field2), np.mean(pot2)))
elif rot and pot:
if not quiet:
print("mean total field energy {} mean summed component energy {}".
format(np.mean(field2), np.mean(rot2 + pot2)))
# Check div and curl approximate/equal zero.
if pot:
if not quiet:
print("Max {} and mean {} of abs(curl(pot field))".format(
max(grid.dx, grid.dy, grid.dz) * amp_field_1 *
np.sqrt(dot2(curl(pot_field, grid.dx, grid.dy,
grid.dz)))[nghost:-nghost, nghost:-nghost,
nghost:-nghost].max(),
max(grid.dx, grid.dy, grid.dz) * amp_field_1 *
np.sqrt(dot2(curl(pot_field, grid.dx, grid.dy,
grid.dz)))[nghost:-nghost, nghost:-nghost,
nghost:-nghost].mean(),
))
if rot:
if not quiet:
print("Max {} and mean {} of abs(div(rot field))".format(
max(grid.dx, grid.dy, grid.dz) * amp_field_1 *
np.abs(div(rot_field, grid.dx, grid.dy,
grid.dz))[nghost:-nghost, nghost:-nghost,
nghost:-nghost].max(),
max(grid.dx, grid.dy, grid.dz) * amp_field_1 *
np.abs(div(rot_field, grid.dx, grid.dy,
grid.dz))[nghost:-nghost, nghost:-nghost,
nghost:-nghost].mean(),
))
if rot and not pot:
ret_opt = rot_field
elif pot and not rot:
ret_opt = pot_field
elif rot and pot:
ret_opt = [rot_field, pot_field]
else:
print("pot and/or rot must be True, returning ones")
ret_opt = np.ones_like(tot_field)
return ret_opt
def var2vtk(var_file='var.dat',
datadir='data',
proc=-1,
variables=None,
b_ext=False,
magic=[],
destination='work',
quiet=True,
trimall=True,
ti=-1,
tf=-1):
"""
Convert data from PencilCode format to vtk.
call signature::
var2vtk(var_file='', datadir='data', proc=-1,
variables='', b_ext=False,
destination='work', quiet=True, trimall=True, ti=-1, tf=-1)
Read *var_file* and convert its content into vtk format. Write the result
in *destination*.
Keyword arguments:
*var_file*:
The original var_file.
*datadir*:
Directory where the data is stored.
*proc*:
Processor which should be read. Set to -1 for all processors.
*variables*:
List of variables which should be written. If None all.
*b_ext*:
Add the external magnetic field.
*destination*:
Destination file.
*quiet*:
Keep quiet when reading the var files.
*trimall*:
Trim the data cube to exclude ghost zones.
*ti, tf*:
Start and end index for animation. Leave negative for no animation.
Overwrites variable var_file.
"""
import numpy as np
import sys
from pencil import read
from pencil import math
# Determine of we want an animation.
if ti < 0 or tf < 0:
animation = False
else:
animation = True
# If no variables specified collect all by default
if not variables:
variables = []
indx = read.index()
for key in indx.__dict__.keys():
if 'keys' not in key:
variables.append(key)
if 'uu' in variables:
magic.append('vort')
variables.append('vort')
if 'rho' in variables or 'lnrho' in variables:
if 'ss' in variables:
magic.append('tt')
variables.append('tt')
magic.append('pp')
variables.append('pp')
if 'aa' in variables:
magic.append('bb')
variables.append('bb')
magic.append('jj')
variables.append('jj')
variables.append('ab')
variables.append('b_mag')
variables.append('j_mag')
else:
# Convert single variable string into length 1 list of arrays.
if (len(variables) > 0):
if (len(variables[0]) == 1):
variables = [variables]
if 'tt' in variables:
magic.append('tt')
if 'pp' in variables:
magic.append('pp')
if 'bb' in variables:
magic.append('bb')
if 'jj' in variables:
magic.append('jj')
if 'vort' in variables:
magic.append('vort')
if 'b_mag' in variables and not 'bb' in magic:
magic.append('bb')
if 'j_mag' in variables and not 'jj' in magic:
magic.append('jj')
if 'ab' in variables and not 'bb' in magic:
magic.append('bb')
for t_idx in range(ti, tf + 1):
if animation:
var_file = 'VAR' + str(t_idx)
# Read the PencilCode variables and set the dimensions.
var = read.var(var_file=var_file,
datadir=datadir,
proc=proc,
magic=magic,
trimall=True,
quiet=quiet)
grid = read.grid(datadir=datadir, proc=proc, trim=trimall, quiet=True)
params = read.param(quiet=True)
# Add external magnetic field.
if (b_ext == True):
B_ext = np.array(params.b_ext)
var.bb[0, ...] += B_ext[0]
var.bb[1, ...] += B_ext[1]
var.bb[2, ...] += B_ext[2]
dimx = len(grid.x)
dimy = len(grid.y)
dimz = len(grid.z)
dim = dimx * dimy * dimz
dx = (np.max(grid.x) - np.min(grid.x)) / (dimx - 1)
dy = (np.max(grid.y) - np.min(grid.y)) / (dimy - 1)
dz = (np.max(grid.z) - np.min(grid.z)) / (dimz - 1)
# Write the vtk header.
if animation:
fd = open(destination + str(t_idx) + '.vtk', 'wb')
else:
fd = open(destination + '.vtk', 'wb')
fd.write('# vtk DataFile Version 2.0\n'.encode('utf-8'))
fd.write('VAR files\n'.encode('utf-8'))
fd.write('BINARY\n'.encode('utf-8'))
fd.write('DATASET STRUCTURED_POINTS\n'.encode('utf-8'))
fd.write('DIMENSIONS {0:9} {1:9} {2:9}\n'.format(dimx, dimy,
dimz).encode('utf-8'))
fd.write('ORIGIN {0:8.12} {1:8.12} {2:8.12}\n'.format(
grid.x[0], grid.y[0], grid.z[0]).encode('utf-8'))
fd.write('SPACING {0:8.12} {1:8.12} {2:8.12}\n'.format(
dx, dy, dz).encode('utf-8'))
fd.write('POINT_DATA {0:9}\n'.format(dim).encode('utf-8'))
# Write the data.
for v in variables:
print('Writing {0}.'.format(v))
# Prepare the data to the correct format.
if v == 'ab':
data = math.dot(var.aa, var.bb)
elif v == 'b_mag':
data = np.sqrt(math.dot2(var.bb))
elif v == 'j_mag':
data = np.sqrt(math.dot2(var.jj))
else:
data = getattr(var, v)
if sys.byteorder == 'little':
data = data.astype(np.float32).byteswap()
else:
data = data.astype(np.float32)
# Check if we have vectors or scalars.
if data.ndim == 4:
data = np.moveaxis(data, 0, 3)
fd.write('VECTORS {0} float\n'.format(v).encode('utf-8'))
else:
fd.write('SCALARS {0} float\n'.format(v).encode('utf-8'))
fd.write('LOOKUP_TABLE default\n'.encode('utf-8'))
fd.write(data.tobytes())
del (var)
fd.close()
def derive_stats(sim_path,
src,
dst,
stat_keys=['Rm', 'uu', 'Ms'],
par=[],
comm=None,
overwrite=False,
rank=0,
size=1,
nghost=3,
status='a',
chunksize=1000.0,
quiet=True,
nmin=32,
lmask=False,
mask_key='hot'):
if comm:
overwrite = False
if isinstance(par, list):
os.chdir(sim_path)
par = read.param(quiet=True, conflicts_quiet=True)
#get data dimensions
nx, ny, nz = src['settings']['nx'][0],\
src['settings']['ny'][0],\
src['settings']['nz'][0]
mx, my, mz = src['settings']['mx'][0],\
src['settings']['my'][0],\
src['settings']['mz'][0]
#split data into manageable memory chunks
dstchunksize = 8 * nx * ny * nz / 1024 * 1024
if dstchunksize > chunksize:
nchunks = cpu_optimal(nx,
ny,
nz,
quiet=quiet,
mvar=src['settings/mvar'][0],
maux=src['settings/maux'][0],
MBmin=chunksize,
nmin=nmin,
size=size)[1]
else:
nchunks = [1, 1, 1]
print('nchunks {}'.format(nchunks))
# for mpi split chunks across processes
if size > 1:
locindx = np.array_split(np.arange(nx) + nghost, nchunks[0])
locindy = np.array_split(np.arange(ny) + nghost, nchunks[1])
locindz = np.array_split(np.arange(nz) + nghost, nchunks[2])
indx = [
locindx[np.mod(
rank + int(rank / nchunks[2]) + int(rank / nchunks[1]),
nchunks[0])]
]
indy = [locindy[np.mod(rank + int(rank / nchunks[2]), nchunks[1])]]
indz = [locindz[np.mod(rank, nchunks[2])]]
allchunks = 1
else:
locindx = np.array_split(np.arange(nx) + nghost, nchunks[0])
locindy = np.array_split(np.arange(ny) + nghost, nchunks[1])
locindz = np.array_split(np.arange(nz) + nghost, nchunks[2])
indx = np.array_split(np.arange(nx) + nghost, nchunks[0])
indy = np.array_split(np.arange(ny) + nghost, nchunks[1])
indz = np.array_split(np.arange(nz) + nghost, nchunks[2])
allchunks = nchunks[0] * nchunks[1] * nchunks[2]
# ensure derived variables are in a list
if isinstance(stat_keys, list):
stat_keys = stat_keys
else:
stat_keys = [stat_keys]
# initialise group
group = group_h5(dst,
'stats',
status='a',
overwrite=overwrite,
comm=comm,
rank=rank,
size=size)
for key in stat_keys:
mean_stat = list()
stdv_stat = list()
mean_mask = list()
stdv_mask = list()
nmask_msk = list()
mean_nmsk = list()
stdv_nmsk = list()
nmask_nmk = list()
for ichunk in range(allchunks):
for iz in [indz[np.mod(ichunk, nchunks[2])]]:
n1, n2 = iz[ 0],\
iz[-1]+1
for iy in [
indy[np.mod(ichunk + int(ichunk / nchunks[2]),
nchunks[1])]
]:
m1, m2 = iy[ 0],\
iy[-1]+1
for ix in [
indx[np.mod(
ichunk + int(ichunk / nchunks[2]) +
int(ichunk / nchunks[1]), nchunks[0])]
]:
l1, l2 = ix[ 0],\
ix[-1]+1
if key in src['data'].keys():
var = src['data'][key][n1:n2, m1:m2, l1:l2]
elif key == 'uu' or key == 'aa':
tmp = np.array([
src['data'][key[0] + 'x'][n1:n2, m1:m2, l1:l2],
src['data'][key[0] + 'y'][n1:n2, m1:m2, l1:l2],
src['data'][key[0] + 'z'][n1:n2, m1:m2, l1:l2]
])
var = np.sqrt(dot2(tmp))
else:
if key in dst['data'].keys():
if is_vector(key):
var = np.sqrt(
dot2(dst['data'][key][:, n1:n2, m1:m2,
l1:l2]))
else:
var = dst['data'][key][n1:n2, m1:m2, l1:l2]
else:
print('stats: ' + key + ' does not exist in ',
src, 'or', dst)
continue
if lmask:
mask = dst['masks'][mask_key][0, n1:n2, m1:m2,
l1:l2]
Nmask = mask[mask == False].size
if Nmask > 0:
mean_mask.append(var[mask == False].mean() *
Nmask)
stdv_mask.append(var[mask == False].std() *
Nmask)
else:
mean_mask.append(0)
stdv_mask.append(0)
nmask_msk.append(Nmask)
nmask = mask[mask == True].size
if nmask > 0:
mean_nmsk.append(var[mask == True].mean() *
nmask)
stdv_nmsk.append(var[mask == True].std() *
nmask)
else:
mean_nmsk.append(0)
stdv_nmsk.append(0)
nmask_nmk.append(nmask)
mean_stat.append(var.mean())
stdv_stat.append(var.std())
if comm:
if lmask:
mean_mask = comm.gather(mean_mask, root=0)
stdv_mask = comm.gather(stdv_mask, root=0)
mean_mask = comm.bcast(mean_mask, root=0)
stdv_mask = comm.bcast(stdv_mask, root=0)
mean_nmsk = comm.gather(mean_nmsk, root=0)
stdv_nmsk = comm.gather(stdv_nmsk, root=0)
mean_nmsk = comm.bcast(mean_nmsk, root=0)
stdv_nmsk = comm.bcast(stdv_nmsk, root=0)
nmask_msk = comm.gather(nmask_msk, root=0)
nmask_nmk = comm.gather(nmask_nmk, root=0)
nmask_msk = comm.bcast(nmask_msk, root=0)
nmask_nmk = comm.bcast(nmask_nmk, root=0)
mean_stat = comm.gather(mean_stat, root=0)
stdv_stat = comm.gather(stdv_stat, root=0)
mean_stat = comm.bcast(mean_stat, root=0)
stdv_stat = comm.bcast(stdv_stat, root=0)
if lmask:
summk = np.sum(nmask_msk)
if summk > 0:
meanm = np.sum(mean_mask) / summk
stdvm = np.sum(stdv_mask) / summk
else:
meanm = 0
stdvm = 0
sumnk = np.sum(nmask_nmk)
if sumnk > 0:
meann = np.sum(mean_nmsk) / sumnk
stdvn = np.sum(stdv_nmsk) / sumnk
else:
meann = 0
stdvn = 0
print(mask_key + '-' + key + '-mean = {}, '.format(meanm) +
mask_key + '-' + key + '-std = {}'.format(stdvm))
print('not-' + mask_key + '-' + key +
'-mean = {}, '.format(meann) + 'not-' + mask_key + '-' +
key + '-std = {}'.format(stdvn))
dataset_h5(group,
mask_key + '-' + key + '-mean',
status=status,
data=meanm,
comm=comm,
size=size,
rank=rank,
overwrite=True)
dataset_h5(group,
mask_key + '-' + key + '-std',
status=status,
data=stdvm,
comm=comm,
size=size,
rank=rank,
overwrite=True)
dataset_h5(group,
'not-' + mask_key + '-' + key + '-mean',
status=status,
data=meann,
comm=comm,
size=size,
rank=rank,
overwrite=True)
dataset_h5(group,
'not-' + mask_key + '-' + key + '-std',
status=status,
data=stdvn,
comm=comm,
size=size,
rank=rank,
overwrite=True)
mstat = np.mean(mean_stat)
dstat = np.mean(stdv_stat)
print(key + '-mean = {}, '.format(mstat) + key +
'-std = {}'.format(dstat))
dataset_h5(group,
key + '-mean',
status=status,
data=mstat,
comm=comm,
size=size,
rank=rank,
overwrite=True)
dataset_h5(group,
key + '-std',
status=status,
data=dstat,
comm=comm,
size=size,
rank=rank,
overwrite=True)