From 247d8512e4ab6b16c0b290ea0c57f290b5d5c0de Mon Sep 17 00:00:00 2001
From: Guilhem Lavaux <lavaux@iap.fr>
Date: Wed, 18 Jun 2014 16:13:16 +0200
Subject: [PATCH] Added help strings and ksz support

---
 python/_project.pyx                           |   4 +
 python_sample/build_2lpt_skymap.py            |   9 +-
 .../build_nbody_ksz_from_galaxies.py          | 138 ++++++++++++++++
 python_sample/gen_2lpt_density.py             |   2 +-
 python_sample/ksz/__init__.py                 |   3 +
 python_sample/ksz/constants.py                |  31 ++++
 python_sample/ksz/gal_prof.py                 | 149 ++++++++++++++++++
 7 files changed, 333 insertions(+), 3 deletions(-)
 create mode 100644 python_sample/build_nbody_ksz_from_galaxies.py
 create mode 100644 python_sample/ksz/__init__.py
 create mode 100644 python_sample/ksz/constants.py
 create mode 100644 python_sample/ksz/gal_prof.py

diff --git a/python/_project.pyx b/python/_project.pyx
index efe66bd..a8ff9dc 100644
--- a/python/_project.pyx
+++ b/python/_project.pyx
@@ -143,6 +143,10 @@ def interp3d(x not None, y not None,
              z not None,
              npx.ndarray[DTYPE_t, ndim=3] d not None, DTYPE_t Lbox,
              bool periodic=False):
+    """ interp3d(x,y,z,d,Lbox,periodic=False) -> interpolated values
+    
+    Compute the tri-linear interpolation of the given field (d) at the given position (x,y,z). It assumes that they are box-centered coordinates. So (x,y,z) == (0,0,0) is equivalent to the pixel at (Nx/2,Ny/2,Nz/2) with Nx,Ny,Nz = d.shape. If periodic is set, it assumes the box is periodic 
+"""
     cdef npx.ndarray[DTYPE_t] out
     cdef npx.ndarray[DTYPE_t] ax, ay, az
     cdef int i
diff --git a/python_sample/build_2lpt_skymap.py b/python_sample/build_2lpt_skymap.py
index 102d6ef..dfc2acb 100644
--- a/python_sample/build_2lpt_skymap.py
+++ b/python_sample/build_2lpt_skymap.py
@@ -18,6 +18,7 @@ parser.add_argument('--maxval', type=float, default=60)
 parser.add_argument('--depth_min', type=float, default=10)
 parser.add_argument('--depth_max', type=float, default=60)
 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"
@@ -61,8 +62,12 @@ for i in xrange(args.start,args.end,args.step):
   plt.clf()
   d = np.load(args.base_cic % i)
   proj = build_sky_proj(1+d, dmin=args.depth_min,dmax=args.depth_max,iid=args.iid)
+  proj /= (args.depth_max-args.depth_min)
 
-  hp.mollview(proj, fig=1, coord='CG', cmap=plt.cm.coolwarm, title='Sample %d' % i, min=args.minval, max=args.maxval)
-  hp.projscatter(b[idx], l[idx], lw=0, color='g', s=2.0, alpha=0.7)
+  hp.write_map("skymaps/proj_map_%d.fits" % i, proj)
+
+  hp.mollview(proj, fig=1, coord='CG', cmap=plt.cm.copper, title='Sample %d' % i, min=args.minval, max=args.maxval)
+  if args.proj_cat:
+    hp.projscatter(b[idx], l[idx], lw=0, color=[0.1,0.8,0.8], s=2.0, alpha=0.7)
 
   ff.savefig(args.base_fig % i)
