<jupyter_start><jupyter_text>Load Data<jupyter_code>x,y,x_test,y_test,y_plot,n,d = load_traintest_hier(1,"radon",100)
K = np.shape(np.unique(x[:,1]))[0]
print(K)<jupyter_output>85<jupyter_text>Conformal<jupyter_code>beta_post = jnp.array(np.load("samples/beta_post_hier_{}.npy".format('radon')))[0]
intercept_post = jnp.array(np.load("samples/intercept_post_hier_{}.npy".format('radon')))[0]
sigma_post = jnp.array(np.load("samples/sigma_post_hier_{}.npy".format('radon')))[0]<jupyter_output>WARNING:absl:No GPU/TPU found, falling back to CPU. (Set TF_CPP_MIN_LOG_LEVEL=0 and rerun for more info.)<jupyter_text>Define x_test<jupyter_code>x_test = np.zeros((2*K,2))
for k in range(K):
    x_test[2*k:2*k + 2,1] = k
    x_test[2*k,0] = 0
    x_test[2*k+1,0]= 1

n_test = np.shape(x_test)[0]

pi_cb = np.zeros((n_test,np.shape(y_plot)[0]))
region_cb = np.zeros((n_test,np.shape(y_plot)[0]))
length_cb = np.zeros(n_test)
band_bayes = np.zeros((n_test,2))
length_bayes = np.zeros(n_test)

dy = y_plot[1]- y_plot[0]<jupyter_output><empty_output><jupyter_text>Compute rank plots<jupyter_code>groups_train = np.unique(x[:,1]).astype('int')
K_train = np.size(groups_train)
n_groups = np.zeros(K_train)

for k in (range(K_train)):
    ind_group = (x[:,1] == groups_train[k])
    n_groups[k] =np.sum(ind_group)
    
alpha_k = 1.1*(1/(n_groups + 1))
print(alpha_k)
#Define likelihood from posterior samples
start = time.time()
@jit
def normal_likelihood_cdf(y,x):
    group = x[:,-1].astype('int32') 
    x_0 = x[:,0] 
    return norm.cdf(y,loc = beta_post[:,group] * x_0.transpose()+ intercept_post[:,group]\
                    ,scale = sigma_post) #compute likelihood samples
#Precompute cdfs
cdf_test =  normal_likelihood_cdf(y_plot.reshape(-1,1,1),x_test)

#Compute conformal regions
for i in tqdm(range(n_test)):
    group_ind = x_test[i,1].astype('int')
    
    band_bayes[i] = bmcmc.compute_bayes_band_MCMC(alpha_k[group_ind],y_plot,cdf_test[:,:,i])
    length_bayes[i] = np.abs(band_bayes[i,1]- band_bayes[i,0]) 
end = time.time()
print('Bayes took {}'.format(end - start))
   
start = time.time()
@jit
def normal_loglikelihood(y,x):
    group = x[:,-1].astype('int32') 
    x_0 = x[:,0] 
    return norm.logpdf(y,loc = beta_post[:,group]* x_0.transpose() + intercept_post[:,group]\
                       ,scale = sigma_post) #compute likelihood samples

#Compute loglikelihood across groups
logp_samp_n = []

for k in (range(K_train)):
    ind_group = (x[:,1] == groups_train[k])
    logp_samp_n.append(normal_loglikelihood(y[ind_group],x[ind_group]))

logwjk = normal_loglikelihood(y_plot.reshape(-1,1,1),x_test)

#Compute conformal regions
for i in tqdm(range(n_test)):
    group_ind = x_test[i,1].astype('int') 
    pi_cb[i] = cb.compute_rank_IS(logp_samp_n[group_ind],logwjk[:,:,i])
    region_cb[i] = cb.compute_cb_region_IS(alpha_k[group_ind],logp_samp_n[group_ind],logwjk[:,:,i])
    length_cb[i] = np.sum(region_cb[i])*dy
end = time.time()
print('Conformal took {}'.format(end-start))
length_bayes_grp = np.zeros(K_train)
length_cb_grp = np.zeros(K_train)
for k in range(K_train):
    length_bayes_grp[k] = np.mean(length_bayes[2*k:2*k+2])
    length_cb_grp[k] = np.mean(length_cb[2*k:2*k+2])
print(np.mean(length_bayes_grp))
print(np.mean(length_cb_grp))<jupyter_output>2.1739750038175023
2.6609625668449177<jupyter_text>Analyze rank plots<jupyter_code>np.where(n_groups ==1)
f = plt.figure(figsize = (12,3))
plt.subplot(1,2,1)
k = 0
i = 2*k
county = x_test[i,1].astype('int')
n_k = n_groups[county].astype('int')
plt.plot(y_plot,pi_cb[i]/(n_k+1),label = r'$x ={}$'.format(x_test[i,0].astype('int')))

