diff --git a/ice_ocean_SIS2/OM4_025/diag_table.MOM6 b/ice_ocean_SIS2/OM4_025/diag_table.MOM6 index 97cadbd276..bc14e146c2 100644 --- a/ice_ocean_SIS2/OM4_025/diag_table.MOM6 +++ b/ice_ocean_SIS2/OM4_025/diag_table.MOM6 @@ -71,6 +71,8 @@ # Generally save annuals, and sometimes monthly and daily. "ocean_model_z", "volcello", "volcello", "ocean_annual_z", "all", "mean", "none",2 # Cell measure for 3d data "ocean_model_z", "volcello", "volcello", "ocean_month_z", "all", "mean", "none",2 # Cell measure for 3d data + "ocean_model_rho2", "volcello", "volcello", "ocean_annual_rho2", "all", "mean", "none",2 # Cell measure for 3d data + "ocean_model_rho2", "volcello", "volcello", "ocean_month_rho2", "all", "mean", "none",2 # Cell measure for 3d data "ocean_model", "volcello", "volcello", "ocean_annual", "all", "mean", "none",2 # Cell measure for 3d data #"ocean_model", "volcello", "volcello", "ocean_month", "all", "mean", "none",2 # Cell measure for 3d data "ocean_model", "pbo", "pbo", "ocean_annual", "all", "mean", "none",2 @@ -82,6 +84,8 @@ "ocean_model", "masso", "masso", "ocean_scalar_month", "all", "mean", "none",2 # global sum masscello "ocean_model", "masso", "masso", "ocean_scalar_annual", "all", "mean", "none",2 # global sum masscello "ocean_model", "thkcello", "thkcello", "ocean_annual", "all", "mean", "none",2 # Only needed in native space a static field needs to be provided for CMIP6 + "ocean_model_rho2", "thkcello", "thkcello", "ocean_annual_rho2", "all", "mean", "none",2 # Only needed in native space a static field needs to be provided for CMIP6 + "ocean_model_rho2", "thkcello", "thkcello", "ocean_month_rho2", "all", "mean", "none",2 # Only needed in native space a static field needs to be provided for CMIP6 #"ocean_model", "thkcello", "thkcello", "ocean_month", "all", "mean", "none",2 "ocean_model", "volo", "volo", "ocean_scalar_month", "all", "mean", "none",2 # global sum thkcello "ocean_model", "volo", "volo", "ocean_scalar_annual", "all", "mean", "none",2 # global sum thkcello @@ -163,7 +167,7 @@ #"ocean_model", "vo", "vo", "ocean_month", "all", "mean", "none",2 #"ocean_model", "umo", "umo", "ocean_annual", "all", "mean", "none",2 "ocean_model_z","umo", "umo", "ocean_annual_z", "all", "mean", "none",2 -#"ocean_model_z","umo", "umo", "ocean_month_z", "all", "mean", "none",2 + "ocean_model_z","umo", "umo", "ocean_month_z", "all", "mean", "none",2 #"ocean_model", "vmo", "vmo", "ocean_annual", "all", "mean", "none",2 "ocean_model_z","vmo", "vmo", "ocean_annual_z", "all", "mean", "none",2 "ocean_model_z","vmo", "vmo", "ocean_month_z", "all", "mean", "none",2 @@ -185,16 +189,16 @@ "ocean_model_z","vh", "vh", "ocean_annual_z", "all", "mean", "none",2 "ocean_model_z","T_adx", "T_adx", "ocean_annual_z", "all", "mean", "none",2 "ocean_model_z","T_ady", "T_ady", "ocean_annual_z", "all", "mean", "none",2 - "ocean_model", "T_adx_2d", "T_adx_2d", "ocean_annual", "all", "mean", "none",2 - "ocean_model", "T_ady_2d", "T_ady_2d", "ocean_annual", "all", "mean", "none",2 + "ocean_model", "T_adx_2d", "T_adx_2d", "ocean_month", "all", "mean", "none",2 + "ocean_model", "T_ady_2d", "T_ady_2d", "ocean_month", "all", "mean", "none",2 "ocean_model_z","S_adx", "S_adx", "ocean_annual_z", "all", "mean", "none",2 "ocean_model_z","S_ady", "S_ady", "ocean_annual_z", "all", "mean", "none",2 - "ocean_model", "S_adx_2d", "S_adx_2d", "ocean_annual", "all", "mean", "none",2 - "ocean_model", "S_ady_2d", "S_ady_2d", "ocean_annual", "all", "mean", "none",2 - "ocean_model", "ndiff_tracer_trans_x_2d_T","ndiff_tracer_trans_x_2d_T", "ocean_annual", "all", "mean", "none",2 - "ocean_model", "ndiff_tracer_trans_y_2d_T","ndiff_tracer_trans_y_2d_T", "ocean_annual", "all", "mean", "none",2 - "ocean_model", "ndiff_tracer_trans_x_2d_S","ndiff_tracer_trans_x_2d_S", "ocean_annual", "all", "mean", "none",2 - "ocean_model", "ndiff_tracer_trans_y_2d_S","ndiff_tracer_trans_y_2d_S", "ocean_annual", "all", "mean", "none",2 + "ocean_model", "S_adx_2d", "S_adx_2d", "ocean_month", "all", "mean", "none",2 + "ocean_model", "S_ady_2d", "S_ady_2d", "ocean_month", "all", "mean", "none",2 + "ocean_model", "ndiff_tracer_trans_x_2d_T","ndiff_tracer_trans_x_2d_T", "ocean_month", "all", "mean", "none",2 + "ocean_model", "ndiff_tracer_trans_y_2d_T","ndiff_tracer_trans_y_2d_T", "ocean_month", "all", "mean", "none",2 + "ocean_model", "ndiff_tracer_trans_x_2d_S","ndiff_tracer_trans_x_2d_S", "ocean_month", "all", "mean", "none",2 + "ocean_model", "ndiff_tracer_trans_y_2d_S","ndiff_tracer_trans_y_2d_S", "ocean_month", "all", "mean", "none",2 # Density space diagnostics (not necessarily using CMIP names but needed to generate CMIP output in post-processing) "ocean_model_rho2","umo", "umo", "ocean_annual_rho2", "all", "mean", "none",2 diff --git a/tools/analysis/MOM6_refineDiag.csh b/tools/analysis/MOM6_refineDiag.csh index 25d2794493..5b2fa068b5 100755 --- a/tools/analysis/MOM6_refineDiag.csh +++ b/tools/analysis/MOM6_refineDiag.csh @@ -2,23 +2,23 @@ #------------------------------------------------------------------------------ # MOM6_refineDiag.csh # -# DESCRIPTION: This is a script that is inteded to drive all the +# DESCRIPTION: This is a script that is inteded to drive all the # pre-postprocessing stages of data manipulations on analysis nodes. -# It is intended to be called by FRE at the "refineDiag" stage +# It is intended to be called by FRE at the "refineDiag" stage # which happens just before the components are post processed by frepp. -# To make this happen a path to the (would be) script should appear -# in the tag of the xmls, e.g., +# To make this happen a path to the (would be) script should appear +# in the tag of the xmls, e.g., # # Note that the above script should exist when frepp is called. -# This could be achieved by cloning the mom6 git repo in the section of the setup block -# in the corresponding the gfdl platfrom. E.g., +# This could be achieved by cloning the mom6 git repo in the section of the setup block +# in the corresponding the gfdl platfrom. E.g., # -zCellTop] = 1 return wet +def MOCpsi(vh, vmsk=None): + """Sums 'vh' zonally and cumulatively in the vertical to yield an overturning stream function, psi(y,z).""" + shape = list(vh.shape); shape[-3] += 1 + psi = np.zeros(shape[:-1]) + if len(shape)==3: + for k in range(shape[-3]-1,0,-1): + if vmsk is None: psi[k-1,:] = psi[k,:] - vh[k-1].sum(axis=-1) + else: psi[k-1,:] = psi[k,:] - (vmsk*vh[k-1]).sum(axis=-1) + else: + for n in range(shape[0]): + for k in range(shape[-3]-1,0,-1): + if vmsk is None: psi[n,k-1,:] = psi[n,k,:] - vh[n,k-1].sum(axis=-1) + else: psi[n,k-1,:] = psi[n,k,:] - (vmsk*vh[n,k-1]).sum(axis=-1) + return psi + + +def moc_maskedarray(vh,mask=None): + if mask is not None: + _mask = np.ma.masked_where(np.not_equal(mask,1.),mask) + else: + _mask = 1. + _vh = vh * _mask + _vh_btm = np.ma.expand_dims(_vh[:,-1,:,:]*0.,axis=1) + _vh = np.ma.concatenate((_vh,_vh_btm),axis=1) + _vh = np.ma.sum(_vh,axis=-1) * -1. + _vh = _vh[:,::-1] # flip z-axis so running sum is from ocean floor to surface + _vh = np.ma.cumsum(_vh,axis=1) + _vh = _vh[:,::-1] # flip z-axis back to original order + return _vh + def nearestJI(x, y, xy0): """ Find (j,i) of cell with center nearest to (x0,y0). diff --git a/tools/analysis/refineDiag_ocean_month.py b/tools/analysis/refineDiag_ocean_month.py new file mode 100755 index 0000000000..741e2b9924 --- /dev/null +++ b/tools/analysis/refineDiag_ocean_month.py @@ -0,0 +1,232 @@ +#!/usr/bin/env python + +import argparse +import m6toolbox +import netCDF4 as nc +import numpy as np +import os +import sys +import matplotlib.pyplot as plt + +##-- RefineDiag Script for CMIP6 +## +## 2-D variables we intend to provide: +## +## hfy -> T_ady_2d + ndiff_tracer_trans_y_2d_T * T_ady_2d needs to be converted to Watts (Cp = 3992.) +## ndiff_tracer_trans_y_2d_T already in Watts +## advective term in both 0.25 and 0.5 resolutions +## diffusive term only in 0.5 resolution +## +## hfx -> same recipie as above, expect for x-dimension +## hfbasin -> summed line of hfy in each basin +## +## CMIP variables that will NOT be provided: +## +## hfbasinpmadv, hfbasinpsmadv, hfbasinpmdiff, hfbasinpadv +## (We advect the tracer with the residual mean velocity; individual terms cannot be diagnosed) +## +## htovgyre, htovovrt, sltovgyre, sltovovrt +## (Code for offline calculation not ready.) +## +##-- + +def run(): + parser = argparse.ArgumentParser(description='''CMIP6 RefineDiag Script for OM4''') + parser.add_argument('infile', type=str, help='''Input file''') + parser.add_argument('-b','--basinfile', type=str, default='', required=True, help='''File containing OM4 basin masks''') + parser.add_argument('-o','--outfile', type=str, default=None, help='''Output file name''') + parser.add_argument('-r','--refineDiagDir', type=str, default=None, help='''Path to refineDiagDir defined by FRE workflow)''') + args = parser.parse_args() + main(args) + +def heat_trans_by_basin(x,mask=None,lat=None,minlat=None): + if mask is not None: + varmask = np.sum(mask,axis=-1) + varmask = np.expand_dims(varmask,0) + varmask = np.where(np.equal(varmask,0.),True,False) + else: + mask = 1. + varmask = False + result = np.ma.sum((x*mask),axis=-1) + result.mask = varmask + if (minlat is not None) and (lat is not None): + if len(lat.shape) == 1: + lat = np.tile(lat,0) + lat = np.tile(lat,(len(x),1)) + result = np.ma.masked_where(np.less(lat,minlat),result) + return result + +def main(args): + nc_misval = 1.e20 + #-- Define Regions and their associated masks + # Note: The Atlantic should include other smaller bays/seas that are + # included in the definition used in meridional_overturning.py + + region = np.array(['atlantic_arctic_ocean','indian_pacific_ocean','global_ocean']) + + #-- Read basin masks + f_basin = nc.Dataset(args.basinfile) + basin_code = f_basin.variables['basin'][:] + + atlantic_arctic_mask = basin_code * 0. + atlantic_arctic_mask[(basin_code==2) | (basin_code==4) | (basin_code==6) | (basin_code==7) | (basin_code==8)] = 1. + + indo_pacific_mask = basin_code * 0. + indo_pacific_mask[(basin_code==3) | (basin_code==5)] = 1. + + #-- Read model data + f_in = nc.Dataset(args.infile) + + #-- Read in existing dimensions from history netcdf file + yq = f_in.variables['yq'] + xh = f_in.variables['xh'] + tax = f_in.variables['time'] + + #-- hfy + advective = f_in.variables['T_ady_2d'][:] + if 'ndiff_tracer_trans_y_2d_T' in f_in.variables.keys(): + diffusive = f_in.variables['ndiff_tracer_trans_y_2d_T'][:] + else: + print("Warning: diffusive term 'ndiff_tracer_trans_y_2d_T' not found. Check if this experiment is running with neutral diffusion.") + diffusive = advective * 0. + hfy = advective + diffusive + hfy.long_name = 'Ocean Heat Y Transport' + hfy.units = 'W' + hfy.cell_methods = 'yq:point xh:mean time:mean' + hfy.time_avg_info = 'average_T1,average_T2,average_DT' + hfy.standard_name = 'ocean_heat_y_transport' + + #-- hfx + advective = f_in.variables['T_adx_2d'][:] + if 'ndiff_tracer_trans_x_2d_T' in f_in.variables.keys(): + diffusive = f_in.variables['ndiff_tracer_trans_x_2d_T'][:] + else: + print("Warning: diffusive term 'ndiff_tracer_trans_x_2d_T' not found. Check if this experiment is running with neutral diffusion.") + diffusive = advective * 0. + hfx = advective + diffusive + hfx.long_name = 'Ocean Heat X Transport' + hfx.units = 'W' + hfx.cell_methods = 'yh:mean xq:point time:mean' + hfx.time_avg_info = 'average_T1,average_T2,average_DT' + hfx.standard_name = 'ocean_heat_x_transport' + + #-- hfbasin + hfbasin = np.ma.ones((len(tax),3,len(yq)))*0. + hfbasin[:,0,:] = heat_trans_by_basin(hfy,mask=atlantic_arctic_mask) + hfbasin[:,1,:] = heat_trans_by_basin(hfy,mask=indo_pacific_mask,minlat=-34,lat=yq) + hfbasin[:,2,:] = heat_trans_by_basin(hfy) + hfbasin[hfbasin.mask] = nc_misval + hfbasin = np.ma.array(hfbasin,fill_value=nc_misval) + hfbasin.long_name = 'Northward Ocean Heat Transport' + hfbasin.units = 'W' + hfbasin.coordinates = 'region' + hfbasin.cell_methods = 'yq:point time:mean' + hfbasin.comment = 'Indo-Pacific heat transport begins at 34 S' + hfbasin.time_avg_info = 'average_T1,average_T2,average_DT' + hfbasin.standard_name = 'northward_ocean_heat_transport' + + #-- Read time bounds + nv = f_in.variables['nv'] + average_T1 = f_in.variables['average_T1'] + average_T2 = f_in.variables['average_T2'] + average_DT = f_in.variables['average_DT'] + time_bnds = f_in.variables['time_bnds'] + + #-- Generate output filename + if args.outfile is None: + if hasattr(f_in,'filename'): + args.outfile = f_in.filename + else: + args.outfile = os.path.basename(args.infile) + args.outfile = args.outfile.split('.') + args.outfile[-2] = args.outfile[-2]+'_refined' + args.outfile = '.'.join(args.outfile) + + if args.refineDiagDir is not None: + args.outfile = args.refineDiagDir + '/' + args.outfile + + #-- Write output file + try: + os.remove(args.outfile) + except: + pass + + if os.path.exists(args.outfile): + raise IOError('Output netCDF file already exists.') + exit(1) + + f_out = nc.Dataset(args.outfile, 'w', format='NETCDF3_CLASSIC') + f_out.setncatts(f_in.__dict__) + f_out.filename = os.path.basename(args.outfile) + f_out.delncattr('associated_files') # not needed for these fields + + time_dim = f_out.createDimension('time', size=None) + basin_dim = f_out.createDimension('basin', size=3) + strlen_dim = f_out.createDimension('strlen', size=21) + yq_dim = f_out.createDimension('yq', size=len(yq[:])) + xh_dim = f_out.createDimension('xh', size=len(xh[:])) + nv_dim = f_out.createDimension('nv', size=len(nv[:])) + + time_out = f_out.createVariable('time', np.float64, ('time')) + yq_out = f_out.createVariable('yq', np.float64, ('yq')) + region_out = f_out.createVariable('region', 'c', ('basin', 'strlen')) + xh_out = f_out.createVariable('xh', np.float64, ('xh')) + nv_out = f_out.createVariable('nv', np.float64, ('nv')) + + hfy_out = f_out.createVariable('hfy', np.float32, ('time', 'yq', 'xh'), fill_value=1.e20) + hfy_out.missing_value = 1.e20 + + hfx_out = f_out.createVariable('hfx', np.float32, ('time', 'yq', 'xh'), fill_value=1.e20) + hfx_out.missing_value = 1.e20 + + hfbasin_out = f_out.createVariable('hfbasin', np.float32, ('time', 'basin', 'yq'), fill_value=1.e20) + hfbasin_out.missing_value = 1.e20 + + average_T1_out = f_out.createVariable('average_T1', np.float64, ('time')) + average_T2_out = f_out.createVariable('average_T2', np.float64, ('time')) + average_DT_out = f_out.createVariable('average_DT', np.float64, ('time')) + time_bnds_out = f_out.createVariable('time_bnds', np.float64, ('time', 'nv')) + + time_out.setncatts(tax.__dict__) + yq_out.setncatts(yq.__dict__) + xh_out.setncatts(xh.__dict__) + nv_out.setncatts(nv.__dict__) + + for k in hfy.__dict__.keys(): + if k[0] != '_': hfy_out.setncattr(k,hfy.__dict__[k]) + + for k in hfx.__dict__.keys(): + if k[0] != '_': hfx_out.setncattr(k,hfx.__dict__[k]) + + for k in hfbasin.__dict__.keys(): + if k[0] != '_': hfbasin_out.setncattr(k,hfbasin.__dict__[k]) + + region_out.setncattr('standard_name','region') + + average_T1_out.setncatts(average_T1.__dict__) + average_T2_out.setncatts(average_T2.__dict__) + average_DT_out.setncatts(average_DT.__dict__) + time_bnds_out.setncatts(time_bnds.__dict__) + + time_out[:] = np.array(tax[:]) + yq_out[:] = np.array(yq[:]) + xh_out[:] = np.array(xh[:]) + nv_out[:] = np.array(nv[:]) + + hfy_out[:] = np.ma.array(hfy[:]) + hfx_out[:] = np.ma.array(hfx[:]) + hfbasin_out[:] = np.ma.array(hfbasin[:]) + + average_T1_out[:] = average_T1[:] + average_T2_out[:] = average_T2[:] + average_DT_out[:] = average_DT[:] + time_bnds_out[:] = time_bnds[:] + + region_out[:] = nc.stringtochar(region) + + f_out.close() + + exit(0) + +if __name__ == '__main__': + run() diff --git a/tools/analysis/refineDiag_ocean_month_rho2.py b/tools/analysis/refineDiag_ocean_month_rho2.py new file mode 100755 index 0000000000..84002ec12b --- /dev/null +++ b/tools/analysis/refineDiag_ocean_month_rho2.py @@ -0,0 +1,190 @@ +#!/usr/bin/env python + +import argparse +import m6toolbox +import netCDF4 as nc +import numpy as np +import os +import sys +import matplotlib.pyplot as plt + +##-- RefineDiag Script for CMIP6 +## +## +## Variables we intend to provide in rho-coordinates: +## (potenital density referenced to 2000 m) +## +## msftyrho -> vmo +## msftyrhompa -> vhGM * applies only to 0.5 resolution +## +##-- + +def run(): + parser = argparse.ArgumentParser(description='''CMIP6 RefineDiag Script for OM4''') + parser.add_argument('infile', type=str, help='''Input file''') + parser.add_argument('-b','--basinfile', type=str, default='', required=True, help='''File containing OM4 basin masks''') + parser.add_argument('-o','--outfile', type=str, default=None, help='''Output file name''') + parser.add_argument('-r','--refineDiagDir', type=str, default=None, help='''Path to refineDiagDir defined by FRE workflow)''') + args = parser.parse_args() + main(args) + +def main(args): + #-- Define Regions and their associated masks + # Note: The Atlantic should include other smaller bays/seas that are + # included in the definition used in meridional_overturning.py + + region = np.array(['atlantic_arctic_ocean','indian_pacific_ocean','global_ocean']) + + #-- Read basin masks + f_basin = nc.Dataset(args.basinfile) + basin_code = f_basin.variables['basin'][:] + + atlantic_arctic_mask = basin_code * 0. + atlantic_arctic_mask[(basin_code==2) | (basin_code==4) | (basin_code==6) | (basin_code==7) | (basin_code==8)] = 1. + + indo_pacific_mask = basin_code * 0. + indo_pacific_mask[(basin_code==3) | (basin_code==5)] = 1. + + #-- Read model data + f_in = nc.Dataset(args.infile) + + #-- Read in existing dimensions from history netcdf file + yq = f_in.variables['yq'] + rho2_l = f_in.variables['rho2_l'] + rho2_i = f_in.variables['rho2_i'] + tax = f_in.variables['time'] + + #-- msftyrho + varname = 'vmo' + msftyrho = np.ma.ones((len(tax),3,len(rho2_i),len(yq)))*0. + msftyrho[:,0,:,:] = m6toolbox.moc_maskedarray(f_in.variables[varname][:],mask=atlantic_arctic_mask) + msftyrho[:,1,:,:] = m6toolbox.moc_maskedarray(f_in.variables[varname][:],mask=indo_pacific_mask) + msftyrho[:,2,:,:] = m6toolbox.moc_maskedarray(f_in.variables[varname][:]) + msftyrho[msftyrho.mask] = 1.e20 + msftyrho = np.ma.array(msftyrho,fill_value=1.e20) + msftyrho.long_name = 'Ocean Y Overturning Mass Streamfunction' + msftyrho.units = 'kg s-1' + msftyrho.coordinates = 'region' + msftyrho.cell_methods = 'rho2_i:point yq:point time:mean' + msftyrho.time_avg_info = 'average_T1,average_T2,average_DT' + msftyrho.standard_name = 'ocean_y_overturning_mass_streamfunction' + + #-- msftyrhompa + varname = 'vhGM' + msftyrhompa = np.ma.ones((len(tax),3,len(rho2_i),len(yq)))*0. + msftyrhompa[:,0,:,:] = m6toolbox.moc_maskedarray(f_in.variables[varname][:],mask=atlantic_arctic_mask) + msftyrhompa[:,1,:,:] = m6toolbox.moc_maskedarray(f_in.variables[varname][:],mask=indo_pacific_mask) + msftyrhompa[:,2,:,:] = m6toolbox.moc_maskedarray(f_in.variables[varname][:]) + msftyrhompa[msftyrhompa.mask] = 1.e20 + msftyrhompa = np.ma.array(msftyrhompa,fill_value=1.e20) + msftyrhompa.long_name = 'ocean Y overturning mass streamfunction due to parameterized mesoscale advection' + msftyrhompa.units = 'kg s-1' + msftyrhompa.coordinates = 'region' + msftyrhompa.cell_methods = 'rho2_i:point yq:point time:mean' + msftyrhompa.time_avg_info = 'average_T1,average_T2,average_DT' + msftyrhompa.standard_name = 'ocean_y_overturning_mass_streamfunction_due_to_parameterized_'+\ + 'mesoscale_advection' + + + #-- Read time bounds + nv = f_in.variables['nv'] + average_T1 = f_in.variables['average_T1'] + average_T2 = f_in.variables['average_T2'] + average_DT = f_in.variables['average_DT'] + time_bnds = f_in.variables['time_bnds'] + + #-- Generate output filename + if args.outfile is None: + if hasattr(f_in,'filename'): + args.outfile = f_in.filename + else: + args.outfile = os.path.basename(args.infile) + args.outfile = args.outfile.split('.') + args.outfile[-2] = args.outfile[-2]+'_refined' + args.outfile = '.'.join(args.outfile) + + if args.refineDiagDir is not None: + args.outfile = args.refineDiagDir + '/' + args.outfile + + #-- Write output file + try: + os.remove(args.outfile) + except: + pass + + if os.path.exists(args.outfile): + raise IOError('Output netCDF file already exists.') + exit(1) + + f_out = nc.Dataset(args.outfile, 'w', format='NETCDF3_CLASSIC') + f_out.setncatts(f_in.__dict__) + f_out.filename = os.path.basename(args.outfile) + f_out.delncattr('associated_files') # not needed for these fields + + time_dim = f_out.createDimension('time', size=None) + basin_dim = f_out.createDimension('basin', size=3) + strlen_dim = f_out.createDimension('strlen', size=21) + yq_dim = f_out.createDimension('yq', size=len(yq[:])) + rho2_l_dim = f_out.createDimension('rho2_l', size=len(rho2_l[:])) + rho2_i_dim = f_out.createDimension('rho2_i', size=len(rho2_i[:])) + nv_dim = f_out.createDimension('nv', size=len(nv[:])) + + time_out = f_out.createVariable('time', np.float64, ('time')) + yq_out = f_out.createVariable('yq', np.float64, ('yq')) + region_out = f_out.createVariable('region', 'c', ('basin', 'strlen')) + rho2_l_out = f_out.createVariable('rho2_l', np.float64, ('rho2_l')) + rho2_i_out = f_out.createVariable('rho2_i', np.float64, ('rho2_i')) + nv_out = f_out.createVariable('nv', np.float64, ('nv')) + + msftyrho_out = f_out.createVariable('msftyrho', np.float32, ('time', 'basin', 'rho2_i', 'yq'), fill_value=1.e20) + msftyrho_out.missing_value = 1.e20 + + msftyrhompa_out = f_out.createVariable('msftyrhompa', np.float32, ('time', 'basin', 'rho2_i', 'yq'), fill_value=1.e20) + msftyrhompa_out.missing_value = 1.e20 + + region_out.setncattr('standard_name','region') + + average_T1_out = f_out.createVariable('average_T1', np.float64, ('time')) + average_T2_out = f_out.createVariable('average_T2', np.float64, ('time')) + average_DT_out = f_out.createVariable('average_DT', np.float64, ('time')) + time_bnds_out = f_out.createVariable('time_bnds', np.float64, ('time', 'nv')) + + time_out.setncatts(tax.__dict__) + yq_out.setncatts(yq.__dict__) + rho2_l_out.setncatts(rho2_l.__dict__) + rho2_i_out.setncatts(rho2_i.__dict__) + nv_out.setncatts(nv.__dict__) + + for k in msftyrho.__dict__.keys(): + if k[0] != '_': msftyrho_out.setncattr(k,msftyrho.__dict__[k]) + + for k in msftyrhompa.__dict__.keys(): + if k[0] != '_': msftyrhompa_out.setncattr(k,msftyrhompa.__dict__[k]) + + average_T1_out.setncatts(average_T1.__dict__) + average_T2_out.setncatts(average_T2.__dict__) + average_DT_out.setncatts(average_DT.__dict__) + time_bnds_out.setncatts(time_bnds.__dict__) + + time_out[:] = np.array(tax[:]) + yq_out[:] = np.array(yq[:]) + rho2_l_out[:] = np.array(rho2_l[:]) + rho2_i_out[:] = np.array(rho2_i[:]) + nv_out[:] = np.array(nv[:]) + + msftyrho_out[:] = np.ma.array(msftyrho[:]) + msftyrhompa_out[:] = np.ma.array(msftyrhompa[:]) + + average_T1_out[:] = average_T1[:] + average_T2_out[:] = average_T2[:] + average_DT_out[:] = average_DT[:] + time_bnds_out[:] = time_bnds[:] + + region_out[:] = nc.stringtochar(region) + + f_out.close() + + exit(0) + +if __name__ == '__main__': + run() diff --git a/tools/analysis/refineDiag_ocean_month_z.py b/tools/analysis/refineDiag_ocean_month_z.py new file mode 100755 index 0000000000..609c7a9b4e --- /dev/null +++ b/tools/analysis/refineDiag_ocean_month_z.py @@ -0,0 +1,306 @@ +#!/usr/bin/env python + +import argparse +import m6toolbox +import netCDF4 as nc +import numpy as np +import os +import sys +import matplotlib.pyplot as plt +from verticalvelocity import calc_w_from_convergence +from strait_transport_timeseries import sum_transport_in_straits +##-- RefineDiag Script for CMIP6 +## +## Variables we intend to provide in z-coordinates: +## +## msftyyz -> vmo (ocean_z) * both 0.25 and 0.5 resolutions +## msftyzmpa -> vhGM (ocean_z) * applies only to 0.5 resolution +## msftyzsmpa -> vhml (ocean_z) * both 0.25 and 0.5 resolutions +## +## +## Variables we intend to provide in rho-coordinates: +## (potenital density referenced to 2000 m) +## +## msftyrho -> vmo +## msftyrhompa -> vhGM * applies only to 0.5 resolution +## +## +## 2-D variables we intent to provide: +## +## hfy -> T_ady_2d + ndiff_tracer_trans_y_2d_T * T_ady_2d needs to be converted to Watts (Cp = 3992.) +## ndiff_tracer_trans_y_2d_T already in Watts +## advective term in both 0.25 and 0.5 resolutions +## diffusive term only in 0.5 resolution +## +## hfx -> same recipie as above, expect for x-dimension +## hfbasin -> summed line of hfy in each basin +## +## +## Outstanding issues +## 1.) regirdding of vh, vhGM to rho-corrdinates +## 2.) vhGM and vhML units need to be in kg s-1 +## 2.) save out RHO_0 and Cp somewhere in netCDF files to key off of +## +## +## CMIP variables that will NOT be provided: +## +## hfbasinpmadv, hfbasinpsmadv, hfbasinpmdiff, hfbasinpadv +## (We advect the tracer with the residual mean velocity; individual terms cannot be diagnosed) +## +## htovgyre, htovovrt, sltovgyre, sltovovrt +## (Code for offline calculation not ready.) +## +##-- + +def run(): + parser = argparse.ArgumentParser(description='''CMIP6 RefineDiag Script for OM4''') + parser.add_argument('infile', type=str, help='''Input file''') + parser.add_argument('-b','--basinfile', type=str, default='', required=True, help='''File containing OM4 basin masks''') + parser.add_argument('-s','--straitdir', type=str, default='', required=True, help='''Directory containing output for straits''') + parser.add_argument('-o','--outfile', type=str, default=None, help='''Output file name''') + parser.add_argument('-r','--refineDiagDir', type=str, default=None, help='''Path to refineDiagDir defined by FRE workflow)''') + args = parser.parse_args() + main(args) + +def main(args): + nc_misval = 1.e20 + #-- Define Regions and their associated masks + # Note: The Atlantic should include other smaller bays/seas that are + # included in the definition used in meridional_overturning.py + + region = np.array(['atlantic_arctic_ocean ','indian_pacific_ocean','global_ocean']) + + #-- Read basin masks + f_basin = nc.Dataset(args.basinfile) + basin_code = f_basin.variables['basin'][:] + + atlantic_arctic_mask = basin_code * 0. + atlantic_arctic_mask[(basin_code==2) | (basin_code==4) | (basin_code==6) | (basin_code==7) | (basin_code==8)] = 1. + + indo_pacific_mask = basin_code * 0. + indo_pacific_mask[(basin_code==3) | (basin_code==5)] = 1. + + #-- Read model data + f_in = nc.Dataset(args.infile) + + #-- Read in existing dimensions from history netcdf file + xh = f_in.variables['xh'] + yh = f_in.variables['yh'] + yq = f_in.variables['yq'] + z_l = f_in.variables['z_l'] + z_i = f_in.variables['z_i'] + tax = f_in.variables['time'] + + #-- Note: based on conversations with @adcroft, the overturning should be reported on the interfaces, z_i. + # Also, the nominal latitude is insufficient for the basin-average fields. Based on the methods in + # meridional_overturning.py, the latitude dimension should be: + # + # y = netCDF4.Dataset(cmdLineArgs.gridspec+'/ocean_hgrid.nc').variables['y'][::2,::2] + # yy = y[1:,:].max(axis=-1)+0*z + # + # The quanity 'yy' above is numerically-equivalent to 'yq' + + #-- msftyyz + varname = 'vmo' + msftyyz = np.ma.ones((len(tax),3,len(z_i),len(yq)))*0. + msftyyz[:,0,:,:] = m6toolbox.moc_maskedarray(f_in.variables[varname][:],mask=atlantic_arctic_mask) + msftyyz[:,1,:,:] = m6toolbox.moc_maskedarray(f_in.variables[varname][:],mask=indo_pacific_mask) + msftyyz[:,2,:,:] = m6toolbox.moc_maskedarray(f_in.variables[varname][:]) + msftyyz[msftyyz.mask] = nc_misval + msftyyz = np.ma.array(msftyyz,fill_value=nc_misval) + msftyyz.long_name = 'Ocean Y Overturning Mass Streamfunction' + msftyyz.units = 'kg s-1' + msftyyz.coordinates = 'region' + msftyyz.cell_methods = 'z_i:point yq:point time:mean' + msftyyz.time_avg_info = 'average_T1,average_T2,average_DT' + msftyyz.standard_name = 'ocean_y_overturning_mass_streamfunction' + + #-- msftyzmpa + varname = 'vhGM' + msftyzmpa = np.ma.ones((len(tax),3,len(z_i),len(yq)))*0. + msftyzmpa[:,0,:,:] = m6toolbox.moc_maskedarray(f_in.variables[varname][:],mask=atlantic_arctic_mask) + msftyzmpa[:,1,:,:] = m6toolbox.moc_maskedarray(f_in.variables[varname][:],mask=indo_pacific_mask) + msftyzmpa[:,2,:,:] = m6toolbox.moc_maskedarray(f_in.variables[varname][:]) + msftyzmpa[msftyzmpa.mask] = nc_misval + msftyzmpa = np.ma.array(msftyzmpa,fill_value=nc_misval) + msftyzmpa.long_name = 'ocean Y overturning mass streamfunction due to parameterized mesoscale advection' + msftyzmpa.units = 'kg s-1' + msftyzmpa.coordinates = 'region' + msftyzmpa.cell_methods = 'z_i:point yq:point time:mean' + msftyzmpa.time_avg_info = 'average_T1,average_T2,average_DT' + msftyzmpa.standard_name = 'ocean_y_overturning_mass_streamfunction_due_to_parameterized_'+\ + 'mesoscale_advection' + + #-- msftyzsmpa + varname = 'vhml' + msftyzsmpa = np.ma.ones((len(tax),3,len(z_i),len(yq)))*0. + msftyzsmpa[:,0,:,:] = m6toolbox.moc_maskedarray(f_in.variables[varname][:],mask=atlantic_arctic_mask) + msftyzsmpa[:,1,:,:] = m6toolbox.moc_maskedarray(f_in.variables[varname][:],mask=indo_pacific_mask) + msftyzsmpa[:,2,:,:] = m6toolbox.moc_maskedarray(f_in.variables[varname][:]) + msftyzsmpa[msftyzsmpa.mask] = nc_misval + msftyzsmpa = np.ma.array(msftyzsmpa,fill_value=nc_misval) + msftyzsmpa.long_name = 'ocean Y overturning mass streamfunction due to parameterized submesoscale advection' + msftyzsmpa.units = 'kg s-1' + msftyzsmpa.coordinates = 'region' + msftyzsmpa.cell_methods = 'z_i:point yq:point time:mean' + msftyzsmpa.time_avg_info = 'average_T1,average_T2,average_DT' + msftyzsmpa.standard_name = 'ocean_meridional_overturning_mass_streamfunction_due_to_parameterized_'+\ + 'submesoscale_advection' + + #-- wmo + varname = 'wmo' + wmo = calc_w_from_convergence(f_in.variables['umo'], f_in.variables['vmo']) + wmo[wmo.mask] = nc_misval + wmo = np.ma.array(wmo,fill_value=nc_misval) + wmo.long_name = 'Upward mass transport from resolved and parameterized advective transport' + wmo.units = 'kg s-1' + wmo.cell_methods = 'z_i:point xh:sum yh:sum time:mean' + wmo.time_avg_info = 'average_T1,average_T2,average_DT' + wmo.standard_name = 'upward_ocean_mass_transport' + wmo.cell_measures = 'area:areacello' + + #-- mfo + _, mfo, straits = sum_transport_in_straits(args.straitdir, monthly_average = True) + #mfo[mfo.mask] = nc_misval + mfo = np.ma.array(mfo, fill_value=nc_misval) + strait_names = np.array( [strait.cmor_name for strait in straits] ) + mfo.long_name = 'Sea Water Transport' + mfo.units = 'kg s-1' + mfo.coordinates = 'strait' + mfo.cell_methods = 'time:mean' + mfo.time_avg_info = 'average_T1,average_T2,average_DT' + mfo.standard_name = 'sea_water_transport_across_line' + + #-- Read time bounds + nv = f_in.variables['nv'] + average_T1 = f_in.variables['average_T1'] + average_T2 = f_in.variables['average_T2'] + average_DT = f_in.variables['average_DT'] + time_bnds = f_in.variables['time_bnds'] + + #-- Generate output filename + if args.outfile is None: + if hasattr(f_in,'filename'): + args.outfile = f_in.filename + else: + args.outfile = os.path.basename(args.infile) + args.outfile = args.outfile.split('.') + args.outfile[-2] = args.outfile[-2]+'_refined' + args.outfile = '.'.join(args.outfile) + + if args.refineDiagDir is not None: + args.outfile = args.refineDiagDir + '/' + args.outfile + + #-- Write output file + try: + os.remove(args.outfile) + except: + pass + + if os.path.exists(args.outfile): + raise IOError('Output netCDF file already exists.') + exit(1) + + f_out = nc.Dataset(args.outfile, 'w', format='NETCDF3_CLASSIC') + f_out.setncatts(f_in.__dict__) + f_out.filename = os.path.basename(args.outfile) + + time_dim = f_out.createDimension('time', size=None) + basin_dim = f_out.createDimension('basin', size=3) + strait_dim = f_out.createDimension('strait', size=len(straits)) + strlen_dim = f_out.createDimension('strlen', size=31) + xh_dim = f_out.createDimension('xh', size=len(xh[:])) + yh_dim = f_out.createDimension('yh', size=len(yh[:])) + yq_dim = f_out.createDimension('yq', size=len(yq[:])) + z_l_dim = f_out.createDimension('z_l', size=len(z_l[:])) + z_i_dim = f_out.createDimension('z_i', size=len(z_i[:])) + nv_dim = f_out.createDimension('nv', size=len(nv[:])) + + time_out = f_out.createVariable('time', np.float64, ('time')) + xh_out = f_out.createVariable('xh', np.float64, ('xh')) + yh_out = f_out.createVariable('yh', np.float64, ('yh')) + yq_out = f_out.createVariable('yq', np.float64, ('yq')) + region_out = f_out.createVariable('region', 'c', ('basin', 'strlen')) + strait_out = f_out.createVariable('strait', 'c', ('strait', 'strlen')) + z_l_out = f_out.createVariable('z_l', np.float64, ('z_l')) + z_i_out = f_out.createVariable('z_i', np.float64, ('z_i')) + nv_out = f_out.createVariable('nv', np.float64, ('nv')) + + msftyyz_out = f_out.createVariable('msftyyz', np.float32, ('time', 'basin', 'z_i', 'yq'), fill_value=nc_misval) + msftyyz_out.missing_value = nc_misval + + msftyzsmpa_out = f_out.createVariable('msftyzsmpa', np.float32, ('time', 'basin', 'z_i', 'yq'), fill_value=nc_misval) + msftyzsmpa_out.missing_value = nc_misval + + msftyzmpa_out = f_out.createVariable('msftyzmpa', np.float32, ('time', 'basin', 'z_i', 'yq'), fill_value=nc_misval) + msftyzmpa_out.missing_value = nc_misval + + wmo_out = f_out.createVariable('wmo', np.float32, ('time', 'z_i', 'yh', 'xh'), fill_value=nc_misval) + wmo_out.missing_value = nc_misval + + mfo_out = f_out.createVariable('mfo', np.float32, ('time', 'strait'), fill_value=nc_misval) + mfo_out.missing_value = nc_misval + + average_T1_out = f_out.createVariable('average_T1', np.float64, ('time')) + average_T2_out = f_out.createVariable('average_T2', np.float64, ('time')) + average_DT_out = f_out.createVariable('average_DT', np.float64, ('time')) + time_bnds_out = f_out.createVariable('time_bnds', np.float64, ('time', 'nv')) + + time_out.setncatts(tax.__dict__) + xh_out.setncatts(xh.__dict__) + yh_out.setncatts(yh.__dict__) + yq_out.setncatts(yq.__dict__) + z_l_out.setncatts(z_l.__dict__) + z_i_out.setncatts(z_i.__dict__) + nv_out.setncatts(nv.__dict__) + + for k in msftyyz.__dict__.keys(): + if k[0] != '_': msftyyz_out.setncattr(k,msftyyz.__dict__[k]) + + for k in msftyzsmpa.__dict__.keys(): + if k[0] != '_': msftyzsmpa_out.setncattr(k,msftyzsmpa.__dict__[k]) + + for k in msftyzmpa.__dict__.keys(): + if k[0] != '_': msftyzmpa_out.setncattr(k,msftyzmpa.__dict__[k]) + + for k in mfo.__dict__.keys(): + if k[0] != '_': mfo_out.setncattr(k,mfo.__dict__[k]) + + for k in wmo.__dict__.keys(): + if k[0] != '_': wmo_out.setncattr(k,wmo.__dict__[k]) + + region_out.setncattr('standard_name','region') + strait_out.setncattr('standard_name','region') + + average_T1_out.setncatts(average_T1.__dict__) + average_T2_out.setncatts(average_T2.__dict__) + average_DT_out.setncatts(average_DT.__dict__) + time_bnds_out.setncatts(time_bnds.__dict__) + + time_out[:] = np.array(tax[:]) + xh_out[:] = np.array(xh[:]) + yh_out[:] = np.array(yh[:]) + yq_out[:] = np.array(yq[:]) + z_l_out[:] = np.array(z_l[:]) + z_i_out[:] = np.array(z_i[:]) + nv_out[:] = np.array(nv[:]) + + msftyyz_out[:] = np.ma.array(msftyyz[:]) + msftyzsmpa_out[:] = np.ma.array(msftyzsmpa[:]) + msftyzmpa_out[:] = np.ma.array(msftyzmpa[:]) + wmo_out[:] = np.ma.array(wmo[:]) + mfo_out[:] = np.array(mfo[:]) + + average_T1_out[:] = average_T1[:] + average_T2_out[:] = average_T2[:] + average_DT_out[:] = average_DT[:] + time_bnds_out[:] = time_bnds[:] + + region_out[:] = nc.stringtochar(region) + strait_out[:] = nc.stringtochar(strait_names) + f_out.close() + + exit(0) + +if __name__ == '__main__': + run() diff --git a/tools/analysis/strait_transport_timeseries.py b/tools/analysis/strait_transport_timeseries.py new file mode 100644 index 0000000000..1fe65698c9 --- /dev/null +++ b/tools/analysis/strait_transport_timeseries.py @@ -0,0 +1,124 @@ +#!/usr/bin/env python +# This script sums the transport for all the straits requested for CMIP6/OMIP that are output in MOM6. If requested, it also +# average the output from daily to monthly +from netCDF4 import Dataset +import numpy as np +from fnmatch import filter as fn_filter +from os import listdir + +def sum_transport_in_straits(runpath, monthly_average = False): + + strait = set_strait_info() + nstraits = len(strait) + # Calculate transport in each of the straits + for sidx in range(0,nstraits): + strait[sidx].transport = 0. + if strait[sidx].is_zonal and strait[sidx].is_meridional: + u_file = fn_filter(listdir(runpath), '*' + strait[sidx].mom6_name + '_U.nc') + v_file = fn_filter(listdir(runpath), '*' + strait[sidx].mom6_name + '_V.nc') + if (len(u_file)==0 and len(v_file)==0): + print("Warning: File not found for %s" % strait[sidx].mom6_name) + continue + u_vargroup = Dataset(runpath+'/'+u_file[0]).variables + v_vargroup = Dataset(runpath+'/'+v_file[0]).variables + elif strait[sidx].is_zonal: + v_file = fn_filter(listdir(runpath), '*' + strait[sidx].mom6_name + '*.nc') + if (len(v_file)==0): + print("Warning: File not found for %s" % strait[sidx].mom6_name) + continue + v_vargroup = Dataset(runpath+'/'+v_file[0]).variables + elif strait[sidx].is_meridional: + u_file = fn_filter(listdir(runpath), '*' + strait[sidx].mom6_name + '*.nc') + if (len(u_file)==0): + print("Warning: File not found for %s" % strait[sidx].mom6_name) + continue + u_vargroup = Dataset(runpath+'/'+u_file[0]).variables + + # Need to find the first interface deeper than or equal to the requested z-limit. If deeper, then we'll need to scale the + # bottommost part of the column + if strait[sidx].is_zonal: + strait[sidx].time = v_vargroup['time'][:] + if strait[sidx].zlim > 0.: + z_i = v_vargroup['z_i'][:] + zidx = np.sum(z_i strait[sidx].zlim: + frac = min(1., (z_i[zidx] - strait[sidx].zlim)/(z_i[zidx]-z_i[zidx-1])) + vmo[:,-1,:] = vmo[:,-1,:]*frac + else: + vmo = v_vargroup['vmo'][:,:,:,:] + strait[sidx].transport += vmo.sum(axis=(1,2,3)) + Dataset(runpath+'/'+v_file[0]).close() + + if strait[sidx].is_meridional: + strait[sidx].time = u_vargroup['time'][:] + if strait[sidx].zlim > 0.: + z_i = u_vargroup['z_i'][:] + zidx = np.sum(z_i strait[sidx].zlim: + frac = min(1., (z_i[zidx] - strait[sidx].zlim)/(z_i[zidx]-z_i[zidx-1])) + umo[:,-1,:] = umo[:,-1,:]*frac + else: + umo = u_vargroup['umo'][:,:,:,:] + strait[sidx].transport += umo.sum(axis=(1,2,3)) + Dataset(runpath+'/'+u_file[0]).close() + if monthly_average: + strait[sidx].transport = make_monthly_averages(strait[sidx].transport) + strait[sidx].time = make_monthly_averages(strait[sidx].time) + ntime = strait[sidx].time.size + time = strait[sidx].time + + transport_array = np.zeros((ntime, nstraits)) + for sidx in range(0,nstraits): + transport_array[:,sidx] = strait[sidx].transport + return time, transport_array, strait + +def make_monthly_averages(data): + # This fundamentally assumes that the output is only one year long, any longer and we'd have to know the actual year to properly + # handle leap years + idx = daily_indices(data.size == 366) + idx_r = [range(0,idx[0])] ; + [idx_r.append(range(idx[t-1],idx[t])) for t in range(1,12)] + return np.array([np.mean(data[idx_r[t]]) for t in range(0,12)]) + +def daily_indices(leap = False): + days_in_months = np.array([31, 28, 31, 30, 31, 30, 31, 31, 30 ,31, 30 ,31]) + if leap: # Change number of days in Feb if leapyear + days_in_months[1] = 29 + return days_in_months.cumsum() + + +def set_strait_info(): + + strait = [] + # The standard list of straits from CMIP/OMIP request + strait.append(strait_type("barents_opening", "Barents_opening", is_meridional = True)) + strait.append(strait_type("bering_strait", "Bering_Strait" , is_zonal = True)) + strait.append(strait_type("carribean_windward_passage", "Windward_Passage", is_zonal = True)) + strait.append(strait_type("davis_strait", "Davis_Strait", is_zonal = True)) + strait.append(strait_type("denmark_strait", "Denmark_Strait", is_zonal = True)) + strait.append(strait_type("drake_passage", "Drake_Passage", is_meridional = True)) + strait.append(strait_type("english_channel", "English_Channel", is_meridional = True)) + strait.append(strait_type("faroe_scotland_channel", "Faroe_Scotland", is_meridional = True)) + strait.append(strait_type("florida_bahamas_strait", "Florida_Bahamas", is_zonal = True)) + strait.append(strait_type("fram_strait", "Fram_Strait", is_zonal = True)) + strait.append(strait_type("gibraltar_strait", "Gibraltar_Strait", is_meridional = True)) + strait.append(strait_type("iceland_faroe_channel", "Iceland_Faroe", is_meridional = True, is_zonal = True)) + strait.append(strait_type("mozambique_channel", "Mozambique_Channel", is_zonal = True)) + strait.append(strait_type("pacific_equatorial_undercurrent", "Pacific_undercurrent", is_meridional = True, zlim = 350)) + strait.append(strait_type("taiwan_and_luzon_straits", "Taiwan_Luzon", is_meridional = True)) + strait.append(strait_type("agulhas_section", "Agulhas_section", is_meridional = True)) + strait.append(strait_type("indonesian_throughflow", "Indonesian_Throughflow", is_zonal = True)) + + return strait + +class strait_type: + def __init__(self, cmor_name, mom6_name, is_meridional = False, is_zonal = False, zlim = -1): + self.cmor_name = cmor_name + self.mom6_name = mom6_name + self.is_meridional = is_meridional + self.is_zonal = is_zonal + self.zlim = zlim + self.time = None + self.transport = None + + diff --git a/tools/analysis/verticalvelocity.py b/tools/analysis/verticalvelocity.py new file mode 100755 index 0000000000..451c1c4c07 --- /dev/null +++ b/tools/analysis/verticalvelocity.py @@ -0,0 +1,82 @@ +#!/usr/bin/env python +# Estimate vertical mass transport (wmo) by the divergence of the horizontal mass transports. The vertical mass transport across an +# interface is the cumulative integral starting from the bottom of a water column. The sign convention is w>0 is an upward transport, +# (i.e., towards the surface of the ocean). By this convention, then div(u,v)<0 implies a negative (downward) transport and vice versa. +# Written by Andrew Shao (andrew.shao@noaa.gov) 29 September 2017 +import numpy as np + +def calc_w_from_convergence(u_var, v_var, wrapx = True, wrapy = False): + + tmax = u_var.shape[0] + + ntime, nk, nlat, nlon = u_var.shape + w = np.ma.zeros( (ntime, nk+1, nlat, nlon) ) + # Work timelevel by timelevel + for tidx in range(0,tmax): + # Get and process the u component + u_dat = u_var[tidx,:,:,:] + h_mask = np.logical_or(np.ma.getmask(u_dat), np.ma.getmask(np.roll(u_dat,1,axis=-1))) + u_dat = u_dat.filled(0.) + + # Get and process the v component + v_dat = v_var[tidx,:,:,:] + h_mask = np.logical_or(h_mask,np.ma.getmask(v_dat)) + h_mask = np.logical_or(h_mask,np.ma.getmask(np.roll(v_dat,1,axis=-2))) + v_dat = v_dat.filled(0.) + + # Order of subtraction based on upwind sign convention and desire for w>0 to correspond with upwards velocity + w[tidx,:-1,:,:] += np.roll(u_dat,1,axis=-1)-u_dat + if not wrapx: # If not wrapping, then convergence on westernmost side is simply so subtract back the rolled value + w[tidx,:-1,:,0] += -u_dat[:-1,:,-1] + w[tidx,:-1,:,:] += np.roll(v_dat,1,axis=-2)-v_dat + if not wrapy: # If not wrapping, convergence on westernmost side is v + w[tidx,:-1,0,:] += -v_dat[:,-1,:] + w[tidx,-1,:,:] = 0. + # Do a double-flip so that we integrate from the bottom + w[tidx,:-1,:,:] = w[tidx,-2::-1,:,:].cumsum(axis=0)[::-1,:,:] + # Mask if any of u[i-1], u[i], v[j-1], v[j] are not masked + w[tidx,:-1,:,:] = np.ma.masked_where(h_mask, w[tidx,:-1,:,:]) + # Bottom should always be zero, mask applied wherever the top interface is a valid value + w[tidx,-1,:,:] = np.ma.masked_where(h_mask[-2,:,:], w[tidx,-1,:,:]) + + return w + +if __name__ == "__main__": + try: import argparse + except: raise Exception('This version of python is not new enough. python 2.7 or newer is required.') + import netCDF4 + + parser = argparse.ArgumentParser(description= + '''Script for calculating vertical velocity from divergence of horizontal mass transports.''') + parser.add_argument('infile', type=str, help='''Daily file containing ssu,ssv.''') + parser.add_argument('--uname', type=str, default='umo', help='''Name of the u-component of mass transport''') + parser.add_argument('--vname', type=str, default='vmo', help='''Name of the v-component of mass transport''') + parser.add_argument('--wrapx', type=bool, default=True, help='''True if the x-component is reentrant''') + parser.add_argument('--wrapy', type=bool, default=False, help='''True if the y-component is reentrant''') + parser.add_argument('--gridspec', type=str, default ='', + help='''File containing variables wet,areacello. Usually the ocean_static.nc from diag_table.''') + parser.add_argument('--plot', type=bool, default=False, help='''Plot w at the 8th interface (kind of random)''') + cmdLineArgs = parser.parse_args() + + # Check for presence of variables + rootGroup = netCDF4.Dataset( cmdLineArgs.infile ) + if cmdLineArgs.uname not in rootGroup.variables: + raise Exception('Could not find "%s" in file "%s"' % (cmdLineArgs.uname, cmdLineArgs.infile)) + if cmdLineArgs.vname not in rootGroup.variables: + raise Exception('Could not find "%s" in file "%s"' % (cmdLineArgs.vname, cmdLineArgs.infile)) + + u_var = netCDF4.Dataset(cmdLineArgs.infile).variables[cmdLineArgs.uname] + v_var = netCDF4.Dataset(cmdLineArgs.infile).variables[cmdLineArgs.vname] + + w = calc_w_from_convergence( u_var, v_var, + cmdLineArgs.wrapx, cmdLineArgs.wrapy ) + if len(cmdLineArgs.gridspec) > 0: + area = netCDF4.Dataset(cmdLineArgs.gridspec).variables['areacello'][:,:] + w = w/(area*1035.0) + # Plot if requested + if cmdLineArgs.plot: + import matplotlib.pyplot as plt + plt.pcolormesh(w[0,7,:,:],vmin=-2e-5,vmax=2e-5,cmap=plt.cm.coolwarm) + plt.colorbar() + plt.show() +