Update cosmotool 2nd part

external/cosmotool/python_sample/build_2lpt_ksz.py

@@ -0,0 +1,111 @@
+import healpy as hp
+import numpy as np
+import cosmotool as ct
+import argparse
+import h5py as h5
+from matplotlib import pyplot as plt
+
+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)
+
+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)) 
+
+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_sky_proj(density, dmax=60.,dmin=0,iid=0):
+
+  N = density.shape[0]
+  ix = (np.arange(N)-0.5)*L/N - 0.5 * L
+
+
+  dist2 = (ix[:,None,None]**2 + ix[None,:,None]**2 + ix[None,None,:]**2)
+
+  flux = density.transpose().astype(ct.DTYPE) # / dist2
+  dmax=N*dmax/L
+  dmin=N*dmin/L
+  if iid == 0:
+    shifter = np.array([0.5,0.5,0.5])
+  else:
+    shifter = np.array([0.,0.,0.])
+    
+  projsky1 = ct.spherical_projection(Nside, flux, dmin, dmax, integrator_id=iid, shifter=shifter, progress=1)
+
+  return projsky1*L/N
+
+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
+
+for i in xrange(args.start,args.end,args.step):
+  ff=plt.figure(1)
+  plt.clf()
+  with h5.File(args.base_h5 % i, mode="r") as f:
+    p = f['velocity'][:]
+    davg = np.average(np.average(np.average(f['density'][:],axis=0),axis=0),axis=0)
+    p /= davg # Now we have momentum scaled to the mean density
+
+  # Rescale by Omega_b / Omega_m
+  p = p.astype(np.float64)
+  print p.max(), p.min(), ksz_normalization, rho_mean_matter
+  p *= ksz_normalization*rho_mean_matter
+  p *= 1e6 # Put it in uK
+  p *= -1 # ksz has a minus
+  
+  ux,uy,uz = build_unit_vectors(p.shape[0])
+  pr = p[:,:,:,0] * ux + p[:,:,:,1] * uy + p[:,:,:,2] * uz 
+  print p.max(), p.min()
+  print pr.max()*Mpc, pr.min()*Mpc
+  
+  @ct.timeit_quiet
+  def run_proj():
+    return build_sky_proj(pr*Mpc, dmin=args.depth_min,dmax=args.depth_max,iid=args.iid)
+
+  proj = run_proj()
+
+  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)

external/cosmotool/python_sample/build_2lpt_skymap.py

@@ -0,0 +1,76 @@
+import matplotlib
+matplotlib.use('Agg')
+import healpy as hp
+import numpy as np
+import cosmotool as ct
+import argparse
+import h5py as h5
+from matplotlib import pyplot as plt
+
+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_cic', 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)
+parser.add_argument('--maxval', type=float, default=4)
+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"
+INDATA="2m++.npy"
+tmpp = np.load(INDATA)
+
+L = args.boxsize
+Nside = args.Nside
+
+def build_sky_proj(density, dmax=60.,dmin=0,iid=0):
+
+  N = density.shape[0]
+  ix = (np.arange(N)-0.5)*L/N - 0.5 * L
+
+#  dist2 = (ix[:,None,None]**2 + ix[None,:,None]**2 + ix[None,None,:]**2)
+
+  flux = density.transpose().astype(ct.DTYPE) # / dist2
+  dmax=N*dmax/L
+  dmin=N*dmin/L
+  if iid == 0:
+    shifter = np.array([0.5,0.5,0.5])
+  else:
+    shifter = np.array([0.,0.,0.])
+    
+  projsky1 = ct.spherical_projection(Nside, flux, dmin, dmax, integrator_id=iid, shifter=shifter)
+
+  return projsky1*L/N
+
+l,b = tmpp['gal_long'],tmpp['gal_lat']
+
+l = np.radians(l)
+b = np.pi/2 - np.radians(b)
+
+dcmb = tmpp['velcmb']/100.
+
+idx = np.where((dcmb>args.depth_min)*(dcmb<args.depth_max))
+
+
+for i in xrange(args.start,args.end,args.step):
+  ff=plt.figure(1)
+  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.write_map("skymaps/proj_map_%d.fits" % i, proj)
+
+  print proj.min(), proj.max()
+  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)

external/cosmotool/python_sample/build_dipole_ksz_from_galaxies.py

