plot samples from each regime, post_train.py

examples/monkey_reaching/post_train.py

@@ -1,9 +1,10 @@
 import torch
+from matplotlib import cm
+import matplotlib.pyplot as plt
 import pytorch_lightning as lightning
 
 from hydra import compose, initialize
-from pytorch_lightning.loggers import CSVLogger
-from pytorch_lightning.callbacks import ModelCheckpoint
+from xfads.plot_utils import plot_z_samples
 from xfads.smoothers.lightning_trainers import LightningMonkeyReaching
 from xfads.ssm_modules.prebuilt_models import create_xfads_poisson_log_link
 
@@ -42,15 +43,74 @@ def main():
     test_dataloader = torch.utils.data.DataLoader(y_test_dataset, batch_size=y_valid_obs.shape[0], shuffle=False)
 
     """create ssm"""
-    ssm = create_xfads_poisson_log_link(cfg, n_neurons_obs, train_dataloader)
+    ssm = create_xfads_poisson_log_link(cfg, n_neurons_obs, train_dataloader, model_type='n')
 
     """lightning"""
     model_ckpt_path = 'results/noncausal_model.ckpt'
     seq_vae = LightningMonkeyReaching.load_from_checkpoint(model_ckpt_path, ssm=ssm, cfg=cfg,
                                                            n_time_bins_enc=n_time_bins_enc, n_time_bins_bhv=n_bins_bhv,
                                                            strict=False)
     """extract trained ssm from lightning module"""
-    ssm = seq_vae.ssm
+    seq_vae.ssm = seq_vae.ssm.to(cfg.device)
+    seq_vae.ssm.eval()
+
+    z_s_train = []
+    z_s_valid = []
+    z_f_valid = []
+    z_p_valid = []
+
+    for batch in train_dataloader:
+        loss, z, stats = seq_vae.ssm(batch[0], cfg.n_samples)
+        z_s_train.append(z)
+
+    for batch in valid_dataloader:
+        z_f, stats = seq_vae.ssm.fast_filter_1_to_T(batch[0], cfg.n_samples)
+        loss, z, stats = seq_vae.ssm(batch[0], cfg.n_samples)
+        z_p = seq_vae.ssm.predict_forward(z_f[:, :, 10], cfg.n_samples)
+        z_p = torch.cat([z_f[:, :, :10], z_p], dim=2)
+        z_f_valid.append(z_f)
+        z_p_valid.append(z_p)
+        z_s_valid.append(z)
+
+    U, S, V = torch.svd(seq_vae.ssm.likelihood_pdf.readout_fn[-1].weight.data)
+    V = S.unsqueeze(-1) * V
+
+    z_s_train = torch.cat(z_s_train, dim=1)
+    z_s_valid = torch.cat(z_s_valid, dim=1)
+    z_f_valid = torch.cat(z_f_valid, dim=1)
+    z_p_valid = torch.cat(z_p_valid, dim=1)
+
+    z_s_train = z_s_train[..., :cfg.n_latents_read] @ V
+    z_s_test = z_s_valid[..., :cfg.n_latents_read] @ V
+    z_f_test = z_f_valid[..., :cfg.n_latents_read] @ V
+    z_p_test = z_p_valid[..., :cfg.n_latents_read] @ V
+
+    """colors"""
+    blues = cm.get_cmap("winter", z_s_test.shape[0])
+    reds = cm.get_cmap("summer", z_s_test.shape[0])
+    springs = cm.get_cmap("spring", z_s_test.shape[0])
+
+    """plot sample regimes"""
+    with torch.no_grad():
+        trial_list = [28, 202, 8, 285]
+        color_map_list = [blues, reds, springs]
+
+        fig, axs = plt.subplots(len(trial_list), 1, figsize=(4, 4))
+        plot_z_samples(fig, axs, z_s_test[:, trial_list, ..., :3], color_map_list)
+        fig.savefig('plots/z_s_test.pdf', bbox_inches='tight', transparent=True)
+        plt.show()
+
+        fig, axs = plt.subplots(len(trial_list), 1, figsize=(4, 4))
+        plot_z_samples(fig, axs, z_f_test[:, trial_list, ..., :3], color_map_list)
+        fig.savefig('plots/z_f_test.pdf', bbox_inches='tight', transparent=True)
+        plt.show()
+
+        fig, axs = plt.subplots(len(trial_list), 1, figsize=(4, 4))
+        [axs[i].axvline(10, linestyle='--', color='red') for i in range(len(trial_list))]
+        plot_z_samples(fig, axs, z_p_test[:, trial_list, ..., :3], color_map_list)
+        fig.savefig('plots/z_prd_test.pdf', bbox_inches='tight', transparent=True)
+        plt.show()
+
 
 
 if __name__ == '__main__':

xfads/plot_utils.py

@@ -18,3 +18,16 @@ def plot_two_d_vector_field(dynamics_fn, axs, min_xy=-3, max_xy=3, n_pts=500, de
         v = s[:, 1].reshape(Y.shape[0], Y.shape[1])
 
         axs.streamplot(X, Y, u, v, color='black', linewidth=0.5, arrowsize=0.5)
+
+
+def plot_z_samples(fig, axs, samples, color_map_list):
+    n_samples, n_trials, n_bins, n_neurons = samples.shape
+    fig.subplots_adjust(hspace=0)
+    [axs[i].axvline(12, linestyle='--', color='gray') for i in range(n_trials)]
+    [axs[i].axis('off') for i in range(n_trials)]
+    [axs[i].plot(samples[j, i, :, n], color=color_map_list[n](j), linewidth=0.5, alpha=0.4)
+     for i in range(n_trials) for j in range(samples.shape[0]) for n in range(n_neurons)]
+
+    [axs[i].set_xlim(0, n_bins) for i in range(n_trials)]
+    [axs[i].set_ylim(-10, 10) for i in range(n_trials)]
+    fig.tight_layout()