CubicGalileonEmu

GP emulator for boost factor in cubic Galileon gravity model

Modified Gravity emulator for boost in the dark matter power spectra

Install (for development, not production)

git clone https://github.com/nesar/CubicGalileonEmu.git
cd CubicGalileonEmu/
pip install -e '.[dev]'

Basic rundown

A few imports

from CubicGalileonEmu.load import *
from CubicGalileonEmu.viz import *
from CubicGalileonEmu.pca import *
from CubicGalileonEmu.gp import *
from CubicGalileonEmu.emu import *
from CubicGalileonEmu.mcmc import *

import pandas as pd
import numpy as np
import matplotlib.pylab as plt


if_train_all = False ## Re-train all the models. Time-consuming. 
if_mcmc_all = False  ## Full MCMC run. Time-consuming. 
if_savefig = False

Loading files

Bk_all, Bk_all_smooth, k_all, z_all = load_boost_data()
Bk_lin_all, _, _ = load_boost_data_lin()
p_all = load_params()

# Load dedicated validation datasets
Bk_val, Bk_val_smooth, k_val, z_val = load_boost_data(LIBRARY_BK_FILE_VAL, LIBRARY_ZK_FILE_VAL)
p_val = load_params(LIBRARY_PARAM_FILE_VAL)

Bk_edge, Bk_edge_smooth, k_edge, z_edge = load_boost_data_edges()
p_edge = load_params(LIBRARY_PARAM_FILE_VAL_EDGES)[:Bk_edge.shape[0]]

A few plotting routines

Experimental design

df_all = pd.DataFrame(
    np.vstack([p_all, p_val, p_edge]),
    columns=PARAM_NAME
)
colors = (
    ['b'] * p_all.shape[0] +    # training
    ['r'] * p_val.shape[0] +    # generic validation
    ['g'] * p_edge.shape[0]     # edge validation
)
plot_scatter_matrix(df_all, colors);
# blue=training (50), red=generic validation (25), green=edge validation

Boost metrics colored by cosmology parameters

color_by_index = 4
z_index = 0

f, ax = plt.subplots(1, 2, figsize=(14, 4))

plot_lines_with_param_color(p_all[:, color_by_index], 
                            k_all, 
                            Bk_all[:, z_index, :], 
                            'Training data, z=' + str(z_all[z_index]), 
                            r'$k [h/Mpc]$', 
                            r'$B(k)$', 
                            PARAM_NAME[color_by_index], ax=ax[0]);

plot_lines_with_param_color(p_all[:, color_by_index], 
                            k_all, 
                            Bk_all[:, z_index, :]/Bk_lin_all[:, z_index, :], 
                            'Training data, z=' + str(z_all[z_index]), 
                            r'$k [h/Mpc]$', 
                            r'$Q(k)$', 
                            PARAM_NAME[color_by_index], ax=ax[1]);

color_by_index = 3
z_index = 21

f, ax = plt.subplots(1, 2, figsize=(14, 4))


plot_lines_with_param_color(p_all[:, color_by_index], 
                            k_all, 
                            Bk_all[:, z_index, :], 
                            'Training data, z=' + str(z_all[z_index]), 
                            r'$k [h/Mpc]$', 
                            r'$B(k)$', 
                            PARAM_NAME[color_by_index],
                            ax=ax[0]);

plot_lines_with_param_color(p_all[:, color_by_index], 
                            k_all, 
                            Bk_all[:, z_index, :]/Bk_lin_all[:, z_index, :], 
                            'Training data, z=' + str(z_all[z_index]), 
                            r'$k [h/Mpc]$', 
                            r'$Q(k)$', 
                            PARAM_NAME[color_by_index],
                            ax=ax[1]);

z_index = 0

f, ax = plt.subplots(1, 2, figsize=(14, 4))


f = plot_lines_with_param_color(z_all, 
                            k_all, 
                            Bk_all[16, :, :],
                            'Training data', 
                            r'$k [h/Mpc]$', 
                            r'$B(k)$', 
                            'redshift',
                            ax=ax[0]);

f = plot_lines_with_param_color(z_all, 
                            k_all, 
                            Bk_all[16, :, :]/Bk_lin_all[16, :, :],
                            'Training data', 
                            r'$k [h/Mpc]$', 
                            r'$Q(k)$', 
                            'redshift',
                            ax=ax[1]);

