Aquila-Consortium/csiborgtools
0
0
Fork 0
mirror of https://github.com/Richard-Sti/csiborgtools.git synced 2024-12-22 12:18:01 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions

Forward model fields to RSP (#66)

Browse source 
* Add radial velocity field

* Add overdensity plot

* Flip velocities too

* Add field calculations

* Add RSP mapping

* Add potential in RSP

* Add projected field plotting
This commit is contained in:
Richard Stiskalek 2023-06-05 17:24:20 +02:00 committed by GitHub
parent f7b8b782a0
commit 63b6cdbe72
No known key found for this signature in database
GPG key ID: 4AEE18F83AFDEB23
7 changed files with 371 additions and 96 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

3
csiborgtools/field/__init__.py
View file

@ -18,7 +18,8 @@ try:
import MAS_library as MASL # noqa
from .density import DensityField, PotentialField, VelocityField # noqa
from .interp import evaluate_cartesian, evaluate_sky, make_sky # noqa
from .interp import (evaluate_cartesian, evaluate_sky, field2rsp, # noqa
make_sky)
from .utils import nside2radec, smoothen_field # noqa
except ImportError:
warn("MAS_library not found, `DensityField` will not be available", UserWarning) # noqa

37
csiborgtools/field/density.py
View file

@ -20,8 +20,10 @@ TODO:
"""
from abc import ABC
import MAS_library as MASL
import numpy
import MAS_library as MASL
from numba import jit
from tqdm import trange
from ..read.utils import real2redshift
@ -220,6 +222,36 @@ class VelocityField(BaseField):
self.box = box
self.MAS = MAS
@staticmethod
@jit(nopython=True)
def radial_velocity(rho_vel):
"""
Calculate the radial velocity field around the observer in the centre
of the box.
Parameters
----------
rho_vel : 4-dimensional array of shape `(3, grid, grid, grid)`.
Velocity field along each axis.
Returns
-------
radvel : 3-dimensional array of shape `(grid, grid, grid)`.
Radial velocity field.
"""
grid = rho_vel.shape[1]
radvel = numpy.zeros((grid, grid, grid), dtype=numpy.float32)
for i in range(grid):
px = i - 0.5 * (grid - 1)
for j in range(grid):
py = j - 0.5 * (grid - 1)
for k in range(grid):
pz = k - 0.5 * (grid - 1)
vx, vy, vz = rho_vel[:, i, j, k]
radvel[i, j, k] = ((px * vx + py * vy + pz * vz)
/ numpy.sqrt(px**2 + py**2 + pz**2))
return radvel
def __call__(self, parts, grid, mpart, flip_xz=True, nbatch=30,
verbose=True):
"""
@ -245,7 +277,7 @@ class VelocityField(BaseField):
Returns
-------
rho_vel : 3-dimensional array of shape `(3, grid, grid, grid)`.
rho_vel : 4-dimensional array of shape `(3, grid, grid, grid)`.
Velocity field along each axis.
References
@ -272,6 +304,7 @@ class VelocityField(BaseField):
mass = force_single_precision(mass, "particle_mass")
if flip_xz:
pos[:, [0, 2]] = pos[:, [2, 0]]
vel[:, [0, 2]] = vel[:, [2, 0]]
vel *= mass.reshape(-1, 1) / mpart
for i in range(3):

100
csiborgtools/field/interp.py
View file

@ -17,9 +17,11 @@ Tools for interpolating 3D fields at arbitrary positions.
"""
import MAS_library as MASL
import numpy
from numba import jit
from scipy.ndimage import gaussian_filter
from tqdm import trange
from ..read.utils import radec_to_cartesian
from ..read.utils import radec_to_cartesian, real2redshift
from .utils import force_single_precision
@ -137,3 +139,99 @@ def make_sky(field, angpos, dist, box, volume_weight=True, verbose=True):
evaluate_sky(field, pos=dir_loop, box=box, isdeg=True))
out *= dx
return out
@jit(nopython=True)
def divide_nonzero(field0, field1):
"""
Divide two fields where the second one is not zero. If the second field
is zero, the first one is left unchanged. Operates in-place.
Parameters
----------
field0 : 3-dimensional array of shape `(grid, grid, grid)`
Field to be divided.
field1 : 3-dimensional array of shape `(grid, grid, grid)`
Field to divide by.
Returns
-------
field0 : 3-dimensional array of shape `(grid, grid, grid)`
Field divided by the second one.
"""
assert field0.shape == field1.shape
imax, jmax, kmax = field0.shape
for i in range(imax):
for j in range(jmax):
for k in range(kmax):
if field1[i, j, k] != 0:
field0[i, j, k] /= field1[i, j, k]
def field2rsp(field, parts, box, nbatch=30, flip_partsxz=True, init_value=0.,
verbose=True):
"""
Forward model real space scalar field to redshift space. Attaches field
values to particles, those are then moved to redshift space and from their
positions reconstructs back the field on a regular grid by NGP
interpolation. This by definition produces a discontinuous field.
Parameters
----------
field : 3-dimensional array of shape `(grid, grid, grid)`
Real space field to be evolved to redshift space.
parts_pos : 2-dimensional array of shape `(n_parts, 3)`
Particle positions in real space.
parts_vel : 2-dimensional array of shape `(n_parts, 3)`
Particle velocities in real space.
box : :py:class:`csiborgtools.read.CSiBORGBox`
The simulation box information and transformations.
nbatch : int, optional
Number of batches to use when moving particles to redshift space.
Particles are assumed to be lazily loaded to memory.
flip_partsxz : bool, optional
Whether to flip the x and z coordinates of the particles. This is
because of a BORG bug.
init_value : float, optional
Initial value of the RSP field on the grid.
verbose : bool, optional
Verbosity flag.
Returns
-------
rsp_fields : (list of) 3-dimensional array of shape `(grid, grid, grid)`
"""
rsp_field = numpy.full(field.shape, init_value, dtype=numpy.float32)
cellcounts = numpy.zeros(rsp_field.shape, dtype=numpy.float32)
# We iterate over the fields and in the inner loop over the particles. This
# is slower than iterating over the particles and in the inner loop over
# the fields, but it is more memory efficient. Typically we will only have
# one field.
nparts = parts.shape[0]
batch_size = nparts // nbatch
start = 0
for k in trange(nbatch + 1) if verbose else range(nbatch + 1):
end = min(start + batch_size, nparts)
pos = parts[start:end]
pos, vel = pos[:, :3], pos[:, 3:6]
if flip_partsxz:
pos[:, [0, 2]] = pos[:, [2, 0]]
vel[:, [0, 2]] = vel[:, [2, 0]]
# Evaluate the field at the particle positions in real space.
values = evaluate_cartesian(field, pos=pos)
# Move particles to redshift space.
rsp_pos = real2redshift(pos, vel, [0.5, 0.5, 0.5], box,
in_box_units=True, periodic_wrap=True,
make_copy=True)
# Assign particles' values to the grid.
MASL.MA(rsp_pos, rsp_field, 1., "NGP", W=values)
# Count the number of particles in each grid cell.
MASL.MA(rsp_pos, cellcounts, 1., "NGP")
if end == nparts:
break
start = end
# Finally divide by the number of particles in each cell and smooth.
divide_nonzero(rsp_field, cellcounts)
return rsp_field

5
csiborgtools/read/paths.py
View file

@ -363,7 +363,8 @@ class Paths:
Parameters
----------
kind : str
Field type. Must be one of: `density`, `velocity`, `potential`.
Field type. Must be one of: `density`, `velocity`, `potential`,
`radvel`.
MAS : str
Mass-assignment scheme.
grid : int
@ -378,7 +379,7 @@ class Paths:
path : str
"""
fdir = join(self.postdir, "environment")
assert kind in ["density", "velocity", "potential"]
assert kind in ["density", "velocity", "potential", "radvel"]
if not isdir(fdir):
makedirs(fdir)
warn(f"Created directory `{fdir}`.", UserWarning, stacklevel=1)

78
scripts/field_derived.py
View file

@ -1,78 +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 calculate density field-derived fields in the CSiBORG
simulations' final snapshot.
"""
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("--kind", type=str, choices=["potential", "velocity"],
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("--in_rsp", type=lambda x: bool(strtobool(x)),
help="Calculate from the RSP density field?")
args = parser.parse_args()
paths = csiborgtools.read.Paths(**csiborgtools.paths_glamdring)
if args.ics is None or args.ics[0] == -1:
ics = paths.get_ics("csiborg")
else:
ics = args.ics
for i in csiborgtools.fits.split_jobs(len(ics), nproc)[rank]:
nsim = ics[i]
if verbose:
print(f"{datetime.now()}: rank {rank} working on simulation {nsim}.",
flush=True)
nsnap = max(paths.get_snapshots(nsim))
box = csiborgtools.read.CSiBORGBox(nsnap, nsim, paths)
density_gen = csiborgtools.field.DensityField(box, args.MAS)
rho = numpy.load(paths.field("density", args.MAS, args.grid, nsim,
args.in_rsp))
rho = density_gen.overdensity_field(rho)
if args.kind == "potential":
gen = csiborgtools.field.PotentialField(box, args.MAS)
else:
raise RuntimeError(f"Field {args.kind} is not implemented yet.")
field = gen(rho)
fout = paths.field("potential", args.MAS, args.grid, nsim,
args.in_rsp)
print(f"{datetime.now()}: rank {rank} saving output to `{fout}`.")
numpy.save(fout, field)

169
scripts/field_prop.py Normal file
View file

@ -0,0 +1,169 @@
# 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 calculate density field-derived fields in the CSiBORG
simulations' final snapshot.
"""
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
from taskmaster import work_delegation
from utils import get_nsims
###############################################################################
# Density field #
###############################################################################
def density_field(nsim, parser_args):
paths = csiborgtools.read.Paths(**csiborgtools.paths_glamdring)
nsnap = max(paths.get_snapshots(nsim))
box = csiborgtools.read.CSiBORGBox(nsnap, nsim, paths)
parts = csiborgtools.read.read_h5(paths.particles(nsim))["particles"]
gen = csiborgtools.field.DensityField(box, parser_args.MAS)
field = gen(parts, parser_args.grid, in_rsp=parser_args.in_rsp,
verbose=parser_args.verbose)
fout = paths.field("density", parser_args.MAS, parser_args.grid,
nsim, parser_args.in_rsp)
print(f"{datetime.now()}: saving output to `{fout}`.")
numpy.save(fout, field)
###############################################################################
# Velocity field #
###############################################################################
def velocity_field(nsim, parser_args):
if parser_args.in_rsp:
raise NotImplementedError("Velocity field in RSP is not implemented.")
paths = csiborgtools.read.Paths(**csiborgtools.paths_glamdring)
mpart = 1.1641532e-10 # Particle mass in CSiBORG simulations.
nsnap = max(paths.get_snapshots(nsim))
box = csiborgtools.read.CSiBORGBox(nsnap, nsim, paths)
parts = csiborgtools.read.read_h5(paths.particles(nsim))["particles"]
gen = csiborgtools.field.VelocityField(box, parser_args.MAS)
field = gen(parts, parser_args.grid, mpart, verbose=parser_args.verbose)
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)
###############################################################################
# Potential field #
###############################################################################
def potential_field(nsim, parser_args):
paths = csiborgtools.read.Paths(**csiborgtools.paths_glamdring)
nsnap = max(paths.get_snapshots(nsim))
box = csiborgtools.read.CSiBORGBox(nsnap, nsim, paths)
# Load the real space overdensity field
density_gen = csiborgtools.field.DensityField(box, parser_args.MAS)
rho = numpy.load(paths.field("density", parser_args.MAS, parser_args.grid,
nsim, in_rsp=False))
rho = density_gen.overdensity_field(rho)
# Calculate the real space potentiel field
gen = csiborgtools.field.PotentialField(box, parser_args.MAS)
field = gen(rho)
if parser_args.in_rsp:
parts = csiborgtools.read.read_h5(paths.particles(nsim))["particles"]
field = csiborgtools.field.field2rsp(field, parts=parts, box=box,
verbose=parser_args.verbose)
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)
###############################################################################
# Radial velocity field #
###############################################################################
def radvel_field(nsim, parser_args):
if parser_args.in_rsp:
raise NotImplementedError("Radial vel. field in RSP not implemented.")
paths = csiborgtools.read.Paths(**csiborgtools.paths_glamdring)
nsnap = max(paths.get_snapshots(nsim))
box = csiborgtools.read.CSiBORGBox(nsnap, nsim, paths)
vel = numpy.load(paths.field("velocity", parser_args.MAS, parser_args.grid,
nsim, parser_args.in_rsp))
gen = csiborgtools.field.VelocityField(box, parser_args.MAS)
field = gen.radial_velocity(vel)
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)
###############################################################################
# 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", "velocity", "radvel", "potential"],
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("--in_rsp", type=lambda x: bool(strtobool(x)),
help="Calculate in RSP?")
parser.add_argument("--verbose", type=lambda x: bool(strtobool(x)),
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":
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)
else:
raise RuntimeError(f"Field {parser_args.kind} is not implemented.")
work_delegation(main, nsims, comm, master_verbose=True)

