source: TOOLS/CPLRESTART/FillOceRestart.py @ 4603

### ===========================================================================
###
### Fill a coupler restart, ocean side
###
### ===========================================================================
##
##  Warning, to install, configure, run, use any of Olivier Marti's
##  software or to read the associated documentation you'll need at least
##  one (1) brain in a reasonably working order. Lack of this implement
##  will void any warranties (either express or implied).
##  O. Marti assumes no responsability for errors, omissions,
##  data loss, or any other consequences caused directly or indirectly by
##  the usage of his software by incorrectly or partially configured
##  personal.
##
## SVN information
__Author__   = "$Author$"
__Date__     = "$Date$"
__Revision__ = "$Revision$"
__Id__       = "$Id$"
__HeadURL    = "$HeadURL$"
##
import netCDF4
import numpy.ma as np
from scipy import ndimage
import nemo
import shutil, getopt, sys

def MyFill (InputData=None) :
    """
    Replace the value of masked 'InputData' cells by the value of the nearest valid data cell
    From https://stackoverflow.com/questions/5551286/filling-gaps-in-a-numpy-array
    
    Input:  InputData : numpy.ma array of any dimension.

    Output: Return a filled array. 
    """    

    Invalid = np.where ( InputData[:,:].mask, True, False)
    Indices = ndimage.distance_transform_bf(Invalid, metric='euclidean', sampling=None, return_distances=False,
                                            return_indices=True, distances=None )
    FilledData = InputData[tuple(Indices)]
    return FilledData

def usage () :
    __help__ = """%(prog)s usage :
python %(prog)s [-d] -i <sstoce file> [-n <perio>] [-v variables] [-o <output file>]
 -d           | --debug                : debug
 -i <file>    | --input=<file>         : input file  (default: none)
 -o <file>    | --output=<file>        : output file (default : build a name form input file)
 -r           | --replace              : replace input file by new file with filled variables
 -v <varlist> | --variable=<variables> : list of variable to fill (default: all variable in file)
 -x           | --exclude              : fills all variable in files, except those given in -v|--variable
 -n <perio>   | --perio=<perio> : periodicity type (default: try to guess)
    1, 4, 6 : Cyclic on i dimension (generaly longitudes)
    2       : Obsolete (was symmetric condition at southern boundary ?)
    3, 4    : North fold T-point pivot (legacy ORCA2)
    5, 6    : North fold F-point pivot (ORCA1, ORCA025, ORCA2 with new grid for paleo)
    If <perio> is not specified, %(prog)s will try to guess it from the grid dimensions
example :
    python %(prog)s -n 4 -i sstoce_ORCA2.3.nc
    python %(prog)s -n 4 -i sstoce_ORCA2.3.nc -v O_SSTSST,OIceFrc
    """
    print ( __help__ % { 'prog':sys.argv[0] } )

## Default input parameters
InFile      = None
OuFile      = None
replace     = False
ListVarName = None
nperio      = None
Debug       = False
ListExclude = None ; Exclude=False

## Command line options
try:
    myopts, myargs = getopt.getopt ( sys.argv[1:], 'i:o:rv:xp:hd',
                                         [ 'input=', 'output=', 'replace', 'exclude', 'variable=', 'variables=', 'perio=', 'debug', 'help' ] )
except getopt.GetoptError as cmdle :
    print ( "Command line error : "+str(cmdle)+"\n" )
    usage ()
    sys.exit(1)

