Merge branch 'master' of bitbucket.org:glavaux/cosmotool

python_sample/build_2lpt_skymap.py

@@ -1,3 +1,5 @@
+import matplotlib
+matplotlib.use('Agg')
 import healpy as hp
 import numpy as np
 import cosmotool as ct
@@ -21,7 +23,8 @@ parser.add_argument('--iid', type=int, default=0)
 parser.add_argument('--proj_cat', type=bool, default=False)
 args = parser.parse_args()
 
-INDATA="/nethome/lavaux/Copy/PlusSimulation/BORG/Input_Data/2m++.npy"
+#INDATA="/nethome/lavaux/Copy/PlusSimulation/BORG/Input_Data/2m++.npy"
+INDATA="2m++.npy"
 tmpp = np.load(INDATA)
 
 L = args.boxsize

python_sample/build_nbody_ksz_from_galaxies.py

@@ -3,6 +3,8 @@ import numpy as np
 import cosmotool as ct
 import argparse
 import h5py as h5
+import matplotlib
+matplotlib.use('Agg')
 from matplotlib import pyplot as plt
 import ksz
 from ksz.constants import *
@@ -54,48 +56,60 @@ def build_unit_vectors(N):
   return ux,uy,uz
 
 def build_radial_v(v):
+  N = v[0].shape[0]
   u = build_unit_vectors(N)
-  vr = v[0] * u[0]
+  vr = v[0] * u[2]
   vr += v[1] * u[1]
-  vr += v[2] * u[2]
+  vr += v[2] * u[0]
-  return vr
 
-def generate_from_catalog(vfield,Boxsize):
+  return vr.transpose()
+
+def generate_from_catalog(vfield,Boxsize,dmin,dmax):
   import progressbar as pbar
   
   cat = np.load("2m++.npy")
 
-  profiler = KSZ_Isothermal(2.37)
 
   cat['distance'] = cat['best_velcmb']
   cat = cat[np.where((cat['distance']>100*dmin)*(cat['distance']<dmax*100))]
 
   deg2rad = np.pi/180
+  Npix = 12*Nside**2
   xp,yp,zp = hp.pix2vec(Nside, np.arange(Npix))
   N2 = np.sqrt(xp**2+yp**2+zp**2)
 
-  ksz_template = np.zeros(12*Nside**2, dtype=np.float64)
+  ksz_template = np.zeros(Npix, dtype=np.float64)
-  ksz_mask = np.zeros(12**Nside**2, dtype=np.uint8)
+  ksz_mask = np.zeros(Npix, dtype=np.uint8)
 
-  pb = pbar.ProgressBar(maxval = cat.size, widgets=[pb.Bar(), pb.ETA()]).start()
+  pb = pbar.ProgressBar(maxval = cat.size, widgets=[pbar.Bar(), pbar.ETA()]).start()
 
   for k,i in np.ndenumerate(cat):
-    pb.update(k)
+    pb.update(k[0])
     l,b=i['gal_long'],i['gal_lat']
+    ra,dec=i['ra'],i['dec']
 
     l *= deg2rad
     b *= deg2rad
 
-    x0 = np.cos(l)*np.cos(b)
+    ra *= deg2rad
-    y0 = np.sin(l)*np.cos(b)
+    dec *= deg2rad
-    z0 = np.sin(b)
+
+#    x0 = np.cos(l)*np.cos(b)
+#    y0 = np.sin(l)*np.cos(b)
+#    z0 = np.sin(b)
+
+    x0 =  xra = np.cos(ra)*np.cos(dec)
+    y0 = yra = np.sin(ra)*np.cos(dec)
+    z0 = zra = np.sin(dec)
     
     DA =i['distance']/100
     Lgal = DA**2*10**(0.4*(tmpp_cat['Msun']-i['K2MRS']+25))
 
-    idx0 = hp.query_disc(Nside, (x0,y0,z0), 3*rGalaxy/DA)
+    profiler = ksz.KSZ_Isothermal(Lgal, 2.37)
 
-    vr = interp3d(DA * x0, DA * y0, DA * z0, vfield, Boxsize)
+    idx0 = hp.query_disc(Nside, (x0,y0,z0), 3*profiler.rGalaxy/DA)
+
+    vr = interp3d(DA * xra, DA * yra, DA * zra, vfield, Boxsize)
 
     xp1 = xp[idx0]
     yp1 = yp[idx0]
@@ -119,18 +133,33 @@ for i in xrange(args.start,args.end,args.step):
   ff=plt.figure(1)
   plt.clf()
   v=[]
