ECMWFCode4Earth / sketchbook-earth

Sketchbook Earth: Illustrated Climate Chronicles
Apache License 2.0
8 stars 3 forks source link

Animation #15

Closed nicrie closed 1 year ago

nicrie commented 1 year ago

Animation showing monthly resolved surface temperatures for a user-defined country

Data

import regionmask

COUNTRIES = regionmask.defined_regions.natural_earth.countries_50.mask_3D(t2m)

region_dict = dict(zip(range(len(COUNTRIES.names)), COUNTRIES.names.values))
print(region_dict)
REGION = int(input())
REGION = region_dict[REGION]
mask = COUNTRIES.isel(region=(COUNTRIES.names == REGION)).squeeze()
nicrie commented 1 year ago

For creating the animation:


import numpy as np
import xarray as xr
import matplotlib.pyplot as plt
from matplotlib.gridspec import GridSpec
import seaborn as sns
from matplotlib.dates import DateFormatter
from matplotlib.animation import FuncAnimation
from matplotlib.collections import LineCollection
import numpy as np
import xarray as xr
import cmocean.cm as cmo
import pandas as pd

from datetime import datetime

plt.style.use('copernicus.mplstyle')
# %%

with xr.open_dataset('../sketchbook-earth/tutorials/01_greenhouse_gases/data/satellites_CO2.nc') as data:
    pass
# %%
weights = np.cos(np.deg2rad(data['lat']))
xco2 = data['xco2'].weighted(weights).mean(('lat', 'lon'))
sigma = data['xco2_stddev'].weighted(weights).mean(('lat', 'lon'))

xco2 *= 1e6

xco2 = xco2.resample({'time': '2W'}).interpolate()
xco2 = xco2.interpolate_na('time', method='nearest')
xco2 = xco2.dropna('time')
# xco2 = xco2.sel(time=slice(None, None, 8))

xco2 = xco2.isel(time=slice(0, 4*12*20))

cmap = sns.color_palette("Spectral_r", as_cmap=True)
cmap_colors = sns.color_palette("Spectral_r", as_cmap=False, n_colors=20)

# Convert time dimension to a period ranging from 0 to 2pi, representing the months January to December
time_rad = xco2['time'].dt.dayofyear / 365. * 2 * np.pi

fig = plt.figure()
gs = GridSpec(5, 10, figure=fig)
bax = plt.subplot(gs[1:4, :5])
ax = plt.subplot(gs[1:4, 5:], polar=True)  # Create a polar axis

# Convert our radial axis to show months instead of radians
ax.set_xticks(np.linspace(0, 2*np.pi, num=12, endpoint=False))
ax.set_xticklabels(['Jan', 'Feb', 'Mar', 'Apr', 'May', 'Jun', 'Jul', 'Aug', 'Sep', 'Oct', 'Nov', 'Dec'])

# LineCollection does not support blitting, we cannot use blit=True in FuncAnimation
collection = LineCollection([], cmap=cmap, norm=plt.Normalize(0, 1))
ax.add_collection(collection)

txt_title_kws = dict(size=20, weight=500, ha='left', va='top', transform=bax.transAxes)
year_title = bax.text(.0, 1.3, '2003', **txt_title_kws)

ax.fill_between(
    np.linspace(0, 2.5*np.pi, num=30, endpoint=False), 30*[300], 30*[170], 
    color=cmap_colors[0], alpha=.2
)
ax.set_theta_offset(np.pi/2)  # so Jan is at top
ax.set_rmin(360)
ax.set_rmax(440)
ax.set_rticks([360, 400, 440])

dates = xco2.coords['time'].values
dates = pd.to_datetime(dates)

date_format = DateFormatter('%Y')
bax.xaxis.set_major_formatter(date_format)
bax.set_xlim(min(dates), max(dates))
xticks = pd.date_range(min(dates), max(dates), freq='3Y')
bax.set_xticks(xticks.tolist())
bax.set_xticklabels([t.strftime('%Y') for t in xticks], rotation=45, ha='right')
bax.set_ylim(360, 440)
bax.set_title('CO$_2$ concentration [ppm]', loc='left', pad=15)
sns.despine(ax=bax, trim=True, offset=5)

line, = bax.plot([], [], '-')

plt.tight_layout()

def init():
    collection.set_segments([])
    return collection,

normalize = plt.Normalize(vmin=xco2.min(), vmax=xco2.max())

def update(i):
    # update the title year
    current_year = xco2.coords['time'].dt.year[i].item()
    year_title.set_text(current_year)

    # update time series
    line.set_data(dates[:i], xco2.values[:i])

    # # update seasonal curve
    points = np.array([time_rad[:i], xco2[:i]]).T.reshape(-1, 1, 2)
    segments = np.concatenate([points[:-1], points[1:]], axis=1)

    # # Normalize the CO2 concentration to the range 0-1 for color mapping
    lc_norm = normalize(xco2[:i])

    collection.set_segments(segments)
    collection.set_array(lc_norm)

    return line, collection

# Create the animation
ani = FuncAnimation(fig, update, frames=xco2.size, init_func=init)

# To display the animation in a Jupyter notebook
from IPython.display import HTML
HTML(ani.to_jshtml())
NikosMastrantonas commented 1 year ago

Currently on the personal repository, will soon be merged to this one. https://github.com/NikosMastrantonas/sketchbook-earth/blob/notebook5/tutorials/05_animations/05_animations.ipynb