Add analysis script

analysis.py

@@ -0,0 +1,130 @@
+from tqdm import tqdm
+import h5py as h5
+import numpy as np
+import configparser
+import sys
+from contextlib import contextmanager
+import os
+import Pk_library as PKL
+import glob
+
+@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
+
+
+def get_mcmc_steps(dirname, nframe, iter_max=-1, iter_min=-1):
+    """
+    Obtain evenly-spaced sample of MCMC steps
+
+    Args:
+        dirname (str): Directory containing MCMC files
+        nframe (int): Number of frames to sample
+        iter_max (int): Maximum MCMC step to sample
+        iter_min (int): Minimum MCMC step to sample
+
+    Returns:
+        all_mcmc (list): List of MCMC steps to sample
+
+    """
+    
+    # Get all MCMC steps
+    all_mcmc = glob.glob(dirname + '/mcmc_*.h5')
+    all_mcmc = np.sort([int(os.path.basename(m)[len('mcmc_'):-3]) for m in all_mcmc])
+
+    # Get in range
+    if iter_max >= 0:
+        all_mcmc = all_mcmc[all_mcmc <= iter_max]
+    all_mcmc = all_mcmc[all_mcmc >= iter_min]
+    
+    # Subsample
+    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 compute_ensemble_mean_field(ini_name, dirname, mcmc_steps, which_field='BORG_final_density'):
+    """
+    Compute the mean and std deviation of the inferred density field
+    """
+
+    print('Computing ensemble mean field')
+
+    #COMPUTE THE MEAN-DENSITY FIELD
+    for i in tqdm(range(len(mcmc_steps))):
+        idx = mcmc_steps[i]
+        with h5.File(dirname + "/mcmc_%d.h5" % idx,'r') as mcmc_file:
+            temp_field = np.array(mcmc_file[f'scalars/{which_field}'][...],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(mcmc_steps))
+    std_field = std_field/np.float64(len(mcmc_steps)) # < delta^2 >
+    std_field = np.sqrt(std_field - mean_field **2) # < delta^2 > - < delta >^2
+        
+    return mean_field, std_field
+
+
+def get_mock_field(ini_name, dirname, which_field='BORG_final_density'):
+    
+    with h5.File(f'{dirname}/mock_data.h5', 'r') as f:
+        dens = f[f'scalars/{which_field}'][:]
+
+    config = configparser.ConfigParser()
+    config.read(ini_name)
+    L = float(config['system']['L0'])
+
+    return dens, L
+
+
+def get_spectra(ini_name, dirname, mcmc_steps, which_field='BORG_final_density', MAS=''):
+    
+    if MAS == '':
+        if which_field == 'BORG_final_density':
+            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)
+        
+    # Turn to overdensity
+    if which_field == 'BORG_final_density':
+        delta1 = (1. + delta1) / (1. + delta1).mean() - 1.
+        
+    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(mcmc_steps), len(k)))
+    all_r = np.zeros((len(mcmc_steps), len(k)))
+    for i in tqdm(range(len(mcmc_steps))):
+        idx = mcmc_steps[i]
+
+        with h5.File(dirname + "/mcmc_%d.h5" % idx,'r') as mcmc_file:
+            delta2 = np.array(mcmc_file['scalars/{which_field}'][...],dtype=np.float64)
+
+        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