-  with h5.File(args.base_h5 % i, mode="r") as f:
+  fname = args.base_h5 % i
-    p = wrapper_impulsion(f)
+  if False:
-    v.append(p[i] / f['density'][:])
+    print("Opening %s..." % fname)
+    with h5.File(fname, mode="r") as f:
+      p = wrapper_impulsion(f)
+      for j in xrange(3):
+        v.append(p[j] / f['density'][:])
+
+    print("Building radial velocities...")
 
   # Rescale by Omega_b / Omega_m
-  vr = build_radial_v(v)
+    vr = build_radial_v(v)
-  vr *= -ksz_normalization*rho_mean_matter*1e6
+#  _,_,vr = build_unit_vectors(128)
-  del v
+#  vr *= 1000 * 500
+    vr *= ksz_normalization*1e6
+    del v
+#    np.save("vr.npy", vr)
+  else:
+    print("Loading vrs...")
+    vr = np.load("vr.npy")
 
-  proj = generate_from_catalog(vfield,Boxsize)
+  print("Generating map...")
+
+  proj,mask = generate_from_catalog(vr,args.boxsize,args.depth_min,args.depth_max)
 
   hp.write_map(args.ksz_map % i, proj)
+  hp.write_map((args.ksz_map % i) + "_mask", mask)
 
   hp.mollview(proj, fig=1, coord='CG', cmap=plt.cm.coolwarm, title='Sample %d' % i, min=args.minval, 
               max=args.maxval)

python_sample/icgen/borgicgen.py

@@ -52,8 +52,8 @@ def compute_ref_power(L, N, cosmo, bins=10, range=(0,1), func='HU_WIGGLES'):
   p.normalize(cosmo['SIGMA8'])
 
   return bin_power(p.compute(k)*cosmo['h']**3, L, bins=bins, range=range)
-                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                  
+
-def run_generation(input_borg, a_borg, a_ic, cosmo, supersample=1, do_lpt2=True, shiftPixel=False, needvel=True):
+def run_generation(input_borg, a_borg, a_ic, cosmo, supersample=1, do_lpt2=True, shiftPixel=False, psi_instead=False, needvel=True):
     """ Generate particles and velocities from a BORG snapshot. Returns a tuple of 
     (positions,velocities,N,BoxSize,scale_factor)."""
 