for myopt, myval in myopts :
    if   myopt in [ '-h', '--help'    ] : usage () ; sys.exit (0) ;
    elif myopt in [ '-d', '--debug'   ] : Debug    = True  ;
    elif myopt in [ '-i', '--input'   ] : InFile   = myval ;
    elif myopt in [ '-o', '--output'  ] :
        if replace :
            print ("Error : you can not specify both -r|--replace and -o|--output" )
            usage ()
            sys.exit(1)
        else :
            OuFile   = myval ;
            if Debug : print ("Out file set to " + OuFile)
    elif myopt in [ '-r', '--replace' ] :
        if OuFile != None :
            print ("Error : you can not specify both -r|--replace and -o|--output" )
            usage ()
            sys.exit(1)
        else :
            if Debug : print ("Out file set to input file")
            OuFile   = InFile ;
    elif myopt in [ '-p', '--perio'     ] : nperio = int(myval) ;
    elif myopt in [ '-v', '--variable', '--variables'  ] :
        if Exclude : ListExclude = myval.split(',')          
        else :       ListVarName = myval.split(',')
    elif myopt in [ '-x', '--exclude'  ] :
        Exclude = True
        if ListVarName != None :
            ListExclude = ListVarName
            ListVarName  = None

if OuFile == None :
    print ( 'Definition OuFile' )
    OuFile = InFile.replace ( ".nc", "_filled.nc" )
    print ("Creates output file name: " + OuFile)

# Copy the input file if needed
if OuFile != InFile : shutil.copyfile ( InFile, OuFile )

print ("Output file: " + OuFile )

# Open file
OuFile = netCDF4.Dataset ( OuFile , "r+" )

# Try to guess periodicity type
if nperio == None :
    print ( "Trying to guess nperio parameter" )
    jpoi = OuFile.dimensions["x"].size
    jpoj = OuFile.dimensions["y"].size
    print ("Grid dimensions: ("+str(jpoj)+", "+str(jpoi)+")")
    if   'ORCA2' in InFile :
        nperio=4
        print ("ORCA 2 grid found from file name: nperio may vary for this configuration")
        print ("Choosen nperio=" + str(nperio) )
    elif 'ORCA1' in InFile :
        if 'eORCA1' in InFile :
            nperio=6
            print ("eORCA 1 grid found from file name, nperio=" + str(nperio))
        else :
            nperio=6
            print ("ORCA 1 grid found from file name, nperio=" + str(nperio))
    elif (jpoj, jpoi) == (149, 182) :
        nperio = 4
        print ("ORCA 2 grid found from dimension: nperio may vary for this configuration")
        print ("Choosen nperio=" + str(nperio))
    elif (jpoj, jpoi) == (332, 292) :
        nperio = 6
        print ("ORCA1 grid found from dimensions, nperio" + str(nperio) )
    elif (jpoj, jpoi) == (332, 362) :
        nperio = 6
        print ("eORCA1 grid found from dimensions, nperio" + str(nperio) )
    elif (jpoj, jpoi) == (1021, 1442) :
        nperio = 6
        print ("ORCA025 grid found from dimensions, nperio" + str(nperio) )
    elif (jpoj, jpoi) == (1207, 1442) :
        nperio = 6
        print ("eORCA025 grid found from dimensions, nperio=" + str(nperio) )
        
if nperio == None :
    print ("%(prog)s couldn't guess the periodicity type of your file")
    print ("Please specify -p|--perio")
    usage ()
    sys.exit(1)
    
# Get variables from file is need
if ListVarName == None : ListVarName = list ( OuFile.variables.keys() )

# Exclude some var if needed
if Exclude : 
    for Var in ListExclude : 
        if Var in ListVarName : ListVarName.remove(Var)

# Loop on variables
for VarName in ListVarName :
    Var    = OuFile.variables[VarName]
    if 'mask' in dir(Var[...]) :
        print ( "Working on " + VarName )
        NewVar = MyFill ( InputData = Var[...] )
        NewVar = nemo.lbc (NewVar, nperio=nperio, cd_type='T', psgn=1.0)
        OuFile.variables[VarName][:,:] = NewVar[:,:]
    else :
        print ( VarName + " is not masked" )
        
# Close file : writes update variables.    
OuFile.close()

## ===========================================================================
##
##                               That's all folk's !!!
##
## ===========================================================================