Untitled
python
2 months ago
3.9 kB
26
Indexable
Never
import matplotlib.pyplot as plt import matplotlib.colors as mcolors import matplotlib.dates as mdates import cartopy.crs as ccrs import cartopy.feature as cfeature from cartopy.mpl.gridliner import LONGITUDE_FORMATTER, LATITUDE_FORMATTER import pandas as pd import numpy as np import xarray as xr from scipy.ndimage.filters import gaussian_filter from matplotlib.colors import LinearSegmentedColormap from mpl_toolkits.axes_grid1 import make_axes_locatable # Load and preprocess the data df = pd.read_csv("/home/users/kieran/incompass/public/kieran/track_data/WD/wd_era5_T42_42p5N.csv") df['time'] = pd.to_datetime(df['time']) df['month'] = df['time'].dt.month df['year'] = df['time'].dt.year fig = plt.figure(figsize=(10,5)) ax = fig.add_subplot(111, projection=ccrs.PlateCarree()) fig.subplots_adjust(bottom=0.25) orofile = xr.open_dataset("/home/users/kieran/ncas/wd-review/trajectory/era5-orography.nc") z = orofile.z.values[0] z = gaussian_filter(z, (2,2)) bins_lat = np.arange(10,60,2.5) bins_lon = np.arange(20,110,2.5) year_list = np.arange(df.year.min()+1,df.year.max()) period = [12, 1, 2, 3] grid = np.zeros((len(year_list), len(bins_lat) - 1, len(bins_lon) - 1)) df_period = df[df['month'].isin(period)] for i, year in enumerate(year_list): df_year = df_period[df_period['year'] == year] density, _, _ = np.histogram2d(df_year['lat'], df_year['lon'], bins=[bins_lat, bins_lon]) grid[i, :, :] = density/len(period) vmin, vmax = 0, 7.5 colors = [(0, "white"), (0.15, "lightgray"), (0.3, "darkgray"), (0.45, "dimgrey"), (0.45001, "yellow"), (0.7, "orange"), (1, "red")] cmap = mcolors.LinearSegmentedColormap.from_list("", colors, N=15) cs = ax.pcolormesh(bins_lon[:-1], bins_lat[:-1], np.mean(grid, axis=0), cmap=cmap, transform=ccrs.PlateCarree(), vmin=vmin, vmax=vmax) ax.add_feature(cfeature.COASTLINE) ax.contour(orofile.longitude, orofile.latitude, z, levels=[20000,], colors='k') df_dec_2022 = df[(df['year'] == 2022) & (df['month'] == 12)] # Group by track_id and get max vort for each track grouped = df_dec_2022.groupby('track_id').agg({'vort': 'max'}) top_3_tracks = grouped.nlargest(3, 'vort').index color_map = {top_3_tracks[0]: 'red', top_3_tracks[1]: 'green', top_3_tracks[2]: 'blue'} for track_id in top_3_tracks: track_df = df_dec_2022[df_dec_2022['track_id'] == track_id] # Make a thicker black line first ax.plot(track_df['lon'], track_df['lat'], color='black', transform=ccrs.Geodetic(), linewidth=4) # Plot a thinner line in the color of choice on top of the black line ax.plot(track_df['lon'], track_df['lat'], color=color_map[track_id], transform=ccrs.Geodetic(), linewidth=2) # Mark and label 00UTC points track_df_00UTC = track_df[track_df['time'].dt.hour == 0] for _, row in track_df_00UTC.iterrows(): if row['lon'] <20: continue ax.scatter(row['lon'], row['lat'], facecolor='white', edgecolor=color_map[track_id], transform=ccrs.PlateCarree(), s=150, zorder=10) ax.text(row['lon'], row['lat'], row['time'].day, color=color_map[track_id], transform=ccrs.PlateCarree(), ha='center', va='center', zorder=20, fontsize='x-small') ax.set_xticks(np.arange(-180, 181, 20), crs=ccrs.PlateCarree()) ax.set_yticks(np.arange(-90, 91, 10), crs=ccrs.PlateCarree()) ax.xaxis.set_major_formatter(LONGITUDE_FORMATTER) ax.yaxis.set_major_formatter(LATITUDE_FORMATTER) ax.set_xlim([20,110]) ax.set_ylim([10,60]) ax.set_title("Strong WDs in December 2022") l, b, w, h = ax.get_position().bounds cax = fig.add_axes([l, b-0.125, w, 0.05]) cbar = fig.colorbar(cs, cax=cax, orientation='horizontal', extend='max') cbar.set_label(u'Climatological track density (WDs per 2.5\u00B0 gridbox per month)') plt.show()