Add Spherical lensing example

jaxpm/lensing.py

@@ -1,82 +1,181 @@
 import jax
 import jax.numpy as jnp
 import jax_cosmo
+import jax_cosmo as jc
 import jax_cosmo.constants as constants
 from jax.scipy.ndimage import map_coordinates
 
-from jaxpm.painting import cic_paint_2d
+from jaxpm.distributed import uniform_particles
+from jaxpm.painting import cic_paint, cic_paint_2d, cic_paint_dx
+from jaxpm.spherical import paint_spherical
 from jaxpm.utils import gaussian_smoothing
 
 
-def density_plane(positions,
-                  box_shape,
-                  center,
-                  width,
-                  plane_resolution,
-                  smoothing_sigma=None):
-    """ Extacts a density plane from the simulation
+def density_plane_fn(box_shape,
+                     box_size,
+                     density_plane_width,
+                     density_plane_npix,
+                     sharding=None):
+
+    def f(t, y, args):
+        positions = y[0]
+        cosmo = args
+        nx, ny, nz = box_shape
+
+        # Converts time t to comoving distance in voxel coordinates
+        w = density_plane_width / box_size[2] * box_shape[2]
+        center = jc.background.radial_comoving_distance(
+            cosmo, t) / box_size[2] * box_shape[2]
+        positions = uniform_particles(box_shape) + positions
+        xy = positions[..., :2]
+        d = positions[..., 2]
+
+        # Apply 2d periodic conditions
+        xy = jnp.mod(xy, nx)
+
+        # Rescaling positions to target grid
+        xy = xy / nx * density_plane_npix
+        # Selecting only particles that fall inside the volume of interest
+        weight = jnp.where((d > (center - w / 2)) & (d <= (center + w / 2)),
+                           1.0, 0.0)
+        # Painting density plane
+        zero_mesh = jnp.zeros([density_plane_npix, density_plane_npix])
+        # Apply sharding in order to recover sharding when taking gradients
+        if sharding is not None:
+            xy = jax.lax.with_sharding_constraint(xy, sharding)
+        # Apply CIC painting
+        density_plane = cic_paint_2d(zero_mesh, xy, weight)
+
+        # Apply density normalization
+        density_plane = density_plane / ((nx / density_plane_npix) *
+                                         (ny / density_plane_npix) * w)
+        return density_plane
+
+    return f
+
+
+def spherical_density_fn(box_shape,
+                         box_size,
+                         nside,
+                         fov,
+                         center_radec,
+                         observer_position,
+                         d_R,
+                         sharding=None):
+
+    def f(t, y, args):
+        positions = y[0]
+        nx, ny, nz = box_shape
+        bx, by, bz = box_size
+        cosmo = args
+        # Converts time t to comoving distance in voxel coordinates
+        w = d_R / box_size[2] * box_shape[2]
+        center = ((jc.background.radial_comoving_distance(cosmo, t)) / bz) * nz
+
+        # Apply sharding in order to recover sharding when taking gradients
+        if sharding is not None:
+            positions = jax.lax.with_sharding_constraint(positions, sharding)
+
+        density_mesh = cic_paint_dx(positions)
+        # Project to spherical map
+        spherical_map = paint_spherical(density_mesh, nside, fov, center_radec,
+                                        observer_position, box_size, center,
+                                        d_R)
+        return spherical_map
+
+    return f
+
+
+# ==========================================================
+# Weak Lensing Born Approximation
+# ==========================================================
+def convergence_Born(cosmo, density_planes, r, a, dx, dz, coords, z_source):
     """
-    nx, ny, nz = box_shape
-    xy = positions[..., :2]
-    d = positions[..., 2]
+    Compute Born-approximation lensing convergence maps.
 
-    # Apply 2d periodic conditions
-    xy = jnp.mod(xy, nx)
+    Parameters
+    ----------
+    cosmo : jc.Cosmology
+        Cosmology object.
+    density_planes : ndarray
+        3D array of lensing density planes [nx, ny, n_planes].
+    r, a : ndarray
+        Comoving distances and scale factors per plane.
+    dx : float
+        Pixel scale.
+    dz : float
+        Redshift bin width.
+    coords : ndarray
+        Angular coordinates grid [2, N, 2] in radians.
+    z_source : ndarray
+        Source redshifts.
 
