import numpy as np
from matplotlib import pyplot as plt
import dapper as dpr
from dapper.mods.LorenzUV.lorenz96 import LUV
from dapper.tools.utils import progbar
from dapper.tools.viz import setup_wrapping
nU, J = LUV.nU, LUV.J
dt = 0.005
K = int(4/dt)
step_1 = dpr.with_rk4(LUV.dxdt, autonom=True)
step_K = dpr.with_recursion(step_1, prog='Simulating')
xx = step_K(LUV.x0, K, np.nan, dt)
# Grab parts of state vector
ii, wrapU = setup_wrapping(nU)
jj, wrapV = setup_wrapping(nU*J)
# Animate linear
plt.figure()
lhU = plt.plot(ii, wrapU(xx[-1, :nU]), 'b', lw=3)[0]
lhV = plt.plot(jj/J, wrapV(xx[-1, nU:]), 'g', lw=2)[0]
for k in progbar(range(K), 'Plotting'):
lhU.set_ydata(wrapU(xx[k, :nU]))
lhV.set_ydata(wrapV(xx[k, nU:]))
plt.pause(0.001)
def init(fignum, stats, key0, plot_u, E, P, **kwargs):
xx, yy, mu = stats.xx, stats.yy, stats.mu
# Set parameters (kwargs takes precedence over params_orig)
p = DotDict(
**{kw: kwargs.get(kw, val)
for kw, val in params_orig.items()})
if not p.dims:
M = xx.shape[-1]
p.dims = arange(M)
else:
M = len(p.dims)
# Make periodic wrapper
ii, wrap = viz.setup_wrapping(M, p.periodicity)
# Set up figure, axes
fig, ax = place.freshfig(fignum, figsize=(8, 5))
fig.suptitle("1d amplitude plot")
# Nans
nan1 = wrap(nan * ones(M))
if E is None and p.conf_mult is None:
p.conf_mult = 2
# Init plots
if p.conf_mult:
lines_s = ax.plot(ii,
nan1,
"b-",
lw=1,
label=(str(p.conf_mult) + r'$\sigma$ conf'))
lines_s += ax.plot(ii, nan1, "b-", lw=1)
line_mu, = ax.plot(ii, nan1, 'b-', lw=2, label='DA mean')
else:
nanE = nan * ones((stats.xp.N, M))
lines_E = ax.plot(ii,
wrap(nanE[0]),
**p.ens_props,
lw=1,
label='Ensemble')
lines_E += ax.plot(ii, wrap(nanE[1:]).T, **p.ens_props, lw=1)
# Truth, Obs
(line_x, ) = ax.plot(ii, nan1, 'k-', lw=3, label='Truth')
if p.obs_inds is not None:
p.obs_inds = np.asarray(p.obs_inds)
(line_y, ) = ax.plot(p.obs_inds,
nan * p.obs_inds,
'g*',
ms=5,
label='Obs')
# Tune plot
ax.set_ylim(*viz.xtrema(xx))
ax.set_xlim(viz.stretch(ii[0], ii[-1], 1))
# Xticks
xt = ax.get_xticks()
xt = xt[abs(xt % 1) < 0.01].astype(int) # Keep only the integer ticks
xt = xt[xt >= 0]
xt = xt[xt < len(p.dims)]
ax.set_xticks(xt)
ax.set_xticklabels(p.dims[xt])
ax.set_xlabel('State index')
ax.set_ylabel('Value')
ax.legend(loc='upper right')
text_t = ax.text(0.01,
0.01,
format_time(None, None, None),
transform=ax.transAxes,
family='monospace',
ha='left')
# Init visibility (must come after legend):
if p.obs_inds is not None:
line_y.set_visible(False)
def update(key, E, P):
k, ko, faus = key
if p.conf_mult:
sigma = mu[key] + p.conf_mult * stats.spread[key] * [[1], [-1]]
lines_s[0].set_ydata(wrap(sigma[0, p.dims]))
lines_s[1].set_ydata(wrap(sigma[1, p.dims]))
line_mu.set_ydata(wrap(mu[key][p.dims]))
else:
for n, line in enumerate(lines_E):
line.set_ydata(wrap(E[n, p.dims]))
update_alpha(key, stats, lines_E)
line_x.set_ydata(wrap(xx[k, p.dims]))
text_t.set_text(format_time(k, ko, stats.HMM.tseq.tt[k]))
if 'f' in faus:
if p.obs_inds is not None:
line_y.set_ydata(yy[ko])
line_y.set_zorder(5)
line_y.set_visible(True)
if 'u' in faus:
if p.obs_inds is not None:
line_y.set_visible(False)
return
return update