def ax_writers(params, g):
observations = gen_single_run() # A conslist
stacked_observations = imap(lambda x: np.stack([x] * num_samples),
observations)
actions = repeat(None)
update, render, recognize = build_conditionals(params)
past_future_dists = look_both_ways(update, render, recognize, flat(),
actions, stacked_observations)
write_params = chain([plot_params], repeat(False))
return imap(write_axes_(params, g, render, observations), write_params,
past_future_dists[:num_frames])
def plot_weights_stats(ax, weights, h1_labels):
h2_labels = cl.imap("_{}".format, cl.count())
h3_labels = ['_W', '_b']
named_weights = flatten_with_names([h1_labels, h2_labels, h3_labels],
weights)
all_labels, all_weights = zip(*named_weights)
ax.boxplot(all_weights, labels=all_labels)
def flatten_with_names(name_lists, tree):
if not name_lists:
return [("", tree)]
cur_names, remaining_names = name_lists[0], name_lists[1:]
descendents_by_child = map(flatten_with_names(remaining_names), tree)
children_with_names = cl.imap(prepend_name, descendents_by_child,
cur_names)
return list_concat(children_with_names)
def write_axes(vect_to_frame, skip, steps_ahead, params, g, render,
observations, write_params, past_future_dists, axes):
rs = np.random.RandomState(1)
past_zs_dist, future_zxs_dist = past_future_dists
past_zs = past_zs_dist.sample(rs)
past_xs = map(lambda z: render(z).sample(rs), past_zs)
future_zs, future_xs = iunzip(future_zxs_dist.sample(rs))
xs = chain(past_xs, future_xs)
zs = chain(past_zs, future_zs)
total_time_steps = steps_ahead
mean_xs = imap(lambda x: np.mean(x, axis=0), xs)
std_xs = imap(lambda x: np.std(x, axis=0), xs)
observation_frames = map(vect_to_frame,
observations[:total_time_steps:skip])
x_mean_stills = map(vect_to_frame, mean_xs[:total_time_steps:skip])
x_std_stills = map(vect_to_frame, std_xs[:total_time_steps:skip])
ax_observations, ax_x_mean, ax_x_std, ax_z = axes[:4]
ax_observations.cla(), ax_x_mean.cla(), ax_x_std.cla(), ax_z.cla()
show_mat(ax_observations, observation_frames, 'Actual rollout')
show_mat(ax_x_mean, x_mean_stills, 'Predictive mean')
show_mat(ax_x_std, x_std_stills, 'Predictive uncertainty')
time_fraction = float(len(past_zs)) / total_time_steps
for ax in axes[:4]:
draw_time(ax, time_fraction)
ax_z.set_ylabel('Latent states')
ax_z.set_xticks([])
ax_z.set_yticks([])
plot_lines_with_uncertainty(ax_z, np.stack(zs[:total_time_steps]))
if write_params:
ax_weights, ax_weight_grads = axes[4:]
ax_weights.cla(), ax_weight_grads.cla()
plot_weights_stats(ax_weights, params,
['dynamics', 'rendering', 'recognition'])
ax_weights.set_ylabel('weights')
plot_weights_stats(ax_weight_grads, g,
['dynamics', 'rendering', 'recognition'])
ax_weight_grads.set_ylabel('grads')
def running_filter(update, recognize, p_z0, actions, observations):
x0, x_rest = observations[0], observations[1:]
z_marginal_0 = multiply_gaussians(p_z0, recognize(x0))
z_marginal_rest, backwards_conditionals = iunzip(
scanl(filtering_step(update, recognize), (z_marginal_0, None), actions,
x_rest))
z_marginals = cons(z_marginal_0, z_marginal_rest)
prev_z_dists = imap(chain_backwards, inits(backwards_conditionals),
z_marginals)
return prev_z_dists, z_marginals
def look_both_ways(update, render, recognize, p_z0, actions, observations):
past_dists, z_marginals = running_filter(update, recognize, p_z0, actions,
observations)
future_dists = imap(generative_dist(update, render), z_marginals,
tails(actions))
return izip(past_dists, future_dists)
def generative_dist(update, render, p_z0, actions):
p_zx_0 = compose(render, p_z0)
return chain(imap(generative_step(update, render), actions), p_zx_0)