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

Density field calculation (#58)

Browse source 
* 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

* Edit documentation

* Edit docs

* update nb
This commit is contained in:
Richard Stiskalek 2023-05-08 14:02:20 +01:00 committed by GitHub
parent 4265315b47
commit cb5d35b1be
No known key found for this signature in database
GPG key ID: 4AEE18F83AFDEB23
8 changed files with 227 additions and 241 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

101
csiborgtools/field/density.py
View file

@ -16,16 +16,14 @@
Density field and cross-correlation calculations.
"""
from abc import ABC
from warnings import warn
import MAS_library as MASL
import numpy
import Pk_library as PKL
import smoothing_library as SL
from tqdm import trange
from .utils import force_single_precision
from ..read.utils import radec_to_cartesian
from ..read.utils import radec_to_cartesian, real2redshift
class BaseField(ABC):
@ -189,11 +187,6 @@ class DensityField(BaseField):
Parameters
----------
pos : 2-dimensional array of shape `(N, 3)`
Particle position array. Columns must be ordered as `['x', 'y', 'z']`.
The positions are assumed to be in box units, i.e. :math:`\in [0, 1 ]`.
mass : 1-dimensional array of shape `(N,)`
Particle mass array. Assumed to be in box units.
box : :py:class:`csiborgtools.read.BoxUnits`
The simulation box information and transformations.
MAS : str
@ -205,52 +198,11 @@ class DensityField(BaseField):
----------
[1] https://pylians3.readthedocs.io/
"""
_pos = None
_mass = None
def __init__(self, pos, mass, box, MAS):
self.pos = pos
self.mass = mass
def __init__(self, box, MAS):
self.box = box
self.MAS = MAS
@property
def pos(self):
"""
Particle position array.
Returns
-------
particles : 2-dimensional array
"""
return self._particles
@pos.setter
def pos(self, pos):
assert pos.ndim == 2
warn("Flipping the `x` and `z` coordinates of the particle positions.",
UserWarning, stacklevel=1)
pos[:, [0, 2]] = pos[:, [2, 0]]
pos = force_single_precision(pos, "particle_position")
self._pos = pos
@property
def mass(self):
"""
Particle mass array.
Returns
-------
mass : 1-dimensional array
"""
return self._mass
@mass.setter
def mass(self, mass):
assert mass.ndim == 1
mass = force_single_precision(mass, "particle_mass")
self._mass = mass
def smoothen(self, field, smooth_scale, threads=1):
"""
Smooth a field with a Gaussian filter.
@ -294,17 +246,28 @@ class DensityField(BaseField):
delta -= 1
return delta
def __call__(self, grid, smooth_scale=None, verbose=True):
def __call__(self, parts, grid, in_rsp, flip_xz=True, nbatch=30,
verbose=True):
"""
Calculate the density field using a Pylians routine [1, 2].
Iteratively loads the particles into memory, flips their `x` and `z`
coordinates. Particles are assumed to be in box units, with positions
in [0, 1] and observer in the centre of the box if RSP is applied.
Parameters
----------
parts : 2-dimensional array of shape `(n_parts, 7)`
Particle positions, velocities and masses.
Columns are: `x`, `y`, `z`, `vx`, `vy`, `vz`, `M`.
grid : int
Grid size.
smooth_scale : float, optional
Gaussian kernal scale to smoothen the density field, in box units.
verbose : bool
in_rsp : bool
Whether to calculate the density field in redshift space.
flip_xz : bool, optional
Whether to flip the `x` and `z` coordinates.
nbatch : int, optional
Number of batches to split the particle loading into.
verbose : bool, optional
Verbosity flag.
Returns
@ -318,12 +281,32 @@ class DensityField(BaseField):
[2] https://github.com/franciscovillaescusa/Pylians3/blob/master
/library/MAS_library/MAS_library.pyx
"""
# Pre-allocate and do calculations
rho = numpy.zeros((grid, grid, grid), dtype=numpy.float32)
MASL.MA(self.pos, rho, self.boxsize, self.MAS, W=self.mass,
verbose=verbose)
if smooth_scale is not None:
rho = self.smoothen(rho, smooth_scale)
nparts = parts.shape[0]
batch_size = nparts // nbatch
start = 0
for __ in trange(nbatch + 1) if verbose else range(nbatch + 1):
end = min(start + batch_size, nparts)
pos = parts[start:end]
pos, vel, mass = pos[:, :3], pos[:, 3:6], pos[:, 6]
pos = force_single_precision(pos, "particle_position")
vel = force_single_precision(vel, "particle_velocity")
mass = force_single_precision(mass, "particle_mass")
if flip_xz:
pos[:, [0, 2]] = pos[:, [2, 0]]
vel[:, [0, 2]] = vel[:, [2, 0]]
if in_rsp:
pos = real2redshift(pos, vel, [0.5, 0.5, 0.5], self.box,
in_box_units=True, periodic_wrap=True,
make_copy=False)
MASL.MA(pos, rho, self.boxsize, self.MAS, W=mass, verbose=False)
if end == nparts:
break
start = end
return rho

2
csiborgtools/read/__init__.py
View file

@ -21,7 +21,7 @@ from .overlap_summary import (NPairsOverlap, PairOverlap, # noqa
binned_resample_mean)
from .paths import CSiBORGPaths # noqa
from .pk_summary import PKReader # noqa
from .readsim import (MmainReader, ParticleReader, halfwidth_select, # noqa
from .readsim import (MmainReader, ParticleReader, halfwidth_mask, # noqa
load_clump_particles, load_parent_particles, read_initcm)
from .tpcf_summary import TPCFReader # noqa
from .utils import (cartesian_to_radec, cols_to_structured, # noqa

10
csiborgtools/read/paths.py
View file

@ -329,16 +329,18 @@ class CSiBORGPaths:
fname = f"parts_{str(nsim).zfill(5)}.h5"
return join(fdir, fname)
def density_field_path(self, mas, nsim):
def density_field_path(self, MAS, nsim, in_rsp):
"""
Path to the files containing the calculated density fields.
Parameters
----------
mas : str
MAS : str
Mass-assignment scheme. Currently only SPH is supported.
nsim : int
IC realisation index.
in_rsp : bool
Whether the density field is calculated in redshift space.
Returns
-------
@ -348,7 +350,9 @@ class CSiBORGPaths:
if not isdir(fdir):
makedirs(fdir)
warn(f"Created directory `{fdir}`.", UserWarning, stacklevel=1)
fname = f"density_{mas}_{str(nsim).zfill(5)}.npy"
fname = f"density_{MAS}_{str(nsim).zfill(5)}.npy"
if in_rsp:
fname = fname.replace("density", "density_rsp")
return join(fdir, fname)
def knnauto_path(self, run, nsim=None):

25
csiborgtools/read/readsim.py
View file

@ -534,34 +534,23 @@ def read_initcm(nsim, srcdir, fname="clump_{}_cm.npy"):
return None
def halfwidth_select(hw, particles):
def halfwidth_mask(pos, hw):
"""
Select particles that in a cube of size `2 hw`, centered at the origin.
Note that this directly modifies the original array and throws away
particles outside the central region.
Mask of particles in a region of width `2 hw, centered at the origin.
Parameters
----------
pos : 2-dimensional array of shape `(nparticles, 3)`
Particle positions, in box units.
hw : float
Central region halfwidth.
particles : structured array
Particle array with keys `x`, `y`, `z`.
Central region half-width.
Returns
-------
particles : structured array
The modified particle array.
mask : 1-dimensional boolean array of shape `(nparticles, )`
"""
assert 0 < hw < 0.5
mask = ((0.5 - hw < particles['x']) & (particles['x'] < 0.5 + hw)
& (0.5 - hw < particles['y']) & (particles['y'] < 0.5 + hw)
& (0.5 - hw < particles['z']) & (particles['z'] < 0.5 + hw))
# Subselect the particles
particles = particles[mask]
# Rescale to range [0, 1]
for p in ('x', 'y', 'z'):
particles[p] = (particles[p] - 0.5 + hw) / (2 * hw)
return particles
return numpy.all((0.5 - hw < pos) & (pos < 0.5 + hw), axis=1)
def load_clump_particles(clid, particles, clump_map, clid2map):

12
csiborgtools/read/utils.py
View file

@ -80,7 +80,8 @@ def radec_to_cartesian(X, isdeg=True):
return dist * numpy.vstack([x, y, z]).T
def real2redshift(pos, vel, origin, box, in_box_units, make_copy=True):
def real2redshift(pos, vel, origin, box, in_box_units, periodic_wrap=True,
make_copy=True):
r"""
Convert real-space position to redshift space position.
@ -99,6 +100,9 @@ def real2redshift(pos, vel, origin, box, in_box_units, make_copy=True):
to be in :math:`\mathrm{Mpc}`, velocity in
:math:`\mathrm{km} \mathrm{s}^{-1}` and math:`h=0.705`, or otherwise
matching the box.
periodic_wrap : bool, optional
Whether to wrap around the box, particles may be outside the default
bounds once RSD is applied.
make_copy : bool, optional
Whether to make a copy of `pos` before modifying it.
@ -122,6 +126,12 @@ def real2redshift(pos, vel, origin, box, in_box_units, make_copy=True):
for i in range(3):
pos[:, i] += origin[i]
if periodic_wrap:
boxsize = 1. if in_box_units else box.box2mpc(1.)
# Wrap around the box: x > 1 -> x - 1, x < 0 -> x + 1
pos[pos > boxsize] -= boxsize
pos[pos < 0] += boxsize
return pos

122
notebooks/fits.ipynb
View file

File diff suppressed because one or more lines are too long

69
scripts/field_density.py Normal file
View file

@ -0,0 +1,69 @@
# 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 the density fields on CSiBORG simulations in the 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("--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)

127
scripts/field_prop.py
View file

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