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

New RSP density distinction (#72)

Browse source 
* Edit docs

* Fix kernel units

* Add NGP interpolation

* Add Vobs but not implemented

* Add BORG density option

* Add BORG mcmc path

* Organise imports

* Add new density field distinction
This commit is contained in:
Richard Stiskalek 2023-06-26 20:41:07 +01:00 committed by GitHub
parent cd6d448874
commit de7def61d5
No known key found for this signature in database
GPG key ID: 4AEE18F83AFDEB23
7 changed files with 141 additions and 74 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

2
csiborgtools/__init__.py
View file

@ -17,6 +17,8 @@ from csiborgtools import clustering, field, fits, match, read # noqa
# Arguments to csiborgtools.read.Paths.
paths_glamdring = {"srcdir": "/mnt/extraspace/hdesmond/",
"postdir": "/mnt/extraspace/rstiskalek/CSiBORG/",
"BORG_dir": "/mnt/extraspace/rstiskalek/BORG/",
"BORG_final_density": "/users/hdesmond/BORG_final/",
"quijote_dir": "/mnt/extraspace/rstiskalek/Quijote",
}

5
csiborgtools/field/__init__.py
View file

@ -17,9 +17,10 @@ from warnings import warn
try:
import MAS_library as MASL # noqa
from .density import DensityField, PotentialField, VelocityField, TidalTensorField # noqa
from .density import (DensityField, PotentialField, # noqa
TidalTensorField, VelocityField)
from .interp import (evaluate_cartesian, evaluate_sky, field2rsp, # noqa
make_sky, fill_outside) # noqa
fill_outside, make_sky, observer_vobs)
from .utils import nside2radec, smoothen_field # noqa
except ImportError:
warn("MAS_library not found, `DensityField` will not be available", UserWarning) # noqa

19
csiborgtools/field/density.py
View file

@ -17,13 +17,13 @@ Density field and cross-correlation calculations.
"""
from abc import ABC
import numpy
import MAS_library as MASL
import numpy
from numba import jit
from tqdm import trange
from ..read.utils import real2redshift
from .interp import divide_nonzero
from .utils import force_single_precision
@ -249,7 +249,7 @@ class VelocityField(BaseField):
/ numpy.sqrt(px**2 + py**2 + pz**2))
return radvel
def __call__(self, parts, grid, mpart, flip_xz=True, nbatch=30,
def __call__(self, parts, grid, flip_xz=True, nbatch=30,
verbose=True):
"""
Calculate the velocity field using a Pylians routine [1, 2].
@ -263,8 +263,6 @@ class VelocityField(BaseField):
Columns are: `x`, `y`, `z`, `vx`, `vy`, `vz`, `M`.
grid : int
Grid size.
mpart : float
Particle mass.
flip_xz : bool, optional
Whether to flip the `x` and `z` coordinates.
nbatch : int, optional
@ -287,6 +285,7 @@ class VelocityField(BaseField):
rho_vely = numpy.zeros((grid, grid, grid), dtype=numpy.float32)
rho_velz = numpy.zeros((grid, grid, grid), dtype=numpy.float32)
rho_vel = [rho_velx, rho_vely, rho_velz]
cellcounts = numpy.zeros((grid, grid, grid), dtype=numpy.float32)
nparts = parts.shape[0]
batch_size = nparts // nbatch
@ -302,16 +301,22 @@ class VelocityField(BaseField):
if flip_xz:
pos[:, [0, 2]] = pos[:, [2, 0]]
vel[:, [0, 2]] = vel[:, [2, 0]]
vel *= mass.reshape(-1, 1) / mpart
vel *= mass.reshape(-1, 1)
for i in range(3):
MASL.MA(pos, rho_vel[i], self.boxsize, self.MAS, W=vel[:, i],
verbose=False)
MASL.MA(pos, cellcounts, self.boxsize, self.MAS, W=mass,
verbose=False)
if end == nparts:
break
start = end
return numpy.stack(rho_vel)
for i in range(3):
divide_nonzero(rho_vel[i], cellcounts)
return numpy.stack([rho_velx, rho_vely, rho_velz])
###############################################################################

