# 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 utils import get_nsims def open_galaxy_positions(survey_name, comm): """ 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. 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 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) 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) 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: fout = paths.field_interpolated( parser_args.survey, parser_args.simname, nsim, parser_args.kind, parser_args.MAS, parser_args.grid) # The survey above had some cuts, however for compatibility we want # the same shape as the `uncut` survey 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"], # noqa help="Simulation name") parser.add_argument("--survey", type=str, required=True, choices=["SDSS", "SDSSxALFALFA", "GW170817"], 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.") args = parser.parse_args() paths = csiborgtools.read.Paths(**csiborgtools.paths_glamdring) nsims = get_nsims(args, paths) pos = open_galaxy_positions(args.survey, MPI.COMM_WORLD) def _main(nsim): main(nsim, args, pos, verbose=MPI.COMM_WORLD.Get_size() == 1) work_delegation(_main, nsims, MPI.COMM_WORLD)