import xarray as xr
from matplotlib.pyplot import colorbar
import matplotlib.pyplot as plt
import cartopy.crs as ccrs
import cartopy.feature as cf
import numpy as np
import pandas as pd
# helper function
def make_cmap(colors, position=None):
'''
make_cmap takes a list of tuples which contain RGB values.
The RGB values are [0 to 255]
make_cmap returns a cmap with equally spaced colors.
Arrange your tuples so that the first color is the lowest value for the
colorbar and the last is the highest.
position contains values from 0 to 1 to dictate the location of each color.
'''
import matplotlib as mpl
import numpy as np
import sys
bit_rgb = np.linspace(0,1,256)
if position == None:
position = np.linspace(0,1,len(colors))
else:
if len(position) != len(colors):
sys.exit("position length must be the same as colors")
elif position[0] != 0 or position[-1] != 1:
sys.exit("position must start with 0 and end with 1")
for i in range(len(colors)):
colors[i] = (bit_rgb[colors[i][0]],
bit_rgb[colors[i][1]],
bit_rgb[colors[i][2]])
cdict = {'red':[], 'green':[], 'blue':[]}
for pos, color in zip(position, colors):
cdict['red'].append((pos, color[0], color[0]))
cdict['green'].append((pos, color[1], color[1]))
cdict['blue'].append((pos, color[2], color[2]))
cmap = mpl.colors.LinearSegmentedColormap('my_colormap',cdict,256)
return cmap
thmin1 = xr.open_dataset('thmin1.nc')
thmin2 = xr.open_dataset('thmin2.nc')
thmin3 = xr.open_dataset('thmin3.nc')
thmax1 = xr.open_dataset('thmax1.nc')
thmax2 = xr.open_dataset('thmax2.nc')
thmax3 = xr.open_dataset('thmax3.nc')
ds = xr.open_dataset('eud_2019013000.nc')
mean_vsw = ds.vsw.mean('step') # Mean Volumetric Soil Moisture for all time steps
th1 = mean_vsw[0,:,:] # Mean Volumetric Soil Moisture for Layer 1
th2 = mean_vsw[1,:,:] # Mean Volumetric Soil Moisture for Layer 2
sd1 = ds.sod[0,0,:,:].values
sd2 = ds.sod[0,1,:,:].values - ds.sod[0,0,:,:].values
data=(thmin1.wiltingpoint.values * sd1 + thmin2.wiltingpoint.values * sd2)/(sd1 + sd2)
thmin = xr.DataArray(data,dims=(ds.y.name,ds.x.name),name='thmin')
ds['thmin'] = thmin
data=(thmax2.fieldcapacity.values * sd1 + thmax2.fieldcapacity.values * sd2)/(sd1 + sd2)
thmax = xr.DataArray(data,dims=(ds.y.name,ds.x.name),name='thmax')
ds['thmax'] = thmax
data=((th1 * sd1) + (th2 * sd2))/(sd1 +sd2)
smtot = xr.DataArray(data,dims=(ds.y.name,ds.x.name),name='smtot')
ds['smtot'] = smtot
data=(((smtot - thmin)/(thmax - thmin)))
SM = xr.DataArray(data,dims=(ds.y.name,ds.x.name),name='SM')
ds['SM'] = SM
colors = [(64,0,3), (133,8,3), (249,0,0), (250,106,0), (249,210,0), (189,248,255), (92,138,255), (0,98,255), (0,0,255), (0,0,88)]
### Create an array or list of positions from 0 to 1.
position = [0, 0.2, 0.3,0.4,0.5,0.6,0.7,0.8,0.9,1]
cmap = make_cmap(colors, position=position)
crs = ccrs.LambertAzimuthalEqualArea(central_longitude=10,central_latitude=52,false_easting=4321000,false_northing=3210000)
fig, ax = plt.subplots(1,1,subplot_kw={'projection': crs}, figsize=(15,15) )
ax.gridlines(crs=crs, linestyle="-")
ax.coastlines()
ax.add_feature(cf.BORDERS)
sc = ds["SM"].plot(ax=ax,cmap=cmap,add_colorbar=False)
cbar = plt.colorbar(sc, shrink=.5,)
cbar.set_label("SM") |