<jupyter_start><jupyter_text>1. Load Test data<jupyter_code>dat_dir = '/scratch/network/chhahn/simbig/'
k_test = np.loadtxt(os.path.join(dat_dir, 'simbig.cmass_sgc.test.k.dat'), skiprows=1)
y_test = np.loadtxt(os.path.join(dat_dir, 'simbig.cmass_sgc.test.nuis.theta.dat'), skiprows=1)
x_test = np.loadtxt(os.path.join(dat_dir, 'simbig.cmass_sgc.test.nuis.p0k.dat'), skiprows=1)<jupyter_output><empty_output><jupyter_text>1. Load ANPE<jupyter_code>lower_bounds = torch.tensor([0.1, 0.03, 0.5, 0.8, 0.6, 13., 0.32, 13.14, 14.0, 0.58, -1e4])
upper_bounds = torch.tensor([0.4, 0.07, 0.9, 1.2, 1., 13.06, 0.44, 13.4, 14.14, 0.94, 1e4])

prior = Ut.BoxUniform(low=lower_bounds, high=upper_bounds, device='cpu')
nhidden = [299, 415, 373, 259]
nblocks = [3, 4, 9, 3]

phats = [] 
for i, kmax in enumerate([0.1, 0.2, 0.3, 0.4]): 
    fanpe = os.path.join(dat_dir, f'anpe.{nhidden[i]}x{nblocks[i]}.p0k.kmax{kmax}.pt')
    
    anpe = Inference.SNPE(prior=prior, 
                      density_estimator=Ut.posterior_nn('maf', hidden_features=nhidden[i], num_transforms=nblocks[i]), 
                      device='cpu')
    # load ANPE
    klim = (k_test < kmax)
    anpe.append_simulations(
        torch.as_tensor(y_test.astype(np.float32)), 
        torch.as_tensor(x_test[:,klim].astype(np.float32)))

    phat = anpe._build_neural_net(
        torch.as_tensor(y_test.astype(np.float32)), 
        torch.as_tensor(x_test[:,klim].astype(np.float32)))
    phat.load_state_dict(torch.load(fanpe, map_location=torch.device('cpu')))

    anpe._x_shape = Ut.x_shape_from_simulation(torch.as_tensor(y_test.astype(np.float32)))

    phats.append(anpe.build_posterior(phat))<jupyter_output><empty_output><jupyter_text>4. Deploy on SIMBIG CMASS test mocksAs a validation of our neural posterior estimation, we deploy on the SIMBIG CMASS test galaxy catalogs, where we know the true cosmology and HOD parameter values. The SIMBIG CMASS test mocks are constructed at the fiducial cosmology with a limited HOD range.<jupyter_code>pp_thetas_kmax = [] 
rank_thetas_kmax = []
ranks_kmax = []

for ii, kmax in enumerate([0.1, 0.2, 0.3, 0.4]): 
    print(f'k_max = {kmax}')
    klim = (k_test < kmax)
    
    pp_thetas, rank_thetas, ranks = [], [], []
    for i in trange(x_test.shape[0]): 
        _x_test = x_test[i][klim]
        _y_test = y_test[i]

        # sample posterior p(theta | x_test_i)
        mcmc_npe = phats[ii].sample((10000,), x=torch.as_tensor(_x_test.astype(np.float32)).to('cpu'), show_progress_bars=False)
        mcmc_npe = np.array(mcmc_npe.detach().to('cpu'))

        if i == 0: 
            fig = DFM.corner(mcmc_npe[:,:5], color='C0', 
                             levels=[0.64, 0.95], 
                             range=[(low, high) for low, high in zip(np.array(lower_bounds)[:5], np.array(upper_bounds))], 
                             labels=[r'$\Omega_m$', r'$\Omega_b$', r'$h$', r'$n_s$', r'$\sigma_8$'])
            DFM.overplot_lines(fig, _y_test[:5], color='C1')
            DFM.overplot_points(fig, [_y_test[:5]], color='C1', markersize=20)
            plt.show()

        # calculate percentile score and rank
        pp_theta, rank_theta = [], []
        for itheta in range(5): 
            pp_theta.append(float(stats.percentileofscore(mcmc_npe[:,itheta], _y_test[itheta]))/100.)
            rank_theta.append(np.sum(mcmc_npe[:,itheta] < _y_test[itheta]))
        pp_thetas.append(pp_theta)
        rank_thetas.append(rank_theta)

        # calculate log p(theta | x_test_i)
        logp_samples = phats[ii].log_prob(torch.as_tensor(mcmc_npe.astype(np.float32)).to('cpu'), 
                                                x=torch.as_tensor(_x_test.astype(np.float32)).to('cpu'))
        logp_truth = phats[ii].log_prob(_y_test.astype(np.float32).reshape(1, -1), 
                                              x=torch.as_tensor(_x_test.astype(np.float32)).to('cpu'))
        ranks.append(np.sum(np.array(logp_samples) < np.array(logp_truth)))

    pp_thetas_kmax.append(np.array(pp_thetas))
    rank_thetas_kmax.append(np.array(rank_thetas))
    ranks_kmax.append(np.array(ranks))