@@ -0,0 +1,287 @@
+import numexpr as ne
+import healpy as hp
+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 *
+from cosmotool import interp3d
+import numpy as np
+from scipy import ndimage
+from scipy.special import sinc
+
+
+def move_direction_new(d_theta, d_phi, theta, phi):
+  cos=np.cos
+  sin=np.sin
+  sqrt=np.sqrt
+
+  amplitude = sqrt(d_theta*d_theta + d_phi*d_phi);
+  cos_alpha = d_theta/amplitude;
+  sin_alpha = d_phi/amplitude;
+
+  if (amplitude == 0):
+    return theta,phi
+
+  cos_d = cos(amplitude);
+  sin_d = sin(amplitude);
+
+  cos_theta = cos(theta);
+  sin_theta = sin(theta);
+
+  cos_phi = cos(phi);
+  sin_phi = sin(phi);
+
+  basis = [
+     [ cos_phi * sin_theta, sin_phi * sin_theta, cos_theta ],
+     [ cos_phi * cos_theta, sin_phi * cos_theta, -sin_theta ],
+     [ -sin_phi, cos_phi, 0 ]
+   ]
+
+  np0 = [ cos_d, cos_alpha*sin_d, sin_alpha*sin_d ]
+  np1 = [ sum([basis[j][i] * np0[j] for j in xrange(3)]) for i in xrange(3) ]
+  dnp = sqrt(sum([np1[i]**2 for i in xrange(3)]))
+
+  theta = np.arccos(np1[2]/dnp);
+  phi = np.arctan2(np1[1], np1[0]) % (2*np.pi);
+
+  return theta,phi
+
+def move_direction_new2(delta_theta, delta_phi, theta, phi):
+	cos,sin,sqrt=np.cos,np.sin,np.sqrt
+	
+	grad_len = sqrt(delta_theta**2 + delta_phi**2)
+	if grad_len==0:
+		return theta,phi
+
+	cth0 = cos(theta)
+	sth0 = sin(theta)
+	
+	topbottom = 1 if (theta < 0.5*np.pi) else -1
+	sinc_grad_len = sinc(grad_len)
+
+	cth = topbottom*cos(grad_len) * cth0 - sinc_grad_len*sth0*delta_theta
+	sth = max(1e-10, sqrt((1.0-cth)*(1.0+cth)) )
+
+	phi = phi + np.arcsin(delta_phi * sinc_grad_len / sth)
+	theta = np.arccos(cth)
+	return theta,phi
+
+move_direction = move_direction_new
+
+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()
+
+
+def build_unit_vectors(N):
+  ii = np.arange(N,dtype=np.float64)/N - 0.5
+  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 compute_vcmb(l, b):
+    # Motion is obtained from Tully (2007): sun_vs_LS + LS_vs_CMB
+    motion = [-25.,-246.,277.];
+
+    x = np.cos(l*np.pi/180) * np.cos(b*np.pi/180)
+    y = np.sin(l*np.pi/180) * np.cos(b*np.pi/180)
+    z = np.sin(b*np.pi/180)
+
+    return x*motion[0] + y*motion[1] + z*motion[2]
+
+
+def compute_vlg(l,b):
+
+    motion = [-79,296,-36]; # [-86, 305, -33];
+
+    x = np.cos(l*np.pi/180) * np.cos(b*np.pi/180)
+    y = np.sin(l*np.pi/180) * np.cos(b*np.pi/180)
+    z = np.sin(b*np.pi/180)
+
+    return x*motion[0] + y*motion[1] + z*motion[2]
+
+
+def generate_from_catalog(dmin,dmax,Nside,perturb=0.0,y=0.0,do_random=False,do_hubble=False,x=2.37,bright=-np.inf,bright_list=[],use_vlg=True,sculpt=-1):
+  import progressbar as pbar
+  
+  cat = np.load("2m++.npy")
+
+  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(Npix, dtype=np.float64)
+  ksz_mask = np.ones(Npix, dtype=np.uint8)
+  if do_hubble:
+    ksz_hubble_template = np.zeros(ksz_template.size, dtype=np.float64)
+
+  for i in pbar.ProgressBar(maxval = cat.size, widgets=[pbar.Bar(), pbar.ETA()])(cat):
+    # Skip too point sources
+    if i['name'] in bright_list:
+      print("Object %s is in bright list" % i['name'])
+      continue
+
+    if do_random:
+      l = np.random.rand()*360
+      b = np.arcsin(2*np.random.rand()-1)*180/np.pi
+    else:
+      l0,b0=i['gal_long'],i['gal_lat']
+
+    l=ne.evaluate('l0*deg2rad')
+    b=ne.evaluate('b0*deg2rad')
+
+    dtheta,dphi = np.random.randn(2)*perturb
+    theta,l=move_direction(dtheta,dphi,0.5*np.pi - b, l)
+	
+    b = 0.5*np.pi-theta
+
+    x0 = np.cos(l)*np.cos(b)
+    y0 = np.sin(l)*np.cos(b)
+    z0 = np.sin(b)
+
+    if use_vlg:
+      vlg = i['best_velcmb'] - compute_vcmb(l0, b0) + compute_vlg(l0, b0)
+      DA = vlg/100
+    else:
+      DA = i['best_velcmb'] / 100
+
+    if DA < dmin or DA > dmax:
+      continue
+
+    Lgal = DA**2*10**(0.4*(tmpp_cat['Msun']-i['K2MRS']+25))
+
+    M_K=i['K2MRS']-5*np.log10(DA)-25
+    # Skip too bright galaxies
+    if M_K < bright:
+      continue
+
+    profiler = ksz.KSZ_Isothermal(Lgal, x, y=y, sculpt=sculpt)
+
+    idx0 = hp.query_disc(Nside, (x0,y0,z0), 3*profiler.rGalaxy/DA)
+
+    xp1 = xp[idx0]
+    yp1 = yp[idx0]
+    zp1 = zp[idx0]
+    N2_1 = N2[idx0]
+
+    cos_theta = ne.evaluate('(x0*xp1+y0*yp1+z0*zp1)/(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] += m
+    ksz_mask[idx_masked] = 0
+    if do_hubble:
+      ksz_hubble_template[idx] += m*DA
+
+  ne.evaluate('ksz_template*ksz_normalization', out=ksz_template)
+
+  result =ksz_template, ksz_mask
+  if do_hubble:
+    ne.evaluate('ksz_hubble_template*ksz_normalization', out=ksz_hubble_template)
+    return result + ( ksz_hubble_template,)
+  else:
+    return result
+
+def get_args():
+  parser=argparse.ArgumentParser(description="Generate Skymaps from CIC maps")
+  parser.add_argument('--Nside', type=int, default=128)
+  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('--ksz_map', type=str, required=True)
+  parser.add_argument('--base_fig', type=str, default="kszfig.png")
+  parser.add_argument('--build_dipole', action='store_true')
+  parser.add_argument('--degrade', type=int, default=-1)
+  parser.add_argument('--y',type=float,default=0.0)
+  parser.add_argument('--x',type=float,default=2.37)
+  parser.add_argument('--random', action='store_true')
+  parser.add_argument('--perturb', type=float, default=0)
+  parser.add_argument('--hubble_monopole', action='store_true')
+  parser.add_argument('--remove_bright', type=float, default=-np.inf)
+  parser.add_argument('--bright_file', type=str)
+  parser.add_argument('--lg', action='store_true')
+  parser.add_argument('--sculpt_beam', type=float, default=-1)
+  return parser.parse_args()
+
+def main():
+
+  args = get_args()
+
+  ff=plt.figure(1)
+  plt.clf()
+  v=[]
+
+  print("Generating map...")
+
+  with open("crap.txt", mode="r") as f:
+    bright_list = [l.split('#')[0].strip(" \t\n\r") for l in f]
+
+  if args.bright_file:
+    with open(args.bright_file, mode="r") as f:
+      idx_name = f.readline().split(',').index('name_2')
+      bright_list = bright_list + [l.split(',')[idx_name] for l in f]
+
+  print("Built bright point source list: " + repr(bright_list))
+  r = generate_from_catalog(args.depth_min,args.depth_max,args.Nside,perturb=args.perturb,y=args.y,do_random=args.random,do_hubble=args.hubble_monopole,x=args.x,bright=args.remove_bright,use_vlg=args.lg,bright_list=bright_list,sculpt=args.sculpt_beam)
+  hubble_map = None
+  if args.hubble_monopole:
+    proj,mask,hubble_map = r
+  else:
+    proj,mask = r
+
+  if args.degrade > 0:
+    proj *= mask
+    proj = hp.ud_grade(proj, nside_out=args.degrade)
+    if hubble_map is not None:
+      hubble_map *= mask
+      hubble_map = hp.ud_grade(hubble_map, nside_out=args.degrade)
+    mask = hp.ud_grade(mask, nside_out=args.degrade)
+    Nside = args.degrade
+  else:
+    Nside = args.Nside
+
+  hp.write_map(args.ksz_map + ".fits", proj)
+  hp.write_map(args.ksz_map + "_mask.fits", mask)
+
+  if args.build_dipole:
+    x,y,z=hp.pix2vec(Nside, np.arange(hp.nside2npix(Nside)))
+    hp.write_map(args.ksz_map + "_x.fits", proj*x)
+    hp.write_map(args.ksz_map + "_y.fits", proj*y)
+    hp.write_map(args.ksz_map + "_z.fits", proj*z)
+
+  if args.hubble_monopole:
+    hp.write_map(args.ksz_map + "_hubble.fits", hubble_map)
+
+  hp.mollview(proj*100*1e6, fig=1, coord='GG', cmap=plt.cm.coolwarm, title='', min=args.minval, 
+              max=args.maxval)
+
+  ff.savefig(args.base_fig)
+  
+main()

external/cosmotool/python_sample/build_nbody_ksz.py

@@ -0,0 +1,109 @@
+import healpy as hp
+import numpy as np
+import cosmotool as ct
+import argparse
+import h5py as h5
+from matplotlib import pyplot as plt
+
+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)
+
+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)) 
+
+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_sky_proj(density, dmax=60.,dmin=0,iid=0):
+
+  N = density.shape[0]
+  ix = (np.arange(N)-0.5)*L/N - 0.5 * L
+
+
+  dist2 = (ix[:,None,None]**2 + ix[None,:,None]**2 + ix[None,None,:]**2)
+
+  flux = density.transpose().astype(ct.DTYPE) # / dist2
+  dmax=N*dmax/L
+  dmin=N*dmin/L
+  if iid == 0:
+    shifter = np.array([0.5,0.5,0.5])
+  else:
+    shifter = np.array([0.,0.,0.])
+    
+  projsky1 = ct.spherical_projection(Nside, flux, dmin, dmax, integrator_id=iid, shifter=shifter)
+
+  return projsky1*L/N
+
+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
+
+for i in xrange(args.start,args.end,args.step):
+  ff=plt.figure(1)
+  plt.clf()
+  with h5.File(args.base_h5 % i, mode="r") as f:
+    p = f['velocity'][:]
+    davg = np.average(np.average(np.average(f['density'][:],axis=0),axis=0),axis=0)
+    p /= davg # Now we have momentum scaled to the mean density
+
+  # Rescale by Omega_b / Omega_m
+  p = p.astype(np.float64)
+  print p.max(), p.min(), ksz_normalization, rho_mean_matter
+  p *= -ksz_normalization*rho_mean_matter*1e6
+  
+  ux,uy,uz = build_unit_vectors(p.shape[0])
+  pr = p[:,:,:,0] * ux + p[:,:,:,1] * uy + p[:,:,:,2] * uz 
+  print p.max(), p.min()
+  print pr.max()*Mpc, pr.min()*Mpc
+  
+  @ct.timeit_quiet
+  def run_proj():
+    return build_sky_proj(pr*Mpc, dmin=args.depth_min,dmax=args.depth_max,iid=args.iid)
+
+  run_proj()
+
+  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)

