csiborgtools/scripts/field_density.py
Richard Stiskalek 98d0578fa7
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
2023-05-08 13:58:12 +01:00

71 lines
2.5 KiB
Python

# 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)