def __init__(self, grid, param):
for key in ['shape', 'incr', 'procs']:
setattr(self, key, grid[key])
for key in ['omega', 'ndeepest', 'myrank']:
setattr(self, key, param[key])
nh = 1
procs0 = [i*p for i, p in zip(self.incr, self.procs)]
# print("---------> procs0=", procs0, self.myrank)
loc = topo.rank2loc(self.myrank, procs0)
neighbours = topo.get_neighbours(loc, procs0, incr=self.incr)
size, domainindices = topo.get_variable_shape(
self.shape, neighbours, nh)
N = np.prod(size)
self.domainindices = domainindices
self.size = size
self.N = N
self.nh = nh
self.neighbours = neighbours
self.x = np.zeros((self.N,))
self.y = np.zeros((self.N,))
self.b = np.zeros((self.N,))
self.r = np.zeros((self.N,))
self.msk = np.zeros(self.size)
k0, k1, j0, j1, i0, i1 = self.domainindices
self.msk[k0:k1, j0:j1, i0:i1] = 1.
# self.xx = self.x
# self.xx.shape = self.size
# self.bb = self.b
# self.bb.shape = self.size
# self.rr = self.r
# self.rr.shape = self.size
# self.xx = self.x.reshape(self.size)
# self.bb = self.b.reshape(self.size)
# self.rr = self.r.reshape(self.size)
# self.xx = np.reshape(self.x, self.size)
# self.bb = np.reshape(self.b, self.size)
# self.rr = np.reshape(self.r, self.size)
# self.xx = np.zeros(self.size)
# self.bb = np.zeros(self.size)
# self.rr = np.zeros(self.size)
# self.x = self.xx.ravel()
# self.b = self.bb.ravel()
# self.r = self.rr.ravel()
self.halo = halo.Halo({"nh": nh, "size": size, "neighbours": neighbours,
"domainindices": domainindices, "shape": self.shape})
def __init__(self, user_param):
# Check, freeze, and get user parameters
user_param.check()
user_param.freeze()
param = user_param.view_parameters()
self.param = param
npx = param["npx"]
npy = param["npy"]
npz = param["npz"]
# Add parameters that are automatically set
param["nx"] = param["global_nx"] // npx
param["ny"] = param["global_ny"] // npy
param["nz"] = param["global_nz"] // npz
# Set up MPI
topo.topology = param["geometry"]
procs = [npz, npy, npx]
myrank = mpitools.get_myrank(procs)
loc = topo.rank2loc(myrank, procs)
neighbours = topo.get_neighbours(loc, procs)
param["procs"] = procs
param["myrank"] = myrank
param["neighbours"] = neighbours
param["loc"] = loc
self.myrank = myrank
if self.myrank == 0:
print("-" * 80)
self.banner()
# Load the grid with the extended parameters
self.grid = grid.Grid(param)
# Load the IO; only the parameters modifiable by the user are saved
self.IO = nylesIO.NylesIO(param)
# redirect the x11 to output.txt
#sys.stdout = Logger(self.IO.output_directory+'/output.txt')
# save the param dictionary in a pkl file
if self.myrank == 0:
with open(self.IO.output_directory + "/param.pkl", "wb") as fid:
pickle.dump(param, fid)
# Initiate the model and needed variables
self.initiate(param)
fid = open('%s/stats.pkl' % path, 'bw')
pickle.dump(self.stats, fid)
if __name__ == '__main__':
import numpy as np
from matplotlib import pyplot as plt
import topology as topo
procs = [1, 1, 1]
topo.topology = 'closed'
myrank = 0
nh = 2
loc = topo.rank2loc(myrank, procs)
neighbours = topo.get_neighbours(loc, procs)
nz, ny, nx = 32, 32, 128
Lx, Ly, Lz = 1.0, 1.0, 1.0
dx, dy, dz = Lx/nx, Ly/ny, Lz/nz
param = {
'nx': nx, 'ny': ny, 'nz': nz, 'nh': nh,
'Lx': Lx, 'Ly': Ly, 'Lz': Lz,
'neighbours': neighbours,
# Choose a timestepping method
'timestepping': 'LFAM3',
# 'timestepping': 'EF',
# Set the order of the upwind scheme
'orderA': 5,
# 'orderA': 3,
# 'orderA': 1,
def test_111_domain():
topo.topology = 'perio_xyz'
extension = 26
nz = 1
ny = 1
nx = 1
nh = 1
loc = topo.rank2loc(myrank, procs)
neighbours = topo.get_neighbours(loc, procs, extension=extension)
param = {'nx': nx, 'ny': ny, 'nz': nz, 'nh': nh, 'neighbours': neighbours}
state = var.get_state(param)
b = state.get('b')
# we can define the halo toolbox
grid = {
'shape': [nz, ny, nx],
'size': b.size,
'nh': nh,
'neighbours': neighbours,
'domainindices': b.domainindices,
'extension': extension
}
halo = Halo(grid)
# for r, buf in halo.rbuf.items():
# if r in neighbours.keys():
# buf[:] = neighbours[r]*1.
# else:
# buf[:] = -1.
# set the field equal to myrank
# making clear that the halo is wrong
b.activeview = 'i'
# if myrank == 0:
# b.view('i')[:, :, :] = 999
# else:
# b.view('i')[:, :, :] = 0.
x = b.view('i')
x[:, :, :] = myrank
halo.fill(b)
# now in the halo we'd have the rank of the neighbour
y = np.asarray(x.copy(), dtype=int)
msg = 'pb withneighbours'
for direc, yourrank in neighbours.items():
dk, dj, di = direc
assert (y[1 + dk, 1 + dj, 1 + di] == yourrank), msg
if myrank == 0:
print('fill halo of Scalar is okay')
l = nz + 2 * nh
m = ny + 2 * nh
n = nx + 2 * nh
x = np.zeros((l, m, n))
grid = {
'shape': [nz, ny, nx],
'size': np.shape(x),
'nh': nh,
'neighbours': neighbours,
'domainindices': b.domainindices,
'extension': extension
}
halox = Halo(grid)
x[:] = myrank
halox.fill(x)
y = np.asarray(x.copy(), dtype=int)
for direc, yourrank in neighbours.items():
dk, dj, di = direc
#print(myrank, y[1+dk, 1+dj, 1+di] , yourrank)
assert (y[1 + dk, 1 + dj, 1 + di] == yourrank), msg
if myrank == 0:
print('fill halo of ndarray is okay')
procs = [1, 3, 3]
msg = 'use mpirun -np %i python ' % np.prod(procs) + ' '.join(sys.argv)
assert comm.Get_size() == np.prod(procs), msg
#test_111_domain()
topo.topology = 'closed'
extension = 18
nz, ny, nx = 1, 2, 2
shape = [nz, ny, nx]
nh = 2
loc = topo.rank2loc(myrank, procs)
print('myrank=%i / loc=' % myrank, loc)
neighbours = topo.get_neighbours(loc, procs, extension=extension)
size, domainindices = topo.get_variable_shape([nz, ny, nx], neighbours, nh)
grid = {
'shape': [nz, ny, nx],
'size': size,
'nh': nh,
'neighbours': neighbours,
'domainindices': domainindices,
'extension': extension
}
check_halo_width(procs, shape, nh)
halox = Halo(grid)
x = np.zeros(size)
x[:] = myrank
halox.fill(x)
def set_subdomains(allgrids):
grid = allgrids[0]
procs = grid['procs'].copy()
procs0 = procs.copy()
incr = [1, 1, 1]
nprocs = np.prod(procs)
ranks = np.arange(nprocs)
loc = topo.rank2loc(ranks, procs)
loc0 = loc.copy()
k0, j0, i0 = loc0
gathers = []
doms = []
first = True
tags = ranks * 0
subdom_partition = []
for lev, grid in enumerate(allgrids):
c = grid['coarsen']
keys = c.keys()
if any(['g' in k for k in keys]) or lev == 0:
g = []
levs = [lev]
for n, direc in enumerate('kji'):
if 'g' + direc in keys:
factor = 2 * c['g' + direc]
incr[n] *= factor
procs[n] //= factor
loc[n] //= factor
g += [direc]
if g != '':
gathers += [g]
else:
pass
k, j, i = loc
nz, ny, nx = procs0
incz, incy, incx = incr
dom = [0] * nprocs
for r in ranks:
l = [k[r], j[r], i[r]]
# d is the index of the subdomain
# rank with same d do the same computation
d = topo.myloc(l, incr, [0, 0, 0], procs0)
dom[r] = d
doms += [dom]
dd = set(dom)
neighbours = [{}] * len(ranks)
domloc = {}
glued = {}
for d in dd:
family = [r for r in ranks if dom[r] == d]
# in a family, all ranks have the same loc
# because they all compute the same subdomain
locs = [(k0[r], j0[r], i0[r]) for r in family]
# tags is defined from the previous domain decomposition
fam = family[0]
glued[d] = [[r for r in family if tags[r] == t]
for t in set(tags[family])]
if len(glued[d]) == 1:
glued[d] = glued[d][0]
# print(' - subdomain : ', set([d]))
# print(' located at ', locs[0])
# print(' computed by ' % d, family)
# if not(first):
# print(' obtained by gluing ranks:', glued)
# here is the list of neighbours
#
# the print has been discarded to lighten the output
# but clearly this information is very important
# note that the neighbours relationships for a given
# rank depends on the subdomain decomposition
# the neighbours get further and further as
# subdomains are glued. This is accounted by
# 'incr'. incr == 1 on the finest grid, then
# whenever there is gluing in one direction
# incr is doubled in that direction, until
# the point where the is only one subdomain is that
# direction. Each rank is then neighbour with itself
#
# print(' neighbours:')
for r in family:
l = [k0[r], j0[r], i0[r]]
nghbs = topo.get_neighbours(l,
procs0,
incr=incr,
extension=26)
ngbs = fixneighbours(nghbs, r, d)
# print('rank : %i / ' % r, ngbs)
neighbours[r] = ngbs
r0 = d # family[0]
domloc[d] = [loc0[0][r0], loc0[1][r0], loc0[2][r0]]
# print('domain: %i' % d, domloc[d])
localpartition = {
'procs': procs.copy(),
'incr': incr.copy(),
'gluing': g.copy(),
'glued': glued.copy(),
'domloc': domloc.copy(),
'dom': dom.copy(),
'levels': levs,
'allneighbours': neighbours.copy()
}
# print(levs, domloc)
# if not(first):
localpartition['tags'] = tags.copy()
subdom_partition += [localpartition]
#
# tags is used to identify which ranks are glued together
#
# glued ranks had identical tag in the previous
# subdomain decomposition
#
for d in dd:
family = [r for r in ranks if dom[r] == d]
for k, r in enumerate(family):
tags[r] = k
first = False
else:
levs += [lev]
return subdom_partition
def join(param):
outfile = "%s.nc" % varname
procs = param["procs"]
nh = param["nh"]
nx = param["nx"]
ny = param["ny"]
nz = param["nz"]
topo.topology = param["geometry"]
global_nx = param["global_nx"]
global_ny = param["global_ny"]
global_nz = param["global_nz"]
size = [nz, ny, nx]
template = param["expname"]+"_%02i_hist.nc"
with nc.Dataset(outfile, "w") as fid, nc.Dataset(template % 0, "r") as fin:
nt = len(fin.dimensions["t"])
dims = fin.dimensions
var = fin.variables
assert varname in var
fid.createDimension("t", nt)
fid.createDimension("x", global_nx)
fid.createDimension("y", global_ny)
fid.createDimension("z", global_nz)
v = fid.createVariable(varname, "f", ("t", "z", "y", "x"))
v.long_name = var[varname].long_name
v.units = var[varname].units
for d in dims:
v = fid.createVariable(d, "f", (d, ))
v.long_name = var[d].long_name
v.units = var[d].units
nranks = np.prod(procs)
print("subdomains partition :", procs)
print("found %i snapshots" % nt)
print("start joining '%s'" % varname)
with nc.Dataset(outfile, "r+") as fid:
with nc.Dataset(template % 0, "r") as fin:
for kt in range(nt):
fid["t"][kt] = fin["t"][kt]
for k in range(procs[0]):
j, i = 0, 0
loc = [k, j, i]
r = topo.loc2rank(loc, procs)
ngs = topo.get_neighbours(loc, procs)
shape, domainindices = topo.get_variable_shape(
size, ngs, nh)
k0, k1, j0, j1, i0, i1 = domainindices
ka, kb = k*nz, (k+1)*nz
with nc.Dataset(template % r, "r") as fin:
fid["z"][ka:kb] = fin["z"][k0:k1]
for j in range(procs[1]):
k, i = 0, 0
loc = [k, j, i]
r = topo.loc2rank(loc, procs)
ngs = topo.get_neighbours(loc, procs)
shape, domainindices = topo.get_variable_shape(
size, ngs, nh)
k0, k1, j0, j1, i0, i1 = domainindices
ja, jb = j*ny, (j+1)*ny
with nc.Dataset(template % r, "r") as fin:
fid["y"][ja:jb] = fin["y"][j0:j1]
for i in range(procs[2]):
k, j = 0, 0
loc = [k, j, i]
r = topo.loc2rank(loc, procs)
ngs = topo.get_neighbours(loc, procs)
shape, domainindices = topo.get_variable_shape(
size, ngs, nh)
k0, k1, j0, j1, i0, i1 = domainindices
ia, ib = i*nx, (i+1)*nx
with nc.Dataset(template % r, "r") as fin:
fid["x"][ia:ib] = fin["x"][i0:i1]
for k in range(procs[0]):
for j in range(procs[1]):
for i in range(procs[2]):
loc = [k, j, i]
rank = topo.loc2rank(loc, procs)
ncfile = template % rank
ngs = topo.get_neighbours(loc, procs)
shape, domainindices = topo.get_variable_shape(
size, ngs, nh)
k0, k1, j0, j1, i0, i1 = domainindices
ka, kb = k*nz, (k+1)*nz
ja, jb = j*ny, (j+1)*ny
ia, ib = i*nx, (i+1)*nx
#print(rank, i0, i1, ia, ib, ngs)
with nc.Dataset(ncfile, "r") as fin:
for kt in range(nt):
print("\r %3i/%3i - %3i/%3i"
% (rank, nranks-1, kt, nt-1), end="")
var = fin[varname][kt][:, :, :]
z2d = var[k0:k1, j0:j1, i0:i1]
fid[varname][kt, ka:kb, ja:jb, ia:ib] = z2d
print()
print("%s has been joined into '%s'" % (varname, outfile))
def __getitem__(self, elem=slice(None)):
varname = self.varname
sizes = self.infos["sizes"]
dims = self.infos["dims"]
halow = self.infos["halow"]
tiles = self.infos["tiles"]
nctemplate = self.infos["nctemplate"]
missing = self.infos["missing"]
if elem == slice(None):
# read all
outdims = dims
iidx = [slice(None)] * len(dims)
elif isinstance(elem, int):
# assume it's time
outdims = dims[1:]
iidx = [slice(None)] * len(dims)
iidx[0] = elem
else:
raise ValueError("pb with the slice")
shape = [sizes[d] for d in outdims]
multiple_files = any([isinstance(s, list) for s in shape])
varshape = shape_from_sizes(sizes, outdims, halow)
var = np.zeros(varshape)
oidx = [slice(None)] * len(varshape)
if multiple_files:
oidx[-1] = slice(halow, -halow)
oidx[-2] = slice(halow, -halow)
if debug:
print(shape)
j0 = halow
ny = shape[-2]
nx = shape[-1]
for j in range(len(ny)):
i0 = halow
for i in range(len(nx)):
if missing[j, i]:
pass
else:
tile = tiles[j, i]
ncfile = getncfile(
nctemplate, tile, **self.kwargs
) #nctemplate(tile)#.format(tile=tile, **self.kwargs)
oidx[-1] = slice(i0, i0 + nx[i])
oidx[-2] = slice(j0, j0 + ny[j])
with Dataset(ncfile) as nc:
var[tuple(oidx)] = nc.variables[varname][iidx]
i0 += nx[i]
j0 += ny[j]
else:
tile = get_one_tile(tiles, missing)
ncfile = getncfile(
nctemplate, tile, **self.kwargs
) #nctemplate(tile)#.format(tile=tile, **self.kwargs)
with Dataset(ncfile) as nc:
var[tuple(oidx)] = nc.variables[varname][iidx]
if (halow > 0) and multiple_files:
# fill the halo
partition = self.infos["blocks"]["partition"]
npy = len(ny)
npx = len(nx)
for j in range(npy):
for i in range(npx):
tile0 = tiles[j, i]
neighbours = topo.get_neighbours(tile0, partition)
direc = []
# loop over directions
for neighb in neighbours.keys():
dj, di = neighb
# check if there is a halo to be filled
# in this direction
if (0 <= (j + dj) < npy) and (0 <= (i + di) < npx):
# this neighbour is an interior tile
pass
else:
direc += [(dj, di)]
if debug:
print(f"{tile} directions to use {direc}")
for dj, di in direc:
shape = (ny[j], nx[i])
tile, hiidx, hoidx = topo.get_haloinfos(
tile0, partition, shape, halow, (dj, di))
hoidx = (shiftslice(hoidx[0], sum(ny[:j])),
shiftslice(hoidx[1], sum(nx[:i])))
ncfile = getncfile(nctemplate, tile, **self.kwargs)
if os.path.isfile(ncfile):
with Dataset(ncfile) as nc:
oidx[-2:] = hoidx
iidx[-2:] = hiidx
if debug:
print(tile0, tile, hoidx, hiidx, oidx,
iidx, var.shape)
var[tuple(oidx)] = nc.variables[varname][tuple(
iidx)]
return var