<jupyter_start><jupyter_text>A notebook where we synthesize regression models, when the true systems have an m function which is a Quadratic surface.<jupyter_code>%load_ext autoreload
%autoreload 2
import copy
import math

import matplotlib.pyplot as plt
import numpy as np
import torch

from janelia_core.ml.extra_torch_modules import DenseLNLNet
from janelia_core.ml.extra_torch_modules import QuadSurf
from janelia_core.ml.utils import list_torch_devices

from probabilistic_model_synthesis.gaussian_nonlinear_regression import compare_weight_prior_dists
from probabilistic_model_synthesis.gaussian_nonlinear_regression import align_low_d_spaces
from probabilistic_model_synthesis.gaussian_nonlinear_regression import fit_with_hypercube_priors
from probabilistic_model_synthesis.gaussian_nonlinear_regression import Fitter
from probabilistic_model_synthesis.gaussian_nonlinear_regression import generate_hypercube_prior_collection
from probabilistic_model_synthesis.gaussian_nonlinear_regression import GNLRMdl
from probabilistic_model_synthesis.simulation import generate_sum_of_bump_fcns_dist
from probabilistic_model_synthesis.simulation import sample_proj_data_from_cone
from probabilistic_model_synthesis.visualization import assign_colors_to_pts
from probabilistic_model_synthesis.visualization import plot_three_dim_pts
%matplotlib notebook<jupyter_output><empty_output><jupyter_text>Parameters go here<jupyter_code># Range of the number of input variables we observe from each individual - the actual number of variables we 
# observe from an individual will be pulled uniformly from this range (inclusive)
n_input_var_range = [10000, 11000]

# Range of the number of samples we observe from each individual - the actual number we observe from each individual
# will be unformly from this range (inclusive)
n_smps_range = [12000, 15000]

# ===============================================================================================
# Parameters for the true priors 

true_hc_params = {'n_divisions_per_dim': [10, 10], 
                  'dim_ranges': np.asarray([[-.1, 1.1],
                                       [-.1, 1.1]]),
                   'n_div_per_hc_side_per_dim': [1, 1]}

# The full options for setting up the prior on weights
true_w_prior_opts = {'mn_hc_params': true_hc_params, 'std_hc_params': true_hc_params, 
                     'min_std': .000001, 'mn_init': 0.0, 'std_init': .3}

# Options for the prior distributions on the scales, biases and psi
true_s_in_prior_opts = {'mn_mn': 1.0/np.sqrt(n_input_var_range[0]), 'mn_std': .00000001, 
                     'std_iv': .00001, 'std_lb': .0000001}

true_s_out_prior_opts = {'mn_mn': 1.0, 'mn_std': .00000001, 
                     'std_iv': .00001, 'std_lb': .0000001}

true_b_prior_opts = {'mn_mn': 0.0, 'mn_std': .00000001, 'std_iv': .001}
true_psi_prior_opts = {'conc_iv': 10.0, 'rate_iv': 1000.0, 'rate_ub': 10000.0}

# Parameters for generating shared m-module we use for fitting
m_n_layers = 5 #2
m_growth_rate = 5 #2

# Specify the number of intermediate variables and number of variables we predict
p = 2
d_pred = 3

# ===============================================================================================
# Parameters for the fit models

# The full options for setting up the prior on weights
fit_hc_params = {'n_divisions_per_dim': [80, 80], 
                 'dim_ranges': np.asarray([[-.1, 1.1],
                                       [-.1, 1.1]]),
                 'n_div_per_hc_side_per_dim': [1, 1]}

# The full options for setting up the prior on weights
fit_w_prior_opts = {'mn_hc_params': fit_hc_params, 'std_hc_params': fit_hc_params, 
                     'min_std': .000001, 'mn_init': 0.0, 'std_init': .3}

# Options for the prior distributions on the scales, biases and psi
fit_s_in_prior_opts = true_s_in_prior_opts
fit_s_out_prior_opts = true_s_out_prior_opts
fit_b_prior_opts = true_b_prior_opts
fit_psi_prior_opts = true_psi_prior_opts

