More robust refineDiag scripts

tools/analysis/refineDiag_ocean_month.py

@@ -18,13 +18,13 @@
 ##                                                       diffusive term only in 0.5 resolution
 ##
 ##     hfx  -> same recipie as above, expect for x-dimension
-##     hfbasin -> summed line of hfy in each basin 
-## 
+##     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.)
 ##
@@ -59,11 +59,11 @@ def heat_trans_by_basin(x,mask=None,lat=None,minlat=None):
 def main(args):
     nc_misval = 1.e20
     #-- Define Regions and their associated masks
-    #   Note: The Atlantic should include other smaller bays/seas that are 
+    #   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'][:]
@@ -73,160 +73,177 @@ def main(args):
 
     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'][:]
+    if 'T_ady_2d' in f_in.variables.keys():
+      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'
+      do_hfy = True
     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'
+      do_hfy = False
 
     #-- 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'][:]
+    if 'T_adx_2d' in f_in.variables.keys():
+      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'
+      do_hfx = True
     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 
+      do_hfx = False
+
+    #-- hfbasin
+    if do_hfy:
+      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'
+      do_hfbasin = True
+    else:
+      do_hfbasin = False
+
+    #-- 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 any([do_hfx,do_hfy,do_hfbasin]):
+      #-- 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'))
+
+      if do_hfy:
+        hfy_out = f_out.createVariable('hfy', np.float32, ('time', 'yq', 'xh'), fill_value=1.e20)
+        hfy_out.missing_value = 1.e20
+        for k in hfy.__dict__.keys():
+          if k[0] != '_': hfy_out.setncattr(k,hfy.__dict__[k])
+
+      if do_hfx:
+        hfx_out = f_out.createVariable('hfx', np.float32, ('time', 'yq', 'xh'), fill_value=1.e20)
+        hfx_out.missing_value = 1.e20
+        for k in hfx.__dict__.keys():
+          if k[0] != '_': hfx_out.setncattr(k,hfx.__dict__[k])
+
+      if do_hfbasin:
+        hfbasin_out = f_out.createVariable('hfbasin', np.float32, ('time', 'basin', 'yq'), fill_value=1.e20)
+        hfbasin_out.missing_value = 1.e20
+        for k in hfbasin.__dict__.keys():
+          if k[0] != '_': hfbasin_out.setncattr(k,hfbasin.__dict__[k])
+
+      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__)
+
+      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[:])
+
+      if do_hfy:     hfy_out[:] = np.ma.array(hfy[:])
+      if do_hfx:     hfx_out[:] = np.ma.array(hfx[:])
+      if do_hfbasin: 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)
+
+    else:
+      print('RefineDiag for ocean_month yielded no output.')
+      exit(1)
+
 
 if __name__ == '__main__':
   run()