Merge pull request #62 from DifferentiableUniverseInitiative/lensing

Adds lensing lightcone computation using the Born approximation

.gitignore

@@ -107,3 +107,4 @@ venv.bak/
 *.ipynb
 models
 *.png
+snapshot*

examples/mesh_lpt_TPU.py

@@ -7,6 +7,7 @@
 
 import numpy as np
 import tensorflow.compat.v1 as tf
+
 tf.disable_v2_behavior()
 from tensorflow.python.lib.io import file_io
 import mesh_tensorflow as mtf

examples/mesh_lpt_local.py

@@ -2,6 +2,7 @@
 import os
 import math
 import tensorflow.compat.v1 as tf
+
 tf.disable_v2_behavior()
 import mesh_tensorflow as mtf
 
@@ -18,6 +19,7 @@
 from flowpm import linear_field, lpt_init, nbody, cic_paint
 from scipy.interpolate import InterpolatedUnivariateSpline as iuspline
 from matplotlib import pyplot as plt
+
 cosmology = Planck15
 
 tf.flags.DEFINE_integer("gpus_per_node", 8, "Number of GPU on each node")

examples/mesh_nbody_SLURM.py

@@ -5,6 +5,7 @@
 from matplotlib import pyplot as plt
 
 import tensorflow.compat.v1 as tf
+
 tf.disable_v2_behavior()
 import tensorflow_probability as tfp
 import mesh_tensorflow as mtf

examples/mesh_nbody_TPU.py

@@ -7,6 +7,7 @@
 
 import numpy as np
 import tensorflow.compat.v1 as tf
+
 tf.disable_v2_behavior()
 from tensorflow.python.lib.io import file_io
 import mesh_tensorflow as mtf

examples/pyramid_lpt_local.py

@@ -2,6 +2,7 @@
 import os
 import math
 import tensorflow.compat.v1 as tf
+
 tf.disable_v2_behavior()
 import mesh_tensorflow as mtf
 
@@ -18,6 +19,7 @@
 from flowpm import linear_field, lpt_init, nbody, cic_paint
 from scipy.interpolate import InterpolatedUnivariateSpline as iuspline
 from matplotlib import pyplot as plt
+
 cosmology = Planck15
 
 tf.flags.DEFINE_integer("gpus_per_node", 8, "Number of GPU on each node")

examples/pyramid_nbody_SLURM.py

@@ -2,6 +2,7 @@
 import os
 import math, time
 import tensorflow.compat.v1 as tf
+
 tf.disable_v2_behavior()
 import mesh_tensorflow as mtf
 import flowpm.mesh_ops as mpm

flowpm/__init__.py

@@ -1,3 +1,6 @@
 from .utils import cic_paint, cic_readout
 from .tfpm import nbody, linear_field, lpt_init
 import flowpm.cosmology as cosmology
+import flowpm.tfbackground as background
+import flowpm.io as io
+import flowpm.raytracing as raytracing

flowpm/experimental/lenstf.py

