dapper.mods.QG.illust_obs

 1"""Stream function and observation time series for QG (quasi-geostrophic) model."""
 2
 3if __name__ == "__main__":  # dont run if pdoc (sample may not be avail/generate-able)
 4    import numpy as np
 5    from matplotlib import pyplot as plt
 6
 7    import dapper as dpr
 8    from dapper.mods.QG.demo import show
 9    from dapper.mods.QG.sakov2008 import HMM, obs_inds
10    from dapper.tools.progressbar import progbar
11
12    # Load (or generate) time-series data of a simulated state AND obs.
13    # (Alternative: re-use sample of QG/__init__.py and observe it)
14    fname = dpr.rc.dirs.data / "QG-ts.npz"
15    try:
16        with np.load(fname) as data:
17            xx = data["xx"][1:]
18            yy = data["yy"]
19    except FileNotFoundError:
20        xx, yy = HMM.simulate()
21        np.savez(fname, xx=xx, yy=yy)
22
23    # Insert obs on the same "grid" as the state vector
24    # Allocate the storage on the parent state dimension
25    yy_xx = np.full_like(xx, np.NaN)
26    for _k, ko, t, _dt in progbar(HMM.tseq.ticker, "Truth & Obs"):
27        if ko is not None:
28            indx = obs_inds(t)
29            yy_xx[ko, indx] = yy[ko, :]
30
31    # Create figure
32    fig, (ax1, ax2) = plt.subplots(ncols=2, sharex=True, sharey=True, figsize=(12, 6))
33    for ax in (ax1, ax2):
34        ax.set_aspect("equal", "box")
35    ax1.set_title(r"Stream function: $\psi$")
36    ax2.set_title(r"Noisy observations: $\mathcal{H}(\psi) + \epsilon$")
37
38    # Define plot updating functions
39    setter1 = show(xx[0], psi=True, ax=ax1)
40    setter2 = show(yy_xx[0], psi=True, ax=ax2)
41
42    # Create double iterable for the animation
43    ts = zip(xx, yy_xx)
44
45    # Animate
46    for k, (xx, yy_xx) in progbar(list(enumerate(ts)), "Animating"):
47        if k % 2 == 0:
48            fig.suptitle("k: " + str(k))
49            setter1(xx)
50            setter2(yy_xx)
51            plt.pause(0.01)