csiborgtools/scripts/field_sample.py

# Copyright (C) 2023 Richard Stiskalek
# This program is free software; you can redistribute it and/or modify it
# under the terms of the GNU General Public License as published by the
# Free Software Foundation; either version 3 of the License, or (at your
# option) any later version.
#
# This program is distributed in the hope that it will be useful, but
# WITHOUT ANY WARRANTY; without even the implied warranty of
# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the GNU General
# Public License for more details.
#
# You should have received a copy of the GNU General Public License along
# with this program; if not, write to the Free Software Foundation, Inc.,
# 51 Franklin Street, Fifth Floor, Boston, MA  02110-1301, USA.
"""
Script to sample a CSiBORG field at galaxy positions and save the result.
Supports additional smoothing of the field as well.
"""
from argparse import ArgumentParser
from os.path import join

import csiborgtools
import numpy
from astropy.cosmology import FlatLambdaCDM
from h5py import File
from mpi4py import MPI
from taskmaster import work_delegation
from tqdm import tqdm
from numba import jit

from utils import get_nsims


@jit(nopython=True, fastmath=True, boundscheck=False)
def scatter_along_radial_direction(pos, scatter, boxsize):
    """
    Scatter galaxy positions along the radial direction. Enforces that the
    radial position is always on the same side of the box and that the galaxy
    is still inside the box.

    Parameters
    ----------
    pos : 2-dimensional array
        Galaxy positions in the form of (distance, RA, DEC).
    scatter : float
        Scatter to add to the radial positions of galaxies in same units as
        `distance` (Mpc / h).
    boxsize : float
        Box size in `Mpc / h`.
    """
    pos_new = numpy.copy(pos)

    for i in range(len(pos)):
        r0, ra, dec = pos[i]
        # Convert to radians
        ra *= numpy.pi / 180
        dec *= numpy.pi / 180

        # Convert to normalized Cartesian coordinates
        xnorm = numpy.cos(dec) * numpy.cos(ra)
        ynorm = numpy.cos(dec) * numpy.sin(ra)
        znorm = numpy.sin(dec)

        while True:
            rnew = numpy.random.normal(r0, scatter)
            if rnew < 0:
                continue

            xnew = rnew * xnorm + boxsize / 2
            ynew = rnew * ynorm + boxsize / 2
            znew = rnew * znorm + boxsize / 2

            if 0 <= xnew < boxsize and 0 <= ynew < boxsize and 0 <= znew < boxsize:  # noqa
                pos_new[i, 0] = rnew
                break

    return pos_new


def open_galaxy_positions(survey_name, comm, scatter=None):
    """
    Load the survey's galaxy positions , broadcasting them to all ranks.

    Parameters
    ----------
    survey_name : str
        Name of the survey.
    comm : mpi4py.MPI.Comm
        MPI communicator.
    scatter : float
        Scatter to add to the radial positions of galaxies, supportted only in
        TNG300-1.

    Returns
    -------
    pos : 2-dimensional array
        Galaxy positions in the form of (distance, RA, DEC).
    """
    rank, size = comm.Get_rank(), comm.Get_size()

    if rank == 0:
        if survey_name == "SDSS":
            survey = csiborgtools.SDSS()()
            pos = numpy.vstack([survey["DIST"],
                                survey["RA"],
                                survey["DEC"]],
                               ).T
            pos = pos.astype(numpy.float32)
        elif survey_name == "SDSSxALFALFA":
            survey = csiborgtools.SDSSxALFALFA()()
            pos = numpy.vstack([survey["DIST"],
                                survey["RA_1"],
                                survey["DEC_1"]],
                               ).T
            pos = pos.astype(numpy.float32)
        elif survey_name == "GW170817":
            samples = File("/mnt/extraspace/rstiskalek/GWLSS/H1L1V1-EXTRACT_POSTERIOR_GW170817-1187008600-400.hdf", 'r')["samples"]  # noqa
            cosmo = FlatLambdaCDM(H0=100, Om0=0.3175)
            pos = numpy.vstack([
                cosmo.comoving_distance(samples["redshift"][:]).value,
                samples["ra"][:] * 180 / numpy.pi,
                samples["dec"][:] * 180 / numpy.pi],
                               ).T
        elif survey_name == "TNG300-1":
            with File("/mnt/extraspace/rstiskalek/TNG300-1/postprocessing/subhalo_catalogue_099.hdf5", 'r') as f:  # noqa
                pos = numpy.vstack([f["SubhaloPos"][:, 0],
                                    f["SubhaloPos"][:, 1],
                                    f["SubhaloPos"][:, 2]],
                                   ).T
                boxsize = csiborgtools.simname2boxsize("TNG300-1")
                pos -= boxsize / 2
                pos = csiborgtools.cartesian_to_radec(pos)
                if scatter is not None:
                    if scatter < 0:
                        raise ValueError("Scatter must be positive.")
                    if scatter > 0:
                        print(f"Adding scatter of {scatter} Mpc / h.",
                              flush=True)
                        pos = scatter_along_radial_direction(pos, scatter,
                                                             boxsize)

        else:
            raise NotImplementedError(f"Survey `{survey_name}` not "
                                      "implemented.")
    else:
        pos = None

    comm.Barrier()

    if size > 1:
        pos = comm.bcast(pos, root=0)

    return pos


