subpolar gyre mean for all models (as job)

get_subpolar_gyre_all_models.py

+# # Get subpolar gyre mean for each ensemble member
+
+import numpy as np
+import netCDF4 as nc
+import matplotlib.pyplot as plt
+from jupyterthemes import jtplot
+import matplotlib as mpl
+import pandas as pd
+import datetime
+from mpl_toolkits.axes_grid1 import make_axes_locatable
+import cartopy.crs as ccrs
+import cartopy
+import os
+import glob
+import xarray as xr
+
+
+# Region
+lat_max = 61
+lat_min = 46
+lon_max = 360-20
+lon_min = 360-55
+lon_max_W = -20
+lon_min_W = -55
+
+# these seem to be the three latitude/longitude keys
+lat_keys = ['lat', 'latitude', 'nav_lat']
+lon_keys = ['lon', 'longitude', 'nav_lon']
+
+
+# ## Functions
+
+# function to get the longitude / latitude keys from the model
+def get_lon_lat_keys(data, location_keys):
+    str_loc = [key for key in data.coords.keys() if key in location_keys]
+    if len(str_loc) != 1:
+        print('CAUTION: For model ' + model + ' there exist the keys: ' + str(str_loc) + '. We must take care of that separately!')
+    return str_loc[0]
+
+# let all datasets have the same representation of time
+def time_to_year_month(data):
+    return data.assign_coords(dict(time = [np.datetime64(d, 'M') for d in data.time.values]))
+
+# function to extract subpolar gyre mean
+def subpolargyre(data, str_lon, str_lat, print_info=False):
+    mean_dims = [d for d in data.dims if d != "time"]
+    if print_info:
+        print("   The mean is taken over the dimensions ", mean_dims)
+    if np.sum(data[str_lon].values<0) > 0:
+        lon_max_tmp = lon_max_W
+        lon_min_tmp = lon_min_W
+    else:
+        lon_max_tmp = lon_max
+        lon_min_tmp = lon_min
+    spg = (
+        data
+        .where(data[str_lon] >= lon_min_tmp)
+        .where(data[str_lon] <= lon_max_tmp)
+        .where(data[str_lat] >= lat_min)
+        .where(data[str_lat] <= lat_max)
+    )
+    spg = spg.mean(dim = mean_dims)
+    gl = data.mean(dim = mean_dims)
+    spg = spg - gl
+    # spg = spg.groupby('time.year').mean() # that would produce yearly means
+    return spg
+
+# wrapper to get the spg mean for an ensemble member
+def get_subpolar_gyre(model, emember, print_info=False):
+
+    paths = paths_all_models[model][emember]
+
+    if len(paths) != 1:
+        # there is one model where there are 2 paths for the same ensemble member...
+        print("CAUTION: Model '" + model + "', ensemble member '" + emember + "' was not added!")
+        print("         (For this ensemble-member, " + str(len(paths)) + " versions are provided, so the version must be selected manullay!")
+        return None
+    else:
+
+        data = xr.open_mfdataset(paths[0], concat_dim="time").tos
+
+        str_lon = get_lon_lat_keys(data, lon_keys)
+        str_lat = get_lon_lat_keys(data, lat_keys)
+        if print_info:
+            print('   + Longitude: ' + str_lon)
+            print('   + Latitude:  ' + str_lat)
+
+        spg = subpolargyre(data, str_lon, str_lat, print_info)
+        spg = time_to_year_month(spg)
+
+        if print_info: print('  Ensemble-members added:')
+        print('    - ' +emember)
+
+        return spg
+
+
+# ## Get paths as dict of models and ensemble_members
+
+path_base = '/p/tmp/mayayami/SYNDA/data/CMIP6/CMIP/'
+institutes = os.listdir(path_base)
+paths_all_models = dict()
+
+for i in institutes:
+    pi = path_base + i
+    models = os.listdir(pi)
+    for m in models:
+        pim = pi + '/' + m + '/historical/'
+        emembers = os.listdir(pim)
+        tmp = dict()
+        for v in emembers:
+            paths = [x[0] for x in os.walk(pim+v)]
+            data_paths = []
+            for path in paths:
+                if glob.glob("{}/tos_Omon_*.nc".format(path)):
+                    data_paths.append(glob.glob("{}/tos_Omon_*.nc".format(path)))
+            if (len(data_paths)) != 1:
+                print(m)
+                print(data_paths)
+            tmp[v] = data_paths
+        paths_all_models[m] = tmp
+
+
+# # Loop over all models and ensemble members
+
+amoc_index = xr.open_dataarray("CMIP6_amoc_index.nc")
+amoc_index = time_to_year_month(amoc_index)
+
+for model, emembers in paths_all_models.items():
+    print('------------------------------------------')
+    print(model)
+    print_info = True
+    for emember, paths in emembers.items():
+        # the models were only f2 or f3 are available still have the "f1" coordinates in the amoc_index file
+        emember_in_amoc_index = emember[:-1]+"1" 
+
+        spg = get_subpolar_gyre(model, emember, print_info)
+        
+        if not spg is None:
+            print_info = False
+            amoc_index.loc[dict(models=model, ensemble_members=emember_in_amoc_index)] = spg
+            
+amoc_index.to_netcdf('CMIP6_amoc_index.nc')
+
+

get_subpolar_gyre_all_models_job.sh

+#!/bin/bash
+
+
+#SBATCH --job-name=sbpg
+#SBATCH --output=%x-%j.out 
+#SBATCH --ntasks=4
+#SBATCH --cpus-per-task=16
+#SBATCH --account=tipes
+#SBATCH --time=00-24:00:00
+#SBATCH --error=%x-%j.err
+#SBATCH --qos=priority
+
+module load anaconda
+source activate /p/tmp/mayayami/mayaenv
+python get_subpolar_gyre_all_models.py