plot_traj.py

#!/usr/bin/env python3

import report_graph
import util_eddies
import matplotlib.pyplot as plt
import numpy as np
import json
import cartopy.crs as ccrs
import random
import argparse

parser = argparse.ArgumentParser()
parser.add_argument("expanded_traj", help = "JSon file")
parser.add_argument("e_overestim", help = "text file")
parser.add_argument("SHPC", help = "directory")
parser.add_argument("orientation", choices = ["Anticyclones", "Cyclones"])
parser.add_argument("--save", metavar = "format",
                    help = "Save file to specified format")
args = parser.parse_args()

with open(args.expanded_traj) as f:
    trajectories = json.load(f)

print("Number of trajectories:", len(trajectories))

with open(args.e_overestim) as f:
    e_overestim = int(f.read())

SHPC = util_eddies.SHPC_class(args.SHPC, def_orient = args.orientation)
src_crs = ccrs.PlateCarree()
projection = ccrs.PlateCarree()
fig, ax = plt.subplots(subplot_kw = {"projection": projection})

for traj in trajectories:
    xy = []

    for node in traj:
        date_index, eddy_index = report_graph.node_to_date_eddy(node,
                                                                e_overestim)
        i_slice = SHPC.get_slice(date_index)
        ishape = SHPC.comp_ishape(date_index, eddy_index, i_slice)
        shape = SHPC.get_reader(i_slice, layer = "extremum").shape(ishape)
        xy.append(shape.points[0])

    xy = np.array(xy)
    ax.plot(xy[:, 0], xy[:, 1], color = "red", transform = src_crs)
    ax.plot(xy[0, 0], xy[0, 1], marker = "s", color = "black",
            transform = src_crs)
    ax.annotate(str(traj[0]), 
                ax.projection.transform_point(xy[0, 0], xy[0, 1], src_crs),
                xytext = (3 * random.random(), 3 * random.random()),
                textcoords = "offset points")

ax.coastlines()
ax.gridlines(draw_labels = True)

if args.save:
    plt.savefig(f"plot_traj.{args.save}")
    print(f'Created "plot_traj.{args.save}".')
else:
    plt.show()