source: TOOLS/MOSAIX/check_conserv.py @ 6764

### ===========================================================================
###
### Checks integrals
###
### ===========================================================================
##
##  MOSAIX is under CeCILL_V2 licence. See "Licence_CeCILL_V2-en.txt"
##  file for an english version of the licence and
##  "Licence_CeCILL_V2-fr.txt" for a french version.
##
##  Permission is hereby granted, free of charge, to any person or
##  organization obtaining a copy of the software and accompanying
##  documentation covered by this license (the "Software") to use,
##  reproduce, display, distribute, execute, and transmit the
##  Software, and to prepare derivative works of the Software, and to
##  permit third-parties to whom the Software is furnished to do so,
##  all subject to the following:
##
##  Warning, to install, configure, run, use any of MOSAIX 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).  Authors 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
##
#
# python3 -i check_conserv.py --file dia_tORCA2.3_to_tLMD9695_1stOrder_Normalized_Surfacic_NoArea.nc --grids grids_ORCA2.3xLMD9695.nc --areas areas_ORCA2.3xLMD9695.nc --masks masks_ORCA2.3xLMD9695.nc --frac LMD9695_grid_maskFrom_ORCA2.3.nc
#
## SVN information
__Author__   = "$Author$"
__Date__     = "$Date$"
__Revision__ = "$Revision$"
__Id__       = "$Id:  $"
__HeadURL    = "$HeadURL$"

import netCDF4
import numpy as np
import glob
import sys, argparse, textwrap
import nemo
import os

# Creating a parser
parser = argparse.ArgumentParser (
    description = """Compute integral of matching coupling fields on atmosphere and ocean side.""",
    epilog='-------- This is the end of the help message --------')

# Adding arguments
parser.add_argument ( '--file'      , help='Input file, e.g. dia_tICO40_to_teORCA1.2_1storder_false.nc', type=str, default='dia_tico_to_torc_1storder_false.nc' )
parser.add_argument ( '--grids'     , help='grids file', type=str, default='grids.nc' )
parser.add_argument ( '--areas'     , help='areas file', type=str, default='areas.nc' )
parser.add_argument ( '--masks'     , help='masks file', type=str, default='masks.nc' )
parser.add_argument ( '--fracs'     , help='fracs file', type=str, default=None )
parser.add_argument ( '--guess'     , help='If set, build grids, areas and masks file name from the dia file name', action="store_true", default=False )
parser.add_argument ( '--src_int'   , help='Source field is integrated'      , action="store_true" )
parser.add_argument ( '--dst_int'   , help='Destination field is integrated' , action="store_true" )
parser.add_argument ( '--src_nofrac', help='Do not apply frac on source'     , action="store_true" )
parser.add_argument ( '--dst_nofrac', help='Do not apply frac on destination', action="store_true" )

# Parse command line
myargs = parser.parse_args()

#
name_input = myargs.file

# Get dir and file
InDir  =  os.path.dirname  (os.path.abspath(myargs.file))
InFile =  os.path.basename (os.path.abspath(myargs.file))

print ('InDir  : ' + InDir )
print ('InFile : ' + InFile )

Elements = InFile.split('_')

# Get models
Name_src = Elements[1][1:] ; Name_dst = Elements[3][1:]
grid_src = Elements[1][0]  ; grid_dst = Elements[3][0]

Grid_src = grid_src.upper() ; Grid_dst = grid_dst.upper()
if Grid_src in ( 'A', 'C') : Grid_src = 'T'
if Grid_dst in ( 'A', 'C') : Grid_dst = 'T'

# src side
if Name_src.count('ico') != 0 or Name_src.count('lmd') != 0 : name_src = Name_src ; ShortNames = True
else :
    if Name_src.count('ICO') : name_src = 'ico' ; ShortNames = False
    if Name_src.count('LMD') : name_src = 'lmd' ; ShortNames = False

if Name_dst.count('orc') != 0    : name_src = Name_src ; ShortNames = True
else :
    if Name_src.count('ORC')!= 0 : name_src = 'orc'    ; ShortNames = False
        
# dst side
if Name_dst.count('ico')!= 0 or Name_dst.count('lmd') != 0 : name_dst = Name_dst ; ShortNames = True
else :
    if Name_dst.count('ICO') : name_dst = 'ico' ; ShortNames = False
    if Name_dst.count('LMD') : name_dst = 'lmd' ; ShortNames = False

if Name_dst.count('orc') != 0 : name_dst = Name_dst ; ShortNames = True
else :
    if Name_dst.count('ORC') != 0 : name_dst = 'orc' ; ShortNames = False
    
if Name_src.count('ORC') : CplModel = Name_src + 'x' + Name_dst
if Name_dst.count('ORC') : CplModel = Name_dst + 'x' + Name_src

print ('CplModel   : ' + CplModel )
print ('Input file : ' + name_input )
     
# Coordonnees et masques
if myargs.fracs == None :
    if Name_src.count('ORC') : name_frc = InDir + '/' + Name_dst + '_grid_maskFrom_' + Name_src + '.nc'
    if Name_dst.count('ORC') : name_frc = InDir + '/' + Name_src + '_grid_maskFrom_' + Name_dst + '.nc'
else:
    name_frc = myargs.fracs

if myargs.guess :
    n_grids =  InDir + '/grids_' + CplModel + '.nc'
    n_masks =  InDir + '/masks_' + CplModel + '.nc'
    n_areas =  InDir + '/areas_' + CplModel + '.nc'
else :
    n_grids = myargs.grids
    n_masks = myargs.masks
    n_areas = myargs.areas

