def get_basemap_from_hdf(file_path="", resolution="l"):
"""
:param file_path:
:return: lons(2d), lats(2), basemap - corresponding to the data in the file
"""
h = tb.open_file(file_path)
# Extract 2d longitudes and latitudes
lons = h.getNode("/", "longitude")[:]
lats = h.getNode("/", "latitude")[:]
rotpoletable = h.getNode("/", "rotpole")
assert isinstance(rotpoletable, tb.Table)
params = {}
for row in rotpoletable:
print(row["name"], row["value"])
params[row["name"].decode()] = row["value"].decode() if isinstance(row["value"], bytes) else row["value"]
rotpoletable.close()
basemap = Crcm5ModelDataManager.get_rotpole_basemap_using_lons_lats(
lons2d=lons, lats2d=lats,
lon_1=params["lon1"], lon_2=params["lon2"],
lat_1=params["lat1"], lat_2=params["lat2"], resolution=resolution
)
h.close()
return lons, lats, basemap
def main():
path = "/skynet3_rech1/huziy/from_guillimin/new_outputs/quebec_0.1_crcm5-hcd-rl-intfl_spinup2/Samples/quebec_crcm5-hcd-rl-intfl_197901/pm1979010100_00000000p"
rObj = RPN(path)
sani = rObj.get_first_record_for_name("SANI")
lons, lats = rObj.get_longitudes_and_latitudes_for_the_last_read_rec()
basemap = Crcm5ModelDataManager.get_rotpole_basemap_using_lons_lats(lons2d=lons, lats2d=lats)
x, y = basemap(lons, lats)
sani = np.ma.masked_where(sani < 2, sani)
levels = [2, 3, 4, 5, 6, 7, 8, 10, 15, 20, 25, 30, 40]
cmap = cm.get_cmap("jet", len(levels) - 1)
bn = BoundaryNorm(levels, cmap.N)
fig = plt.figure()
basemap.contourf(x, y, sani, levels=levels, cmap=cmap, norm=bn)
basemap.drawcoastlines()
basemap.colorbar(ticks=levels)
fig.tight_layout()
fig.savefig("soil_anis.jpeg")
pass
def main():
path = "/skynet3_rech1/huziy/from_guillimin/new_outputs/quebec_0.1_crcm5-hcd-rl-intfl_spinup2/Samples/quebec_crcm5-hcd-rl-intfl_197901/pm1979010100_00000000p"
rObj = RPN(path)
sani = rObj.get_first_record_for_name("SANI")
lons, lats = rObj.get_longitudes_and_latitudes_for_the_last_read_rec()
basemap = Crcm5ModelDataManager.get_rotpole_basemap_using_lons_lats(
lons2d=lons, lats2d=lats)
x, y = basemap(lons, lats)
sani = np.ma.masked_where(sani < 2, sani)
levels = [2, 3, 4, 5, 6, 7, 8, 10, 15, 20, 25, 30, 40]
cmap = cm.get_cmap("jet", len(levels) - 1)
bn = BoundaryNorm(levels, cmap.N)
fig = plt.figure()
basemap.contourf(x, y, sani, levels=levels, cmap=cmap, norm=bn)
basemap.drawcoastlines()
basemap.colorbar(ticks=levels)
fig.tight_layout()
fig.savefig("soil_anis.jpeg")
pass
def main():
rpn_folder = "/home/huziy/skynet3_rech1/from_guillimin/new_outputs/quebec_0.1_crcm5-r_spinup"
nc_db_folder = "/home/huziy/skynet3_rech1/crcm_data_ncdb"
start_year = 1979
end_year = 1988
sim_name = "crcm5-r"
#dmManager = Crcm5ModelDataManager(samples_folder_path=rpn_folder, file_name_prefix="dm", all_files_in_samples_folder=True)
pmManager = Crcm5ModelDataManager(samples_folder_path=rpn_folder,
file_name_prefix="pm",
all_files_in_samples_folder=True)
#plot results
assert isinstance(pmManager, Crcm5ModelDataManager)
lons, lats = pmManager.lons2D, pmManager.lats2D
stfl_mask = pmManager.cbf < 0
basemap = Crcm5ModelDataManager.get_rotpole_basemap_using_lons_lats(
lons2d=lons, lats2d=lats)
x, y = basemap(lons, lats)
print(x.shape)
month_dates = [
datetime(year, month, 15) for year in range(start_year, end_year + 1)
for month in range(1, 13)
]
print(len(month_dates), " month of animation ")
nc_data_folder = os.path.join(nc_db_folder, sim_name)
dsDict = {}
var_names = ["STFL", "PR", "TT"]
for i, vName in enumerate(var_names):
path = os.path.join(nc_data_folder, "{0}.nc4".format(vName))
dsDict[vName] = Dataset(path)
aniObj = Animator(var_names, dsDict, basemap, x, y, stfl_mask=stfl_mask)
#fa = animation.FuncAnimation(fig, aniObj.animate, month_dates, interval=50)
temp_folder = "for_anim_{0}".format(sim_name)
if os.path.isdir(temp_folder):
shutil.rmtree(temp_folder)
os.mkdir(temp_folder)
for d in month_dates:
aniObj.animate(d)
aniObj.saveFrame(tmp_folder=temp_folder)
def main():
rpn_folder = "/home/huziy/skynet3_rech1/from_guillimin/new_outputs/quebec_0.1_crcm5-r_spinup"
nc_db_folder = "/home/huziy/skynet3_rech1/crcm_data_ncdb"
start_year = 1979
end_year = 1988
sim_name = "crcm5-r"
#dmManager = Crcm5ModelDataManager(samples_folder_path=rpn_folder, file_name_prefix="dm", all_files_in_samples_folder=True)
pmManager = Crcm5ModelDataManager(samples_folder_path=rpn_folder, file_name_prefix="pm", all_files_in_samples_folder=True)
#plot results
assert isinstance(pmManager, Crcm5ModelDataManager)
lons, lats = pmManager.lons2D, pmManager.lats2D
stfl_mask = pmManager.cbf < 0
basemap = Crcm5ModelDataManager.get_rotpole_basemap_using_lons_lats(
lons2d=lons, lats2d = lats
)
x, y = basemap(lons, lats)
print(x.shape)
month_dates = [ datetime(year, month, 15) for year in range(start_year, end_year + 1) for month in range(1,13) ]
print(len(month_dates), " month of animation ")
nc_data_folder = os.path.join(nc_db_folder, sim_name)
dsDict = {}
var_names = [ "STFL", "PR", "TT"]
for i, vName in enumerate(var_names):
path = os.path.join(nc_data_folder, "{0}.nc4".format(vName))
dsDict[vName] = Dataset(path)
aniObj = Animator(var_names, dsDict, basemap, x, y, stfl_mask = stfl_mask)
#fa = animation.FuncAnimation(fig, aniObj.animate, month_dates, interval=50)
temp_folder = "for_anim_{0}".format(sim_name)
if os.path.isdir(temp_folder):
shutil.rmtree(temp_folder)
os.mkdir(temp_folder)
for d in month_dates:
aniObj.animate(d)
aniObj.saveFrame(tmp_folder=temp_folder)
def main():
AFRIC = 1
QUEBEC = 2
varname = "drainage_density"
region = QUEBEC
if region == QUEBEC:
data_path = "/home/huziy/skynet3_rech1/Netbeans Projects/Java/DDM/directions_with_drainage_density/directions_qc_dx0.1deg_5.nc"
out_path = "qc_{0}_0.1deg.pdf".format(varname)
elif region == AFRIC:
data_path = "/home/huziy/skynet3_rech1/Netbeans Projects/Java/DDM/directions_africa_dx0.44deg.v3.nc"
out_path = "af_{0}_0.44deg.pdf".format(varname)
else:
raise Exception("Unknown region...")
#
ds = Dataset(data_path)
data = ds.variables[varname][20:-20, 20:-20]
lons = ds.variables["lon"][20:-20, 20:-20]
lats = ds.variables["lat"][20:-20, 20:-20]
slope = ds.variables["slope"][20:-20, 20:-20]
fig = plt.figure()
print(data.min(), data.max())
ax = plt.gca()
data = np.ma.masked_where(slope < 0, data)
basemap = Crcm5ModelDataManager.get_rotpole_basemap_using_lons_lats(lons2d=lons, lats2d=lats)
lons[lons > 180] -= 360
x, y = basemap(lons, lats)
data = maskoceans(lons, lats, data, inlands=False)
img = basemap.contourf(x, y, data, cmap=cm.get_cmap("jet", 10))
ax.set_title("Drainage density")
divider = make_axes_locatable(ax)
cax = divider.append_axes("right", "5%", pad="3%")
cb = fig.colorbar(img, cax=cax)
cax.set_title("(km**-1)")
basemap.drawcoastlines(ax=ax)
fig.tight_layout()
fig.savefig(out_path)
def get_omerc_basemap_quebec(lons, lats, lon1=-68, lat1=52, lon2=16.65, lat2=0):
return Crcm5ModelDataManager.get_rotpole_basemap_using_lons_lats(lons, lats,
lon_1=lon1, lat_1=lat1, lon_2=lon2, lat_2=lat2)
def plot_flow_directions_and_basin_boundaries(ax, s, sim_name_to_station_to_model_point,
sim_name_to_manager=None):
assert isinstance(ax, Axes)
assert isinstance(s, Station)
assert isinstance(sim_name_to_station_to_model_point, dict)
mp_list = list(sim_name_to_station_to_model_point.items())[0][1][s]
#selecting only one (the first model point)
mp = mp_list[0]
manager = list(sim_name_to_manager.items())[0][1]
assert isinstance(manager, Crcm5ModelDataManager)
flow_in_mask = manager.get_mask_for_cells_upstream(mp.ix, mp.jy)
lons2d, lats2d = manager.lons2D, manager.lats2D
i_upstream, j_upstream = np.where(flow_in_mask == 1)
nx_total, ny_total = lons2d.shape
imin, imax = np.min(i_upstream), np.max(i_upstream)
jmin, jmax = np.min(j_upstream), np.max(j_upstream)
margin = 8
imin = max(0, imin - margin)
jmin = max(0, jmin - margin)
imax = min(nx_total - 1, imax + margin)
jmax = min(ny_total - 1, jmax + margin)
sub_lons2d = lons2d[imin:imax + 1, jmin:jmax + 1]
sub_lats2d = lats2d[imin:imax + 1, jmin:jmax + 1]
sub_flow_directions = manager.flow_directions[imin:imax + 1, jmin:jmax + 1]
sub_flow_in_mask = flow_in_mask[imin:imax + 1, jmin:jmax + 1]
sub_i_upstream, sub_j_upstream = np.where(sub_flow_in_mask == 1)
basemap = Crcm5ModelDataManager.get_rotpole_basemap_using_lons_lats(
lons2d=sub_lons2d, lats2d=sub_lats2d, resolution="h"
)
#plot all stations
#stations = sim_name_to_station_to_model_point.items()[0][1].keys()
x_list = [the_station.longitude for the_station in (s,)]
y_list = [the_station.latitude for the_station in (s,)]
basemap_big = Crcm5ModelDataManager.get_rotpole_basemap_using_lons_lats(
lons2d=lons2d, lats2d=lats2d
)
x_list, y_list = basemap_big(x_list, y_list)
basemap_big.scatter(x_list, y_list, c="r", s=40, linewidths=0, ax=ax)
basemap_big.drawcoastlines(ax=ax)
basemap_big.drawrivers(ax=ax)
ax.annotate(s.id, xy=basemap(s.longitude, s.latitude), xytext=(3, 3), textcoords='offset points',
font_properties=FontProperties(weight="bold"), bbox=dict(facecolor='w', alpha=1),
ha="left", va="bottom", zorder=2)
x_big, y_big = basemap_big(manager.lons2D, manager.lats2D)
####zoom to the area of interest
#axins = zoomed_inset_axes(ax, 3, loc=2)
displayCoords = ax.transData.transform((x_big[imin, jmin], y_big[imin, jmin]))
x_shift_fig, y_shift_fig = ax.figure.transFigure.inverted().transform(displayCoords)
print("After transData", displayCoords)
print("xshift and yshift", x_shift_fig, y_shift_fig)
#ax.annotate("yuyu", xy = ( 0.733264985153, 0.477182994408), xycoords = "figure fraction" )
rect = [0.1, y_shift_fig + 0.1, 0.4, 0.4]
axins = ax.figure.add_axes(rect)
#assert isinstance(axins, Axes)
x, y = basemap(sub_lons2d, sub_lats2d)
x1d_start = x[sub_flow_in_mask == 1]
y1d_start = y[sub_flow_in_mask == 1]
fld1d = sub_flow_directions[sub_flow_in_mask == 1]
from util import direction_and_value
ishift, jshift = direction_and_value.flowdir_values_to_shift(fld1d)
sub_i_upstream_next = sub_i_upstream + ishift
sub_j_upstream_next = sub_j_upstream + jshift
u = x[sub_i_upstream_next, sub_j_upstream_next] - x1d_start
v = y[sub_i_upstream_next, sub_j_upstream_next] - y1d_start
u2d = np.ma.masked_all_like(x)
v2d = np.ma.masked_all_like(y)
u2d[sub_i_upstream, sub_j_upstream] = u
v2d[sub_i_upstream, sub_j_upstream] = v
basemap.quiver(x, y, u2d, v2d, angles="xy", scale_units="xy", scale=1, ax=axins)
x_list = [the_station.longitude for the_station in (s,)]
y_list = [the_station.latitude for the_station in (s,)]
x_list, y_list = basemap(x_list, y_list)
basemap.scatter(x_list, y_list, c="r", s=40, linewidths=0)
basemap.drawcoastlines(ax=axins)
basemap.drawrivers(ax=axins)
#read waterbase file, and plot only the polygons which contain at least one upstream point
shp_path = "/skynet3_exec1/huziy/Netbeans Projects/Java/DDM/data/shape/waterbase/na_bas_ll_r500m/na_bas_ll_r500m.shp"
c = fiona.open(shp_path)
hits = c.filter(bbox=(sub_lons2d[0, 0], sub_lats2d[0, 0], sub_lons2d[-1, -1], sub_lats2d[-1, -1]))
points = [Point(the_x, the_y) for the_x, the_y in zip(x1d_start, y1d_start)]
selected_polygons = []
for feature in hits:
new_coords = []
old_coords = feature["geometry"]["coordinates"]
#transform to the map coordinates
for ring in old_coords:
x1 = [tup[0] for tup in ring]
y1 = [tup[1] for tup in ring]
x2, y2 = basemap(x1, y1)
new_coords.append(list(zip(x2, y2)))
feature["geometry"]["coordinates"] = new_coords
poly = shape(feature["geometry"])
#print poly, type(poly)
#print feature.keys()
#print feature["properties"]
prep_poly = prep.prep(poly)
hits = list(filter(prep_poly.contains, points))
if len(hits) > 2:
selected_polygons.append(feature["geometry"])
for p in selected_polygons:
axins.add_patch(PolygonPatch(p, fc="none", ec="b", lw=1.5))
zoom_lines_color = "#6600FF"
#set color of the frame
for child in axins.get_children():
if isinstance(child, Spine):
child.set_color(zoom_lines_color)
child.set_linewidth(3)
# draw a bbox of the region of the inset axes in the parent axes and
# connecting lines between the bbox and the inset axes area
mark_inset(ax, axins, loc1=1, loc2=3, fc="none", ec=zoom_lines_color, lw=3)
#basemap.readshapefile(, "basin")
pass
def main():
sim1 = "crcm5-r"
sim2 = "crcm5-hcd-r"
start_date = datetime(1985,1,1)
end_date = datetime(1985,12,31)
sims = [sim1, sim2]
season_months = [4,5,6]
ncdb_path = "/home/huziy/skynet3_rech1/crcm_data_ncdb"
fname_pattern = "{0}_all.nc4"
varname = "TT"
figure = plt.figure()
file_paths = [ os.path.join(ncdb_path, the_sim, fname_pattern.format(varname)) for the_sim in sims ]
dsList = [ Dataset(fPath) for fPath in file_paths ]
lons2d, lats2d, time = dsList[0].variables["lon"][:], dsList[0].variables["lat"][:], dsList[0].variables["time"][:]
basemap = Crcm5ModelDataManager.get_rotpole_basemap_using_lons_lats(lons2d = lons2d, lats2d = lats2d)
varsDict = {}
time = num2date(time, dsList[0].variables["time"].units)
sel_times = np.where([t.month in season_months for t in time])[0]
sim_name_to_mean = {}
for sim, ds in zip( sims, dsList ):
varsDict["{0} ({1})".format(varname, sim)] = ds.variables[varname]
data = ds.variables[varname][sel_times,0,:,:].mean(axis = 0)
sim_name_to_mean[sim] = data
x, y = basemap(lons2d, lats2d)
#plot clickable field
ax = plt.gca()
basemap.contourf(x, y, sim_name_to_mean[sim2] - sim_name_to_mean[sim1])
basemap.colorbar()
basemap.drawcoastlines()
ax.set_title("({0})-({1})".format(sim2, sim1).upper())
TimeSeriesPlotter(ax, basemap, lons2d, lats2d, varsDict, time, start_date, end_date)
plt.show()
pass
def main():
dmManager = Crcm5ModelDataManager(samples_folder_path=rpn_folder, file_name_prefix="dm", all_files_in_samples_folder=True)
pmManager = Crcm5ModelDataManager(samples_folder_path=rpn_folder, file_name_prefix="pm", all_files_in_samples_folder=True)
#export monthly means to netcdf files if necessary
if export_to_nc:
for varname, prefix in zip( field_names, file_name_prefixes ):
manager = None
if prefix == "dm":
manager = dmManager
elif prefix == "pm":
manager = pmManager
level = -1
level_kind = -1
if varname == "TT":
level = 1
level_kind = level_kinds.HYBRID
if varname in ["TRAF", "TDRA"]:
level = 1
level_kind = level_kinds.ARBITRARY
if varname == "STFA": continue
export_monthly_means_to_ncdb(manager, varname, level= level,
level_kind= level_kind, rewrite = True)
#plot results
assert isinstance(pmManager, Crcm5ModelDataManager)
lons, lats = pmManager.lons2D, pmManager.lats2D
basemap = Crcm5ModelDataManager.get_rotpole_basemap_using_lons_lats(
lons2d=lons, lats2d = lats
)
x, y = basemap(lons, lats)
nc_data_folder = os.path.join(nc_db_folder, sim_name)
import matplotlib.pyplot as plt
all_axes = []
ax_to_levels = {}
imgs = []
gs = gridspec.GridSpec(2,3, height_ratios=[1,1], width_ratios=[1,1,1])
fig = plt.figure()
#fig.suptitle("({0} - {1})".format(start_year, end_year))
#plot Temp
varname = "TT"
levels = [-30, -25, -10, -5, -2, 0, 2, 5,10, 15, 20, 25]
cmap = cm.get_cmap("jet", len(levels) - 1)
bn = BoundaryNorm(levels, cmap.N)
ds = Dataset(os.path.join(nc_data_folder, "{0}.nc4".format(varname)))
years = ds.variables["year"][:]
sel = np.where((start_year <= years) & (years <= end_year))[0]
tt = ds.variables[varname][sel,:,:,:].mean(axis = 0).mean(axis = 0)
ds.close()
ax = fig.add_subplot(gs[0,0])
ax.set_title("Temperature (${\\rm ^\circ C}$)")
img = basemap.contourf(x, y, tt, levels = levels, cmap = cmap, norm = bn)
all_axes.append(ax)
imgs.append(img)
ax_to_levels[ax] = levels
#plot precip
varname = "PR"
levels = np.arange(0, 6.5, 0.5)
cmap = cm.get_cmap("jet_r", len(levels) - 1)
bn = BoundaryNorm(levels, cmap.N)
ds = Dataset(os.path.join(nc_data_folder, "{0}.nc4".format(varname)))
years = ds.variables["year"][:]
sel = np.where((start_year <= years) & (years <= end_year))[0]
pr = ds.variables[varname][sel,:,:,:].mean(axis = 0).mean(axis = 0)
convert_factor = 1000.0 * 24 * 60 * 60 #m/s to mm/day
pr *= convert_factor
ds.close()
ax = fig.add_subplot(gs[0,1])
ax.set_title("Precip (mm/day)")
img = basemap.contourf(x, y, pr, levels = levels, cmap = cmap, norm = bn)
all_axes.append(ax)
imgs.append(img)
ax_to_levels[ax] = levels
#plot AH
varname = "AH"
ds = Dataset(os.path.join(nc_data_folder, "{0}.nc4".format(varname)))
years = ds.variables["year"][:]
sel = np.where((start_year <= years) & (years <= end_year))[0]
ah = ds.variables[varname][sel,:,:,:].mean(axis = 0).mean(axis = 0)
ds.close()
levels = np.arange(-60, 160, 20)# np.linspace(au.min(), au.max(), 10)
cmap = cm.get_cmap("jet", len(levels) - 1)
bn = BoundaryNorm(levels, cmap.N)
ax = fig.add_subplot(gs[1,0])
ax.set_title("Sensible heat flux (${\\rm W/m^2}$)")
img = basemap.contourf(x, y, ah, cmap = cmap, norm = bn, levels = levels)
all_axes.append(ax)
imgs.append(img)
ax_to_levels[ax] = levels
#plot AV
varname = "AV"
ds = Dataset(os.path.join(nc_data_folder, "{0}.nc4".format(varname)))
years = ds.variables["year"][:]
sel = np.where((start_year <= years) & (years <= end_year))[0]
av = ds.variables[varname][sel,:,:,:].mean(axis = 0).mean(axis = 0)
ds.close()
coef = 3600* 3
levels = np.arange(-40, 220, 20)
cmap = cm.get_cmap("jet", len(levels) - 1)
bn = BoundaryNorm(levels, cmap.N)
ax = fig.add_subplot(gs[1,1])
ax.set_title("Latent heat flux (${\\rm W/m^2}$)")
img = basemap.contourf(x, y, av * coef, levels = levels, cmap = cmap, norm = bn)
all_axes.append(ax)
imgs.append(img)
ax_to_levels[ax] = levels
#plot stfl
varname = "STFL"
ds = Dataset(os.path.join(nc_data_folder, "{0}.nc4".format(varname)))
years = ds.variables["year"][:]
sel = np.where((start_year <= years) & (years <= end_year))[0]
stfl = ds.variables[varname][sel,:,:,:].mean(axis = 0).mean(axis = 0)
ds.close()
levels = [0,50,100,200,300,500,750,1000, 1500,2000,5000,10000,15000]
stfl = np.ma.masked_where(stfl < 0.01, stfl)
cmap = cm.get_cmap("jet", len(levels) - 1)
bn = BoundaryNorm(levels, cmap.N)
ax = fig.add_subplot(gs[0,2])
ax.set_title("Streamflow (${\\rm m^3/s}$)")
img = basemap.contourf(x, y, stfl, levels = levels, cmap = cmap, norm = bn)
all_axes.append(ax)
imgs.append(img)
ax_to_levels[ax] = levels
sf = ScalarFormatter(useMathText=True)
sf.set_powerlimits([-3,4])
#draw coast lines
for the_ax, the_img in zip(all_axes, imgs):
basemap.drawcoastlines(ax = the_ax)
divider = make_axes_locatable(the_ax)
cax = divider.append_axes("right", "5%", pad="3%")
cb = plt.colorbar(the_img, cax = cax, ticks = ax_to_levels[the_ax])
assert isinstance(cax, Axes)
title = cax.get_title()
fig.tight_layout()
fig.savefig("{0}-mean-annual-fields.pdf".format(sim_name))
def main():
dmManager = Crcm5ModelDataManager(samples_folder_path=rpn_folder,
file_name_prefix="dm",
all_files_in_samples_folder=True)
pmManager = Crcm5ModelDataManager(samples_folder_path=rpn_folder,
file_name_prefix="pm",
all_files_in_samples_folder=True)
#export monthly means to netcdf files if necessary
if export_to_nc:
for varname, prefix in zip(field_names, file_name_prefixes):
manager = None
if prefix == "dm":
manager = dmManager
elif prefix == "pm":
manager = pmManager
level = -1
level_kind = -1
if varname == "TT":
level = 1
level_kind = level_kinds.HYBRID
if varname in ["TRAF", "TDRA"]:
level = 1
level_kind = level_kinds.ARBITRARY
if varname == "STFA": continue
export_monthly_means_to_ncdb(manager,
varname,
level=level,
level_kind=level_kind,
rewrite=True)
#plot results
assert isinstance(pmManager, Crcm5ModelDataManager)
lons, lats = pmManager.lons2D, pmManager.lats2D
basemap = Crcm5ModelDataManager.get_rotpole_basemap_using_lons_lats(
lons2d=lons, lats2d=lats)
x, y = basemap(lons, lats)
nc_data_folder = os.path.join(nc_db_folder, sim_name)
import matplotlib.pyplot as plt
all_axes = []
ax_to_levels = {}
imgs = []
gs = gridspec.GridSpec(2, 3, height_ratios=[1, 1], width_ratios=[1, 1, 1])
fig = plt.figure()
#fig.suptitle("({0} - {1})".format(start_year, end_year))
#plot Temp
varname = "TT"
levels = [-30, -25, -10, -5, -2, 0, 2, 5, 10, 15, 20, 25]
cmap = cm.get_cmap("jet", len(levels) - 1)
bn = BoundaryNorm(levels, cmap.N)
ds = Dataset(os.path.join(nc_data_folder, "{0}.nc4".format(varname)))
years = ds.variables["year"][:]
sel = np.where((start_year <= years) & (years <= end_year))[0]
tt = ds.variables[varname][sel, :, :, :].mean(axis=0).mean(axis=0)
ds.close()
ax = fig.add_subplot(gs[0, 0])
ax.set_title("Temperature (${\\rm ^\circ C}$)")
img = basemap.contourf(x, y, tt, levels=levels, cmap=cmap, norm=bn)
all_axes.append(ax)
imgs.append(img)
ax_to_levels[ax] = levels
#plot precip
varname = "PR"
levels = np.arange(0, 6.5, 0.5)
cmap = cm.get_cmap("jet_r", len(levels) - 1)
bn = BoundaryNorm(levels, cmap.N)
ds = Dataset(os.path.join(nc_data_folder, "{0}.nc4".format(varname)))
years = ds.variables["year"][:]
sel = np.where((start_year <= years) & (years <= end_year))[0]
pr = ds.variables[varname][sel, :, :, :].mean(axis=0).mean(axis=0)
convert_factor = 1000.0 * 24 * 60 * 60 #m/s to mm/day
pr *= convert_factor
ds.close()
ax = fig.add_subplot(gs[0, 1])
ax.set_title("Precip (mm/day)")
img = basemap.contourf(x, y, pr, levels=levels, cmap=cmap, norm=bn)
all_axes.append(ax)
imgs.append(img)
ax_to_levels[ax] = levels
#plot AH
varname = "AH"
ds = Dataset(os.path.join(nc_data_folder, "{0}.nc4".format(varname)))
years = ds.variables["year"][:]
sel = np.where((start_year <= years) & (years <= end_year))[0]
ah = ds.variables[varname][sel, :, :, :].mean(axis=0).mean(axis=0)
ds.close()
levels = np.arange(-60, 160, 20) # np.linspace(au.min(), au.max(), 10)
cmap = cm.get_cmap("jet", len(levels) - 1)
bn = BoundaryNorm(levels, cmap.N)
ax = fig.add_subplot(gs[1, 0])
ax.set_title("Sensible heat flux (${\\rm W/m^2}$)")
img = basemap.contourf(x, y, ah, cmap=cmap, norm=bn, levels=levels)
all_axes.append(ax)
imgs.append(img)
ax_to_levels[ax] = levels
#plot AV
varname = "AV"
ds = Dataset(os.path.join(nc_data_folder, "{0}.nc4".format(varname)))
years = ds.variables["year"][:]
sel = np.where((start_year <= years) & (years <= end_year))[0]
av = ds.variables[varname][sel, :, :, :].mean(axis=0).mean(axis=0)
ds.close()
coef = 3600 * 3
levels = np.arange(-40, 220, 20)
cmap = cm.get_cmap("jet", len(levels) - 1)
bn = BoundaryNorm(levels, cmap.N)
ax = fig.add_subplot(gs[1, 1])
ax.set_title("Latent heat flux (${\\rm W/m^2}$)")
img = basemap.contourf(x, y, av * coef, levels=levels, cmap=cmap, norm=bn)
all_axes.append(ax)
imgs.append(img)
ax_to_levels[ax] = levels
#plot stfl
varname = "STFL"
ds = Dataset(os.path.join(nc_data_folder, "{0}.nc4".format(varname)))
years = ds.variables["year"][:]
sel = np.where((start_year <= years) & (years <= end_year))[0]
stfl = ds.variables[varname][sel, :, :, :].mean(axis=0).mean(axis=0)
ds.close()
levels = [
0, 50, 100, 200, 300, 500, 750, 1000, 1500, 2000, 5000, 10000, 15000
]
stfl = np.ma.masked_where(stfl < 0.01, stfl)
cmap = cm.get_cmap("jet", len(levels) - 1)
bn = BoundaryNorm(levels, cmap.N)
ax = fig.add_subplot(gs[0, 2])
ax.set_title("Streamflow (${\\rm m^3/s}$)")
img = basemap.contourf(x, y, stfl, levels=levels, cmap=cmap, norm=bn)
all_axes.append(ax)
imgs.append(img)
ax_to_levels[ax] = levels
sf = ScalarFormatter(useMathText=True)
sf.set_powerlimits([-3, 4])
#draw coast lines
for the_ax, the_img in zip(all_axes, imgs):
basemap.drawcoastlines(ax=the_ax)
divider = make_axes_locatable(the_ax)
cax = divider.append_axes("right", "5%", pad="3%")
cb = plt.colorbar(the_img, cax=cax, ticks=ax_to_levels[the_ax])
assert isinstance(cax, Axes)
title = cax.get_title()
fig.tight_layout()
fig.savefig("{0}-mean-annual-fields.pdf".format(sim_name))
