Merge pull request #26 from xylar/fix_hadgem

Fix HadGEM2-ES preprocessing script

ismip6_ocean_forcing/preprocess/hadgem2_es/process_hadgem2_es_historical.py

@@ -8,24 +8,23 @@
 
 parser = argparse.ArgumentParser(
     description=__doc__, formatter_class=argparse.RawTextHelpFormatter)
-parser.add_argument('-o', dest='out_dir', metavar='DIR',
+parser.add_argument('-o', dest='out_dir', metavar='DIR', default='.',
                     type=str, help='output directory')
 args = parser.parse_args()
 
 
-def compute_yearly_mean(inFileName, outFileName):
+def compute_yearly_mean(inFileNames):
     # crop to below 48 S and take annual mean from monthly data
-    if os.path.exists(outFileName):
-        return
-    print('{} to {}'.format(inFileName, outFileName))
+    print('{}...{} to {}'.format(inFileNames[0], inFileNames[-1], outFileName))
 
-    ds = xarray.open_dataset(inFileName)
+    ds = xarray.open_mfdataset(inFileNames, combine='nested', concat_dim='time')
+    dsFirst = xarray.open_dataset(inFileNames[0])
 
     # crop to Southern Ocean
     ds = ds.isel(lat=slice(0, 43))
 
     for coord in ['lev_bnds', 'lon_bnds', 'lat_bnds']:
-        ds.coords[coord] = ds[coord]
+        ds.coords[coord] = dsFirst[coord]
 
     # annual mean
     with warnings.catch_warnings():
@@ -42,12 +41,12 @@ def compute_yearly_mean(inFileName, outFileName):
     ds.time.attrs['long_name'] = "time"
     ds.time.attrs['standard_name'] = "time"
 
-    timeBounds = numpy.zeros((ds.sizes['time'], ds.sizes['bnds']))
+    timeBounds = numpy.zeros((ds.sizes['time'], 2))
     timeBounds[:, 0] = 365.0*ds.time.values
     timeBounds[:, 1] = 365.0*(ds.time.values+1)
     ds['time_bnds'] = (('time', 'bnds'), timeBounds)
 
-    ds.to_netcdf(outFileName)
+    return ds
 
 
 dates = ['185912-186911',
@@ -66,39 +65,32 @@ def compute_yearly_mean(inFileName, outFileName):
          '198912-199911',
          '199912-200512']
 
-for date in dates:
-    for field in ['so', 'thetao']:
+fileNames = {}
+for field in ['so', 'thetao']:
+    fileNames[field] = []
+    for date in dates:
         inFileName = '{}/{}_Omon_HadGEM2-ES_historical_r1i1p1_{}.nc'.format(
             args.out_dir, field, date)
-
-        outFileName = '{}/{}_annual_HadGEM2-ES_historical_r1i1p1_{}.nc'.format(
-            args.out_dir, field, date)
-
-        compute_yearly_mean(inFileName, outFileName)
+        fileNames[field].append(inFileName)
 
 dates = ['200512-201511']
 
-for date in dates:
-    for field in ['so', 'thetao']:
+for field in ['so', 'thetao']:
+    for date in dates:
         inFileName = '{}/{}_Omon_HadGEM2-ES_rcp85_r1i1p1_{}.nc'.format(
             args.out_dir, field, date)
+        fileNames[field].append(inFileName)
 
-        outFileName = '{}/{}_annual_HadGEM2-ES_rcp85_r1i1p1_{}.nc'.format(
-            args.out_dir, field, date)
-
-        compute_yearly_mean(inFileName, outFileName)
 
 for field in ['so', 'thetao']:
     outFileName = \
         '{}/{}_annual_HadGEM2-ES_rcp85_r1i1p1_186001-201412.nc'.format(
             args.out_dir, field)
-    if not os.path.exists(outFileName):
-        print(outFileName)
-
-        # combine it all into a single data set
-        ds = xarray.open_mfdataset(
-            '{}/{}_annual_HadGEM2-ES_*_r1i1p1_*.nc'.format(
-                args.out_dir, field),
-            combine='nested', concat_dim='time', decode_times=False)
-        ds = ds.isel(time=slice(1, ds.sizes['time']-1))
-        ds.to_netcdf(outFileName)
+    if os.path.exists(outFileName):
+        continue
+
+    ds = compute_yearly_mean(fileNames[field])
+
+    # we don't want 1859 or 2015
+    ds = ds.isel(time=slice(1, ds.sizes['time']-1))
+    ds.to_netcdf(outFileName)