mirror of
https://github.com/Richard-Sti/csiborgtools.git
synced 2024-10-18 22:25:05 +02:00
225 lines
8.5 KiB
Python
225 lines
8.5 KiB
Python
# 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.
|
|
"""
|
|
Sample a CSiBORG field at galaxy positions and save the result to disk.
|
|
"""
|
|
from argparse import ArgumentParser
|
|
from distutils.util import strtobool
|
|
from os.path import join
|
|
|
|
import numpy
|
|
from astropy.cosmology import FlatLambdaCDM
|
|
from h5py import File
|
|
from mpi4py import MPI
|
|
from taskmaster import work_delegation
|
|
from tqdm import tqdm
|
|
|
|
import csiborgtools
|
|
from utils import get_nsims
|
|
|
|
MPC2BOX = 1 / 677.7
|
|
|
|
|
|
def steps(cls, survey_name):
|
|
"""Make a list of selection criteria to apply to a survey."""
|
|
if survey_name == "SDSS":
|
|
return [
|
|
# (lambda x: cls[x], ("IN_DR7_LSS",)),
|
|
# (lambda x: cls[x] < 17.6, ("ELPETRO_APPMAG_r", )),
|
|
(lambda x: cls[x] < 155.5, ("DIST", ))
|
|
]
|
|
else:
|
|
raise NotImplementedError(f"Survey `{survey_name}` not implemented.")
|
|
|
|
|
|
def open_galaxy_positions(survey_name, comm):
|
|
"""
|
|
Load the survey galaxy positions and indices, broadcasting them to all
|
|
ranks.
|
|
"""
|
|
rank, size = comm.Get_rank(), comm.Get_size()
|
|
|
|
if rank == 0:
|
|
if survey_name == "SDSS":
|
|
survey = csiborgtools.SDSS()()
|
|
pos = numpy.vstack([survey["DIST_UNCORRECTED"],
|
|
survey["RA"],
|
|
survey["DEC"]],
|
|
).T
|
|
pos = pos.astype(numpy.float32)
|
|
indxs = survey["INDEX"]
|
|
if survey_name == "SDSSxALFALFA":
|
|
survey = csiborgtools.SDSSxALFALFA()()
|
|
pos = numpy.vstack([survey["DIST_UNCORRECTED"],
|
|
survey["RA_1"],
|
|
survey["DEC_1"]],
|
|
).T
|
|
pos = pos.astype(numpy.float32)
|
|
indxs = survey["INDEX"]
|
|
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
|
|
indxs = numpy.arange(pos.shape[0])
|
|
else:
|
|
raise NotImplementedError(f"Survey `{survey_name}` not "
|
|
"implemented.")
|
|
else:
|
|
pos = None
|
|
indxs = None
|
|
|
|
comm.Barrier()
|
|
|
|
if size > 1:
|
|
pos = comm.bcast(pos, root=0)
|
|
indxs = comm.bcast(indxs, root=0)
|
|
|
|
return pos, indxs
|
|
|
|
|
|
def evaluate_field(field, pos, nrand, smooth_scales=None, seed=42,
|
|
verbose=True):
|
|
"""
|
|
Evaluate the field at the given sky positions. Additionally, evaluate the
|
|
field at `nrand` random positions.
|
|
"""
|
|
if smooth_scales is None:
|
|
smooth_scales = [0.]
|
|
|
|
nsample = pos.shape[0]
|
|
nsmooth = len(smooth_scales)
|
|
|
|
val = numpy.full((nsample, nsmooth), numpy.nan, dtype=field.dtype)
|
|
if nrand > 0:
|
|
rand_val = numpy.full((nsample, nsmooth, nrand), numpy.nan,
|
|
dtype=field.dtype)
|
|
else:
|
|
rand_val = None
|
|
|
|
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)
|
|
|
|
if nrand == 0:
|
|
continue
|
|
|
|
for j in range(nrand):
|
|
gen = numpy.random.default_rng(seed + j)
|
|
pos_rand = numpy.vstack([
|
|
gen.permutation(pos[:, 0]),
|
|
gen.uniform(0, 360, nsample),
|
|
90 - numpy.rad2deg(numpy.arccos(gen.uniform(-1, 1, nsample))),
|
|
]).T
|
|
|
|
rand_val[:, i, j] = csiborgtools.field.evaluate_sky(
|
|
field_smoothed, pos=pos_rand, mpc2box=MPC2BOX)
|
|
|
|
return val, rand_val, smooth_scales
|
|
|
|
|
|
def match_to_no_selection(val, rand_val, parser_args):
|
|
if parser_args.survey == "SDSSxALFALFA":
|
|
survey = csiborgtools.SDSSxALFALFA()(apply_selection=False)
|
|
else:
|
|
raise NotImplementedError(
|
|
f"Survey `{parser_args.survey}` not implemented for matching to no selection.") # noqa
|
|
|
|
if val is not None:
|
|
val = csiborgtools.read.match_array_to_no_masking(val, survey)
|
|
if rand_val is not None:
|
|
rand_val = csiborgtools.read.match_array_to_no_masking(rand_val,
|
|
survey)
|
|
|
|
return val, rand_val
|
|
|
|
|
|
def main(nsim, parser_args, pos, indxs, paths, verbose):
|
|
"""Load the field, interpolate it and save it to disk."""
|
|
fpath_field = paths.field(parser_args.kind, parser_args.MAS,
|
|
parser_args.grid, nsim, parser_args.in_rsp)
|
|
field = numpy.load(fpath_field)
|
|
|
|
val, rand_val, smooth_scales = evaluate_field(
|
|
field, pos, nrand=parser_args.nrand,
|
|
smooth_scales=parser_args.smooth_scales, verbose=verbose)
|
|
|
|
if parser_args.survey == "GW170817":
|
|
kind = parser_args.kind
|
|
kind = kind + "_rsp" if parser_args.in_rsp else kind
|
|
|
|
fout = join(
|
|
"/mnt/extraspace/rstiskalek/GWLSS/",
|
|
f"{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.kind,
|
|
parser_args.MAS, parser_args.grid,
|
|
nsim, parser_args.in_rsp)
|
|
# The survey above had some cuts, however for compatibility we want
|
|
# the same shape as the `uncut` survey
|
|
val, rand_val = match_to_no_selection(val, rand_val, parser_args)
|
|
|
|
if verbose:
|
|
print(f"Saving to ... `{fout}`.")
|
|
numpy.savez(fout, val=val, rand_val=rand_val, indxs=indxs,
|
|
smooth_scales=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("--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"],
|
|
help="Mass assignment scheme.")
|
|
parser.add_argument("--grid", type=int, help="Grid resolution.")
|
|
parser.add_argument("--in_rsp", type=lambda x: bool(strtobool(x)),
|
|
help="Field in RSP?")
|
|
parser.add_argument("--nrand", type=int, required=True,
|
|
help="Number of rand. positions to evaluate the field")
|
|
parser.add_argument("--simname", type=str, default="csiborg1",
|
|
choices=["csiborg1"], help="Simulation name")
|
|
args = parser.parse_args()
|
|
|
|
paths = csiborgtools.read.Paths(**csiborgtools.paths_glamdring)
|
|
nsims = get_nsims(args, paths)
|
|
|
|
pos, indxs = open_galaxy_positions(args.survey, MPI.COMM_WORLD)
|
|
|
|
def _main(nsim):
|
|
main(nsim, args, pos, indxs, paths,
|
|
verbose=MPI.COMM_WORLD.Get_size() == 1)
|
|
|
|
work_delegation(_main, nsims, MPI.COMM_WORLD)
|