%pip install seaborn
import sys
from netCDF4 import Dataset
import numpy as np
import scipy as sp
import scipy.io as sp_io
import matplotlib.pyplot as plt
import intake_esm, intake
import seaborn as sns
from matplotlib.colors import ListedColormap
import xarray as xr
colours=sns.color_palette('colorblind', 10)
my_cmap = ListedColormap(colours)
kCluster6 = "/home/jovyan/zouberou/DNN4Cli/areamask/GFDL-ESM4-Historical-187201-201112-EnsembleMLP.npy"
#col_url = "https://cmip6-nc.s3.us-east-2.amazonaws.com/esgf-world.json"
col_url = "https://raw.githubusercontent.com/aradhakrishnanGFDL/gfdl-aws-analysis/master/esm-collection-spec-examples/esgf-world.json"
col = intake.open_esm_datastore(col_url)
esmcol_data = col.esmcol_data
def latest_version(cat):
"""
input
cat: esmdatastore
output
esmdatastore with latest DRS versions
"""
latest_cat = cat.df.sort_values(by=['version','path']).drop_duplicates(['temporal_subset','source_id','table_id',
'institution_id','variable_id','member_id',
'grid_label','experiment_id'],keep='last')
return latest_cat
query_Ofx = dict(experiment_id=['abrupt-4xCO2','1pctCO2','historical'],
table_id=['Ofx'],
# member_id=["r1i1p1f1","r1i1p1f2"],
source_id=['GFDL-ESM4'],
grid_label=['gn'],
variable_id=["areacello"])
cat_Ofx = col.search(**query_Ofx)
cat_Ofx_lat = latest_version(cat_Ofx)
cat_Ofx_lat = intake.open_esm_datastore(cat_Ofx_lat,esmcol_data=esmcol_data)
dset_dict_Ofx = cat_Ofx_lat.to_dataset_dict(storage_options=dict(anon=True),cdf_kwargs={'decode_times': True,'chunks': {}})
ds2 = dset_dict_Ofx["CMIP6.NOAA-GFDL.GFDL-ESM4.historical.Ofx"] #xarray dataset object to access Ofx areacello dataset used to calculate the weighted average
--> The keys in the returned dictionary of datasets are constructed as follows: 'project.institution_id.source_id.experiment_id.table_id'
areaGlobal = ds2.areacello
labels2=np.load(kCluster6)
labels2.shape
(7, 576, 720)
#plt.imshow(np.flip(labels2[1][:][:]==5))
Plot all regime data
# First we check that they look OK, and have the same orientation.
#plt.imshow(np.flip(labels2[0][:][:]))
#plt.figure()
def creaCalc(data, areaData):
Get data and areaData into numpy arrays
Calculate area for the data for all the regimes
(This will be an array where only the areas that are covered by the different dynamical regimes are non-nan)
return areaDataRegime0, areaDataRegime1, areaDataRegime2, areaDataRegime3, areaDataRegime4, areaDataRegime5
def extractionAreaCalc(data, area,chunk=None):
"""Code to extract area of different regimes. In principle this is applicable to any field, not just the area."""
'''
input:
data is the labelled regime data
area is the total area
output:
mask matrix, indexed by regime labels 0 to 5
'''
regimeData=np.array(data) #May have to add selection of specific data
areaData=np.array(area) #May have to add selection of specific data
# print(areaData.shape)
#Create dummy variable
Msk=np.zeros((6,chunk, data.shape[0], data.shape[1]))
for chunk in np.arange(0,1): #just one chunk for now
for regime in np.arange(0,6): # Loop over regimes
regimeIndexes=np.where(regimeData!=regime) # Determine where regimes are *NOT* present
areaSelected=np.copy(areaData[0])
areaSelected[regimeIndexes]=-9999 # Set areas there the regime is not present to zero. Can be np.nan too, but I'm not sure what will caouse the least trouble!
Msk[regime][chunk]=areaSelected # Slot into results matrix
# plt.figure() # Can be commented out
# plt.title("Chunk-"+(str)(chunk)+"-- Area of regime "+(str)(regime))
#plt.imshow(np.rot90(Msk[regime][chunk])) #*land,1)) # Can be commented out
#print(areaSelected[(areaSelected>0)])
#print("test",Msk[regime][chunk][Msk[regime][chunk] > 0])
return Msk # Decided to return one matric instead of 6 fields
#todo concatenate chunks in regimeData
areaData=np.array(areaGlobal)
Mask=extractionAreaCalc(labels2[0][:][:], areaGlobal,chunk=1)
labels2[0][:][:].shape
(576, 720)
Mask.shape
(6, 1, 576, 720)
new_labels2=np.zeros((1, labels2[0][:][:].shape[0], labels2[0][:][:].shape[1]))
new_labels2[0] = labels2[0][:]
new_labels2.shape
(1, 576, 720)
da_regimeData = xr.DataArray(
data=new_labels2,
dims=["chunk","y", "x"],
)
ds = xr.Dataset(
data_vars=dict(
regimeData=(["chunk","y","x"], da_regimeData),
areaGlobal=(ds2.areacello),
mask=(["regimeIndex","chunk","y", "x"], Mask),
),
coords=dict(
x=(ds2.x),
y=(ds2.y),
chunk=[0]
)
)
Mask.shape
(6, 1, 576, 720)
ds.regimeData[0].where(ds.regimeData[0] == 1).plot(cmap=my_cmap)
<matplotlib.collections.QuadMesh at 0x7f6990a0d220>
ds.mask[1][0].where(ds.mask[1][0] != -9999).plot(cmap=my_cmap)# mask=(["regimeIndex","chunk","y", "x"], Mask),
<matplotlib.collections.QuadMesh at 0x7f699092b880>
ds.regimeData[0].where(ds.regimeData[0] == 1).plot(cmap=my_cmap)
<matplotlib.collections.QuadMesh at 0x7f69908cba30>
regime_nino3 = ds.regimeData.sel(y = slice(-5,5), x = slice(-150, -90))
regime_nino3[0].plot(cmap=my_cmap)
<matplotlib.collections.QuadMesh at 0x7f69907ed790>
mask_nino3 = ds.mask.sel(y = slice(-5,5), x = slice(-150, -90))
mask_nino3.shape
(6, 1, 39, 120)
#ds.mask[1][0].where(ds.mask[1][0] != -1).sel(y = slice(-5,5), x = slice(-150, -90)).plot()
for regime in np.arange(0,6):
plt.figure()
mask_nino3_regime = ds.mask[regime][0].where(ds.mask[regime][0] != -9999).sel(y = slice(-5,5), x = slice(-150, -90))
mask_nino3_regime.plot()
avgarea = mask_nino3_regime.mean(dim=('x', 'y'))
plt.title("Regime:"+(str)(regime)+"Average area:"+(str)(avgarea))
import matplotlib.pyplot as plt
import cartopy.crs as ccrs
lon_nino3 = ds2["x"][(ds2["x"] > -150) & (ds2["x"] < -90)]
lat_nino3 = ds2["y"][(ds2["y"] > -5) & (ds2["y"] < 5)]
plt.figure(figsize=(6, 3))
plt.title('gfdl-esm4-regimes')
ax2 = plt.axes(projection=ccrs.PlateCarree())
ax2.set_global()
ax2.coastlines()
ax2.contourf(lon_nino3,lat_nino3,regime_nino3[0],cmap='cmo.phase') # didn't use transform, but looks ok...
plt.show()
ax2.set_global()
ax2.coastlines();
lon_nino3.shape
(120,)