# Options for posterior distribtions 
s_in_post_opts = {'mn_mn': 1.0/np.sqrt(n_input_var_range[0]), 'mn_std': .000001, 
               'std_iv': .00001, 'std_lb': .0000001}
s_out_post_opts = {'mn_mn': 1.0, 'mn_std': .000001, 
               'std_iv': .00001, 'std_lb': .0000001}

b_post_opts = {'mn_mn': 0.0, 'mn_std': .01}
psi_post_opts = {'rate_ub': 10000.0}

# Options for the densenet which makes up the shared-m module
dense_net_opts = {'n_layers': 2, 'growth_rate': 5, 'bias': True}

# ======================================================================================================
# Parameters for fitting - should be entered as lists, each entry corresponding to one round of fitting
sp_fit_opts = [{'n_epochs': 2000, 'milestones': [1500], 'update_int': 100, 'init_lr': .01}]
ip_fit_opts = [{'n_epochs': 20000, 'milestones': [100], 'update_int': 100, 'init_lr': .01}]<jupyter_output><empty_output><jupyter_text>Create true distributions that govern how systems under study are generated<jupyter_code>true_priors = generate_hypercube_prior_collection(p=p, d_pred=d_pred, 
                                                  w_prior_opts=true_w_prior_opts, 
                                                  s_in_prior_opts=true_s_in_prior_opts, 
                                                  b_in_prior_opts=true_b_prior_opts, 
                                                  s_out_prior_opts=true_s_out_prior_opts,
                                                  b_out_prior_opts=true_b_prior_opts, 
                                                  psi_prior_opts=true_psi_prior_opts)

# We replace the prior over weights 
true_priors.w_prior = generate_sum_of_bump_fcns_dist(n_bump_fcns=10, d_in=2, p=2, 
                                                     mn_m_std=1.0, std_m_std=.1, bump_w=.2)

#for d in range(2):
#    true_priors.w_prior.dists[d].mn_f.b_m.data[:] = torch.randn(true_priors.w_prior.dists[d].mn_f.b_m.data.shape)
#    true_priors.w_prior.dists[d].std_f[0].b_m.data[:] -= 1.0*torch.rand(true_priors.w_prior.dists[d].std_f[0].b_m.shape)
#    true_priors.w_prior.dists[d].std_f[0].b_m.data[torch.arange(0, 100, 3, dtype=torch.long)] = 1.0<jupyter_output><empty_output><jupyter_text>Define the true non-linear function relating projections of input variables to the mean of output variables<jupyter_code>class Selector(torch.nn.Module):
    
    def forward(self, x):
        return x[:,2].unsqueeze(-1)
m_true = torch.nn.Sequential(QuadSurf(torch.tensor([0.0, 0.0]), torch.tensor([1.0, -1.0])), Selector())
m_true = QuadSurf(torch.tensor([0.0, 0.0]), torch.tensor([1.0, -1.0]))
#m_true = torch.nn.Identity()<jupyter_output><empty_output><jupyter_text>Generate data  Generate properties<jupyter_code>n_individuals = 4
ind_n_vars = np.random.randint(n_input_var_range[0], n_input_var_range[1]+1, n_individuals)
ind_props = [torch.rand(size=[n_vars,2]) for n_vars in ind_n_vars]<jupyter_output><empty_output><jupyter_text>Generate true models for each individual<jupyter_code>with torch.no_grad():
    ind_true_mdls = [GNLRMdl(m=m_true, 
                     w=true_priors.w_prior.form_standard_sample(true_priors.w_prior.sample(props)),
                     s_in=true_priors.s_in_prior.form_standard_sample(true_priors.s_in_prior.sample(props)).squeeze(),
                     b_in=true_priors.b_in_prior.form_standard_sample(true_priors.b_in_prior.sample(props)).squeeze(),
                     s_out=true_priors.s_out_prior.form_standard_sample(true_priors.s_out_prior.sample(props)),
                     b_out=true_priors.b_out_prior.form_standard_sample(true_priors.b_out_prior.sample(props)),
                     psi=true_priors.psi_prior.form_standard_sample(true_priors.psi_prior.sample(props)))
                     for props in ind_props]
    
    if d_pred > 1:
        for mdl in ind_true_mdls:
            mdl.s_out.data = mdl.s_out.data.squeeze()
            mdl.b_out.data = mdl.b_out.data.squeeze()
            mdl.psi.data = mdl.psi.data.squeeze()<jupyter_output><empty_output><jupyter_text>Generate observations from true models<jupyter_code>ind_n_smps = np.random.randint(n_smps_range[0], n_smps_range[1]+1, n_individuals)