Training involves: PCA, GP fitting.

## Data prep
z_index = 8
# y_vals = Bk_all[:, z_index, :]/Bk_lin_all[:, z_index, :] ## Unsmoothed Q(k)
y_vals = Bk_all_smooth[:, z_index, :] ## for B ### Using smooth values for emulation
# y_ind = np.arange(0, y_vals.shape[1])
y_ind = k_all

# Find matching z index in validation grids (z_val/z_edge have 101 z values vs training's 51)
z_target = z_all[z_index]
z_val_index = np.argmin(np.abs(z_val - z_target))
z_edge_index = np.argmin(np.abs(z_edge - z_target))

# Train on ALL 50 simulations
train_indices = list(np.arange(50))
p_all_train = p_all[train_indices]
# y_vals_train = Bk_all[:, z_index, :][train_indices]/Bk_lin_all[:, z_index, :][train_indices] ## For Q
y_vals_train = Bk_all_smooth[train_indices, z_index, :]

# Generic validation set - select up to 5 samples
n_val_select = min(5, p_val.shape[0])
val_select_indices = np.arange(n_val_select)
input_params_val = p_val[val_select_indices]
# target_vals_val = Bk_val_smooth[:, z_val_index, :][val_select_indices]  ## smoothed
target_vals_val = Bk_val[:, z_val_index, :][val_select_indices]

# Edge validation set - select up to 5 samples
n_edge_select = min(5, p_edge.shape[0])
edge_select_indices = np.arange(n_edge_select)
input_params_edge = p_edge[edge_select_indices]
# target_vals_edge = Bk_edge_smooth[:, z_edge_index, :][edge_select_indices]  ## smoothed
target_vals_edge = Bk_edge[:, z_edge_index, :][edge_select_indices]

print("Training on %d simulations" % len(train_indices))
print("Generic validation: %d samples, z_val_index=%d (z=%.4f)" % (n_val_select, z_val_index, z_val[z_val_index]))
print("Edge validation: %d samples, z_edge_index=%d (z=%.4f)" % (n_edge_select, z_edge_index, z_edge[z_edge_index]))
print("Redshift: z=%.4f (training z_index=%d)" % (z_all[z_index], z_index))

Training on 50 simulations
Generic validation: 5 samples, z_val_index=16 (z=0.1860)
Edge validation: 5 samples, z_edge_index=16 (z=0.1860)
Redshift: z=0.1860 (training z_index=8)

sepia_data = sepia_data_format(p_all_train, y_vals_train, y_ind)
# sepia_data = sepia_data_by_redshift(redshift=0.01)

print(p_all_train.shape, y_vals_train.shape, y_ind.shape)


print(sepia_data)
model_filename = '../CubicGalileonEmu/model/multivariate_model_z_index' + str(z_index) 

# sepia_model = do_pca(sepia_data, exp_variance=0.95)
sepia_model_pca = do_pca(sepia_data, exp_variance=0.999)

(50, 5) (50, 768) (768,)
This SepiaData instance implies the following:
This is a simulator (eta)-only model, y dimension 768
m  =    50 (number of simulated data)
p  =     1 (number of inputs)
q  =     5 (number of additional simulation inputs)
pu NOT SET (transformed response dimension); call method create_K_basis 

sepia_model = do_gp_train(sepia_model_pca, model_filename)
plot_train_diagnostics(sepia_model)

Starting tune_step_sizes...
Default step sizes:
betaU
[[0.1 0.1 0.1 0.1 0.1]
 [0.1 0.1 0.1 0.1 0.1]
 [0.1 0.1 0.1 0.1 0.1]
 [0.1 0.1 0.1 0.1 0.1]
 [0.1 0.1 0.1 0.1 0.1]
 [0.1 0.1 0.1 0.1 0.1]]
lamUz
[[5. 5. 5. 5. 5.]]
lamWs
[[100. 100. 100. 100. 100.]]
lamWOs
[[100.]]

Step size tuning: 100%|██████████| 50/50 [05:24<00:00,  6.50s/it]

Done with tune_step_size.
Selected step sizes:
betaU
[[0.22597448 0.34265107 0.23392414 0.46468998 0.24461205]
 [0.09432116 0.24376015 3.09893801 1.62011445 1.41086578]
 [0.05127119 0.09112937 3.86650993 0.3726364  0.29293188]
 [0.11113476 0.25862219 0.72659085 0.10684409 0.94793404]
 [0.08170817 0.24372759 0.78772743 0.84790403 1.31999935]
 [1.04092879 0.77135154 4.44962602 0.5536803  2.75215399]]
