def plot_tracking_values(observed, filtered, cov_hist, signal_label, ax):
"""
observed: array(nsteps, 2)
Array of observed values
filtered: array(nsteps, state_size)
Array of latent (hidden) values. We consider only the first
two dimensions of the latent values
cov_hist: array(nsteps, state_size, state_size)
History of the retrieved (filtered) covariance matrices
ax: matplotlib AxesSubplot
"""
timesteps, _ = observed.shape
ax.plot(observed[:, 0],
observed[:, 1],
marker="o",
linewidth=0,
markerfacecolor="none",
markeredgewidth=2,
markersize=8,
label="observed",
c="tab:green")
ax.plot(*filtered[:, :2].T,
label=signal_label,
c="tab:red",
marker="x",
linewidth=2)
for t in range(0, timesteps, 1):
covn = cov_hist[t][:2, :2]
pml.plot_ellipse(covn,
filtered[t, :2],
ax,
n_std=2.0,
plot_center=False)
ax.axis("equal")
ax.legend()
def plot_inference(sample_obs, mean_hist, Sigma_hist):
fig, ax = plt.subplots()
ax.scatter(*sample_obs.T, marker="+", color="tab:green")
ax.plot(*mean_hist.T, c="tab:orange", label="filtered")
ax.scatter(*mean_hist[0], c="black", zorder=3)
plt.legend()
collection = [(mut, Vt) for mut, Vt in zip(mean_hist[::4], Sigma_hist[::4])
if Vt[0, 0] > 0 and Vt[1, 1] > 0 and abs(Vt[1, 0] - Vt[0, 1]) < 7e-4]
for mut, Vt in collection:
pml.plot_ellipse(Vt, mut, ax, plot_center=False, alpha=0.9, zorder=3)
plt.axis("equal")
def plot_collection(obs, ax, means=None, covs=None, **kwargs):
n_samples, n_steps, _ = obs.shape
for nsim in range(n_samples):
X = obs[nsim]
if means is not None:
mean = means[nsim]
ax.scatter(*mean[0, :2], marker="o", s=20, c="black", zorder=2)
ax.plot(*mean[:, :2].T,
marker="o",
markersize=2,
**kwargs,
zorder=1)
if covs is not None:
cov = covs[nsim]
for t in range(1, n_steps, 3):
pml.plot_ellipse(cov[t][:2, :2],
mean[t, :2],
ax,
plot_center=False,
alpha=0.7)
ax.scatter(*X.T, marker="+", s=60)
step = 0.1
vmin, vmax = -1.5, 1.5 + step
X = np.mgrid[-1:1.5:step, vmin:vmax:step][::-1]
X_dot = jnp.einsum("ij,jxy->ixy", A, X)
fig, ax = plt.subplots()
ax.plot(*sample_state.T, label="state space")
ax.scatter(*sample_obs.T, marker="+", c="tab:green", s=60, label="observations")
ax.scatter(*sample_state[0], c="black", zorder=3)
field = ax.streamplot(*X, *X_dot, density=1.1, color="#ccccccaa")
ax.legend()
plt.axis("equal")
ax.set_title("State Space")
pml.savefig("kf-circle-state.pdf")
fig, ax = plt.subplots()
ax.plot(*mu_hist.T, c="tab:orange", label="Filtered")
ax.scatter(*sample_obs.T, marker="+", s=60, c="tab:green", label="observations")
ax.scatter(*mu_hist[0], c="black", zorder=3)
for mut, Vt in zip(mu_hist[::4], V_hist[::4]):
pml.plot_ellipse(Vt, mut, ax, plot_center=False, alpha=0.9, zorder=3)
plt.legend()
field = ax.streamplot(*X, *X_dot, density=1.1, color="#ccccccaa")
ax.legend()
plt.axis("equal")
ax.set_title("Approximate Space")
pml.savefig("kf-circle-filtered.pdf")
plt.show()
def plot_uncertainty_ellipses(means, covs):
timesteps = len(means)
for t in range(timesteps):
pml.plot_ellipse(covs[t], means[t], ax, plot_center=False, alpha=0.7)
