<jupyter_start><jupyter_text>1. Validate ANPE using p-p plot and Simulation-Based Calibration Load test data<jupyter_code># x = theta_sps
# y = [u, g, r, i, z, sigma_u, sigma_g, sigma_r, sigma_i, sigma_z, z]
_x_test, y_test = Train.load_data('test', version=1, sample='toy', params='thetas_unt')
x_test = np.zeros((_x_test.shape[0], _x_test.shape[1]+1))
x_test[:,0] = _x_test[:,0]

# transform back to dirichlet space
x_test[:,1:5] = dirichlet_transform(_x_test[:,1:4])
x_test[:,5:] = _x_test[:,4:]

# log gamma1, gamma2
x_test[:,7] = np.log10(x_test[:,7])
x_test[:,8] = np.log10(x_test[:,8])<jupyter_output><empty_output><jupyter_text>load samples from `SEDflow` ANPE<jupyter_code>#arch = '500x10.4'
arch = '500x15.2'
anpe_samples = np.load('/scratch/network/chhahn/sedflow/anpe_thetaunt_magsigz.toy.%s.samples.npy' % arch)<jupyter_output><empty_output><jupyter_text>calculate the percentile score and rank of the true values<jupyter_code>pp_thetas, rank_thetas = [], [] 
for igal in np.arange(1000): 
    _mcmc_anpe = anpe_samples[igal,:,:]
    _mcmct_anpe = np.zeros((_mcmc_anpe.shape[0], _mcmc_anpe.shape[1]+1))
    _mcmct_anpe[:,0] = _mcmc_anpe[:,0]

    # transform back to dirichlet space
    _mcmct_anpe[:,1:5] = dirichlet_transform(_mcmc_anpe[:,1:4])
    _mcmct_anpe[:,5:] = _mcmc_anpe[:,4:]
    
    pp_theta, rank_theta = [], []
    for itheta in range(_mcmct_anpe.shape[1]): 
        pp_theta.append(stats.percentileofscore(_mcmct_anpe[:,itheta], x_test[igal,itheta])/100.)
        rank_theta.append(np.sum(np.array(_mcmct_anpe[:,itheta]) < x_test[igal,itheta]))
    pp_thetas.append(pp_theta)
    rank_thetas.append(rank_theta)
    
pp_thetas = np.array(pp_thetas)
rank_thetas = np.array(rank_thetas)<jupyter_output><empty_output><jupyter_text>calculate percentile score and rank for test MCMC from arcoiris, for reference<jupyter_code>dat_dir = '/scratch/network/chhahn/arcoiris/sedflow/'
x_mcmc_test = np.load(os.path.join(dat_dir, 'test.gold.thetas_sps.toy_noise.npy')) 
x_mcmc_test[:,7:9] = np.log10(x_mcmc_test[:,7:9])

_mags = np.load(os.path.join(dat_dir, 'test.gold.mags.toy_noise.npy'))
_sigs = np.load(os.path.join(dat_dir, 'test.gold.sigs.toy_noise.npy'))
_zred = np.load(os.path.join(dat_dir, 'test.gold.zred.toy_noise.npy')) 
y_mcmc_test = np.concatenate([_mags, _sigs, _zred], axis=1)


pp_thetas_mcmc, rank_thetas_mcmc = [], []
for igal in np.arange(100):
    _mcmc_test = np.load(os.path.join('/scratch/network/chhahn/arcoiris/sedflow/mcmc_test_redo/', 'mcmc.test.toy.gold.%i.npy' % igal))
    _mcmc_test = Train.flatten_chain(_mcmc_test[2000:])
    
    _mcmct_test = np.zeros((_mcmc_test.shape[0], _mcmc_test.shape[1]+1))
    _mcmct_test[:,0] = _mcmc_test[:,0]

    # transform back to dirichlet space
    _mcmct_test[:,1:5] = dirichlet_transform(_mcmc_test[:,1:4])
    _mcmct_test[:,5:] = _mcmc_test[:,4:]
    
    pp_theta_mcmc, rank_theta_mcmc = [], []
    for itheta in range(_mcmct_test.shape[1]):
        pp_theta_mcmc.append(stats.percentileofscore(_mcmct_test[:,itheta], x_mcmc_test[igal,itheta])/100.)
        rank_theta_mcmc.append(np.sum(np.array(_mcmct_test[:,itheta]) < x_mcmc_test[igal,itheta])/float(_mcmct_test.shape[0]))
    pp_thetas_mcmc.append(pp_theta_mcmc)
    rank_thetas_mcmc.append(rank_theta_mcmc)

