-
Lionel GUEZ authored
The corresponding modules are not used in the functions.
Lionel GUEZ authoredThe corresponding modules are not used in the functions.
plot_eddy_contours.py 8.32 KiB
#!/usr/bin/env python3
"""Plots outermost contours and max-speed contours.
Red for anti-cyclones, blue for cyclones.
Without the --light option. Empty circles for extrema. Crosses on
outermost contour for an eddy with no distinct max-speed
contour. Filled circles on outermost contour and max-speed contour for
an eddy with a distinct max-speed contour.
With the --light option. Crosses for extrema.
This script takes about 30 s of CPU for a 90° by 90° window, about 4
mn for a 180° by 180° window.
"""
import shapefile
import numpy as np
from cartopy.mpl import geoaxes
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):
"""Select ishapes at date d and in window. We assume date d is in
slice i_slice of the SHPC.
"""
ishape_list = SHPC.ishape_range(d, i_slice, orientation)
if window is None:
ishape_list_filt = ishape_list
else:
ishape_list_filt = []
reader = SHPC.get_reader(i_slice, orientation, layer = "extremum")
for ishape in ishape_list:
points = reader.shape(ishape).points[0]
if util_eddies.in_window(points, window):
ishape_list_filt.append(ishape)
return ishape_list_filt
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.
dashed: boolean. ax should be an instance of GeoAxes.
"""
if orientation == "Cyclones":
color_outer = "blue"
color_m_s = "cyan"
else:
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")
for ishape in ishape_list:
shape_rec_extr = reader_extr.shapeRecord(ishape)
shape_outer = reader_outer.shape(ishape)
if reader_m_s is None:
shape_m_s = None
else:
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)
if light:
lines[0].set_marker("+")
else:
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")
points = np.array(shape_outer.points)
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:
lines[0].set_marker("x")
else:
lines[0].set_marker("o")
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")
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)
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"])
if window is not None:
llcrnrlon, llcrnrlat, urcrnrlon, urcrnrlat = window
longitude += np.ceil((llcrnrlon - longitude) / 360) * 360
# (in [llcrnrlon, llcrnrlon + 2 pi[)
lon_mask = longitude <= urcrnrlon
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_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
import argparse
import f90nml
from os import path
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")
args = parser.parse_args()
if args.window is not None:
llcrnrlon, llcrnrlat, urcrnrlon, urcrnrlat = args.window
if urcrnrlon - llcrnrlon > 360:
sys.exit("bad values of urcrnrlon and llcrnrlon")
if args.window is None or args.grid:
file = path.join(args.shpc_dir, "grid_nml.txt")
try:
grid_nml = f90nml.read(file)["grid_nml"]
except FileNotFoundError:
print("grid_nml.txt not found. Will not plot bounding box nor grid "
"points.")
grid_nml = None
fig = plt.figure()
src_crs = ccrs.PlateCarree()
projection = ccrs.PlateCarree()
##projection = ccrs.NorthPolarStereo()
ax = plt.axes(projection = projection)
if args.grid: plot_grid_points(grid_nml, args.window, ax, src_crs)
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]
print("No date option, plotting first date of first slice:", d)
i_slice = 0
else:
d = args.date
i_slice = SHPC.get_slice(d)
for orientation in ["Anticyclones", "Cyclones"]:
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)
ax.coastlines()
if args.save:
plt.savefig(f"plot_eddy_contours.{args.save}")
print(f'Created "plot_eddy_contours.{args.save}".')
else:
plt.show()