def plot_log_likelihood(out_dir: str, config: ExperimentConfig, result: ExperimentResult, metrics: ExperimentMetrics):
domain = np.arange(result.iterations)
log_likelihood = metrics.log_likelihood
with SubplotsAndSave(out_dir, 'log-likelihood') as (fig, ax):
show_debug_info(fig, config, result)
ax.plot(domain, log_likelihood[:result.iterations], label='Log-Likelihood')
ax.set_title('Log-Likelihood (%s), %d Iterations' % (config.title, result.iterations))
ax.set_xlabel('Iterations')
ax.set_ylabel('Log-Likelihood')
ax.legend(loc='lower right')
def _plot_lunar_lander(out_dir: str, config: ExperimentConfig,
result: ExperimentResult, N: int,
observation_trajectories: List[np.ndarray]):
learned_initial_observation = result.g_numpy(result.m0)
if result.y_shift is not None and result.y_scale is not None:
learned_initial_observation = result.y_shift + learned_initial_observation * result.y_scale
observation_trajectories_smoothed = result.g_numpy(result.estimations_latents.transpose((0, 2, 1)).reshape(-1, config.latent_dim)) \
.reshape((N, config.T_train, config.observation_dim))
with SubplotsAndSave(out_dir,
'lunar-lander',
1,
N,
sharex='col',
sharey='row',
figsize=figsize(1, N),
squeeze=False) as (fig, axs):
show_debug_info(fig, config, result)
for n, (ax, observation_trajectory,
observation_trajectory_smoothed) in enumerate(
zip(axs.flatten(), observation_trajectories[:N],
observation_trajectories_smoothed[:N])):
truth_trajectory_x_train = result.observations[n, :config.T_train,
0]
truth_trajectory_y_train = result.observations[n, :config.T_train,
1]
truth_trajectory_x_test = result.observations[n,
config.T_train - 1:,
0]
truth_trajectory_y_test = result.observations[n,
config.T_train - 1:,
1]
observation_trajectory_x_train = observation_trajectory[:config.
T_train, 0]
observation_trajectory_y_train = observation_trajectory[:config.
T_train, 1]
observation_trajectory_x_test = observation_trajectory[
config.T_train - 1:, 0]
observation_trajectory_y_test = observation_trajectory[
config.T_train - 1:, 1]
observation_trajectory_smoothed_x = observation_trajectory_smoothed[:,
0]
observation_trajectory_smoothed_y = observation_trajectory_smoothed[:,
1]
# Ground truth.
ax.scatter(truth_trajectory_x_train,
truth_trajectory_y_train,
s=1,
color='tuda:black',
label='Truth')
ax.scatter(truth_trajectory_x_test,
truth_trajectory_y_test,
s=1,
color='tuda:gray',
label='Truth (Test)')
# Smoothed trajectory.
ax.scatter(observation_trajectory_smoothed_x,
observation_trajectory_smoothed_y,
s=1,
color='tuda:orange',
label='Smoothed')
# Rollout w/ control inputs.
ax.scatter(observation_trajectory_x_train,
observation_trajectory_y_train,
s=1,
color='tuda:blue',
label='Rollout')
ax.scatter(observation_trajectory_x_test,
observation_trajectory_y_test,
s=1,
color='tuda:purple',
label='Rollout (Prediction)')
# Prediction boundary and learned initial value.
ax.scatter(learned_initial_observation[0],
learned_initial_observation[1],
marker='*',
color='tuda:green',
label='Learned Initial Value')
ax.set_title('Sequence %d' % (n + 1))
ax.set_xlabel(r'$x$')
ax.set_ylabel(r'$y$')
ax.legend()
def _plot_latent_rollout(
out_dir: str, config: ExperimentConfig, result: ExperimentResult,
N: int, latent_rollout: List[np.ndarray],
latent_covariances: List[np.ndarray],
latent_rollout_without_control: List[Optional[np.ndarray]],
latent_pred_rollout: List[np.ndarray],
latent_pred_covariances: List[np.ndarray]):
domain = np.arange(config.T) * config.h
domain_train = domain[:config.T_train]
domain_test = domain[config.T_train - 1:]
for dim in range(config.latent_dim):
for n, item in enumerate(
zip(latent_rollout[:N], latent_covariances[:N],
latent_rollout_without_control[:N],
result.estimations_latents[:N], latent_pred_rollout[:N],
latent_pred_covariances[:N])):
latent_trajectory, latent_covariance, latent_trajectory_without_control, latent_trajectory_smoothed, latent_pred_trajectory, latent_pred_covariance = item
with SubplotsAndSave(out_dir,
f'rollout-latents-N{n}-D{dim}',
place_legend_outside=True) as (fig, ax):
show_debug_info(fig, config, result)
_plot_latent_rollout_to_ax(
ax, config, dim, domain, domain_test, domain_train,
latent_covariance, latent_pred_covariance,
latent_pred_trajectory, latent_trajectory,
latent_trajectory_smoothed,
latent_trajectory_without_control, n, result)
if N > 1:
ax.set_title('Sequence %d' % (n + 1))
ax.set_xlabel('$t$')
for n, item in enumerate(
zip(latent_rollout[:N], latent_covariances[:N],
latent_rollout_without_control[:N],
result.estimations_latents[:N], latent_pred_rollout[:N],
latent_pred_covariances[:N])):
latent_trajectory, latent_covariance, latent_trajectory_without_control, latent_trajectory_smoothed, latent_pred_trajectory, latent_pred_covariance = item
ncols, nrows, check = 1, None, 0
while nrows is None or nrows > ncols + check:
ncols += 1
nrows = int(np.ceil(config.latent_dim / ncols))
check += 1
with SubplotsAndSave(out_dir,
f'rollout-latents-N{n}',
nrows=nrows,
ncols=ncols,
place_legend_outside=True,
sharex='col',
squeeze=False) as (fig, axs):
show_debug_info(fig, config, result)
for dim, ax in zip(range(config.latent_dim), axs.flatten()):
_plot_latent_rollout_to_ax(
ax, config, dim, domain, domain_test, domain_train,
latent_covariance, latent_pred_covariance,
latent_pred_trajectory, latent_trajectory,
latent_trajectory_smoothed,
latent_trajectory_without_control, n, result)
if N > 1:
ax.set_title('Sequence %d' % (n + 1))
if dim == config.latent_dim - 1 or dim == config.latent_dim - 2:
ax.set_xlabel('$t$')
def _plot_observations_rollout(
out_dir: str, config: ExperimentConfig, result: ExperimentResult,
N: int, observation_trajectories: List[np.ndarray],
observation_covariances: List[np.ndarray],
observation_trajectories_without_control: List[Optional[np.ndarray]],
observation_covariances_without_control: List[np.ndarray],
observation_pred_rollout: List[np.ndarray],
observation_pred_covariances: List[np.ndarray]):
domain = np.arange(config.T) * config.h
domain_train = domain[:config.T_train]
domain_test = domain[config.T_train - 1:]
learned_initial_observation = result.g_numpy(result.m0)
if result.y_shift is not None and result.y_scale is not None:
learned_initial_observation = result.y_shift + learned_initial_observation * result.y_scale
if result.V_hat is None:
observation_trajectories_smoothed = result.g_numpy(result.estimations_latents.transpose((0, 2, 1)).reshape(-1, config.latent_dim)) \
.reshape((N, config.T_train, config.observation_dim))
observation_covariances_smoothed = [None] * N
else:
observation_trajectories_smoothed, observation_covariances_smoothed = zip(
*[
compute_observations(
config, result, result.estimations_latents[n].T,
result.V_hat[n, :, :, :].transpose((2, 0, 1)))
for n in range(N)
])
for dim, dim_name in enumerate(config.observation_dim_names):
for n, item in enumerate(
zip(observation_trajectories[:N], observation_covariances[:N],
observation_trajectories_without_control[:N],
observation_trajectories_smoothed[:N],
observation_covariances_smoothed[:N],
observation_covariances_without_control[:N],
observation_pred_rollout[:N],
observation_pred_covariances[:N])):
observation_trajectory, observation_covariance, observation_trajectory_without_control, observation_trajectory_smoothed, _, _, _, _ = item
_, _, _, _, observation_covariance_smoothed, observation_covariance_without_control, observation_pred_trajectory, observation_pred_covariance = item
with SubplotsAndSave(out_dir,
f'rollout-observations-N{n}-D{dim}',
place_legend_outside=True) as (fig, ax):
show_debug_info(fig, config, result)
_plot_observations_rollout_to_ax(
ax, config, dim, dim_name, domain, domain_test,
domain_train, learned_initial_observation, n,
observation_covariance, observation_covariance_smoothed,
observation_covariance_without_control,
observation_pred_covariance, observation_pred_trajectory,
observation_trajectory, observation_trajectory_smoothed,
observation_trajectory_without_control, result)
if N > 1:
ax.set_title('Sequence %d' % (n + 1))
ax.set_xlabel('$t$')
for n, item in enumerate(
zip(observation_trajectories[:N], observation_covariances[:N],
observation_trajectories_without_control[:N],
observation_trajectories_smoothed[:N],
observation_covariances_smoothed[:N],
observation_covariances_without_control[:N],
observation_pred_rollout[:N],
observation_pred_covariances[:N])):
observation_trajectory, observation_covariance, observation_trajectory_without_control, observation_trajectory_smoothed, _, _, _, _ = item
_, _, _, _, observation_covariance_smoothed, observation_covariance_without_control, observation_pred_trajectory, observation_pred_covariance = item
with SubplotsAndSave(out_dir,
f'rollout-observations-N{n}',
nrows=int(np.ceil(config.observation_dim / 2)),
ncols=min(config.observation_dim, 2),
place_legend_outside=True,
sharex='col',
squeeze=False) as (fig, axs):
show_debug_info(fig, config, result)
for dim, (ax, dim_name) in enumerate(
zip(axs.flatten(), config.observation_dim_names)):
_plot_observations_rollout_to_ax(
ax, config, dim, dim_name, domain, domain_test,
domain_train, learned_initial_observation, n,
observation_covariance, observation_covariance_smoothed,
observation_covariance_without_control,
observation_pred_covariance, observation_pred_trajectory,
observation_trajectory, observation_trajectory_smoothed,
observation_trajectory_without_control, result)
if N > 1:
ax.set_title('Sequence %d' % (n + 1))
if dim == config.observation_dim - 1 or dim == config.observation_dim - 2:
ax.set_xlabel('$t$')