Commit bc8bd6bc authored by Lana B's avatar Lana B
Browse files

clean up: export functions to separate file

("get_subpolar_gyre_functions.py")
parent 57b6920b
......@@ -13,118 +13,43 @@ import cartopy
import os
import glob
import xarray as xr
from get_subpolar_gyre_functions import *
# Region
lat_max = 61
lat_min = 46
lon_max = 360-20
lon_min = 360-55
lon_max_W = -20
lon_min_W = -55
# Regional variables
latitude_maximum = 61
latitude_minimum = 46
longitude_maximum = 360-20
longitude_minimum = 360-55
set_boundaries(
longitude_maximum = longitude_maximum,
longitude_minimum = longitude_minimum,
latitude_maximum = latitude_maximum,
latitude_minimum = latitude_minimum
)
# these seem to be the three latitude/longitude keys
lat_keys = ['lat', 'latitude', 'nav_lat']
lon_keys = ['lon', 'longitude', 'nav_lon']
set_lon_lat_keys(
longitudes=lon_keys,
latitudes =lat_keys
)
# ## 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
get_values()
# # Loop over all models and ensemble members
# Get paths as dict of models and ensemble_members
paths_all_models = get_paths_dict()
# Load shared amoc index file
amoc_index = xr.open_dataarray("CMIP6_amoc_index.nc")
amoc_index = time_to_year_month(amoc_index)
# # Loop over all models and ensemble members
for model, emembers in paths_all_models.items():
print('------------------------------------------')
print(model)
......@@ -133,12 +58,18 @@ for model, emembers in paths_all_models.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)
spg = get_subpolar_gyre(
model = model,
emember = emember,
print_info = print_info,
paths_all_models = paths_all_models)
if not spg is None:
print_info = False
amoc_index.loc[dict(models=model, ensemble_members=emember_in_amoc_index)] = spg
# Save the amoc index file with its new values
amoc_index.to_netcdf('CMIP6_amoc_index.nc')
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 os
import glob
import xarray as xr
# Global variables
# 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 to specify/edit the global variables
# set parameters
def set_boundaries(
longitude_maximum = 360-20,
longitude_minimum = 360-55,
latitude_maximum = 61,
latitude_minimum = 46):
global lon_max, lon_min, lat_max, lat_min, lon_max_W, lon_min_W
lon_max = longitude_maximum
lon_min = longitude_minimum
lat_max = latitude_maximum
lat_min = latitude_minimum
lon_max_W = lon_max-360
lon_min_W = lon_min-360
def set_lon_lat_keys(
longitudes=['lon', 'longitude', 'nav_lon'],
latitudes =['lat', 'latitude', 'nav_lat']):
global lon_keys, lat_keys
lon_keys = longitudes
lat_keys = latitudes
def get_values():
print("lon_min: "+str(lon_min))
print("lon_max: "+str(lon_max))
print("lat_min: "+str(lat_min))
print("lat_max: "+str(lat_max))
print("lon_keys: "+str(lon_keys))
print("lat_keys: "+str(lat_keys))
# 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: There exist the keys: ' + str(str_loc) + '. The first one is used!')
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_all_models=None):
if paths_all_models is None:
paths_all_models = get_paths_dict()
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: The model '" + model + "' provides " + str(len(paths)) + " paths for the ensemble member '" + emember + "'. It is NOT added to CMIP6_amoc_index.nc, this must be done manually!")
return None
if model == 'EC-Earth3-Veg' and emember == 'r10i1p1f1':
print(" Ensemble-member '"+emember+ "' does not provide all years. It also does not occur in Matthew's data, so its not included here.")
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
# make plots of world
def make_nice_plot(ds, str_lon, str_lat):
fig, axis = plt.subplots(
1, 1, subplot_kw=dict(projection=ccrs.Orthographic(central_longitude=320))
)
ds.plot(
x = str_lon,
y = str_lat,
ax = axis,
transform = ccrs.PlateCarree()
)
axis.coastlines(linewidth=2,alpha=0.6)
return fig
# make nice plots of spg area
def make_nice_plot_of_spg(data, str_lon, str_lat, model=None):
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)
)
fig, axis = plt.subplots(
1, 1, subplot_kw=dict(projection=ccrs.Orthographic(central_longitude=320))
)
spg.plot(
x = str_lon,
y = str_lat,
ax = axis,
transform = ccrs.PlateCarree()
)
axis.coastlines(linewidth=2,alpha=0.6)
if not model is None:
fig.suptitle(model)
return fig
def get_paths_dict():
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("Several paths provided for model '"+m+"' and ensemble member '"+v+"'!")
print(data_paths)
tmp[v] = data_paths
paths_all_models[m] = tmp
return paths_all_models
This diff is collapsed.
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