lamUz
[[1.04561809 0.77052503 1.13455342 1.3466235  1.29424591]]
lamWs
[[ 610.75855548  274.93623072 4351.91290499 3379.77215584 4295.47826176]]
lamWOs
[[132.99667668]]

MCMC sampling: 100%|██████████| 1000/1000 [05:20<00:00,  3.12it/s]

Model saved to ../CubicGalileonEmu/model/multivariate_model_z_index8.pkl
No thetas to plot

Load existing model

sepia_model = gp_load(sepia_model_pca, model_filename)

WARNING: make sure this model was instantiated with the same input data as the model corresonding to this saved model info.

Single-redshift emulation for new cosmological parameters

# --- Generic validation set ---
pred_mean_val, pred_std_val = emulate(sepia_model, sepia_data, input_params_val)
print("Generic validation (z=%.4f):" % z_val[z_val_index])
f = validation_plot(k_all, target_vals_val, pred_mean_val, pred_std_val, xy_lims=[2e-2, 1e1, 0.98, 1.06]);

# --- Edge validation set ---
pred_mean_edge, pred_std_edge = emulate(sepia_model, sepia_data, input_params_edge)
print("Edge validation (z=%.4f):" % z_edge[z_edge_index])
f_edge = validation_plot(k_all, target_vals_edge, pred_mean_edge, pred_std_edge, xy_lims=[2e-2, 1e1, 0.98, 1.35]);

Generic validation (z=0.1860):
Edge validation (z=0.1860):

Sensitivity analysis from the emulator

f = sensitivity_plot(k_all, p_all, sepia_model, sepia_data, emulate, PARAM_NAME, xy_lims=[2e-2, 1e1, 0.98, 1.35]);
# f = sensitivity_plot(k_all, p_all, sepia_model, sepia_data, emulate, PARAM_NAME, xy_lims=[2e-2, 1e1, 0.9, 1.1]); ## Unsmoothed Q(k)

Multi-redshift emulation

Train all the models

if if_train_all:
    
    do_gp_train_multiple(model_dir='../CubicGalileonEmu/model/', 
                        p_train_all = p_all[train_indices],
                        # y_vals_all = Bk_all[train_indices]/Bk_lin_all[train_indices] ## Unsmoothed Q(k)
                        y_vals_all = Bk_all_smooth[train_indices],
                        y_ind_all = k_all,
                        z_index_range=range(49))

Load all trained models

sepia_model_list, sepia_data_list = load_model_multiple(model_dir='../CubicGalileonEmu/model/', 
                                        p_train_all=p_all[train_indices],
                                        # y_vals_all=Bk_all[train_indices]/Bk_lin_all[train_indices], ## Unsmoothed Q(k)
                                        y_vals_all=Bk_all_smooth[train_indices],
                                        y_ind_all=k_all,
                                        z_index_range=range(49)
                                        );

Number of models loaded: 49 from: ../CubicGalileonEmu/model/

#### Emulator uncertainty across parameter range


input_params0 = input_params_val[0]
z_inputs = 0.05
input_params_and_redshift = np.append(input_params0, z_inputs)
print(input_params_and_redshift[np.newaxis, :])

# sepia_data_select = sepia_data_by_redshift(redshift=z_inputs)
# sepia_data_select1 = sepia_data_by_redshift(redshift=z_all[2])
# sepia_data_select2 = sepia_data_by_redshift(redshift=z_all[3])



sepia_data_select = sepia_data_list[23]
sepia_data_select1 = sepia_data_list[22]
sepia_data_select2 = sepia_data_list[24]

emulated_with_redshift, emulated_with_redshift_err = emu_redshift(input_params_and_redshift[np.newaxis, :], sepia_model_list, sepia_data_list, z_all)
## There is an unknown issue with z_index=5 model, (sepia_model_list[6])
# emulate(sepia_model_list[6], sepia_data_select, input_params_and_redshift[:-1])[0]

