plot_snapshot.py

#!/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()