source: TOOLS/CPLRESTART/FillOceRestart.py @ 6755

### ===========================================================================
###
### 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$"
'''
 All documentation available at https://forge.ipsl.jussieu.fr/igcmg/wiki/IPSLCM6/CPLRESTART

Extraction :
> svn co http://forge.ipsl.jussieu.fr/igcmg/svn/TOOLS/CPLRESTART CPLRESTART
'''

## SVN information
__SVN__ = ({
    'Author'   : "$Author$",
    'Date'     : "$Date$",
    'Revision' : "$Revision$",
    'Id'       : "$Id$",
    'HeadURL'  : "$HeadURL$",
    })
##
import shutil
import getopt
import sys
import numpy.ma as np
import netCDF4
from scipy import ndimage
import nemo

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 is not 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 is not None :
            ListExclude = ListVarName
            ListVarName  = None

if OuFile is 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
jpoi = OuFile.dimensions["x"].size
jpoj = OuFile.dimensions["y"].size
nperio = nemo.__guess_nperio__ (jpoj, jpoi, nperio=nperio, out='nperio')

if nperio is 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 is None : ListVarName = list ( OuFile.variables.keys() )

# Exclude some var if needed
for Var in ['time_centered', 'time_counter', 'time_centered_bounds', 'time_counter_bounds' ] :
    if Var in ListVarName : ListVarName.remove(Var)
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 !!!
##
## ===========================================================================