py: check VELOCIraptor density grids totals against theory

check_density_totals.py

#!/usr/bin/env python

"""Test the Genesis grids totals against cosmology.  This is to help K Ren (see email 2020-03-17)."""

__author__ = "Simon Mutch"
__date__ = "2020-03-18"

from pathlib import Path
import logging

import click
import numpy as np
import h5py as h5
from astropy import cosmology
from astropy import units as U
import coloredlogs
import dask.array as da

logger = logging.getLogger()
coloredlogs.install()


@click.command()
@click.argument("grids_path", type=click.Path(exists=True))
@click.argument("snapshot", type=click.INT)
@click.option("loglevel", "-l", type=click.STRING, default="WARNING")
def main(grids_path: str, snapshot: int, loglevel: str = "WARNING"):
    logger.setLevel(loglevel.upper())
    grids_path = Path(grids_path)

    if not grids_path.is_dir():
        raise ValueError("Expect grids_path to be directory!")

    flist = sorted(
        list(grids_path.glob(f"snapshot_{snapshot:03d}.den.[0-9]*")), key=lambda p: int(p.name.split(".")[-1])
    )

    with h5.File(flist[0], "r") as fp:
        dim = fp.attrs["Ngrid_X"]
        #  Ode0 = fp.attrs['Omega_Lambda']
        Ob0 = fp.attrs["Omega_b"]
        Odm0 = fp.attrs["Omega_cdm"]
        hubble_h = fp.attrs["h"]

    logger.info(f"dim = {dim}")
    logger.info(f"Ob0 = {Ob0}")
    logger.info(f"Odm0 = {Odm0}")
    logger.info(f"hubble_h = {hubble_h}")

    ix = 0
    slices = []
    fps = []  # file pointers
    for file in flist:
        fp = h5.File(file, "r")
        fps.append(fp)
        local_nx = fp.attrs["Local_nx"]
        slices.append(da.from_array(fp["Density"]))
        ix += local_nx

    grid = da.concatenate(slices)
    grid.map_blocks(np.sort)
    grid_total = (grid.sum() / dim ** 3).compute() * U.Unit('1e10 Msun Mpc-3')
    print(f"grid total = {grid_total}")

    for fp in fps:
        fp.close()

    cosmo = cosmology.FlatLambdaCDM(hubble_h * 100.0, Odm0 + Ob0, Ob0=Ob0, name="Genesis")

    expected_total = cosmo.critical_density0.to(grid_total.unit) * cosmo.Om0
    print(f"expected total = {expected_total}")


if __name__ == "__main__":
    main()