dapper.mods.QG.demo

 1"""Demonstrate the QG (quasi-geostrophic) model."""
 2
 3import numpy as np
 4import scipy.ndimage.filters as filters
 5from matplotlib import pyplot as plt
 6
 7from dapper.mods.QG import default_prms, nx, sample_filename, square
 8from dapper.tools.progressbar import progbar
 9
10
11def show(x0, psi=True, ax=None):
12    # Whether to show psi (streamfun) or q (potential vorticity)
13    def psi_or_q(x):
14        return x if psi else compute_q(x)
15
16    # Create fig if necessary
17    if ax is None:
18        fig, ax = plt.subplots()
19    # Init ax
20    im = ax.imshow(psi_or_q(square(x0)))
21    if psi:
22        im.set_clim(-30, 30)
23    else:
24        im.set_clim(-28e4, 25e4)
25
26    # Define plot update fun, which we return
27    def update(x):
28        im.set_data(psi_or_q(square(x)))
29
30    return update
31
32
33# Although psi is the state variable, q looks cooler.
34# q = Nabla^2(psi) - F*psi.
35dx = 1 / (nx - 1)
36
37
38def compute_q(psi):
39    Lapl = filters.laplace(psi, mode="constant") / dx**2
40    # mode='constant' coz BCs are: psi = nabla psi = nabla^2 psi = 0
41    return Lapl - default_prms["F"] * psi
42
43
44if __name__ == "__main__":
45    fig, (ax1, ax2) = plt.subplots(ncols=2, sharex=True, sharey=True, figsize=(8, 4))
46    for ax in (ax1, ax2):
47        ax.set_aspect("equal", "box")
48    ax1.set_title(r"$\psi$")
49    ax2.set_title("$q$")
50
51    xx = np.load(sample_filename)["sample"]
52    setter1 = show(xx[0], psi=True, ax=ax1)
53    setter2 = show(xx[0], psi=False, ax=ax2)
54
55    for k, x in progbar(list(enumerate(xx)), "Animating"):
56        if k % 2 == 0:
57            fig.suptitle("k: " + str(k))
58            setter1(x)
59            setter2(x)
60            plt.pause(0.01)