Untitled

mail@pastecode.io avatarunknown
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()