<jupyter_start><jupyter_text>A notebook to generate plots from the post processed results of an across condition analysis<jupyter_code>%load_ext autoreload
%autoreload 2
from pathlib import Path
import pickle
import re

from IPython.display import display
import matplotlib
import matplotlib.pyplot as plt
import numpy as np
import pandas as pd
import torch

from janelia_core.stats.regression import r_squared
from janelia_core.visualization.matrix_visualization import colorized_tbl
%matplotlib notebook<jupyter_output><empty_output><jupyter_text>Parameters go here<jupyter_code># Top-level directory holding the results
base_dir = r'/groups/bishop/bishoplab/projects/probabilistic_model_synthesis/results/real_data/gnldr/across_cond_transfer_analysis/v0'

# Name used for all files containing post-processed results
pp_file = r'pp_test_results.pt'

# Specify the type of models we assess performance for - can be 'sp' or 'ip'
mdl_type = 'ip'

# Specify conditions we test on 
test_conds = ['omr_forward', 'omr_right', 'omr_left']

# Specify folder to save images in 
save_folder = r'/groups/bishop/bishoplab/projects/probabilistic_model_synthesis/results/real_data/gnldr/across_cond_transfer_analysis/v0/imgs'<jupyter_output><empty_output><jupyter_text>Load all results<jupyter_code>train_conds = [cond.name for cond in Path(base_dir).iterdir() if cond.is_dir()]
test_subjs = np.sort([int(re.search('.*_(\d+)', subj.name)[1]) 
                      for subj in (Path(base_dir) / train_conds[0]).iterdir()])
FIT_TYPES = ['single_cond', 'multi_cond']
rs = dict()
for cond in train_conds:
    rs[cond] = dict()
    for subj in test_subjs:
        rs[cond][subj] = dict()
        for fit_type in FIT_TYPES:
            fit_type_dir = Path(base_dir) / cond / ('subj_' + str(subj)) / fit_type
            fit_type_file = fit_type_dir / pp_file
            
            rs[cond][subj][fit_type] = torch.load(fit_type_file)
            
            # Print some diagonstic information 
            print('Results for subject ' + str(subj) + ', train condition: ' + cond + 
                  ', fit type: ' + fit_type + ' ***')
            print('Best CP Ind: ' + str(rs[cond][subj][fit_type][mdl_type]['early_stopping']['best_cp_ind']))<jupyter_output>Results for subject 8, train condition: omr_f_ns, fit type: single_cond ***
Best CP Ind: 20
Results for subject 8, train condition: omr_f_ns, fit type: multi_cond ***
Best CP Ind: 20
Results for subject 9, train condition: omr_f_ns, fit type: single_cond ***
Best CP Ind: 2
Results for subject 9, train condition: omr_f_ns, fit type: multi_cond ***
Best CP Ind: 7
Results for subject 11, train condition: omr_f_ns, fit type: single_cond ***
Best CP Ind: 1
Results for subject 11, train condition: omr_f_ns, fit type: multi_cond ***
Best CP Ind: 1
Results for subject 8, train condition: omr_r_ns, fit type: single_cond ***
Best CP Ind: 2
Results for subject 8, train condition: omr_r_ns, fit type: multi_cond ***
Best CP Ind: 2
Results for subject 9, train condition: omr_r_ns, fit type: single_cond ***
Best CP Ind: 2
Results for subject 9, train condition: omr_r_ns, fit type: multi_cond ***
Best CP Ind: 2
Results for subject 11, train condition: omr_r_ns, fit type: single_cond ***
Best CP Ind: 8[...]<jupyter_text>Now assess test performance<jupyter_code>subj_elbos = dict()
for subj in test_subjs:
    subj_elbos[subj] = {'multi_cond': pd.DataFrame(np.zeros([3,3]), index=train_conds, columns=test_conds), 
                        'single_cond': pd.DataFrame(np.zeros([3,3]), index=train_conds, columns=test_conds)}
for ts_i, test_cond in enumerate(test_conds):
    print('**** Test Condition: ' + test_cond + ' ****')
    for s_i, subj in enumerate(test_subjs):
        print('Subject: ' + str(subj))
        print('Single ELBO, Multi ELBO, Delta, Training Condition')
        for tr_i, train_cond in enumerate(train_conds):
            single_period_elbo_vls = rs[train_cond][subj]['single_cond'][mdl_type]['period_elbo_vls'][subj]
            multi_period_elbo_vls = rs[train_cond][subj]['multi_cond'][mdl_type]['period_elbo_vls'][subj]
            
            if single_period_elbo_vls[test_cond] is not None:
                single_elbo = single_period_elbo_vls[test_cond]['vl']['elbo'].item()/single_period_elbo_vls[test_cond]['n_smps']
                multi_elbo = multi_period_elbo_vls[test_cond]['vl']['elbo'].item()/multi_period_elbo_vls[test_cond]['n_smps']
                delta_elbo = multi_elbo - single_elbo
                print('{:.2E}'.format(single_elbo) + ', ' + '{:.2E}'.format(multi_elbo) + ', ' + 
                      '{:.2E}'.format(delta_elbo) + ', ' + train_cond)
                
                subj_elbos[subj]['single_cond'][test_cond][train_cond] = single_elbo
                subj_elbos[subj]['multi_cond'][test_cond][train_cond] = multi_elbo
