Add galaxy sampling (#88)

* Improve calculations

* Improve flags

* Add smoothed options

* Remove some old comments

* Edit little things

* Save smoothed

* Move files

* Edit imports

* Edit imports

* Renaming imports

* Renaming imports

* Sort imports

* Sort files

* Sorting

* Optionally make copies of the field

* Add quijote backup check

* Add direct field smoothing

* Shorten stupid documentation

* Shorten stupid docs

* Update conversion

* Add particles to ASCII conversion

* Add a short comment

* Add SDSS uncorrected distance

* Adjust comment

* Add FITS index to galaxies

* Remove spare space

* Remove a stupid line

* Remove blank line

* Make space separated

* Add interpolated field path

* Add field sampling

* Sort imports

* Return density in cells

* Clear out observer velocity

* Add 170817 sampling

* Fix normalization

* Update plot
This commit is contained in:
Richard Stiskalek 2023-09-01 16:29:50 +01:00 committed by GitHub
parent 0af925e26a
commit eccd8e3507
No known key found for this signature in database
GPG key ID: 4AEE18F83AFDEB23
26 changed files with 610 additions and 365 deletions

View file

@ -23,16 +23,9 @@ from gc import collect
import numpy
from mpi4py import MPI
try:
import csiborgtools
except ModuleNotFoundError:
import sys
sys.path.append("../")
import csiborgtools
from taskmaster import work_delegation
import csiborgtools
from utils import get_nsims
###############################################################################
@ -60,6 +53,10 @@ def density_field(nsim, parser_args, to_save=True):
radvel_field = numpy.load(paths.field(
"radvel", parser_args.MAS, parser_args.grid, nsim, False))
if parser_args.verbose:
print(f"{datetime.now()}: converting density field to RSP.",
flush=True)
field = csiborgtools.field.field2rsp(field, radvel_field, box,
parser_args.MAS)
@ -187,6 +184,10 @@ def environment_field(nsim, parser_args, to_save=True):
density_gen = csiborgtools.field.DensityField(box, parser_args.MAS)
rho = density_gen.overdensity_field(rho)
if parser_args.smooth_scale > 0.0:
rho = csiborgtools.field.smoothen_field(
rho, parser_args.smooth_scale, box.box2mpc(1.))
gen = csiborgtools.field.TidalTensorField(box, parser_args.MAS)
field = gen(rho)
@ -217,7 +218,7 @@ def environment_field(nsim, parser_args, to_save=True):
if to_save:
fout = paths.field("environment", parser_args.MAS, parser_args.grid,
nsim, parser_args.in_rsp)
nsim, parser_args.in_rsp, parser_args.smooth_scale)
print(f"{datetime.now()}: saving output to `{fout}`.")
numpy.save(fout, env)
return env
@ -241,6 +242,8 @@ if __name__ == "__main__":
parser.add_argument("--grid", type=int, help="Grid resolution.")
parser.add_argument("--in_rsp", type=lambda x: bool(strtobool(x)),
help="Calculate in RSP?")
parser.add_argument("--smooth_scale", type=float, default=0.0,
help="Smoothing scale in Mpc / h. Only used for the environment field.") # noqa
parser.add_argument("--verbose", type=lambda x: bool(strtobool(x)),
help="Verbosity flag for reading in particles.")
parser.add_argument("--simname", type=str, default="csiborg",

194
scripts/field_sample.py Normal file
View file

@ -0,0 +1,194 @@
# 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.read.SDSS(
h=1, sel_steps=lambda cls: steps(cls, survey_name))
pos = numpy.vstack([survey["DIST_UNCORRECTED"],
survey["RA"],
survey["DEC"]],
).T
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 = 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)
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 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)
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", "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")
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)

View file