external/cosmotool/python_sample/build_nbody_ksz_from_galaxies.py

@@ -0,0 +1,167 @@
+import healpy as hp
+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 *
+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):
+  N = v[0].shape[0]
+  u = build_unit_vectors(N)
+  vr = v[0] * u[2]
+  vr += v[1] * u[1]
+  vr += v[2] * u[0]
+
+  return vr.transpose()
+
+def generate_from_catalog(vfield,Boxsize,dmin,dmax):
+  import progressbar as pbar
+  
+  cat = np.load("2m++.npy")
+
+
+  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(Npix, dtype=np.float64)
+  ksz_mask = np.zeros(Npix, dtype=np.uint8)
+
+  pb = pbar.ProgressBar(maxval = cat.size, widgets=[pbar.Bar(), pbar.ETA()]).start()
+
+  for k,i in np.ndenumerate(cat):
+    pb.update(k[0])
+    l,b=i['gal_long'],i['gal_lat']
+    ra,dec=i['ra'],i['dec']
+
+    l *= deg2rad
+    b *= deg2rad
+
+    ra *= deg2rad
+    dec *= deg2rad
+
+#    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))
+
+    profiler = ksz.KSZ_Isothermal(Lgal, 2.37)
+
+    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]
+    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=[]
+  fname = args.base_h5 % i
+  if False:
+    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_unit_vectors(128)
+#  vr *= 1000 * 500
+    vr *= ksz_normalization*1e6
+    del v
+#    np.save("vr.npy", vr)
+  else:
+    print("Loading vrs...")
+    vr = np.load("vr.npy")
+
+  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)
+
+  ff.savefig(args.base_fig % i)

external/cosmotool/python_sample/build_nbody_skymap.py

@@ -0,0 +1,58 @@
+import healpy as hp
+import numpy as np
+import cosmotool as ct
+import h5py as h5
+from matplotlib import pyplot as plt
+
+L=600.
+Nside=128
+
+INDATA="/nethome/lavaux/Copy/PlusSimulation/BORG/Input_Data/2m++.npy"
+tmpp = np.load(INDATA)
+
+def build_sky_proj(density, dmax=60.,dmin=0):
+
+  N = density.shape[0]
+  ix = (np.arange(N)-0.5)*L/N - 0.5 * L
+
+
+  dist2 = (ix[:,None,None]**2 + ix[None,:,None]**2 + ix[None,None,:]**2)
+
+  flux = density.transpose().astype(ct.DTYPE) # / dist2
+  dmax=N*dmax/L
+  dmin=N*dmin/L
+  projsky1 = ct.spherical_projection(Nside, flux, dmin, dmax, integrator_id=1)
+#  projsky0 = ct.spherical_projection(Nside, flux, 0, 52, integrator_id=0)
+
+  return projsky1*L/N#,projsky0
+
+l,b = tmpp['gal_long'],tmpp['gal_lat']
+
+l = np.radians(l)
+b = np.pi/2 - np.radians(b)
+
+dcmb = tmpp['velcmb']/100.
+
+idx = np.where((dcmb>10)*(dcmb<60))
+
+plt.figure(1)
+plt.clf()
+if True:
+  with h5.File("fields.h5", mode="r") as f:
+    d = f["density"][:].transpose()
+    d /= np.average(np.average(np.average(d,axis=0),axis=0),axis=0)
+    proj = build_sky_proj(d, dmin=10,dmax=60.)
+  proj0 = proj1 = proj
+else:
+  d = np.load("icgen/dcic0.npy")
+  proj0 = build_sky_proj(1+d, dmin=10,dmax=60.)
+  d = np.load("icgen/dcic1.npy")
+  proj1 = build_sky_proj(1+d, dmin=10,dmax=60.)
+
+hp.mollview(proj0, fig=1, coord='CG', max=60, cmap=plt.cm.coolwarm)
+hp.projscatter(b[idx], l[idx], lw=0, color='g', s=5.0, alpha=0.8)
+
+plt.figure(2)
+plt.clf()
+hp.mollview(proj1, fig=2, coord='CG', max=60, cmap=plt.cm.coolwarm)
+hp.projscatter(b[idx], l[idx], lw=0, color='g', s=5.0, alpha=0.8)

external/cosmotool/python_sample/gen_2lpt_asmooth.py

@@ -0,0 +1,48 @@
+import os
+import h5py as h5
+import numpy as np
+import cosmotool as ct
+import icgen as bic
+import icgen.cosmogrowth as cg
+import sys
+import argparse
+
+#ADAPT_SMOOTH="/home/bergeron1NS/lavaux/Software/cosmotool/build/sample/simple3DFilter"
+ADAPT_SMOOTH="/home/guilhem/PROJECTS/cosmotool/build/sample/simple3DFilter"
+cosmo={'omega_M_0':0.3175, 'h':0.6711}
+cosmo['omega_lambda_0']=1-cosmo['omega_M_0']
+cosmo['omega_k_0'] = 0
+cosmo['omega_B_0']=0.049
+cosmo['SIGMA8']=0.8344
+cosmo['ns']=0.9624
+N0=256
+
+doSimulation=True
+
+astart=1.
+
+parser=argparse.ArgumentParser(description="Generate CIC density from 2LPT")
+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('--base', type=str, required=True)
+parser.add_argument('--N', type=int, default=256)
+parser.add_argument('--output', type=str, default="fields_%d.h5")
+parser.add_argument('--supersample', type=int, default=1)
+args = parser.parse_args()
+
+
+
+for i in [4629]:#xrange(args.start, args.end, args.step):
+  print i
+
+  pos,vel,density,N,L,_,_ = bic.run_generation("%s/initial_density_%d.dat" % (args.base,i), 0.001, astart, 
+                                           cosmo, supersample=args.supersample, do_lpt2=True)
+
+  q = pos + vel + [np.ones(vel[0].shape[0])]
+
+  with h5.File("particles.h5", mode="w") as f:
+    f.create_dataset("particles", data=np.array(q).transpose())
+  
+  os.system(ADAPT_SMOOTH + " %s %lg %lg %d %lg %lg %lg" % ("particles.h5", 3000000, L, args.N, 0, 0, 0))
+  os.rename("fields.h5", args.output % i)