plt.figure(433)
plt.plot(k_all, emulated_with_redshift[:, 0], label='interp at z=%.4f'%input_params_and_redshift[-1], lw=5, ls='--')
plt.plot(k_all, emulate(sepia_model_list[22], sepia_data_select1, input_params_and_redshift[:-1])[0], label='grid z=%.4f'%z_all[2])
plt.plot(k_all, emulate(sepia_model_list[24], sepia_data_select2, input_params_and_redshift[:-1])[0], label='grid z=%.4f'%z_all[3])
plt.legend()
plt.title('Comparison of redshift-space interpolation')
# plt.plot(k_all, emulate(sepia_model_list[0], input_params))
# plt.plot(k_all, emulate(sepia_model_list[0], input_params))

plt.show()

[[0.321  0.9626 2.092  0.6688 0.03   0.05  ]]

len(emulate(sepia_model_list[8], sepia_data, input_params_and_redshift[:-1]))

### Will give an error:
### ValueError: matmul: Input operand 1 has a mismatch in its core dimension 0, with gufunc signature (n?,k),(k,m?)->(n?,m?) (size 4 is different from 5)
### This happens when the sepia_model and sepia_data do not match

2

Emulator confidence across parameter range

# Parameter settings
steps = 50  # Number of steps in the grid for each parameter
param_name_extended = np.append(PARAM_NAME, 'Redshift')
red_min = 0
red_max = 3
red_mean = 1.0

param_min = np.append(p_all.min(axis=0), red_min)
param_max = np.append(p_all.max(axis=0), red_max)
param_mean = np.append(p_all.mean(axis=0), red_mean)

# Compute outputs and errors for a range of parameter values
def compute_errors(param_grid):
    # print(param_grid.shape)
    errors = np.zeros(shape=(param_grid.shape[0], ))
    # errors = np.array([np.mean(emu_redshift(params[np.newaxis, :], sepia_model_list, z_all)[1][:, 0, :]**2) for params in param_grid])
    for par_indx in range(errors.shape[0]):
        emu_pred = np.array(emu_redshift(param_grid[par_indx][np.newaxis, :], sepia_model_list, sepia_data_list, z_all))
        errors_emu = emu_pred[1, :, 0]
        # print(errors_emu.shape)

        errors[par_indx] = np.max(errors_emu)
        # print(errors_emu)

    return errors.reshape(steps, steps)


# Example usage:
param_indices = [4, 2]  # Indices of parameters to vary
fixed_indices = [i for i in range(len(param_name_extended)) if i not in param_indices]
fixed_params = {param_name_extended[i]: param_mean[i] for i in fixed_indices}

param_grid = generate_param_grid_with_fixed(param_name_extended, param_indices, fixed_params, param_min, param_max, steps)

errors = compute_errors(param_grid)

f = plot_error_heatmap( errors, 
                       [param_name_extended[i] for i in param_indices], 
                       [(param_min[param_indices[0]], param_max[param_indices[0]]), (param_min[param_indices[1]], param_max[param_indices[1]])]
                       )

if if_savefig: 
    f.savefig('../../../Plots/heatmap_params_4_5.png', bbox_inches='tight')

Parameter inference via MCMC using the emulator

ndim = 5
nwalkers = 100  # 500
nrun_burn = 100  # 300
nrun = 700  # 1000

Creating mock observational data

target_indx = 0 #0 ,14, 35
z_index = 2
L = 32
fake_obs_data_index_every = 1

redshift = z_all[z_index]
k_cond = k_all < 1
x = k_all[::fake_obs_data_index_every][k_cond]
y = Bk_all[:, z_index, :][target_indx][::fake_obs_data_index_every][k_cond]
y = y + 1e-4*np.random.normal(0.0, 1.0, size=y.shape)
yerr = np.sqrt( Bk_all[:, z_index, :][target_indx][::fake_obs_data_index_every][k_cond]*(L**3))/(L**3)

x_grid = k_all
params_calib = p_all[target_indx][:, np.newaxis].T
print('redshift: ', redshift)

redshift:  0.041

f, a = plt.subplots(1,1, figsize = (8, 5)) 
input_params_and_redshift = np.append(params_calib, redshift)
bk_target, err_target = emu_redshift(input_params_and_redshift[np.newaxis, :], sepia_model_list, sepia_data_list, z_all)
a.plot(k_all, bk_target[:, 0], label='Emulated at target params', lw=5, ls='--')
a.errorbar(x, y, yerr, label='Target mock observations', ls='none', lw=1, color = "r")
a.scatter(x, y, s = 5, marker = "h", color = "r")


