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 rhs_data(sim_path, src, dst, magic=["uxb","etadel2a"], par=[], comm=None,
gd=[], grp_overwrite=False, overwrite=False,
rank=0, size=1, nghost=3,status="a",
chunksize = 1000.0, dtype=np.float64, quiet=True, nmin=32,
Reynolds_shock=False, lmix=False
):
if comm:
overwrite = False
if isinstance(par, list):
os.chdir(sim_path)
par = read.param(quiet=True,conflicts_quiet=True)
if isinstance(gd, list):
os.chdir(sim_path)
gd = read.grid(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]
# save time
dataset_h5(dst, "time", status=status, data=src["time"][()],
comm=comm, size=size, rank=rank,
overwrite=overwrite, dtype=dtype)
# ensure derived variables are in a list
if isinstance(magic, list):
magic = magic
else:
magic = [magic]
# confirm exists group
group_h5(dst, "data", status="a", overwrite=grp_overwrite,
comm=comm, rank=rank, size=size)
# initialise group
group = group_h5(dst, "calc", status="a", overwrite=grp_overwrite,
comm=comm, rank=rank, size=size)
for key in magic:
if is_vector(key):
dataset_h5(group, key, status=status, shape=[3,mz,my,mx],
comm=comm, size=size, rank=rank,
overwrite=overwrite, dtype=dtype)
print("writing "+key+" shape {}".format([3,mz,my,mx]))
else:
dataset_h5(group, key, status=status, shape=[mz,my,mx],
comm=comm, size=size, rank=rank,
overwrite=overwrite, dtype=dtype)
print("writing "+key+" shape {}".format([mz,my,mx]))
for ichunk in range(allchunks):
for iz in [indz[np.mod(ichunk,nchunks[2])]]:
n1, n2 = iz[ 0]-nghost,\
iz[-1]+nghost+1
n1out = n1+nghost
n2out = n2-nghost
varn1 = nghost
varn2 = -nghost
if iz[0] == locindz[0][0]:
n1out = 0
varn1 = 0
if iz[-1] == locindz[-1][-1]:
n2out = n2
varn2 = n2
for iy in [indy[np.mod(ichunk+
int(ichunk/nchunks[2]),nchunks[1])]]:
m1, m2 = iy[ 0]-nghost,\
iy[-1]+nghost+1
m1out = m1+nghost
m2out = m2-nghost
varm1 = nghost
varm2 = -nghost
if iy[0] == locindy[0][0]:
m1out = 0
varm1 = 0
if iy[-1] == locindy[-1][-1]:
m2out = m2
varm2 = m2
for ix in [indx[np.mod(ichunk+int(ichunk/nchunks[2])
+int(ichunk/nchunks[1]),nchunks[0])]]:
l1, l2 = ix[ 0]-nghost,\
ix[-1]+nghost+1
l1out = l1+nghost
l2out = l2-nghost
varl1 = nghost
varl2 = -nghost
if ix[0] == locindx[0][0]:
l1out = 0
varl1 = 0
if ix[-1] == locindx[-1][-1]:
l2out = l2
varl2 = l2
if not quiet:
print("remeshing "+key+" chunk {}".format(
[iz,iy,ix]))
var = calc_rhs_data(src, dst,
key, par, gd, l1, l2, m1, m2, n1, n2,
nghost=nghost, Reynolds_shock=Reynolds_shock,
lmix=lmix)
if is_vector(key):
dst["calc"][key][:,n1out:n2out,
m1out:m2out,
l1out:l2out] = dtype(var[:,
varn1:varn2,
varm1:varm2,
varl1:varl2])
else:
dst["calc"][key][n1out:n2out,
m1out:m2out,
l1out:l2out] = dtype(var[
varn1:varn2,
varm1:varm2,
varl1:varl2])
def load_dataset(src, key, lindx, lindy, lindz, nghost):
if is_vector(key):
var = np.empty([3, lindz.size, lindy.size, lindx.size])
else:
var = np.empty([lindz.size, lindy.size, lindx.size])
nz, ny, nx = src[key][()].shape[-3], src[key][()].shape[-2], src[key][(
)].shape[-1]
if var.shape == src[key][()].shape:
var = src[key][()]
else:
n3, n4, m3, m4, l3, l4 = None, None, None, None, None, None
# define array indices in z direction
if var.shape[-3] == nz:
n1, n2 = lindz[0], lindz[-1] + 1
nn1, nn2 = 0, lindz.size
else:
if lindz[0] < lindz[-1]:
n1, n2 = lindz[0], lindz[-1] + 1
nn1, nn2 = 0, lindz.size
else:
if lindz[-1] > nz:
n1 = lindz[0]
n2 = nz - nghost
n3 = nghost
n4 = lindz[-1] + 1
nn1 = 0
nn2 = n2 - n1
nn3 = n2 - n1
nn4 = lindz.size
else:
n3 = nghost
n4 = lindz[-1] + 1
n1 = lindz[0]
n2 = nz - nghost
nn1 = 0
nn2 = n2 - n1
nn3 = n2 - n1
nn4 = lindz.size
# define array indices in y direction
if var.shape[-2] == ny:
m1, m2 = lindy[0], lindy[-1] + 1
mm1, mm2 = 0, lindy.size
else:
if lindy[0] < lindy[-1]:
m1, m2 = lindy[0], lindy[-1] + 1
mm1, mm2 = 0, lindy.size
else:
if lindy[-1] > ny:
m1 = lindy[0]
m2 = ny - nghost
m3 = nghost
m4 = lindy[-1] + 1
mm1 = 0
mm2 = m2 - m1
mm3 = m2 - m1
mm4 = lindy.size
else:
m3 = nghost
m4 = lindy[-1] + 1
m1 = lindy[0]
m2 = ny - nghost
mm1 = 0
mm2 = m2 - m1
mm3 = m2 - m1
mm4 = lindy.size
# define array indices in x direction
if var.shape[-1] == nx:
l1, l2 = lindx[0], lindx[-1] + 1
ll1, ll2 = 0, lindx.size
else:
if lindx[0] < lindx[-1]:
l1, l2 = lindx[0], lindx[-1] + 1
ll1, ll2 = 0, lindx.size
else:
if lindx[-1] > nx:
l1 = lindx[0]
l2 = nx - nghost
l3 = nghost
l4 = lindx[-1] + 1
ll1 = 0
ll2 = l2 - l1
ll3 = l2 - l1
ll4 = lindx.size
else:
l3 = nghost
l4 = lindx[-1] + 1
l1 = lindx[0]
l2 = nx - nghost
ll1 = 0
ll2 = l2 - l1
ll3 = l2 - l1
ll4 = lindx.size
var[:, nn1:nn2, mm1:mm2, ll1:ll2] = src["bb"][:, n1:n2, m1:m2, l1:l2]
if l3:
var[:, nn1:nn2, mm1:mm2, ll3:ll4] = src["bb"][:, n1:n2, m1:m2,
l3:l4]
if m3:
var[:, nn1:nn2, mm3:mm4, ll3:ll4] = src["bb"][:, n1:n2, m3:m4,
l3:l4]
var[:, nn1:nn2, mm3:mm4, ll1:ll2] = src["bb"][:, n1:n2, m3:m4,
l1:l2]
if n3:
var[:, nn3:nn4, mm3:mm4, ll3:ll4] = src["bb"][:, n3:n4,
m3:m4, l3:l4]
var[:, nn3:nn4, mm1:mm2, ll3:ll4] = src["bb"][:, n3:n4,
m1:m2, l3:l4]
var[:, nn3:nn4, mm3:mm4, ll1:ll2] = src["bb"][:, n3:n4,
m3:m4, l1:l2]
var[:, nn3:nn4, mm1:mm2, ll1:ll2] = src["bb"][:, n3:n4,
m1:m2, l1:l2]
elif m3:
var[:, nn1:nn2, mm3:mm4, ll1:ll2] = src["bb"][:, n1:n2, m3:m4,
l1:l2]
if n3:
var[:, nn3:nn4, mm3:mm4, ll1:ll2] = src["bb"][:, n3:n4, m3:m4,
l1:l2]
var[:, nn3:nn4, mm1:mm2, ll1:ll2] = src["bb"][:, n3:n4, m1:m2,
l1:l2]
else:
if n3:
var[:, nn3:nn4, mm1:mm2, ll1:ll2] = src["bb"][:, n3:n4, m1:m2,
l1:l2]
return var
def kernel_smooth(
sim_path,
src,
dst,
magic=["meanuu"],
par=[],
comm=None,
gd=[],
grp_overwrite=False,
overwrite=False,
rank=0,
size=1,
nghost=3,
kernel=1.,
status="a",
chunksize=1000.0,
dtype=np.float64,
quiet=True,
nmin=32,
typ='piecewise',
mode=list(),
):
if comm:
overwrite = False
if isinstance(par, list):
os.chdir(sim_path)
par = read.param(quiet=True, conflicts_quiet=True)
if isinstance(gd, list):
os.chdir(sim_path)
gd = read.grid(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],
)
# extend gost zones to include up to 1.5 * kernel length)
dx = max(src['grid/dx'][()], src['grid/dy'][()], src['grid/dz'][()])
nkernel = np.int(2.5 * kernel / dx)
sigma = kernel / dx
print('sigma {:.2f}, kernel {:.2f}, dx {:.2f}'.format(sigma, kernel, dx))
# 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]
if 1 in nchunks:
mode = ["reflect", "reflect", "reflect"]
for ich in range(3):
if nchunks[ich] == 1:
mode[2 - ich] = "wrap"
if mode[2 - ich] == "reflect":
typ = "piecewise"
else:
typ = "all"
print('mode:', mode, 'typ:', typ)
# save time
dataset_h5(
dst,
"time",
status=status,
data=src["time"][()],
comm=comm,
size=size,
rank=rank,
overwrite=overwrite,
dtype=dtype,
)
# ensure derived variables are in a list
if isinstance(magic, list):
magic = magic
else:
magic = [magic]
# initialise group
group = group_h5(
dst,
"data",
status="a",
overwrite=grp_overwrite,
comm=comm,
rank=rank,
size=size,
)
for key in magic:
if is_vector(key):
dataset_h5(
group,
key + str(nkernel),
status=status,
shape=[3, mz, my, mx],
comm=comm,
size=size,
rank=rank,
overwrite=overwrite,
dtype=dtype,
)
print("writing " + key + " shape {}".format([3, mz, my, mx]))
else:
dataset_h5(
group,
key + str(nkernel),
status=status,
shape=[mz, my, mx],
comm=comm,
size=size,
rank=rank,
overwrite=overwrite,
dtype=dtype,
)
print("writing " + key + " shape {}".format([mz, my, mx]))
for ichunk in range(allchunks):
for iz in [indz[np.mod(ichunk, nchunks[2])]]:
if nchunks[2] == 1:
zextra = nghost
else:
zextra = nkernel + nghost
n1, n2 = iz[0] - zextra, iz[-1] + zextra + 1
lindz = np.arange(n1, n2)
n1out = n1 + zextra
n2out = n2 - zextra
varn1 = zextra
varn2 = -zextra
if iz[0] == locindz[0][0]:
n1out = 0
varn1 = zextra - nghost
if iz[-1] == locindz[-1][-1]:
n2out = n2 - zextra + nghost
varn2 = n2 - n1 - zextra + nghost
if n1 < 0:
lindz[np.where(lindz < nghost)[0]] += nz
if n2 > mz - 1:
lindz[np.where(lindz > mz - 1 - nghost)[0]] -= nz
print('n1out {},n2out {},varn1 {},varn2 {},zextra {}'.format(
n1out, n2out, varn1, varn2, zextra))
for iy in [
indy[np.mod(ichunk + int(ichunk / nchunks[2]),
nchunks[1])]
]:
if nchunks[1] == 1:
yextra = nghost
else:
yextra = nkernel + nghost
m1, m2 = iy[0] - yextra, iy[-1] + yextra + 1
lindy = np.arange(m1, m2)
m1out = m1 + yextra
m2out = m2 + 1 - yextra
varm1 = yextra
varm2 = -yextra
if iy[0] == locindy[0][0]:
m1out = 0
varm1 = yextra - nghost
if iy[-1] == locindy[-1][-1]:
m2out = m2 - yextra + nghost
varm2 = m2 - m1 - yextra + nghost
if m1 < 0:
lindy[np.where(lindy < 0)[0]] += ny
if m2 > my - 1:
lindy[np.where(lindy > my - 1)[0]] -= ny
print(
'm1out {},m2out {},varm1 {},varm2 {},yextra {}'.format(
m1out, m2out, varm1, varm2, yextra))
for iy in [
indy[np.mod(ichunk + int(ichunk / nchunks[2]),
nchunks[1])]
]:
for ix in [
indx[np.mod(
ichunk + int(ichunk / nchunks[2]) +
int(ichunk / nchunks[1]),
nchunks[0],
)]
]:
if nchunks[1] == 1:
xextra = nghost
else:
xextra = nkernel + nghost
l1, l2 = ix[0] - xextra, ix[-1] + xextra + 1
lindx = np.arange(l1, l2)
l1out = l1 + xextra
l2out = l2 + 1 - xextra
varl1 = xextra
varl2 = -xextra
if ix[0] == locindx[0][0]:
l1out = 0
varl1 = xextra - nghost
if ix[-1] == locindx[-1][-1]:
l2out = l2 - xextra + nghost
varl2 = l2 - l1 - xextra + nghost
if l1 < 0:
lindx[np.where(lindx < 0)[0]] += nx
if l2 > mx - 1:
lindx[np.where(lindx > mx - 1)[0]] -= nx
print('l1out {},l2out {},varl1 {},varl2 {},xextra {}'.
format(l1out, l2out, varl1, varl2, xextra))
if not quiet:
print("remeshing " + key +
" chunk {}".format([iz, iy, ix]))
print('sending ichunk {} with index ranges {}'.format(
ichunk, [n1, n2, m1, m2, l1, l2]))
var = smoothed_data(src["data"], dst["data"], key, par,
gd, lindx, lindy, lindz, nghost,
sigma, typ, mode)
print(
'ichunk {}, var min {:.1e}, var max {:.1e}'.format(
ichunk, var.min(), var.max()))
# print('var shape {}'.format(var.shape))
# if not quiet:
# print('writing '+key+
# ' shape {} chunk {}'.format(
# var.shape, [iz,iy,ix]))
print('ichunk: out indices {}'.format(
[n1out, n2out, m1out, m2out, l1out, l2out]))
if is_vector(key):
dst["data"][key +
str(nkernel)][:, n1out:n2out,
m1out:m2out,
l1out:l2out] = dtype(
var[:, varn1:varn2,
varm1:varm2,
varl1:varl2])
else:
dst["data"][key +
str(nkernel)][n1out:n2out, m1out:m2out,
l1out:l2out] = dtype(
var[varn1:varn2,
varm1:varm2,
varl1:varl2])
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)