external/cosmotool/python_sample/gen_2lpt_density.py

@@ -0,0 +1,55 @@
+import numexpr as ne
+import numpy as np
+import cosmotool as ct
+import icgen as bic
+import icgen.cosmogrowth as cg
+import sys
+import argparse
+
+cosmo={'omega_M_0':0.3175, 'h':0.6711}
+cosmo['omega_lambda_0']=1-cosmo['omega_M_0']
+cosmo['omega_k_0'] = 0
+cosmo['omega_B_0']=0.049
+cosmo['SIGMA8']=0.8344
+cosmo['ns']=0.9624
+N0=256
+
+doSimulation=True
+
+astart=1.
+
+parser=argparse.ArgumentParser(description="Generate CIC density from 2LPT")
+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('--base', type=str, required=True)
+parser.add_argument('--N', type=int, default=256)
+parser.add_argument('--output', type=str, default="dcic_%d.npy")
+parser.add_argument('--supersample', type=int, default=1)
+parser.add_argument('--rsd', action='store_true') 
+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,vel,density,N,L,_,_ = bic.run_generation("%s/initial_density_%d.dat" % (args.base,i), 0.001, astart, 
+                                           cosmo, supersample=args.supersample, do_lpt2=True, needvel=True)
+
+  if args.rsd:
+    inv_r2 = ne.evaluate('1/sqrt(x**2+y**2+z**2)', 
+        local_dict={'x':pos[0], 'y':pos[1], 'z':pos[2]})
+    rsd = lambda p,v: ne.evaluate('x + (x*vx)*inv_r2 / H', 
+        local_dict={'x':p, 'inv_r2':inv_r2,
+         'vx':v, 'H':100.0}, out=p, casting='unsafe')
+
+    rsd(pos[0], vel[0])
+    rsd(pos[1], vel[1])
+    rsd(pos[2], vel[2])
+
+  dcic = ct.cicParticles(pos, L, args.N)
+  dcic /= np.average(np.average(np.average(dcic, axis=0), axis=0), axis=0)
+  dcic -= 1
+
+  np.save(args.output % i, dcic)

external/cosmotool/python_sample/icgen/__init__.py

@@ -0,0 +1,2 @@
+from borgicgen import *
+import cosmogrowth

external/cosmotool/python_sample/icgen/borgadaptor.py

@@ -0,0 +1,55 @@
+import numpy as np
+
+def fourier_analysis(borg_vol):
+  L = (borg_vol.ranges[1]-borg_vol.ranges[0])
+  N = borg_vol.density.shape[0]
+    
+  return np.fft.rfftn(borg_vol.density)*(L/N)**3, L, N
+
+def borg_upgrade_sampling(dhat, supersample):
+    N = dhat.shape[0]
+    N2 = N * supersample
+    dhat_new = np.zeros((N2, N2, N2/2+1), dtype=np.complex128)
+
+    hN = N/2
+    dhat_new[:hN, :hN, :hN+1] = dhat[:hN, :hN, :]
+    dhat_new[:hN, (N2-hN):N2, :hN+1] = dhat[:hN, hN:, :]
+    dhat_new[(N2-hN):N2, (N2-hN):N2, :hN+1] = dhat[hN:, hN:, :]
+    dhat_new[(N2-hN):N2, :hN, :hN+1] = dhat[hN:, :hN, :]
+
+    return dhat_new, N2
+
+def half_pixel_shift(borg, doshift=False):
+
+    dhat,L,N = fourier_analysis(borg)
+    if not doshift:
+      return dhat, L
+
+    return bare_half_pixel_shift(dhat, L, N)
+
+def bare_half_pixel_shift(dhat, L, N, doshift=False):
+
+#    dhat_new,N2 = borg_upgrade_sampling(dhat, 2)
+#    d = (np.fft.irfftn(dhat_new)*(N2/L)**3)[1::2,1::2,1::2]
+#    del dhat_new
+#    dhat = np.fft.rfftn(d)*(L/N)**3
+#    return dhat, L
+
+#    dhat2 = np.zeros((N,N,N),dtype=np.complex128)
+#    dhat2[:,:,:N/2+1] = dhat
+#    dhat2[N:0:-1, N:0:-1, N:N/2:-1] = np.conj(dhat[1:,1:,1:N/2])
+#    dhat2[0, N:0:-1, N:N/2:-1] = np.conj(dhat[0, 1:, 1:N/2])
+#    dhat2[N:0:-1, 0, N:N/2:-1] = np.conj(dhat[1:, 0, 1:N/2])
+#    dhat2[0,0,N:N/2:-1] = np.conj(dhat[0, 0, 1:N/2])
+
+    ik = np.fft.fftfreq(N,d=L/N)*2*np.pi
+    phi = 0.5*L/N*(ik[:,None,None]+ik[None,:,None]+ik[None,None,:(N/2+1)])
+#    phi %= 2*np.pi
+    phase = np.cos(phi)+1j*np.sin(phi)
+    dhat = dhat*phase    
+    dhat[N/2,:,:] = 0
+    dhat[:,N/2,:] = 0
+    dhat[:,:,N/2] = 0
+    
+    return dhat, L
+

external/cosmotool/python_sample/icgen/borgicgen.py

