def pause():
"""Loop until user decision is made."""
ch = read1()
while True:
# Set state (pause, skipping, ipdb)
if ch in ENTERs:
self.paused = False
elif ch == CHAR_I:
self.run_ipdb = True
# If keypress valid, resume execution
if ch in ENTERs + [SPACE, CHAR_I]:
break
ch = read1()
# Pause to enable zoom, pan, etc. of mpl GUI
plot_pause(0.01) # Don't use time.sleep()!
def finalize_init(ax, lines, mm):
# Rm lines that only contain NaNs
for name in list(lines):
ln = lines[name]
stat = deep_getattr(stats, name)
if not stat.were_changed:
ln['handle'].remove() # rm from axes
del lines[name] # rm from dict
# Add legends
if lines:
ax.legend(loc='upper left',
bbox_to_anchor=(1.01, 1),
borderaxespad=0)
if mm:
ax.annotate(star + ": mean of\nmarginals",
xy=(0, -1.5 / len(lines)),
xycoords=ax.get_legend().get_frame(),
bbox=dict(alpha=0.0),
fontsize='small')
# coz placement of annotate needs flush sometimes:
plot_pause(0.01)
########################
# ACF
########################
# NB: Won't work with QG (too big, and BCs==0).
fig, ax = freshfig(4)
if "ii" not in locals():
ii = np.arange(min(100, Nx))
if "nlags" not in locals():
nlags = min(100, K-1)
ax.plot(tt[:nlags], np.nanmean(
series.auto_cov(xx[:nlags, ii], nlags=nlags-1, corr=1),
axis=1))
ax.set_xlabel('Time (t)')
ax.set_ylabel('Auto-corr')
viz.plot_pause(0.1)
########################
# "Linearized" forecasting
########################
LL = np.zeros((K, N)) # Local (in time) Lyapunov exponents
E = x + eps*np.eye(Nx)[:N] # Init E
for k, t in enumerate(progbar(tt, "Ens (≈TLM)")):
# if t%10.0==0: print(t)
x = xx[k+1] # = step(x,t,dt) # f.cast reference
E = step(E, t, dt) # f.cast ens (i.e. perturbed ref)
E = (E-x).T/eps # Compute f.cast perturbations
def init(fignum, stats, key0, plot_u, E, P, **kwargs):
xx, yy, mu, spread, tseq = \
stats.xx, stats.yy, stats.mu, stats.spread, stats.HMM.tseq
# Set parameters (kwargs takes precedence over params_orig)
p = DotDict(
**{kw: kwargs.get(kw, val)
for kw, val in params_orig.items()})
# Lag settings:
T_lag, K_lag, a_lag = validate_lag(p.Tplot, tseq)
K_plot = comp_K_plot(K_lag, a_lag, plot_u)
# Extend K_plot forther for adding blanks in resampling (PartFilt):
has_w = hasattr(stats, 'w')
if has_w:
K_plot += a_lag
# Chose marginal dims to plot
if not p.dims:
Nx = min(10, xx.shape[-1])
DimsX = linspace_int(xx.shape[-1], Nx)
else:
Nx = len(p.dims)
DimsX = p.dims
# Pre-process obs dimensions
# Rm inds of obs if not in DimsX
iiY = [i for i, m in enumerate(p.obs_inds) if m in DimsX]
# Rm obs_inds if not in DimsX
DimsY = [m for i, m in enumerate(p.obs_inds) if m in DimsX]
# Get dim (within y) of each x
DimsY = [DimsY.index(m) if m in DimsY else None for m in DimsX]
Ny = len(iiY)
# Set up figure, axes
fig, axs = place.freshfig(fignum,
figsize=(5, 7),
nrows=Nx,
sharex=True)
if Nx == 1:
axs = [axs]
# Tune plots
axs[0].set_title("Marginal time series")
for ix, (m, ax) in enumerate(zip(DimsX, axs)):
# ax.set_ylim(*viz.stretch(*viz.xtrema(xx[:, m]), 1/p.zoomy))
if not p.labels:
ax.set_ylabel("$x_{%d}$" % m)
else:
ax.set_ylabel(p.labels[ix])
axs[-1].set_xlabel('Time (t)')
plot_pause(0.05)
plt.tight_layout()
# Allocate
d = DotDict() # data arrays
h = DotDict() # plot handles
# Why "if True" ? Just to indent the rest of the line...
if True:
d.t = RollingArray((K_plot, ))
if True:
d.x = RollingArray((K_plot, Nx))
h.x = []
if True:
d.y = RollingArray((K_plot, Ny))
h.y = []
if E is not None:
d.E = RollingArray((K_plot, len(E), Nx))
h.E = []
if P is not None:
d.mu = RollingArray((K_plot, Nx))
h.mu = []
if P is not None:
d.s = RollingArray((K_plot, 2, Nx))
h.s = []
# Plot (invisible coz everything here is nan, for the moment).
for ix, (_m, iy, ax) in enumerate(zip(DimsX, DimsY, axs)):
if True:
h.x += ax.plot(d.t, d.x[:, ix], 'k')
if iy != None:
h.y += ax.plot(d.t, d.y[:, iy], 'g*', ms=10)
if 'E' in d:
h.E += [ax.plot(d.t, d.E[:, :, ix], **p.ens_props)]
if 'mu' in d:
h.mu += ax.plot(d.t, d.mu[:, ix], 'b')
if 's' in d:
h.s += [ax.plot(d.t, d.s[:, :, ix], 'b--', lw=1)]
def update(key, E, P):
k, ko, faus = key
EE = duplicate_with_blanks_for_resampled(E, DimsX, key, has_w)
# Roll data array
ind = k if plot_u else ko
for Ens in EE: # If E is duplicated, so must the others be.
if 'E' in d:
d.E.insert(ind, Ens)
if 'mu' in d:
d.mu.insert(ind, mu[key][DimsX])
if 's' in d:
d.s.insert(
ind, mu[key][DimsX] + [[1], [-1]] * spread[key][DimsX])
if True:
d.t.insert(ind, tseq.tt[k])
if True:
d.y.insert(
ind, yy[ko, iiY] if ko is not None else nan * ones(Ny))
if True:
d.x.insert(ind, xx[k, DimsX])
# Update graphs
for ix, (_m, iy, ax) in enumerate(zip(DimsX, DimsY, axs)):
sliding_xlim(ax, d.t, T_lag, True)
if True:
h.x[ix].set_data(d.t, d.x[:, ix])
if iy != None:
h.y[iy].set_data(d.t, d.y[:, iy])
if 'mu' in d:
h.mu[ix].set_data(d.t, d.mu[:, ix])
if 's' in d:
[h.s[ix][b].set_data(d.t, d.s[:, b, ix]) for b in [0, 1]]
if 'E' in d:
[
h.E[ix][n].set_data(d.t, d.E[:, n, ix])
for n in range(len(E))
]
if 'E' in d:
update_alpha(key, stats, h.E[ix])
# TODO 3: fixup. This might be slow?
# In any case, it is very far from tested.
# Also, relim'iting all of the time is distracting.
# Use d_ylim?
if 'E' in d:
lims = d.E
elif 'mu' in d:
lims = d.mu
lims = np.array(viz.xtrema(lims[..., ix]))
if lims[0] == lims[1]:
lims += [-.5, +.5]
ax.set_ylim(*viz.stretch(*lims, 1 / p.zoomy))
return
return update
def update(self, key, E, P):
"""Update liveplots"""
# Check if there are still open figures
if self.any_figs:
open_figns = plt.get_figlabels()
live_figns = set(self.figures.keys())
self.any_figs = bool(live_figns.intersection(open_figns))
else:
return
# Playback control
SPACE = b' '
CHAR_I = b'i'
ENTERs = [b'\n', b'\r'] # Linux + Windows
def pause():
"""Loop until user decision is made."""
ch = read1()
while True:
# Set state (pause, skipping, ipdb)
if ch in ENTERs:
self.paused = False
elif ch == CHAR_I:
self.run_ipdb = True
# If keypress valid, resume execution
if ch in ENTERs + [SPACE, CHAR_I]:
break
ch = read1()
# Pause to enable zoom, pan, etc. of mpl GUI
plot_pause(0.01) # Don't use time.sleep()!
# Enter pause loop
if self.paused:
pause()
else:
if key == (0, None, 'u'):
# Skip read1 for key0 (coz it blocks)
pass
else:
ch = read1()
if ch == SPACE:
# Pause
self.paused = True
self.skipping = False
pause()
elif ch in ENTERs:
# Toggle skipping
self.skipping = not self.skipping
elif ch == CHAR_I:
# Schedule debug
# Note: The reason we dont set_trace(frame) right here is:
# - I could not find the right frame, even doing
# > frame = inspect.stack()[0]
# > while frame.f_code.co_name != "assimilate":
# > frame = frame.f_back
# - It just restarts the plot.
self.run_ipdb = True
# Update figures
if not self.skipping:
faus = key[-1]
if faus != 'u' or self.plot_u:
for name, (updater) in self.figures.items():
if plt.fignum_exists(name) and \
getattr(updater, 'is_active', 1):
_ = plt.figure(name)
updater(key, E, P)
plot_pause(self.params['pause_' + faus])
if self.run_ipdb:
self.run_ipdb = False
import inspect
import ipdb
print("Entering debug mode (ipdb).")
print("Type '?' (and Enter) for usage help.")
print("Type 'c' to continue the assimilation.")
ipdb.set_trace(inspect.stack()[2].frame)
