def topo_main(filename, savename, varname, topname,
latname, lonname, depname, timname):
"""
Creates a mask of are values for a file or a set of splitted files.
If one file is used this can be automatically splitted using ind_out.
Parameters
----------
filename : str
Name of the file to be open.
savename : str
Name of the file to be saved.
latname : str
Name of the latitude variable.
lonname : str
Name of the longitude variable.
Nproc : int, optional
If bigger than 1 will run in parallel. The default is 1.
ind : tuple, optional
If the file(s) to be loaded come(s) from a partition from breakds
a tuple with the index must be provided (Z, Y, X). Where X, Y, Z
are the respective index in each dimension (floats or array like).
In that case the loaded data will be appended.
The default is None.
Nb : tuple of 3 int, optional
Number of partitions to be made in the z, y and x axis.
Only used if ind is None and Nb not None. The default is None.
Returns
-------
None.
"""
[bathy], lat, lon, tim, dep = load_main(filename, [varname], latname,
lonname, depname, timname)
sh = bathy.shape
if len(sh) == 2:
bathy = bathy[None, :, :]
sh = bathy.shape
topo = np.empty_like(bathy)
for i in range(sh[1]):
for j in range(sh[2]):
topo[:, i, j] = dep[int(bathy[0, i, j])]
print("Main depth: ", np.mean(topo))
ds = xr.Dataset({timname:(('t'), tim),
lonname: (('y', 'x'), lon),
latname: (('y', 'x'), lat),
topname: (('t','y', 'x'), topo)})
ds.to_netcdf(savename+'.nc')
def stra_1file(file_den, file_str, bpnt, denname, latname, lonname, depname,
strname, timname, H):
#############
# LOAD DATA #
#############
[den], lats, lons, tim, dep = load_main(file_den, [denname], latname,
lonname, depname, timname)
#dind = np.logical_and(dep > 500, dep < 2000)
#den, dep = den[:, dind], dep[dind]
dim = den.shape
strval = np.empty((dim[0], len(bpnt[0]), dim[2], dim[3]))
for k, (h, p) in enumerate(zip(bpnt[0], bpnt[1])):
dind = np.logical_and(dep > h - 250, dep < h + 250)
denk, depk = den[:, dind], dep[dind]
for t in range(dim[0]):
for j in range(dim[2]):
for i in range(dim[3]):
nonanval = ~np.isnan(denk[t, :, j, i])
denkji = denk[t, nonanval, j, i]
depkji = depk[nonanval]
if len(depkji) == 0\
or p < np.min(denkji)\
or p > np.max(denkji):
strval[t, k, j, i] = np.nan
print("nan value encountered")
else:
strval[t, k, j, i] =\
interp1d(denkji, depkji, kind='cubic')(p)
ds = {
lonname: (('y', 'x'), lons),
latname: (('y', 'x'), lats),
strname: (('t', 'z', 'y', 'x'), strval)
}
if timname is not None:
ds[timname] = (('t'), tim)
ds = xr.Dataset(ds)
ds.to_netcdf(file_str + '.nc')
def stra_main(file_den,
file_str,
bpnt,
denname,
latname,
lonname,
depname,
strname,
timname=None,
Nproc=1,
ind=None,
H=5000):
if ind is None:
stra_1file(file_den, file_str, bpnt, denname, latname, lonname,
depname, strname, timname, H)
else:
def str_kji(kji):
k, j, i = kji
fden = file_den + '_' + str(k) + '_' + str(j) + '_' + str(i)
fstr = file_str + '_' + str(k) + '_' + str(j) + '_' + str(i)
stra_1file(fden, fstr, bpnt, denname, latname, lonname, depname,
strname, timname, H)
return 1
# Get iterables for the 3 index and call combinations
indk = int2iterable(ind[0])
indj = int2iterable(ind[1])
indi = int2iterable(ind[2])
for k in indk:
kji_com = [(k, j, i) for j in indj for i in indi]
totl = len(kji_com)
# Run in parallel
if Nproc > 1:
kji_com = split_iterable(kji_com, Nproc)
print("Processing {} files with {} cores".format(totl, Nproc))
totl = len(kji_com)
for i, kji_ in enumerate(kji_com):
print("\t{:.2f}%".format(100. * i * (k + 1) / totl))
output = []
for kji in kji_:
run_paral = dask.delayed(str_kji)(kji)
output.append(run_paral)
total = dask.delayed(sum)(output)
total.compute()
# Run in series
else:
print("Processing {} files".format(totl))
for i, kji_ in enumerate(kji_com):
print("\t{:.2f}%".format(100. * i * (k + 1) / totl))
str_kji(kji_)
[strval], lats, lons, tim, _ = load_main(file_str, [strname], latname,
lonname, None, timname, ind)
ds = {
lonname: (('y', 'x'), lons),
latname: (('y', 'x'), lats),
strname: (('t', 'z', 'y', 'x'), strval)
}
if timname is not None:
ds[timname] = (('t'), tim)
ds = xr.Dataset(ds)
ds.to_netcdf(file_str + '.nc')
def den_1file(filename_tem, filename_sal, savename, temname,
salname, depname, denname, latname, lonname,
pressure_lvl, mlat, timname):
"""
Computes the density profile from a given file.
Check the documentation of den_main for more information.
"""
#############
# LOAD DATA #
#############
if filename_tem == filename_sal:
# Launch reading routine for 1 file
[tem, sal], lats, lons, tim, deps = load_main(filename_tem,
[temname, salname],
latname, lonname,
depname, timname)
else:
# Launch reading routine for 2 different files
[tem], _, _, _, _ = load_main(filename_tem, [temname],
latname, lonname,
depname, timname)
[sal], lats, lons, tim, deps = load_main(filename_sal, [salname],
latname, lonname,
depname, timname)
pdens = np.empty_like(sal)
dim = lats.shape
index = np.logical_and(sal == 0, tem == 0)
sal[index], tem[index] = np.nan, np.nan
######################
# GET DENSITY VALUES #
######################
# Depth variable is already a pressure level
if pressure_lvl:
for j in range(dim[0]):
for i in range(dim[1]):
for t in range(sal.shape[0]):
pdens[t, :, j, i] = pden(sal[t, :, j, i],
tem[t, :, j, i],
deps)
# Compute the pressure levels with a mean latitude value
elif mlat is not None:
pre = pres(deps, mlat)
for j in range(dim[0]):
for i in range(dim[1]):
for t in range(sal.shape[0]):
pdens[t, :, j, i] = pden(sal[t, :, j, i],
tem[t, :, j, i],
pre)
del pre
# Compute pressure levels with all the latitudes
else:
for j in range(dim[0]):
for i in range(dim[1]):
pre = pres(deps, lats[j, i])
for t in range(sal.shape[0]):
pdens[t, :, j, i] = pden(sal[t, :, j, i],
tem[t, :, j, i],
pre)
#############
# SAVE DATA #
#############
ds = {lonname: (('y', 'x'), lons),
latname: (('y', 'x'), lats),
depname: (('z'), deps),
denname: (('t', 'z', 'y', 'x'), pdens)}
if timname is not None:
ds[timname] = (('t'), tim)
ds = xr.Dataset(ds)
ds.to_netcdf(savename+'.nc')
def int_1file(filename, varnames, latname, lonname, depname, timname, lat_out,
lat_outr, lon_out, lon_outr, N, method, Nproc, parallel, ind):
#############
# LOAD DATA #
#############
print("Loading data from " + filename + ".nc")
vars_in, lat_in, lon_in, tim, deps = load_main(filename, varnames, latname,
lonname, depname, timname,
ind)
NT = vars_in[0].shape[0]
NL = []
for var in vars_in:
if len(var.shape) == 4:
NL.append(var.shape[1])
else:
NL.append(0)
vars_out = []
for nl in NL:
if nl > 0:
vars_out.append(np.empty((NT, nl, N, N)))
else:
vars_out.append(np.empty((NT, N, N)))
lat_inr = lat_in.ravel()
lon_inr = lon_in.ravel()
#################
# INTERPOLATION #
#################
def int_v(v):
print("\tInterpolating " + varnames[v])
vars_out[v] = int_var(vars_in[v], lat_inr, lon_inr, NT, NL[v],
lat_outr, lon_outr, N, method, Nproc, parallel)
return 1
if (Nproc > 1) and (len(vars_out) > 1) and (parallel == 'var'):
output = []
for v in range(len(varnames)):
run_paral = dask.delayed(int_v)(v)
output.append(run_paral)
total = dask.delayed(sum)(output)
total.compute()
else:
for v in range(len(varnames)):
int_v(v)
#############
# SAVE DATA #
#############
ds = {lonname: (('y', 'x'), lon_out), latname: (('y', 'x'), lat_out)}
if depname is not None:
ds[depname] = (('z'), deps)
if timname is not None:
ds[timname] = (('t'), tim)
for nl, var, varn in zip(NL, vars_out, varnames):
if nl > 0:
ds[varn] = (('t', 'z', 'y', 'x'), var)
else:
ds[varn] = (('t', 'y', 'x'), var)
ds = xr.Dataset(ds)
ds.to_netcdf(filename + '_' + str(N) + 'x' + str(N) + '.nc')