county = x_test[i+1,1].astype('int')
n_k = n_groups[county].astype('int')
plt.plot(y_plot,pi_cb[i+1]/(n_k+1),label = r'$x ={}$'.format(x_test[i+1,0].astype('int')))
plt.xlabel('$y$',fontsize = 12)
plt.ylabel(r'$\pi_j(y)$')
plt.ylim(-0.05,1.1)
plt.title('County {}: $n_j = {}$'.format(county,n_k))
plt.legend()

plt.subplot(1,2,2)
k = 1
i = 2*k
county = x_test[i,1].astype('int')
n_k = n_groups[county].astype('int')
plt.plot(y_plot,pi_cb[i]/(n_k+1),label = r'$x ={}$'.format(x_test[i,0].astype('int')))

county = x_test[i+1,1].astype('int')
n_k = n_groups[county].astype('int')
plt.plot(y_plot,pi_cb[i+1]/(n_k+1),label = r'$x ={}$'.format(x_test[i+1,0].astype('int')))
plt.xlabel('$y$',fontsize = 12)
plt.ylabel(r'$\pi_j(y)$')
plt.ylim(-0.05,1.1)
plt.title('County {}: $n_j = {}$'.format(county,n_k))
f.savefig('plots/radon_big.pdf', bbox_inches='tight')
f = plt.figure(figsize = (12,3))
plt.subplot(1,2,1)
k = 41
i = 2*k
county = x_test[i,1].astype('int')
n_k = n_groups[county].astype('int')
plt.plot(y_plot,pi_cb[i]/(n_k+1),label = r'$x ={}$'.format(x_test[i,0].astype('int')))

county = x_test[i+1,1].astype('int')
n_k = n_groups[county].astype('int')
plt.plot(y_plot,pi_cb[i+1]/(n_k+1),label = r'$x ={}$'.format(x_test[i+1,0].astype('int')))
plt.xlabel('$y$',fontsize = 12)
plt.ylabel(r'$\pi_j(y)$')
plt.ylim(-0.05,1.1)
plt.title('County {}: $n_j = {}$'.format(county,n_k))
plt.legend()

plt.subplot(1,2,2)
k = 49
i = 2*k
county = x_test[i,1].astype('int')
n_k = n_groups[county].astype('int')
plt.plot(y_plot,pi_cb[i]/(n_k+1),label = r'$x ={}$'.format(x_test[i,0].astype('int')))

county = x_test[i+1,1].astype('int')
n_k = n_groups[county].astype('int')
plt.plot(y_plot,pi_cb[i+1]/(n_k+1),label = r'$x ={}$'.format(x_test[i+1,0].astype('int')))
plt.xlabel('$y$',fontsize = 12)
plt.ylabel(r'$\pi_j(y)$')
plt.ylim(-0.05,1.1)
plt.title('County {}: $n_j = {}$'.format(county,n_k))
f.savefig('plots/radon_small.pdf', bbox_inches='tight')<jupyter_output><empty_output><jupyter_text>LMER<jupyter_code>import statsmodels.formula.api as smf
d = {'log_radon': y, 'floor': x[:,0],'county':x[:,1]}
data = pd.DataFrame(data=d)

md = smf.mixedlm("log_radon ~ floor", data, groups=data["county"], re_formula="~floor")
%time mdf = md.fit(method=["lbfgs"])
print(mdf.summary())<jupyter_output>CPU times: user 1.52 s, sys: 250 ms, total: 1.77 s
Wall time: 301 ms
           Mixed Linear Model Regression Results
============================================================
Model:              MixedLM  Dependent Variable:  log_radon 
No. Observations:   919      Method:              REML      
No. Groups:         85       Scale:               0.5138    
Min. group size:    1        Log-Likelihood:      -1047.2140
Max. group size:    116      Converged:           Yes       
Mean group size:    10.8                                    
------------------------------------------------------------
                  Coef.  Std.Err.   z    P>|z| [0.025 0.975]
------------------------------------------------------------
Intercept          1.495    0.052 28.576 0.000  1.392  1.597
floor             -0.655    0.086 -7.650 0.000 -0.823 -0.487
Group Var          0.114    0.047                           
Group x floor Cov -0.041    0.055                           
floor Var          0.108    [...]
{'fraction_non_alphanumeric': 0.22394549934773156, 'fraction_numerical': 0.04478909986954631, 'mean_word_length': 4.184072126220887, 'pattern_counts': {'":': 0, '<': 21, '<?xml version=': 0, '>': 22, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 1, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'chain_length': '6', 'dir': 'jupyter-structured-clean-dedup', 'id': '512552', 'license': 'MIT', 'n_tokens_mistral': 3199, 'n_tokens_neox': 3021, 'n_words': 411, 'path': 'Radon.ipynb', 'repo_name': 'edfong/conformal_bayes', 'size': '6899'}