-    # Rescaling positions to target grid
-    xy = xy / nx * plane_resolution
-
-    # Selecting only particles that fall inside the volume of interest
-    weight = jnp.where(
-        (d > (center - width / 2)) & (d <= (center + width / 2)), 1., 0.)
-    # Painting density plane
-    density_plane = cic_paint_2d(
-        jnp.zeros([plane_resolution, plane_resolution]), xy, weight)
-
-    # Apply density normalization
-    density_plane = density_plane / ((nx / plane_resolution) *
-                                     (ny / plane_resolution) * (width))
-
-    # Apply Gaussian smoothing if requested
-    if smoothing_sigma is not None:
-        density_plane = gaussian_smoothing(density_plane, smoothing_sigma)
-
-    return density_plane
-
-
-def convergence_Born(cosmo, density_planes, coords, z_source):
+    Returns
+    -------
+    convergence : ndarray
+        2D convergence map for each source redshift.
     """
-  Compute the Born convergence
-  Args:
-    cosmo: `Cosmology`, cosmology object.
-    density_planes: list of dictionaries (r, a, density_plane, dx, dz), lens planes to use
-    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.
-  """
-    # Compute constant prefactor:
     constant_factor = 3 / 2 * cosmo.Omega_m * (constants.H0 / constants.c)**2
     # Compute comoving distance of source galaxies
-    r_s = jax_cosmo.background.radial_comoving_distance(
-        cosmo, 1 / (1 + z_source))
+    r_s = jc.background.radial_comoving_distance(cosmo, 1 / (1 + z_source))
+    n_planes = len(r)
 
-    convergence = 0
-    for entry in density_planes:
-        r = entry['r']
-        a = entry['a']
-        p = entry['plane']
-        dx = entry['dx']
-        dz = entry['dz']
-        # Normalize density planes
-        density_normalization = dz * r / a
+    def scan_fn(carry, i):
+        density_planes, a, r = carry
+
+        p = density_planes[:, :, i]
+        density_normalization = dz * r[i] / a[i]
         p = (p - p.mean()) * constant_factor * density_normalization
 
         # Interpolate at the density plane coordinates
-        im = map_coordinates(p, coords * r / dx - 0.5, order=1, mode="wrap")
+        im = map_coordinates(p, coords * r[i] / dx - 0.5, order=1, mode="wrap")
 
-        convergence += im * jnp.clip(1. -
-                                     (r / r_s), 0, 1000).reshape([-1, 1, 1])
+        return carry, im * jnp.clip(1.0 -
+                                    (r[i] / r_s), 0, 1000).reshape([-1, 1, 1])
 
-    return convergence
+    _, convergence = jax.lax.scan(scan_fn, (density_planes, a, r),
+                                  jnp.arange(n_planes))
+    return convergence.sum(axis=0)
+
+
+def spherical_convergence_Born(cosmo, density_planes, r, a, nside, z_source):
+    """
+    Compute Born-approximation lensing convergence maps on a sphere.
+
+    Parameters
+    ----------
+    cosmo : jc.Cosmology
+        Cosmology object.
+    density_planes : ndarray
+        3D array of lensing density planes [n_planes, npix].
+    r, a : ndarray
+        Comoving distances and scale factors per plane.
+    nside : int
+        Healpix nside parameter.
+    z_source : ndarray
+        Source redshifts.
+
+    Returns
+    -------
+    convergence : ndarray
+        2D convergence map for each source redshift.
+    """
+    constant_factor = 3 / 2 * cosmo.Omega_m * (constants.H0 / constants.c)**2
+    # Compute comoving distance of source galaxies
+    r_s = jc.background.radial_comoving_distance(cosmo, 1 / (1 + z_source))
+    n_planes = len(r)
+
+    def scan_fn(carry, i):
+        density_planes, a, r = carry
+
+        p = density_planes[i, :]
+        density_normalization = r[i] / a[
+            i]  # This normalization needs to be checked
+        p = (p - p.mean()) * constant_factor * density_normalization
+
+        return carry, p * jnp.clip(1.0 -
+                                   (r[i] / r_s), 0, 1000).reshape([-1, 1])
+
+    _, convergence = jax.lax.scan(scan_fn, (density_planes, a, r),
+                                  jnp.arange(n_planes))
+    return convergence.sum(axis=0)