Adds basic infrastructure with prototype code for PM pipeline #9
Conversation
…-pipelines-cosmology into u/EiffL/infrastructure
Co-authored-by: Alexandre Boucaud <[email protected]>
bpcosmo/pm.py
Outdated
|
|
||
| # Extract the lensplanes | ||
| density_planes = [] | ||
| for i in range(len(a_center)): |
There was a problem hiding this comment.
| for i in range(len(a_center)): | |
| for i in range(n_lens): |
There was a problem hiding this comment.
line 123: is it the same 0.01 as in line 84? If so, maybe worth defining a new variable?
bpcosmo/pm.py
Outdated
| dz = density_plane_width | ||
| plane = density_plane(res[0][::-1][i], | ||
| nc, | ||
| (i+0.5)*density_plane_width/box_size[-1]*nc[-1], |
There was a problem hiding this comment.
What is the "+0.5"? Some corner assignment scheme?
| density_plane_smoothing: Gaussian scale of plane smoothing | ||
| box_size: [sx,sy,sz] size in Mpc/h of the simulation volume | ||
| nc: number of particles/voxels in the PM scheme | ||
| neural_spline_params: optional parameters for neural correction of PM scheme |
There was a problem hiding this comment.
Do we need a reference cited for the neural spline?
There was a problem hiding this comment.
yes indeed, unfortunately we don't have one yet ^^' this is some very recent trick that @dlanzieri is still working on. We might have an arxiv ref in few weeks
bpcosmo/pm.py
Outdated
| neural_spline_params: optional parameters for neural correction of PM scheme | ||
| Returns: | ||
| list of [r, a, plane], slices through the lightcone along with their | ||
| comoving distance and scale factors. |
There was a problem hiding this comment.
Being pedantic here, but maybe it's worth clarifying: "list of [r, a, plane], slices through the lightcone along with their comoving distance (r) and scale factors (a), with 'plane' being the index of the density slice.", or something along those lines.
There was a problem hiding this comment.
plane would range from 0 to nc[0]?
There was a problem hiding this comment.
yep thank you, added more clarification
bpcosmo/pm.py
Outdated
| a_center = jc.background.a_of_chi(cosmology, r_center) | ||
|
|
||
| # Create a small function to generate the matter power spectrum | ||
| k = jnp.logspace(-4, 1, 128) |
There was a problem hiding this comment.
I've upped to to 256 to be safe. These are points used to interpolate the matter power spectrum, needed when building initial conditions. It shouldn't be very sensitive to this parameter choice.
bpcosmo/pm.py
Outdated
|
|
||
| # Initial displacement | ||
| cosmology._workspace = {} # FIX ME: this a temporary fix | ||
| dx, p, f = lpt(cosmology, initial_conditions, particles, 0.01) |
There was a problem hiding this comment.
where does 0.01 come from?
There was a problem hiding this comment.
Ah yes, 0.01 is the the initial scale factor at which we initialize the simulation.
There was a problem hiding this comment.
I added explicitly the keyword
bpcosmo/pm.py
Outdated
| @jax.jit | ||
| def neural_nbody_ode(state, a, cosmo, params): | ||
| """ | ||
| state is a tuple (position, velocities) |
There was a problem hiding this comment.
hummm yeah good question... I think what I added there is correct
bpcosmo/pm.py
Outdated
| # Computes gravitational potential | ||
| pot_k = delta_k * laplace_kernel(kvec) * longrange_kernel(kvec, r_split=0) | ||
|
|
||
| # Apply a correction filter |
There was a problem hiding this comment.
why do we need this correction? what filter is this?
There was a problem hiding this comment.
This is the magic sauce from this spline correction function that @dlanzieri is building. I hope we'll soon have a paper to describe it.
There was a problem hiding this comment.
it corrects power on small scales, to effectively increase the resolution of the simulation
|
I've tried to address all your comments @elts6570, let me know what you think and if you approve this PR :-) |
|
@EiffL looks great! looking forward to the magic splines paper! :) |
This PR adds some basic infra, like adds a setup.py. It also adds some code based on jaxpm to generate density planes from an nbody simulation, which can be used in a forward model.
A demo notebook is provided here
It generates DC2-like maps:
This provides a few first elements towards solving #8