borg_velocity/notebooks/Analyse_chain.ipynb

import numpy as np import matplotlib.pyplot as plt import configparser import glob from tqdm import tqdm import h5py as h5 import ast import os import corner import Pk_library as PKL import MAS_library as MASL from contextlib import contextmanager import sys import symbolic_pofk.linear @contextmanager def suppress_stdout(): with open(os.devnull, 'w') as devnull: old_stdout = sys.stdout sys.stdout = devnull try: yield finally: sys.stdout = old_stdout from analysis import ( get_mcmc_steps, load_param_samples, get_truths, crop_field, compute_ensemble_mean_field, get_mock_field, get_spectra, get_both_fields, get_likelihood_values)

In [4]:

Plot the non-ic parameters¶

In [25]:

# ini_name = '../conf/basic_ini.ini'
# dirname = '/data101/bartlett/fsigma8/borg_velocity/basic_run_cosmo_model'
# dirname = '/data101/bartlett/fsigma8/borg_velocity/supranta_pars_N64'
# dirname = '/data101/bartlett/fsigma8/borg_velocity/supranta_N64_all'
dirname = '/data101/bartlett/fsigma8/borg_velocity/supranta_N64_not_ic_not_lam'
dirname = '/data101/bartlett/fsigma8/borg_velocity/test_dir'
dirname = '/data101/bartlett/fsigma8/borg_velocity/test_dir_mvslice'
dirname = '/data101/bartlett/fsigma8/borg_velocity/test_dir_mvslice_modelpar'
# dirname = '/data101/bartlett/fsigma8/borg_velocity/test_dir_hmc_bulk'; iter_max = 3601
# dirname = '/data101/bartlett/fsigma8/borg_velocity/test_dir_hmc_modelpar'; iter_max = 1696
dirname = '/data101/bartlett/fsigma8/borg_velocity/test_dir_transformedblackjax'; iter_max = 234
dirname = '/data101/bartlett/fsigma8/borg_velocity/test_dir_blackjax'; iter_max = -1
dirname = '/data101/bartlett/fsigma8/borg_velocity/test_dir_blackjax_model_par'; iter_max = 667
dirname = '/data101/bartlett/fsigma8/borg_velocity/transformedblackjax_model_par'; iter_max = 1093
dirname = '/data101/bartlett/fsigma8/borg_velocity/blackjax_model_cosmo'; iter_max = -1
dirname = '/data101/bartlett/fsigma8/borg_velocity/blackjax_model_ic_v2'; iter_max = 1906
dirname = '/data101/bartlett/fsigma8/borg_velocity/blackjax_model_ic_v3'; iter_max = -1
dirname = '/data101/bartlett/fsigma8/borg_velocity/blackjax_model_ic_bf_prior'; iter_max = -1
dirname = '/data101/bartlett/fsigma8/borg_velocity/mvslice_model'; iter_max = 757
ini_name = f'{dirname}/ini_file.ini'
nframe = 700
iter_min = -1
# iter_max = -1

names, steps, samples = load_param_samples(ini_name, dirname, nframe, iter_max, iter_min)
print(samples.shape)
truths = get_truths(ini_name, names)
# print(samples[-1,:])

TO SAMPLE ['sig_v']
MY SAMPLER IS mvslice
['sig_v']

100%|██████████| 758/758 [00:27<00:00, 27.64it/s]

(1, 758)

Trace plot

In [26]:

ncol = 4
nrow = int(np.ceil(len(names) / ncol))

fig, axs = plt.subplots(nrow, ncol, figsize=(ncol*4, nrow*3))
axs = np.atleast_2d(axs)
# for i in range(ncol):
    # axs[-1,i].set_xlabel('Step Number')
axs = axs.flatten()
for i in range(len(names)):
    axs[i].plot(steps, samples[i,:])
    axs[i].set_title(names[i])
    axs[i].axhline(truths[i], color='k')
    axs[i].set_xlabel('Step Number')
