Untitled

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()