Skip to content
Snippets Groups Projects
Commit b7895279 authored by Lionel GUEZ's avatar Lionel GUEZ
Browse files

Blacken

parent 4810537e
No related branches found
No related tags found
No related merge requests found
Hide whitespace changes
Inline Side-by-side
Inst_eddies/Analysis/plot_eddy_contours.py
+ 132
− 71
View file @ b7895279
...@@ -23,8 +23,14 @@ import cartopy.crs as ccrs ...@@ -23,8 +23,14 @@ import cartopy.crs as ccrs
from matplotlib import patches from matplotlib import patches
import util_eddies import util_eddies
def select_ishapes(d: int, SHPC: util_eddies.SHPC_class, i_slice: int,
orientation: str, window: list = None): def select_ishapes(
d: int,
SHPC: util_eddies.SHPC_class,
i_slice: int,
orientation: str,
window: list = None,
):
"""Select ishapes at date d and in window. We assume date d is in """Select ishapes at date d and in window. We assume date d is in
slice i_slice of the SHPC. slice i_slice of the SHPC.
...@@ -36,7 +42,7 @@ def select_ishapes(d: int, SHPC: util_eddies.SHPC_class, i_slice: int, ...@@ -36,7 +42,7 @@ def select_ishapes(d: int, SHPC: util_eddies.SHPC_class, i_slice: int,
ishape_list_filt = ishape_list ishape_list_filt = ishape_list
else: else:
ishape_list_filt = [] ishape_list_filt = []
reader = SHPC.get_reader(i_slice, orientation, layer = "extremum") reader = SHPC.get_reader(i_slice, orientation, layer="extremum")
for ishape in ishape_list: for ishape in ishape_list:
points = reader.shape(ishape).points[0] points = reader.shape(ishape).points[0]
...@@ -45,9 +51,18 @@ def select_ishapes(d: int, SHPC: util_eddies.SHPC_class, i_slice: int, ...@@ -45,9 +51,18 @@ def select_ishapes(d: int, SHPC: util_eddies.SHPC_class, i_slice: int,
return ishape_list_filt return ishape_list_filt
def snapshot(ax: geoaxes.GeoAxesSubplot, ishape_list, SHPC, i_slice,
orientation, *, dashed = False, light = False, def snapshot(
src_crs = ccrs.PlateCarree()): ax: geoaxes.GeoAxesSubplot,
ishape_list,
SHPC,
i_slice,
orientation,
*,
dashed=False,
light=False,
src_crs=ccrs.PlateCarree(),
):
"""Plots extrema, outermost contours and max-speed contours from a """Plots extrema, outermost contours and max-speed contours from a
given SHPC, for slice i_slice, for a given orientation, for given SHPC, for slice i_slice, for a given orientation, for
ishapes given by ishape_list. ishapes given by ishape_list.
...@@ -63,11 +78,13 @@ def snapshot(ax: geoaxes.GeoAxesSubplot, ishape_list, SHPC, i_slice, ...@@ -63,11 +78,13 @@ def snapshot(ax: geoaxes.GeoAxesSubplot, ishape_list, SHPC, i_slice,
color_outer = "red" color_outer = "red"
color_m_s = "magenta" color_m_s = "magenta"
reader_extr = SHPC.get_reader(i_slice, orientation, layer = "extremum") reader_extr = SHPC.get_reader(i_slice, orientation, layer="extremum")
reader_outer = SHPC.get_reader(i_slice, orientation, reader_outer = SHPC.get_reader(
layer = "outermost_contour") i_slice, orientation, layer="outermost_contour"
reader_m_s = SHPC.get_reader(i_slice, orientation, )
layer = "max_speed_contour") reader_m_s = SHPC.get_reader(
i_slice, orientation, layer="max_speed_contour"
)
for ishape in ishape_list: for ishape in ishape_list:
shape_rec_extr = reader_extr.shapeRecord(ishape) shape_rec_extr = reader_extr.shapeRecord(ishape)
...@@ -79,9 +96,14 @@ def snapshot(ax: geoaxes.GeoAxesSubplot, ishape_list, SHPC, i_slice, ...@@ -79,9 +96,14 @@ def snapshot(ax: geoaxes.GeoAxesSubplot, ishape_list, SHPC, i_slice,
shape_m_s = reader_m_s.shape(ishape) shape_m_s = reader_m_s.shape(ishape)
points = shape_rec_extr.shape.points[0] points = shape_rec_extr.shape.points[0]
lines = ax.plot(points[0], points[1], markersize = 10, lines = ax.plot(
color = color_outer, fillstyle = "none", points[0],
transform = src_crs) points[1],
markersize=10,
color=color_outer,
fillstyle="none",
transform=src_crs,
)
if light: if light:
lines[0].set_marker("+") lines[0].set_marker("+")
...@@ -89,52 +111,67 @@ def snapshot(ax: geoaxes.GeoAxesSubplot, ishape_list, SHPC, i_slice, ...@@ -89,52 +111,67 @@ def snapshot(ax: geoaxes.GeoAxesSubplot, ishape_list, SHPC, i_slice,
lines[0].set_marker("o") lines[0].set_marker("o")
if not light: if not light:
ax.annotate(str(shape_rec_extr.record.eddy_index), ax.annotate(
ax.projection.transform_point(points[0], points[1], str(shape_rec_extr.record.eddy_index),
src_crs), ax.projection.transform_point(points[0], points[1], src_crs),
xytext = (3, 3), textcoords = "offset points") xytext=(3, 3),
textcoords="offset points",
)
points = np.array(shape_outer.points) points = np.array(shape_outer.points)
lines = ax.plot(points[:, 0], points[:, 1], color = color_outer, lines = ax.plot(
transform = src_crs) points[:, 0], points[:, 1], color=color_outer, transform=src_crs
)
if not light: if not light:
if shape_m_s == None \ if shape_m_s == None or shape_m_s.shapeType == shapefile.NULL:
or shape_m_s.shapeType == shapefile.NULL:
lines[0].set_marker("x") lines[0].set_marker("x")
else: else:
lines[0].set_marker("o") lines[0].set_marker("o")
if dashed: lines[0].set_linestyle("dashed") if dashed:
lines[0].set_linestyle("dashed")
if shape_m_s != None and shape_m_s.shapeType != shapefile.NULL: if shape_m_s != None and shape_m_s.shapeType != shapefile.NULL:
points = np.array(shape_m_s.points) points = np.array(shape_m_s.points)
lines = ax.plot(points[:, 0], points[:, 1], color = color_m_s, lines = ax.plot(
transform = src_crs) points[:, 0], points[:, 1], color=color_m_s, transform=src_crs
if not light: lines[0].set_marker("o") )
if dashed: lines[0].set_linestyle("dashed") if not light:
lines[0].set_marker("o")
if dashed:
lines[0].set_linestyle("dashed")
def plot_grid_bb(grid_nml, ax): def plot_grid_bb(grid_nml, ax):
"""Grid bounding box.""" """Grid bounding box."""
if grid_nml is not None: if grid_nml is not None:
step = grid_nml["STEP_DEG"] step = grid_nml["STEP_DEG"]
rect = patches.Rectangle(grid_nml["corner_deg"], rect = patches.Rectangle(
(grid_nml["nlon"] - 1) * step[0], grid_nml["corner_deg"],
(grid_nml["nlat"] - 1) * step[1], (grid_nml["nlon"] - 1) * step[0],
edgecolor="black", fill=False) (grid_nml["nlat"] - 1) * step[1],
edgecolor="black",
fill=False,
)
ax.add_patch(rect) ax.add_patch(rect)
def plot_grid_points(grid_nml, window, ax, src_crs): def plot_grid_points(grid_nml, window, ax, src_crs):
if grid_nml is not None: if grid_nml is not None:
width = (grid_nml["NLON"] - 1) * grid_nml["STEP_DEG"][0] width = (grid_nml["NLON"] - 1) * grid_nml["STEP_DEG"][0]
height = (grid_nml["NLat"] - 1) * grid_nml["STEP_DEG"][1] height = (grid_nml["NLat"] - 1) * grid_nml["STEP_DEG"][1]
longitude = np.linspace(grid_nml["CORNER_DEG"][0], longitude = np.linspace(
grid_nml["CORNER_DEG"][0] + width, grid_nml["CORNER_DEG"][0],
grid_nml["NLON"]) grid_nml["CORNER_DEG"][0] + width,
latitude = np.linspace(grid_nml["CORNER_DEG"][1], grid_nml["NLON"],
grid_nml["CORNER_DEG"][1] + height, )
grid_nml["NLat"]) latitude = np.linspace(
grid_nml["CORNER_DEG"][1],
grid_nml["CORNER_DEG"][1] + height,
grid_nml["NLat"],
)
if window is not None: if window is not None:
llcrnrlon, llcrnrlat, urcrnrlon, urcrnrlat = window llcrnrlon, llcrnrlat, urcrnrlon, urcrnrlat = window
...@@ -142,16 +179,23 @@ def plot_grid_points(grid_nml, window, ax, src_crs): ...@@ -142,16 +179,23 @@ def plot_grid_points(grid_nml, window, ax, src_crs):
# (in [llcrnrlon, llcrnrlon + 2 pi[) # (in [llcrnrlon, llcrnrlon + 2 pi[)
lon_mask = longitude <= urcrnrlon lon_mask = longitude <= urcrnrlon
lat_mask = np.logical_and(latitude >= llcrnrlat, lat_mask = np.logical_and(
latitude <= urcrnrlat) latitude >= llcrnrlat, latitude <= urcrnrlat
)
else: else:
lon_mask = np.ones(len(longitude), dtype = bool) lon_mask = np.ones(len(longitude), dtype=bool)
lat_mask = np.ones(len(latitude), dtype = bool) lat_mask = np.ones(len(latitude), dtype=bool)
lon_2d, lat_2d = np.meshgrid(longitude[lon_mask], latitude[lat_mask])
ax.plot(
lon_2d.reshape(-1),
lat_2d.reshape(-1),
transform=src_crs,
marker="+",
color="gray",
linestyle="None",
)
lon_2d, lat_2d = np.meshgrid(longitude[lon_mask],
latitude[lat_mask])
ax.plot(lon_2d.reshape(-1), lat_2d.reshape(-1), transform = src_crs,
marker = "+", color = "gray", linestyle = "None")
if __name__ == "__main__": if __name__ == "__main__":
import matplotlib.pyplot as plt import matplotlib.pyplot as plt
...@@ -161,19 +205,27 @@ if __name__ == "__main__": ...@@ -161,19 +205,27 @@ if __name__ == "__main__":
import sys import sys
parser = argparse.ArgumentParser() parser = argparse.ArgumentParser()
parser.add_argument("-d", "--date", type = int, help = "date index") parser.add_argument("-d", "--date", type=int, help="date index")
parser.add_argument("-g", "--grid", help = "plot grid points", parser.add_argument(
action = "store_true") "-g", "--grid", help="plot grid points", action="store_true"
parser.add_argument("-w", "--window", help = "choose a limited plot window", )
type = float, nargs = 4, parser.add_argument(
metavar = ("llcrnrlon", "llcrnrlat", "urcrnrlon", "-w",
"urcrnrlat")) "--window",
parser.add_argument("-l", "--light", help = "lighter plot", help="choose a limited plot window",
action = "store_true") type=float,
parser.add_argument("shpc_dir", help = "directory containing the " nargs=4,
"collection of shapefiles") metavar=("llcrnrlon", "llcrnrlat", "urcrnrlon", "urcrnrlat"),
parser.add_argument("--save", metavar = "format", )
help = "Save file to specified format") parser.add_argument(
"-l", "--light", help="lighter plot", action="store_true"
)
parser.add_argument(
"shpc_dir", help="directory containing the " "collection of shapefiles"
)
parser.add_argument(
"--save", metavar="format", help="Save file to specified format"
)
args = parser.parse_args() args = parser.parse_args()
if args.window is not None: if args.window is not None:
...@@ -186,20 +238,23 @@ if __name__ == "__main__": ...@@ -186,20 +238,23 @@ if __name__ == "__main__":
try: try:
grid_nml = f90nml.read(file)["grid_nml"] grid_nml = f90nml.read(file)["grid_nml"]
except FileNotFoundError: except FileNotFoundError:
print("grid_nml.txt not found. Will not plot bounding box nor grid " print(
"points.") "grid_nml.txt not found. Will not plot bounding box nor grid "
"points."
)
grid_nml = None grid_nml = None
fig = plt.figure() fig = plt.figure()
projection = ccrs.PlateCarree() projection = ccrs.PlateCarree()
##projection = ccrs.NorthPolarStereo() ##projection = ccrs.NorthPolarStereo()
ax = plt.axes(projection = projection) ax = plt.axes(projection=projection)
if args.grid: if args.grid:
plot_grid_points(grid_nml, args.window, ax, ccrs.PlateCarree()) plot_grid_points(grid_nml, args.window, ax, ccrs.PlateCarree())
if args.window is None: plot_grid_bb(grid_nml, ax) if args.window is None:
SHPC = util_eddies.SHPC_class(args.shpc_dir, def_orient = "Anticyclones") plot_grid_bb(grid_nml, ax)
SHPC = util_eddies.SHPC_class(args.shpc_dir, def_orient="Anticyclones")
if args.date is None: if args.date is None:
d = SHPC.d_init[0] d = SHPC.d_init[0]
...@@ -210,17 +265,23 @@ if __name__ == "__main__": ...@@ -210,17 +265,23 @@ if __name__ == "__main__":
i_slice = SHPC.get_slice(d) i_slice = SHPC.get_slice(d)
for orientation in ["Anticyclones", "Cyclones"]: for orientation in ["Anticyclones", "Cyclones"]:
ishape_list = select_ishapes(d, SHPC, i_slice, ishape_list = select_ishapes(
orientation = orientation, d, SHPC, i_slice, orientation=orientation, window=args.window
window = args.window) )
if len(ishape_list) == 0: if len(ishape_list) == 0:
print(f"{orientation}: No eddy found") print(f"{orientation}: No eddy found")
else: else:
snapshot(ax, ishape_list, SHPC, i_slice, orientation = orientation, snapshot(
light = args.light) ax,
ishape_list,
ax.set_title(f"d = {d}", y = 1.05) SHPC,
ax.gridlines(draw_labels = True) i_slice,
orientation=orientation,
light=args.light,
)
ax.set_title(f"d = {d}", y=1.05)
ax.gridlines(draw_labels=True)
ax.coastlines() ax.coastlines()
if args.save: if args.save:
......
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