diff --git a/python_sample/build_nbody_ksz_from_galaxies.py b/python_sample/build_nbody_ksz_from_galaxies.py
new file mode 100644
index 0000000..c2d999b
--- /dev/null
+++ b/python_sample/build_nbody_ksz_from_galaxies.py
@@ -0,0 +1,138 @@
+import healpy as hp
+import numpy as np
+import cosmotool as ct
+import argparse
+import h5py as h5
+from matplotlib import pyplot as plt
+import ksz
+from ksz.constants import *
+from cosmotool import interp3d
+
+def wrapper_impulsion(f):
+
+  class _Wrapper(object):
+    def __init__(self):
+      pass
+      
+    def __getitem__(self,direction):
+    
+      if 'velocity' in f:
+        return f['velocity'][:,:,:,direction]
+      
+      n = "p%d" % direction
+      return f[n]
+
+  return _Wrapper()
+
+parser=argparse.ArgumentParser(description="Generate Skymaps from CIC maps")
+parser.add_argument('--boxsize', type=float, required=True)
+parser.add_argument('--Nside', type=int, default=128)
+parser.add_argument('--base_h5', type=str, required=True)
+parser.add_argument('--base_fig', type=str, required=True)
+parser.add_argument('--start', type=int, required=True)
+parser.add_argument('--end', type=int, required=True)
+parser.add_argument('--step', type=int, required=True)
+parser.add_argument('--minval', type=float, default=-0.5)
+parser.add_argument('--maxval', type=float, default=0.5)
+parser.add_argument('--depth_min', type=float, default=10)
+parser.add_argument('--depth_max', type=float, default=60)
+parser.add_argument('--iid', type=int, default=0)
+parser.add_argument('--ksz_map', type=str, required=True)
+args = parser.parse_args()
+
+L = args.boxsize
+Nside = args.Nside
+
+def build_unit_vectors(N):
+  ii = np.arange(N)*L/N - 0.5*L
+  d = np.sqrt(ii[:,None,None]**2 + ii[None,:,None]**2 + ii[None,None,:]**2)
+  d[N/2,N/2,N/2] = 1
+  ux = ii[:,None,None] / d
+  uy = ii[None,:,None] / d
+  uz = ii[None,None,:] / d
+
+  return ux,uy,uz
+
+def build_radial_v(v):
+  u = build_unit_vectors(N)
+  vr = v[0] * u[0]
+  vr += v[1] * u[1]
+  vr += v[2] * u[2]
+  return vr
+
+def generate_from_catalog(vfield,Boxsize):
+  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
+  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_mask = np.zeros(12**Nside**2, dtype=np.uint8)
+
+  pb = pbar.ProgressBar(maxval = cat.size, widgets=[pb.Bar(), pb.ETA()]).start()
+
+  for k,i in np.ndenumerate(cat):
+    pb.update(k)
+    l,b=i['gal_long'],i['gal_lat']
+
+    l *= deg2rad
+    b *= deg2rad
+
+    x0 = np.cos(l)*np.cos(b)
+    y0 = np.sin(l)*np.cos(b)
+    z0 = np.sin(b)
+    
+    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)
+
+    vr = interp3d(DA * x0, DA * y0, DA * z0, vfield, Boxsize)
+
+    xp1 = xp[idx0]
+    yp1 = yp[idx0]
+    zp1 = zp[idx0]
+    N2_1 = N2[idx0]
+
+    cos_theta = x0*xp1+y0*yp1+z0*zp1
+    cos_theta /= np.sqrt(x0**2+y0**2+z0**2)*(N2_1)
+
+    idx,idx_masked,m = profiler.projected_profile(cos_theta, DA)
+    idx = idx0[idx]
+    idx_masked = idx0[idx_masked]
+    ksz_template[idx] += vr * m
+    ksz_mask[idx_masked] = 0
+
+  pb.finish()
+
+  return ksz_template, ksz_mask
+
+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:
+    p = wrapper_impulsion(f)
+    v.append(p[i] / f['density'][:])
+
+  # Rescale by Omega_b / Omega_m
+  vr = build_radial_v(v)
+  vr *= -ksz_normalization*rho_mean_matter*1e6
+  del v
+
+  proj = generate_from_catalog(vfield,Boxsize)
+
+  hp.write_map(args.ksz_map % i, proj)
+
+  hp.mollview(proj, fig=1, coord='CG', cmap=plt.cm.coolwarm, title='Sample %d' % i, min=args.minval, 
+              max=args.maxval)
+
+  ff.savefig(args.base_fig % i)
diff --git a/python_sample/gen_2lpt_density.py b/python_sample/gen_2lpt_density.py
index 33c1260..2af9d0a 100644
--- a/python_sample/gen_2lpt_density.py
+++ b/python_sample/gen_2lpt_density.py
@@ -32,7 +32,7 @@ args = parser.parse_args()
 for i in xrange(args.start, args.end, args.step):
   print i
 #  pos,_,density,N,L,_ = bic.run_generation("/nethome/lavaux/remote/borg_2m++_128/initial_density_%d.dat" % i, 0.001, astart, cosmo, supersample=2, do_lpt2=True)