@ -81,7 +81,7 @@ def find_neighbour(args, nsim, cats, paths, comm, save_kind):
numpy.savez(fout, **out)
paths = csiborgtools.read.Paths(**csiborgtools.paths_glamdring)
reader = csiborgtools.read.NearestNeighbourReader(
reader = csiborgtools.summary.NearestNeighbourReader(
paths=paths, **csiborgtools.neighbour_kwargs)
counts = numpy.zeros((reader.nbins_radial, reader.nbins_neighbour),
dtype=numpy.float32)

View file

@ -68,7 +68,7 @@ def pair_match_max(nsim0, nsimx, simname, min_logmass, mult, verbose):
else:
raise ValueError(f"Unknown simulation `{simname}`.")
reader = csiborgtools.read.PairOverlap(cat0, catx, paths, min_logmass,
reader = csiborgtools.summary.PairOverlap(cat0, catx, paths, min_logmass,
maxdist=maxdist)
out = csiborgtools.match.matching_max(
cat0, catx, mass_kind, mult=mult, periodic=periodic,

View file

@ -0,0 +1,82 @@
# 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.
"""Convert the HDF5 CSiBORG particle file to an ASCII file."""
from argparse import ArgumentParser
import csiborgtools
import h5py
from mpi4py import MPI
from utils import get_nsims
from tqdm import trange
from taskmaster import work_delegation
def h5_to_ascii(nsim, paths, chunk_size=50_000, verbose=True):
"""
Convert the HDF5 CSiBORG particle file to an ASCII file. Outputs only
particle positions in Mpc / h. Ignores the unequal particle masses.
"""
fname = paths.particles(nsim, args.simname)
boxsize = 677.7
fname_out = fname.replace(".h5", ".txt")
with h5py.File(fname, 'r') as f:
dataset = f["particles"]
total_size = dataset.shape[0]
if verbose:
print(f"Number of rows to write: {total_size}")
with open(fname_out, 'w') as out_file:
# Write the header
out_file.write("#px py pz\n")
# Loop through data in chunks
for i in trange(0, total_size, chunk_size,
desc=f"Writing to ... `{fname_out}`",
disable=not verbose):
end = i + chunk_size
if end > total_size:
end = total_size
data_chunk = dataset[i:end]
# Convert to positions Mpc / h
data_chunk = data_chunk[:, :3] * boxsize
chunk_str = "\n".join([f"{x:.4f} {y:.4f} {z:.4f}"
for x, y, z in data_chunk])
out_file.write(chunk_str + "\n")
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="csiborg",
choices=["csiborg"],
help="Simulation name")
args = parser.parse_args()
paths = csiborgtools.read.Paths(**csiborgtools.paths_glamdring)
nsims = get_nsims(args, paths)
def main(nsim):
h5_to_ascii(nsim, paths, verbose=MPI.COMM_WORLD.Get_size() == 1)
work_delegation(main, nsims, MPI.COMM_WORLD)

View file

@ -37,22 +37,14 @@ except ModuleNotFoundError:
def get_nsims(args, paths):
"""
Get simulation indices from the command line arguments.
Parameters
----------
args : argparse.Namespace
Command line arguments. Must include `nsims` and `simname`. If `nsims`
is `None` or `-1`, all simulations in `simname` are used.
paths : :py:class`csiborgtools.paths.Paths`
Paths object.
Returns
-------
nsims : list of int
Simulation indices.
"""
try:
from_quijote_backup = args.from_quijote_backup
except AttributeError:
from_quijote_backup = False
if args.nsims is None or args.nsims[0] == -1:
nsims = paths.get_ics(args.simname, args.from_quijote_backup)
nsims = paths.get_ics(args.simname, from_quijote_backup)
else:
nsims = args.nsims
return list(nsims)
@ -81,8 +73,7 @@ def read_single_catalogue(args, config, nsim, run, rmax, paths, nobs=None):
Returns
-------
cat : csiborgtools.read.CSiBORGHaloCatalogue or csiborgtools.read.QuijoteHaloCatalogue # noqa
Halo catalogue with selection criteria applied.
`csiborgtools.read.CSiBORGHaloCatalogue` or `csiborgtools.read.QuijoteHaloCatalogue` # noqa
"""
selection = config.get(run, None)
if selection is None: