Initial import

scripts/ares_tools/__init__.py

@@ -0,0 +1,16 @@
+#+
+#   ARES/HADES/BORG Package -- ./scripts/ares_tools/__init__.py
+#   Copyright (C) 2014-2020 Guilhem Lavaux <guilhem.lavaux@iap.fr>
+#   Copyright (C) 2009-2020 Jens Jasche <jens.jasche@fysik.su.se>
+#
+#   Additional contributions from:
+#      Guilhem Lavaux <guilhem.lavaux@iap.fr> (2023)
+#   
+#+
+from .read_all_h5 import *
+from .analysis import analysis, IndexTracker
+
+try:
+  from .visu import vtktools 
+except ImportError:
+  print("Skipping VTK tools")

scripts/ares_tools/analysis/__init__.py

@@ -0,0 +1,10 @@
+#+
+#   ARES/HADES/BORG Package -- ./scripts/ares_tools/analysis/__init__.py
+#   Copyright (C) 2014-2020 Guilhem Lavaux <guilhem.lavaux@iap.fr>
+#   Copyright (C) 2009-2020 Jens Jasche <jens.jasche@fysik.su.se>
+#
+#   Additional contributions from:
+#      Guilhem Lavaux <guilhem.lavaux@iap.fr> (2023)
+#   
+#+
+from .analysis import analysis, IndexTracker

scripts/ares_tools/analysis/analysis.py

@@ -0,0 +1,673 @@
+#+
+#   ARES/HADES/BORG Package -- ./scripts/ares_tools/analysis/analysis.py
+#   Copyright (C) 2014-2020 Guilhem Lavaux <guilhem.lavaux@iap.fr>
+#   Copyright (C) 2009-2020 Jens Jasche <jens.jasche@fysik.su.se>
+#
+#   Additional contributions from:
+#      Guilhem Lavaux <guilhem.lavaux@iap.fr> (2023)
+#   
+#+
+from ..read_all_h5 import explore_chain, rebuild_spliced_h5
+import errno
+import h5py as h5
+import numpy as np
+import healpy as hp
+import numexpr as ne
+import os
+import math
+from pylab import *
+import matplotlib.pyplot as plt
+from matplotlib.pyplot import figure, show
+
+#import numba
+'''
+@numba.jit
+def _special_spectrum_builder(P, PP, tmp, alpha, N_i, Pk_i):
+
+    for i in range(Pk_i.size):
+        for j in range(PP.size):
+            v = PP[j]**(-alpha)*exp(- N_i[i]*Pk_i[i] / (2 * PP[j]))
+            tmp[j] = v
+            total += v
+
+        for j in range(PP.size):
+            P[j] += tmp[j] / total
+ '''
+
+
+ #ARES/HADES/BORG image scanning class
+class IndexTracker:
+    def __init__(self, ax, X):
+        self.ax = ax
+        self.ax.set_title('use up/down keys to navigate images')
+
+        self.X = X
+        rows,cols,self.slices = X.shape
+        self.ind  = self.slices/2
+
+        cmax=X.max()
+        cmin=X.min()
+
+        self.im = self.ax.imshow(self.X[:,self.ind,:],vmax=cmax,vmin=cmin)
+        self.update()
+
+    def onscroll(self, event):
+        #print ("%s " % (event.key))
+        if event.key=='up':
+            self.ind = np.clip(self.ind+1, 0, self.slices-1)
+        else:
+            self.ind = np.clip(self.ind-1, 0, self.slices-1)
+
+
+        self.update()
+
+    def update(self):
+        self.im.set_data(self.X[:,self.ind,:])
+        self.ax.set_ylabel('slice %s'%self.ind)
+        self.im.axes.figure.canvas.draw()
+
+
+def detect_ncpus(path):
+  ncpu = 0
+  try:
+    while True:
+      with open("%s_%d" % (path,ncpu), mode= "rb") as f:
+        ncpu += 1
+  except IOError as e:
+    if e.errno != errno.ENOENT:
+      raise e
+
+  return ncpu
+
+
+
+#ARES/HADES/BORG analysis class
+class analysis:
+
+
+    def __init__(self, chain_path='.', LSS_framework='ARES', start=0, step=1):
+        self.chain_path         = chain_path
+        self.LSS_framework      = LSS_framework
+
+        self.description = "This Class is part of the ARES/HADES/BORG analysis framework"
+        self.author = "Copyright (C) 2009-2017 Jens Jasche \n Copyright (C) 2014-2017 Guilhem Lavaux"
+
+        #get chain setup
+
+        self.L0=0
+        self.L1=0
+        self.L2=0
+        self.N0=0
+        self.N1=0
+        self.N2=0
+        self.x0=0
+        self.x1=0
+        self.x2=0
+
+        self.ncpus = detect_ncpus(os.path.join(self.chain_path, "restart.h5"))
+
+        self.mcmc_list=[]
+        Fmax=start
+        while True:
+            try:
+              os.stat("mcmc_%d.h5" % Fmax)
+            except:
+              break
+            self.mcmc_list.append(Fmax)
+            Fmax += step
+        
+
+        with h5.File(os.path.join(self.chain_path, "restart.h5_0"), mode="r") as f:
+            #print markov.keys()
+            #print info.keys()
+            #print f['scalars'].keys()
+
+            self.L0 = f['scalars']['L0'][0]
+            self.L1 = f['scalars']['L1'][0]
+            self.L2 = f['scalars']['L2'][0]
+            self.N0 = int(f['scalars']['N0'][:])
+            self.N1 = int(f['scalars']['N1'][:])
+            self.N2 = int(f['scalars']['N2'][:])
+            self.xmin0 = int(f['scalars']['corner0'][:])
+            self.xmin1 = int(f['scalars']['corner1'][:])
+            self.xmin2 = int(f['scalars']['corner2'][:])
+            self.ncat = int(f['scalars']['NCAT'][:])
+
+            if(LSS_framework!='BORG'):
+                self.kmodes = f['/scalars/k_modes'][:]
+                self.nmodes = len(self.kmodes)
+
+
+        #get brefs
+            #self.k_keys = rebuild_spliced_h5(os.path.join(self.chain_path, "restart.h5"), ["scalars.k_keys"], self.ncpus)["scalars.k_keys"]
+
+            bref=[]
+            if LSS_framework != 'VIRBIUS' and LSS_framework != 'LYA':
+              for i in range(self.ncat):
+                  bref.append(f['scalars']['galaxy_bias_ref_'+str(i)][:])
+              self.bias_ref=np.array(bref)
+
+    def _internal_power(self, P, Nbins, range, unit=False):
+        if not hasattr(self, '_k'):
+          ik0 = np.fft.fftfreq(self.N0, d=self.L0/self.N0)*2*np.pi
+          ik1 = np.fft.fftfreq(self.N1, d=self.L1/self.N1)*2*np.pi
+          ik2 = np.fft.fftfreq(self.N2, d=self.L2/self.N2)*2*np.pi
+          k = self._k = np.sqrt(ik0[:,None,None]**2+ik1[None,:,None]**2+ik2[None,None,:(self.N2//2+1)]**2)
+          self._Pw, _ = np.histogram(k, bins=Nbins, range=range)
+          Pw = self._Pw
+        else:
+          k = self._k 
+          Pw = self._Pw
+
+        P, b = np.histogram(k, weights=P, bins=Nbins, range=range)
+        if not unit:
+          P /= self.L0*self.L1*self.L2
+        cond = Pw > 0
+        P[cond] /= Pw[cond]
+        return P, Pw, b 
+
+    def rebin_power_spectrum(self, chain_id, Nbins=100, range=(0,1), unit=False):
+        with h5.File(os.path.join(self.chain_path, "mcmc_%d.h5" % (chain_id,)), mode="r") as f:
+          P = f['/scalars/powerspectrum'][...] * (self.L0 * self.L1 * self.L2)
+
+        return self._internal_power(P[self.k_keys], Nbins, range, unit=unit)
+
+
+    def compute_power_spectrum_mock(self, Nbins=100, unit=False, range=(0,1)):
+        with h5.File(os.path.join(self.chain_path, "mock_data.h5"), mode="r") as f:
+          shat = f['/scalars/s_hat_field'][...]
+        return self._internal_power(ne.evaluate('real(s)**2+imag(s)**2', dict(s=shat)), Nbins, range, unit=unit)
+
+    def compute_power_spectrum_galaxydata(self, Nbins=100, range=(0,1)):
+        with h5.File(os.path.join(self.chain_path, "mock_data.h5"), mode="r") as f:
+          # FFT galaxy data  
+          dV = self.L0*self.L1*self.L2/(self.N0*self.N1*self.N2)
+          fd = np.fft.rfftn(f['/scalars/galaxy_data_0'][...])*dV
+          # remove zero mode
+          fd[0][0][0] = 0. + 0.j
+        return self._internal_power(ne.evaluate('real(s)**2+imag(s)**2', dict(s=fd)), Nbins, range)
+
+    def compute_power_shat_spectrum(self, chain_id, Nbins=100, unit=False, range=(0,1)):
+
+        dV = self.L0*self.L1*self.L2/(self.N0*self.N1*self.N2)
+
+        with h5.File(os.path.join(self.chain_path, "mcmc_%d.h5" % (chain_id,)), mode="r") as f:
+          if '/scalars/s_hat_field' in f:
+            shat = f['/scalars/s_hat_field'][...]
+          else:
+            shat = np.fft.rfftn(f['/scalars/s_field'][...])*dV
+
+        return self._internal_power(ne.evaluate('real(s)**2+imag(s)**2', dict(s=shat)), Nbins, range, unit=unit)
+
+    def compute_power_shat_cross_spectrum(self, chain_id1, chain_id2,
+                                          Nbins=100, unit=False, range=(0,1)):
+
+        dV = self.L0*self.L1*self.L2/(self.N0*self.N1*self.N2)
+
+        with h5.File(os.path.join(self.chain_path, "mcmc_%d.h5" % (chain_id1,)), mode="r") as f:
+          if '/scalars/s_hat_field' in f:
+            shat1 = f['/scalars/s_hat_field'][...]
+          else:
+            shat1 = np.fft.rfftn(f['/scalars/s_field'][...])*dV
+
+        with h5.File(os.path.join(self.chain_path, "mcmc_%d.h5" % (chain_id2,)), mode="r") as f:
+          if '/scalars/s_hat_field' in f:
+            shat2 = f['/scalars/s_hat_field'][...]
+          else:
+            shat2 = np.fft.rfftn(f['/scalars/s_field'][...])*dV
+
+        return self._internal_power(ne.evaluate('real(s1)*real(s2)+imag(s1)*imag(s2)', dict(s1=shat1,s2=shat2)), Nbins, range, unit=unit)
+    
+        #return self._internal_power(ne.evaluate('real(s)**2+imag(s)**2', dict(s=shat)), Nbins, range, unit=unit)
+
+    # maybe 'unit' argument is not sensible here...
+    def compute_power_spectrum_finaldensity(self, chain_id, Nbins=100, unit=False, range=(0,1)):
+
+        dV = self.L0*self.L1*self.L2/(self.N0*self.N1*self.N2)
+
+        with h5.File(os.path.join(self.chain_path, "mcmc_%d.h5" % (chain_id,)), mode="r") as f:
+            fd = np.fft.rfftn(f['/scalars/BORG_final_density'][...]) * dV
+
+        return self._internal_power(ne.evaluate('real(s)**2+imag(s)**2', dict(s=fd)), Nbins, range, unit=unit)
+
+    # compute power spectrum of real-space field from given path and field name
+    def compute_power_shat_spectrum_file(self, path, fieldname='phases', Nbins=100, unit=False, range=(0,1)):
+
+        dV = self.L0*self.L1*self.L2/(self.N0*self.N1*self.N2)
+
+        with h5.File(path, mode="r") as f:
+          if fieldname in f:
+              # do FFT
+              shat = np.fft.rfftn(f[fieldname][...])*dV
+              # remove zero mode
+              shat[0][0][0] = 0. + 0.j
+          else:
+              print("No field '%s' found in file." % fieldname)
+
+        return self._internal_power(ne.evaluate('real(s)**2+imag(s)**2', dict(s=shat)), Nbins, range, unit=unit)
+
+    def check_biasref(self):
+        return self.bias_ref
+
+    def get_ncat(self):
+        return self.ncat
+
+    def get_mask(self,msknr):
+        selkey = "scalars.galaxy_sel_window_%s"%msknr
+        return \
+             rebuild_spliced_h5(
+                 os.path.join(
+                    self.chain_path,"restart.h5"
+                 ),
+                 [selkey],
+                 self.ncpus
+             )[selkey] 
+
+    def get_data(self,datnr):
+        datkey = "scalars.galaxy_data_%s"%datnr
+        return \
+             rebuild_spliced_h5(
+                 os.path.join(
+                    self.chain_path,"restart.h5"
+                 ),
+                 [datkey],
+                 self.ncpus
+             )[datkey]
+
+    def scan_datacube(self,data):
+        fig = figure()
+        ax = fig.add_subplot(111)
+        plt.jet()
+        tracker = IndexTracker(ax, data)
+        fig.canvas.mpl_connect('key_press_event', tracker.onscroll)
+        show()
+
+
+    def get_2d_marginal(self,attribute_a='s_field',attribute_b='s_field',id_a=None,id_b=None, first_sample=0,last_sample=1000):
+        print( '-'*60)
+        print( 'Estimate 2d marginals for parameters ', attribute_a, ' and ', attribute_b , ' for ' , self.LSS_framework, ' run!')
+        print( '-'*60)
+
+        if(id_a==None or id_b==None):
+            print( "Error: no index chosen")
+            return -1
+
+        #2) collect chain
+        samples_a = []
+        samples_b = []
+        for i,a in explore_chain(self.chain_path, first_sample,last_sample, 1):
+            d   = a[attribute_a][:]
+            e   = a[attribute_b][:]
+
+            samples_a.append(d[id_a])
+            samples_b.append(e[id_b])
+
+        H, xedges, yedges = np.histogram2d(samples_a, samples_b)
+
+        return xedges,yedges, H
+
+    def get_cross_corcoeff(self,attribute_a='s_field',attribute_b='s_field',id_a=None,id_b=None, first_sample=0,last_sample=1000):
+        print( '-'*60)
+        print( 'Estimate 2d marginals for parameters ', attribute_a, ' and ', attribute_b , ' for ' , self.LSS_framework, ' run!')
+        print( '-'*60)
+
+        if(id_a==None or id_b==None):
+            print("Error: no index chosen")
+            return -1
+
+        #2) collect chain
+        samples_a = []
+        samples_b = []
+
+        nelements_a = len(id_a[0])
+        nelements_b = len(id_b[0])
+
+        mu_a    = np.zeros(nelements_a)
+        var_a   = np.zeros(nelements_a)
+
+        mu_b    = np.zeros(nelements_b)
+        var_b   = np.zeros(nelements_b)
+        nn=1
+        for i,a in explore_chain(self.chain_path, first_sample,last_sample, 1):
+            d   = a[attribute_a][:]
+            e   = a[attribute_b][:]
+
+            aux_a   = d[id_a]
+            aux_b   = e[id_b]
+
+            mu_a  = (nn-1.)/float(nn)*mu_a +1./float(nn)*aux_a
+            if(nn>1):
+                aux     = (mu_a-aux_a)**2
+                var_a     = (nn-1.)/nn*var_a+1./(nn-1)*aux
+
+            mu_b  = (nn-1.)/float(nn)*mu_b +1./float(nn)*aux_b
+            if(nn>1):
+                aux     = (mu_b-aux_b)**2
+                var_b     = (nn-1.)/nn*var_b+1./(nn-1)*aux
+
+            samples_a.append(aux_a)
+            samples_b.append(aux_b)
+            nn+=1
+
+        pc= np.zeros((nelements_a,nelements_b))
+        cnt=0
+        for n in range(nn-1):
+            x=samples_a[n]
+            y=samples_b[n]
+            pc += np.multiply.outer(x-mu_a, y-mu_b)
+            cnt+=1
+
+        return pc/float(cnt) #/np.sqrt(var_a*var_b)
+
+
+    def get_trace(self,attribute='s_field',element_id=None, first_sample=0,last_sample=1000):
+        print( '-'*60)
+        print( 'Record trace for parameters ', attribute , ' for ' , self.LSS_framework, ' run!')
+        print( '-'*60)
+
+        '''
+        if(element_id==None):
+            print "Error: no list of indices provided"
+            return -1
+        '''
+
+        #1) collect chain
+        samples = []
+
+        nn=1
+        for i,a in explore_chain(self.chain_path, first_sample,last_sample, 1):
+            d   = a[attribute][:]
+            if (element_id!=None):
+                samples.append(d[element_id])
+            else:
+                samples.append(d)
+            nn+=1
+
+        return samples
+
+
+    def get_corrlength(self,attribute='s_field',element_id=None,nlength=100, first_sample=0,last_sample=1000):
+        print( '-'*60)
+        print( 'Estimate correlation length for parameters ', attribute , ' for ' , self.LSS_framework, ' run!')
+        print( '-'*60)
+
+        if(element_id==None):
+            print( "Error: no list of indices provided")
+            return -1
+
+        if(nlength>last_sample-first_sample):
+            print("Warning: Chain setting not long enough set nlength to last_sample")
+            nlength = last_sample-first_sample -1
+
+        nelements = len(element_id[0])
+
+        #1) calculate mean and variance
+        mu  = np.zeros(nelements)
+        var = np.zeros(nelements)
+
+        #2) collect chain
+        samples = []
+
+        nn=1
+        for i,a in explore_chain(self.chain_path, first_sample,last_sample, 1):
+            d   = a[attribute][:]
+
+            print( np.shape(d))
+
+            mu  = (nn-1.)/float(nn)*mu +1./float(nn)*d[element_id]
+            if(nn>1):
+                aux     = (mu-d[element_id])**2
+                var     = (nn-1.)/nn*var+1./(nn-1)*aux
+
+            samples.append(d[element_id])
+
+            nn+=1
+
+
+        cl      = np.zeros((nlength,nelements))
+        cl_count= np.zeros(nlength)
+
+        for i in range(nlength):
+            for j in range(len(samples)-i):
+                cl[i]+= (samples[j]-mu)*(samples[j+i]-mu)/var
+                cl_count[i] +=1.;
+
+        for i in range(nlength):
+            cl[i]/=cl_count[i]
+
+        return np.array(range(nlength)), cl
+
+    def print_job(self,msg):
+        print('-'*60)
+        print(msg)
+        print('-'*60)
+
+
+    def build_power_spectrum_chain(Nbins=256):
+      opts=dict(Nbins=Nbins,range=(0,self.kmodes.max()))
+      #FIXME: Do not use the first element
+      Pref = self.rebin_power_spectrum(self.mcmc_list[0], **opts)
+      try:
+        data = np.load("power_%s.npz" % suffix)
+        loc_names = names[len(data['P']):]
+        PP = list(data['P'])
+      except:
+        PP = []
+        loc_names = list(names)
+      print(loc_names)
+      if len(loc_names) == 0:
+        return
+      
+      for i in pb.ProgressBar()(loc_names):
+        PP.append(ss.compute_power_shat_spectrum(i, **opts))
+      
+      bins = 0.5*(Pref[2][1:]+Pref[2][:-1])
+      
+      np.savez("power_%s.npz" % suffix, bins=bins, P=PP, startMC=startMC, Fmax=Fmax, Pref=Pref)
+
+
+    def spectrum_pdf(self, first_sample=0, last_sample=-1, sample_steps=10, gridsize=1000, Pmin=None, Pmax=None):
+
+        P = np.zeros((gridsize, Npk), dtype=np.float64)
+
+        if Pmin is None or Pmax is None:
+            P0m,P0M = np.inf,0
+            for i,a in explore_chain(self.chain_path, first_sample,last_sample, sample_steps):
+                P0m = a['/scalars/powerspectrum'].min()
+                P0M = a['/scalars/powerspectrum'].max()
+                Pb_m,Pb_M = min(P0m, Pb_m),max(P0M,Pb_M)
+            if Pmin is None:
+                Pmin = Pb_m
+            if Pmax is None:
+                Pmax = Pb_M
+
+        PP = Pmin*np.exp(np.arange(gridsize)*np.log(Pmax/Pmin))
+        N=0
+        prior=0
+        N_ib = 0.5*(self.Nk+prior)[None,:]
+        for i,a in explore_chain(self.chain_path, first_sample,last_sample, sample_steps):
+            Pk_i = a['/scalars/powerspectrum'][:]
+            N_i = self.Nk
+
+            _special_spectrum_builder(P, PP, tmp_PP, N_ib, N_i, Pk_i)
+
+            N += 1
+
+        P /= N
+
+        return P
+
+    def get_spherical_slice(self,vdata,nside=32, observer=np.array([0,0,0]),rslice = 150.):
+
+        def RenderSphere(VolumeData3D,image,rslice,observer,Larr,Narr):
+            print( "Rendering Sphere...")
+
+            NSIDE=hp.npix2nside(len(image))
+
+            idx=Larr[0]/Narr[0]
+            idy=Larr[1]/Narr[1]
+            idz=Larr[2]/Narr[2]
+
+            for ipix in range(len(image)):
+                #get direction of pixel and calculate unit vectors
+                dx,dy,dz=hp.pix2vec(NSIDE, ipix)
+                d = math.sqrt(dx * dx + dy * dy + dz * dz)
+                dx = dx / d; dy = dy / d; dz = dz / d # ray unit vector
+
+                rayX = observer[0]+rslice*dx; rayY = observer[1]+rslice*dy; rayZ = observer[2]+rslice*dz
+                rayX /= idx; rayY /= idy; rayZ /= idz
+                #find voxel inside box
+                ix = int(round(rayX))
+                iy = int(round(rayY))
+                iz = int(round(rayZ))
+
+                image[ipix]=np.nan
+                if ix > -1 and ix < Narr[0] \
+                        or iy > -1 and iy < Narr[1] \
+                        or iz > -1 and iz < Narr[2]:
+
+                    jx = (ix+1) % Narr[0];
+                    jy = (iy+1) % Narr[1];
+                    jz = (iz+1) % Narr[2];
+                    rx = (rayX - ix);
+                    ry = (rayY - iy);
+                    rz = (rayZ - iz);
+
+                    qx = 1-rx;
+                    qy = 1-ry;
+                    qz = 1-rz;
+                    val = VolumeData3D[ix,iy,iz] * qx * qy * qz +VolumeData3D[ix,iy,jz] * qx * qy * rz +VolumeData3D[ix,jy,iz] * qx * ry * qz +VolumeData3D[ix,jy,jz] * qx * ry * rz +VolumeData3D[jx,iy,iz] * rx * qy * qz +VolumeData3D[jx,iy,jz] * rx * qy * rz +VolumeData3D[jx,jy,iz] * rx * ry * qz +VolumeData3D[jx,jy,jz] * rx * ry * rz;
+                    image[ipix]=val
+
+                print( '\r'+str(100 * ipix / (len(image) - 1)).zfill(3) + "%")
+
+
+        obs = np.array([observer[0]-self.xmin0,observer[1]-self.xmin1,observer[2]-self.xmin2])
+
+        Larr=np.array([self.L0,self.L1,self.L2])
+        Narr=np.array([self.N0,self.N1,self.N2])
+        image = np.zeros(hp.nside2npix(nside))
+        RenderSphere(vdata,image,rslice,obs,Larr,Narr)
+        return image
+
+    def mean_var_density(self, first_sample=0,last_sample=-1,sample_steps=10):
+        self.print_job('Estimate mean and variance of density fields for %s run!' % self.LSS_framework)
+
+
+        if(self.LSS_framework=='ARES'):
+            mu_i    = np.zeros((self.N0,self.N1,self.N2))
+
+            var_i     = np.zeros((self.N0,self.N1,self.N2))
+
+            nn=1
+            for i,a in explore_chain(self.chain_path, first_sample,last_sample, sample_steps):
+                d = a['s_field'][:]
+                mu_i    = (nn-1.)/float(nn)*mu_i +1./float(nn)*d
+
+                if(nn>1):
+                    aux = (mu_i-d)**2
+                    var_i   = (nn-1.)/nn*var_i+1./(nn-1)*aux
+
+                nn+=1
+
+            return mu_i, var_i
+
+        elif(self.LSS_framework=='HADES'):
+            mu_i    = np.zeros((self.N0,self.N1,self.N2))
+
+            var_i     = np.zeros((self.N0,self.N1,self.N2))
+
+            nn=1
+            for i,a in explore_chain(self.chain_path, first_sample,last_sample, sample_steps):
+                d = a['s_field'][:]
+                mu_i    = (nn-1.)/float(nn)*mu_i +1./float(nn)*d
+
+                if(nn>1):
+                    aux = (mu_i-d)**2
+                    var_i   = (nn-1.)/nn*var_i+1./(nn-1)*aux
+
+                nn+=1
+
+            return mu_i, var_i
+
+        else:
+            mu_i    = np.zeros((self.N0,self.N1,self.N2))
+            mu_f    = np.zeros((self.N0,self.N1,self.N2))
+
+            var_i     = np.zeros((self.N0,self.N1,self.N2))
+            var_f     = np.zeros((self.N0,self.N1,self.N2))
+
+            nn=1
+            for i,a in explore_chain(self.chain_path, first_sample,last_sample, sample_steps):
+                d = a['s_field'][:]
+                mu_i    = (nn-1.)/float(nn)*mu_i +1./float(nn)*d
+
+                if(nn>1):
+                    aux = (mu_i-d)**2
+                    var_i   = (nn-1.)/nn*var_i+1./(nn-1)*aux
+
+                d = a['BORG_final_density'][:]
+                mu_f    = (nn-1.)/float(nn)*mu_f +1./float(nn)*d
+
+                if(nn>1):
+                    aux = (mu_f-d)**2
+                    var_f   = (nn-1.)/nn*var_f+1./(nn-1)*aux
+
+                nn+=1
+
+
+            return mu_i, var_i, mu_f, var_f
+
+    def mean_var_spec(self, first_sample=0,last_sample=-1,sample_steps=10):
+        self.print_job('Estimate mean and variance of density fields for %s run!' % self.LSS_framework)
+
+        if(self.LSS_framework=='ARES'):
+            mu    = np.zeros(self.nmodes)
+            var   = np.zeros(self.nmodes)
+
+            nn=1
+            for i,a in explore_chain(self.chain_path, first_sample,last_sample, sample_steps):
+                d = a['/scalars/powerspectrum'][:]
+                mu    = (nn-1.)/float(nn)*mu +1./float(nn)*d
+
+                if(nn>1):
+                    aux = (mu-d)**2
+                    var   = (nn-1.)/nn*var+1./(nn-1)*aux
+
+                nn+=1
+
+            return self.kmodes,mu, var
+
+        elif(self.LSS_framework=='HADES'):
+            mu    = np.zeros(self.nmodes)
+            var   = np.zeros(self.nmodes)
+
+            nn=1
+            for i,a in explore_chain(self.chain_path, first_sample,last_sample, sample_steps):
+                d = a['/scalars/powerspectrum'][:]
+                mu    = (nn-1.)/float(nn)*mu +1./float(nn)*d
+
+                if(nn>1):
+                    aux = (mu-d)**2
+                    var   = (nn-1.)/nn*var+1./(nn-1)*aux
+
+                nn+=1
+
+            return self.kmodes,mu, var
+
+        else:
+            mu    = np.zeros(self.nmodes)
+            var   = np.zeros(self.nmodes)
+
+            nn=1
+            for i,a in explore_chain(self.chain_path, first_sample,last_sample, sample_steps):
+                d = a['/scalars/powerspectrum'][:]
+                mu    = (nn-1.)/float(nn)*mu +1./float(nn)*d
+
+                if(nn>1):
+                    aux = (mu-d)**2
+                    var   = (nn-1.)/nn*var+1./(nn-1)*aux
+
+                nn+=1
+
+            return self.kmodes,mu, var

scripts/ares_tools/read_all_h5.py

@@ -0,0 +1,495 @@
+#+
+#   ARES/HADES/BORG Package -- ./scripts/ares_tools/read_all_h5.py
+#   Copyright (C) 2014-2020 Guilhem Lavaux <guilhem.lavaux@iap.fr>
+#   Copyright (C) 2009-2020 Jens Jasche <jens.jasche@fysik.su.se>
+#
+#   Additional contributions from:
+#      Guilhem Lavaux <guilhem.lavaux@iap.fr> (2023)
+#   
+#+
+import os
+import numpy as np
+import numexpr as ne
+import h5py as h5
+
+def isnotebook():
+    try:
+        shell = get_ipython().__class__.__name__
+        if shell == 'ZMQInteractiveShell':
+            return True   # Jupyter notebook or qtconsole
+        elif shell == 'TerminalInteractiveShell':
+            return False  # Terminal running IPython
+        else:
+            return False  # Other type (?)
+    except NameError:
+        return False      # Probably standard Python interpreter
+
+
+try:
+    import tqdm
+
+    if isnotebook():
+      u_tqdm = tqdm.tqdm_notebook
+    else:
+      u_tqdm = tqdm.tqdm
+
+    def progress(iterator):
+        L = list(iterator)
+        return u_tqdm(L)
+except:
+    def progress(iterator):
+        for i, q in enumerate(iterator):
+            if ((i % 100) == 0):
+                print(i)
+            yield q
+
+
+def default_slicer(x): return x[...]
+
+
+class Bunch(object):
+    def __init__(self, **kwds):
+        self.__dict__.update(**kwds)
+
+    def __del__(self):
+        if hasattr(self, '_lazy') and self._lazy:
+            # close the file/group
+            if isinstance(self._group, h5.File):
+                print("Closing HDF5 file")
+                try:
+                    self._group.close()
+                except:
+                    # Eat all exceptions
+                    pass
+
+    def insert(self, aname, data):
+        if len(aname) == 1:
+            self.__dict__.update({aname[0]: data})
+        else:
+            if aname[0] in self.__dict__:
+                k = self.__dict__[aname[0]]
+            else:
+                k = Bunch()
+                self.__dict__.update({aname[0]: k})
+            k.insert(aname[1:], data)
+
+
+def read_group(g, lazy=False):
+    m = {}
+    for k in g:
+        if hasattr(g[k], 'keys'):
+            m[k] = read_group(g[k], lazy=lazy)
+        else:
+            if lazy:
+                m[k] = g[k]
+            else:
+                m[k] = g[k][:]
+    if lazy:
+        m['_group'] = g
+    m['_lazy'] = lazy
+
+    return Bunch(**m)
+
+
+def read_attr_group(g, ename, bunch, slicer=default_slicer):
+    for a in ename:
+        g = g[a]
+    if bunch == None:
+        bunch = Bunch()
+
+    bunch.insert(ename, slicer(g))
+    return bunch
+
+
+def read_all_h5(fname, lazy=False):
+    if lazy:
+        f = h5.File(fname, mode="r")
+        return read_group(f, lazy=True)
+    else:
+        with h5.File(fname, mode="r") as f:
+            return read_group(f, lazy=False)
+
+
+def read_attr_h5(fname, egroup, slicer=default_slicer):
+    with h5.File(fname, mode="r") as f:
+        b = None
+        for e in egroup:
+            b = read_attr_group(f, e, b, slicer=slicer)
+    return b
+
+
+def grabber(obj, e):
+    if len(e) == 1:
+        return getattr(obj, e[0])
+    else:
+        return grabber(getattr(obj, e[0]), e[1:])
+
+
+def rebuild_spliced_h5(path, element_list, ncpu, verbose=False, flex_cpu=False):
+    """Rebuild a set of fields which are spliced across different hdf5 files.
+    The use of this function is typically dedicated to the analysis of MPI run.
+
+    Parameters
+    ----------
+    path         : string
+                   base path for the set of hdf5 file. A suffix "_[CPUID]" will be appended.
+    element_list : list of string
+                   list of elements to rebuild from the set of files
+    ncpu         : int
+                   number of cpus for the run
+    verbose      : boolean
+                   if True the code will run with verbose output
+    flex_cpu     : boolean
+                   if True, ncpu is understood as the maximum number of cpus. If the code does
+                   find a file then it stops the rebuilding without failing.
+
+    Returns
+    -------
+       dictionnary
+       each element name is a key, and the value is a numpy array. The arrays are concatenated
+       according to their first axis.
+    """
+    b = [None for e in element_list]
+    egrab = [e.split('.') for e in element_list]
+
+    for cpu in range(ncpu):
+        fname = path + "_%d" % cpu
+        if verbose:
+            print("Loading CPU file '%s'" % fname)
+        try:
+            a = read_attr_h5(fname, egrab)
+        except OSError:
+            if flex_cpu:
+                break
+            raise
+        for j, e in enumerate(egrab):
+            if b[j] is None:
+                b[j] = grabber(a, e)
+            else:
+                b[j] = np.append(b[j], grabber(a, e), axis=0)
+
+    dtype = [(e, t[0].dtype, t[0].shape) for e, t in zip(element_list, b)]
+    arr = {}
+    for e, q in zip(element_list, b):
+        arr[e] = q
+    return arr
+
+
+def chain_iterator(path, start=0, step=1, err_up=0, end=-1, prefix="mcmc", need_id=False):
+    import os
+
+    i = start
+    while True:
+        fname = os.path.join(path, "%s_%d.h5" % (prefix,i,))
+        try:
+            os.stat(fname)
+        except IOError:
+            if (i >= err_up):
+                return
+            else:
+                i += step
+                continue
+
+        if need_id:
+          yield (i,fname)
+        else:
+          yield fname
+        i += step
+        if end > 0 and i > end:
+          break
+
+
+def read_chain_h5(path, element_list, start=0, step=1, err_up=0, slicer=default_slicer, prefix="mcmc", flexible=True):
+    """
+    read_chain_h5(path,element_list,start=0,step=1,err_up=0,slicer=default_slicer)
+
+    Arguments:
+      * path: path where you have the chain (mcmc_* files)
+      * element_list: list of strings where the MCMC objects are stored. For example, you
+                      have scalars.galaxy_nmean_0 for the nbar parameter of the first catalog.
+      * start: the first element of the chain to consider
+      * step: if you want to thin the chain by "step"
+      * err_up: whether to accept I/O errors when opening files up to the specified MCMC id
+      * slicer: a lambda function that can only take a subset of the array of the
+      specified MCMC object. For example, it can be lambda d: d[:,:,64], to indicate only
+       the plane 64 in the 3d grid that is being loaded. [64,...]
+
+    Returns:
+      * a columned numpy array "a". You have one column for each element_id of the
+      element_list.  You can access one of the column like this:
+      a["scalars.galaxy_nmean_0"]
+
+  """
+    i = start
+    b = [[] for e in element_list]
+    egrab = [e.split('.') for e in element_list]
+    for fname in progress(chain_iterator(path, start=start, step=step, err_up=err_up, prefix=prefix)):
+        try:
+          a = read_attr_h5(fname, egrab, slicer=slicer)
+        except OSError:
+          if not flexible:
+              raise
+          else:
+              break
+        for j, e in enumerate(egrab):
+            b[j].append(grabber(a, e))
+
+    dtype = [(e, t[0].dtype, t[0].shape) for e, t in zip(element_list, b)]
+
+    arr = np.empty(len(b[0]), dtype=dtype)
+    for e, q in zip(element_list, b):
+        arr[e] = q
+    return arr
+
+
+def chain_compute_xcor(path, start=0, Nbins=100):
+    import cosmotool as ct
+    import numexpr as ne
+
+    i = 0
+
+    with h5.File(os.path.join(path, "restart.h5_0"), mode="r") as f:
+        L0 = f['scalars']['L0'][:]
+        L1 = f['scalars']['L1'][:]
+        L2 = f['scalars']['L2'][:]
+        N0 = int(f['scalars']['N0'][:])
+        N1 = int(f['scalars']['N1'][:])
+        N2 = int(f['scalars']['N2'][:])
+
+    ix = np.fft.fftfreq(N0, d=1.0/L0)[:, None, None]
+    iy = np.fft.fftfreq(N1, d=1.0/L1)[None, :, None]
+    iz = np.fft.fftfreq(N2, d=1.0/L2)[None, None, :]
+    r2 = ne.evaluate('sqrt(ix**2+iy**2+iz**2)')
+    rmax = r2.max()
+    ir = (r2 * Nbins / rmax).astype(np.int32).ravel()
+    xi = []
+    W = np.bincount(ir, minlength=Nbins)
+
+    fft = ct.CubeFT(L0, N0)
+
+    while True:
+        try:
+            if i % 10 == 0:
+                print(i)
+            fname = os.path.join(path, "mcmc_%d.h5" % (i+start))
+            with h5.File(fname, mode="r") as f:
+                fft.density = f["scalars"]["s_field"][:]
+
+            fft.rfft()
+            ne.evaluate("complex(real(d)**2 + imag(d)**2, 0)",
+                        local_dict={'d': fft.dhat}, out=fft.dhat, casting='unsafe')
+            fft.irfft()
+
+            xi.append(np.bincount(
+                ir, weights=fft.density.ravel(), minlength=Nbins))
+            i += 1
+
+        except Exception as e:
+            print(repr(e))
+            break
+
+    xi = np.array(xi) / W
+
+    r = np.arange(Nbins) * rmax / Nbins
+    return r, xi
+
+
+def explore_chain(path, start, end=-1, step=1, quiet=True):
+    """
+    Arguments:
+      * path
+      * start
+      * end
+      * step
+
+    Returns:
+      * iterator with hdf5 object. Example:
+           for i in explore_chain(".", 0):
+              mean = i['/scalars/galaxy_nmean_0'][0]
+              # Then do stuff with "mean"
+  """
+    n = int(start)
+    nmax = int(end)
+    step = int(step)
+    k = 0
+
+    while (nmax == -1) or (n < nmax):
+        p = path + "/mcmc_%d.h5" % n
+        if not quiet and (k % 100) == 0:
+            print("%d" % n)
+        try:
+            f = h5.File(p, mode="r")
+        except Exception as e:
+            print(e)
+            break
+
+        try:
+            yield n, f['scalars']
+        finally:
+            f.close()
+
+        n += step
+        k += 1
+
+
+def build_power_histogram(path, start=0, step=1, Nhisto=100, quiet=True, logP=True, Prange=(0.1, 1e5)):
+    """
+    Use the scalars.powerspectrum mcmc element to build the PDF of the posterior
+    powerspectrum
+
+    Arguments:
+       * path
+       * start
+       * step
+       * Nhisto: number of bins for each k mode of the P(k) histogram
+       * quiet:
+       * logP: whether you want to use a log scale for plotting
+       * Prange: a tuple for giving the entire P range to represent
+
+     Returns:
+       * a tuple: t=(kmodes, Pgrid, Pk_pdf_values) which is directly usable in pcolormesh
+         like pcolormesh(*t)
+  """
+    # print path+ "/restart.h5_0"
+    for _, scalars in explore_chain(path, start, end=start+1, step=1, quiet=quiet):
+        Nk = scalars["powerspectrum"].size
+    with h5.File(path + "/restart.h5_0", mode="r") as f:
+        k_mode = f["/scalars/k_modes"][:]
+
+    Phisto = np.zeros((Nk, Nhisto), dtype=np.int)
+    if logP:
+        logPmin = np.log10(Prange[0])
+        DeltaLog = np.log10(Prange[1]/Prange[0])/(Nhisto)
+
+        def transform(P): return ((np.log10(P)-logPmin)/DeltaLog)
+    else:
+        Pmin = Prange[0]
+        DeltaP = (Prange[1]-Prange[0])/(Nhisto)
+
+        def transform(P): return (P-Pmin)/DeltaP
+
+    for n, scalars in explore_chain(path, start, step=step, quiet=quiet):
+        P = scalars["powerspectrum"][:]
+        iP = np.floor(transform(P))
+        ik = np.where((np.isnan(iP) == False)*(iP >= 0)*(iP < Nhisto))
+        iP = iP[ik].astype(np.int)
+        for i, j in zip(ik[0], iP):
+            Phisto[i, j] = Phisto[i, j]+1
+
+    k_mode = k_mode[:, None].repeat(Nhisto, axis=1)
+    if logP:
+        Pg = 10**((np.arange(Nhisto)+0.5)*DeltaLog + logPmin)
+    else:
+        Pg = ((np.arange(Nhisto)+0.5)*DeltaP + Pmin)
+
+    Pg = Pg[None, :].repeat(Nk, axis=0)
+    return k_mode, Pg, Phisto
+
+
+def read_chain_complex_avg_dev(path, op, start=0, end=-1, do_dev=False, step=1, slicer=default_slicer, prefix="mcmc"):
+    """
+    Compute mean and standard deviation of the given element_list
+
+    Arguments:
+      * path
+      * op:
+      * element_list
+      * start
+      * do_dev: boolean for computing the standard deviation (or not)
+      * slicer:
+
+
+     Returns:
+      * a columned numpy array. Each column has a name that corresponds to an element with an additional dimension. For example, a['scalars.galaxy_nmean_0'][0] -> mean,
+      a['scalars.galaxy_nmean_0'][1] is the standard deviation.
+  """
+    i = 0
+    b = None
+    bd = None
+    try:
+      for fname in progress(chain_iterator(path, start=start, step=step, end=end, prefix=prefix)):
+        with h5.File(fname, mode="r") as ff:
+            if b is None:
+                b = op(ff)
+                if do_dev:
+                    bd = np.zeros(b.shape)
+            else:
+                data = op(ff)
+                ne.evaluate('r*a+k*c',
+                            dict(k=1/float(i+1),
+                                 r=(float(i)/float(i+1)),
+                                 a=b,
+                                 c=data),
+                            out=b)
+            if do_dev and i > 1:
+                ne.evaluate('k*(xn-mun)**2 + f*bdn',
+                            dict(k=1/float(i),
+                                 f=float(i)/float(i+1),
+                                 xn=data,
+                                 mun=b,
+                                 bdn=bd),
+                            out=bd)
+            i += 1
+
+    except OSError:
+      pass
+
+    bd = np.sqrt(bd)
+    return b, bd
+
+
+def read_chain_avg_dev(path, element_list, start=0, end=-1, do_dev=False, operator=lambda x: x, step=1, slicer=default_slicer, prefix="mcmc", err_up=0):
+    """
+    Compute mean and standard deviation of the given element_list
+
+    Arguments:
+      * path
+      * element_list
+      * start
+      * do_dev: boolean for computing the standard deviation (or not)
+      * operator: applies the operator on all the elements before computing the mean and
+        standard deviation.
+      * slicer:
+
+
+     Returns:
+      * a columned numpy array. Each column has a name that corresponds to an element with an additional dimension. For example, a['scalars.galaxy_nmean_0'][0] -> mean,
+      a['scalars.galaxy_nmean_0'][1] is the standard deviation.
+  """
+    i = 0
+    b = [None for e in element_list]
+    bd = [None for e in element_list]
+    egrab = [e.split('.') for e in element_list]
+    for fname in progress(chain_iterator(path, start=start, step=step, end=end, err_up=err_up, prefix=prefix)):
+            a = read_attr_h5(fname, egrab, slicer=slicer)
+            for j, e in enumerate(egrab):
+                if b[j] is None:
+                    b[j] = operator(grabber(a, e))
+                    if do_dev:
+                        bd[j] = np.zeros(b[j].shape)
+                else:
+                    data = operator(grabber(a, e))
+                    ne.evaluate('r*a+k*c',
+                                dict(k=1/float(i+1),
+                                     r=(float(i)/float(i+1)),
+                                     a=b[j],
+                                     c=data),
+                                out=b[j])
+                    if do_dev and i > 0:
+                        ne.evaluate('k*(xn-mun)**2 + f*bdn',
+                                    dict(k=1/float(i),
+                                         f=float(i)/float(i+1),
+                                         xn=data,
+                                         mun=b[j],
+                                         bdn=bd[j]),
+                                    out=bd[j])
+                i+=1
+
+    dtype = [(e, t.dtype, t.shape) for e, t in zip(element_list, b)]
+
+    arr = np.empty(2 if do_dev else 1, dtype=dtype)
+    for e, q, q2 in zip(element_list, b, bd):
+        arr[e][0] = q
+        if do_dev:
+            arr[e][1] = np.sqrt(q2)
+    return arr

scripts/ares_tools/visu/__init__.py

@@ -0,0 +1,9 @@
+#+
+#   ARES/HADES/BORG Package -- ./scripts/ares_tools/visu/__init__.py
+#   Copyright (C) 2014-2020 Guilhem Lavaux <guilhem.lavaux@iap.fr>
+#   Copyright (C) 2009-2020 Jens Jasche <jens.jasche@fysik.su.se>
+#
+#   Additional contributions from:
+#      Guilhem Lavaux <guilhem.lavaux@iap.fr> (2023)
+#   
+#+

scripts/ares_tools/visu/vtktools.py

@@ -0,0 +1,109 @@
+#+
+#   ARES/HADES/BORG Package -- ./scripts/ares_tools/visu/vtktools.py
+#   Copyright (C) 2014-2020 Guilhem Lavaux <guilhem.lavaux@iap.fr>
+#   Copyright (C) 2009-2020 Jens Jasche <jens.jasche@fysik.su.se>
+#
+#   Additional contributions from:
+#      Guilhem Lavaux <guilhem.lavaux@iap.fr> (2023)
+#   
+#+
+import numpy as np
+from vtk.util.numpy_support import get_vtk_array_type, create_vtk_array, get_numpy_array_type
+
+def numpy_scalar_to_vtk(np_array):
+    """This function converts a numpy scalar array to a VTK array.
+    
+    Args:
+      np_array (np.array): A numpy array
+      
+    Returns:
+      vtk.vtkArray: An array of the closest possible type of numpy array. The array is deep
+                    copied to avoid SEGV.
+    
+    """
+    atype = get_vtk_array_type(np_array.dtype)    
+    array = create_vtk_array(atype)
+    array.SetNumberOfComponents(1)
+    ntype = get_numpy_array_type(atype)
+    adata = np.ravel(np_array).astype(ntype)
+    array.SetVoidArray(adata, len(adata), 1)
+    
+    copy = array.NewInstance()
+    copy.DeepCopy(array)
+    return copy
+
+def numpy_vector_to_vtk(np_array):
+    """This function converts a numpy scalar array to a VTK array.
+    
+    Args:
+      np_array (np.array): A numpy array
+      
+    Returns:
+      vtk.vtkArray: An array of the closest possible type of numpy array. The array is deep
+                    copied to avoid SEGV.
+    
+    """
+    if np_array.shape[3] != 3:
+      raise ValueError()
+      
+    atype = get_vtk_array_type(np_array.dtype)    
+    array = create_vtk_array(atype)
+    array.SetNumberOfComponents(3)
+    ntype = get_numpy_array_type(atype)
+    adata = np.ravel(np_array).astype(ntype)
+    array.SetVoidArray(adata, len(adata), 1)
+    
+    copy = array.NewInstance()
+    copy.DeepCopy(array)
+    return copy
+
+
+def smooth_array(a, L=[1.0,1.0,1.0], R=0.1):
+    a_hat = np.fft.rfftn(a)
+    ik = [np.fft.fftfreq(iN, d=iL/iN)*2*np.pi for iN,iL in zip(a.shape,L)]
+    k2 = ik[0][:,None,None]**2 + ik[1][None,:,None]**2 + ik[2][None,None,:a.shape[2]/2+1]**2
+    
+    a_hat *= np.exp(-0.5*k2*R**2)
+    return np.fft.irfftn(a_hat)
+
+def displacement_array(a, L=[1.0,1.0,1.0], R=0.1):
+    a_hat = np.fft.rfftn(a)
+    ik = [np.fft.fftfreq(iN, d=iL/iN)*2*np.pi for iN,iL in zip(a.shape,L)]
+    k2 = ik[0][:,None,None]**2 + ik[1][None,:,None]**2 + ik[2][None,None,:a.shape[2]/2+1]**2
+    
+    b = np.empty(a.shape + (3,), dtype=np.float32)
+    
+    b_hat = -a_hat * 1j*ik[0][:,None,None]/k2
+    b_hat[0,0,0]=0
+    b[...,0] = np.fft.irfftn(b_hat)
+
+    b_hat = -a_hat * 1j*ik[1][None,:,None]/k2
+    b_hat[0,0,0]=0
+    b[...,1] = np.fft.irfftn(b_hat)
+
+    b_hat = -a_hat * 1j*ik[2][None,None,:a.shape[2]/2+1]/k2
+    b_hat[0,0,0]=0
+    b[...,2] = np.fft.irfftn(b_hat)
+    
+    return b
+
+
+def setupImageData3D(img_data, np_array, dims=[1.0,1.0,1.0], name="numpy array"):
+    """This function setups a 3D image data object.
+    
+    """
+    shape = np_array.shape
+    
+    dx = (d0/(N0-1) for d0,N0 in zip(dims,shape))
+    img_data.SetOrigin(*(-d0/2 for d0 in dims)) # default values
+    img_data.SetSpacing(*dx)
+    img_data.SetDimensions(*shape) # number of points in each direction
+    
+    array = numpy_scalar_to_vtk(np_array)
+    
+    img_data.GetPointData().AddArray(array)
+    array.SetName(name)
+    
+    return array
+    
+