a.plot(k_all, Bk_all[:, z_index, :].T, 'k', alpha=0.1)


plt.plot(k_all, emulate(sepia_model_list[z_index], sepia_data_list[z_index], input_params_and_redshift[:-1])[0], label='grid z=%.4f'%z_all[z_index])
plt.plot(k_all, emulate(sepia_model_list[z_index+1], sepia_data_list[z_index + 1], input_params_and_redshift[:-1])[0], label='grid z=%.4f'%z_all[z_index + 1])

string_print0 = 'Target Params \n\n' 
string_print1 = PARAM_NAME[0] + '= %.4f'%input_params_and_redshift[0] + '\n'
string_print2 = PARAM_NAME[1] + '= %.4f'%input_params_and_redshift[1] + '\n'
string_print3 = PARAM_NAME[2] + '= %.4f'%input_params_and_redshift[2] + '\n'
string_print4 = PARAM_NAME[3] + '= %.4f'%input_params_and_redshift[3] + '\n'
string_print5 = PARAM_NAME[4] + '= %.4f'%input_params_and_redshift[4] + '\n'
string_print6 = 'redshift' + '= %.4f'%input_params_and_redshift[5] 


string_print = string_print0 + string_print1 + string_print2 + string_print3 + string_print4 + string_print5 + string_print6

props = dict(boxstyle='round', facecolor='gray', alpha=0.2)
plt.text(1.02, 0.1, string_print, transform=a.transAxes, fontsize=12, bbox=props)


a.set_xscale('log')
plt.title('pre-MCMC')
a.set_xlabel(r'$k [h/Mpc]$')
a.set_ylabel(r'$B(k)$')
plt.legend()

allMax = np.max(p_all, axis = 0)
allMin = np.min(p_all, axis = 0)

param1 = [PARAM_NAME[0], params_calib[0][0], allMin[0], allMax[0]] 
param2 = [PARAM_NAME[1], params_calib[0][1], allMin[1], allMax[1]]
param3 = [PARAM_NAME[2], params_calib[0][2], allMin[2], allMax[2]]
param4 = [PARAM_NAME[3], params_calib[0][3], allMin[3], allMax[2]]
param5 = [PARAM_NAME[4], params_calib[0][4], allMin[4], allMax[4]]

params_list = [param1, param2, param3, param4, param5]

pos0 = chain_init(params_list, ndim, nwalkers)
sampler = define_sampler(redshift, ndim, nwalkers, params_list, x_grid, sepia_model_list, sepia_data_list, z_all, x, y, yerr)

MCMC run - first burn, then full.

pos, prob, state, samples, sampler, autocorr, index = do_mcmc(sampler, pos0, nrun_burn, ndim, if_burn=True)

if if_mcmc_all: # Full MCMC-run, will be slow
    pos, prob, state, samples, sampler, autocorr, index = do_mcmc(sampler, pos, nrun, ndim, if_burn=False)

p_mcmc = mcmc_results(samples)

fig = plot_mcmc(samples, params_list, if_truth_know=True)
if if_savefig: 
    plt.savefig('../../../Plots/mcmc_plot.png', bbox_inches='tight')

Burn-in phase
time (minutes): 15.385760943094889

100%|██████████| 100/100 [15:37<00:00,  9.38s/it]

mcmc results: 0.3141152383554347 1.0814831836429077 2.543808187521769 0.6824603082006169 0.47648316627575504

f, a = plt.subplots(1,1, figsize = (8, 5)) 
input_params_and_redshift = np.append(p_mcmc, redshift)
bk_mcmc, err_mcmc = emu_redshift(input_params_and_redshift[np.newaxis, :], sepia_model_list, sepia_data_list, z_all)
a.plot(k_all, bk_mcmc[:, 0], label='Emulated at best MCMC', lw=3, ls='--')
a.errorbar(x, y, yerr, label='Mock target', ls='none', lw=1, color = "r")
a.scatter(x, y, s = 5, marker = "h", color = "r", alpha=0.5)

a.plot(k_all, Bk_all[:, z_index, :].T, 'k', alpha=0.1)


# plt.plot(k_all, emulate(sepia_model_list[z_index], input_params_and_redshift[:-1])[0], label='Z1')
# plt.plot(k_all, emulate(sepia_model_list[z_index+1], input_params_and_redshift[:-1])[0], label='Z2')

