LSSTDESC / bayesian-pipelines-cosmology

Bayesian Cosmological Inference from Tomographic Maps
MIT License
4 stars 1 forks source link

Intrinsic Alignment #4

Open EiffL opened 2 years ago

EiffL commented 2 years ago

A first pass at adding an IA model would be to reimplement the simple NLA-based prescription @dlanzieri is using in her project. We can consider more complex models afterwards.

Please add here any suggestions for other models.

@dlanzieri could you describe here your strategy? maybe also with links to useful bits and pieces of code?

elts6570 commented 2 years ago
  1. Formulate a likelihood or likelihood-free inference? For the former see https://arxiv.org/abs/2112.04484. @EiffL any suggestion for the latter?
  2. Do we want to consider the NLA + environment-dependence model? Kindly provided by @dlanzieri :) https://ui.adsabs.harvard.edu/abs/2022MNRAS.509.3868H/abstract
  3. Joint inference of WL and IA: I have a block sampling MCMC framework to jointly sample the NLA IA amplitude and random noise given a tidal field on a galaxy-by-galaxy basis, but Denise's code seems more suitable since we are not operating at the individual-object basis.
EiffL commented 2 years ago

Here is the link to the few lines of code needed to implement simple IA https://github.com/DifferentiableUniverseInitiative/flowpm/blob/8f9415be8866f507fe40259c4dd3fc981583ee41/flowpm/NLA_IA.py#L69

EiffL commented 2 years ago

Here is some very very simple prototype code of how we could implement IA.

  # Definining IA ampliture
  Aia = numpyro.sample('Aia', dist.Uniform(0.0, 10.0))

  def resampling_func(entry, coord):
    r = entry['r']; a = entry['a']; p = entry['plane']
    dx = entry['dx']; dz = entry['dz']

    norm_factor = -Aia * jc.constants.C_1 * cosmo.Omega_m * jc.constants.rhocrit / jc.background.D1(
      cosmo, a)

    # Normalize density planes
    constant_factor = 3 / 2 * cosmo.Omega_m * (constants.H0 / constants.c)**2
    density_normalization = dz * r / a
    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")
    return im * norm_factor

  ia_maps = [ resampling_func(plane) for plane in density_planes]

  # To be refined
  ia_maps = [simps(lambda z: nz(z).reshape([-1,1,1]) * 
                              ia_maps(z), 0., 2.5, N=64) 
                      for nz in nzs_s_sys]

  # Generate convergence maps by integrating over nz and source planes
  convergence_maps = [simps(lambda z: nz(z).reshape([-1,1,1]) * 
                              convergence_Born(cosmo, density_planes, coords, z), 0., 2.5, N=64) 
                      for nz in nzs_s_sys]

  # Apply noise to the maps (this defines the likelihood)
  observed_maps = [numpyro.sample('kappa_%d'%i, 
                                  dist.Normal(k+k_ia, sigma_e/jnp.sqrt(10**2*galaxy_density))) # assumes pixel size of 10 arcmin
                   for i,(k,k_ia) in enumerate(zip(convergence_maps, ia_maps))]

  return observed_maps
EiffL commented 2 years ago

A good illustration of what the simulation volume looks like: https://arxiv.org/pdf/2010.09731.pdf

image

https://arxiv.org/pdf/2107.08041.pdf

jecampagne commented 1 year ago

I would be pleased to participate to a LFI/SBI code. I have nicely test the sbi library (written in (Py)Torch) with a JAX-based simulator running on GPU.

dlanzieri commented 1 year ago

Hey @nataliaporqueres this is the very simple prototype we saw last time. There are still a few things I have to solve, like figuring out what it’s happening for the auto power-spectrum II in the bin 3, but this is the general idea!

EiffL commented 1 year ago

Just for safe keeping, here is a demo notebook to compute Q/U maps in jax, useful for IA stuff

https://colab.research.google.com/drive/1upWB7r2MBnRirEQNJT8UN8KXb06gce_H?usp=sharing