@@ -0,0 +1,228 @@
+import cosmotool as ct
+import numpy as np
+import cosmolopy as cpy
+from cosmogrowth import *
+import borgadaptor as ba
+
+@ct.timeit
+def gen_posgrid(N, L, delta=1, dtype=np.float32):
+    """ Generate an ordered lagrangian grid"""
+
+    ix = (np.arange(N)*(L/N*delta)).astype(dtype)
+
+    x = ix[:,None,None].repeat(N, axis=1).repeat(N, axis=2)
+    y = ix[None,:,None].repeat(N, axis=0).repeat(N, axis=2)
+    z = ix[None,None,:].repeat(N, axis=0).repeat(N, axis=1)
+
+    return x.reshape((x.size,)), y.reshape((y.size,)), z.reshape((z.size,))
+
+def bin_power(P, L, bins=20, range=(0,1.), dev=False):
+
+  N = P.shape[0]
+  ik = np.fft.fftfreq(N, d=L/N)*2*np.pi
+
+  k = np.sqrt(ik[:,None,None]**2 + ik[None,:,None]**2 + ik[None,None,:(N/2+1)]**2)
+
+  H,b = np.histogram(k, bins=bins, range=range)
+  Hw,b = np.histogram(k, bins=bins, weights=P, range=range)
+  
+  if dev:
+    return Hw/(H-1), 0.5*(b[1:]+b[0:bins]), 1.0/np.sqrt(H)
+  else:
+    return Hw/(H-1), 0.5*(b[1:]+b[0:bins])
+
+def compute_power_from_borg(input_borg, a_borg, cosmo, bins=10, range=(0,1)):
+  borg_vol = ct.read_borg_vol(input_borg)
+  N = borg_vol.density.shape[0]
+
+  cgrowth = CosmoGrowth(**cosmo)
+  D1 = cgrowth.D(1)
+  D1_0 = D1/cgrowth.D(a_borg)
+  print("D1_0=%lg" % D1_0)
+
+  density_hat, L = ba.half_pixel_shift(borg_vol)
+
+  return bin_power(D1_0**2*np.abs(density_hat)**2/L**3, L, bins=bins, range=range)
+  
+def compute_ref_power(L, N, cosmo, bins=10, range=(0,1), func='HU_WIGGLES'):
+  ik = np.fft.fftfreq(N, d=L/N)*2*np.pi
+
+  k = np.sqrt(ik[:,None,None]**2 + ik[None,:,None]**2 + ik[None,None,:(N/2+1)]**2)
+  p = ct.CosmologyPower(**cosmo)
+  p.setFunction(func)
+  p.normalize(cosmo['SIGMA8'])
+
+  return bin_power(p.compute(k)*cosmo['h']**3, L, bins=bins, range=range)
+
+
+def do_supergenerate(density, density_out=None, mulfac=None,zero_fill=False,Pk=None,L=None,h=None):
+ 
+  N = density.shape[0]
+  
+  if density_out is None:
+    assert mulfac is not None
+    Ns = mulfac*N
+    density_out = np.zeros((Ns,Ns,Ns/2+1), dtype=np.complex128)
+    density_out[:] = np.nan
+  elif mulfac is None:
+    mulfac = density_out.shape[0] / N
+    Ns = density_out.shape[0]
+    assert (density_out.shape[0] % N) == 0
+  	
+  assert len(density_out.shape) == 3
+  assert density_out.shape[0] == density_out.shape[1]
+  assert density_out.shape[2] == (density_out.shape[0]/2+1)
+  
+  hN = N/2
+  density_out[:hN,               :hN, :hN+1] = density[:hN, :hN, :]
+  density_out[:hN,        (Ns-hN):Ns, :hN+1] = density[:hN, hN:, :]
+  density_out[(Ns-hN):Ns, (Ns-hN):Ns, :hN+1] = density[hN:, hN:, :]
+  density_out[(Ns-hN):Ns,        :hN, :hN+1] = density[hN:, :hN, :]
+
+  if mulfac > 1:
+
+    cond=np.isnan(density_out)
+    if zero_fill:
+      density_out[cond] = 0
+    else:
+
+      if Pk is not None:
+        assert L is not None and h is not None
+
+      @ct.timeit_quiet
+      def build_Pk():
+         ik = np.fft.fftfreq(Ns, d=L/Ns)*2*np.pi
+         k = ne.evaluate('sqrt(kx**2 + ky**2 + kz**2)', {'kx':ik[:,None,None], 'ky':ik[None,:,None], 'kz':ik[None,None,:(Ns/2+1)]})
+         return Pk.compute(k)*L**3
+
+      print np.where(np.isnan(density_out))[0].size
+      Nz = np.count_nonzero(cond)
+      amplitude = np.sqrt(build_Pk()[cond]/2) if Pk is not None else (1.0/np.sqrt(2))
+      density_out.real[cond] = np.random.randn(Nz) * amplitude
+      density_out.imag[cond] = np.random.randn(Nz) * amplitude
+      print np.where(np.isnan(density_out))[0].size
+    
+      # Now we have to fix the Nyquist plane
+      hNs = Ns/2
+      nyquist = density_out[:, :, hNs]
+      Nplane = nyquist.size
+      nyquist.flat[:Nplane/2] = np.sqrt(2.0)*nyquist.flat[Nplane:Nplane/2:-1].conj()
+
+
+  return density_out
+
+@ct.timeit_quiet
+def run_generation(input_borg, a_borg, a_ic, cosmo, supersample=1, supergenerate=1, do_lpt2=True, shiftPixel=False, psi_instead=False, needvel=True, func='HU_WIGGLES'):
+    """ Generate particles and velocities from a BORG snapshot. Returns a tuple of 
+    (positions,velocities,N,BoxSize,scale_factor)."""
+
+    borg_vol = ct.read_borg_vol(input_borg)
+    N = borg_vol.density.shape[0]
+
+    cgrowth = CosmoGrowth(**cosmo)
+
+    density, L = ba.half_pixel_shift(borg_vol, doshift=shiftPixel)
+
+
+    # Compute LPT scaling coefficient
+    D1 = cgrowth.D(a_ic)
+    D1_0 = D1/cgrowth.D(a_borg)
+    Dborg = cgrowth.D(a_borg)/cgrowth.D(1.0)
+    print "D1_0=%lg" % D1_0
+
+    if supergenerate>1:
+      print("Doing supergeneration (factor=%d)" % supergenerate)
+      p = ct.CosmologyPower(**cosmo)
+      p.setFunction(func)
+      p.normalize(cosmo['SIGMA8']*Dborg)
+      density = do_supergenerate(density,mulfac=supergenerate,Pk=p,L=L,h=cosmo['h'])
+
+    lpt = LagrangianPerturbation(-density, L, fourier=True, supersample=supersample)
+
+    # Generate grid
+    posq = gen_posgrid(N*supersample, L)
+    vel= []
+    posx = []
+    
+    velmul = cgrowth.compute_velmul(a_ic) if not psi_instead else 1
+    
+    D2 = -3./7 * D1_0**2
+
+    if do_lpt2:
+      psi2 = lpt.lpt2('all')
+    for j in xrange(3):
+        # Generate psi_j (displacement along j)
+        print("LPT1 axis=%d" % j)
+        psi = D1_0*lpt.lpt1(j)
+        psi = psi.reshape((psi.size,))
+        if do_lpt2:
+          print("LPT2 axis=%d" % j)
+          psi += D2 * psi2[j].reshape((psi2[j].size,))
+        # Generate posx
+        posx.append(((posq[j] + psi)%L).astype(np.float32))
+        # Generate vel
+        if needvel:
+          vel.append((psi*velmul).astype(np.float32)) 
+
+    print("velmul=%lg" % (cosmo['h']*velmul))
+
+    lpt.cube.dhat = lpt.dhat
+    density = lpt.cube.irfft()
+    density *= (cgrowth.D(1)/cgrowth.D(a_borg))
+
+    return posx,vel,density,N*supergenerate*supersample,L,a_ic,cosmo
+
+
+@ct.timeit_quiet
+def whitify(density, L, cosmo, supergenerate=1, zero_fill=False, func='HU_WIGGLES'):
+
+  N = density.shape[0]
+  p = ct.CosmologyPower(**cosmo)
+  p.setFunction(func)
+  p.normalize(cosmo['SIGMA8'])
+
+  @ct.timeit_quiet
+  def build_Pk():
+    ik = np.fft.fftfreq(N, d=L/N)*2*np.pi
+    k = np.sqrt(ik[:,None,None]**2 + ik[None,:,None]**2 + ik[None,None,:(N/2+1)]**2)
+    return p.compute(k)*L**3
+    
+  Pk = build_Pk()
+  Pk[0,0,0]=1
+
+  cube = CubeFT(L, N)
+  cube.density = density
+  density_hat = cube.rfft()
+  density_hat /= np.sqrt(Pk)
+
+  Ns = N*supergenerate
+
+  density_hat_super = do_supergenerate(density_hat, mulfac=supergenerate)
+
+  cube = CubeFT(L, Ns)
+  cube.dhat = density_hat_super
+  return np.fft.irfftn(density_hat_super)*Ns**1.5
+
+
+    
+def write_icfiles(*generated_ic, **kwargs):
+  """Write the initial conditions from the tuple returned by run_generation"""
+  
+  supergenerate=kwargs.get('supergenerate', 1)
+  zero_fill=kwargs.get('zero_fill', False)
+  posx,vel,density,N,L,a_ic,cosmo = generated_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(["z","y","x"]):
+    ct.writeGrafic("Data/ic_velc%s" % c, vel[i].reshape((N,N,N)), L, a_ic, **cosmo)
+
+  ct.writeGrafic("Data/ic_deltab", density, L, a_ic, **cosmo)
+  
+  ct.writeWhitePhase("Data/white.dat", whitify(density, L, cosmo, supergenerate=supergenerate,zero_fill=zero_fill))
+  
+  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)
+    f.write("2\n\n0\nwhite.dat\n0\npadding_white.dat\n")
+