print ('Opening grids file: ' + n_grids )
f_grids = netCDF4.Dataset ( n_grids )
print ('Opening masks file: ' + n_masks )
f_masks = netCDF4.Dataset  ( n_masks )
print ('Opening areas file: ' + n_areas )
f_areas = netCDF4.Dataset  ( n_areas )
print ('Opening fracs file: ' + name_frc )   
f_frac  = netCDF4.Dataset  ( name_frc     )

#
msk_src = np.float64(1.0) - f_masks.variables[ grid_src + name_src + '.msk'][:].squeeze()
lon_src = f_grids.variables [grid_src + name_src + '.lon'][:].squeeze()
lat_src = f_grids.variables [grid_src + name_src + '.lat'][:].squeeze()
if myargs.src_int : srf_src = np.ones ( (msk_src.shape), dtype=np.float64 )[:]
else              : srf_src = f_areas.variables [grid_src + name_src + '.srf'][:].squeeze()

#
msk_dst = np.float64(1.0) - f_masks.variables[grid_dst + name_dst + '.msk'][:]
lon_dst = f_grids.variables [grid_dst + name_dst + '.lon'][:].squeeze()
lat_dst = f_grids.variables [grid_dst + name_dst + '.lat'][:].squeeze()
if myargs.dst_int : srf_dst = np.ones ( (msk_dst.shape), dtype=np.float64)[:]
else              : srf_dst = f_areas.variables [grid_dst + name_dst + '.srf'][:].squeeze()

#
if name_src in ( 'ico', 'lmd' ) :
    if myargs.src_nofrac : frc_src = np.ones ( (msk_src.shape), dtype=np.float64)
    else                 : frc_src = f_frac.variables ['OceFrac'][:].squeeze()
        
if name_src == 'orc'     : frc_src = np.ones ( (msk_src.shape), dtype=np.float64)
    
if name_dst in ( 'ico', 'lmd' ) :
    if myargs.dst_nofrac : frc_dst = np.ones ( (msk_dst.shape), dtype=np.float64)
    else                 : frc_dst = f_frac.variables['OceFrac'][:].squeeze()
        
if name_dst == 'orc'     : frc_dst = np.ones ( (msk_dst.shape), dtype=np.float64)

# Periodicity
nperio_src = 0 ; nperio_dst = 0

if Name_src == 'ORCA2.3'  : nperio_src = 4
if Name_src == 'eORCA1.2' : nperio_src = 6
if Name_src == 'eORCA025' : nperio_src = 6

if Name_dst == 'ORCA2.3'  : nperio_dst = 4
if Name_dst == 'eORCA1.2' : nperio_dst = 6
if Name_dst == 'eORCA025' : nperio_dst = 6

# NEMO periodicity.
if Name_src in ('ORCA2.3', 'eORCA1.2', 'eORCA025') : msk_src = nemo.lbc_mask ( msk_src, nperio=nperio_src, cd_type=Grid_src )
if Name_dst in ('ORCA2.3', 'eORCA1.2', 'eORCA025') : msk_dst = nemo.lbc_mask ( msk_dst, nperio=nperio_dst, cd_type=Grid_dst )
    
# Surfaces
area_src = np.sum ( srf_src * msk_src * frc_src )
area_dst = np.sum ( srf_dst * msk_dst * frc_dst )

print ( Name_src )
print ( 'mask: {:12.3} {:12.3} {:12.3}'.format( np.min(msk_src), np.max(msk_src), np.sum(msk_src) ) )
print ( 'frac: {:12.3} {:12.3} {:12.3}'.format( np.min(frc_src), np.max(frc_src), np.sum(frc_src) ) )
print ( 'area: {:12.3} {:12.3} {:12.3}'.format( np.min(srf_src), np.max(srf_src), np.sum(srf_src) ) )

print ( Name_dst )
print ( 'mask: {:12.3} {:12.3} {:12.3}'.format( np.min(msk_dst), np.max(msk_dst), np.sum(msk_dst) ) )
print ( 'frac: {:12.3} {:12.3} {:12.3}'.format( np.min(frc_dst), np.max(frc_dst), np.sum(frc_dst) ) )
print ( 'area: {:12.3} {:12.3} {:12.3}'.format( np.min(srf_dst), np.max(srf_dst), np.sum(srf_dst) ) )
print (' ')
print ( "Surfaces   : %14.6e %14.6e %11.3e"%( area_src, area_dst, (area_src-area_dst)/(area_src+area_dst)*0.5 ) )
print (' ')

##
f_input = netCDF4.Dataset( name_input )
# Loop over fields
for num in np.arange (1, 7) :
    name_src = 'field{:02}_src'.format(num)
    name_dst = 'field{:02}_dst'.format(num)
    print ( '{:02}'.format(num) + ':' + name_src + ':' + name_dst)

    # Reading field
    v_src = f_input [ name_src ][:]
    v_dst = f_input [ name_dst ][:]

    # Compute integrals
    sum_src = np.sum (v_src * srf_src * msk_src * frc_src)
    sum_dst = np.sum (v_dst * srf_dst * msk_dst * frc_dst)
    
    # Compute average
    mean_src = sum_src / area_src
    mean_dst = sum_dst / area_dst

    #
    print ( "Integrals : %14.6e %14.6e %11.3e"%( sum_src , sum_dst , (sum_src - sum_dst)/(sum_src + sum_dst)*0.5 ) )
    print ( "Averages  : %14.6e %14.6e %11.3e"%( mean_src, mean_dst, (mean_src-mean_dst)/(mean_src+mean_dst)*0.5 ) )
    print ( " " )
#=======================================