ind_data = [None]*n_individuals

start_ang = -math.pi
ang_inc = math.pi/2
cur_range = np.asarray([start_ang, start_ang+ang_inc])
ind_ang_ranges = np.zeros([4,2])
for i in range(n_individuals):
    ind_ang_ranges[i,:] = cur_range
    x_i = torch.tensor(sample_proj_data_from_cone(n_smps=ind_n_smps[i], 
                                     w=ind_true_mdls[i].w.detach().numpy(), ang_range=cur_range),
                       dtype=torch.float)
    with torch.no_grad():
        y_i = ind_true_mdls[i].sample(x=x_i)
    cur_range += ang_inc
    ind_data[i] = (x_i, y_i)
    
    print('Done generating data for individual ' + str(i) + '.')<jupyter_output>Done generating data for individual 0.
Done generating data for individual 1.
Done generating data for individual 2.
Done generating data for individual 3.<jupyter_text>Synthesize models (fit models to the data together)<jupyter_code>fit_rs = fit_with_hypercube_priors(data=ind_data, props=ind_props, p=2, 
                                   w_prior_opts=fit_w_prior_opts, 
                                   s_in_prior_opts=fit_s_in_prior_opts,
                                   b_in_prior_opts=fit_b_prior_opts,
                                   s_out_prior_opts=fit_s_out_prior_opts,
                                   b_out_prior_opts=fit_b_prior_opts,
                                   psi_prior_opts=fit_psi_prior_opts,
                                   s_in_post_opts=s_in_post_opts,
                                   b_in_post_opts=b_post_opts,
                                   s_out_post_opts=s_out_post_opts,
                                   b_out_post_opts=b_post_opts,
                                   psi_post_opts=psi_post_opts,
                                   dense_net_opts=dense_net_opts, 
                                   sp_fit_opts=sp_fit_opts, 
                                   ip_fit_opts=ip_fit_opts)<jupyter_output>Found 1 GPUs
Beginning SP fitting.

=========== EPOCH 0 COMPLETE ===========
Obj: 5.34e+06
----------------------------------------
NELL: 6.22e+04, 6.93e+04, 6.89e+04, 6.16e+04
W KL: 0.00e+00, 0.00e+00, 0.00e+00, 0.00e+00
S_in KL: 5.07e+05, 5.06e+05, 5.07e+05, 5.07e+05
B_in KL: 3.08e+02, 1.25e+02, 1.34e+02, 2.81e+02
S_out KL: 7.60e+05, 7.60e+05, 7.60e+05, 7.60e+05
B_out KL: 1.96e+02, 3.86e+02, 2.28e+02, 5.24e+02
Psi KL: 2.79e+03, 2.82e+03, 2.82e+03, 2.79e+03
----------------------------------------
LR: 0.01
Elapsed time (secs): 0.40133213996887207
----------------------------------------
CPU cur memory used (GB): 4.83e+00
GPU_0 cur memory used (GB): 2.22e+00, max memory used (GB): 2.22e+00

