test_mean_speed.py

#!/usr/bin/env python3

import netCDF4
import cartopy.crs as ccrs
from scipy import interpolate
import shapefile
import matplotlib.pyplot as plt
import numpy as np
import sys
import math
from numpy import ma

def plot_wind(lon, lat, u, v, ax):
    if lat.ndim == 1 and u.ndim == 2:
        u_src_crs = u / np.cos(lat[:, np.newaxis] / 180 * np.pi)
    else:
        u_src_crs = u / np.cos(lat / 180 * np.pi)

    v_src_crs = v
    magnitude = ma.sqrt(u**2 + v**2)
    magn_src_crs = ma.sqrt(u_src_crs**2 + v_src_crs**2)
    q = ax.quiver(lon, lat, u_src_crs * magnitude / magn_src_crs,
              v_src_crs * magnitude / magn_src_crs,
              transform = ccrs.PlateCarree(), angles = "xy")
    ax.quiverkey(q, 1.1, 1, 0.5, "0.5 m s-1")

with netCDF4.Dataset("uv.nc") as f:
    u = f.variables["ugos"][:].squeeze()
    v = f.variables["vgos"][:].squeeze()
    lon = f.variables["lon"][:]
    lat = f.variables["lat"][:]

lon_rad = lon * np.pi / 180
lat_rad = lat * np.pi / 180

values = np.stack((u, v), axis = -1)

with shapefile.Reader("SHPC/max_speed_contour") as sf:
    p = sf.shape(0)
    
p = np.array(p.points)
lon_interp = p[:, 0]
lat_interp = p[:, 1]

lon_interp_rad = lon_interp[:- 1] / 180 * np.pi
lat_interp_rad = lat_interp[:- 1] / 180 * np.pi

xi = np.column_stack((lat_interp_rad, lon_interp_rad))
values_x = interpolate.interpn((lat_rad, lon_rad), values, xi)

ui = values_x[:, 0]
vi = values_x[:, 1]

with shapefile.Reader("SHPC/extremum") as sf:
    p = sf.shape(0)
    
lon_0, lat_0 = p.points[0]
lon_0_rad = lon_0 * np.pi/180
lat_0_rad = lat_0 * np.pi/180

x = np.cos(lat_0_rad) * (lon_interp_rad - lon_0_rad)
y = lat_interp_rad  - lat_0_rad

v_azim = (x * vi - y * ui) / np.sqrt(x**2 + y**2)
print("mean azimuthal speed = ", np.mean(v_azim), "m s-1")

src_crs = ccrs.PlateCarree()
projection = ccrs.Stereographic(central_latitude = lat_0,
                                central_longitude = lon_0)

plt.figure()
ax = plt.axes(projection = projection)
ax.plot(lon_interp, lat_interp, transform = src_crs)
ax.gridlines(draw_labels = True)
ax.scatter(lon_0, lat_0, transform = src_crs)
plot_wind(lon_interp[:- 1], lat_interp[:- 1], ui, vi, ax)
plt.title("interpolated velocity")

plt.figure()
ax = plt.axes(projection = projection)
ax.plot(lon_interp, lat_interp, transform = src_crs)
ax.gridlines(draw_labels = True)
plot_wind(lon, lat, u, v, ax)
ax.scatter(lon_0, lat_0, transform = src_crs)
plt.title("velocity at grid points")

plt.show()