48
csiborgtools/field/interp.py
View file

@ -24,10 +24,9 @@ from ..read.utils import radec_to_cartesian, real2redshift
from .utils import force_single_precision
def evaluate_cartesian(*fields, pos):
def evaluate_cartesian(*fields, pos, interp="CIC"):
"""
Evaluate a scalar field at Cartesian coordinates using CIC
interpolation.
Evaluate a scalar field at Cartesian coordinates.
Parameters
----------
@ -36,19 +35,29 @@ def evaluate_cartesian(*fields, pos):
pos : 2-dimensional array of shape `(n_samples, 3)`
Positions to evaluate the density field. Assumed to be in box
units.
interp : str, optional
Interpolation method. Can be either `CIC` or `NGP`.
Returns
-------
interp_fields : (list of) 1-dimensional array of shape `(n_samples,).
"""
assert interp in ["CIC", "NGP"]
boxsize = 1.
pos = force_single_precision(pos, "pos")
nsamples = pos.shape[0]
interp_fields = [numpy.full(nsamples, numpy.nan, dtype=numpy.float32)
for __ in range(len(fields))]
for i, field in enumerate(fields):
MASL.CIC_interp(field, boxsize, pos, interp_fields[i])
if interp == "CIC":
for i, field in enumerate(fields):
MASL.CIC_interp(field, boxsize, pos, interp_fields[i])
else:
pos = numpy.floor(pos * fields[0].shape[0]).astype(numpy.int32)
for i, field in enumerate(fields):
for j in range(nsamples):
interp_fields[i][j] = field[pos[j, 0], pos[j, 1], pos[j, 2]]
if len(fields) == 1:
return interp_fields[0]
@ -168,8 +177,30 @@ def divide_nonzero(field0, field1):
field0[i, j, k] /= field1[i, j, k]
def field2rsp(*fields, parts, box, nbatch=30, flip_partsxz=True, init_value=0.,
verbose=True):
def observer_vobs(velocity_field):
"""
Calculate the observer velocity from a velocity field. Assumes the observer
is in the centre of the box.
Parameters
----------
velocity_field : 4-dimensional array of shape `(3, grid, grid, grid)`
Velocity field to calculate the observer velocity from.
Returns
-------
vobs : 1-dimensional array of shape `(3,)`
Observer velocity in units of `velocity_field`.
"""
pos = numpy.asanyarray([0.5, 0.5, 0.5]).reshape(1, 3)
vobs = numpy.full(3, numpy.nan, dtype=numpy.float32)
for i in range(3):
vobs[i] = evaluate_cartesian(velocity_field[i, ...], pos=pos)[0]
return vobs
def field2rsp(*fields, parts, vobs, box, nbatch=30, flip_partsxz=True,
init_value=0., verbose=True):
"""
Forward model real space scalar fields to redshift space. Attaches field
values to particles, those are then moved to redshift space and from their
@ -182,6 +213,8 @@ def field2rsp(*fields, parts, box, nbatch=30, flip_partsxz=True, init_value=0.,
parts : 2-dimensional array of shape `(n_parts, 6)`
Particle positions and velocities in real space. Must be organised as
`x, y, z, vx, vy, vz`.
vobs : 1-dimensional array of shape `(3,)`
Observer velocity in units matching `parts`.
box : :py:class:`csiborgtools.read.CSiBORGBox`
The simulation box information and transformations.
nbatch : int, optional
@ -199,6 +232,7 @@ def field2rsp(*fields, parts, box, nbatch=30, flip_partsxz=True, init_value=0.,
-------
rsp_fields : (list of) 3-dimensional array of shape `(grid, grid, grid)`
"""
raise NotImplementedError("Figure out what to do with Vobs.")
nfields = len(fields)
# Check that all fields have the same shape.
if nfields > 1:

44
csiborgtools/read/paths.py
View file

