#!/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
with a valid outermost contour, empty squares for other
extrema. Squares on outermost contour for a well-defined but invalid
outermost contour. Crosses on outermost contour for a valid outermost
contour but with no distinct max-speed contour. Filled circles on
outermost contour and max-speed contour for a valid outermost contour
with a distinct max-speed contour.
With the --light option. Crosses for extrema with a valid outermost
contour.
This scripts 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
import matplotlib.pyplot as plt
import netCDF4
import math
import itertools
import cartopy.crs as ccrs
from os import path
def snaphot(directory, dashed = False, light = False, new_figure = False):
"""Plots extrema, outermost contours and max-speed contours.
directory: containing the three shapefiles.
dashed: boolean."""
reader_extr = shapefile.Reader(path.join(directory, "extremum"))
reader_outer = shapefile.Reader(path.join(directory, "outermost_contour"))
try:
reader_m_s = shapefile.Reader(path.join(directory, "max_speed_contour"))
except shapefile.ShapefileException:
print("Shapefile max_speed_contour not found. "
"Max-speed contours will not be plotted.")
m_s_iterShapes = itertools.repeat(None)
else:
m_s_iterShapes = reader_m_s.iterShapes()
if new_figure:
plt.figure()
ax = plt.axes(projection = projection)
else:
ax = plt.gca()
for shape_rec_extr, shape_outer, shape_m_s \
in zip(reader_extr, reader_outer.iterShapes(), m_s_iterShapes):
points = shape_rec_extr.shape.points[0]
if args.window:
points[0] += math.ceil((llcrnrlon - points[0]) / 360) * 360
# (in [llcrnrlon, llcrnrlon + 2 pi[)
if not args.window or (points[0] <= urcrnrlon
and llcrnrlat <= points[1] <= urcrnrlat):
if shape_rec_extr.record.cyclone == 0:
# Anti-cyclone
color = "red"
else:
color = "blue"
lines = ax.plot(points[0], points[1], markersize = 10,
color = color, fillstyle = "none",
transform = src_crs)
if shape_rec_extr.record.valid == 1:
if light:
lines[0].set_marker("+")
else:
lines[0].set_marker("o")
elif not light:
# Invalid outermost contour
lines[0].set_marker("s")
if not light:
ax.annotate(str(shape_rec_extr.record.eddy_index),
projection.transform_point(points[0], points[1],
src_crs),
xytext = (3, 3), textcoords = "offset points")
if shape_outer.shapeType != shapefile.NULL and not light \
or shape_rec_extr.record.valid:
points = np.array(shape_outer.points)
lines = ax.plot(points[:, 0], points[:, 1], color = color,
transform = src_crs)
if not light:
if shape_rec_extr.record.valid == 0:
# Invalid outermost contour
lines[0].set_marker("s")
lines[0].set_fillstyle("none")
elif 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)
if shape_rec_extr.record.cyclone == 0:
# Anti-cyclone
color = "magenta"
else:
color = "cyan"
lines = ax.plot(points[:, 0], points[:, 1], color = color,
transform = src_crs)
if not light: lines[0].set_marker("o")
if dashed: lines[0].set_linestyle("dashed")
ax.gridlines(draw_labels = True)
ax.coastlines()
if __name__ == "__main__":
import argparse
import sys
parser = argparse.ArgumentParser()
parser.add_argument("-v", "--velocity", help = "plot velocity field",
action = "store_true")
parser.add_argument("-s", "--scale", default = 20, type = float,
help = "scale of arrows for the velocity field")
parser.add_argument("-g", "--grid", help = "plot grid",
action = "store_true")
parser.add_argument("-w", "--window", help = "choose a limited plot window",
action = "store_true")
parser.add_argument("-l", "--light", help = "lighter plot",
action = "store_true")
parser.add_argument("-d", "--dashed", action = "store_true",
help = "dashed linestyle, useful for a second snapshot")
parser.add_argument("directory", help = "containing the three shapefiles")
args = parser.parse_args()
if args.grid or args.velocity:
with netCDF4.Dataset("h.nc") as f:
longitude = f.variables["longitude"][:]
latitude = f.variables["latitude"][:]
if args.window:
l = input("llcrnrlon, llcrnrlat, urcrnrlon, "
"urcrnrlat = ? ").split(sep = ",")
llcrnrlon, llcrnrlat, urcrnrlon, urcrnrlat = [float(v) for v in l]
if urcrnrlon - llcrnrlon > 360:
sys.exit("bad values of urcrnrlon and llcrnrlon")
if args.grid or args.velocity:
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)
elif args.grid or args.velocity:
lon_mask = np.ones(len(longitude), dtype = bool)
lat_mask = np.ones(len(latitude), dtype = bool)
plt.figure()
src_crs = ccrs.PlateCarree()
projection = ccrs.PlateCarree()
##projection = ccrs.NorthPolarStereo()
ax = plt.axes(projection = projection)
if args.grid:
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")
snaphot(args.directory, args.dashed, args.light)
if args.velocity:
with netCDF4.Dataset("uv.nc") as f:
quiver_return = ax.quiver(longitude[lon_mask], latitude[lat_mask],
f["ugos"][0, lat_mask][:, lon_mask],
f["vgos"][0, lat_mask][:, lon_mask],
scale = args.scale, scale_units = "width",
transform = src_crs)
plt.quiverkey(quiver_return, 0.9, 0.9, 1, r"1 m s$^{-1}$",
coordinates = "figure")
plt.show()