theta_lbls = [r'$\Omega_m$', r'$\Omega_b$', r'$h$', r'$n_s$', r'$\sigma_8$']

fig = plt.figure(figsize=(15,3))
for kmax, rank_thetas in zip([0.1, 0.2, 0.3, 0.4], rank_thetas_kmax): 

    for i in range(5): 
        sub = fig.add_subplot(1,5,i+1)
        sub.hist(rank_thetas[:,i]/10000., density=True, histtype='step', linewidth=2, label=theta_lbls[i])
        sub.plot([0., 1.], [1., 1.], c='k', ls='--')

        sub.text(0.05, 0.05, theta_lbls[i], fontsize=20, transform=sub.transAxes)
        sub.set_xlim(0, 1.)
        sub.set_ylim(0., 2.)
        sub.set_yticklabels([])
        sub.set_xticklabels([])
fig = plt.figure(figsize=(5,5))
sub = fig.add_subplot(111)
for kmax, ranks in zip([0.1, 0.2, 0.3, 0.4], ranks_kmax): 
    sub.hist(ranks/10000., density=True, histtype='step', linewidth=2, label=r'$k_{\rm max} = %.1f$'% kmax)
    sub.plot([0., 1.], [1., 1.], c='k', ls='--')
sub.legend(loc='upper right', fontsize=20)
sub.set_xlim(0, 1.)
sub.set_ylim(0., 2.)
sub.set_yticklabels([])
sub.set_xticklabels([])<jupyter_output><empty_output><jupyter_text>5. Deploy on observationsFinally, lets deploy it on the observed CMASS-SGC galaxy catalog!<jupyter_code>k_obs = np.loadtxt(os.path.join(dat_dir, 'obs.cmass_sgc.k.dat'), skiprows=1)
x_obs = np.loadtxt(os.path.join(dat_dir, 'obs.cmass_sgc.p0k.w_nofc.dat'), skiprows=1)
for kmax, phat in zip([0.1, 0.2, 0.3, 0.4], phats): 
    klim_obs = (k_obs < kmax)
    
    mcmc_npe = phat.sample((10000,), x=torch.as_tensor(x_obs[klim_obs].astype(np.float32)).to('cpu'))
    fig = DFM.corner(np.array(mcmc_npe[:,:5].detach().to('cpu')), color='C0', 
                     levels=[0.68, 0.95], 
                     range=[(low, high) for low, high in zip(np.array(lower_bounds)[:5], np.array(upper_bounds))], 
                     labels=[r'$\Omega_m$', r'$\Omega_b$', r'$h$', r'$n_s$', r'$\sigma_8$'])

    DFM.overplot_lines(fig, [0.302, None, 0.658, None, None], label="d'Amico+(2020) CMASS SGC", color='C1')

    # Ivanov+(2020) Table 10
    DFM.overplot_lines(fig, [0.309, None, 0.665, None, 0.646], label="Ivanov+(2020) high-z SGC", color='k')

    # Kobayashi+(2021) Table 4
    DFM.overplot_lines(fig, [0.3088, None, None, None, 0.7838], label="Kobayashi+(2021) CMASSLOWZTOT high-z", color='C2')
    plt.legend()
    fig.suptitle('Compressed $P_0(k < 0.5; z=%.1f)$ CMASS-SGC' % kmax, fontsize=25)
    plt.show()<jupyter_output><empty_output>
{'fraction_non_alphanumeric': 0.1815242494226328, 'fraction_numerical': 0.04680523479599692, 'mean_word_length': 3.3509711989283324, 'pattern_counts': {'":': 0, '<': 23, '<?xml version=': 0, '>': 17, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 4, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'chain_length': '4', 'dir': 'jupyter-structured-clean-dedup', 'id': '418836', 'license': 'MIT', 'n_tokens_mistral': 2867, 'n_tokens_neox': 2662, 'n_words': 415, 'path': 'cmass/nb/plk_6_valid_sbi_p0k_kmax.ipynb', 'repo_name': 'changhoonhahn/boss-sbi', 'size': '6495'}