mirror of
https://github.com/Richard-Sti/csiborgtools_public.git
synced 2025-05-14 06:31:11 +00:00
Run density field estimator (#56)
* Add iterative density field generation * Edit particle halfwidth selection * Update import * Remove old file * Add position wrapping * Add RSD support * Add density field calculation * Edit paths to the density field * Flip argument order
This commit is contained in:
Richard Stiskalek
GitHub
parent
1d871e7109
commit
98d0578fa7
7 changed files with 138 additions and 209 deletions
70
scripts/field_density.py
Normal file
70
scripts/field_density.py
Normal file
|
|
@ -0,0 +1,70 @@
|
|||
# Copyright (C) 2022 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.
|
||||
"""
|
||||
MPI script to evaluate field properties at the galaxy positions.
|
||||
|
||||
NOTE THAT ONLY MAX SNAP
|
||||
"""
|
||||
from argparse import ArgumentParser
|
||||
from datetime import datetime
|
||||
from distutils.util import strtobool
|
||||
|
||||
import numpy
|
||||
from mpi4py import MPI
|
||||
|
||||
try:
|
||||
import csiborgtools
|
||||
except ModuleNotFoundError:
|
||||
import sys
|
||||
sys.path.append("../")
|
||||
import csiborgtools
|
||||
|
||||
comm = MPI.COMM_WORLD
|
||||
rank = comm.Get_rank()
|
||||
nproc = comm.Get_size()
|
||||
verbose = nproc == 1
|
||||
|
||||
parser = ArgumentParser()
|
||||
parser.add_argument("--ics", type=int, nargs="+", default=None,
|
||||
help="IC realisations. If `-1` processes all simulations.")
|
||||
parser.add_argument("--grid", type=int, help="Grid resolution.")
|
||||
parser.add_argument("--in_rsp", type=lambda x: bool(strtobool(x)),
|
||||
help="Calculate the density field in redshift space?")
|
||||
parser.add_argument("--MAS", type=str, choices=["NGP", "CIC", "TSC", "PCS"])
|
||||
args = parser.parse_args()
|
||||
paths = csiborgtools.read.CSiBORGPaths(**csiborgtools.paths_glamdring)
|
||||
|
||||
if args.ics is None or args.ics[0] == -1:
|
||||
ics = paths.get_ics(tonew=False)
|
||||
else:
|
||||
ics = args.ics
|
||||
|
||||
|
||||
for i in csiborgtools.fits.split_jobs(len(ics), nproc)[rank]:
|
||||
nsim = ics[i]
|
||||
print(f"{datetime.now()}: rank {rank} working on simulation {nsim}.",
|
||||
flush=True)
|
||||
|
||||
nsnap = max(paths.get_snapshots(nsim))
|
||||
box = csiborgtools.read.BoxUnits(nsnap, nsim, paths)
|
||||
parts = csiborgtools.read.read_h5(paths.particles_path(nsim))["particles"]
|
||||
density_generator = csiborgtools.field.DensityField(box, args.MAS)
|
||||
|
||||
rho = density_generator(parts, args.grid, in_rsp=args.in_rsp,
|
||||
verbose=verbose)
|
||||
|
||||
fout = paths.density_field_path(args.MAS, nsim, args.in_rsp)
|
||||
print(f"{datetime.now()}: rank {rank} saving output to `{fout}`.")
|
||||
numpy.save(fout, rho)
|
||||
|
|
@ -1,127 +0,0 @@
|
|||
# Copyright (C) 2022 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.
|
||||
"""
|
||||
MPI script to evaluate field properties at the galaxy positions.
|
||||
"""
|
||||
from argparse import ArgumentParser
|
||||
from datetime import datetime
|
||||
from os import remove
|
||||
from os.path import join
|
||||
|
||||
import numpy
|
||||
from mpi4py import MPI
|
||||
|
||||
try:
|
||||
import csiborgtools
|
||||
except ModuleNotFoundError:
|
||||
import sys
|
||||
sys.path.append("../")
|
||||
import csiborgtools
|
||||
|
||||
import utils
|
||||
|
||||
dumpdir = "/mnt/extraspace/rstiskalek/csiborg/"
|
||||
parser = ArgumentParser()
|
||||
parser.add_argument("--survey", type=str, choices=["SDSS"])
|
||||
parser.add_argument("--grid", type=int)
|
||||
parser.add_argument("--MAS", type=str, choices=["NGP", "CIC", "TSC", "PCS"])
|
||||
parser.add_argument("--halfwidth", type=float)
|
||||
parser.add_argument("--smooth_scale", type=float, default=None)
|
||||
args = parser.parse_args()
|
||||
# Smooth scale of 0 means no smoothing. Note that this is in Mpc/h
|
||||
args.smooth_scale = None if args.smooth_scale == 0 else args.smooth_scale
|
||||
|
||||
# Get MPI things
|
||||
comm = MPI.COMM_WORLD
|
||||
rank = comm.Get_rank()
|
||||
nproc = comm.Get_size()
|
||||
|
||||
# Galaxy positions
|
||||
survey = utils.surveys[args.survey]()()
|
||||
pos = numpy.vstack([survey[p] for p in ("DIST", "RA", "DEC")]).T
|
||||
pos = pos.astype(numpy.float32)
|
||||
|
||||
# File paths
|
||||
fname = "out_{}_{}_{}_{}_{}".format(
|
||||
survey.name, args.grid, args.MAS, args.halfwidth, args.smooth_scale)
|
||||
ftemp = join(dumpdir, "temp_fields", fname + "_{}.npy")
|
||||
fperm = join(dumpdir, "fields", fname + ".npy")
|
||||
|
||||
# Edit depending on what is calculated
|
||||
dtype = {"names": ["delta", "phi"], "formats": [numpy.float32] * 2}
|
||||
|
||||
# CSiBORG simulation paths
|
||||
paths = csiborgtools.read.CSiBORGPaths(**csiborgtools.paths_glamdring)
|
||||
ics = paths.get_ics(tonew=False)
|
||||
nsims = len(ics)
|
||||
|
||||
for n in csiborgtools.utils.split_jobs(nsims, nproc)[rank]:
|
||||
print("Rank {}@{}: working on {}th IC.".format(rank, datetime.now(), n),
|
||||
flush=True)
|
||||
nsim = ics[n]
|
||||
nsnap = max(paths.get_snapshots(nsim))
|
||||
reader = csiborgtools.read.ParticleReader(paths)
|
||||
box = csiborgtools.read.BoxUnits(nsnap, nsim, paths)
|
||||
|
||||
# Read particles and select a subset of them
|
||||
particles = reader.read_particle(nsnap, nsim, ["x", "y", "z", "M"],
|
||||
verbose=False)
|
||||
if args.halfwidth < 0.5:
|
||||
particles = csiborgtools.read.halfwidth_select(
|
||||
args.halfwidth, particles)
|
||||
length = box.box2mpc(2 * args.halfwidth) * box.h # Mpc/h
|
||||
else:
|
||||
length = box.box2mpc(1) * box.h # Mpc/h
|
||||
|
||||
# Initialise the field object and output array
|
||||
field = csiborgtools.field.DensityField(particles, length, box, args.MAS)
|
||||
out = numpy.full(pos.shape[0], numpy.nan, dtype=dtype)
|
||||
|
||||
# Calculate the overdensity field and interpolate at galaxy positions
|
||||
feval = field.overdensity_field(args.grid, args.smooth_scale,
|
||||
verbose=False)
|
||||
out["delta"] = field.evaluate_sky(feval, pos=pos, isdeg=True)[0]
|
||||
|
||||
# Potential
|
||||
feval = field.potential_field(args.grid, args.smooth_scale, verbose=False)
|
||||
out["phi"] = field.evaluate_sky(feval, pos=pos, isdeg=True)[0]
|
||||
|
||||
# Calculate the remaining fields
|
||||
# ...
|
||||
# ...
|
||||
|
||||
# Dump the results
|
||||
with open(ftemp.format(nsim), "wb") as f:
|
||||
numpy.save(f, out)
|
||||
|
||||
# Wait for all ranks to finish
|
||||
comm.Barrier()
|
||||
if rank == 0:
|
||||
print("Collecting files...", flush=True)
|
||||
|
||||
out = numpy.full((nsims, pos.shape[0]), numpy.nan, dtype=dtype)
|
||||
|
||||
for n in range(nsims):
|
||||
nsim = ics[n]
|
||||
with open(ftemp.format(nsim), "rb") as f:
|
||||
fin = numpy.load(f, allow_pickle=True)
|
||||
for name in dtype["names"]:
|
||||
out[name][n, ...] = fin[name]
|
||||
# Remove the temporary file
|
||||
remove(ftemp.format(nsim))
|
||||
|
||||
print("Saving results to `{}`.".format(fperm), flush=True)
|
||||
with open(fperm, "wb") as f:
|
||||
numpy.save(f, out)
|
||||
Loading…
Add table
Add a link
Reference in a new issue