borg_velocity/notebooks/Analyse_chain.ipynb
2024-10-25 11:40:58 +02:00

1.3 MiB

In [1]:
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
In [2]:
def get_mcmc_steps(dirname, nframe, iter_max, iter_min=0):
    """
    Obtain evenly-spaced sample of MCMC steps to make movie from
    """
    all_mcmc = glob.glob(dirname + '/mcmc_*.h5')
    x = [m[len(dirname + '/mcmc_'):-3] for m in all_mcmc]
    all_mcmc = np.sort([int(m[len(dirname + '/mcmc_'):-3]) for m in all_mcmc])
    if iter_max >= 0:
        all_mcmc = all_mcmc[all_mcmc <= iter_max]
    all_mcmc = all_mcmc[all_mcmc >= iter_min]
    if nframe > 0:
        max_out = max(all_mcmc)
        min_out = min(all_mcmc)
        step = max(int((max_out - min_out+1) / nframe), 1)
        all_mcmc = all_mcmc[::step]
        if max_out not in all_mcmc:
            all_mcmc = np.concatenate([all_mcmc, [max_out]])
    return all_mcmc

def load_param_samples(ini_name, dirname, nframe, iter_max, iter_min):
    
    config = configparser.ConfigParser()
    config.read(ini_name)
    to_sample = []
    for k,v in config['block_loop'].items():
        if v.strip() == 'false':
            i = k.index('_sampler')
            if k[:i] not in ['hades', 'bias', 'nmean']:
                to_sample.append(k[:i])
                
    print("TO SAMPLE", to_sample)
    nsamp = int(config['run']['nsamp'])
    new_to_sample = []
    for s in to_sample:
        if s in ['omega_m', 'sigma8', 'sig_v']:
            new_to_sample.append(s)
        elif s == 'bulk_flow':
            for d in ['_x', '_y', '_z']:
                new_to_sample.append(f'{s}{d}')
        else:
            for i in range(nsamp):
                new_to_sample.append(f'{s}{i}')
    
    # This is desired list to sample
    to_sample = new_to_sample
    
    # Which steps to use
    all_mcmc = get_mcmc_steps(dirname, nframe, iter_max, iter_min=iter_min)
    
    sampler = config['sampling']['algorithm'].lower()
    samples = np.empty((len(to_sample),len(all_mcmc)))
    
    print('MY SAMPLER IS', sampler)
    
    if sampler == 'slice':       

        for i in tqdm(range(len(all_mcmc))):
            with h5.File(f'{dirname}/mcmc_{all_mcmc[i]}.h5', 'r') as f:
                for j, s in enumerate(to_sample):
                    if 'model_params_' + s in f['scalars'].keys():
                        samples[j,i] = f['scalars/model_params_' + s][:][0]
                    elif 'model_params_cosmology.' + s in f['scalars'].keys():
                        samples[j,i] = f['scalars/model_params_cosmology.' + s][:][0]
                    elif s == 'sig_v':
                        samples[j,i] = float(config['model'][s])
                    elif s.startswith('bulk_flow'):
                        if s[-1] == 'x':
                            samples[j,i] = np.array(ast.literal_eval(config['model']['bulk_flow']))[0]
                        elif s[-1] == 'y':
                            samples[j,i] = np.array(ast.literal_eval(config['model']['bulk_flow']))[1]
                        elif s[-1] == 'z':
                            samples[j,i] = np.array(ast.literal_eval(config['model']['bulk_flow']))[2]
                        else:
                            raise NotImplementedError
                    else:
                        if s in config[f'cosmology'].keys():
                            samples[j,i] = float(config['cosmology'][s])
                        else:
                            print("NOT THERE")
                            samples[j,i] = float(config[f'sample_{s[-1]}'][s[:-1]]) 
    
    elif sampler in ['hmc', 'mvslice']:
        
        if sampler == 'hmc':
            key_name = 'attributes'
            key_name = 'model_params'
        elif sampler == 'mvslice':
            key_name = 'model_paramsattributes'
        
        # Get order in which model parameters are stored
        if os.path.isfile(f'{dirname}/model_params.txt'):
            with open(f'{dirname}/model_params.txt', 'r') as file:
                model_params = [line.strip() for line in file]
        else:
            model_params = []
            
        print(model_params)
        
        for i in tqdm(range(len(all_mcmc))):
            with h5.File(f'{dirname}/mcmc_{all_mcmc[i]}.h5', 'r') as f:
                if key_name in f['scalars'].keys():
                    data = f[f'scalars/{key_name}'][:]
                else:
                    data = None
                for j, s in enumerate(to_sample):
                    if s in model_params:
                        samples[j,i] = data[model_params.index(s)]
                    elif 'model_params_cosmology.' + s in f['scalars'].keys():
                        samples[j,i] = f['scalars/model_params_cosmology.' + s][:][0]
                    elif s == 'sig_v':
                        samples[j,i] = float(config['model'][s])
                    elif s in config[f'cosmology'].keys():
                        samples[j,i] = float(config['cosmology'][s])
                    elif s.startswith('bulk_flow'):
                        idx = {'x':0, 'y':1, 'z':2}
                        idx = idx[s[-1]]
                        samples[j,i] = np.array(ast.literal_eval(config['model']['bulk_flow']))[idx]
                    else:
                        samples[j,i] = float(config[f'sample_{s[-1]}'][s[:-1]]) 
    else:
        raise NotImplementedError

    return to_sample, samples