for subj in test_subjs:
    subj_elbos[subj]['multi_cond'] = subj_elbos[subj]['multi_cond'].rename(
        columns={'omr_forward': 'F', 'omr_right': 'R', 'omr_left': 'L'}, 
        index={'omr_f_ns': 'F', 'omr_r_ns': 'R', 'omr_l_ns': 'L'})
    
    subj_elbos[subj]['single_cond'] = subj_elbos[subj]['single_cond'].rename(
        columns={'omr_forward': 'F', 'omr_right': 'R', 'omr_left': 'L'}, 
        index={'omr_f_ns': 'F', 'omr_r_ns': 'R', 'omr_l_ns': 'L'})<jupyter_output><empty_output><jupyter_text>Make tables of improvements in Normalized ELBO values when models are synthesized using data of different behaviors across fish<jupyter_code>vl_range = 15000
cmap = matplotlib.colors.LinearSegmentedColormap.from_list(name='red_to_green', colors=[(0,  [0.0, 0.0, 1.0, 1.0]),
                                                               (.5, [0.0, 0.0, 0.0, 0.0]),
                                                               (1.0, [1.0, 0.0, 0.0, 1.0])], N=1024)
for subj in test_subjs:
    elbo_improvements = subj_elbos[subj]['multi_cond'] - subj_elbos[subj]['single_cond']  
    fig = plt.figure()
    ax = plt.subplot(1,1,1)
    colorized_tbl(tbl=elbo_improvements.to_numpy(), cmap='bwr', vmin=-vl_range, vmax=vl_range, 
                  dim_0_lbls=elbo_improvements.index.to_list(), 
                  dim_1_lbls=elbo_improvements.columns.to_list(), 
                  tbl_fontsize=16, label_fontsize=16,
                  ax=ax)
    
    save_file = 'elbo_improvements_subj_' + str(subj) + '.eps'
    save_path = Path(save_folder) / save_file
    plt.savefig(save_path)<jupyter_output><empty_output><jupyter_text>Generate scatter plots<jupyter_code>subj_clrs = {8: [1, 0, 0], 
             9: [0, 1, 0], 
            11: [0, 0, 1]}

cond_markers = {'F': 'o', 'R': 's', 'L': 'd'}

plt_conds = ['F', 'R', 'L']
plt.figure()
ax = plt.subplot(1,1,1)
for subj in test_subjs:
    for test_cond in plt_conds:
        diff_test_conds = list(set(plt_conds) - set([test_cond]))
        
        multi_cond_diff_test_conds = .5*(subj_elbos[subj]['multi_cond'][test_cond][diff_test_conds[0]] + 
                                          subj_elbos[subj]['multi_cond'][test_cond][diff_test_conds[1]])
        
        multi_cond_same_test_cond =  subj_elbos[subj]['multi_cond'][test_cond][test_cond]
        
        single_cond_diff_test_conds = .5*(subj_elbos[subj]['single_cond'][test_cond][diff_test_conds[0]] + 
                                          subj_elbos[subj]['single_cond'][test_cond][diff_test_conds[1]])
        
        single_cond_same_test_cond =  subj_elbos[subj]['single_cond'][test_cond][test_cond]
        
        plt.plot(single_cond_diff_test_conds, multi_cond_diff_test_conds, marker=cond_markers[test_cond], 
                 color=np.asarray(subj_clrs[subj] + [1]))
        
        plt.plot(single_cond_same_test_cond, multi_cond_same_test_cond, marker=cond_markers[test_cond], 
                 color=np.asarray(subj_clrs[subj] + [.25]))


plt.plot([-200000, -8000], [-200000, -8000], 'k--')
ax.axis('equal')<jupyter_output><empty_output>
{'fraction_non_alphanumeric': 0.1324143164284811, 'fraction_numerical': 0.016441128114329077, 'mean_word_length': 3.5136986301369864, 'pattern_counts': {'":': 0, '<': 24, '<?xml version=': 0, '>': 24, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 8, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'chain_length': '6', 'dir': 'jupyter-structured-clean-dedup', 'id': '645196', 'license': 'MIT', 'n_tokens_mistral': 2925, 'n_tokens_neox': 2796, 'n_words': 559, 'path': 'real_data/gnldr/across_cond_transfer_analysis/make_across_cond_plots.ipynb', 'repo_name': 'wbishopJanelia/probabilistic_model_synthesis', 'size': '7907'}