From b78952797c4b6adee1246a08aae1575ab767d0d4 Mon Sep 17 00:00:00 2001
From: Lionel GUEZ <guez@lmd.ens.fr>
Date: Tue, 7 May 2024 17:57:45 +0200
Subject: [PATCH] Blacken
---
Inst_eddies/Analysis/plot_eddy_contours.py | 203 ++++++++++++++-------
1 file changed, 132 insertions(+), 71 deletions(-)
diff --git a/Inst_eddies/Analysis/plot_eddy_contours.py b/Inst_eddies/Analysis/plot_eddy_contours.py
index 091eb2cd..b3bcc0b9 100755
--- a/Inst_eddies/Analysis/plot_eddy_contours.py
+++ b/Inst_eddies/Analysis/plot_eddy_contours.py
@@ -23,8 +23,14 @@ import cartopy.crs as ccrs
from matplotlib import patches
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
slice i_slice of the SHPC.
@@ -36,7 +42,7 @@ def select_ishapes(d: int, SHPC: util_eddies.SHPC_class, i_slice: int,
ishape_list_filt = ishape_list
else:
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:
points = reader.shape(ishape).points[0]
@@ -45,9 +51,18 @@ def select_ishapes(d: int, SHPC: util_eddies.SHPC_class, i_slice: int,
return ishape_list_filt
-def snapshot(ax: geoaxes.GeoAxesSubplot, ishape_list, SHPC, i_slice,
- orientation, *, dashed = False, light = False,
- src_crs = ccrs.PlateCarree()):
+
+def snapshot(
+ 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
given SHPC, for slice i_slice, for a given orientation, for
ishapes given by ishape_list.
@@ -63,11 +78,13 @@ def snapshot(ax: geoaxes.GeoAxesSubplot, ishape_list, SHPC, i_slice,
color_outer = "red"
color_m_s = "magenta"
- reader_extr = SHPC.get_reader(i_slice, orientation, layer = "extremum")
- reader_outer = SHPC.get_reader(i_slice, orientation,
- layer = "outermost_contour")
- reader_m_s = SHPC.get_reader(i_slice, orientation,
- layer = "max_speed_contour")
+ reader_extr = SHPC.get_reader(i_slice, orientation, layer="extremum")
+ reader_outer = SHPC.get_reader(
+ i_slice, orientation, layer="outermost_contour"
+ )
+ reader_m_s = SHPC.get_reader(
+ i_slice, orientation, layer="max_speed_contour"
+ )
for ishape in ishape_list:
shape_rec_extr = reader_extr.shapeRecord(ishape)
@@ -79,9 +96,14 @@ def snapshot(ax: geoaxes.GeoAxesSubplot, ishape_list, SHPC, i_slice,
shape_m_s = reader_m_s.shape(ishape)
points = shape_rec_extr.shape.points[0]
- lines = ax.plot(points[0], points[1], markersize = 10,
- color = color_outer, fillstyle = "none",
- transform = src_crs)
+ lines = ax.plot(
+ points[0],
+ points[1],
+ markersize=10,
+ color=color_outer,
+ fillstyle="none",
+ transform=src_crs,
+ )
if light:
lines[0].set_marker("+")
@@ -89,52 +111,67 @@ def snapshot(ax: geoaxes.GeoAxesSubplot, ishape_list, SHPC, i_slice,
lines[0].set_marker("o")
if not light:
- ax.annotate(str(shape_rec_extr.record.eddy_index),
- ax.projection.transform_point(points[0], points[1],
- src_crs),
- xytext = (3, 3), textcoords = "offset points")
+ ax.annotate(
+ str(shape_rec_extr.record.eddy_index),
+ ax.projection.transform_point(points[0], points[1], src_crs),
+ xytext=(3, 3),
+ textcoords="offset points",
+ )
points = np.array(shape_outer.points)
- lines = ax.plot(points[:, 0], points[:, 1], color = color_outer,
- transform = src_crs)
+ lines = ax.plot(
+ points[:, 0], points[:, 1], color=color_outer, transform=src_crs
+ )
if not light:
- if shape_m_s == None \
- or shape_m_s.shapeType == shapefile.NULL:
+ if shape_m_s == None or shape_m_s.shapeType == shapefile.NULL:
lines[0].set_marker("x")
else:
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:
points = np.array(shape_m_s.points)
- lines = ax.plot(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")
+ lines = ax.plot(
+ 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")
+
def plot_grid_bb(grid_nml, ax):
"""Grid bounding box."""
if grid_nml is not None:
step = grid_nml["STEP_DEG"]
- rect = patches.Rectangle(grid_nml["corner_deg"],
- (grid_nml["nlon"] - 1) * step[0],
- (grid_nml["nlat"] - 1) * step[1],
- edgecolor="black", fill=False)
+ rect = patches.Rectangle(
+ grid_nml["corner_deg"],
+ (grid_nml["nlon"] - 1) * step[0],
+ (grid_nml["nlat"] - 1) * step[1],
+ edgecolor="black",
+ fill=False,
+ )
ax.add_patch(rect)
+
def plot_grid_points(grid_nml, window, ax, src_crs):
if grid_nml is not None:
width = (grid_nml["NLON"] - 1) * grid_nml["STEP_DEG"][0]
height = (grid_nml["NLat"] - 1) * grid_nml["STEP_DEG"][1]
- longitude = np.linspace(grid_nml["CORNER_DEG"][0],
- grid_nml["CORNER_DEG"][0] + width,
- grid_nml["NLON"])
- latitude = np.linspace(grid_nml["CORNER_DEG"][1],
- grid_nml["CORNER_DEG"][1] + height,
- grid_nml["NLat"])
+ longitude = np.linspace(
+ grid_nml["CORNER_DEG"][0],
+ grid_nml["CORNER_DEG"][0] + width,
+ grid_nml["NLON"],
+ )
+ latitude = np.linspace(
+ grid_nml["CORNER_DEG"][1],
+ grid_nml["CORNER_DEG"][1] + height,
+ grid_nml["NLat"],
+ )
if window is not None:
llcrnrlon, llcrnrlat, urcrnrlon, urcrnrlat = window
@@ -142,16 +179,23 @@ def plot_grid_points(grid_nml, window, ax, src_crs):
# (in [llcrnrlon, llcrnrlon + 2 pi[)
lon_mask = longitude <= urcrnrlon
- lat_mask = np.logical_and(latitude >= llcrnrlat,
- latitude <= urcrnrlat)
+ lat_mask = np.logical_and(
+ latitude >= llcrnrlat, latitude <= urcrnrlat
+ )
else:
- lon_mask = np.ones(len(longitude), dtype = bool)
- lat_mask = np.ones(len(latitude), dtype = bool)
+ lon_mask = np.ones(len(longitude), 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__":
import matplotlib.pyplot as plt
@@ -161,19 +205,27 @@ if __name__ == "__main__":
import sys
parser = argparse.ArgumentParser()
- parser.add_argument("-d", "--date", type = int, help = "date index")
- parser.add_argument("-g", "--grid", help = "plot grid points",
- action = "store_true")
- parser.add_argument("-w", "--window", help = "choose a limited plot window",
- type = float, nargs = 4,
- metavar = ("llcrnrlon", "llcrnrlat", "urcrnrlon",
- "urcrnrlat"))
- 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")
+ parser.add_argument("-d", "--date", type=int, help="date index")
+ parser.add_argument(
+ "-g", "--grid", help="plot grid points", action="store_true"
+ )
+ parser.add_argument(
+ "-w",
+ "--window",
+ help="choose a limited plot window",
+ type=float,
+ nargs=4,
+ metavar=("llcrnrlon", "llcrnrlat", "urcrnrlon", "urcrnrlat"),
+ )
+ 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()
if args.window is not None:
@@ -186,20 +238,23 @@ if __name__ == "__main__":
try:
grid_nml = f90nml.read(file)["grid_nml"]
except FileNotFoundError:
- print("grid_nml.txt not found. Will not plot bounding box nor grid "
- "points.")
+ print(
+ "grid_nml.txt not found. Will not plot bounding box nor grid "
+ "points."
+ )
grid_nml = None
fig = plt.figure()
projection = ccrs.PlateCarree()
##projection = ccrs.NorthPolarStereo()
- ax = plt.axes(projection = projection)
+ ax = plt.axes(projection=projection)
if args.grid:
plot_grid_points(grid_nml, args.window, ax, ccrs.PlateCarree())
- if args.window is None: plot_grid_bb(grid_nml, ax)
- SHPC = util_eddies.SHPC_class(args.shpc_dir, def_orient = "Anticyclones")
+ if args.window is None:
+ plot_grid_bb(grid_nml, ax)
+ SHPC = util_eddies.SHPC_class(args.shpc_dir, def_orient="Anticyclones")
if args.date is None:
d = SHPC.d_init[0]
@@ -210,17 +265,23 @@ if __name__ == "__main__":
i_slice = SHPC.get_slice(d)
for orientation in ["Anticyclones", "Cyclones"]:
- ishape_list = select_ishapes(d, SHPC, i_slice,
- orientation = orientation,
- window = args.window)
+ ishape_list = select_ishapes(
+ d, SHPC, i_slice, orientation=orientation, window=args.window
+ )
if len(ishape_list) == 0:
print(f"{orientation}: No eddy found")
else:
- snapshot(ax, ishape_list, SHPC, i_slice, orientation = orientation,
- light = args.light)
-
- ax.set_title(f"d = {d}", y = 1.05)
- ax.gridlines(draw_labels = True)
+ snapshot(
+ ax,
+ ishape_list,
+ SHPC,
+ i_slice,
+ orientation=orientation,
+ light=args.light,
+ )
+
+ ax.set_title(f"d = {d}", y=1.05)
+ ax.gridlines(draw_labels=True)
ax.coastlines()
if args.save:
--
GitLab