=========== EPOCH 100 COMPLETE ===========
Obj: -1.09e+04
----------------------------------------
NELL: -1.76e+03, -3.10e+03, -9.25e+03, 4.41e+02
W KL: 0.00e+00, 0.00e+00, 0.00e+00, 0.00e+00
S_in KL: 2.80e+02, 1.26e+03, 2.44e+02, 6.77e+01
B_in KL: 1.79e+01, 6.87e+01, 1.18e+02, 9.88e+01
S_out KL: 2.[...]<jupyter_text>Look at model fits<jupyter_code>exam_ind = 0
exam_type = 'ip'
if exam_type == 'sp':
    exam_priors = fit_rs['sp']['priors']
    exam_posts = [coll.posteriors for coll in fit_rs['sp']['vi_collections']]
    exam_vi_collections = fit_rs['sp']['vi_collections']
    exam_mdls = [coll.mdl for coll in fit_rs['sp']['vi_collections']]
    exam_logs = fit_rs['sp']['logs']
else:
    exam_priors = fit_rs['ip']['priors']
    exam_posts = [coll.posteriors for coll in fit_rs['ip']['vi_collections']]
    exam_vi_collections = fit_rs['ip']['vi_collections']
    exam_mdls = [coll.mdl for coll in fit_rs['ip']['vi_collections']]
    exam_logs = fit_rs['ip']['logs']<jupyter_output><empty_output><jupyter_text>Look at fitting logs<jupyter_code>for log in exam_logs:
    Fitter.plot_log(log)
exam_w = exam_posts[exam_ind].w_post(ind_props[exam_ind])
exam_s_in = exam_posts[exam_ind].s_in_post(ind_props[exam_ind]).squeeze()
exam_b_in = exam_posts[exam_ind].b_in_post(ind_props[exam_ind]).squeeze()
exam_s_out = exam_posts[exam_ind].s_out_post(ind_props[exam_ind]).squeeze()
exam_b_out = exam_posts[exam_ind].b_out_post(ind_props[exam_ind]).squeeze()
exam_psi = exam_posts[exam_ind].psi_post.dists[0].mode(ind_props[exam_ind]).squeeze()

exam_mdl = GNLRMdl(m = exam_vi_collections[exam_ind].mdl.m, w=exam_w, s_in=exam_s_in, b_in=exam_b_in, 
                  s_out=exam_s_out, b_out=exam_b_out, psi=exam_psi)

plt.figure()
GNLRMdl.compare_mdls(ind_true_mdls[exam_ind], exam_mdl)<jupyter_output><empty_output><jupyter_text>Look at predicted mean values over the full range of input values for the exam_ind individual<jupyter_code>n_exam_smps = 10000
x_exam = torch.randn([n_exam_smps, ind_n_vars[exam_ind]])
#x_exam = ind_data[exam_ind][0]
with torch.no_grad():
    
    # Determine which of our test points are in and out of the training distribution for this individual 
    x_true_proj = ind_true_mdls[exam_ind].project(x_exam).numpy()
    x_true_angles = np.asarray([math.atan2(vl[0], vl[1]) for vl in x_true_proj])
    x_within_train_dist = np.logical_and(x_true_angles > ind_ang_ranges[exam_ind, 0], 
                                         x_true_angles < ind_ang_ranges[exam_ind, 1])
    x_outside_train_dist = np.logical_not(x_within_train_dist)
    
    # Get true mean and predicted mean for each data point for this individual 
    true_mns = ind_true_mdls[exam_ind].m(ind_true_mdls[exam_ind].project(x_exam)).numpy()
    
    x_pred_proj = exam_mdls[exam_ind].project(x=x_exam, w=exam_w, s_in=exam_s_in, b_in=exam_b_in)
    pred_mns = exam_mdls[exam_ind].m(x_pred_proj).numpy()<jupyter_output><empty_output><jupyter_text>Look at point-by-point correspondance<jupyter_code>cmp_dim = 2