@@ -0,0 +1,121 @@
+import tensorflow as tf
+import tensorflow_probability as tfp
+from nbodykit import cosmology
+import numpy as np
+from flowpm.utils import r2c2d, c2r2d
+
+
+def generate(
+    di,
+    df,
+    ds,
+    boxsize,
+    boxsize2D,
+):
+  """ Returns a list of rotation matrices and shifts that are applied to the box
+    di : distance to inital redshift (before taking fastpm step)
+    df : distance to final redshift (after taking fastpm step)
+    ds : Source distance
+    boxsize: size of computational box
+    """
+  # Generate the possible rotation matrices
+  x = np.asarray([1, 0, 0], dtype=int)
+  y = np.asarray([0, 1, 0], dtype=int)
+  z = np.asarray([0, 0, 1], dtype=int)
+
+  # shuffle directions, only 90 deg rotations, makes a total of 6
+  M_matrices = [np.asarray([x,y,z],dtype=int), np.asarray([x,z,y],dtype=int),np.asarray([z,y,x],dtype=int),np.asarray([z,x,y],dtype=int), \
+                       np.asarray([y,x,z],dtype=int), np.asarray([y,z,x],dtype=int)]
+  # Generate possible xy shifts
+  I = np.zeros(3)
+  fac = 0.5
+  xyshifts = [
+      np.asarray([fac, fac, 0.], dtype=float),
+      np.asarray([-fac, fac, 0.], dtype=float),
+      np.asarray([-fac, -fac, 0.], dtype=float),
+      np.asarray([fac, -fac, 0.], dtype=float)
+  ]
+
+  # Find the maximum number of rotations needed in the z direction
+  vert_num = ds * boxsize2D / boxsize[-1]
+
+  print('rotations available: %d' % len(M_matrices))
+  print('rotations required: %d' % np.ceil(ds / boxsize[-1]))
+
+  try:
+    assert (len(M_matrices) * boxsize[-1] > ds)
+    print('sufficient number of rotations to fill lightcone.')
+  except:
+    print('insufficient number of rotations to fill the lightcone.')
+
+  if df > ds:
+    return 0, 0
+  else:
+    shift_ini = np.floor(max(di, ds) / boxsize[-1])
+    shift_end = np.floor(df / boxsize[-1])
+    M = []
+    if vert_num == 1:
+      for ii in np.arange(shift_end, shift_ini + 1, dtype=int):
+        M.append((M_matrices[ii % len(M_matrices)], I + ii * z))
+    elif vert_num == 2:
+      for ii in np.arange(shift_end, shift_ini + 1, dtype=int):
+        for jj in range(4):
+          M.append(
+              (M_matrices[ii % len(M_matrices)], I + ii * z + xyshifts[jj]))
+    else:
+      raise ValueError('vertical number of boxes must be 1 or 2, but is %d' %
+                       vert_num)
+
+    return M
+
+
+def rotate(x, M, boxsize, boxshift, name='rotate'):
+  """
+    rotates, shift, and separates particle coordinates into distance and xy position
+    x:        particle positions
+    M:        rotation matrix
+    boxshift: shift vector
+    """
+  with tf.name_scope(name):
+    y = tf.einsum('ij,kj->ki', M, x)
+    y = tf.add(y, boxsize * boxshift)
+    d = tf.gather(y, 2, axis=1, name='gather_z')
+    xy = tf.gather(y, (0, 1), axis=1, name='gather_xy')
+    return xy, d
+
+
+def z_chi(d, cosmo, name='z_chi'):
+  with tf.name_scope(name):
+    # redshift as a function of comsoving distance for underlying cosmology
+    z_int = np.logspace(-12, np.log10(1500), 40000)
+    chis = cosmo.comoving_distance(z_int)  #Mpc/h
+    z = tfp.math.interp_regular_1d_grid(d,
+                                        1e-12,
+                                        1.5e3,
+                                        tf.convert_to_tensor(chis,
+                                                             dtype='float'),
+                                        name='interpolation')
+    return z
+
+
+def wlen(d, ds, cosmo, boxsize, boxsize2D, mesh2D, name='efficiency_kernel'):
+  """
+    returns the correctly weighted lensing efficiency kernel
+    d:   particle distance (assuming parllel projection)
+    ds:  source redshift
+    """
+  with tf.name_scope(name):
+    # Find the redshift at ds
+    z = z_chi(d, cosmo)
+    # Find the number density of simulated particles
+    nbar = tf.shape(d) / boxsize[0]**3
+
+    #Find angular pixel area for projection
+    A = mesh2D[0]**2 / boxsize2D**2
+    columndens = tf.multiply(tf.pow(d, 2), float(
+        nbar *
+        A))  #particles/Volume*angular pixel area* distance^2 -> 1/L units
+    w = tf.divide(
+        tf.multiply(tf.multiply(tf.subtract(ds, d), tf.divide(d, ds)),
+                    (1. + z)), columndens)  #distance
+    return w

flowpm/raytracing.py