string_print0 = 'Target Params \n\n' 
string_print1 = PARAM_NAME[0] + '= %.3f'%params_calib[0][0] + '\n'
string_print2 = PARAM_NAME[1] + '= %.3f'%params_calib[0][1] + '\n'
string_print3 = PARAM_NAME[2] + '= %.3f'%params_calib[0][2] + '\n'
string_print4 = PARAM_NAME[3] + '= %.3f'%params_calib[0][3] + '\n'
string_print5 = PARAM_NAME[4] + '= %.3f'%params_calib[0][4] + '\n'
string_print6 = 'redshift' + '= %.3f'%redshift


string_print = string_print0 + string_print1 + string_print2 + string_print3 + string_print4 + string_print5 + string_print6

props = dict(boxstyle='round', facecolor='gray', alpha=0.2)
plt.text(1.02, 0.5, string_print, transform=a.transAxes, fontsize=12, bbox=props)

string_print0_mcmc = 'Optimized Params \n\n' 
string_print1_mcmc = PARAM_NAME[0] + '= %.3f'%p_mcmc[0] + '\n'
string_print2_mcmc = PARAM_NAME[1] + '= %.3f'%p_mcmc[1] + '\n'
string_print3_mcmc = PARAM_NAME[2] + '= %.3f'%p_mcmc[2] + '\n'
string_print4_mcmc = PARAM_NAME[3] + '= %.3f'%p_mcmc[3] + '\n'
string_print5_mcmc = PARAM_NAME[4] + '= %.3f'%p_mcmc[4] 

string_print_mcmc = string_print0_mcmc + string_print1_mcmc + string_print2_mcmc + string_print3_mcmc + string_print4_mcmc + string_print5_mcmc

props = dict(boxstyle='round', facecolor='blue', alpha=0.2)
plt.text(1.02, 0.05, string_print_mcmc, transform=a.transAxes, fontsize=12, bbox=props)



a.set_xscale('log')
plt.title('B(k) at MCMC constraints')
a.set_xlabel(r'$k [h/Mpc]$')
a.set_ylabel(r'$B(k)$')
plt.legend()

if if_savefig: 
    plt.savefig('../../../Plots/mcmc_results_Bk.png', bbox_inches='tight')

if if_mcmc_all:
    tau = sampler.get_autocorr_time(tol=0)
    print(tau)

    plt.figure(43)
    plt.plot(prob)
    if if_savefig: 
        plt.savefig('../../../Plots/prob_plot.png', bbox_inches='tight')

    selected_indices_for_plot = [0, 2, 4]
    fig = plot_mcmc(samples[:, selected_indices_for_plot], [params_list[i] for i in selected_indices_for_plot], if_truth_know=True)

    if if_savefig: 
        plt.savefig('../../../Plots/mcmc_plot_reduced_params.png', bbox_inches='tight')

def plot_convergence(sampler, params_list, nrun, ndim, nwalkers):
    n_params = len(params_list)  # Number of parameters
    fig, ax = plt.subplots(n_params, 1, figsize=(20, 2 * n_params), sharex=True)
    ax[-1].set_xlabel('steps')

    for i, param in enumerate(params_list):
        ax[i].plot(np.arange(nrun), sampler.chain[:, :, i].T, lw=0.2, alpha=0.9)
        ax[i].text(0.9, 0.9, param[0], horizontalalignment='center', verticalalignment='center', transform=ax[i].transAxes, fontsize=12)

    # fig.savefig('plots/convergence_mcmc_ndim{}_nwalk{}_run{}_{}-{}.png'.format(ndim, nwalkers, nrun, summary_stat, design), dpi=100)

    return fig

if False:
    # Example usage
    plot_convergence(sampler, params_list, nrun, ndim, nwalkers)

if if_mcmc_all:

    n = 100 * np.arange(1, index + 1)
    y = autocorr[:index]
    plt.plot(n, n / 100.0, "--k")
    plt.plot(n, y)
    plt.xlim(0, n.max())
    plt.ylim(0, y.max() + 0.1 * (y.max() - y.min()))
    plt.xlabel("number of steps")
    plt.ylabel(r"mean $\hat{\tau}$");
    plt.show()

if if_mcmc_all:
    plt.plot(autocorr)
    plt.xscale('log')
    # plt.yscale('log')