73
scripts_plots/plot_data.py
View file

@ -162,29 +162,73 @@ def plot_hmf(pdf=False):
plt.close()
###############################################################################
# Sky distribution #
###############################################################################
@cache_to_disk(7)
def load_field(kind, nsim, grid, MAS, in_rsp=False):
paths = csiborgtools.read.Paths(**csiborgtools.paths_glamdring)
print(paths.field(kind, MAS, grid, nsim, in_rsp=in_rsp))
return numpy.load(paths.field(kind, MAS, grid, nsim, in_rsp=in_rsp))
###############################################################################
# Projected field #
###############################################################################
def plot_projected_field(kind, nsim, grid, in_rsp, MAS="PCS", pdf=False):
print(f"Plotting projected field `{kind}`. ", flush=True)
paths = csiborgtools.read.Paths(**csiborgtools.paths_glamdring)
nsnap = max(paths.get_snapshots(nsim))
box = csiborgtools.read.CSiBORGBox(nsnap, nsim, paths)
if kind == "overdensity":
field = load_field("density", nsim, grid, MAS=MAS, in_rsp=in_rsp)
density_gen = csiborgtools.field.DensityField(box, MAS)
field = density_gen.overdensity_field(field) + 2
else:
field = load_field(kind, nsim, grid, MAS=MAS, in_rsp=in_rsp)
print(field)
with plt.style.context(utils.mplstyle):
fig, ax = plt.subplots(figsize=(3.5 * 2, 2.625), ncols=3, sharey=True,
sharex=True)
fig.subplots_adjust(hspace=0, wspace=0)
for i in range(3):
ax[i].imshow(numpy.sum(field, axis=i))
fig.tight_layout(h_pad=0, w_pad=0)
for ext in ["png"] if pdf is False else ["png", "pdf"]:
fout = join(utils.fout, f"field_{kind}_{nsim}_rsp{in_rsp}.{ext}")
print(f"Saving to `{fout}`.")
fig.savefig(fout, dpi=utils.dpi, bbox_inches="tight")
plt.close()
###############################################################################
# Sky distribution #
###############################################################################
def get_sky_label(kind, volume_weight):
if volume_weight:
if kind == "density":
label = r"$\log \int_{0}^{R} r^2 \delta(r, \mathrm{RA}, \mathrm{dec}) \mathrm{d} r$" # noqa
label = r"$\log \int_{0}^{R} r^2 \rho(r, \mathrm{RA}, \mathrm{dec}) \mathrm{d} r$" # noqa
if kind == "overdensity":
label = r"$\log \int_{0}^{R} r^2 \left[\delta(r, \mathrm{RA}, \mathrm{dec}) + 2\right] \mathrm{d} r$" # noqa
elif kind == "potential":
label = r"$\int_{0}^{R} r^2 \phi(r, \mathrm{RA}, \mathrm{dec}) \mathrm{d} r$" # noqa
elif kind == "radvel":
label = r"$\int_{0}^{R} r^2 v_r(r, \mathrm{RA}, \mathrm{dec}) \mathrm{d} r$" # noqa
else:
label = None
else:
if kind == "density":
label = r"$\log \int_{0}^{R} \delta(r, \mathrm{RA}, \mathrm{dec}) \mathrm{d} r$" # noqa
label = r"$\log \int_{0}^{R} \rho(r, \mathrm{RA}, \mathrm{dec}) \mathrm{d} r$" # noqa
if kind == "overdensity":
label = r"$\log \int_{0}^{R} \left[\delta(r, \mathrm{RA}, \mathrm{dec}) + 2\right] \mathrm{d} r$" # noqa
elif kind == "potential":
label = r"$\int_{0}^{R} \phi(r, \mathrm{RA}, \mathrm{dec}) \mathrm{d} r$" # noqa
elif kind == "radvel":
label = r"$\int_{0}^{R} v_r(r, \mathrm{RA}, \mathrm{dec}) \mathrm{d} r$" # noqa
else:
label = None
return label
@ -200,6 +244,11 @@ def plot_sky_distribution(kind, nsim, grid, nside, MAS="PCS", plot_groups=True,
nsnap = max(paths.get_snapshots(nsim))
box = csiborgtools.read.CSiBORGBox(nsnap, nsim, paths)
if kind == "overdensity":
field = load_field("density", nsim, grid, MAS=MAS, in_rsp=False)
density_gen = csiborgtools.field.DensityField(box, MAS)
field = density_gen.overdensity_field(field) + 2
else:
field = load_field(kind, nsim, grid, MAS=MAS, in_rsp=False)
angpos = csiborgtools.field.nside2radec(nside)
@ -209,7 +258,7 @@ def plot_sky_distribution(kind, nsim, grid, nside, MAS="PCS", plot_groups=True,
with plt.style.context(utils.mplstyle):
label = get_sky_label(kind, volume_weight)
if kind == "density":
if kind in ["density", "overdensity"]:
out = numpy.log10(out)
healpy.mollview(out, fig=0, title="", unit=label)
@ -248,7 +297,7 @@ if __name__ == "__main__":
parser.add_argument('-c', '--clean', action='store_true')
args = parser.parse_args()
cached_funcs = []
cached_funcs = ["load_field"]
if args.clean:
for func in cached_funcs:
print(f"Cleaning cache for function {func}.")
@ -258,6 +307,8 @@ if __name__ == "__main__":
# plot_mass_vs_normcells(7444 + 24 * 4, pdf=False)
# plot_mass_vs_ncells(7444, pdf=True)
# plot_hmf(pdf=True)
plot_sky_distribution("potential", 7444, 256, nside=64, plot_groups=False,
dmin=50, dmax=100, plot_halos=5e13,
volume_weight=True)
# plot_sky_distribution("radvel", 7444, 256, nside=64,
# plot_groups=False, dmin=50, dmax=100,
# plot_halos=5e13, volume_weight=False)
plot_projected_field("potential", 7444, 256, in_rsp=True)