ESS-Visualization/vis_methods.py at master · bparazin/ESS-Visualization · GitHub

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119

import matplotlib.pyplot as plt
import cartopy
import cartopy.crs as ccrs
import env_methods as em
import numpy as np
import geopandas as gpd
import math

#This is a basic method to plot a given 2d numpy array of scalar data, along with it's associated latitude and longitude values.
#Optional parameters to limit the visualization to a rectangle with minimum and maximum longitude and latitude values
def plot_data(lon, lat, data, scalebar, title, lon_min = 0, lon_max = 359, lat_min = -90, lat_max = 90, has_scale_bounds = False, scale_min = 0, scale_max = 0, figsize_x = 20, figsize_y = 15, colormap = "coolwarm", has_rivers = True, file_name = None):
    data = np.asarray(data)

    #make the figure
    fig = plt.figure(figsize = (figsize_x, figsize_y))
    ax = plt.axes(projection = ccrs.PlateCarree(central_longitude = 0))
    ax.coastlines()
    ax.add_feature(cartopy.feature.LAKES, alpha=1)
    if has_rivers: ax.add_feature(cartopy.feature.RIVERS, alpha = 1)

    #The latitudes break because of how they're set up in the NOAA renalysis data unless you do this
    lat1 = em.find_closest_val(lat_min, lat)
    lat2 = em.find_closest_val(lat_max, lat)
    lon1 = em.find_closest_val(lon_min, lon)
    lon2 = em.find_closest_val(lon_max, lon)
    print(lat1)
    print(lat2)

    #label axes and the graph
    ax.set_extent((lon_min, lon_max, lat_min, lat_max))
    ax.set_xticks(lon[math.ceil(lon1 / 10) * 10:lon2][::10], crs = ccrs.PlateCarree())
    ax.set_yticks(lat[math.ceil(lat1 / 10) * 10:lat2][::10], crs = ccrs.PlateCarree())
    plt.xlabel('lon')
    plt.ylabel('lat')
    plt.title(title)

    #
    if not has_scale_bounds:
        mesh = plt.pcolormesh(lon[lon1:lon2],
                              lat[lat1:lat2],
                              data[lat1:lat2, lon1:lon2],
                              cmap=colormap)
    else:
        mesh = plt.pcolormesh(lon[lon1:lon2],
                              lat[lat1:lat2],
                              data[lat1:lat2, lon1:lon2],
                              cmap=colormap, vmin = scale_min, vmax = scale_max)
    cbar = plt.colorbar(mesh)
    cbar.set_label(scalebar)
    if file_name != None: plt.savefig(file_name)

#This is a basic method to plot a given collection of vector data (with the x and y componenets given separately) along with associated
#lat and lon values, vectors colored according to magnitude. Optional parameters to limit bounds of visualization, the scale of the vectors, and, since these plot can get
#messy, how many degrees of lat & lon should separate the vectors
def plot_vector_data(lon, lat, x_data, y_data, title, scale_title,
                     lon_min = 0, lon_max = 359, lat_min = -90, lat_max = 90, detail=3, v_scale = 400,
                     figsize_x = 20, figsize_y = 15, colormap = "GnBu"):
    #make the figure
    fig = plt.figure(figsize = (figsize_x, figsize_y))
    ax = plt.axes(projection = ccrs.PlateCarree(central_longitude = 0))
    ax.coastlines()
    ax.add_feature(cartopy.feature.LAKES, alpha=1)
    ax.add_feature(cartopy.feature.RIVERS, alpha=1)

    lat1 = em.find_closest_val(lat_min, lat)
    lat2 = em.find_closest_val(lat_max, lat)

    #set labels and axes up
    ax.set_extent((lon_min, lon_max, lat_min, lat_max))
    ax.set_xticks(lon[math.ceil(lon_min / 10) * 10:lon_max][::10], crs = ccrs.PlateCarree())
    ax.set_yticks(lat[math.ceil(lat1 / 10) * 10:lat2][::10], crs = ccrs.PlateCarree())
    plt.xlabel('lon')
    plt.ylabel('lat')
    plt.title(title)

    color = np.hypot(x_data[lat1:lat2, lon_min:lon_max][::detail, ::detail],
                     y_data[lat1:lat2, lon_min:lon_max][::detail, ::detail])
    v_field = plt.quiver(lon[lon_min:lon_max][::detail],
                          lat[lat1:lat2][::detail],
                        x_data[lat1:lat2, lon_min:lon_max][::detail, ::detail],
                        y_data[lat1:lat2, lon_min:lon_max][::detail, ::detail], color, cmap=colormap, scale = v_scale)
    scalebar = plt.colorbar(v_field)
    scalebar.set_label(scale_title)


#This is a basic method to plot a given collection of vector data (with the x and y componenets given separately) along with associated
#lat and lon values. Optional parameters to limit bounds of visualization, and, since these plot can get
#messy, how many degrees should separate the vectors. This plots the data as unit vectors colored corresponding to their magnitude.
def plot_nvector_data(lon, lat, x_data, y_data, title, scale_title,
                     lon_min = 0, lon_max = 359, lat_min = -90, lat_max = 90, detail=1,
                     figsize_x = 20, figsize_y = 15, colormap = "GnBu"):
    fig = plt.figure(figsize = (figsize_x, figsize_y))
    ax = plt.axes(projection = ccrs.PlateCarree(central_longitude = 0))
    ax.coastlines()
    ax.add_feature(cartopy.feature.LAKES, alpha=1)
    ax.add_feature(cartopy.feature.RIVERS, alpha=1)

    lat1 = em.find_closest_val(lat_min, lat)
    lat2 = em.find_closest_val(lat_max, lat)

    ax.set_extent((lon_min, lon_max, lat_min, lat_max))
    ax.set_xticks(lon[math.ceil(lon_min / 10) * 10:lon_max][::10], crs = ccrs.PlateCarree())
    ax.set_yticks(lat[math.ceil(lat1 / 10) * 10:lat2][::10], crs = ccrs.PlateCarree())
    plt.xlabel('lon')
    plt.ylabel('lat')
    plt.title(title)

    color = np.hypot(x_data[lat1:lat2, lon_min:lon_max][::detail, ::detail],
                     y_data[lat1:lat2, lon_min:lon_max][::detail, ::detail])
    v_field = plt.quiver(lon[lon_min:lon_max][::detail],
                          lat[lat1:lat2][::detail],
                        np.divide(x_data[lat1:lat2, lon_min:lon_max][::detail, ::detail], color),
                        np.divide(y_data[lat1:lat2, lon_min:lon_max][::detail, ::detail], color),
                         color, cmap=colormap, scale = 125)
    scalebar = plt.colorbar(v_field)
    scalebar.set_label(scale_title)