def __init__(self, fine, coarse):
assert any(coarse.incr != fine.incr)
comm = MPI.COMM_WORLD
nprocs = comm.Get_size()
myrank = comm.Get_rank()
matshape = tuple(coarse.incr // fine.incr)
xshape = coarse.shape
dummyshape = [j//i for i, j in zip(matshape, xshape)]
ngbs = fine.neighbours
nh = 1
dummysize, domainindices = topo.get_variable_shape(
dummyshape, ngbs, nh)
N = np.prod(dummysize)
self.shape = dummyshape
self.neighbours = ngbs
self.size = dummysize
self.N = N
self.domainindices = domainindices
self.nh = nh
self.x = np.zeros((self.N,))
self.A = []
self.halo = halo.Halo({"nh": nh, "size": dummysize, "neighbours": ngbs,
"domainindices": domainindices, "shape": self.shape})
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 attach_subdomain_to_grids(allgrids, subdomains, myrank):
for isub, subd in enumerate(subdomains):
ngbs = subd['allneighbours'][myrank]
for count, lev in enumerate(subd['levels']):
g = allgrids[lev]
shape = g['shape']
nh = g['nh']
size, domainindices = topo.get_variable_shape(shape, ngbs, nh)
g['subd'] = subd
g['subdomain'] = isub
g['neighbours'] = ngbs
g['N'] = np.prod(size)
g['size'] = size
g['domainindices'] = domainindices
if (count == 0) and (len(subd["gluing"]) > 0):
g["glueflag"] = True
g["gluing"] = subd["gluing"]
dom = subd["dom"]
g["glued"] = subd["glued"]
else:
g["glueflag"] = False
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)
if myrank in [4, 2]:
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 __init__(self,
param,
name,
nickname,
dimension,
prognostic: bool = False):
"""Construct a scalar field in 3D space for a physical quantity.
The arguments of the constructor have the same role as the
correspondent attributes of the class. See in the description
of the class for more information.
"""
# Copy the necessary information of param to a local dictionary;
# this will be replaced by an object of a class Param (like in Fluid2d)
required_param = ['nx', 'ny', 'nz', 'nh', 'neighbours']
self.param = {k: param[k] for k in required_param}
self.name = name
self.nickname = nickname
self.dimension = dimension
self.prognostic = prognostic
# Calculate size needed in every direction, taking into account
# the halo in the direction where there is a neighbour
neighbours = self.param['neighbours']
p = param
nx, ny, nz, nh = p['nx'], p['ny'], p['nz'], p['nh']
shape = [nz, ny, nx]
self.shape = shape
size, domainindices = topo.get_variable_shape(shape, neighbours, nh)
nzl, nyl, nxl = size
self.size = {'i': nxl, 'j': nyl, 'k': nzl}
self.domainindices = domainindices
# define self.mg_idx, the MG array index span
# MG arrays have the same halo policy
# than Scalars (halo only if necessary)
# with the difference that halo width is 1 (nh=1) for MG
# also MG arrays all use the same (k,j,i) convention
#
# if 'field' is a view('i') of a Scalar then
# field[mg_idx] returns an array the exact same size
# as a MG array
k0, k1, j0, j1, i0, i1 = domainindices
startk, endk = max(0, k0 - 1), min(nzl, k1 + 1)
startj, endj = max(0, j0 - 1), min(nyl, j1 + 1)
starti, endi = max(0, i0 - 1), min(nxl, i1 + 1)
kdx = slice(startk, endk)
jdx = slice(startj, endj)
idx = slice(starti, endi)
self.mg_idx = (kdx, jdx, idx)
self.mg_idx2 = np.array([startk, endk, startj, endj, starti, endi],
dtype=int)
# Create arrays extended by the halo;
# it might be smarter to remove the halo
# in the directions where it is not needed
self.data = {
'i': np.zeros((nzl, nyl, nxl)),
'j': np.zeros((nxl, nzl, nyl)),
'k': np.zeros((nyl, nxl, nzl)),
}
self.activeview = 'i'