def evaluate_field(field, pos, boxsize, smooth_scales, verbose=True):
    """
    Evaluate the field at the given galaxy positions.

    Parameters
    ----------
    field : 3-dimensional array
        Cartesian field to be evaluated.
    pos : 2-dimensional array
        Galaxy positions in the form of (distance, RA, DEC).
    boxsize : float
        Box size in `Mpc / h`.
    smooth_scales : list
        List of smoothing scales in `Mpc / h`.
    verbose : bool
        Verbosity flag.

    Returns
    -------
    val : 2-dimensional array
        Evaluated field.
    """
    mpc2box = 1. / boxsize
    val = numpy.full((pos.shape[0], len(smooth_scales)), numpy.nan,
                     dtype=field.dtype)

    for i, scale in enumerate(tqdm(smooth_scales, desc="Smoothing",
                                   disable=not verbose)):
        if scale > 0:
            field_smoothed = csiborgtools.field.smoothen_field(
                field, scale * mpc2box, boxsize=1, make_copy=True)
        else:
            field_smoothed = numpy.copy(field)
        print("Going to evaluate the field....")
        val[:, i] = csiborgtools.field.evaluate_sky(
            field_smoothed, pos=pos, mpc2box=mpc2box)

    return val


def match_to_no_selection(val, parser_args):
    """
    Match the shape of the evaluated field to the shape of the survey without
    any masking. Missing values are filled with `numpy.nan`.

    Parameters
    ----------
    val : n-dimensional array
        Evaluated field.
    parser_args : argparse.Namespace
        Command line arguments.

    Returns
    -------
    n-dimensional array
    """
    if parser_args.survey == "SDSS":
        survey = csiborgtools.SDSS()()
    elif parser_args.survey == "SDSSxALFALFA":
        survey = csiborgtools.SDSSxALFALFA()()
    else:
        raise NotImplementedError(
            f"Survey `{parser_args.survey}` not implemented for matching to no selection.")  # noqa

    return csiborgtools.read.match_array_to_no_masking(val, survey)