pp_thetas_mcmc = np.array(pp_thetas_mcmc)
rank_thetas_mcmc = np.array(rank_thetas_mcmc)<jupyter_output><empty_output><jupyter_text>p-p plot<jupyter_code>theta_lbls = [r'$\log M_*$', r"$\beta_1$", r"$\beta_2$", r"$\beta_3$", r"$\beta_4$", r'$f_{\rm burst}$', r'$t_{\rm burst}$', r'$\log \gamma_1$', r'$\log \gamma_2$', r'$\tau_1$', r'$\tau_2$', r'$n_{\rm dust}$']
mpl.use('PDF')
mpl.rcParams['text.usetex'] = True

fig = plt.figure(figsize=(8,8))
sub = fig.add_subplot(111)
for itheta in range(pp_thetas.shape[1]): 
    # evaluate the histogram
    values, base = np.histogram(pp_thetas[:,itheta], bins=40)
    #evaluate the cumulative
    cumulative = np.cumsum(values) / np.sum(values)
    sub.plot(base[:-1], cumulative, label=theta_lbls[itheta])
    
    # evaluate the histogram
    values, base = np.histogram(pp_thetas_mcmc[:,itheta], bins=40)
    #evaluate the cumulative
    cumulative = np.cumsum(values) / np.sum(values)
    sub.plot(base[:-1], cumulative, c='gray', lw=1, ls=':')
sub.plot([], [], c='gray', ls=':', label='MCMC')
sub.plot([0., 1.], [0., 1.], c='k', ls='--')
sub.set_xlim(0., 1.)
sub.set_ylim(0., 1.)
sub.legend(loc='upper left', fontsize=15)
fig.savefig('paper/figs/ppplot.pdf', bbox_inches='tight')

IFrame("paper/figs/ppplot.pdf", width=600, height=600)<jupyter_output><empty_output><jupyter_text>simulation-based calibrationMetric from Talts+(2020) and uses rank statistic rather than percentile score.<jupyter_code>fig = plt.figure(figsize=(15,9))
for i in range(pp_thetas.shape[1]): 
    sub = fig.add_subplot(3,4,i+1)
    sub.hist(rank_thetas[:,i]/10000., density=True, histtype='step', linewidth=2)
    sub.hist(rank_thetas_mcmc[:,i], density=True, histtype='step', linewidth=0.75, color='gray', linestyle=':')
    sub.plot([0., 1.], [1., 1.], c='k', ls='--')
    
    sub.text(0.05, 0.95, theta_lbls[i], ha='left', va='top', fontsize=20, transform=sub.transAxes)
    sub.set_xlim(0, 1.)
    sub.set_ylim(0., 2.)
    sub.set_yticklabels([])
    sub.set_xticklabels([])
sub.plot([], [], c='C0', label='ANPE')
sub.plot([], [], c='gray', ls=':', label='MCMC')
sub.legend(loc='lower right', fontsize=15, handletextpad=0)
fig.savefig('paper/figs/sbc.pdf', bbox_inches='tight')

IFrame("paper/figs/sbc.pdf", width=600, height=600)<jupyter_output><empty_output><jupyter_text>Comparison of ANPE vs MCMC posterior for single galaxy<jupyter_code>prior_low = [7, 0., 0., 0., 0., 1e-2, np.log10(4.5e-5), np.log10(4.5e-5), 0, 0., -2.]
prior_high = [12.5, 1., 1., 1., 1., 13.27, np.log10(1.5e-2), np.log10(1.5e-2), 3., 3., 1.]
lower_bounds = torch.tensor(prior_low)
upper_bounds = torch.tensor(prior_high)

prior = Ut.BoxUniform(low=lower_bounds, high=upper_bounds, device='cpu')
fanpe = '/scratch/network/chhahn/sedflow/anpe_thetaunt_magsigz.toy.%s.pt' % arch

anpe = Inference.SNPE(prior=prior, density_estimator=Ut.posterior_nn('maf', hidden_features=500, num_transforms=15), device='cpu')
anpe.append_simulations(
    torch.as_tensor(_x_test.astype(np.float32)), 
    torch.as_tensor(y_test.astype(np.float32)))

p_x_y_estimator = anpe._build_neural_net(torch.as_tensor(_x_test.astype(np.float32)), torch.as_tensor(y_test.astype(np.float32)))
p_x_y_estimator.load_state_dict(torch.load(fanpe))

anpe._x_shape = Ut.x_shape_from_simulation(torch.as_tensor(y_test.astype(np.float32)))
hatp_x_y = anpe.build_posterior(p_x_y_estimator)<jupyter_output><empty_output><jupyter_text>Load NSA data and MCMC posteriors<jupyter_code>nsa, _ = Obs.NSA()

flux_nsa = nsa['NMGY'][:,2:]
ivar_nsa = nsa['NMGY_IVAR'][:,2:]

