PARAMOUR_TOOLS/nemo_bdy/generate_new_bdy_geom.py

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()