def get_truths(ini_name, to_sample):
    
    config = configparser.ConfigParser()
    config.read(ini_name)
    
    truths = [None] * len(to_sample)
    
    for i, s in enumerate(to_sample):
        if s in config[f'cosmology'].keys():
            truths[i] = float(config['cosmology'][s])
        elif s == 'sig_v':
            truths[i] = float(config['model'][s])
        elif s.startswith('bulk_flow'):
            idx = {'x':0, 'y':1, 'z':2}
            idx = idx[s[-1]]
            truths[i] = np.array(ast.literal_eval(config['model']['bulk_flow']))[idx]
        else:
            truths[i] = float(config[f'sample_{s[-1]}'][s[:-1]]) 
    
    return truths


def crop_field(ini_name, field):
    
    config = configparser.ConfigParser()
    config.read(ini_name)
    Rmax = float(config['mock']['R_max'])
    xmin = float(config['system']['corner0'])
    L = float(config['system']['L0'])
    N = int(config['system']['N0'])
    x = np.linspace(xmin, xmin+L, N)
    m = np.abs(x) < Rmax
    L = x[m].max() - x[m].min()
    
    return field[m][:, m][:, :, m], L


def compute_ensemble_mean_field(ini_name, dirname, nframe, iter_max, iter_min, cut_field=True):
    """
    Compute the mean and std deviation of the inferred density field
    """

    print('Computing ensemble mean field')
    
    # Which steps to use
    all_mcmc = get_mcmc_steps(dirname, nframe, iter_max, iter_min=iter_min)

    #COMPUTE THE MEAN-DENSITY FIELD
    for i in tqdm(range(len(all_mcmc))):
        idx = all_mcmc[i]
        with h5.File(dirname + "/mcmc_%d.h5" % idx,'r') as mcmc_file:
            temp_field = np.array(mcmc_file['scalars/BORG_final_density'][...],dtype=np.float64)
            if i == 0:
                mean_field = np.array(np.full(temp_field.shape,0),dtype=np.float64)
                std_field = np.array(np.full(temp_field.shape,0),dtype=np.float64)
            mean_field += temp_field
            std_field += temp_field*temp_field
    mean_field = mean_field/np.float64(len(all_mcmc))
    std_field = std_field/np.float64(len(all_mcmc)) # < delta^2 >
    std_field = np.sqrt(std_field - mean_field **2) # < delta^2 > - < delta >^2
    
    # Cut the density field if needed
    if cut_field:
        mean_field, _ = crop_field(ini_name, mean_field)
        std_field, _ = crop_field(ini_name, std_field)
        
    return mean_field, std_field


def get_mock_field(ini_name, dirname, which_field='delta', cut_field=True):
    with h5.File(f'{dirname}/mock_data.h5', 'r') as f:
        if which_field == 'delta':
            dens = f['scalars/BORG_final_density'][:]
        elif which_field == 'ics':
            dens = f['scalars/s_field'][:]
    if cut_field:
        dens, L = crop_field(ini_name, dens)
    else:
        config = configparser.ConfigParser()
        config.read(ini_name)
        L = float(config['system']['L0'])
    return dens, L