def main(nsim, parser_args, pos, verbose):
    """
    Main function to load the field, interpolate (and smooth it) it and save
    the results to the disk.

    Parameters
    ----------
    nsim : int
        IC realisation.
    parser_args : argparse.Namespace
        Command line arguments.
    pos : numpy.ndarray
        Galaxy coordinates in the form of (distance, RA, DEC) where to evaluate
        the field.
    verbose : bool
        Verbosity flag.

    Returns
    -------
    None
    """
    paths = csiborgtools.read.Paths(**csiborgtools.paths_glamdring)
    boxsize = csiborgtools.simname2boxsize(parser_args.simname)

    # Get the appropriate field loader
    if parser_args.simname == "csiborg1":
        freader = csiborgtools.read.CSiBORG1Field(nsim)
    elif "csiborg2" in parser_args.simname:
        kind = parser_args.simname.split("_")[-1]
        freader = csiborgtools.read.CSiBORG2Field(nsim, kind)
    elif parser_args.simname == "TNG300-1":
        freader = csiborgtools.read.TNG300_1Field()
    else:
        raise NotImplementedError(f"Simulation `{parser_args.simname}` is not supported.")  # noqa

    # Get the appropriate field
    if parser_args.kind == "density":
        field = freader.density_field(parser_args.MAS, parser_args.grid)
    else:
        raise NotImplementedError(f"Field `{parser_args.kind}` is not supported.")  # noqa

    val = evaluate_field(field, pos, boxsize, parser_args.smooth_scales,
                         verbose=verbose)

    if parser_args.survey == "GW170817":
        fout = join(
            "/mnt/extraspace/rstiskalek/GWLSS/",
            f"{parser_args.kind}_{parser_args.MAS}_{parser_args.grid}_{nsim}_H1L1V1-EXTRACT_POSTERIOR_GW170817-1187008600-400.npz")  # noqa
    else:
        if parser_args.simname == "TNG300-1":
            scatter = parser_args.scatter
        else:
            scatter = None

        fout = paths.field_interpolated(
            parser_args.survey, parser_args.simname, nsim, parser_args.kind,
            parser_args.MAS, parser_args.grid, scatter)

        # The survey above had some cuts, however for compatibility we want
        # the same shape as the `uncut` survey
        if parser_args.survey != "TNG300-1":
            val = match_to_no_selection(val, parser_args)

    if verbose:
        print(f"Saving to ... `{fout}`.")

    numpy.savez(fout, val=val, smooth_scales=parser_args.smooth_scales)


if __name__ == "__main__":
    parser = ArgumentParser()
    parser.add_argument("--nsims", type=int, nargs="+", default=None,
                        help="IC realisations. If `-1` processes all.")
    parser.add_argument("--simname", type=str, default="csiborg1",
                        choices=["csiborg1", "csiborg2_main", "csiborg2_random", "csiborg2_varysmall", "TNG300-1"],  # noqa
                        help="Simulation name")
    parser.add_argument("--survey", type=str, required=True,
                        choices=["SDSS", "SDSSxALFALFA", "GW170817", "TNG300-1"],  # noqa
                        help="Galaxy survey")
    parser.add_argument("--smooth_scales", type=float, nargs="+", default=None,
                        help="Smoothing scales in Mpc / h.")
    parser.add_argument("--kind", type=str,
                        choices=["density", "rspdensity", "velocity", "radvel",
                                 "potential"],
                        help="What field to interpolate.")
    parser.add_argument("--MAS", type=str,
                        choices=["NGP", "CIC", "TSC", "PCS", "SPH"],
                        help="Mass assignment scheme.")
    parser.add_argument("--grid", type=int, help="Grid resolution.")
    parser.add_argument("--scatter", type=float, default=None,
                        help="Scatter to add to the radial positions of galaxies, supportted only in TNG300-1.")  # noqa
    args = parser.parse_args()

    if args.simname == "TNG300-1" and args.survey != "TNG300-1":
        raise ValueError("TNG300-1 simulation is only supported for TNG300-1 survey.")  # noqa

    paths = csiborgtools.read.Paths(**csiborgtools.paths_glamdring)
    if args.simname == "TNG300-1":
        nsims = [0]
    else:
        nsims = get_nsims(args, paths)

    pos = open_galaxy_positions(args.survey, MPI.COMM_WORLD, args.scatter)

    def _main(nsim):
        main(nsim, args, pos, verbose=MPI.COMM_WORLD.Get_size() == 1)

    work_delegation(_main, nsims, MPI.COMM_WORLD)