for i in range(len(names), len(axs)):
    axs[i].remove()
fig.tight_layout()
fig.savefig('../figs/trace.png')

if 'sigma8' in names and 'omega_m' in names:
    sigma8 = samples[names.index('sigma8')]
    omega_m = samples[names.index('omega_m')]

    config = configparser.ConfigParser()
    config.read(ini_name)
    omega_b = float(config['cosmology']['omega_b'])
    h = float(config['cosmology']['h100'])
    n_s = float(config['cosmology']['n_s'])
    A_s = 1.e-9 * symbolic_pofk.linear.sigma8_to_As(
            sigma8, omega_m, omega_b, h, n_s)
    print('A_s range:', A_s.min(), A_s.max())
    plt.figure()
    plt.hist(A_s)

Corner plot

In [10]:

burn_in = 0
burn_in = 1000

m = steps > burn_in

# Full corner plot
corner.corner(
    samples[:,m].T,
    labels=names,
    truths=truths
);
fig = plt.gcf()
fig.savefig('../figs/corner_full.png', bbox_inches='tight', facecolor='white')

# Cosmology corner plot
if 'sigma8' in names and 'omega_m' in names:
    idx = np.array([names.index('sigma8'), names.index('omega_m')], dtype=int)
    corner.corner(
        samples[idx,m].T,
        labels=[names[i] for i in idx],
        truths=[truths[i] for i in idx]
    );
    fig = plt.gcf()
    fig.savefig('../figs/corner_cosmo.png', bbox_inches='tight', facecolor='white')

WARNING:root:Too few points to create valid contours
WARNING:root:Too few points to create valid contours
WARNING:root:Too few points to create valid contours
WARNING:root:Too few points to create valid contours
WARNING:root:Too few points to create valid contours

Mock data¶

In [7]:

dirname = '/data101/bartlett/fsigma8/borg_velocity/basic_run_testing'
dirname = '/data101/bartlett/fsigma8/borg_velocity/supranta_ics_N64_lam'

with h5.File(f'{dirname}/mock_data.h5', 'r') as f:
    dens = f['scalars/BORG_final_density'][:]
    print('DENS', dens.min(), dens.max(), dens.mean(), dens.std())