@@ -74,7 +74,7 @@ def run_generation(input_borg, a_borg, a_ic, cosmo, supersample=1, do_lpt2=True,
     # Compute LPT scaling coefficient
     D1 = cgrowth.D(a_ic)
     D1_0 = D1/cgrowth.D(a_borg)
-    velmul = cgrowth.compute_velmul(a_ic)
+    velmul = cgrowth.compute_velmul(a_ic) if not psi_instead else 1
     
     D2 = -3./7 * D1_0**2
 
@@ -96,7 +96,7 @@ def run_generation(input_borg, a_borg, a_ic, cosmo, supersample=1, do_lpt2=True,
     print("velmul=%lg" % (cosmo['h']*velmul))
 
     lpt.cube.dhat = lpt.dhat
-    density = lpt.cube.irfft()    
+    density = lpt.cube.irfft()
     density *= (cgrowth.D(1)/cgrowth.D(a_borg))
 
     return posx,vel,density,N*supersample,L,a_ic,cosmo
@@ -119,7 +119,7 @@ def whitify(density, L, cosmo, supergenerate=1, func='HU_WIGGLES'):
   Pk = build_Pk()
   Pk[0,0,0]=1
 
-  cube = CubeFT(N, L)
+  cube = CubeFT(L, N)
   cube.density = density
   density_hat = cube.rfft()
   density_hat /= np.sqrt(Pk)
@@ -170,15 +170,15 @@ def write_icfiles(*generated_ic, **kwargs):
     supergenerate=kwargs['supergenerate']
   posx,vel,density,N,L,a_ic,cosmo = generated_ic
 
-  ct.simpleWriteGadget("borg.gad", posx, velocities=vel, boxsize=L, Hubble=cosmo['h'], Omega_M=cosmo['omega_M_0'], time=a_ic) 
+  ct.simpleWriteGadget("Data/borg.gad", posx, velocities=vel, boxsize=L, Hubble=cosmo['h'], Omega_M=cosmo['omega_M_0'], time=a_ic) 
   for i,c in enumerate(["x","y","z"]):
-    ct.writeGrafic("ic_velc%s" % c, vel[i].reshape((N,N,N)), L, a_ic, **cosmo)
+    ct.writeGrafic("Data/ic_velc%s" % c, vel[i].reshape((N,N,N)), L, a_ic, **cosmo)
 
-  ct.writeGrafic("ic_deltab", density, L, a_ic, **cosmo)
+  ct.writeGrafic("Data/ic_deltab", density, L, a_ic, **cosmo)
   
-  ct.writeWhitePhase("white.dat", whitify(density, L, cosmo, supergenerate=supergenerate))
+  ct.writeWhitePhase("Data/white.dat", whitify(density, L, cosmo, supergenerate=supergenerate))
   
-  with file("white_params", mode="w") as f:
+  with file("Data/white_params", mode="w") as f:
     f.write("4\n%lg, %lg, %lg\n" % (cosmo['omega_M_0'], cosmo['omega_lambda_0'], 100*cosmo['h']))
     f.write("%lg\n%lg\n-%lg\n0,0\n" % (cosmo['omega_B_0'],cosmo['ns'],cosmo['SIGMA8']))
     f.write("-%lg\n1\n0\n\n\n\n\n" % L)

python_sample/icgen/cosmogrowth.py

@@ -1,3 +1,4 @@
+import numexpr as ne
 import multiprocessing
 import pyfftw
 import weakref
@@ -13,14 +14,14 @@ class CubeFT(object):
       self.max_cpu = multiprocessing.cpu_count() if max_cpu < 0 else max_cpu
       self._dhat = pyfftw.n_byte_align_empty((self.N,self.N,self.N/2+1), self.align, dtype='complex64')
       self._density = pyfftw.n_byte_align_empty((self.N,self.N,self.N), self.align, dtype='float32')
-      self.irfft = pyfftw.FFTW(self._dhat, self._density, axes=(0,1,2), direction='FFTW_BACKWARD', threads=self.max_cpu, normalize_idft=False)
+      self._irfft = pyfftw.FFTW(self._dhat, self._density, axes=(0,1,2), direction='FFTW_BACKWARD', threads=self.max_cpu, normalize_idft=False)
-      self.rfft = pyfftw.FFTW(self._density, self._dhat, axes=(0,1,2), threads=self.max_cpu, normalize_idft=False)
+      self._rfft = pyfftw.FFTW(self._density, self._dhat, axes=(0,1,2), threads=self.max_cpu, normalize_idft=False)
 
   def rfft(self):
-      return self.rfft()*(self.L/self.N)**3
+      return ne.evaluate('c*a', local_dict={'c':self._rfft(normalise_idft=False),'a':(self.L/self.N)**3})
       
   def irfft(self):
-      return self.irfft()/self.L**3
+      return ne.evaluate('c*a', local_dict={'c':self._irfft(normalise_idft=False),'a':(1/self.L)**3})
    
   def get_dhat(self):
       return self._dhat
@@ -152,16 +153,18 @@ class LagrangianPerturbation(object):
         k2 = self._get_k2()
         k2[0,0,0] = 1
 
+        potgen0 = lambda i: ne.evaluate('kdir**2*d/k2',local_dict={'kdir':self._kdir(i),'d':self.dhat,'k2':k2} )
+        potgen = lambda i,j: ne.evaluate('kdir0*kdir1*d/k2',local_dict={'kdir0':self._kdir(i),'kdir1':self._kdir(j),'d':self.dhat,'k2':k2} )
+
         if 'lpt2_potential' not in self.cache:
             print("Rebuilding potential...")
             div_phi2 = np.zeros((self.N,self.N,self.N), dtype=np.float64)
             for j in xrange(3):
-                q = self._do_irfft( self._kdir(j)**2*self.dhat / k2 ).copy()
+                q = self._do_irfft( potgen0(j) ).copy()
                 for i in xrange(j+1, 3):
-                    div_phi2 += q * self._do_irfft( self._kdir(i)**2*self.dhat / k2 )
+                    div_phi2 += q * self._do_irfft( potgen0(i) )
-                    div_phi2 -= (self._do_irfft(self._kdir(j)*self._kdir(i)*self.dhat / k2 ) )**2
+                    div_phi2 -= self._do_irfft(potgen(i,j))**2
 
-            div_phi2 *= 1/self.L**6
             phi2_hat = -self._do_rfft(div_phi2) / k2
             #self.cache['lpt2_potential'] = phi2_hat
             del div_phi2

python_sample/icgen/gen_ic_from_borg.py

@@ -9,10 +9,10 @@ cosmo['omega_B_0']=0.049
 cosmo['SIGMA8']=0.8344
 cosmo['ns']=0.9624
 
-supergen=8
+supergen=2
 zstart=50
 astart=1/(1.+zstart)
 halfPixelShift=False
 
 if __name__=="__main__":
-    bic.write_icfiles(*bic.run_generation("initial_condition_borg.dat", 0.001, astart, cosmo, supersample=1, shiftPixel=halfPixelShift, do_lpt2=False), supergenerate=supergen)
+    bic.write_icfiles(*bic.run_generation("initial_condition_borg.dat", 0.001, astart, cosmo, supersample=2, shiftPixel=halfPixelShift, do_lpt2=False), supergenerate=supergen)

python_sample/icgen/test_ic_from_borg.py

@@ -12,8 +12,8 @@ cosmo['SIGMA8']=0.8344
 cosmo['ns']=0.9624
 N0=256
 
-doSimulation=True
+doSimulation=False
-simShift=True
+simShift=False
 
 snap_id=int(sys.argv[1])
 astart=1/100.
@@ -37,7 +37,7 @@ if doSimulation:
   dsim_hat = np.fft.rfftn(dsim)*(L/N0)**3
   Psim, bsim = bic.bin_power(np.abs(dsim_hat)**2/L**3, L, range=(0,1.), bins=150)
 
-pos,_,density,N,L,_,_ = bic.run_generation("initial_density_2588.dat", 0.001, astart, cosmo, supersample=2, do_lpt2=True)
+pos,_,density,N,L,_,_ = bic.run_generation("initial_density_988.dat", 0.001, astart, cosmo, supersample=1, do_lpt2=True)
 
 dcic = ct.cicParticles(pos, L, N0)
 dcic /= np.average(np.average(np.average(dcic, axis=0), axis=0), axis=0)

python_sample/ksz/constants.py

@@ -18,11 +18,15 @@ frac_electron = 1.0 # Hmmmm
 frac_gas_galaxy = 0.14
 mu = 1/(1-0.5*Y)
 
-tmpp_cat={'Msun':3.29,'alpha':-0.7286211634758224,'Mstar':-23.172904033796893,'PhiStar':0.0113246633636846,'lbar':393109973.22508669}
+tmpp_cat={'Msun':3.29,
+    'alpha':-0.7286211634758224,
+    'Mstar':-23.172904033796893,
+    'PhiStar':0.0113246633636846,
+    'lbar':393109973.22508669}
 
 baryon_fraction = Omega_baryon / Omega_matter
 
-ksz_normalization = T_cmb*sigmaT*v_unit/(lightspeed*mu*mp) * baryon_fraction
+ksz_normalization = -T_cmb*sigmaT*v_unit/(lightspeed*mu*mp) * baryon_fraction
 rho_mean_matter = Omega_matter * (3*(100e3/Mpc)**2/(8*np.pi*G)) 
 Lbar = tmpp_cat['lbar'] / Mpc**3
 M_over_L_galaxy = rho_mean_matter / Lbar

python_sample/ksz/gal_prof.py

@@ -1,3 +1,4 @@
+import numpy as np
 from .constants import *
 
 # -----------------------------------------------------------------------------
@@ -11,27 +12,28 @@ class KSZ_Profile(object):
   def __init__(self):
     self.rGalaxy = 1.0
 
-  def evaluate_profile(r):
+  def evaluate_profile(self, r):
     raise NotImplementedError("Abstract function")
 
-  def projected_profile(cos_theta,angularDistance):
+  def projected_profile(self, cos_theta,angularDistance):
 
     idx = np.where(cos_theta > 0)[0]
     tan_theta_2 = 1/(cos_theta[idx]**2) - 1
     tan_theta_2_max = (self.rGalaxy/angularDistance)**2
-    tan_theta_2_min = (R_star/angularDistance)**2
+    tan_theta_2_min = (self.R_star/angularDistance)**2
 
     idx0 = np.where((tan_theta_2 < tan_theta_2_max))
     idx = idx[idx0]
     tan_theta_2 = tan_theta_2[idx0]
     tan_theta = np.sqrt(tan_theta_2)
 
-    r = (tan_theta*DA)
+    r = (tan_theta*angularDistance)
 
     m,idx_mask = self.evaluate_profile(r)
 
     idx_mask = np.append(idx_mask,np.where(tan_theta_2<tan_theta_2_min)[0])
-    idx_mask = np.append(idx_mask,[tan_theta_2.argmin()])
+    if tan_theta_2.size > 0:
+      idx_mask = np.append(idx_mask,[tan_theta_2.argmin()])
 
     return idx,idx_mask,m
      
@@ -75,7 +77,7 @@ class KSZ_Isothermal(KSZ_Profile):
     self._table = x,profile
 
 
-  def evaluate_profile(r):
+  def evaluate_profile(self,r):
     rho0, rGalaxy, rInner = self.rho0, self.rGalaxy, self.rInnerGalaxy
     
     Q=rho0*2/r*np.arctan(np.sqrt((rGalaxy/r)**2 - 1))/Mpc
@@ -140,7 +142,7 @@ class KSZ_NFW(KSZ_Profile):
 
   	return exp(0.971 - 0.094*log(Mvir/(1e12*MassSun)))
 
-  def evaluate_profile(r):
+  def evaluate_profile(self,r):
     cs = self._get_concentration(self.Mvir)
     rs = self.Rvir/cs