Extend charm to resolve subgrid positions and velocities

[maho3]

Extend charm to resolve subgrid positions and velocities

[tlmakinen]

TLDR: modelling halo positions in [x,y,z] as a conditional distribution from existing CHARM pipeline:

$$p(M, [x,y,z]) = \int p([x,y,z] | M_{\rm halos}, N_{\rm halos}) p(M_{\rm halos}, N_{\rm halos} | \rm{CNN+charm\ inputs}) d\ {\rm inputs}\ d M_{\rm halos} d N_{\rm halos} $$

For now, we're fixing M_halos, N_halos to the simulation values, but eventually these will be sampled from CHARM

[shivampcosmo]

experimenting with velocities in CHARM, by implementing a separate autoregressive network that predicts the halo velocity. This is first test that assumes halo position and masses are true, and just learns the correction for velocity on top of the interpolated velocity from FASTPM field.

[tlmakinen]

how many transformations / descriptive is the velocity autoregressive network ? is it the same architecture as is used for the positions ?

[shivampcosmo]

Yes, more or less. It is stack of two normalizing flows going from Gaussian to velocity distribution (difference between FastPM and true).

[shivampcosmo]

first attempt at generalizing for cosmologies. Seems to work much better than getting the velocities from FastPM, but currently only ~10% level.

[shivampcosmo]

CHARM updated training for both masses and velocities. Now the network goes to smaller halo masses (M > 5e12 Msun/h) and also trains much better on multiple GPUs. Also the code is stand-alone installable and pushed here : https://github.com/shivampcosmo/CHARM . Integration with current ltu-cmass pipeline in progress.

[tlmakinen]

some updated objectives after conversations with @shivampcosmo @maho3:

• look at modelling positions on nodes within $ 8^3$ voxels using a cheap fully-connected GNN layer.
• TODO: look at (true) displacement distributions from COM in each voxel
• heavy-tailed studentT MDN might be the best parameterization for the loss $- \ln p(x, \hat{x})$

[tlmakinen]

Updates from LtU meeting:

Predicting positions within voxels might be best done as a residual learning problem. Looking at the difference in position between the heaviest halo $x_0$ and the next heaviest halo $x_1$ (and so on), we see a zero-centered heavy-tailed distribution.

Would be good to model this with an MDN neural studentT e.g. following https://arxiv.org/pdf/2006.06599