Commit 767f270d authored by Lana B's avatar Lana B
Browse files

subpolar gyre mean for all models (as job)

parent de22fa5d
# # 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')
#!/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
Markdown is supported
0% or .
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment