mirror of
https://github.com/Richard-Sti/csiborgtools_public.git
synced 2025-05-13 14:11:11 +00:00
Clean density calculation (#97)
* Get rid of utils * Clean up imports * Move some utils here * Rename file * Add simname to boxsize * Add imports * Delete old files * Update README * Update imports * Add a new draft of the density calculator * Update fields * Draft of new density field calculatiosn * Add snapshot * Add boxsizes * Little updates * Bring back utils * Edit docstrings * Edits imports * Add progress on snapshots * edit improts * add basic snapshot catalogue * Add support for CSiBORG2 snapshot reader * add paths to fofcat for csiborg2 * Add more imports * Add more boxsize * Add more imports * Add field readers * Simplify field paths * Fix typo * Add observer vp * Clean up density field calculation * Add a short note * Edit args * Remove old comments * Edit docs * Remove blank line * Stop flipping RAMSES * Remove comment * Edit desc * Remove normalization * Remove old dist array * Remove non-volume weighting * Remove non-volume weight * Add ignore of flake8 notebooks * Fix path typo * Fix units * Edit paths docs * Update nb
This commit is contained in:
Richard Stiskalek
GitHub
parent
eeff8f0ab9
commit
eb1797e8a9
19 changed files with 1260 additions and 1139 deletions
|
|
@ -1,108 +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.
|
||||
"""
|
||||
Script to calculate the peculiar velocity of an observer in the centre of the
|
||||
CSiBORG box.
|
||||
"""
|
||||
from argparse import ArgumentParser
|
||||
from distutils.util import strtobool
|
||||
|
||||
import numpy
|
||||
from mpi4py import MPI
|
||||
|
||||
from taskmaster import work_delegation
|
||||
from tqdm import tqdm
|
||||
from utils import get_nsims
|
||||
|
||||
try:
|
||||
import csiborgtools
|
||||
except ModuleNotFoundError:
|
||||
import sys
|
||||
sys.path.append("../")
|
||||
import csiborgtools
|
||||
|
||||
|
||||
def observer_peculiar_velocity(nsim, parser_args):
|
||||
"""
|
||||
Calculate the peculiar velocity of an observer in the centre of the box
|
||||
for several smoothing scales.
|
||||
"""
|
||||
pos = numpy.array([0.5, 0.5, 0.5]).reshape(-1, 3)
|
||||
boxsize = 677.7
|
||||
smooth_scales = [0, 0.5, 1.0, 1.5, 2.0, 2.5, 3.0, 3.5, 4.0]
|
||||
|
||||
observer_vp = numpy.full((len(smooth_scales), 3), numpy.nan,
|
||||
dtype=numpy.float32)
|
||||
|
||||
paths = csiborgtools.read.Paths(**csiborgtools.paths_glamdring)
|
||||
field_path = paths.field("velocity", parser_args.MAS, parser_args.grid,
|
||||
nsim, in_rsp=False)
|
||||
field0 = numpy.load(field_path)
|
||||
|
||||
for j, smooth_scale in enumerate(tqdm(smooth_scales,
|
||||
desc="Smoothing the fields",
|
||||
disable=not parser_args.verbose)):
|
||||
if smooth_scale > 0:
|
||||
field = [None, None, None]
|
||||
for k in range(3):
|
||||
field[k] = csiborgtools.field.smoothen_field(
|
||||
field0[k], smooth_scale, boxsize)
|
||||
else:
|
||||
field = field0
|
||||
|
||||
v = csiborgtools.field.evaluate_cartesian(
|
||||
field[0], field[1], field[2], pos=pos)
|
||||
observer_vp[j, 0] = v[0][0]
|
||||
observer_vp[j, 1] = v[1][0]
|
||||
observer_vp[j, 2] = v[2][0]
|
||||
|
||||
fout = paths.observer_peculiar_velocity(parser_args.MAS, parser_args.grid,
|
||||
nsim)
|
||||
if parser_args.verbose:
|
||||
print(f"Saving to ... `{fout}`")
|
||||
numpy.savez(fout, smooth_scales=smooth_scales, observer_vp=observer_vp)
|
||||
return observer_vp
|
||||
|
||||
|
||||
###############################################################################
|
||||
# Command line interface #
|
||||
###############################################################################
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
parser = ArgumentParser()
|
||||
parser.add_argument("--nsims", type=int, nargs="+", default=None,
|
||||
help="IC realisations. `-1` for all simulations.")
|
||||
parser.add_argument("--kind", type=str,
|
||||
choices=["density", "rspdensity", "velocity", "radvel",
|
||||
"potential", "environment"],
|
||||
help="What derived field to calculate?")
|
||||
parser.add_argument("--MAS", type=str,
|
||||
choices=["NGP", "CIC", "TSC", "PCS"])
|
||||
parser.add_argument("--grid", type=int, help="Grid resolution.")
|
||||
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",
|
||||
help="Verbosity flag for reading in particles.")
|
||||
parser_args = parser.parse_args()
|
||||
|
||||
comm = MPI.COMM_WORLD
|
||||
paths = csiborgtools.read.Paths(**csiborgtools.paths_glamdring)
|
||||
nsims = get_nsims(parser_args, paths)
|
||||
|
||||
def main(nsim):
|
||||
return observer_peculiar_velocity(nsim, parser_args)
|
||||
|
||||
work_delegation(main, nsims, comm, master_verbose=True)
|
||||
|
|
@ -12,14 +12,9 @@
|
|||
# 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 calculate density field-derived fields in the CSiBORG
|
||||
simulations' final snapshot.
|
||||
"""
|
||||
"""MPI script to calculate the various fields."""
|
||||
from argparse import ArgumentParser
|
||||
from datetime import datetime
|
||||
from distutils.util import strtobool
|
||||
from gc import collect
|
||||
|
||||
import numpy
|
||||
from mpi4py import MPI
|
||||
|
|
@ -29,55 +24,43 @@ import csiborgtools
|
|||
from utils import get_nsims
|
||||
|
||||
|
||||
###############################################################################
|
||||
# Cosmotool SPH density & velocity field #
|
||||
###############################################################################
|
||||
|
||||
def cosmotool_sph(nsim, parser_args):
|
||||
pass
|
||||
|
||||
|
||||
|
||||
###############################################################################
|
||||
# Density field #
|
||||
###############################################################################
|
||||
|
||||
|
||||
def density_field(nsim, parser_args, to_save=True):
|
||||
"""
|
||||
Calculate the density field in the CSiBORG simulation.
|
||||
"""
|
||||
def density_field(nsim, parser_args):
|
||||
"""Calculate the density field."""
|
||||
if parser_args.MAS == "SPH":
|
||||
raise NotImplementedError("SPH is not implemented here. Use cosmotool")
|
||||
|
||||
paths = csiborgtools.read.Paths(**csiborgtools.paths_glamdring)
|
||||
nsnap = max(paths.get_snapshots(nsim, "csiborg"))
|
||||
box = csiborgtools.read.CSiBORG1Box(nsnap, nsim, paths)
|
||||
fname = paths.processed_output(nsim, "csiborg", "halo_catalogue")
|
||||
nsnap = max(paths.get_snapshots(nsim, parser_args.simname))
|
||||
|
||||
if not parser_args.in_rsp:
|
||||
# TODO I removed this function
|
||||
snap = csiborgtools.read.read_h5(fname)["snapshot_final"]
|
||||
pos = snap["pos"]
|
||||
mass = snap["mass"]
|
||||
|
||||
gen = csiborgtools.field.DensityField(box, parser_args.MAS)
|
||||
field = gen(pos, mass, parser_args.grid, verbose=parser_args.verbose)
|
||||
# Read in the particle coordinates and masses
|
||||
if parser_args.simname == "csiborg1":
|
||||
snapshot = csiborgtools.read.CSIBORG1Snapshot(nsim, nsnap, paths)
|
||||
elif "csiborg2" in parser_args.simname:
|
||||
kind = parser_args.simname.split("_")[-1]
|
||||
snapshot = csiborgtools.read.CSIBORG2Snapshot(nsim, nsnap, paths, kind)
|
||||
elif parser_args.simname == "quijote":
|
||||
snapshot = csiborgtools.read.QuijoteSnapshot(nsim, nsnap, paths)
|
||||
else:
|
||||
field = numpy.load(paths.field(
|
||||
"density", parser_args.MAS, parser_args.grid, nsim, False))
|
||||
radvel_field = numpy.load(paths.field(
|
||||
"radvel", parser_args.MAS, parser_args.grid, nsim, False))
|
||||
raise RuntimeError(f"Unknown simulation name `{parser_args.simname}`.")
|
||||
|
||||
if parser_args.verbose:
|
||||
print(f"{datetime.now()}: converting density field to RSP.",
|
||||
flush=True)
|
||||
pos = snapshot.coordinates()
|
||||
mass = snapshot.masses()
|
||||
|
||||
field = csiborgtools.field.field2rsp(field, radvel_field, box,
|
||||
parser_args.MAS)
|
||||
# Run the field generator
|
||||
boxsize = csiborgtools.simname2boxsize(parser_args.simname)
|
||||
gen = csiborgtools.field.DensityField(boxsize, parser_args.MAS)
|
||||
field = gen(pos, mass, parser_args.grid)
|
||||
|
||||
if to_save:
|
||||
fout = paths.field(parser_args.kind, parser_args.MAS, parser_args.grid,
|
||||
nsim, parser_args.in_rsp)
|
||||
print(f"{datetime.now()}: saving output to `{fout}`.")
|
||||
numpy.save(fout, field)
|
||||
fout = paths.field("density", parser_args.MAS, parser_args.grid,
|
||||
nsim, parser_args.simname)
|
||||
|
||||
print(f"{datetime.now()}: saving output to `{fout}`.")
|
||||
numpy.save(fout, field)
|
||||
return field
|
||||
|
||||
|
||||
|
|
@ -86,31 +69,36 @@ def density_field(nsim, parser_args, to_save=True):
|
|||
###############################################################################
|
||||
|
||||
|
||||
def velocity_field(nsim, parser_args, to_save=True):
|
||||
"""
|
||||
Calculate the velocity field in a CSiBORG simulation.
|
||||
"""
|
||||
if parser_args.in_rsp:
|
||||
raise NotImplementedError("Velocity field in RSP is not implemented.")
|
||||
def velocity_field(nsim, parser_args):
|
||||
"""Calculate the velocity field."""
|
||||
if parser_args.MAS == "SPH":
|
||||
raise NotImplementedError("SPH is not implemented here. Use cosmotool")
|
||||
|
||||
paths = csiborgtools.read.Paths(**csiborgtools.paths_glamdring)
|
||||
nsnap = max(paths.get_snapshots(nsim, "csiborg"))
|
||||
box = csiborgtools.read.CSiBORG1Box(nsnap, nsim, paths)
|
||||
fname = paths.processed_output(nsim, "csiborg", "halo_catalogue")
|
||||
nsnap = max(paths.get_snapshots(nsim, parser_args.simname))
|
||||
|
||||
snap = csiborgtools.read.read_h5(fname)["snapshot_final"]
|
||||
pos = snap["pos"]
|
||||
vel = snap["vel"]
|
||||
mass = snap["mass"]
|
||||
if parser_args.simname == "csiborg1":
|
||||
snapshot = csiborgtools.read.CSIBORG1Snapshot(nsim, nsnap, paths)
|
||||
elif "csiborg2" in parser_args.simname:
|
||||
kind = parser_args.simname.split("_")[-1]
|
||||
snapshot = csiborgtools.read.CSIBORG2Snapshot(nsim, nsnap, paths, kind)
|
||||
elif parser_args.simname == "quijote":
|
||||
snapshot = csiborgtools.read.QuijoteSnapshot(nsim, nsnap, paths)
|
||||
else:
|
||||
raise RuntimeError(f"Unknown simulation name `{parser_args.simname}`.")
|
||||
|
||||
gen = csiborgtools.field.VelocityField(box, parser_args.MAS)
|
||||
field = gen(pos, vel, mass, parser_args.grid, verbose=parser_args.verbose)
|
||||
pos = snapshot.coordinates()
|
||||
vel = snapshot.velocities()
|
||||
mass = snapshot.masses()
|
||||
|
||||
if to_save:
|
||||
fout = paths.field("velocity", parser_args.MAS, parser_args.grid,
|
||||
nsim, in_rsp=False)
|
||||
print(f"{datetime.now()}: saving output to `{fout}`.")
|
||||
numpy.save(fout, field)
|
||||
boxsize = csiborgtools.simname2boxsize(parser_args.simname)
|
||||
gen = csiborgtools.field.VelocityField(boxsize, parser_args.MAS)
|
||||
field = gen(pos, vel, mass, parser_args.grid)
|
||||
|
||||
fout = paths.field("velocity", parser_args.MAS, parser_args.grid,
|
||||
nsim, parser_args.simname)
|
||||
print(f"{datetime.now()}: saving output to `{fout}`.")
|
||||
numpy.save(fout, field)
|
||||
return field
|
||||
|
||||
|
||||
|
|
@ -119,125 +107,62 @@ def velocity_field(nsim, parser_args, to_save=True):
|
|||
###############################################################################
|
||||
|
||||
|
||||
def radvel_field(nsim, parser_args, to_save=True):
|
||||
"""
|
||||
Calculate the radial velocity field in the CSiBORG simulation.
|
||||
"""
|
||||
if parser_args.in_rsp:
|
||||
raise NotImplementedError("Radial vel. field in RSP not implemented.")
|
||||
|
||||
def radvel_field(nsim, parser_args):
|
||||
"""Calculate the radial velocity field."""
|
||||
paths = csiborgtools.read.Paths(**csiborgtools.paths_glamdring)
|
||||
nsnap = max(paths.get_snapshots(nsim, "csiborg"))
|
||||
box = csiborgtools.read.CSiBORG1Box(nsnap, nsim, paths)
|
||||
|
||||
vel = numpy.load(paths.field("velocity", parser_args.MAS, parser_args.grid,
|
||||
nsim, parser_args.in_rsp))
|
||||
observer_velocity = csiborgtools.field.observer_vobs(vel)
|
||||
|
||||
gen = csiborgtools.field.VelocityField(box, parser_args.MAS)
|
||||
field = gen.radial_velocity(vel, observer_velocity)
|
||||
|
||||
if to_save:
|
||||
fout = paths.field("radvel", parser_args.MAS, parser_args.grid,
|
||||
nsim, parser_args.in_rsp)
|
||||
print(f"{datetime.now()}: saving output to `{fout}`.")
|
||||
numpy.save(fout, field)
|
||||
return field
|
||||
|
||||
###############################################################################
|
||||
# Potential field #
|
||||
###############################################################################
|
||||
|
||||
|
||||
def potential_field(nsim, parser_args, to_save=True):
|
||||
"""
|
||||
Calculate the potential field in the CSiBORG simulation.
|
||||
"""
|
||||
paths = csiborgtools.read.Paths(**csiborgtools.paths_glamdring)
|
||||
nsnap = max(paths.get_snapshots(nsim, "csiborg"))
|
||||
box = csiborgtools.read.CSiBORG1Box(nsnap, nsim, paths)
|
||||
|
||||
if not parser_args.in_rsp:
|
||||
rho = numpy.load(paths.field(
|
||||
"density", parser_args.MAS, parser_args.grid, nsim, in_rsp=False))
|
||||
density_gen = csiborgtools.field.DensityField(box, parser_args.MAS)
|
||||
rho = density_gen.overdensity_field(rho)
|
||||
|
||||
gen = csiborgtools.field.PotentialField(box, parser_args.MAS)
|
||||
field = gen(rho)
|
||||
if parser_args.simname == "csiborg1":
|
||||
field = csiborgtools.read.CSiBORG1Field(nsim, paths)
|
||||
elif "csiborg2" in parser_args.simname:
|
||||
kind = parser_args.simname.split("_")[-1]
|
||||
field = csiborgtools.read.CSiBORG2Field(nsim, paths, kind)
|
||||
elif parser_args.simname == "quijote":
|
||||
field = csiborgtools.read.QuijoteField(nsim, paths)
|
||||
else:
|
||||
field = numpy.load(paths.field(
|
||||
"potential", parser_args.MAS, parser_args.grid, nsim, False))
|
||||
radvel_field = numpy.load(paths.field(
|
||||
"radvel", parser_args.MAS, parser_args.grid, nsim, False))
|
||||
raise RuntimeError(f"Unknown simulation name `{parser_args.simname}`.")
|
||||
|
||||
field = csiborgtools.field.field2rsp(field, radvel_field, box,
|
||||
parser_args.MAS)
|
||||
vel = field.velocity_field(parser_args.MAS, parser_args.grid)
|
||||
|
||||
if to_save:
|
||||
fout = paths.field(parser_args.kind, parser_args.MAS, parser_args.grid,
|
||||
nsim, parser_args.in_rsp)
|
||||
print(f"{datetime.now()}: saving output to `{fout}`.")
|
||||
numpy.save(fout, field)
|
||||
observer_velocity = csiborgtools.field.observer_peculiar_velocity(vel)
|
||||
radvel = csiborgtools.field.radial_velocity(vel, observer_velocity)
|
||||
|
||||
fout = paths.field("radvel", parser_args.MAS, parser_args.grid,
|
||||
nsim, parser_args.simname)
|
||||
print(f"{datetime.now()}: saving output to `{fout}`.")
|
||||
numpy.save(fout, radvel)
|
||||
return field
|
||||
|
||||
|
||||
###############################################################################
|
||||
# Environment classification #
|
||||
###############################################################################
|
||||
|
||||
|
||||
def environment_field(nsim, parser_args, to_save=True):
|
||||
def observer_peculiar_velocity(nsim, parser_args):
|
||||
"""
|
||||
Calculate the environmental classification in the CSiBORG simulation.
|
||||
Calculate the peculiar velocity of an observer in the centre of the box
|
||||
for several smoothing scales.
|
||||
"""
|
||||
paths = csiborgtools.read.Paths(**csiborgtools.paths_glamdring)
|
||||
nsnap = max(paths.get_snapshots(nsim, "csiborg"))
|
||||
box = csiborgtools.read.CSiBORG1Box(nsnap, nsim, paths)
|
||||
boxsize = csiborgtools.simname2boxsize(parser_args.simname)
|
||||
# NOTE thevse values are hard-coded.
|
||||
smooth_scales = numpy.array([0, 1.0, 2.0, 3.0, 4.0, 5.0, 6.0, 7.0, 8.0])
|
||||
smooth_scales /= boxsize
|
||||
|
||||
rho = numpy.load(paths.field(
|
||||
"density", parser_args.MAS, parser_args.grid, nsim, in_rsp=False))
|
||||
density_gen = csiborgtools.field.DensityField(box, parser_args.MAS)
|
||||
rho = density_gen.overdensity_field(rho)
|
||||
if parser_args.simname == "csiborg1":
|
||||
field = csiborgtools.read.CSiBORG1Field(nsim, paths)
|
||||
elif "csiborg2" in parser_args.simname:
|
||||
kind = parser_args.simname.split("_")[-1]
|
||||
field = csiborgtools.read.CSiBORG2Field(nsim, paths, kind)
|
||||
elif parser_args.simname == "quijote":
|
||||
field = csiborgtools.read.QuijoteField(nsim, paths)
|
||||
else:
|
||||
raise RuntimeError(f"Unknown simulation name `{parser_args.simname}`.")
|
||||
|
||||
if parser_args.smooth_scale > 0.0:
|
||||
rho = csiborgtools.field.smoothen_field(
|
||||
rho, parser_args.smooth_scale, box.box2mpc(1.))
|
||||
vel = field.velocity_field(parser_args.MAS, parser_args.grid)
|
||||
|
||||
gen = csiborgtools.field.TidalTensorField(box, parser_args.MAS)
|
||||
field = gen(rho)
|
||||
observer_vp = csiborgtools.field.observer_peculiar_velocity(
|
||||
vel, smooth_scales)
|
||||
|
||||
del rho
|
||||
collect()
|
||||
|
||||
if parser_args.in_rsp:
|
||||
radvel_field = numpy.load(paths.field(
|
||||
"radvel", parser_args.MAS, parser_args.grid, nsim, False))
|
||||
args = (radvel_field, box, parser_args.MAS)
|
||||
|
||||
field.T00 = csiborgtools.field.field2rsp(field.T00, *args)
|
||||
field.T11 = csiborgtools.field.field2rsp(field.T11, *args)
|
||||
field.T22 = csiborgtools.field.field2rsp(field.T22, *args)
|
||||
field.T01 = csiborgtools.field.field2rsp(field.T01, *args)
|
||||
field.T02 = csiborgtools.field.field2rsp(field.T02, *args)
|
||||
field.T12 = csiborgtools.field.field2rsp(field.T12, *args)
|
||||
|
||||
del radvel_field
|
||||
collect()
|
||||
|
||||
eigvals = gen.tensor_field_eigvals(field)
|
||||
|
||||
del field
|
||||
collect()
|
||||
|
||||
env = gen.eigvals_to_environment(eigvals)
|
||||
|
||||
if to_save:
|
||||
fout = paths.field("environment", parser_args.MAS, parser_args.grid,
|
||||
nsim, parser_args.in_rsp, parser_args.smooth_scale)
|
||||
print(f"{datetime.now()}: saving output to `{fout}`.")
|
||||
numpy.save(fout, env)
|
||||
return env
|
||||
fout = paths.observer_peculiar_velocity(parser_args.MAS, parser_args.grid,
|
||||
nsim, parser_args.simname)
|
||||
print(f"Saving to ... `{fout}`")
|
||||
numpy.savez(fout, smooth_scales=smooth_scales, observer_vp=observer_vp)
|
||||
return observer_vp
|
||||
|
||||
|
||||
###############################################################################
|
||||
|
|
@ -249,39 +174,124 @@ if __name__ == "__main__":
|
|||
parser = ArgumentParser()
|
||||
parser.add_argument("--nsims", type=int, nargs="+", default=None,
|
||||
help="IC realisations. `-1` for all simulations.")
|
||||
parser.add_argument("--simname", type=str, help="Simulation name.")
|
||||
parser.add_argument("--kind", type=str,
|
||||
choices=["density", "rspdensity", "velocity", "radvel",
|
||||
"potential", "environment"],
|
||||
choices=["density", "velocity", "radvel", "observer_vp"], # noqa
|
||||
help="What derived field to calculate?")
|
||||
parser.add_argument("--MAS", type=str,
|
||||
choices=["NGP", "CIC", "TSC", "PCS"])
|
||||
choices=["NGP", "CIC", "TSC", "PCS", "SPH"],
|
||||
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="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",
|
||||
choices=["csiborg", "csiborg2"],
|
||||
help="Verbosity flag for reading in particles.")
|
||||
parser_args = parser.parse_args()
|
||||
|
||||
comm = MPI.COMM_WORLD
|
||||
paths = csiborgtools.read.Paths(**csiborgtools.paths_glamdring)
|
||||
nsims = get_nsims(parser_args, paths)
|
||||
|
||||
def main(nsim):
|
||||
if parser_args.kind == "density" or parser_args.kind == "rspdensity":
|
||||
if parser_args.kind == "density":
|
||||
density_field(nsim, parser_args)
|
||||
elif parser_args.kind == "velocity":
|
||||
velocity_field(nsim, parser_args)
|
||||
elif parser_args.kind == "radvel":
|
||||
radvel_field(nsim, parser_args)
|
||||
elif parser_args.kind == "potential":
|
||||
potential_field(nsim, parser_args)
|
||||
elif parser_args.kind == "environment":
|
||||
environment_field(nsim, parser_args)
|
||||
elif parser_args.kind == "observer_vp":
|
||||
observer_peculiar_velocity(nsim, parser_args)
|
||||
else:
|
||||
raise RuntimeError(f"Field {parser_args.kind} is not implemented.")
|
||||
|
||||
work_delegation(main, nsims, comm, master_verbose=True)
|
||||
|
||||
|
||||
# def potential_field(nsim, parser_args, to_save=True):
|
||||
# """
|
||||
# Calculate the potential field in the CSiBORG simulation.
|
||||
# """
|
||||
# paths = csiborgtools.read.Paths(**csiborgtools.paths_glamdring)
|
||||
# nsnap = max(paths.get_snapshots(nsim, "csiborg"))
|
||||
# box = csiborgtools.read.CSiBORG1Box(nsnap, nsim, paths)
|
||||
#
|
||||
# if not parser_args.in_rsp:
|
||||
# rho = numpy.load(paths.field(
|
||||
# "density", parser_args.MAS, parser_args.grid, nsim,
|
||||
# in_rsp=False))
|
||||
# density_gen = csiborgtools.field.DensityField(box, parser_args.MAS)
|
||||
# rho = density_gen.overdensity_field(rho)
|
||||
#
|
||||
# gen = csiborgtools.field.PotentialField(box, parser_args.MAS)
|
||||
# field = gen(rho)
|
||||
# else:
|
||||
# field = numpy.load(paths.field(
|
||||
# "potential", parser_args.MAS, parser_args.grid, nsim, False))
|
||||
# radvel_field = numpy.load(paths.field(
|
||||
# "radvel", parser_args.MAS, parser_args.grid, nsim, False))
|
||||
#
|
||||
# field = csiborgtools.field.field2rsp(field, radvel_field, box,
|
||||
# parser_args.MAS)
|
||||
#
|
||||
# if to_save:
|
||||
# fout = paths.field(parser_args.kind, parser_args.MAS,
|
||||
# parser_args.grid,
|
||||
# nsim, parser_args.in_rsp)
|
||||
# print(f"{datetime.now()}: saving output to `{fout}`.")
|
||||
# numpy.save(fout, field)
|
||||
# return field
|
||||
#
|
||||
#
|
||||
# #############################################################################
|
||||
# # Environment classification #
|
||||
# #############################################################################
|
||||
#
|
||||
#
|
||||
# def environment_field(nsim, parser_args, to_save=True):
|
||||
# """
|
||||
# Calculate the environmental classification in the CSiBORG simulation.
|
||||
# """
|
||||
# paths = csiborgtools.read.Paths(**csiborgtools.paths_glamdring)
|
||||
# nsnap = max(paths.get_snapshots(nsim, "csiborg"))
|
||||
# box = csiborgtools.read.CSiBORG1Box(nsnap, nsim, paths)
|
||||
#
|
||||
# rho = numpy.load(paths.field(
|
||||
# "density", parser_args.MAS, parser_args.grid, nsim, in_rsp=False))
|
||||
# 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)
|
||||
#
|
||||
# del rho
|
||||
# collect()
|
||||
#
|
||||
# if parser_args.in_rsp:
|
||||
# radvel_field = numpy.load(paths.field(
|
||||
# "radvel", parser_args.MAS, parser_args.grid, nsim, False))
|
||||
# args = (radvel_field, box, parser_args.MAS)
|
||||
#
|
||||
# field.T00 = csiborgtools.field.field2rsp(field.T00, *args)
|
||||
# field.T11 = csiborgtools.field.field2rsp(field.T11, *args)
|
||||
# field.T22 = csiborgtools.field.field2rsp(field.T22, *args)
|
||||
# field.T01 = csiborgtools.field.field2rsp(field.T01, *args)
|
||||
# field.T02 = csiborgtools.field.field2rsp(field.T02, *args)
|
||||
# field.T12 = csiborgtools.field.field2rsp(field.T12, *args)
|
||||
#
|
||||
# del radvel_field
|
||||
# collect()
|
||||
#
|
||||
# eigvals = gen.tensor_field_eigvals(field)
|
||||
#
|
||||
# del field
|
||||
# collect()
|
||||
#
|
||||
# env = gen.eigvals_to_environment(eigvals)
|
||||
#
|
||||
# if to_save:
|
||||
# fout = paths.field("environment", parser_args.MAS, parser_args.grid,
|
||||
# nsim, parser_args.in_rsp,
|
||||
# parser_args.smooth_scale)
|
||||
# print(f"{datetime.now()}: saving output to `{fout}`.")
|
||||
# numpy.save(fout, env)
|
||||
# return env
|
||||
|
|
|
|||
|
|
@ -30,7 +30,12 @@ import csiborgtools
|
|||
from utils import get_nsims
|
||||
|
||||
# TODO get rid of this.
|
||||
MPC2BOX = 1 / 677.7
|
||||
# MPC2BOX = 1 / 677.7
|
||||
SIM2BOXSIZE = {"csiborg1": 677.7,
|
||||
"csiborg2_main": None,
|
||||
"csiborg2_random": None,
|
||||
"csiborg2_varysmall": None,
|
||||
}
|
||||
|
||||
|
||||
def steps(cls, survey_name):
|
||||
|
|
|
|||
|
|
@ -1,374 +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.
|
||||
r"""
|
||||
Script to process simulation files and create a single HDF5 file, in which
|
||||
particles are sorted by the particle halo IDs.
|
||||
"""
|
||||
from argparse import ArgumentParser
|
||||
from gc import collect
|
||||
|
||||
import h5py
|
||||
import numpy
|
||||
from mpi4py import MPI
|
||||
|
||||
import csiborgtools
|
||||
from csiborgtools import fprint
|
||||
from numba import jit
|
||||
from taskmaster import work_delegation
|
||||
from tqdm import trange, tqdm
|
||||
from utils import get_nsims
|
||||
|
||||
|
||||
@jit(nopython=True, boundscheck=False)
|
||||
def minmax_halo(hid, halo_ids, start_loop=0):
|
||||
"""
|
||||
Find the start and end index of a halo in a sorted array of halo IDs.
|
||||
This is much faster than using `numpy.where` and then `numpy.min` and
|
||||
`numpy.max`.
|
||||
"""
|
||||
start = None
|
||||
end = None
|
||||
|
||||
for i in range(start_loop, halo_ids.size):
|
||||
n = halo_ids[i]
|
||||
if n == hid:
|
||||
if start is None:
|
||||
start = i
|
||||
end = i
|
||||
elif n > hid:
|
||||
break
|
||||
return start, end
|
||||
|
||||
|
||||
def process_snapshot(nsim, simname, halo_finder, verbose):
|
||||
"""
|
||||
Read in the snapshot particles, sort them by their halo ID and dump
|
||||
into a HDF5 file. Stores the first and last index of each halo in the
|
||||
particle array for fast slicing of the array to acces particles of a single
|
||||
halo.
|
||||
"""
|
||||
paths = csiborgtools.read.Paths(**csiborgtools.paths_glamdring)
|
||||
nsnap = max(paths.get_snapshots(nsim, simname))
|
||||
|
||||
if simname == "csiborg":
|
||||
partreader = csiborgtools.read.CSiBORGReader(paths)
|
||||
box = csiborgtools.read.CSiBORGBox(nsnap, nsim, paths)
|
||||
else:
|
||||
partreader = csiborgtools.read.QuijoteReader(paths)
|
||||
box = None
|
||||
|
||||
desc = {"hid": f"Halo finder ID ({halo_finder})of the particle.",
|
||||
"pos": "DM particle positions in box units.",
|
||||
"vel": "DM particle velocity in km / s.",
|
||||
"mass": "DM particle mass in Msun / h.",
|
||||
"pid": "DM particle ID",
|
||||
}
|
||||
|
||||
fname = paths.processed_output(nsim, simname, halo_finder)
|
||||
|
||||
fprint(f"loading HIDs of IC {nsim}.", verbose)
|
||||
hids = partreader.read_halo_id(nsnap, nsim, halo_finder, verbose)
|
||||
collect()
|
||||
|
||||
fprint(f"sorting HIDs of IC {nsim}.")
|
||||
sort_indxs = numpy.argsort(hids)
|
||||
|
||||
with h5py.File(fname, "w") as f:
|
||||
group = f.create_group("snapshot_final")
|
||||
group.attrs["header"] = "Snapshot data at z = 0."
|
||||
|
||||
fprint("dumping halo IDs.", verbose)
|
||||
dset = group.create_dataset("halo_ids", data=hids[sort_indxs])
|
||||
dset.attrs["header"] = desc["hid"]
|
||||
del hids
|
||||
collect()
|
||||
|
||||
fprint("reading, sorting and dumping the snapshot particles.", verbose)
|
||||
for kind in ["pos", "vel", "mass", "pid"]:
|
||||
x = partreader.read_snapshot(nsnap, nsim, kind)[sort_indxs]
|
||||
|
||||
if simname == "csiborg" and kind == "vel":
|
||||
x = box.box2vel(x) if simname == "csiborg" else x
|
||||
|
||||
if simname == "csiborg" and kind == "mass":
|
||||
x = box.box2solarmass(x) if simname == "csiborg" else x
|
||||
|
||||
dset = f["snapshot_final"].create_dataset(kind, data=x)
|
||||
dset.attrs["header"] = desc[kind]
|
||||
del x
|
||||
collect()
|
||||
|
||||
del sort_indxs
|
||||
collect()
|
||||
|
||||
fprint(f"creating a halo map for IC {nsim}.")
|
||||
with h5py.File(fname, "r") as f:
|
||||
part_hids = f["snapshot_final"]["halo_ids"][:]
|
||||
# We loop over the unique halo IDs and remove the 0 halo ID
|
||||
unique_halo_ids = numpy.unique(part_hids)
|
||||
unique_halo_ids = unique_halo_ids[unique_halo_ids != 0]
|
||||
halo_map = numpy.full((unique_halo_ids.size, 3), numpy.nan,
|
||||
dtype=numpy.uint64)
|
||||
start_loop, niters = 0, unique_halo_ids.size
|
||||
for i in trange(niters, disable=not verbose):
|
||||
hid = unique_halo_ids[i]
|
||||
k0, kf = minmax_halo(hid, part_hids, start_loop=start_loop)
|
||||
halo_map[i, :] = hid, k0, kf
|
||||
start_loop = kf
|
||||
|
||||
# Dump the halo mapping.
|
||||
with h5py.File(fname, "r+") as f:
|
||||
dset = f["snapshot_final"].create_dataset("halo_map", data=halo_map)
|
||||
dset.attrs["header"] = """
|
||||
Halo to particle mapping. Columns are HID, start index, end index.
|
||||
"""
|
||||
f.close()
|
||||
|
||||
del part_hids
|
||||
collect()
|
||||
|
||||
# Add the halo finder catalogue
|
||||
with h5py.File(fname, "r+") as f:
|
||||
group = f.create_group("halofinder_catalogue")
|
||||
group.attrs["header"] = f"Original {halo_finder} halo catalogue."
|
||||
cat = partreader.read_catalogue(nsnap, nsim, halo_finder)
|
||||
|
||||
hid2pos = {hid: i for i, hid in enumerate(unique_halo_ids)}
|
||||
|
||||
for key in cat.dtype.names:
|
||||
x = numpy.full(unique_halo_ids.size, numpy.nan,
|
||||
dtype=cat[key].dtype)
|
||||
for i in range(len(cat)):
|
||||
j = hid2pos[cat["index"][i]]
|
||||
x[j] = cat[key][i]
|
||||
group.create_dataset(key, data=x)
|
||||
f.close()
|
||||
|
||||
# Lastly create the halo catalogue
|
||||
with h5py.File(fname, "r+") as f:
|
||||
group = f.create_group("halo_catalogue")
|
||||
group.attrs["header"] = f"{halo_finder} halo catalogue."
|
||||
group.create_dataset("index", data=unique_halo_ids)
|
||||
f.close()
|
||||
|
||||
|
||||
def add_initial_snapshot(nsim, simname, halo_finder, verbose):
|
||||
"""
|
||||
Sort the initial snapshot particles according to their final snapshot and
|
||||
add them to the final snapshot's HDF5 file.
|
||||
"""
|
||||
paths = csiborgtools.read.Paths(**csiborgtools.paths_glamdring)
|
||||
fname = paths.processed_output(nsim, simname, halo_finder)
|
||||
|
||||
if simname == "csiborg":
|
||||
partreader = csiborgtools.read.CSiBORGReader(paths)
|
||||
else:
|
||||
partreader = csiborgtools.read.QuijoteReader(paths)
|
||||
|
||||
fprint(f"processing simulation `{nsim}`.", verbose)
|
||||
if simname == "csiborg":
|
||||
nsnap0 = 1
|
||||
elif simname == "quijote":
|
||||
nsnap0 = -1
|
||||
else:
|
||||
raise ValueError(f"Unknown simulation `{simname}`.")
|
||||
|
||||
fprint("loading and sorting the initial PID.", verbose)
|
||||
sort_indxs = numpy.argsort(partreader.read_snapshot(nsnap0, nsim, "pid"))
|
||||
|
||||
fprint("loading the final particles.", verbose)
|
||||
with h5py.File(fname, "r") as f:
|
||||
sort_indxs_final = f["snapshot_final/pid"][:]
|
||||
f.close()
|
||||
|
||||
fprint("sorting the particles according to the final snapshot.", verbose)
|
||||
sort_indxs_final = numpy.argsort(numpy.argsort(sort_indxs_final))
|
||||
sort_indxs = sort_indxs[sort_indxs_final]
|
||||
|
||||
del sort_indxs_final
|
||||
collect()
|
||||
|
||||
fprint("loading and sorting the initial particle position.", verbose)
|
||||
pos = partreader.read_snapshot(nsnap0, nsim, "pos")[sort_indxs]
|
||||
|
||||
del sort_indxs
|
||||
collect()
|
||||
|
||||
# In Quijote some particles are position precisely at the edge of the
|
||||
# box. Move them to be just inside.
|
||||
if simname == "quijote":
|
||||
mask = pos >= 1
|
||||
if numpy.any(mask):
|
||||
spacing = numpy.spacing(pos[mask])
|
||||
assert numpy.max(spacing) <= 1e-5
|
||||
pos[mask] -= spacing
|
||||
|
||||
fprint(f"dumping particles for `{nsim}` to `{fname}`.", verbose)
|
||||
with h5py.File(fname, "r+") as f:
|
||||
if "snapshot_initial" in f.keys():
|
||||
del f["snapshot_initial"]
|
||||
group = f.create_group("snapshot_initial")
|
||||
group.attrs["header"] = "Initial snapshot data."
|
||||
dset = group.create_dataset("pos", data=pos)
|
||||
dset.attrs["header"] = "DM particle positions in box units."
|
||||
|
||||
f.close()
|
||||
|
||||
|
||||
def calculate_initial(nsim, simname, halo_finder, verbose):
|
||||
"""Calculate the Lagrangian patch centre of mass and size."""
|
||||
paths = csiborgtools.read.Paths(**csiborgtools.paths_glamdring)
|
||||
|
||||
fname = paths.processed_output(nsim, simname, halo_finder)
|
||||
fprint("loading the particle information.", verbose)
|
||||
f = h5py.File(fname, "r")
|
||||
pos = f["snapshot_initial/pos"]
|
||||
mass = f["snapshot_final/mass"]
|
||||
hid = f["halo_catalogue/index"][:]
|
||||
hid2map = csiborgtools.read.make_halomap_dict(
|
||||
f["snapshot_final/halo_map"][:])
|
||||
|
||||
if simname == "csiborg":
|
||||
kwargs = {"box_size": 2048, "bckg_halfsize": 512}
|
||||
else:
|
||||
kwargs = {"box_size": 512, "bckg_halfsize": 256}
|
||||
overlapper = csiborgtools.match.ParticleOverlap(**kwargs)
|
||||
|
||||
lagpatch_pos = numpy.full((len(hid), 3), numpy.nan, dtype=numpy.float32)
|
||||
lagpatch_size = numpy.full(len(hid), numpy.nan, dtype=numpy.float32)
|
||||
lagpatch_ncells = numpy.full(len(hid), numpy.nan, dtype=numpy.int32)
|
||||
|
||||
for i in trange(len(hid), disable=not verbose):
|
||||
h = hid[i]
|
||||
# These are unasigned particles.
|
||||
if h == 0:
|
||||
continue
|
||||
|
||||
parts_pos = csiborgtools.read.load_halo_particles(h, pos, hid2map)
|
||||
parts_mass = csiborgtools.read.load_halo_particles(h, mass, hid2map)
|
||||
|
||||
# Skip if the halo has no particles or is too small.
|
||||
if parts_pos is None or parts_pos.size < 5:
|
||||
continue
|
||||
|
||||
cm = csiborgtools.center_of_mass(parts_pos, parts_mass, boxsize=1.0)
|
||||
sep = csiborgtools.periodic_distance(parts_pos, cm, boxsize=1.0)
|
||||
delta = overlapper.make_delta(parts_pos, parts_mass, subbox=True)
|
||||
|
||||
lagpatch_pos[i] = cm
|
||||
lagpatch_size[i] = numpy.percentile(sep, 99)
|
||||
lagpatch_ncells[i] = csiborgtools.delta2ncells(delta)
|
||||
|
||||
f.close()
|
||||
collect()
|
||||
|
||||
with h5py.File(fname, "r+") as f:
|
||||
grp = f["halo_catalogue"]
|
||||
dset = grp.create_dataset("lagpatch_pos", data=lagpatch_pos)
|
||||
dset.attrs["header"] = "Lagrangian patch centre of mass in box units."
|
||||
|
||||
dset = grp.create_dataset("lagpatch_size", data=lagpatch_size)
|
||||
dset.attrs["header"] = "Lagrangian patch size in box units."
|
||||
|
||||
dset = grp.create_dataset("lagpatch_ncells", data=lagpatch_ncells)
|
||||
dset.attrs["header"] = f"Lagrangian patch number of cells on a {kwargs['box_size']}^3 grid." # noqa
|
||||
|
||||
f.close()
|
||||
|
||||
|
||||
def make_phew_halo_catalogue(nsim, verbose):
|
||||
"""
|
||||
Process the PHEW halo catalogue for a CSiBORG simulation at all snapshots.
|
||||
"""
|
||||
paths = csiborgtools.read.Paths(**csiborgtools.paths_glamdring)
|
||||
snapshots = paths.get_snapshots(nsim, "csiborg")
|
||||
reader = csiborgtools.read.CSiBORGReader(paths)
|
||||
keys_write = ["index", "x", "y", "z", "mass_cl", "parent",
|
||||
"ultimate_parent", "summed_mass"]
|
||||
|
||||
# Create a HDF5 file to store all this.
|
||||
fname = paths.processed_phew(nsim)
|
||||
with h5py.File(fname, "w") as f:
|
||||
f.close()
|
||||
|
||||
for nsnap in tqdm(snapshots, disable=not verbose, desc="Snapshot"):
|
||||
try:
|
||||
data = reader.read_phew_clumps(nsnap, nsim, verbose=False)
|
||||
except FileExistsError:
|
||||
continue
|
||||
|
||||
with h5py.File(fname, "r+") as f:
|
||||
if str(nsnap) in f:
|
||||
print(f"Group {nsnap} already exists. Deleting.", flush=True)
|
||||
del f[str(nsnap)]
|
||||
grp = f.create_group(str(nsnap))
|
||||
for key in keys_write:
|
||||
grp.create_dataset(key, data=data[key])
|
||||
|
||||
grp.attrs["header"] = f"CSiBORG PHEW clumps at snapshot {nsnap}."
|
||||
f.close()
|
||||
|
||||
# Now write the redshifts
|
||||
scale_factors = numpy.full(len(snapshots), numpy.nan, dtype=numpy.float32)
|
||||
for i, nsnap in enumerate(snapshots):
|
||||
box = csiborgtools.read.CSiBORGBox(nsnap, nsim, paths)
|
||||
scale_factors[i] = box._aexp
|
||||
|
||||
redshifts = scale_factors[-1] / scale_factors - 1
|
||||
|
||||
with h5py.File(fname, "r+") as f:
|
||||
grp = f.create_group("info")
|
||||
grp.create_dataset("redshift", data=redshifts)
|
||||
grp.create_dataset("snapshots", data=snapshots)
|
||||
grp.create_dataset("Om0", data=[box.Om0])
|
||||
grp.create_dataset("boxsize", data=[box.boxsize])
|
||||
f.close()
|
||||
|
||||
|
||||
def main(nsim, args):
|
||||
if args.make_final:
|
||||
process_snapshot(nsim, args.simname, args.halofinder, True)
|
||||
|
||||
if args.make_initial:
|
||||
add_initial_snapshot(nsim, args.simname, args.halofinder, True)
|
||||
calculate_initial(nsim, args.simname, args.halofinder, True)
|
||||
|
||||
if args.make_phew:
|
||||
make_phew_halo_catalogue(nsim, True)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
parser = ArgumentParser()
|
||||
parser.add_argument("--simname", type=str, default="csiborg",
|
||||
choices=["csiborg", "quijote"],
|
||||
help="Simulation name")
|
||||
parser.add_argument("--nsims", type=int, nargs="+", default=None,
|
||||
help="IC realisations. If `-1` processes all.")
|
||||
parser.add_argument("--halofinder", type=str, help="Halo finder")
|
||||
parser.add_argument("--make_final", action="store_true", default=False,
|
||||
help="Process the final snapshot.")
|
||||
parser.add_argument("--make_initial", action="store_true", default=False,
|
||||
help="Process the initial snapshot.")
|
||||
parser.add_argument("--make_phew", action="store_true", default=False,
|
||||
help="Process the PHEW halo catalogue.")
|
||||
|
||||
args = parser.parse_args()
|
||||
paths = csiborgtools.read.Paths(**csiborgtools.paths_glamdring)
|
||||
nsims = get_nsims(args, paths)
|
||||
|
||||
def _main(nsim):
|
||||
main(nsim, args)
|
||||
|
||||
work_delegation(_main, nsims, MPI.COMM_WORLD)
|
||||
Loading…
Add table
Add a link
Reference in a new issue