external/cosmotool/python_sample/icgen/cosmogrowth.py

@@ -0,0 +1,202 @@
+import numexpr as ne
+import multiprocessing
+import pyfftw
+import weakref
+import numpy as np
+import cosmolopy as cpy
+import cosmotool as ct
+
+class CubeFT(object):
+  def __init__(self, L, N, max_cpu=-1):
+  
+      self.N = N
+      self.align = pyfftw.simd_alignment
+      self.L = L
+      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._rfft = pyfftw.FFTW(self._density, self._dhat, axes=(0,1,2), threads=self.max_cpu, normalize_idft=False)
+
+  def rfft(self):
+      return ne.evaluate('c*a', local_dict={'c':self._rfft(normalise_idft=False),'a':(self.L/self.N)**3})
+      
+  def irfft(self):
+      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
+  def set_dhat(self, in_dhat):
+      self._dhat[:] = in_dhat 
+  dhat = property(get_dhat, set_dhat, None)
+      
+  def get_density(self):
+      return self._density      
+  def set_density(self, d):
+      self._density[:] = d 
+  density = property(get_density, set_density, None)
+      
+
+class CosmoGrowth(object):
+
+  def __init__(self, **cosmo):
+    self.cosmo = cosmo
+
+  def D(self, a):
+    return cpy.perturbation.fgrowth(1/a-1, self.cosmo['omega_M_0'], unnormed=True)
+
+  def compute_E(self, a):
+    om = self.cosmo['omega_M_0']
+    ol = self.cosmo['omega_lambda_0']
+    ok = self.cosmo['omega_k_0']
+
+    E = np.sqrt(om/a**3 + ol + ok/a**2)
+  
+    H2 = -3*om/a**4  - 2*ok/a**3
+
+    Eprime = 0.5*H2/E
+
+    return E,Eprime
+
+  def Ddot(self, a):
+    E,Eprime = self.compute_E(a)
+    D = self.D(a)
+    Ddot_D = Eprime/E + 2.5 * self.cosmo['omega_M_0']/(a**3*E**2*D)
+    Ddot_D *= a
+    return Ddot_D
+
+  def compute_velmul(self, a):
+    E,_ = self.compute_E(a)
+    velmul = self.Ddot(a)
+    velmul *= 100 * a * E
+    return velmul
+
+
+
+
+
+class LagrangianPerturbation(object):
+
+    def __init__(self,density,L, fourier=False, supersample=1, max_cpu=-1):
+    
+        self.L = L
+        self.N = density.shape[0]
+        
+        self.max_cpu = max_cpu
+        self.cube = CubeFT(self.L, self.N, max_cpu=max_cpu)
+        
+        if not fourier:
+          self.cube.density = density
+          self.dhat = self.cube.rfft().copy()
+        else:
+          self.dhat = density.copy()
+
+        if supersample > 1:
+          self.upgrade_sampling(supersample)
+        self.ik = np.fft.fftfreq(self.N, d=L/self.N)*2*np.pi
+        self._kx = self.ik[:,None,None]
+        self._ky = self.ik[None,:,None]
+        self._kz = self.ik[None,None,:(self.N/2+1)]
+        self.cache = {}#weakref.WeakValueDictionary()
+
+    @ct.timeit_quiet
+    def upgrade_sampling(self, supersample):
+        N2 = self.N * supersample
+        N = self.N
+        dhat_new = np.zeros((N2, N2, N2/2+1), dtype=np.complex128)
+
+        hN = N/2
+        dhat_new[:hN, :hN, :hN+1] = self.dhat[:hN, :hN, :]
+        dhat_new[:hN, (N2-hN):N2, :hN+1] = self.dhat[:hN, hN:, :]
+        dhat_new[(N2-hN):N2, (N2-hN):N2, :hN+1] = self.dhat[hN:, hN:, :]
+        dhat_new[(N2-hN):N2, :hN, :hN+1] = self.dhat[hN:, :hN, :]
+
+        self.dhat = dhat_new
+        self.N = N2
+        self.cube = CubeFT(self.L, self.N, max_cpu=self.max_cpu)
+        
+    @ct.timeit_quiet
+    def _gradient(self, phi, direction):
+        if direction == 'all':
+          dirs = [0,1,2]
+          copy = True
+        else:
+          dirs = [direction]
+          copy = False
+        ret=[]
+        for dir in dirs:
+          ne.evaluate('phi_hat * i * kv / (kx**2 + ky**2 + kz**2)', out=self.cube.dhat,
+              local_dict={'i':-1j, 'phi_hat':phi, 'kv':self._kdir(dir),
+                          'kx':self._kx, 'ky':self._ky, 'kz':self._kz},casting='unsafe')
+#        self.cube.dhat = self._kdir(direction)*1j*phi 
+          self.cube.dhat[0,0,0] = 0
+          x = self.cube.irfft()
+          ret.append(x.copy() if copy else x)
+        return ret[0] if len(ret)==1 else ret
+            
+    @ct.timeit_quiet
+    def lpt1(self, direction=0):
+        return self._gradient(self.dhat, direction)
+        
+    def new_shape(self,direction, q=3, half=False):
+        N0 = (self.N/2+1) if half else self.N
+        return ((1,)*direction) + (N0,) + ((1,)*(q-1-direction))
+
+    def _kdir(self, direction, q=3):
+        if direction != q-1:
+            return self.ik.reshape(self.new_shape(direction, q=q))
+        else:
+            return self.ik[:self.N/2+1].reshape(self.new_shape(direction, q=q, half=True))
+
+    def _get_k2(self, q=3):
+        if 'k2' in self.cache:
+            return self.cache['k2']
+            
+        k2 = self._kdir(0, q=q)**2
+        for d in xrange(1,q):
+            k2 = k2 + self._kdir(d, q=q)**2
+            
+        self.cache['k2'] = k2
+        return k2
+        
+    def _do_irfft(self, array, copy=True):
+        if copy:
+          self.cube.dhat = array
+        return self.cube.irfft()
+        
+    def _do_rfft(self, array, copy=True):
+        if copy:
+          self.cube.density = array
+        return self.cube.rfft()
+
+    @ct.timeit_quiet
+    def lpt2(self, direction=0):
+#        k2 = self._get_k2()
+#        k2[0,0,0] = 1
+
+        inv_k2 = ne.evaluate('1/(kx**2+ky**2+kz**2)', {'kx':self._kdir(0),'ky':self._kdir(1),'kz':self._kdir(2)})
+        inv_k2[0,0,0]=0
+        potgen0 = lambda i: ne.evaluate('kdir**2*d*ik2',out=self.cube.dhat,local_dict={'kdir':self._kdir(i),'d':self.dhat,'ik2':inv_k2}, casting='unsafe' )
+        potgen = lambda i,j: ne.evaluate('kdir0*kdir1*d*ik2',out=self.cube.dhat,local_dict={'kdir0':self._kdir(i),'kdir1':self._kdir(j),'d':self.dhat,'ik2':inv_k2}, casting='unsafe' )
+
+        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( potgen0(j) ).copy()
+                for i in xrange(j+1, 3):
+                    with ct.time_block("LPT2 elemental (%d,%d)" %(i,j)):
+                      ne.evaluate('div + q * pot', out=div_phi2, 
+                            local_dict={'div':div_phi2, 'q':q,'pot':self._do_irfft( potgen0(i), copy=False ) } 
+                            )
+                      ne.evaluate('div - pot**2',out=div_phi2,
+                            local_dict={'div':div_phi2,'pot':self._do_irfft(potgen(i,j), copy=False) }
+                            )
+
+            phi2_hat = self._do_rfft(div_phi2)
+            #self.cache['lpt2_potential'] = phi2_hat
+            del div_phi2
+        else:
+            phi2_hat = self.cache['lpt2_potential']
+            
+        return self._gradient(phi2_hat, direction)