-  pos,_,density,N,L,_ = bic.run_generation("%s/initial_density_%d.dat" % (args.base,i), 0.001, astart, 
+  pos,_,density,N,L,_,_ = bic.run_generation("%s/initial_density_%d.dat" % (args.base,i), 0.001, astart, 
                                            cosmo, supersample=args.supersample, do_lpt2=True)
 
   dcic = ct.cicParticles(pos, L, args.N)
diff --git a/python_sample/ksz/__init__.py b/python_sample/ksz/__init__.py
new file mode 100644
index 0000000..b26b1bb
--- /dev/null
+++ b/python_sample/ksz/__init__.py
@@ -0,0 +1,3 @@
+from .constants import * 
+from .gal_prof import KSZ_Profile, KSZ_Isothermal, KSZ_NFW
+
diff --git a/python_sample/ksz/constants.py b/python_sample/ksz/constants.py
new file mode 100644
index 0000000..9834b3b
--- /dev/null
+++ b/python_sample/ksz/constants.py
@@ -0,0 +1,31 @@
+import numpy as np
+
+Mpc=3.08e22
+rhoc = 1.8864883524081933e-26  # m^(-3)
+sigmaT = 6.6524e-29
+mp = 1.6726e-27
+lightspeed = 299792458.
+v_unit = 1e3  # Unit of 1 km/s
+T_cmb=2.725
+h = 0.71
+Y = 0.245  #The Helium abundance
+Omega_matter = 0.26
+Omega_baryon=0.0445
+
+G=6.67e-11
+MassSun=2e30
+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}
+
+baryon_fraction = Omega_baryon / Omega_matter
+
+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
+
+
+del np
diff --git a/python_sample/ksz/gal_prof.py b/python_sample/ksz/gal_prof.py
new file mode 100644
index 0000000..f7dbed6
--- /dev/null
+++ b/python_sample/ksz/gal_prof.py
@@ -0,0 +1,149 @@
+from .constants import *
+
+# -----------------------------------------------------------------------------
+# Generic profile generator
+# -----------------------------------------------------------------------------
+
+class KSZ_Profile(object):
+  R_star= 0.050 # 15 kpc/h 
+  L_gal0 = 10**(0.4*(tmpp_cat['Msun']-tmpp_cat['Mstar']))
+
+  def __init__(self):
+    self.rGalaxy = 1.0
+
+  def evaluate_profile(r):
+    raise NotImplementedError("Abstract function")
+
+  def projected_profile(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
+
+    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)
+
+    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()])
+
+    return idx,idx_mask,m
+     
+
+# -----------------------------------------------------------------------------
+# Isothermal profile generator
+# -----------------------------------------------------------------------------
+
+class KSZ_Isothermal(KSZ_Profile):
+  sigma_FP=160e3  #m/s
+  R_innergal = 0.226
+
+  def __init__(self, Lgal, x, y=0.0):
+    "Support for Isothermal profile"
+    
+    super(KSZ_Isothermal,self).__init__()
+    
+    self.R_gal = 0.226 * x
+    self.R_innergal *= y
+
+    self.rho0 = self.sigma_FP**2/(2*np.pi*G) # * (Lgal/L_gal0)**(2./3)
+    self.rGalaxy = self.R_gal*(Lgal/self.L_gal0)**(1./3)
+    self.rInnerGalaxy = self.R_innergal*(Lgal/self.L_gal0)**(1./3)
+#    self._prepare()
+
+  def _prepare(self, x_min, x_max):
+    from scipy.integrate import quad
+    from numpy import sqrt, log10
+
+    lmin = log10(x_min)
+    lmax = log10(x_max)
+
+    x = 10**(np.arange(100)*(lmax-lmin)/100.+lmin)
+    profile = np.empty(x.size)
+    
+    nu_tilde = lambda u: (1/(u**2*(1+u)))
+
+    for i in range(x.size):
+        profile[i] = 2*quad(lambda y: (nu_tilde(sqrt(x[i]**2+y**2))), 0, args.x)[0]
+
+    self._table = x,profile
+
+
+  def evaluate_profile(r):
+    rho0, rGalaxy, rInner = self.rho0, self.rGalaxy, self.rInnerGalaxy
+    
+    Q=rho0*2/r*np.arctan(np.sqrt((rGalaxy/r)**2 - 1))/Mpc
+#    Q[r<rInner] = 0
+    return Q,np.where(r<rInner)[0]
+
+
+# -----------------------------------------------------------------------------
+# NFW profile generator
+# -----------------------------------------------------------------------------
+
+class KSZ_NFW(KSZ_Profile):
+  """ Support for NFW profile
+"""
+
+  def __init__(self,x,y=0.0):
+    from numpy import log, pi
+
+    if 'pre_nfw' not in self:
+      self._prepare()
+
+    kiso = KSZ_Isothermal(x,y)
+    r_is = kiso.rGalaxy
+    rho_is = kiso.rho0
+    r_inner = kiso.rInnerGalaxy
+
+    self.Mgal = rho_is*4*pi*(r_is/args.x)*Mpc #Lgal*M_over_L_galaxy
+    self.Rvir = r_is/x
+
+    cs = self._get_concentration(Mgal)
+
+    self.rs = Rvir/cs
+    b = (log(1.+cs)-cs/(1.+cs))
+    
+    self.rho_s = Mgal/(4*pi*b*(rs*Mpc)**3)
+
+
+  def _prepare(self, _x_min=1e-4, _x_max=1e4):
+    from scipy.integrate import quad
+    from numpy import sqrt, log10
+    from scipy.interpolate import interp1d
+
+    lmin = log10(x_min)
+    lmax = log10(x_max)
+
+    x = 10**(np.arange(100)*(lmax-lmin)/100.+lmin)
+    profile = np.empty(x.size)
+    
+    nu_tilde = lambda u: (1/(u*(1+u)**2))
+
+    for i in range(x.size):
+        if x[i] < args.x:
+            profile[i] = 2*quad(lambda y: (nu_tilde(sqrt(x[i]**2+y**2))), 0, np.sqrt((args.x)**2-x[i]**2))[0]
+        else:
+            profile[i] = 0
+
+    # Insert the interpolator into the class definition
+    KSZ_NFW.pre_nfw = self.pre_nfw = interp1d(x,prof)
+
+  def _get_concentration(self, Mvir):
+  	from numpy import exp, log
+
+  	return exp(0.971 - 0.094*log(Mvir/(1e12*MassSun)))
+
+  def evaluate_profile(r):
+    cs = self._get_concentration(self.Mvir)
+    rs = self.Rvir/cs
+    
+    return self.rho_s*rs*Mpc*self.pre_nfw(r/rs),np.array([],dtype=int)
+
+