zred_nsa = nsa['Z']

mags_nsa = Train.flux2mag(flux_nsa)
sigs_nsa = Train.sigma_flux2mag(ivar_nsa**-0.5, flux_nsa)

y_nsa = np.concatenate([mags_nsa, sigs_nsa, zred_nsa[:,None]], axis=1)
i_nsa = 25
print(nsa['NSAID'][i_nsa])
c_light = 2.998e18
jansky_cgs = 1e-23

lambda_sdss = np.array([3543., 4770., 6231., 7625., 9134.])

flux_nsa_conv = flux_nsa[i_nsa] * 1e-9 * 1e17 * c_light / lambda_sdss**2 * (3631. * jansky_cgs) # convert to 10^-17 ergs/s/cm^2/Ang
ivar_nsa_conv = ivar_nsa[i_nsa] * (1e-9 * 1e17 * c_light / lambda_sdss**2 * (3631. * jansky_cgs))**-2.

dir_mcmc = '/scratch/network/chhahn/arcoiris/sedflow/mcmc_nsa_redo/'

_anpe = np.array(hatp_x_y.sample((10000,), x=torch.as_tensor(y_nsa[i_nsa,:]), show_progress_bars=True))
anpe = np.zeros((_anpe.shape[0], _anpe.shape[1]+1))
anpe[:,0] = _anpe[:,0]
# transform back to dirichlet space
anpe[:,1:5] = dirichlet_transform(_anpe[:,1:4])
anpe[:,5:] = _anpe[:,4:]    

_mcmc = np.load(os.path.join(dir_mcmc, 'mcmc.nsa.%i.npy' % i_nsa))
_mcmc = Train.flatten_chain(_mcmc[4000:])
mcmc = np.zeros((_mcmc.shape[0], _mcmc.shape[1]+1))
mcmc[:,0] = _mcmc[:,0]
# transform back to dirichlet space
mcmc[:,1:5] = dirichlet_transform(_mcmc[:,1:4])
mcmc[:,5:] = _mcmc[:,4:]    
from provabgs import models as Models

nsa_filters = Train.photometry_bands()

m_sps = Models.NMF(burst=True, emulator=True)
log_prob_anpe = hatp_x_y.log_prob(_anpe, x=torch.as_tensor(y_nsa[i_nsa,:]))
i_anpe_bf = np.argmax(log_prob_anpe)
# get best fit
_tt_bf = anpe[i_anpe_bf].copy()
_tt_bf[7:9] = 10**_tt_bf[7:9]
w_anpe, f_anpe = m_sps.sed(_tt_bf, zred_nsa[i_nsa])
maggies_anpe = Train.SED_to_maggies(w_anpe, f_anpe, filters=nsa_filters)
mags_anpe = Train.flux2mag(np.array(list(maggies_anpe.as_array()[0])) * 1e9)
fspecs_mcmc, fluxes_mcmc = [], []
for _tt in mcmc[::100]: 
    tt = _tt.copy()
    tt[7:9] = 10**tt[7:9]
    w_mcmc, f_mcmc = m_sps.sed(tt, zred_nsa[i_nsa])
    _maggies = Train.SED_to_maggies(w_mcmc, f_mcmc, filters=nsa_filters)
    fspecs_mcmc.append(f_mcmc)
    fluxes_mcmc.append(np.array(_maggies.as_array()[0].tolist()) * 1e9)

dflux = flux_nsa[i_nsa] - np.array(fluxes_mcmc)
log_prob_mcmc = -0.5 * np.sum(dflux**2 * ivar_nsa[i_nsa])  

mags_mcmc = Train.flux2mag(fluxes_mcmc[np.argmax(log_prob_mcmc)])
f_mcmc = fspecs_mcmc[np.argmax(log_prob_mcmc)]
print(mags_nsa[i_nsa])
print(mags_anpe)
print(mags_mcmc)
_low = [7, 0., 0., 0., 0., 0., 1e-2, np.log10(4.5e-5), np.log10(4.5e-5), 0, 0., -2.]
_high = [12.5, 1., 1., 1., 1., 1., 13.27, np.log10(1.5e-2), np.log10(1.5e-2), 3., 3., 1.]
theta_range = [(l, h) for l, h in zip(_low, _high)]
# read in photometric bandpass filters 
from speclite import filters as specFilter
filter_response = specFilter.load_filters('sdss2010-*')
wave_eff = [filter_response[i].effective_wavelength.value for i in range(len(filter_response))]
theta_lbls = [r'$\log M_*$', r"$\beta_1$", r"$\beta_2$", r"$\beta_3$", r"$\beta_4$", r'$f_{\rm burst}$', r'$t_{\rm burst}$', r'$\log \gamma_1$', r'$\log \gamma_2$', r'$\tau_1$', r'$\tau_2$', r'$n_{\rm dust}$']
mpl.use('PDF')
mpl.rcParams['text.usetex'] = True