def get_spectra(ini_file, dirname, nframe, iter_max, iter_min, which_field='delta', cut_field=True):
    
    # Which steps to use
    all_mcmc = get_mcmc_steps(dirname, nframe, iter_max, iter_min=iter_min)
    
    if which_field == 'delta':
        MAS = "CIC"
    elif which_field == 'ics':
        MAS = None
    else:
        raise NotImplementedError
    
    # Compute original power spectrum
    delta1, boxsize = get_mock_field(ini_name, dirname, which_field=which_field, cut_field=cut_field)
    print("BOXSIZE", boxsize)
    Pk = PKL.Pk(delta1.astype(np.float32), boxsize, axis=0, MAS=MAS, threads=1, verbose=True)
    k      = Pk.k3D
    Pk_true     = Pk.Pk[:,0]
    
    # Get other spectra
    all_pk = np.zeros((len(all_mcmc), len(k)))
    all_r = np.zeros((len(all_mcmc), len(k)))
    for i in tqdm(range(len(all_mcmc))):
        idx = all_mcmc[i]
        with h5.File(dirname + "/mcmc_%d.h5" % idx,'r') as mcmc_file:
            if which_field == 'delta':
                delta2= np.array(mcmc_file['scalars/BORG_final_density'][...],dtype=np.float64)
            elif which_field == 'ics':
                delta2 = np.array(mcmc_file['scalars/s_field'][...],dtype=np.float64)
            else:
                raise NotImplementedError
            if cut_field:
                delta2, _ = crop_field(ini_name, delta2)
        with suppress_stdout():
            Pk = PKL.XPk([delta1.astype(np.float32),delta2.astype(np.float32)], boxsize, axis=0, MAS=[MAS, MAS], threads=1)
        all_pk[i,:] = Pk.Pk[:,0,1]  #monopole of field 2
        all_r[i,:] = Pk.XPk[:,0,0] / np.sqrt(Pk.Pk[:,0,1] * Pk.Pk[:,0,0])
        
    return k, Pk_true, all_pk, all_r


def get_both_fields(ini_file, dirname, step, which_field='delta', cut_field=True):
    
    # Mock
    delta1, boxsize = get_mock_field(ini_name, dirname, which_field=which_field, cut_field=cut_field)
    
    # Step
    with h5.File(dirname + "/mcmc_%d.h5" % step,'r') as mcmc_file:
        if which_field == 'delta':
            delta2= np.array(mcmc_file['scalars/BORG_final_density'][...],dtype=np.float64)
        elif which_field == 'ics':
            delta2 = np.array(mcmc_file['scalars/s_field'][...],dtype=np.float64)
        else:
            raise NotImplementedError
        if cut_field:
            delta2, _ = crop_field(ini_name, delta2)
    
    return delta1, delta2


def get_likelihood_values(dirname, nframe, iter_max, iter_min):
    
    all_mcmc = get_mcmc_steps(dirname, nframe, iter_max, iter_min=iter_min)
    
    all_logL = np.zeros(len(all_mcmc))
    all_logprior = np.zeros(len(all_mcmc))
    for i in tqdm(range(len(all_mcmc))):
        with h5.File(f'{dirname}/mcmc_{all_mcmc[i]}.h5', 'r') as f:
            s_hat = f['scalars/s_hat_field'][:]
            all_logL[i] = f['scalars/hmc_Elh'][:]
            all_logprior[i] = f['scalars/hmc_Eprior'][:]
            
    return all_logL, all_logprior

Plot the non-ic parameters

In [82]:
# 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
ini_name = f'{dirname}/ini_file.ini'
nframe = 700
iter_min = -1
# iter_max = -1

names, 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 ['mua', 'alpha', 'lam', 'sig_v', 'bulk_flow']
MY SAMPLER IS mvslice
['mua0', 'mua1', 'mua2', 'mua3', 'alpha0', 'alpha1', 'alpha2', 'alpha3', 'lam0', 'lam1', 'lam2', 'lam3', 'sig_v', 'bulk_flow_x', 'bulk_flow_y', 'bulk_flow_z']
100%|██████████| 849/849 [00:36<00:00, 23.25it/s]
(16, 849)

Trace plot

In [83]:
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(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)
No description has been provided for this image

Corner plot

In [79]:
burn_in = 0

# Full corner plot
corner.corner(
    samples[:,burn_in:].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,burn_in:].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')
No description has been provided for this image

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
No description has been provided for this image

Analyse density fields

In [6]:
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
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%|██████████| 102/102 [00:04<00:00, 21.45it/s]
No description has been provided for this image
No description has been provided for this image
In [4]:
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

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.00 seconds
100%|██████████| 102/102 [00:01<00:00, 74.97it/s]
BOXSIZE 198.41269841269843

Computing power spectrum of the field...
Time to complete loop = 0.00
Time taken = 0.00 seconds
100%|██████████| 102/102 [00:02<00:00, 45.87it/s]
No description has been provided for this image
No description has been provided for this image
In [5]:
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
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')
No description has been provided for this image

Plot likelihood as a function of step

In [6]:
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
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_808471/3453443321.py:270: 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_808471/3453443321.py:271: 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:07<00:00, 14.00it/s]
No description has been provided for this image

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 [ ]: