Fancy animation

In this example, we show how to make an animation such as the one used in the documentation index.

Let us start by importing all the machinery.

import matplotlib.pyplot as plt
import matplotlib.animation
import cartopy.crs as ccrs
import cmocean
import IPython.display

from pyshqg.backend.numpy_backend import NumpyBackend as Backend
from pyshqg.preprocessing.reference_data import load_test_data
from pyshqg.core.constructors import construct_model, construct_integrator

Now create a model. For this example, let us use the test model with spectral resolution T21 and grid resolution T31.

ds_test = load_test_data(
    internal_truncature=21,
    grid_truncature=31,
)

backend = Backend(floatx='float64')
model = construct_model(backend, ds_test.config)
rk4 = construct_integrator(ds_test.config['rk4_integration'], model)

Then, we use one of the test cases as initial condition. We first run a spinup integration, followed by a full integration.

initial_state = model.model_state(ds_test.isel(ensemble=0).spec_q.to_numpy())
spinup = rk4.run(
    initial_state,
    t_start=0,
    num_snapshots=1344,
    num_steps_per_snapshot=1,
    variables=('spec_q',),
)

state = model.model_state(spinup.to_xarray().isel(time=-1).spec_q.to_numpy())
trajectory = rk4.run(
    state,
    t_start=0,
    num_snapshots=128,
    num_steps_per_snapshot=8,
    variables=('psi',),
)

trajectory_xr = trajectory.to_xarray()

We then use matplolib to create the animation.

fig = plt.figure()
ax = fig.add_subplot(111, projection=ccrs.Robinson())

lon = trajectory_xr.lon.to_numpy()
lat = trajectory_xr.lat.to_numpy()
psi = trajectory_xr.psi.isel(level=-1).to_numpy() / 1e7
time = trajectory_xr.time.to_numpy()

def make_plot(t):
    if t%10 == 0:
        print(f'plotting {t}')
    ax.cla()
    ax.set_title(f'$t=t_0$+{time[t]} [s]')
    im = ax.pcolormesh(
        lon,
        lat,
        psi[t],
        cmap=cmocean.cm.balance,
        vmin=-4,
        vmax=4,
        transform=ccrs.PlateCarree(),
    )
    ax.coastlines()
    return im,

im, = make_plot(0)
plt.colorbar(
    im,
    orientation='horizontal',
    extend='both',
    fraction=0.05,
    shrink=0.75,
    label='Stream function [$\\times 10^{7}$]',
)
anim = matplotlib.animation.FuncAnimation(
    fig,
    make_plot,
    frames=range(len(time)),
    interval=75,
    blit=True,
)
plt.close(fig)

plotting 0
plotting 0
plotting 0

Use the following command to save the animation as a .gif file.

# anim.save('map.gif')

The following command will show the animation in jupyter.

IPython.display.HTML(anim.to_jshtml())

plotting 0
plotting 0
plotting 10
plotting 20
plotting 30
plotting 40
plotting 50
plotting 60
plotting 70
plotting 80
plotting 90
plotting 100
plotting 110
plotting 120

Once Loop Reflect