@ -31,16 +31,21 @@ class Paths:
Path to the folder where the RAMSES outputs are stored.
postdir: str, optional
Path to the folder where post-processed files are stored.
borg_dir : str, optional
Path to the folder where BORG MCMC chains are stored.
quiote_dir : str, optional
Path to the folder where Quijote simulations are stored.
"""
_srcdir = None
_postdir = None
_borg_dir = None
_quijote_dir = None
def __init__(self, srcdir=None, postdir=None, quijote_dir=None):
def __init__(self, srcdir=None, postdir=None, borg_dir=None,
quijote_dir=None):
self.srcdir = srcdir
self.postdir = postdir
self.borg_dir = borg_dir
self.quijote_dir = quijote_dir
@staticmethod
@ -68,6 +73,26 @@ class Paths:
self._check_directory(path)
self._srcdir = path
@property
def borg_dir(self):
"""
Path to the folder where BORG MCMC chains are stored.
Returns
-------
path : str
"""
if self._borg_dir is None:
raise ValueError("`borg_dir` is not set!")
return self._borg_dir
@borg_dir.setter
def borg_dir(self, path):
if path is None:
return
self._check_directory(path)
self._borg_dir = path
@property
def quijote_dir(self):
"""
@ -146,6 +171,21 @@ class Paths:
return nsim
return f"{str(nobs).zfill(2)}{str(nsim).zfill(3)}"
def borg_mcmc(self, nsim):
"""
Path to the BORG MCMC chain file.
Parameters
----------
nsim : int
IC realisation index.
Returns
-------
path : str
"""
return join(self.borg_dir, "mcmc", f"mcmc_{nsim}.h5")
def mmain(self, nsnap, nsim):
"""
Path to the `mmain` CSiBORG files of summed substructure.
@ -373,7 +413,7 @@ class Paths:
IC realisation index.
in_rsp : bool
Whether the calculation is performed in redshift space.
smooth_scale : float
smooth_scale : float, optional
Smoothing scale in :math:`\mathrm{Mpc}/h`
Returns

76
scripts/field_prop.py
View file

@ -61,54 +61,30 @@ def density_field(nsim, parser_args, to_save=True):
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)
if parser_args.kind == "density":
field = gen(parts, parser_args.grid, in_rsp=False,
verbose=parser_args.verbose)
if parser_args.in_rsp:
field = csiborgtools.field.field2rsp(*field, parts=parts, box=box,
verbose=parser_args.verbose)
else:
field = gen(parts, parser_args.grid, in_rsp=parser_args.in_rsp,
verbose=parser_args.verbose)
if parser_args.smooth_scale > 0:
field = csiborgtools.field.smoothen_field(
field, parser_args.smooth_scale, box.boxsize * box.h, threads=1)
if to_save:
fout = paths.field("density", parser_args.MAS, parser_args.grid,
nsim, parser_args.in_rsp)
fout = paths.field(parser_args.kind, 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, field)
return field
def density_field_smoothed(nsim, parser_args, to_save=True):
"""
Calculate the smoothed density field in the CSiBORG simulation. The
unsmoothed density field must already be precomputed.
Parameters
----------
nsim : int
Simulation index.
parser_args : argparse.Namespace
Parsed arguments.
to_save : bool, optional
Whether to save the output to disk.
Returns
-------
smoothed_density : 3-dimensional array
"""
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
rho = numpy.load(paths.field("density", parser_args.MAS, parser_args.grid,
nsim, in_rsp=False))
rho = csiborgtools.field.smoothen_field(rho, parser_args.smooth_scale,
box.boxsize, threads=1)
if to_save:
fout = paths.field("density", 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, rho)
return rho
###############################################################################
# Velocity field #
###############################################################################
@ -185,7 +161,7 @@ def potential_field(nsim, parser_args, to_save=True):
nsim, in_rsp=False))
if parser_args.smooth_scale > 0:
rho = csiborgtools.field.smoothen_field(rho, parser_args.smooth_scale,
box.boxsize, threads=1)
box.boxsize * box.h, threads=1)
rho = density_gen.overdensity_field(rho)
# Calculate the real space potentiel field
gen = csiborgtools.field.PotentialField(box, parser_args.MAS)
@ -281,7 +257,7 @@ def environment_field(nsim, parser_args, to_save=True):
nsim, in_rsp=False))
if parser_args.smooth_scale > 0:
rho = csiborgtools.field.smoothen_field(rho, parser_args.smooth_scale,
box.boxsize, threads=1)
box.boxsize * box.h, threads=1)
rho = density_gen.overdensity_field(rho)
# Calculate the real space tidal tensor field, delete overdensity.
if parser_args.verbose:
@ -339,28 +315,28 @@ if __name__ == "__main__":
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",
"environment"],
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("--in_rsp", type=lambda x: bool(strtobool(x)),
help="Calculate in RSP?")
parser.add_argument("--smooth_scale", type=float, default=0)
parser.add_argument("--smooth_scale", type=float, default=0,
help="Smoothing scale in Mpc/h.")
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):
if parser_args.kind == "density":
if parser_args.smooth_scale > 0:
density_field_smoothed(nsim, parser_args)
else:
density_field(nsim, parser_args)
if parser_args.kind == "density" or parser_args.kind == "rspdensity":
density_field(nsim, parser_args)
elif parser_args.kind == "velocity":
velocity_field(nsim, parser_args)
elif parser_args.kind == "radvel":

21
scripts_plots/plot_data.py
View file

@ -19,6 +19,7 @@ from argparse import ArgumentParser
import matplotlib as mpl
import matplotlib.pyplot as plt
import numpy
from h5py import File
import healpy
import scienceplots # noqa
@ -244,7 +245,8 @@ def load_field(kind, nsim, grid, MAS, in_rsp=False, smooth_scale=None):
def plot_projected_field(kind, nsim, grid, in_rsp, smooth_scale, MAS="PCS",
highres_only=True, slice_find=None, pdf=False):
vel_component=0, highres_only=True, slice_find=None,
pdf=False):
r"""
Plot the mean projected field, however can also plot a single slice.
@ -262,6 +264,8 @@ def plot_projected_field(kind, nsim, grid, in_rsp, smooth_scale, MAS="PCS",
Smoothing scale in :math:`\mathrm{Mpc} / h`.
MAS : str, optional
Mass assignment scheme.
vel_component : int, optional
Which velocity field component to plot.
highres_only : bool, optional
Whether to only plot the high-resolution region.
slice_find : float, optional
@ -283,12 +287,16 @@ def plot_projected_field(kind, nsim, grid, in_rsp, smooth_scale, MAS="PCS",
smooth_scale=smooth_scale)
density_gen = csiborgtools.field.DensityField(box, MAS)
field = density_gen.overdensity_field(field) + 1
elif kind == "borg_density":
field = File(paths.borg_mcmc(nsim), 'r')
field = field["scalars"]["BORG_final_density"][...]
else:
field = load_field(kind, nsim, grid, MAS=MAS, in_rsp=in_rsp,
smooth_scale=smooth_scale)
if kind == "velocity":
field = field[0, ...]
field = field[vel_component, ...]
field = box.box2vel(field)
if highres_only:
csiborgtools.field.fill_outside(field, numpy.nan, rmax=155.5,
@ -552,22 +560,23 @@ if __name__ == "__main__":
if False:
plot_hmf(pdf=False)
if True:
if False:
kind = "overdensity"
grid = 1024
plot_sky_distribution(kind, 7444, grid, nside=64,
plot_groups=False, dmin=45, dmax=60,
plot_halos=5e13, volume_weight=True)
if False:
kind = "density"
if True:
kind = "overdensity"
grid = 256
smooth_scale = 0
# plot_projected_field("overdensity", 7444, grid, in_rsp=True,
# highres_only=False)
plot_projected_field(kind, 7444, grid, in_rsp=False,
smooth_scale=smooth_scale, slice_find=0.5,
highres_only=False)
MAS="PCS",
highres_only=True)
if False:
paths = csiborgtools.read.Paths(**csiborgtools.paths_glamdring)