@@ -0,0 +1,163 @@
+#!/usr/bin/env python3
+# -*- coding: utf-8 -*-
+"""
+Created on Wed Feb  3 12:29:19 2021
+
+@author: Denise Lanzieri
+"""
+import numpy as np
+import tensorflow as tf
+import tensorflow_addons as tfa
+import flowpm
+import flowpm.constants as constants
+
+
+def rotation_matrices():
+  """ Generate the 3D rotation matrices along the 3 axis.
+
+  Returns:
+    list of 6 3x3 float rotation matrices 
+  """
+  x = np.asarray([1, 0, 0], dtype=np.float32)
+  y = np.asarray([0, 1, 0], dtype=np.float32)
+  z = np.asarray([0, 0, 1], dtype=np.float32)
+  M_matrices = [
+      np.asarray([x, y, z], dtype=np.float32),
+      np.asarray([x, z, y], dtype=np.float32),
+      np.asarray([z, y, x], dtype=np.float32),
+      np.asarray([z, x, y], dtype=np.float32),
+      np.asarray([y, x, z], dtype=np.float32),
+      np.asarray([y, z, x], dtype=np.float32)
+  ]
+
+  return M_matrices
+
+
+def random_2d_shift():
+  """ Generates a random xy shift
+
+  Returns:
+    float vector of shape [3] for x,y,z random shifts between 0 and 1.
+  """
+  I = np.zeros(3)
+  facx = np.random.uniform(0, 1)
+  facy = np.random.uniform(0, 1)
+  xyshifts = [np.asarray([facx, facy, 0], dtype=np.float32)]
+  return I + xyshifts
+
+
+def density_plane(state,
+                  nc,
+                  center,
+                  width,
+                  plane_resolution,
+                  rotation=None,
+                  shift=None,
+                  name='density_plane'):
+  """ Extract a slice from the input state vector and
+  project it as a density plane.
+
+  Args:
+    state: `Tensor`, input state tensor.
+    nc: int or list of int, size of simulation volume
+    center: float, center of the slice along the z axis in voxel coordinates
+    width: float, width of the slice in voxel coordinates
+    plane_size: int, number of pixels of the density plane. 
+    rotation: 3x3 float tensor, 3D rotation matrix to apply to the cube before cutting the plane
+    shift: float tensor of shape [3], 3D shift to apply to the cube before cutting the plane
+    name: `string`, name of the operation.
+
+  Returns:
+    `Tensor` of shape [batch_size, plane_size, plane_size], of projected density plane.
+  """
+  with tf.name_scope(name):
+    state = tf.convert_to_tensor(state, name="state")
+    if isinstance(nc, int):
+      nc = [nc, nc, nc]
+    nx, ny, nz = nc
+    pos = state[0]
+
+    shape = tf.shape(pos)
+    batch_size = shape[0]
+
+    # Apply optional shifts and rotations to the cube
+    if rotation is not None:
+      pos = tf.einsum('ij,bkj->bki', rotation, pos)
+    if shift is not None:
+      pos = pos + [nx, ny, nz] * shift
+
+    xy = pos[..., :2]
+    d = pos[..., 2]
+
+    # Apply 2d periodic conditions
+    xy = tf.math.mod(xy, nx)
+
+    # Rescaling positions to target grid
+    xy = xy / nx * plane_resolution
+
+    # Selecting only particles that fall inside the volume of interest
+    mask = (d > (center - width / 2)) & (d <= (center + width / 2))
+
+    # Painting density plane
+    density_plane = tf.zeros([batch_size, plane_resolution, plane_resolution])
+    density_plane = flowpm.utils.cic_paint_2d(density_plane, xy, mask=mask)
+
+    # Apply density normalization
+    density_plane = density_plane / ((nx / plane_resolution) *
+                                     (ny / plane_resolution) * (width))
+
+    return density_plane
+
+
+def convergenceBorn(cosmo,
+                    lensplanes,
+                    dx,
+                    dz,
+                    coords,
+                    z_source,
+                    name="convergenceBorn"):
+  """
+  Compute the Born–approximated convergence
+
+  Args:
+    cosmo: `Cosmology`, cosmology object.
+    lensplanes: list of tuples (r, a, lens_plane), lens planes to use 
+    dx: float, transverse pixel resolution of the lensplanes [Mpc/h]
+    dz: float, width of the lensplanes [Mpc/h]
+    coords: a 3-D array of angular coordinates in radians of N points with shape [batch, N, 2].
+    z_source: 1-D `Tensor` of source redshifts with shape [Nz] .
+    name: `string`, name of the operation.
+
+  Returns:
+    `Tensor` of shape [batch_size, N, Nz], of convergence values.
+  """
+  with tf.name_scope(name):
+    coords = tf.convert_to_tensor(coords, dtype=tf.float32)
+
+    # Compute constant prefactor:
+    constant_factor = 3 / 2 * cosmo.Omega_m * (constants.H0 / constants.c)**2
+    # Compute comoving distance of source galaxies
+    r_s = flowpm.background.rad_comoving_distance(cosmo, 1 / (1 + z_source))
+
+    convergence = 0
+    for r, a, p in lensplanes:
+      density_normalization = dz * r / a
+      p = (p - tf.reduce_mean(p, axis=[1, 2], keepdims=True)
+           ) * constant_factor * density_normalization
+      c = coords * r / dx
+
+      # Applying periodic conditions on lensplane
+      shape = tf.shape(p)
+      c = tf.math.mod(c, tf.cast(shape[1], tf.float32))
+
+      # Shifting pixel center convention
+      c = tf.expand_dims(c, axis=0) - 0.5
+
+      im = tfa.image.interpolate_bilinear(tf.expand_dims(p, -1),
+                                          c,
+                                          indexing='xy')
+
+      convergence += im * tf.reshape(tf.clip_by_value(1. - (r / r_s), 0, 1000),
+                                     [1, 1, -1])
+
+    return convergence

flowpm/tfbackground.py

@@ -616,7 +616,7 @@ def maybe_compute_ODE(cosmo, log10_amin=-2, steps=256):
     # been computed
     cache = cosmo._workspace['cache_ODE']
     # Checking that the stored ODE results have the right lenght
-    assert len(cache['a']) == steps
+    assert cache['a'].shape[0] == steps
   else:
     # Otherwise, we compute it now, and save the results for later
     a = tf.convert_to_tensor(np.logspace(log10_amin, 0., steps),