Open EiffL opened 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
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
A good illustration of what the simulation volume looks like: https://arxiv.org/pdf/2010.09731.pdf
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.
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!
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
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?