external/cosmotool/python_sample/icgen/gen_ic_from_borg.py

@@ -0,0 +1,24 @@
+import pyfftw
+import numpy as np
+import cosmotool as ct
+import borgicgen as bic
+import pickle
+
+with file("wisdom") as f:
+  pyfftw.import_wisdom(pickle.load(f))
+
+cosmo={'omega_M_0':0.3175, 'h':0.6711}
+cosmo['omega_lambda_0']=1-cosmo['omega_M_0']
+cosmo['omega_k_0'] = 0
+cosmo['omega_B_0']=0.049
+cosmo['SIGMA8']=0.8344
+cosmo['ns']=0.9624
+
+supergen=1
+zstart=99
+astart=1/(1.+zstart)
+halfPixelShift=False
+zero_fill=False
+
+if __name__=="__main__":
+    bic.write_icfiles(*bic.run_generation("initial_density_1872.dat", 0.001, astart, cosmo, supersample=1, shiftPixel=halfPixelShift, do_lpt2=False, supergenerate=supergen), supergenerate=1, zero_fill=zero_fill)

external/cosmotool/python_sample/icgen/test_ic_from_borg.py

@@ -0,0 +1,63 @@
+import numpy as np
+import cosmotool as ct
+import borgicgen as bic
+import cosmogrowth as cg
+import sys
+
+cosmo={'omega_M_0':0.3175, 'h':0.6711}
+cosmo['omega_lambda_0']=1-cosmo['omega_M_0']
+cosmo['omega_k_0'] = 0
+cosmo['omega_B_0']=0.049
+cosmo['SIGMA8']=0.8344
+cosmo['ns']=0.9624
+N0=256
+
+doSimulation=False
+simShift=False
+
+snap_id=int(sys.argv[1])
+astart=1/100.
+
+if doSimulation:
+  s = ct.loadRamsesAll("/nethome/lavaux/remote2/borgsim3/", snap_id, doublePrecision=True)
+  astart=s.getTime()
+  L = s.getBoxsize()
+
+  p = s.getPositions()
+  Nsim = int( np.round( p[0].size**(1./3)) )
+  print("Nsim = %d" % Nsim)
+
+  if simShift:
+    p = [(q-0.5*L/Nsim)%L for q in p]
+
+  dsim = ct.cicParticles(p[::-1], L, N0)
+  dsim /= np.average(np.average(np.average(dsim, axis=0), axis=0), axis=0)
+  dsim -= 1
+
+  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_1872.dat", 0.001, astart, cosmo, supersample=1, do_lpt2=False, supergenerate=2)
+
+dcic = ct.cicParticles(pos, L, N0)
+dcic /= np.average(np.average(np.average(dcic, axis=0), axis=0), axis=0)
+dcic -= 1
+
+dcic_hat = np.fft.rfftn(dcic)*(L/N0)**3
+dens_hat = np.fft.rfftn(density)*(L/N0)**3
+
+Pcic, bcic = bic.bin_power(np.abs(dcic_hat)**2/L**3, L, range=(0,4.), bins=150)
+Pdens, bdens = bic.bin_power(np.abs(dens_hat)**2/L**3, L, range=(0,4.), bins=150)
+
+cgrowth = cg.CosmoGrowth(**cosmo)
+D1 = cgrowth.D(astart)
+D1_0 = D1/cgrowth.D(1)#0.001)
+
+Pref, bref = bic.compute_ref_power(L, N0, cosmo, range=(0,4.), bins=150)
+
+Pcic /= D1_0**2
+
+#borg_evolved = ct.read_borg_vol("final_density_1380.dat")
+#dborg_hat = np.fft.rfftn(borg_evolved.density)*L**3/borg_evolved.density.size
+
+#Pborg, bborg = bic.bin_power(np.abs(dborg_hat)**2/L**3, L, range=(0,1.),bins=150)

external/cosmotool/python_sample/icgen/test_whitify.py

@@ -0,0 +1,41 @@
+import numpy as np
+import cosmotool as ct
+import borgicgen as bic
+from matplotlib import pyplot as plt
+
+cosmo={'omega_M_0':0.3175, 'h':0.6711}
+cosmo['omega_lambda_0']=1-cosmo['omega_M_0']
+cosmo['omega_k_0'] = 0
+cosmo['omega_B_0']=0.049
+cosmo['SIGMA8']=0.8344
+cosmo['ns']=0.9624
+
+zstart=50
+astart=1/(1.+zstart)
+halfPixelShift=False
+
+posx,vel,density,N,L,a_ic,cosmo = bic.run_generation("initial_condition_borg.dat", 0.001, astart, cosmo, supersample=1, shiftPixel=halfPixelShift, do_lpt2=False)
+
+w1 = bic.whitify(density, L, cosmo, supergenerate=1)
+w2 = bic.whitify(density, L, cosmo, supergenerate=2)
+
+N = w1.shape[0]
+Ns = w2.shape[0]
+
+w1_hat = np.fft.rfftn(w1)*(L/N)**3
+w2_hat = np.fft.rfftn(w2)*(L/Ns)**3
+
+P1, b1, dev1 = bic.bin_power(np.abs(w1_hat)**2, L, range=(0,3),bins=150,dev=True)
+P2, b2, dev2 = bic.bin_power(np.abs(w2_hat)**2, L, range=(0,3),bins=150,dev=True)
+
+fig = plt.figure(1)
+fig.clf()
+plt.fill_between(b1, P1*(1-dev1), P1*(1+dev1), label='Supergen=1', color='b')
+plt.fill_between(b2, P2*(1-dev2), P2*(1+dev2), label='Supergen=2', color='g', alpha=0.5)
+ax = plt.gca()
+ax.set_xscale('log')
+plt.ylim(0.5,1.5)
+plt.xlim(1e-2,4)
+plt.axhline(1.0, color='red', lw=4.0)
+plt.legend()
+plt.show()

external/cosmotool/python_sample/ksz/__init__.py

@@ -0,0 +1,3 @@
+from .constants import * 
+from .gal_prof import KSZ_Profile, KSZ_Isothermal, KSZ_NFW
+

external/cosmotool/python_sample/ksz/constants.py

@@ -0,0 +1,36 @@
+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
+assert ksz_normalization < 0
+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

external/cosmotool/python_sample/ksz/gal_prof.py

@@ -0,0 +1,175 @@
+import numpy as np
+import numexpr as ne
+from .constants import *
+
+# -----------------------------------------------------------------------------
+# Generic profile generator
+# -----------------------------------------------------------------------------
+
+class KSZ_Profile(object):
+  R_star= 0.0 # 15 kpc/h 
+  L_gal0 = 10**(0.4*(tmpp_cat['Msun']-tmpp_cat['Mstar']))
+
+  def __init__(self,sculpt):
+    """Base class for KSZ profiles
+
+    Arguments:
+      sculpt (float): If negative, do not sculpt. If positive, there will be a 2d 
+                      suppression of the profile with a radius given by sculpt (in arcmins).
+    """
+    self.sculpt = sculpt * np.pi/180/60.
+    self.rGalaxy = 1.0
+
+  def evaluate_profile(self, r):
+    raise NotImplementedError("Abstract function")
+
+  def projected_profile(self, cos_theta,angularDistance):
+
+    idx_base = 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 = (self.R_star/angularDistance)**2
+
+    idx0 = np.where((tan_theta_2 < tan_theta_2_max))
+    idx = idx_base[idx0]
+    tan_theta_2 = tan_theta_2[idx0]
+    tan_theta = np.sqrt(tan_theta_2)
+
+    r = (tan_theta*angularDistance)
+
+    m,idx_mask = self.evaluate_profile(r)
+    idx_mask = idx[idx_mask]
+
+    idx_mask = np.append(idx_mask,idx[np.where(tan_theta_2<tan_theta_2_min)[0]])
+    if tan_theta_2.size > 0:
+      idx_mask = np.append(idx_mask,idx[tan_theta_2.argmin()])
+
+    if self.sculpt > 0:
+      theta = np.arctan(tan_theta)
+      cond = theta < self.sculpt
+      m[cond] *= (theta[cond]/self.sculpt)**2
+
+    return idx,idx_mask,m
+     
+
+# -----------------------------------------------------------------------------
+# Isothermal profile generator
+# -----------------------------------------------------------------------------
+
+class KSZ_Isothermal(KSZ_Profile):
+  sigma_FP=160e3  #m/s
+  R_innergal = 0.030
+
+  def __init__(self, Lgal, x, y=0.0, sculpt=-1):
+    """Support for Isothermal profile
+
+    Arguments:
+      Lgal (float): Galaxy luminosity in solar units
+      x (float): extent of halo in virial radius units
+
+    Keyword arguments:
+      y (float): Inner part where there is no halo
+      sculpt (float): If negative, do not sculpt. If positive, there will be a 2d 
+                      suppression of the profile with a radius given by sculpt (in arcmins).
+    """
+    
+    super(KSZ_Isothermal,self).__init__(sculpt)
+    
+    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):
+    pass
+
+  def evaluate_profile(self,r):
+    rho0, rGalaxy, rInner = self.rho0, self.rGalaxy, self.rInnerGalaxy
+    
+    D = {'rho0':rho0, 'rGalaxy':rGalaxy, 'rInner': rInner, 'Mpc':Mpc }
+    
+    Q = np.zeros(r.size)
+
+    cond = (r<=1e-10)
+    Q[cond] = rho0*2/Mpc * (rGalaxy-rInner)/(rGalaxy*rInner)
+
+    cond = (r>0)*(r <= rInner)
+    D['r'] = r[cond]
+    Q[cond] = ne.evaluate('rho0*2/(Mpc*r) * (arctan(sqrt( (rGalaxy/r)**2 -1 ))  - arctan(sqrt( (rInner/r)**2 - 1 )))',
+                local_dict=D)
+ 
+    cond = (r > rInner)*(r <= rGalaxy)
+    D['r'] = r[cond]
+    Q[cond] = ne.evaluate('rho0*2/(Mpc*r) * arctan(sqrt( (rGalaxy/r)**2 -1 ))',
+                local_dict=D)
+
+    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(self,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)
+
+

external/cosmotool/python_sample/ramsesToArray.py

@@ -0,0 +1,15 @@
+import numpy as np
+import sys
+import h5py as h5
+import cosmotool as ct
+
+s = ct.loadRamsesAll(sys.argv[1], int(sys.argv[2]), doublePrecision=True, loadVelocity=True)
+
+
+q = [p for p in s.getPositions()]
+q += [p for p in s.getVelocities()]
+q += [np.ones(q[0].size,dtype=q[0].dtype)]
+q = np.array(q)
+
+with h5.File("particles.h5", mode="w") as f:
+  f.create_dataset("particles", data=q.transpose())

external/cosmotool/python_sample/test_bispectrum.py

@@ -0,0 +1,28 @@
+import timeit
+import numpy as np
+import cosmotool as ct
+
+def myfun(N):
+  f=0.001
+  L = 1.0
+  d=np.random.normal(size=(N,)*3)   * np.sqrt(float(N))**3 / L**3
+  rho = d + f *(d*d - np.average(d*d))
+  delta = (L/N)**3
+  
+  B = ct.bispectrum(rho * delta, 1, N, fourier=False)
+  P = ct.powerspectrum(rho * delta, 1, N, fourier=False)
+  PP = P[1]/P[0] / L**3 
+  
+  x = PP[:,None,None] * PP[None,:,None] + PP[:,None,None]*PP[None,None,:] + PP[None,:,None]*PP[None,None,:]
+  
+  BB = B[1]/B[0] / L**3
+
+  y = BB/x
+
+  np.savez("bispec_%d.npz" % N, x=x, y=y, d=d,B_nt=B[0], B_r=B[1], P_n=P[0], P=P[1], BB=BB, rho=rho, PP=PP);
+  
+  
+#print( timeit.timeit('from __main__ import myfun; myfun(16)', number=1) )
+#print( timeit.timeit('from __main__ import myfun; myfun(24)', number=1) )
+print( timeit.timeit('from __main__ import myfun; myfun(32)', number=1) )
+#print( timeit.timeit('from __main__ import myfun; myfun(64)', number=1) )

external/cosmotool/python_sample/test_spheric_proj.py

@@ -0,0 +1,20 @@
+import cosmotool as ct
+import numpy as np
+import healpy as hp
+
+d = np.zeros((64,64,64), ct.DTYPE)
+
+d[32,32,32] = 1
+
+ii=np.arange(256)*64/256.-32
+xx = ii[:,None,None].repeat(256,axis=1).repeat(256,axis=2).reshape(256**3)
+yy = ii[None,:,None].repeat(256,axis=0).repeat(256,axis=2).reshape(256**3)
+zz = ii[None,None,:].repeat(256,axis=0).repeat(256,axis=1).reshape(256**3)
+
+d_high = ct.interp3d(xx, yy, zz, d, 64, periodic=True)
+d_high = d_high.reshape((256,256,256))
+
+proj0 = ct.spherical_projection(64, d, 0, 20, integrator_id=0, shifter=np.array([0.5,0.5,0.5]))
+proj1 = ct.spherical_projection(64, d, 0, 20, integrator_id=1)
+
+proj0_high = ct.spherical_projection(256, d_high, 0, 30, integrator_id=0, shifter=np.array([0.5,0.5,0.5]))