N = dens.shape[0]
fig, axs = plt.subplots(1, 3, figsize=(15,5))
axs[0].pcolor(dens[N//2])
axs[1].pcolor(dens[:,N//2,:])
axs[2].pcolor(dens[:,:,N//2])
for ax in axs:
    ax.set_aspect('equal')

DENS -0.9808840727490473 17.530269738764385 -1.6263032587282567e-18 0.8490318268169118

Analyse density fields¶

In [11]:

dirname = '/data101/bartlett/fsigma8/borg_velocity/basic_run_ic'; iter_max = 1748; iter_min = 0
ini_name = '../conf/basic_ini.ini'
# dirname = '/data101/bartlett/fsigma8/borg_velocity/supranta_run_ic_N64'; iter_max = -1; iter_min = 1000
# dirname = '/data101/bartlett/fsigma8/borg_velocity/supranta_run_ic'; iter_max = -1; iter_min = 1000
dirname = '/data101/bartlett/fsigma8/borg_velocity/supranta_ics_N64_lam'; iter_max = -1; iter_min = 1000
dirname = '/data101/bartlett/fsigma8/borg_velocity/supranta_N64_all'; iter_max = -1; iter_min = 1000
dirname = '/data101/bartlett/fsigma8/borg_velocity/blackjax_model_ic_v2'; iter_max = 1773; iter_min = 1000
dirname = '/data101/bartlett/fsigma8/borg_velocity/blackjax_model_ic_v3'; iter_max = 3509; iter_min = 1000
dirname = '/data101/bartlett/fsigma8/borg_velocity/blackjax_model_ic_bf_prior'; iter_max = -1; iter_min = 1000
ini_name = f'{dirname}/ini_file.ini'
nframe = 100

all_true_delta = [None] * 2
all_mean_delta = [None] * 2
all_std_delta = [None] * 2

# Uncropped
all_true_delta[0], _ = get_mock_field(ini_name, dirname, cut_field=False)
all_mean_delta[0], all_std_delta[0] = compute_ensemble_mean_field(ini_name, dirname, nframe, iter_max, iter_min, cut_field=False)

# Cropped
all_true_delta[1], Lcrop = crop_field(ini_name, all_true_delta[0])
all_mean_delta[1], _ = crop_field(ini_name, all_mean_delta[0])
all_std_delta[1], _ = crop_field(ini_name, all_std_delta[0])

for cut_field in [False, True]:
    
    i = int(cut_field)
    true_delta = all_true_delta[i]
    mean_delta = all_mean_delta[i]
    std_delta = all_std_delta[i]

    fig, axs = plt.subplots(3, 3, figsize=(8,7))

    config = configparser.ConfigParser()
    config.read(ini_name)
    Rmax = float(config['mock']['R_max'])
    corner0 = float(config['system']['corner0'])
    L0 = float(config['system']['L0'])
    N0 = int(config['system']['N0'])
    x = np.linspace(corner0, corner0+L0, N0)
    if cut_field:
        m = np.abs(x) < Rmax
        x = x[m]

    for i, field in enumerate([np.log10(2+true_delta), np.log10(2+mean_delta), std_delta]):
        if i == 0:
            vmin = field.min()
            vmax = field.max()
        N = field.shape[0]
        j = N//2
        axs[i,0].pcolor(x, x, field[j], vmin=vmin, vmax=vmax)
        axs[i,1].pcolor(x, x, field[:,j,:], vmin=vmin, vmax=vmax)
        axs[i,2].pcolor(x, x, field[:,:,j], vmin=vmin, vmax=vmax)
        for ax in axs[i,:]:
            circle = plt.Circle((0, 0), Rmax, color='red', fill=False, linestyle='dotted', lw=1.5)
            xlim = ax.get_xlim()
            ylim = ax.get_ylim()
            ax.add_patch(circle)
            ax.set_xlim(xlim)
            ax.set_ylim(ylim)

        for ax in axs[i,:]:
            ax.set_aspect('equal')

    axs[0,1].set_title('Truth')
    axs[1,1].set_title('Ensemble Mean')
    axs[2,1].set_title('Ensemble Std')
            
    fig.tight_layout()
    
    savename = '../figs/ensemble_mean'
    if cut_field:
        savename += '_cropped'
    savename += '.png'
    fig.savefig(savename, bbox_inches='tight', facecolor='w')

Computing ensemble mean field

100%|██████████| 104/104 [00:12<00:00,  8.02it/s]

In [12]:

# ini_name = '../conf/basic_ini.ini'
# dirname = '/data101/bartlett/fsigma8/borg_velocity/basic_run_ic'; iter_max = 1748
# dirname = '/data101/bartlett/fsigma8/borg_velocity/basic_run_cosmo_model'; iter_max = 3780
# dirname = '/data101/bartlett/fsigma8/borg_velocity/supranta_run_ic'; iter_max = -1
# dirname = '/data101/bartlett/fsigma8/borg_velocity/supranta_run_ic_N64'; iter_max = -1
# dirname = '/data101/bartlett/fsigma8/borg_velocity/supranta_ics_N64_lam'; iter_max = -1
# dirname = '/data101/bartlett/fsigma8/borg_velocity/supranta_N64_all'; iter_max = -1
dirname = '/data101/bartlett/fsigma8/borg_velocity/blackjax_model_ic_v2'; iter_max = 1773
dirname = '/data101/bartlett/fsigma8/borg_velocity/blackjax_model_ic_v3'; iter_max = 3509
dirname = '/data101/bartlett/fsigma8/borg_velocity/blackjax_model_ic_bf_prior'; iter_max = -1

ini_name = f'{dirname}/ini_file.ini'
nframe = 100
iter_min = 0
# which_field = 'ics'
# which_field = 'delta'
cut_field = True
# cut_field = False

for which_field in ['ics', 'delta']:

    k, Pk_true, all_pk, all_r = get_spectra(ini_name, dirname, nframe, iter_max, iter_min, which_field=which_field, cut_field=cut_field)
    all_mcmc = get_mcmc_steps(dirname, nframe, iter_max, iter_min=iter_min)

    # Normalize the index for coloring
    norm = plt.Normalize(all_mcmc[0], all_mcmc[-1])
    cmap = plt.cm.spring  # Choose a colormap

    fig, axs = plt.subplots(1, 2, figsize=(15,4))

    # Plot Pk and r
    for i in range(all_pk.shape[0]):
        axs[0].loglog(k, all_pk[i, :], color=cmap(norm(all_mcmc[i])))
        axs[1].semilogx(k, all_r[i, :], color=cmap(norm(all_mcmc[i])))
    axs[0].loglog(k, Pk_true, color='k')
    for i in [-1, 0, 1]:
        axs[1].axhline(i, color='k')

    sm = plt.cm.ScalarMappable(cmap=cmap, norm=norm)
    sm.set_array([])
    for ax in axs:
        cbar = plt.colorbar(sm, ax=ax)
        cbar.set_label('Step')
        ax.set_title(which_field)
        ax.set_xlabel(r'$k$')
    axs[0].set_ylabel(r'$P(k)$')
    axs[1].set_ylabel(r'$r(k)$')
    fig.tight_layout()
    fig.savefig(f'../figs/spectra_{which_field}.png')

BOXSIZE 198.41269841269843

Computing power spectrum of the field...
Time to complete loop = 0.00
Time taken = 0.01 seconds

100%|██████████| 104/104 [00:14<00:00,  7.18it/s]

BOXSIZE 198.41269841269843

Computing power spectrum of the field...
Time to complete loop = 0.00
Time taken = 0.00 seconds

100%|██████████| 104/104 [00:07<00:00, 13.07it/s]

In [13]:

dirname = '/data101/bartlett/fsigma8/borg_velocity/supranta_run_ic_N64'; step = 9554
dirname = '/data101/bartlett/fsigma8/borg_velocity/supranta_run_ic'; step = 1860
dirname = '/data101/bartlett/fsigma8/borg_velocity/supranta_ics_N64_lam'; step = 9624
dirname = '/data101/bartlett/fsigma8/borg_velocity/supranta_N64_all'; step = 8609
dirname = '/data101/bartlett/fsigma8/borg_velocity/blackjax_model_ic_v2'; step = 1773
dirname = '/data101/bartlett/fsigma8/borg_velocity/blackjax_model_ic_v3'; step = 3509
dirname = '/data101/bartlett/fsigma8/borg_velocity/blackjax_model_ic_bf_prior'; step = 4591
ini_name = f'{dirname}/ini_file.ini'
cut_field = True
cut_field = False

fig, axs = plt.subplots(1, 2, figsize=(10,4))

for cut_field in [False, True]:

    d1, d2 = get_both_fields(ini_name, dirname, step, which_field='delta', cut_field=cut_field)
    d1 = np.log10(2 + d1)
    d2 = np.log10(2 + d2)
    
    hist_kwargs = dict(bins=30, density=True, histtype='step')
    i = int(cut_field)
    axs[i].hist(d1.flatten(), label='Mock', **hist_kwargs)
    axs[i].hist(d2.flatten(), label='Sample %i'%step, **hist_kwargs)
    axs[i].legend()
    axs[i].set_xlabel(r'$\log_{10} (2 + \delta)$')
    
axs[0].set_title('Full Box')
axs[1].set_title('Central Region')
fig.tight_layout()
fig.savefig('../figs/density_histogram.png', bbox_inches='tight', facecolor='white')

Plot likelihood as a function of step¶

In [8]:

ini_name = '../conf/basic_ini.ini'
dirname = '/data101/bartlett/fsigma8/borg_velocity/basic_run_ic'; iter_max = 1748
dirname = '/data101/bartlett/fsigma8/borg_velocity/supranta_run_ic_N64'; iter_max = 2421
dirname = '/data101/bartlett/fsigma8/borg_velocity/supranta_ics_N64_lam'; iter_max = -1
dirname = '/data101/bartlett/fsigma8/borg_velocity/supranta_N64_all'; iter_max = -1
dirname = '/data101/bartlett/fsigma8/borg_velocity/blackjax_model_ic_v2'; iter_max = 1773
dirname = '/data101/bartlett/fsigma8/borg_velocity/blackjax_model_ic_v3'; iter_max = 3509
nframe = 100
iter_min = 0

all_mcmc = get_mcmc_steps(dirname, nframe, iter_max, iter_min=iter_min)
logL, logP = get_likelihood_values(dirname, nframe, iter_max, iter_min)
logL -= logL[-1]
logP -= logP[-1]

burnin = 10

fig, axs = plt.subplots(1, 2, figsize=(8,3))
axs[0].plot(all_mcmc[burnin:], logL[burnin:], label='logL')
axs[1].plot(all_mcmc[burnin:], logP[burnin:], label='logPrior')
axs[0].set_ylabel('Log-Likelihood')
axs[1].set_ylabel('Log-Prior')
for ax in axs:
    ax.set_xlabel('Step')
fig.tight_layout()
fig.savefig('../figs/likelihood_trace.png', bbox_inches='tight', facecolor='white')

  0%|          | 0/102 [00:00<?, ?it/s]/tmp/ipykernel_2378/3269883151.py:279: DeprecationWarning: Conversion of an array with ndim > 0 to a scalar is deprecated, and will error in future. Ensure you extract a single element from your array before performing this operation. (Deprecated NumPy 1.25.)
  all_logL[i] = f['scalars/hmc_Elh'][:]
/tmp/ipykernel_2378/3269883151.py:280: DeprecationWarning: Conversion of an array with ndim > 0 to a scalar is deprecated, and will error in future. Ensure you extract a single element from your array before performing this operation. (Deprecated NumPy 1.25.)
  all_logprior[i] = f['scalars/hmc_Eprior'][:]
100%|██████████| 102/102 [00:01<00:00, 67.49it/s]

Mock Data Plot¶

In [ ]:

# dirname = '/data101/bartlett/fsigma8/borg_velocity/supranta_run_ic'
dirname = '/data101/bartlett/fsigma8/borg_velocity/supranta_ics_N64_lam'
use_true = True
# use_true = False

ini_name = f'{dirname}/ini_file.ini'
config = configparser.ConfigParser()
config.read(ini_name)
nsamp = int(config['run']['nsamp'])
rmax = float(config['mock']['R_max'])
L = float(config['system']['L0'])
N = int(config['system']['N0'])

hist_kwargs = dict(bins=30, density=True, histtype='step')
scatter_kwargs = dict(marker='.', ms=1, lw=0)

fig, axs = plt.subplots(1, 4, figsize=(20,4))
fig2, axs2 = plt.subplots(1, 3, figsize=(15,4))
    
all_x = []

with h5.File(dirname + "/tracer_data_mock.h5",'r') as mcmc_file:
    for i in range(nsamp):
        if use_true:
            x = mcmc_file[f'sample_{i}/coord_true'][:]
        else:
            x = mcmc_file[f'sample_{i}/coord_meas'][:]
        all_x.append(x)
        r = np.sqrt(np.sum(x**2, axis=0))
        for j in range(3):
            axs[j].hist(x[j], **hist_kwargs)
        axs[-1].hist(r, label=f'Sample {i}', **hist_kwargs)
    
        axs2[0].plot(x[0], x[1], **scatter_kwargs)
        axs2[1].plot(x[1], x[2], **scatter_kwargs)
        axs2[2].plot(x[2], x[0], **scatter_kwargs)
        
for ax in axs2:
    ax.set_aspect('equal')    
    
fig3, axs3 = plt.subplots(1, 4, figsize=(15,4))
with h5.File(dirname + "/tracer_data_mock.h5",'r') as mcmc_file:
    for i in range(nsamp):
        xtrue = mcmc_file[f'sample_{i}/coord_true'][:]
        xmeas = mcmc_file[f'sample_{i}/coord_meas'][:]
        for j in range(3):
            axs3[j].plot(xtrue, xmeas, **scatter_kwargs)
        rtrue = np.sqrt(np.sum(xtrue**2, axis=0))
        rmeas = np.sqrt(np.sum(xmeas**2, axis=0))
        axs3[-1].plot(rtrue, rmeas, **scatter_kwargs)
    
# construct 3D counts field
all_x = np.hstack(all_x).T.astype(np.float32)
counts = np.zeros((N, N, N), dtype=np.float32)
MASL.MA(all_x, counts, L, 'CIC', verbose=True)
print(counts.shape)

plt.figure()
plt.pcolor(counts[N//2])

Using CIC mass assignment scheme
Time taken = 0.000 seconds

(64, 64, 64)

Out[ ]:

<matplotlib.collections.PolyCollection at 0x7f7ae66b1150>

In [ ]:

In [ ]:

Debug Underflow¶

In [28]:

dirname = '/data101/bartlett/fsigma8/borg_velocity/blackjax_model_ic/'
boxsize = 500
MAS = "CIC"

with h5.File(f'{dirname}/mcmc_95.h5', 'r') as f:
    dens_mcmc = f['scalars/BORG_final_density'][:]
    
with h5.File(f'{dirname}/mock_data.h5', 'r') as f:
    dens_true = f['scalars/BORG_final_density'][:]
    
# Crop
# ini_name = f"{dirname}/ini_file.ini"
# dens_mcmc, boxsize = crop_field(ini_name, dens_mcmc)
# dens_true, boxsize = crop_field(ini_name, dens_true)
    
fig, axs = plt.subplots(2, 3, figsize=(15,8))

for i, f in enumerate([dens_true, dens_mcmc]):
    N = f.shape[0]
    axs[i,0].pcolor(f[:,:,N//2])
    axs[i,1].pcolor(f[:,N//2,:])
    axs[i,2].pcolor(f[N//2,:,:])
for ax in axs.flatten():
    ax.set_aspect('equal')
    
axs[0,1].set_title('True')
axs[1,1].set_title('MCMC')

Pk = PKL.XPk([dens_true.astype(np.float32),dens_mcmc.astype(np.float32)], boxsize, axis=0, MAS=[MAS, MAS], threads=1)
Pk0_1  = Pk.Pk[:,0,0]  #monopole of field 1
Pk0_2  = Pk.Pk[:,0,1]  #monopole of field 2
Pk0_X  = Pk.XPk[:,0,0] #monopole of 1-2 cross P(k)
k      = Pk.k3D
r = Pk0_X / np.sqrt(Pk0_1 * Pk0_2)

fig, axs = plt.subplots(1, 2, figsize=(15,4))
axs[0].loglog(k, Pk0_1, label='True')
axs[0].loglog(k, Pk0_2, label='MCMC')
axs[1].semilogx(k, r, label='TruexMCMC')
for ax in axs:
    ax.legend()
for y in [0, 1]:
    axs[1].axhline(y, color='k')
    
    
fig, ax = plt.subplots(1, 1)
bins = np.linspace(np.log10(2 + min(dens_true.min(), dens_mcmc.min())), np.log10(2 + max(dens_true.max(), dens_mcmc.max())), 30) 
ax.hist(np.log10(2 + dens_true.flatten()), bins=bins, density=True, histtype='step', label='true')
ax.hist(np.log10(2 + dens_mcmc.flatten()), bins=bins, density=True, histtype='step', label='mcmc')
ax.legend()

Computing power spectra of the fields...
Time FFTS = 0.00
Time loop = 0.01
Time taken = 0.01 seconds

Out[28]:

<matplotlib.legend.Legend at 0x149efc0d4110>

In [49]:

dirname = '/data101/bartlett/fsigma8/borg_velocity/blackjax_model_ic/'
component = 0

vel_true = np.load(f'{dirname}/mock_vel.npy')
vel_mcmc = np.load(f'{dirname}/mcmc_vel.npy')
    
fig, axs = plt.subplots(2, 3, figsize=(15,8))

for i, f in enumerate([vel_true, vel_mcmc]):
    N = f.shape[0]
    axs[i,0].pcolor(f[component,:,:,N//2])
    axs[i,1].pcolor(f[component,:,N//2,:])
    axs[i,2].pcolor(f[component,N//2,:,:])
for ax in axs.flatten():
    ax.set_aspect('equal')
    
axs[0,1].set_title('True')
axs[1,1].set_title('MCMC')

fig, axs = plt.subplots(1, 3, figsize=(15,4))
bins = np.linspace(min(vel_true.min(), vel_mcmc.min()), max(vel_true.min(), vel_mcmc.max()), 30)
for i in range(len(axs)):
    axs[i].hist(vel_true[i].flatten(), bins=bins, density=True, histtype='step', label='True')
    axs[i].hist(vel_mcmc[i].flatten(), bins=bins, density=True, histtype='step', label='MCMC')
    axs[i].legend()
    axs[i].set_yscale('log')
    
fig, axs = plt.subplots(1, 3, figsize=(15,4))
for i in range(len(axs)):
    axs[i].hexbin(vel_true[i].flatten(), vel_mcmc[i].flatten(), mincnt=1)

In [86]:

dirname = '/data101/bartlett/fsigma8/borg_velocity/blackjax_model_ic/'
component = 0

lkl_true = np.load(f'{dirname}/mock_lkl_ind.npz')
lkl_mcmc = np.load(f'{dirname}/mcmc_lkl_ind.npz')
nsamp = len(lkl_true.files)

samp_num = None
iworst = None
worst_diff = 0

fig, axs = plt.subplots(1, nsamp, figsize=(4*nsamp, 4))

for i in range(nsamp):
    xtrue = lkl_true[f'arr_{i}']
    xmcmc = lkl_mcmc[f'arr_{i}']
    axs[i].plot(xtrue, xmcmc, '.')
    xlim = axs[i].get_xlim()
    ylim = axs[i].get_ylim()
    xlim = (min(xlim[0], ylim[0]), max(xlim[1], ylim[1]))
    axs[i].plot(xlim, xlim, color='k')
    axs[i].set_xlim(xlim)
    axs[i].set_ylim(xlim)
    axs[i].set_xlabel('True log-likelihood')
    axs[i].set_ylabel('MCMC log-likelihood')
    axs[i].set_title(f'Sample {i}')
    
    # Get worst point
    diff = np.abs(xtrue - xmcmc)
    isort = np.argsort(-diff)
    if diff[isort[0]] > worst_diff:
        worst_diff = diff[isort[0]]
        samp_num = i
        iworst = isort[0]
    print(isort[0], diff[isort[:2]])    
fig.tight_layout()
fig.savefig('../figs/debug_logL.png', facecolor='white', bbox_inches='tight')
    
print('Worst:', samp_num, iworst, worst_diff)

# Find information on this tracer
with h5.File(f'{dirname}/tracer_data_mock.h5', 'r') as f:
    coord_true = np.array(f[f'sample_{samp_num}/coord_true'][:])
    coord_meas = np.array(f[f'sample_{samp_num}/coord_meas'][:])
    sig_mu = f[f'sample_{samp_num}/sig_mu'][()]
    vr_true = np.array(f[f'sample_{samp_num}/vr_true'][:])
    cz_obs = np.array(f[f'sample_{samp_num}/cz_obs'][:])
    r_hMpc = np.sqrt(np.sum(coord_meas ** 2, axis=0))
    print(coord_true.shape, coord_meas.shape)
    print(sig_mu)
    print(vr_true.shape)
    print(cz_obs.shape)
    print(r_hMpc.shape)
    
print('True position:', coord_true[:,iworst])
print('Measured r:', coord_meas[:,iworst])
print('True vr:', vr_true[iworst])
print('Observed cz:', cz_obs[iworst])

xmax = np.amax(np.abs(coord_true), axis=1)
print('Largest true (abs) along each axis:', xmax)
xmax = np.amax(np.abs(coord_meas), axis=1)
print('Largest measured (abs) along each axis:', xmax)

# Plot radial velocity for the worst tracer
vr_true = np.load(f'{dirname}/mock_tracer_vr.npz')[f'arr_{samp_num}'][iworst,:]
vr_mcmc = np.load(f'{dirname}/mcmc_tracer_vr.npz')[f'arr_{samp_num}'][iworst,:]
fig, ax = plt.subplots(1, 1, figsize=(15,4))
ax.plot(vr_true, label='true')
ax.plot(vr_mcmc, label='mcmc')
ax.legend()
print(vr_true.shape)

# Plot the MB positions for the worst tracer
itracer = iworst
itracer = 3
print(coord_true[:,itracer])
MB_pos = np.load(f'{dirname}/MB_pos.npz')[f'arr_{samp_num}'][:,itracer,:]
fig, axs = plt.subplots(1, 3, figsize=(15,4))
for i in range(3):
    axs[i].plot(MB_pos[i], '.', label='MB')
    axs[i].axhline(coord_true[i,itracer], color='k', label='TRUE')
    axs[i].axhline(coord_meas[i,itracer], color='r', label='MEASURED')
    axs[i].legend()
    axs[i].set_xlabel('Index')
    axs[i].set_ylabel('Value')
axs[0].set_title(r'$x$')
axs[1].set_title(r'$y$')
axs[2].set_title(r'$z$')
fig.tight_layout()
# fig.savefig(f'../figs/MB_pos_{samp_num}_{itracer}.png', facecolor='white', bbox_inches='tight')

192 [2.8885622 1.8333287]
976 [3.123437  2.7141156]
129 [1.9811783 1.97686  ]
282 [483.06342     2.1391425]
Worst: 3 282 483.06342
(3, 1225) (3, 1225)
0.15200306866613814
(1225,)
(1225,)
(1225,)
True position: [ 180.07121842 -249.22822852   96.67707877]
Measured r: [ 197.73909905 -273.68152333  106.16265399]
True vr: -269.32922
Observed cz: 32934.527
Largest true (abs) along each axis: [180.07121842 249.22822852 150.8314107 ]
Largest measured (abs) along each axis: [197.73909905 273.68152333 161.08961474]
(201,)
[ -26.74362669  -38.61752179 -147.64106462]

Use here to create warnings to implement in the code

In [90]:

R_lim = 500 / 2

fig, ax = plt.subplots(1, 1, figsize=(10,4))
bins = np.linspace(0, 400, 100)

with h5.File(f'{dirname}/tracer_data_mock.h5', 'r') as f:
    for i in range(4):
        coord_true = np.array(f[f'sample_{i}/coord_true'][:])
        coord_meas = np.array(f[f'sample_{i}/coord_meas'][:])
        r_hMpc = np.sqrt(np.sum(coord_meas ** 2, axis=0))
        ax.hist(r_hMpc, bins=bins, density=True, histtype='step', label=f'sample_{i}')
        m = r_hMpc > R_lim
        print(i, m.sum())
ax.legend()
ax.axvline(R_lim)

7.269984374515526 210.9616513655193 (345,)
6.541561558252167 255.9753565156616 (1682,)
4.3919989345102 262.8841977513585 (556,)
5.479083286248933 353.9390295101536 (1225,)

Out[90]:

<matplotlib.lines.Line2D at 0x149eff55c4d0>

In [ ]: