def __init__(self, file_path = "", var_name = "",
bathymetry_path = "/skynet3_rech1/huziy/NEMO_OFFICIAL/dev_v3_4_STABLE_2012/NEMOGCM/CONFIG/GLK_LIM3_Michigan/EXP00/bathy_meter.nc"):
"""
:param file_path:
:param var_name:
:param bathymetry_path: used to mask land points
"""
self.current_time_frame = -1
self.var_name = var_name
self.cube = iris.load_cube(file_path, constraint=iris.Constraint(cube_func=lambda c: c.var_name == var_name))
self.lons, self.lats = cartography.get_xy_grids(self.cube)
lons2d_gl, lats2d_gl = nemo_commons.get_2d_lons_lats_from_nemo(path=bathymetry_path)
mask_gl = nemo_commons.get_mask(path=bathymetry_path)
xs, ys, zs = lat_lon.lon_lat_to_cartesian(lons2d_gl.flatten(), lats2d_gl.flatten())
xt, yt, zt = lat_lon.lon_lat_to_cartesian(self.lons.flatten(), self.lats.flatten())
tree = cKDTree(list(zip(xs, ys, zs)))
dists, indices = tree.query(list(zip(xt, yt, zt)))
self.mask = mask_gl.flatten()[indices].reshape(self.lons.shape)
self.nt = self.cube.shape[0]
assert isinstance(self.cube, Cube)
print(self.nt)
def main_for_lake(bathy_path = "",
basemap = None,
manager_u = None, manager_v=None,
scalar_manager=None, img_dir=""):
"""
:param bathy_path: file used for land sea mask
:param basemap:
:param lons:
:param lats:
:param manager_u:
:param manager_v:
:return:
"""
the_mask = nemo_commons.get_mask(bathy_path)
print(the_mask.shape)
scalar_levels = np.arange(-25, 30, 5)
for frame in range(manager_u.get_total_number_of_time_frames()):
fig = plt.figure()
# bounds = [0, 0.02, 0.05, 0.08, 0.09, 0.1]
bounds = np.arange(-20, 21, 1)
bn = BoundaryNorm(boundaries=bounds, ncolors=len(bounds) - 1)
# img = plt.pcolormesh(data, vmin = bounds[0], vmax = bounds[-1], cmap = cm.get_cmap("jet", len(bounds) - 1), norm = bn)
b_local, lons, lats = nemo_commons.get_basemap_and_coordinates_from_file(path=bathy_path, resolution="i")
x, y = basemap(lons, lats)
if scalar_manager is not None:
data = scalar_manager.get_next_time_frame_data()
print(data.shape)
img = basemap.pcolormesh(x, y, np.ma.masked_where(~the_mask, data),
vmin=scalar_levels[0], vmax=scalar_levels[-1], zorder=-6)
basemap.colorbar()
u, v = manager_u.get_next_time_frame_data(), manager_v.get_next_time_frame_data()
u = np.ma.masked_where(~the_mask, u)
v = np.ma.masked_where(~the_mask, v)
# qp = basemap.quiver(x, y, u, v, scale=1.5)
c = np.sqrt(u ** 2 + v ** 2)
qp = basemap.streamplot(x, y, u, v, color="k", density=(5, 5), linewidth=3 * c / c.max())
basemap.drawcoastlines()
basemap.drawmapboundary(fill_color="gray")
plt.title(str(manager_u.get_current_time()).split()[0])
print(str(manager_u.get_current_time()).split()[0])
fig.savefig(os.path.join(img_dir, "{0:08d}.png".format(frame)))
plt.close(fig)
def __init__(
self,
file_path="",
var_name="",
bathymetry_path="/skynet3_rech1/huziy/NEMO_OFFICIAL/dev_v3_4_STABLE_2012/NEMOGCM/CONFIG/GLK_LIM3_Michigan/EXP00/bathy_meter.nc"
):
"""
:param file_path:
:param var_name:
:param bathymetry_path: used to mask land points
"""
self.current_time_frame = -1
self.var_name = var_name
self.cube = iris.load_cube(
file_path,
constraint=iris.Constraint(
cube_func=lambda c: c.var_name == var_name))
self.lons, self.lats = cartography.get_xy_grids(self.cube)
lons2d_gl, lats2d_gl = nemo_commons.get_2d_lons_lats_from_nemo(
path=bathymetry_path)
mask_gl = nemo_commons.get_mask(path=bathymetry_path)
xs, ys, zs = lat_lon.lon_lat_to_cartesian(lons2d_gl.flatten(),
lats2d_gl.flatten())
xt, yt, zt = lat_lon.lon_lat_to_cartesian(self.lons.flatten(),
self.lats.flatten())
tree = cKDTree(list(zip(xs, ys, zs)))
dists, indices = tree.query(list(zip(xt, yt, zt)))
self.mask = mask_gl.flatten()[indices].reshape(self.lons.shape)
self.nt = self.cube.shape[0]
assert isinstance(self.cube, Cube)
print(self.nt)
def draw_seasonal_means_panel(path="", var_name="sosstsst"):
cube = iris.load_cube(
path,
constraint=iris.Constraint(cube_func=lambda c: c.var_name == var_name))
assert isinstance(cube, Cube)
#Add month_number coordinate
#coord_categorisation.add_month_number(cube, "time")
coord_categorisation.add_season(cube, "time")
cube_seasonal = cube.aggregated_by("season", analysis.MEAN)
print(cube_seasonal.shape)
#plot results
fig = plt.figure(figsize=(7, 4))
#fig.suptitle(cube.name() + " ({0})".format(cube.units))
nplots = cube_seasonal.shape[0]
ncols = 2
nrows = nplots // ncols if nplots % ncols == 0 else nplots // ncols + 1
b, lons, lats = nemo_commons.get_basemap_and_coordinates_from_file(
T_FILE_PATH, resolution="i")
x, y = b(lons, lats)
gs = gridspec.GridSpec(ncols=ncols + 1,
nrows=nrows,
width_ratios=[1, 1, 0.05],
wspace=0) # +1 for the colorbar
the_mask = nemo_commons.get_mask(
path=os.path.join(EXP_DIR, "bathy_meter.nc"))
vmin = None
vmax = None
for i, season in zip(list(range(nplots)),
cube_seasonal.coord("season").points):
data = cube_seasonal.extract(iris.Constraint(season=season)).data
the_min = data[the_mask].min()
the_max = np.percentile(data[the_mask], 95)
if vmin is None:
vmin, vmax = the_min, the_max
else:
vmin = min(the_min, vmin)
vmax = max(the_max, vmax)
print("{0}: ".format(var_name), vmin, vmax)
cs = None
for i, season in zip(list(range(nplots)),
cube_seasonal.coord("season").points):
print(season)
row = i // ncols
col = i % ncols
ax = fig.add_subplot(gs[row, col])
ax.set_title(season.upper())
data = cube_seasonal.extract(iris.Constraint(season=season)).data
#plot only upper level of the 3d field if given
if data.ndim > 2:
if data.shape[1:] == x.shape:
data = data[0, :, :]
else:
data = data[:, :, 0]
to_plot = np.ma.masked_where(~the_mask, data)
print(to_plot.min(), to_plot.max())
cs = b.pcolormesh(x,
y,
to_plot,
ax=ax,
vmin=vmin,
vmax=vmax,
cmap=cm.get_cmap("jet", 20))
b.drawcoastlines(linewidth=cpp.COASTLINE_WIDTH)
b.drawparallels(np.arange(-90, 90, 2))
b.drawmeridians(np.arange(-180, 180, 2))
plt.colorbar(cs, cax=fig.add_subplot(gs[:, ncols]))
fname = "{0}_{1}.jpeg".format(
cube_seasonal.var_name, "-".join(cube_seasonal.coord("season").points))
if not os.path.isdir(NEMO_IMAGES_DIR):
os.mkdir(NEMO_IMAGES_DIR)
fig.tight_layout()
fig.savefig(os.path.join(NEMO_IMAGES_DIR, fname), dpi=cpp.FIG_SAVE_DPI)
def plot_vector_fields(u_path="",
v_path="",
u_name="vozocrtx",
v_name="vomecrty",
level=0):
name_constraint = iris.Constraint(
cube_func=lambda c: c.var_name == u_name or c.var_name == v_name)
u_cube = iris.load_cube(u_path, constraint=name_constraint)
u_cube = u_cube.extract(
iris.Constraint(model_level_number=u_cube.coord(
"model_level_number").points[level]))
v_cube = iris.load_cube(v_path, constraint=name_constraint)
v_cube = v_cube.extract(
iris.Constraint(model_level_number=v_cube.coord(
"model_level_number").points[level]))
assert isinstance(u_cube, Cube)
#calculate seasonal means
coord_categorisation.add_season(u_cube, "time")
coord_categorisation.add_season(v_cube, "time")
u_cube_seasonal = u_cube.aggregated_by("season", analysis.MEAN)
v_cube_seasonal = v_cube.aggregated_by("season", analysis.MEAN)
#plot results
b, lons, lats = nemo_commons.get_basemap_and_coordinates_from_file(
T_FILE_PATH, resolution="h")
x, y = b(lons, lats)
the_mask = nemo_commons.get_mask(
path=os.path.join(EXP_DIR, "bathy_meter.nc"))
levels = np.arange(0, 0.11, 0.01)
bn = BoundaryNorm(levels, len(levels) - 1)
cmap = cm.get_cmap("Accent", len(levels) - 1)
for season in u_cube_seasonal.coord("season").points:
print(season)
fig = plt.figure(figsize=(8, 4))
ax = fig.add_subplot(111)
ax.set_title(season.upper())
u = u_cube_seasonal.extract(iris.Constraint(season=season)).data
v = v_cube_seasonal.extract(iris.Constraint(season=season)).data
u = np.ma.masked_where(~the_mask, u)
v = np.ma.masked_where(~the_mask, v)
speed = np.sqrt(u**2 + v**2)
cs = b.pcolormesh(x,
y,
speed,
norm=bn,
vmin=levels[0],
vmax=levels[-1],
cmap=cmap)
b.colorbar(cs)
u, v = b.rotate_vector(u, v, lons, lats)
q = b.quiver(x, y, u, v, scale=1.5, width=0.002)
qk = plt.quiverkey(q, 0.15, 0.1, 0.05, '0.05 m/s', labelpos='W')
b.drawcoastlines(linewidth=cpp.COASTLINE_WIDTH)
fname = "{0}-{1}_{2}.jpeg".format(u_cube_seasonal.var_name,
v_cube_seasonal.var_name, season)
if not os.path.isdir(NEMO_IMAGES_DIR):
os.mkdir(NEMO_IMAGES_DIR)
fig.tight_layout()
fig.savefig(os.path.join(NEMO_IMAGES_DIR, fname), dpi=cpp.FIG_SAVE_DPI)
#plot annual mean
fig = plt.figure()
ax = fig.add_subplot(111)
fname = "{0}-{1}_{2}.jpeg".format(u_cube_seasonal.var_name,
v_cube_seasonal.var_name, "annual")
u_annual = u_cube_seasonal.collapsed("season", analysis.MEAN).data
v_annual = v_cube_seasonal.collapsed("season", analysis.MEAN).data
u_annual = np.ma.masked_where(~the_mask, u_annual)
v_annual = np.ma.masked_where(~the_mask, v_annual)
fig.suptitle("Annual")
q = b.quiver(x, y, u_annual, v_annual, scale=1.5, width=0.002, zorder=5)
qk = plt.quiverkey(q, 0.15, 0.1, 0.05, '0.05 m/s', labelpos='W')
levels = np.arange(0, 0.15, 0.01)
bn = BoundaryNorm(levels, len(levels) - 1)
#cmap = my_colormaps.get_cmap_from_ncl_spec_file("colormap_files/wgne15.rgb", len(levels) - 1)
cmap = cm.get_cmap("Paired", len(levels) - 1)
cs = b.pcolormesh(x,
y,
np.sqrt(u_annual**2 + v_annual**2),
cmap=cmap,
norm=bn)
b.drawcoastlines(linewidth=cpp.COASTLINE_WIDTH)
b.colorbar(cs)
fig.tight_layout()
fig.savefig(os.path.join(NEMO_IMAGES_DIR, fname), dpi=cpp.FIG_SAVE_DPI)
def main_for_lake(bathy_path="",
basemap=None,
manager_u=None,
manager_v=None,
scalar_manager=None,
img_dir=""):
"""
:param bathy_path: file used for land sea mask
:param basemap:
:param lons:
:param lats:
:param manager_u:
:param manager_v:
:return:
"""
the_mask = nemo_commons.get_mask(bathy_path)
print(the_mask.shape)
scalar_levels = np.arange(-25, 30, 5)
for frame in range(manager_u.get_total_number_of_time_frames()):
fig = plt.figure()
# bounds = [0, 0.02, 0.05, 0.08, 0.09, 0.1]
bounds = np.arange(-20, 21, 1)
bn = BoundaryNorm(boundaries=bounds, ncolors=len(bounds) - 1)
# img = plt.pcolormesh(data, vmin = bounds[0], vmax = bounds[-1], cmap = cm.get_cmap("jet", len(bounds) - 1), norm = bn)
b_local, lons, lats = nemo_commons.get_basemap_and_coordinates_from_file(
path=bathy_path, resolution="i")
x, y = basemap(lons, lats)
if scalar_manager is not None:
data = scalar_manager.get_next_time_frame_data()
print(data.shape)
img = basemap.pcolormesh(x,
y,
np.ma.masked_where(~the_mask, data),
vmin=scalar_levels[0],
vmax=scalar_levels[-1],
zorder=-6)
basemap.colorbar()
u, v = manager_u.get_next_time_frame_data(
), manager_v.get_next_time_frame_data()
u = np.ma.masked_where(~the_mask, u)
v = np.ma.masked_where(~the_mask, v)
# qp = basemap.quiver(x, y, u, v, scale=1.5)
c = np.sqrt(u**2 + v**2)
qp = basemap.streamplot(x,
y,
u,
v,
color="k",
density=(5, 5),
linewidth=3 * c / c.max())
basemap.drawcoastlines()
basemap.drawmapboundary(fill_color="gray")
plt.title(str(manager_u.get_current_time()).split()[0])
print(str(manager_u.get_current_time()).split()[0])
fig.savefig(os.path.join(img_dir, "{0:08d}.png".format(frame)))
plt.close(fig)
def draw_seasonal_means_panel(path="", var_name="sosstsst"):
cube = iris.load_cube(path, constraint=iris.Constraint(cube_func=lambda c: c.var_name == var_name))
assert isinstance(cube, Cube)
#Add month_number coordinate
#coord_categorisation.add_month_number(cube, "time")
coord_categorisation.add_season(cube, "time")
cube_seasonal = cube.aggregated_by("season", analysis.MEAN)
print(cube_seasonal.shape)
#plot results
fig = plt.figure(figsize=(7, 4))
#fig.suptitle(cube.name() + " ({0})".format(cube.units))
nplots = cube_seasonal.shape[0]
ncols = 2
nrows = nplots // ncols if nplots % ncols == 0 else nplots // ncols + 1
b, lons, lats = nemo_commons.get_basemap_and_coordinates_from_file(T_FILE_PATH, resolution="i")
x, y = b(lons, lats)
gs = gridspec.GridSpec(ncols=ncols + 1, nrows=nrows, width_ratios=[1, 1, 0.05], wspace=0) # +1 for the colorbar
the_mask = nemo_commons.get_mask(path=os.path.join(EXP_DIR, "bathy_meter.nc"))
vmin = None
vmax = None
for i, season in zip(list(range(nplots)), cube_seasonal.coord("season").points):
data = cube_seasonal.extract(iris.Constraint(season=season)).data
the_min = data[the_mask].min()
the_max = np.percentile(data[the_mask], 95)
if vmin is None:
vmin, vmax = the_min, the_max
else:
vmin = min(the_min, vmin)
vmax = max(the_max, vmax)
print("{0}: ".format(var_name), vmin, vmax)
cs = None
for i, season in zip(list(range(nplots)), cube_seasonal.coord("season").points):
print(season)
row = i // ncols
col = i % ncols
ax = fig.add_subplot(gs[row, col])
ax.set_title(season.upper())
data = cube_seasonal.extract(iris.Constraint(season=season)).data
#plot only upper level of the 3d field if given
if data.ndim > 2:
if data.shape[1:] == x.shape:
data = data[0, :, :]
else:
data = data[:, :, 0]
to_plot = np.ma.masked_where(~the_mask, data)
print(to_plot.min(), to_plot.max())
cs = b.pcolormesh(x, y, to_plot, ax=ax, vmin=vmin, vmax=vmax, cmap=cm.get_cmap("jet", 20))
b.drawcoastlines(linewidth=cpp.COASTLINE_WIDTH)
b.drawparallels(np.arange(-90, 90, 2))
b.drawmeridians(np.arange(-180, 180, 2))
plt.colorbar(cs, cax=fig.add_subplot(gs[:, ncols]))
fname = "{0}_{1}.jpeg".format(cube_seasonal.var_name, "-".join(cube_seasonal.coord("season").points))
if not os.path.isdir(NEMO_IMAGES_DIR):
os.mkdir(NEMO_IMAGES_DIR)
fig.tight_layout()
fig.savefig(os.path.join(NEMO_IMAGES_DIR, fname), dpi=cpp.FIG_SAVE_DPI)
def plot_vector_fields(u_path="", v_path="", u_name="vozocrtx", v_name="vomecrty", level=0):
name_constraint = iris.Constraint(cube_func=lambda c: c.var_name == u_name or c.var_name == v_name)
u_cube = iris.load_cube(u_path, constraint=name_constraint)
u_cube = u_cube.extract(iris.Constraint(model_level_number=u_cube.coord("model_level_number").points[level]))
v_cube = iris.load_cube(v_path, constraint=name_constraint)
v_cube = v_cube.extract(iris.Constraint(model_level_number=v_cube.coord("model_level_number").points[level]))
assert isinstance(u_cube, Cube)
#calculate seasonal means
coord_categorisation.add_season(u_cube, "time")
coord_categorisation.add_season(v_cube, "time")
u_cube_seasonal = u_cube.aggregated_by("season", analysis.MEAN)
v_cube_seasonal = v_cube.aggregated_by("season", analysis.MEAN)
#plot results
b, lons, lats = nemo_commons.get_basemap_and_coordinates_from_file(T_FILE_PATH, resolution="h")
x, y = b(lons, lats)
the_mask = nemo_commons.get_mask(path=os.path.join(EXP_DIR, "bathy_meter.nc"))
levels = np.arange(0, 0.11, 0.01)
bn = BoundaryNorm(levels, len(levels) - 1)
cmap = cm.get_cmap("Accent", len(levels) - 1)
for season in u_cube_seasonal.coord("season").points:
print(season)
fig = plt.figure(figsize=(8, 4))
ax = fig.add_subplot(111)
ax.set_title(season.upper())
u = u_cube_seasonal.extract(iris.Constraint(season=season)).data
v = v_cube_seasonal.extract(iris.Constraint(season=season)).data
u = np.ma.masked_where(~the_mask, u)
v = np.ma.masked_where(~the_mask, v)
speed = np.sqrt(u ** 2 + v ** 2)
cs = b.pcolormesh(x, y, speed, norm=bn, vmin=levels[0], vmax=levels[-1], cmap=cmap)
b.colorbar(cs)
u, v = b.rotate_vector(u, v, lons, lats)
q = b.quiver(x, y, u, v, scale=1.5, width=0.002)
qk = plt.quiverkey(q, 0.15, 0.1, 0.05, '0.05 m/s', labelpos='W')
b.drawcoastlines(linewidth=cpp.COASTLINE_WIDTH)
fname = "{0}-{1}_{2}.jpeg".format(u_cube_seasonal.var_name, v_cube_seasonal.var_name, season)
if not os.path.isdir(NEMO_IMAGES_DIR):
os.mkdir(NEMO_IMAGES_DIR)
fig.tight_layout()
fig.savefig(os.path.join(NEMO_IMAGES_DIR, fname), dpi=cpp.FIG_SAVE_DPI)
#plot annual mean
fig = plt.figure()
ax = fig.add_subplot(111)
fname = "{0}-{1}_{2}.jpeg".format(u_cube_seasonal.var_name, v_cube_seasonal.var_name, "annual")
u_annual = u_cube_seasonal.collapsed("season", analysis.MEAN).data
v_annual = v_cube_seasonal.collapsed("season", analysis.MEAN).data
u_annual = np.ma.masked_where(~the_mask, u_annual)
v_annual = np.ma.masked_where(~the_mask, v_annual)
fig.suptitle("Annual")
q = b.quiver(x, y, u_annual, v_annual, scale=1.5, width=0.002, zorder=5)
qk = plt.quiverkey(q, 0.15, 0.1, 0.05, '0.05 m/s', labelpos='W')
levels = np.arange(0, 0.15, 0.01)
bn = BoundaryNorm(levels, len(levels) - 1)
#cmap = my_colormaps.get_cmap_from_ncl_spec_file("colormap_files/wgne15.rgb", len(levels) - 1)
cmap = cm.get_cmap("Paired", len(levels) - 1)
cs = b.pcolormesh(x, y, np.sqrt(u_annual ** 2 + v_annual ** 2), cmap=cmap, norm=bn)
b.drawcoastlines(linewidth=cpp.COASTLINE_WIDTH)
b.colorbar(cs)
fig.tight_layout()
fig.savefig(os.path.join(NEMO_IMAGES_DIR, fname), dpi=cpp.FIG_SAVE_DPI)