from netCDF4 import Dataset import sys import time import numpy as np # Script for generating a NEMO BDY geometry file from a full 2D coordinates file. # Author: C. Pelletier # Note: this script just takes a NEMO grid file and "peels it off" (like a potato) # to get the input BDY file for a geometry. # All nodes are included, regardless of their mask status. # I think it's better to include them all, so that if the mask evolves (e.g. by changing # topography, or changing the minimum bathymetry), the BDY data and geometry will remain # relevant. The only drawback is that there are superfluous data (corresponding to masked cells), # but in practise it only means that data files are slightly larger. Since BDY stored are stored as # small 1D array, I do not consider this drawback a problem. data_folder="/storepelican/pelletie/nemo_bdy/grids" coord_file = "coord_paraso.nc" # must be the full domain and contain glamt, gphit, glamu, gphiu, glamv, gphiv # taking the "mesh mask" or the "coordinates" file from a run is OK is_size_bdy = [False, True, False, False] # boolean array: are the 4 sides of the domain an open bdy? Order: x=1, y=ny, x=nx, y=1 (in FORTRAN indexing convention; W/N/E/S if grid is orientated "normally", i.e. North upwards) nbr_out = 1 # "rim width" of the BDY. NEMO documentation says it should be 8-9, # but at UCL et al. in 2018 - 2021, I've only seen 1 used (Huot, Van Achter, Pelletier). # N. Jourdain (IGE, Grenoble) also uses 1 so... # I've tried setting 8 over my config (circumpolar 0.25°), it ran a bit but was numerically unstable. # I am pretty confident about the "geometric" implementation of nbr_out > 1, though... out_file = "eORCA025_SO_BDY-COORDS_rim%d"%nbr_out+".nc" extra_dbg = False # output extra dbg vars in out_file ### END INPUT coord_nc = Dataset(data_folder+'/'+coord_file, mode='r') grids = ['t', 'u', 'v'] lons = {} lats = {} for grid in grids: lons[grid] = coord_nc.variables['glam'+grid][:] lats[grid] = coord_nc.variables['gphi'+grid][:] coord_nc.close() ny, nx = np.shape(lons['t']) idx_ifort = np.arange(nx) + 1 idx_jfort = np.arange(ny) + 1 idx_ifort_2d = idx_ifort[None,:] * np.ones(shape=[ny], dtype=int)[:,None] idx_jfort_2d = idx_jfort[:,None] * np.ones(shape=[nx], dtype=int)[None,:] idx_ifort_flat = np.ndarray.flatten(idx_ifort_2d) idx_jfort_flat = np.ndarray.flatten(idx_jfort_2d) nbrt_3d_bdy = np.ma.array( data = np.ndarray(shape=[ny,nx,4], dtype=int), mask = False) if(is_size_bdy[0]): nbrt_3d_bdy[:,1:,0] = (idx_ifort[1:])[None,:] * (np.ones(shape=[ny], dtype=int))[:,None] - 1 else: nbrt_3d_bdy.mask[:,:,0] = True if(is_size_bdy[1]): nbrt_3d_bdy[:-1,:,1] = ny - (idx_jfort[:-1])[:,None] * (np.ones(shape=[nx], dtype=int))[None,:] else: nbrt_3d_bdy.mask[:,:,1] = True if(is_size_bdy[2]): nbrt_3d_bdy[:,:-1,2] = nx - (idx_ifort[:-1])[None,:] * (np.ones(shape=[ny], dtype=int))[:,None] else: nbrt_3d_bdy.mask[:,:,2] = True if(is_size_bdy[3]): nbrt_3d_bdy[1:,:,3] = (idx_jfort[1:])[:,None] * (np.ones(shape=[nx], dtype=int))[None,:] - 1 else: nbrt_3d_bdy.mask[:,:,3] = True which_bdy = np.ma.array(data = np.argmin(nbrt_3d_bdy, axis=-1), \ mask = np.all(nbrt_3d_bdy.mask, axis=-1)) nbrt_3d_bdy.mask[0,:,:] = True nbrt_3d_bdy.mask[-1,:,:] = True nbrt_3d_bdy.mask[:,0,:] = True nbrt_3d_bdy.mask[:,-1,:] = True nbr_flats = {} nbrs = {} # prendre j*nx*4 + i*4 + which_bdy[grid] idx_take = np.repeat(np.arange(ny), nx) * nx * 4 + np.tile(np.arange(nx), ny) * 4 + np.ndarray.flatten(which_bdy) nbr_flats['t'] = (np.ndarray.flatten(nbrt_3d_bdy))[idx_take] nbrs['t'] = np.ma.array( data = np.reshape( nbr_flats['t'], (ny,nx)), \ mask = np.all(nbrt_3d_bdy.mask)) nbrs['u'] = np.ma.array(data = np.ndarray(shape=[ny, nx], dtype=int), \ mask = False) nbrs['u'][:,-1].mask = True nbrs['u'][:,:-1] = np.minimum(nbrs['t'][:,:-1], nbrs['t'][:,1:]) nbrs['u'].mask[:,:-1] = np.logical_or(nbrs['t'].mask[:,:-1], nbrs['t'].mask[:,1:]) nbrs['v'] = np.ma.array(data = np.ndarray(shape=[ny, nx], dtype=int), \ mask = False) nbrs['v'][-1,:].mask = True nbrs['v'][:-1,:] = np.minimum(nbrs['t'][:-1,:], nbrs['t'][1:,:]) nbrs['v'].mask[:-1,:] = np.logical_or(nbrs['t'].mask[:-1,:], nbrs['t'].mask[1:,:]) out_nc = Dataset(data_folder+'/'+out_file, mode='w') if(extra_dbg): out_nc.createDimension('y', ny) out_nc.createDimension('x', nx) out_nc.createDimension('bdy', 4) tmp = out_nc.createVariable('which_'+'t', datatype='i2', dimensions=['y', 'x'], fill_value = -10) tmp[:] = which_bdy nbr_3d = out_nc.createVariable("nbr"+'t'+"_3d", datatype='i4', dimensions=['y', 'x', 'bdy'], fill_value = -10) nbr_3d[:] = nbrt_3d_bdy for grid in grids: tmp = out_nc.createVariable('nbr'+grid+'_2d', datatype='i4', dimensions=['y', 'x'], fill_value = -10) tmp[:] = nbrs[grid] n_xbs = {} out_nc.createDimension('yb', 1) for grid in grids: cnt = 0 iidx_tmp = [] jidx_tmp = [] lat_tmp = [] lon_tmp = [] nbr_tmp = [] curr_nbrflat = np.ndarray.flatten(nbrs[grid]) for l in range(1,nbr_out+1): print(l) idx_take = np.nonzero(curr_nbrflat == l)[0].astype(int) idx_ifort_take = (np.ndarray.flatten(idx_ifort_2d))[idx_take] idx_jfort_take = (np.ndarray.flatten(idx_jfort_2d))[idx_take] sort_order = [] n_take = np.size(idx_take) already_in = np.zeros(shape=[n_take], dtype=bool) idx_sorted = [] check = np.logical_and(np.logical_not(already_in), \ idx_ifort_take == np.amin(idx_ifort_take)) if(np.any(check)): where_line = (np.nonzero(check))[0] tmp_idx = np.argsort(idx_jfort_take[where_line]) idx_sorted = np.concatenate(( idx_sorted, (idx_take[where_line])[tmp_idx] )) already_in[where_line] = True check = np.logical_and(np.logical_not(already_in), \ idx_jfort_take == np.amax(idx_jfort_take)) if(np.any(check)): where_line = (np.nonzero(check))[0] tmp_idx = np.argsort(idx_ifort_take[where_line]) idx_sorted = np.concatenate(( idx_sorted, (idx_take[where_line])[tmp_idx] )) already_in[where_line] = True check = np.logical_and(np.logical_not(already_in), \ idx_ifort_take == np.amax(idx_ifort_take)) if(np.any(check)): where_line = (np.nonzero(check))[0] tmp_idx = (np.argsort(idx_jfort_take[where_line]))[::-1] idx_sorted = np.concatenate(( idx_sorted, (idx_take[where_line])[tmp_idx] )) already_in[where_line] = True check = np.logical_and(np.logical_not(already_in), \ idx_jfort_take == np.amin(idx_jfort_take)) if(np.any(check)): where_line = (np.nonzero(check))[0] tmp_idx = (np.argsort(idx_ifort_take[where_line]))[::-1] idx_sorted = np.concatenate(( idx_sorted, (idx_take[where_line])[tmp_idx] )) already_in[where_line] = True check = np.logical_not(already_in) if(np.count_nonzero(check) > 0): print('error in reordering') sys.exit() idx_sorted = idx_sorted.astype(int) iidx_tmp = np.concatenate(( iidx_tmp, idx_ifort_flat[idx_sorted] )) jidx_tmp = np.concatenate(( jidx_tmp, idx_jfort_flat[idx_sorted] )) lon_tmp = np.concatenate(( lon_tmp, (np.ndarray.flatten(lons[grid]))[idx_sorted])) lat_tmp = np.concatenate(( lat_tmp, (np.ndarray.flatten(lats[grid]))[idx_sorted])) curr_size = np.size(idx_take) nbr_tmp = np.concatenate(( nbr_tmp, np.repeat(l, curr_size) )) cnt+=curr_size out_nc.createDimension('xb'+grid.upper(), cnt) tmp = out_nc.createVariable(varname='nbi'+grid, datatype='i4', dimensions=['yb', 'xb'+grid.upper()]) tmp[:] = iidx_tmp tmp = out_nc.createVariable(varname='nbj'+grid, datatype='i4', dimensions=['yb', 'xb'+grid.upper()]) tmp[:] = jidx_tmp tmp = out_nc.createVariable(varname='nbr'+grid, datatype='i4', dimensions=['yb', 'xb'+grid.upper()]) tmp[:] = nbr_tmp tmp = out_nc.createVariable(varname='glam'+grid, datatype='f8', dimensions=['yb', 'xb'+grid.upper()]) tmp[:] = lon_tmp tmp = out_nc.createVariable(varname='gphi'+grid, datatype='f8', dimensions=['yb', 'xb'+grid.upper()]) tmp[:] = lat_tmp out_nc.close()