ndim = len(theta_lbls)

fig = plt.figure(figsize=(15, 18))

gs0 = fig.add_gridspec(nrows=ndim, ncols=ndim, top=0.95, bottom=0.28)
for yi in range(ndim):
    for xi in range(ndim):
        sub = fig.add_subplot(gs0[yi, xi])

_fig = DFM.corner(mcmc, color='k', levels=[0.68, 0.95], range=theta_range,
                  plot_density=False, plot_datapoints=False, hist_kwargs={'density': True}, fig=fig)
_ = DFM.corner(anpe, color='C0', levels=[0.68, 0.95], range=theta_range,
               plot_density=False, plot_datapoints=False, hist_kwargs={'density': True}, 
               labels=theta_lbls, label_kwargs={'fontsize': 25}, fig=fig)


# ndim = int(np.sqrt(len(fig.axes)))
axes = np.array(fig.axes).reshape((ndim, ndim))

ax = axes[2, ndim-4]
ax.fill_between([], [], [], color='k', label='MCMC posterior')
ax.fill_between([], [], [], color='C0', label='ANPE posterior')
ax.legend(handletextpad=0.2, markerscale=10, fontsize=25)

axes = np.array(fig.axes).reshape((ndim, ndim))
for yi in range(1, ndim):
    ax = axes[yi, 0]
    ax.set_ylabel(theta_lbls[yi], fontsize=20, labelpad=30)
    ax.yaxis.set_label_coords(-0.6, 0.5)
for xi in range(ndim): 
    ax = axes[-1, xi]
    ax.set_xlabel(theta_lbls[xi], fontsize=20, labelpad=30)
    ax.xaxis.set_label_coords(0.5, -0.55)
for xi in range(ndim): 
    ax = axes[xi,xi]
    ax.set_xlim(theta_range[xi])

gs1 = fig.add_gridspec(nrows=1, ncols=5, top=0.2, bottom=0.05)
for ii in range(5):
    sub = fig.add_subplot(gs1[0,ii])

    sub.plot(w_anpe, f_anpe, c='C0', lw=1, label='ANPE best-fit')
    sub.plot(w_mcmc, f_mcmc, c='k', lw=1, ls=':', label='MCMC best-fit')
    #for i in range(len(filter_response)): 
    #    #sub.plot(filter_response[i].wavelength, 15 + 30 * filter_response[i].response, ls='--') 
    #    sub.text(filter_response[i].effective_wavelength.value, 1.1 * 30 * np.max(filter_response[i].response), ['u', 'g', 'r', 'i', 'z'][i], fontsize=20, color='C%i' % i)

    sub.errorbar(lambda_sdss, flux_nsa_conv, yerr=ivar_nsa_conv**-0.5, fmt='.C3', markersize=20, elinewidth=4, label='NSA Photometry')
    if ii == 4: sub.legend(loc='upper right', fontsize=18, handletextpad=0)
    if ii == 2: sub.set_xlabel('wavelength [$A$]', fontsize=20) 
    #sub.set_xlim(3e3, 1e4)
    sub.set_xlim(lambda_sdss[ii]-500, lambda_sdss[ii]+500)
    if ii == 0: sub.set_ylabel('flux [$10^{-17} erg/s/cm^2/A$]', fontsize=20, labelpad=15) 
    else: sub.set_yticklabels([])
    if ii == 1: sub.set_xticks([4500, 4750, 5000])
    #if ii == 1: sub.set_xticks([4500, 4750, 5000])
    sub.set_ylim(15., 45.)
    sub.set_title('$%s$' % ['u', 'g', 'r', 'i', 'z'][ii], fontsize=20)

fig.savefig('paper/figs/corner.pdf', bbox_inches='tight')
IFrame("paper/figs/corner.pdf", width=600, height=600)<jupyter_output><empty_output>
{'fraction_non_alphanumeric': 0.19296577946768062, 'fraction_numerical': 0.04404309252217997, 'mean_word_length': 5.33785140562249, 'pattern_counts': {'":': 0, '<': 35, '<?xml version=': 0, '>': 33, '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': '8', 'dir': 'jupyter-structured-clean-dedup', 'id': '463313', 'license': 'MIT', 'n_tokens_mistral': 6062, 'n_tokens_neox': 5723, 'n_words': 835, 'path': 'docs/nb/anpe_vs_mcmc.ipynb', 'repo_name': 'changhoonhahn/SEDflow', 'size': '12624'}