def load_grid(grdfiles, blocks, dimpart, nsigma, **kwargs): """Setup a `grid` by reading `grdfiles` on `blocks` """ ncgrid = nct.MDataset(grdfiles, blocks, dimpart, **kwargs) # dummy time, to be updated later time = ma.Marray(np.arange(10), dims=tdims) lat = nct.readmarray(ncgrid, "lat_rho", hdims) lon = nct.readmarray(ncgrid, "lon_rho", hdims) if ncgrid.halow > 0: halow = ncgrid.halow # extend lon-lat on the halow=1 to infer correctly @ f-points fill_halo(lat, halow) fill_halo(lon, halow) depth = nct.readmarray(ncgrid, "h", hdims) angle = nct.readmarray(ncgrid, "angle", hdims) mask = nct.readmarray(ncgrid, "mask_rho", hdims) pm = nct.readmarray(ncgrid, "pm", hdims) pn = nct.readmarray(ncgrid, "pn", hdims) f = nct.readmarray(ncgrid, "f", hdims) sigma = ma.Marray((np.arange(nsigma)+0.5)/nsigma, dims=vdims) coords = {"t": time, "sigma": sigma, "eta": lat, "xi": lon} return gr.Grid(coords, dims, depth=depth, angle=angle, mask=mask, pm=pm, pn=pn, f=f, **kwargs)
def gridslice(grid, marray, sliceindex): """ To return the surface level sliceindex = {"sigma": -1} """ dims = [d for d in marray.dims if not(d in sliceindex.keys())] out = ma.Marray(marray[-1], dims=dims) return out
def ncread(mdataset, grid, varname, elem=slice(None)): """Read a variable from a MDataset and returns it as a Marray consistent with `grid` """ assert varname in variables, f"this `{varname}` is not implemented" var = mdataset.variables[varname] # data = var[elem] attrs = var.attrs stagg = variables[varname] vardims = list(var.infos["dims"]) if isinstance(elem, int): attrs[vardims[0]] = elem vardims = vardims[1:] convert_netcdf_dim_to_var(vardims, grid.dims, mdataset.mapping) sizes = grid.get_sizes(stagg) shape_horiz_data = var.shape[-2:] shape_horiz_grid = [sizes[vardims[-2]], sizes[vardims[-1]]] #shape[-2] = sizes[vardims[-2]] #shape[-1] = sizes[vardims[-1]] #shape = tuple(shape) if shape_horiz_grid != shape_horiz_data: # fix shape # the fix is only on "eta" and "xi" jj0 = shape_horiz_grid[-2]-shape_horiz_data[-2] ii0 = shape_horiz_grid[-1]-shape_horiz_data[-1] if debug: print( f"grid {shape_horiz_grid} != data {shape_horiz_data} start index {jj0}, {ii0}") if isinstance(elem, int): shape = var.shape[1:-2]+shape_horiz_grid else: shape = var.shape[:-2]+shape_horiz_grid if debug: print(f"target shape {shape}") data = np.zeros(shape) if jj0 < 2: jslice = slice(jj0, None) elif jj0 == 2: jslice = slice(1, -1) if ii0 < 2: islice = slice(ii0, None) elif ii0 == 2: islice = slice(1, -1) data[..., jslice, islice] = var[elem] else: data = var[elem] return ma.Marray(data, attrs=attrs, dims=vardims, stagg=stagg)
def sigma2z(grid, zeta=0, stagg={}): """ compute depth from sigma coordinates """ staggering = {} for d in grid.dims: staggering[d] = (d in stagg) and stagg[d] sizes = grid.get_sizes(staggering) z = ma.Marray(ma.zeros(sizes), dims=grid.dims, attrs={"name": "depth"}, stagg=staggering) staggdims = tuple( [d for d in grid.dims if staggering[d] and d in grid.depth.dims]) if isinstance(zeta, int): timeison = False hinv = ma.Marray(zeta+grid.depth, dims=grid.depth.dims) else: timeison = True hinv = ma.Marray(zeta+grid.depth, dims=zeta.dims) hinv2 = grid.avg(hinv, staggdims) h = ma.Marray(grid.depth, dims=grid.depth.dims) h2 = grid.avg(h, staggdims) sigma = grid.sigma(stagg=staggering) if debug: print(hinv2.shape, h2.shape, staggering, sigma, z.shape) if "t" in grid.sizes: for kt in range(grid.sizes["t"]): if timeison: htot = hinv2[kt] else: htot = hinv for k in range(grid.sizes["sigma"]): z[kt, k] = -h2 + sigma[k]*htot else: for k in range(grid.sizes["sigma"]): z[k] = -h2 + sigma[k]*hinv2 return z
debug = False ma.debug = False dims = ("t", "sigma", "eta", "xi") hdims = dims[2:] vdims = dims[1] tdims = dims[0] sdims = dims[1:] nt, nz, ny, nx = 4, 2, 3, 5 lon, lat = np.meshgrid(range(nx), range(ny)) zeta = np.zeros((ny, nx)) depth = ma.Marray(np.ones((ny, nx)), dims=hdims) lat = ma.Marray(lat, dims=hdims, attrs={"name": "latitude"}) lon = ma.Marray(lon, dims=hdims, attrs={"name": "longitude"}) sigma = ma.Marray((np.arange(nz)+0.5)/nz, dims=vdims) time = ma.Marray(np.arange(nt), dims=tdims) coords = {"t": time, "sigma": sigma, "eta": lat, "xi": lon} g = Grid(coords, dims, depth=depth) space = {d: g.sizes[d] for d in sdims} # define a 3D spatial array (no time) temp = ma.Marray(ma.zeros(space), dims=sdims, attrs={"name": "temperature", "units": "degC"}) temp.flat = range(nz*ny*nx)
def readmarray(mfdataset, varname, dims, elem=slice(None)): """ Read `varname` and returns it as a Marray """ v = mfdataset.variables[varname] data = v[elem] return ma.Marray(data, attrs=v.attrs, dims=dims)