plt.figure()
plt.plot(true_mns[x_within_train_dist, cmp_dim], pred_mns[x_within_train_dist, cmp_dim], 'k.')
plt.plot(true_mns[x_outside_train_dist, cmp_dim], pred_mns[x_outside_train_dist, cmp_dim], 'r.')
plt.legend(['Within Train Dist', 'Outside Train Dist'])<jupyter_output><empty_output><jupyter_text>Look at correspondance of the m functions  Look at shape of surface of true and predict means as a function of the true low-d projections<jupyter_code>clrs = np.zeros([n_exam_smps,4])
clrs[:,-1] = 1.0
clrs[x_outside_train_dist, 0] = 1.0

plt.figure()
a = plt.subplot(1,2,1, projection='3d')
plot_three_dim_pts(np.concatenate([x_true_proj, true_mns[:, cmp_dim:cmp_dim+1]], axis=1), 
                   clrs=clrs, a=a)

a = plt.subplot(1,2,2, projection='3d')
plot_three_dim_pts(np.concatenate([x_true_proj, pred_mns[:, cmp_dim:cmp_dim+1]], axis=1), 
                   clrs=clrs, a=a)<jupyter_output><empty_output><jupyter_text>Look at shape of surface of true and predicted means when predicted means are also a function of the predicted low-d projections.This is a more stringent test of identifiability than the above plots, becuase the learned models could "scramble" the projection to the low-s space and then the learned m-module could sort it out so that mean predictions when viewed as a function from input variables to means is correct, even though the intermediate transformation to low-d variables, and therefore the mapping from low-d variables to means, is not the same as the true models<jupyter_code># Get the estimate low-d projections in the space best aligned to truth

align_vls = align_low_d_spaces(w_0=ind_true_mdls[exam_ind].w.detach().numpy(),
                               s_in_0=ind_true_mdls[exam_ind].s_in.detach().numpy(),
                               b_in_0=ind_true_mdls[exam_ind].b_in.detach().numpy(),
                               w_1=exam_w.detach().numpy(),
                               s_in_1=exam_s_in.detach().numpy(),
                               b_in_1=exam_b_in.detach().numpy(),
                               z_1=x_pred_proj.numpy())

x_pred_proj_aligned = align_vls[-1]
# Generate plots

plt.figure()
a = plt.subplot(1,2,1, projection='3d')
plot_three_dim_pts(np.concatenate([x_true_proj, true_mns[:, cmp_dim:cmp_dim+1]], axis=1), clrs=clrs, a=a)

a = plt.subplot(1,2,2, projection='3d')
plot_three_dim_pts(np.concatenate([x_pred_proj_aligned, pred_mns[:, cmp_dim:cmp_dim+1]], axis=1), clrs=clrs, a=a)<jupyter_output><empty_output><jupyter_text>Compare true and fit prior distributions over weights<jupyter_code>compare_weight_prior_dists(w0_prior=true_priors.w_prior, w1_prior=exam_priors.w_prior, 
                           dim_0_range=[0, 1], dim_1_range=[0, 1], n_pts_per_dim=[200, 200])<jupyter_output>[[ 4.33991276e-01 -7.27448811e-01]
 [ 6.83981162e-01  4.84184673e-01]
 [-1.52066434e-03  7.22798178e-04]]
[[ 0.43399128 -0.72744881]
 [ 0.68398116  0.48418467]]
{'fraction_non_alphanumeric': 0.13721139430284857, 'fraction_numerical': 0.04305847076461769, 'mean_word_length': 3.510684338652962, 'pattern_counts': {'":': 0, '<': 59, '<?xml version=': 0, '>': 60, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 3, 'word_list_counts': {'cursed_substrings.json': 33, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'chain_length': '15', 'dir': 'jupyter-structured-clean-dedup', 'id': '530916', 'license': 'MIT', 'n_tokens_mistral': 6624, 'n_tokens_neox': 5984, 'n_words': 1123, 'path': 'gnlr_simulations/QuadSurface_Synthesis.ipynb', 'repo_name': 'wbishopJanelia/probabilistic_model_synthesis', 'size': '16675'}