mirror of
https://github.com/Richard-Sti/csiborgtools.git
synced 2024-12-22 18:48:01 +00:00
Gravitational potential calculation (#60)
* Rename file * add velocity field * Add velcoity field * Move smoothening * Add verbosity flags * remove blank * Simplify paths * Add potential calculation * Update paths * Add potential field * Add potential calculation * Move away sky matching * Move interpolation functions * Update nbs
This commit is contained in:
Richard Stiskalek
GitHub
parent
26ef1661a4
commit
b3fd14b81f
10 changed files with 12509 additions and 162 deletions
|
|
@ -17,6 +17,8 @@ from warnings import warn
|
|||
try:
|
||||
import MAS_library as MASL # noqa
|
||||
|
||||
from .density import DensityField # noqa
|
||||
from .density import DensityField, PotentialField, VelocityField # noqa
|
||||
from .interp import evaluate_cartesian, evaluate_sky, make_sky # noqa
|
||||
from .utils import smoothen_field # noqa
|
||||
except ImportError:
|
||||
warn("MAS_library not found, `DensityField` will not be available", UserWarning) # noqa
|
||||
|
|
|
|||
|
|
@ -14,16 +14,18 @@
|
|||
# 51 Franklin Street, Fifth Floor, Boston, MA 02110-1301, USA.
|
||||
"""
|
||||
Density field and cross-correlation calculations.
|
||||
|
||||
TODO:
|
||||
- [ ] Project the velocity field along the line of sight.
|
||||
"""
|
||||
from abc import ABC
|
||||
|
||||
import MAS_library as MASL
|
||||
import numpy
|
||||
import smoothing_library as SL
|
||||
from tqdm import trange
|
||||
|
||||
from ..read.utils import real2redshift
|
||||
from .utils import force_single_precision
|
||||
from ..read.utils import radec_to_cartesian, real2redshift
|
||||
|
||||
|
||||
class BaseField(ABC):
|
||||
|
|
@ -80,101 +82,6 @@ class BaseField(ABC):
|
|||
assert MAS in ["NGP", "CIC", "TSC", "PCS"]
|
||||
self._MAS = MAS
|
||||
|
||||
def evaluate_cartesian(self, *fields, pos):
|
||||
"""
|
||||
Evaluate a scalar field at Cartesian coordinates using CIC
|
||||
interpolation.
|
||||
|
||||
Parameters
|
||||
----------
|
||||
field : (list of) 3-dimensional array of shape `(grid, grid, grid)`
|
||||
Fields to be interpolated.
|
||||
pos : 2-dimensional array of shape `(n_samples, 3)`
|
||||
Positions to evaluate the density field. Assumed to be in box
|
||||
units.
|
||||
|
||||
Returns
|
||||
-------
|
||||
interp_fields : (list of) 1-dimensional array of shape `(n_samples,).
|
||||
"""
|
||||
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, self.boxsize, pos, interp_fields[i])
|
||||
|
||||
if len(fields) == 1:
|
||||
return interp_fields[0]
|
||||
return interp_fields
|
||||
|
||||
def evaluate_sky(self, *fields, pos, isdeg=True):
|
||||
"""
|
||||
Evaluate the scalar fields at given distance, right ascension and
|
||||
declination. Assumes an observed in the centre of the box, with
|
||||
distance being in :math:`Mpc`. Uses CIC interpolation.
|
||||
|
||||
Parameters
|
||||
----------
|
||||
fields : (list of) 3-dimensional array of shape `(grid, grid, grid)`
|
||||
Field to be interpolated.
|
||||
pos : 2-dimensional array of shape `(n_samples, 3)`
|
||||
Spherical coordinates to evaluate the field. Columns are distance,
|
||||
right ascension, declination, respectively.
|
||||
isdeg : bool, optional
|
||||
Whether `ra` and `dec` are in degres. By default `True`.
|
||||
|
||||
Returns
|
||||
-------
|
||||
interp_fields : (list of) 1-dimensional array of shape `(n_samples,).
|
||||
"""
|
||||
pos = force_single_precision(pos, "pos")
|
||||
# We first calculate convert the distance to box coordinates and then
|
||||
# convert to Cartesian coordinates.
|
||||
X = numpy.copy(pos)
|
||||
X[:, 0] = self.box.mpc2box(X[:, 0])
|
||||
X = radec_to_cartesian(pos, isdeg)
|
||||
# Then we move the origin to match the box coordinates
|
||||
X -= 0.5
|
||||
return self.evaluate_field(*fields, pos=X)
|
||||
|
||||
def make_sky(self, field, angpos, dist, verbose=True):
|
||||
r"""
|
||||
Make a sky map of a scalar field. The observer is in the centre of the
|
||||
box the field is evaluated along directions `angpos`. Along each
|
||||
direction, the field is evaluated distances `dist_marg` and summed.
|
||||
Uses CIC interpolation.
|
||||
|
||||
Parameters
|
||||
----------
|
||||
field : 3-dimensional array of shape `(grid, grid, grid)`
|
||||
Field to be interpolated
|
||||
angpos : 2-dimensional arrays of shape `(ndir, 2)`
|
||||
Directions to evaluate the field. Assumed to be RA
|
||||
:math:`\in [0, 360]` and dec :math:`\in [-90, 90]` degrees,
|
||||
respectively.
|
||||
dist : 1-dimensional array
|
||||
Radial distances to evaluate the field.
|
||||
verbose : bool, optional
|
||||
Verbosity flag.
|
||||
|
||||
Returns
|
||||
-------
|
||||
interp_field : 1-dimensional array of shape `(n_pos, )`.
|
||||
"""
|
||||
assert angpos.ndim == 2 and dist.ndim == 1
|
||||
# We loop over the angular directions, at each step evaluating a vector
|
||||
# of distances. We pre-allocate arrays for speed.
|
||||
dir_loop = numpy.full((dist.size, 3), numpy.nan, dtype=numpy.float32)
|
||||
ndir = angpos.shape[0]
|
||||
out = numpy.zeros(ndir, numpy.nan, dtype=numpy.float32)
|
||||
for i in trange(ndir) if verbose else range(ndir):
|
||||
dir_loop[1, :] = angpos[i, 0]
|
||||
dir_loop[2, :] = angpos[i, 1]
|
||||
out[i] = numpy.sum(self.evaluate_sky(field, dir_loop, isdeg=True))
|
||||
return out
|
||||
|
||||
|
||||
###############################################################################
|
||||
# Density field calculation #
|
||||
|
|
@ -183,7 +90,7 @@ class BaseField(ABC):
|
|||
|
||||
class DensityField(BaseField):
|
||||
r"""
|
||||
Density field calculations. Based primarily on routines of Pylians [1].
|
||||
Density field calculation. Based primarily on routines of Pylians [1].
|
||||
|
||||
Parameters
|
||||
----------
|
||||
|
|
@ -203,30 +110,6 @@ class DensityField(BaseField):
|
|||
self.box = box
|
||||
self.MAS = MAS
|
||||
|
||||
def smoothen(self, field, smooth_scale, threads=1):
|
||||
"""
|
||||
Smooth a field with a Gaussian filter.
|
||||
|
||||
Parameters
|
||||
----------
|
||||
field : 3-dimensional array of shape `(grid, grid, grid)`
|
||||
Field to be smoothed.
|
||||
smooth_scale : float, optional
|
||||
Gaussian kernal scale to smoothen the density field, in box units.
|
||||
threads : int, optional
|
||||
Number of threads. By default 1.
|
||||
|
||||
Returns
|
||||
-------
|
||||
smoothed_field : 3-dimensional array of shape `(grid, grid, grid)`
|
||||
"""
|
||||
filter_kind = "Gaussian"
|
||||
grid = field.shape[0]
|
||||
# FFT of the filter
|
||||
W_k = SL.FT_filter(self.boxsize, smooth_scale, grid, filter_kind,
|
||||
threads)
|
||||
return SL.field_smoothing(field, W_k, threads)
|
||||
|
||||
def overdensity_field(self, delta):
|
||||
r"""
|
||||
Calculate the overdensity field from the density field.
|
||||
|
|
@ -310,6 +193,97 @@ class DensityField(BaseField):
|
|||
return rho
|
||||
|
||||
|
||||
###############################################################################
|
||||
# Density field calculation #
|
||||
###############################################################################
|
||||
|
||||
|
||||
class VelocityField(BaseField):
|
||||
r"""
|
||||
Velocity field calculation. Based primarily on routines of Pylians [1].
|
||||
|
||||
Parameters
|
||||
----------
|
||||
box : :py:class:`csiborgtools.read.BoxUnits`
|
||||
The simulation box information and transformations.
|
||||
MAS : str
|
||||
Mass assignment scheme. Options are Options are: 'NGP' (nearest grid
|
||||
point), 'CIC' (cloud-in-cell), 'TSC' (triangular-shape cloud), 'PCS'
|
||||
(piecewise cubic spline).
|
||||
|
||||
References
|
||||
----------
|
||||
[1] https://pylians3.readthedocs.io/
|
||||
"""
|
||||
|
||||
def __init__(self, box, MAS):
|
||||
self.box = box
|
||||
self.MAS = MAS
|
||||
|
||||
def __call__(self, parts, grid, mpart, flip_xz=True, nbatch=30,
|
||||
verbose=True):
|
||||
"""
|
||||
Calculate the velocity 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.
|
||||
|
||||
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.
|
||||
mpart : float
|
||||
Particle mass.
|
||||
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
|
||||
-------
|
||||
rho_vel : 3-dimensional array of shape `(3, grid, grid, grid)`.
|
||||
Velocity field along each axis.
|
||||
|
||||
References
|
||||
----------
|
||||
[1] https://pylians3.readthedocs.io/
|
||||
[2] https://github.com/franciscovillaescusa/Pylians3/blob/master
|
||||
/library/MAS_library/MAS_library.pyx
|
||||
"""
|
||||
rho_velx = numpy.zeros((grid, grid, grid), dtype=numpy.float32)
|
||||
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]
|
||||
|
||||
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 *= mass.reshape(-1, 1) / mpart
|
||||
|
||||
for i in range(3):
|
||||
MASL.MA(pos, rho_vel[i], self.boxsize, self.MAS, W=vel[i, :],
|
||||
verbose=False)
|
||||
if end == nparts:
|
||||
break
|
||||
start = end
|
||||
|
||||
return numpy.stack(rho_vel)
|
||||
|
||||
|
||||
###############################################################################
|
||||
# Potential field calculation #
|
||||
###############################################################################
|
||||
|
|
@ -377,6 +351,8 @@ class TidalTensorField(BaseField):
|
|||
Calculate eigenvalues of the tidal tensor field, sorted in increasing
|
||||
order.
|
||||
|
||||
TODO: evaluate this on a grid instead.
|
||||
|
||||
Parameters
|
||||
----------
|
||||
tidal_tensor : :py:class:`MAS_library.tidal_tensor`
|
||||
|
|
|
|||
124
csiborgtools/field/interp.py
Normal file
124
csiborgtools/field/interp.py
Normal file
|
|
@ -0,0 +1,124 @@
|
|||
# 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.
|
||||
"""
|
||||
Tools for interpolating 3D fields at arbitrary positions.
|
||||
"""
|
||||
import MAS_library as MASL
|
||||
import numpy
|
||||
from tqdm import trange
|
||||
|
||||
from ..read.utils import radec_to_cartesian
|
||||
from .utils import force_single_precision
|
||||
|
||||
|
||||
def evaluate_cartesian(*fields, pos):
|
||||
"""
|
||||
Evaluate a scalar field at Cartesian coordinates using CIC
|
||||
interpolation.
|
||||
|
||||
Parameters
|
||||
----------
|
||||
field : (list of) 3-dimensional array of shape `(grid, grid, grid)`
|
||||
Fields to be interpolated.
|
||||
pos : 2-dimensional array of shape `(n_samples, 3)`
|
||||
Positions to evaluate the density field. Assumed to be in box
|
||||
units.
|
||||
|
||||
Returns
|
||||
-------
|
||||
interp_fields : (list of) 1-dimensional array of shape `(n_samples,).
|
||||
"""
|
||||
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 len(fields) == 1:
|
||||
return interp_fields[0]
|
||||
return interp_fields
|
||||
|
||||
|
||||
def evaluate_sky(*fields, pos, box, isdeg=True):
|
||||
"""
|
||||
Evaluate the scalar fields at given distance, right ascension and
|
||||
declination. Assumes an observed in the centre of the box, with
|
||||
distance being in :math:`Mpc`. Uses CIC interpolation.
|
||||
|
||||
Parameters
|
||||
----------
|
||||
fields : (list of) 3-dimensional array of shape `(grid, grid, grid)`
|
||||
Field to be interpolated.
|
||||
pos : 2-dimensional array of shape `(n_samples, 3)`
|
||||
Spherical coordinates to evaluate the field. Columns are distance,
|
||||
right ascension, declination, respectively.
|
||||
box : :py:class:`csiborgtools.read.BoxUnits`
|
||||
The simulation box information and transformations.
|
||||
isdeg : bool, optional
|
||||
Whether `ra` and `dec` are in degres. By default `True`.
|
||||
|
||||
Returns
|
||||
-------
|
||||
interp_fields : (list of) 1-dimensional array of shape `(n_samples,).
|
||||
"""
|
||||
pos = force_single_precision(pos, "pos")
|
||||
# We first calculate convert the distance to box coordinates and then
|
||||
# convert to Cartesian coordinates.
|
||||
X = numpy.copy(pos)
|
||||
X[:, 0] = box.mpc2box(X[:, 0])
|
||||
X = radec_to_cartesian(pos, isdeg)
|
||||
# Then we move the origin to match the box coordinates
|
||||
X -= 0.5
|
||||
return evaluate_cartesian(*fields, pos=X)
|
||||
|
||||
|
||||
def make_sky(field, angpos, dist, verbose=True):
|
||||
r"""
|
||||
Make a sky map of a scalar field. The observer is in the centre of the
|
||||
box the field is evaluated along directions `angpos`. Along each
|
||||
direction, the field is evaluated distances `dist_marg` and summed.
|
||||
Uses CIC interpolation.
|
||||
|
||||
Parameters
|
||||
----------
|
||||
field : 3-dimensional array of shape `(grid, grid, grid)`
|
||||
Field to be interpolated
|
||||
angpos : 2-dimensional arrays of shape `(ndir, 2)`
|
||||
Directions to evaluate the field. Assumed to be RA
|
||||
:math:`\in [0, 360]` and dec :math:`\in [-90, 90]` degrees,
|
||||
respectively.
|
||||
dist : 1-dimensional array
|
||||
Radial distances to evaluate the field.
|
||||
verbose : bool, optional
|
||||
Verbosity flag.
|
||||
|
||||
Returns
|
||||
-------
|
||||
interp_field : 1-dimensional array of shape `(n_pos, )`.
|
||||
"""
|
||||
assert angpos.ndim == 2 and dist.ndim == 1
|
||||
# We loop over the angular directions, at each step evaluating a vector
|
||||
# of distances. We pre-allocate arrays for speed.
|
||||
dir_loop = numpy.full((dist.size, 3), numpy.nan, dtype=numpy.float32)
|
||||
ndir = angpos.shape[0]
|
||||
out = numpy.zeros(ndir, numpy.nan, dtype=numpy.float32)
|
||||
for i in trange(ndir) if verbose else range(ndir):
|
||||
dir_loop[1, :] = angpos[i, 0]
|
||||
dir_loop[2, :] = angpos[i, 1]
|
||||
out[i] = numpy.sum(evaluate_sky(field, dir_loop, isdeg=True))
|
||||
return out
|
||||
|
|
@ -18,6 +18,7 @@ Utility functions for the field module.
|
|||
from warnings import warn
|
||||
|
||||
import numpy
|
||||
import smoothing_library as SL
|
||||
|
||||
|
||||
def force_single_precision(x, name):
|
||||
|
|
@ -40,3 +41,27 @@ def force_single_precision(x, name):
|
|||
warn(f"Converting `{name}` to float32.", UserWarning, stacklevel=1)
|
||||
x = x.astype(numpy.float32)
|
||||
return x
|
||||
|
||||
|
||||
def smoothen_field(field, smooth_scale, boxsize, threads=1):
|
||||
"""
|
||||
Smooth a field with a Gaussian filter.
|
||||
|
||||
Parameters
|
||||
----------
|
||||
field : 3-dimensional array of shape `(grid, grid, grid)`
|
||||
Field to be smoothed.
|
||||
smooth_scale : float, optional
|
||||
Gaussian kernal scale to smoothen the density field, in box units.
|
||||
boxsize : float
|
||||
Size of the box.
|
||||
threads : int, optional
|
||||
Number of threads. By default 1.
|
||||
|
||||
Returns
|
||||
-------
|
||||
smoothed_field : 3-dimensional array of shape `(grid, grid, grid)`
|
||||
"""
|
||||
W_k = SL.FT_filter(boxsize, smooth_scale, field.shape[0], "Gaussian",
|
||||
threads)
|
||||
return SL.field_smoothing(field, W_k, threads)
|
||||
|
|
|
|||
|
|
@ -476,15 +476,21 @@ class NPairsOverlap:
|
|||
List of cross simulation halo catalogues.
|
||||
paths : py:class`csiborgtools.read.CSiBORGPaths`
|
||||
CSiBORG paths object.
|
||||
verbose : bool, optional
|
||||
Verbosity flag for loading the overlap objects.
|
||||
"""
|
||||
_pairs = None
|
||||
|
||||
def __init__(self, cat0, catxs, paths):
|
||||
self._pairs = [PairOverlap(cat0, catx, paths) for catx in catxs]
|
||||
def __init__(self, cat0, catxs, paths, verbose=True):
|
||||
pairs = [None] * len(catxs)
|
||||
for i, catx in enumerate(tqdm(catxs) if verbose else catxs):
|
||||
pairs[i] = PairOverlap(cat0, catx, paths)
|
||||
|
||||
def summed_overlap(self, from_smoothed, verbose=False):
|
||||
self._pairs = pairs
|
||||
|
||||
def summed_overlap(self, from_smoothed, verbose=True):
|
||||
"""
|
||||
Calcualte summed overlap of each halo in the reference simulation with
|
||||
Calculate summed overlap of each halo in the reference simulation with
|
||||
the cross simulations.
|
||||
|
||||
Parameters
|
||||
|
|
@ -503,7 +509,7 @@ class NPairsOverlap:
|
|||
out[i] = pair.summed_overlap(from_smoothed)
|
||||
return numpy.vstack(out).T
|
||||
|
||||
def prob_nomatch(self, from_smoothed, verbose=False):
|
||||
def prob_nomatch(self, from_smoothed, verbose=True):
|
||||
"""
|
||||
Probability of no match for each halo in the reference simulation with
|
||||
the cross simulation.
|
||||
|
|
@ -526,7 +532,7 @@ class NPairsOverlap:
|
|||
|
||||
def counterpart_mass(self, from_smoothed, overlap_threshold=0.,
|
||||
in_log=False, mass_kind="totpartmass",
|
||||
return_full=True, verbose=False):
|
||||
return_full=False, verbose=True):
|
||||
"""
|
||||
Calculate the expected counterpart mass of each halo in the reference
|
||||
simulation from the crossed simulation.
|
||||
|
|
@ -549,7 +555,7 @@ class NPairsOverlap:
|
|||
Whether to return the full results of matching each pair or
|
||||
calculate summary statistics by Gaussian averaging.
|
||||
verbose : bool, optional
|
||||
Verbosity flag. By default `False`.
|
||||
Verbosity flag.
|
||||
|
||||
Returns
|
||||
-------
|
||||
|
|
|
|||
|
|
@ -322,30 +322,35 @@ class CSiBORGPaths:
|
|||
fname = f"parts_{str(nsim).zfill(5)}.h5"
|
||||
return join(fdir, fname)
|
||||
|
||||
def density_field_path(self, MAS, nsim, in_rsp):
|
||||
def field_path(self, kind, MAS, grid, nsim, in_rsp):
|
||||
"""
|
||||
Path to the files containing the calculated density fields.
|
||||
|
||||
Parameters
|
||||
----------
|
||||
kind : str
|
||||
Field type. Must be one of: `density`, `velocity`, `potential`.
|
||||
MAS : str
|
||||
Mass-assignment scheme. Currently only SPH is supported.
|
||||
Mass-assignment scheme.
|
||||
grid : int
|
||||
Grid size.
|
||||
nsim : int
|
||||
IC realisation index.
|
||||
in_rsp : bool
|
||||
Whether the density field is calculated in redshift space.
|
||||
Whether the calculation is performed in redshift space.
|
||||
|
||||
Returns
|
||||
-------
|
||||
path : str
|
||||
"""
|
||||
fdir = join(self.postdir, "environment")
|
||||
assert kind in ["density", "velocity", "potential"]
|
||||
if not isdir(fdir):
|
||||
makedirs(fdir)
|
||||
warn(f"Created directory `{fdir}`.", UserWarning, stacklevel=1)
|
||||
fname = f"density_{MAS}_{str(nsim).zfill(5)}.npy"
|
||||
if in_rsp:
|
||||
fname = fname.replace("density", "density_rsp")
|
||||
kind = kind + "_rsp"
|
||||
fname = f"{kind}_{MAS}_{str(nsim).zfill(5)}_grid{grid}.npy"
|
||||
return join(fdir, fname)
|
||||
|
||||
def knnauto_path(self, run, nsim=None):
|
||||
|
|
|
|||
File diff suppressed because one or more lines are too long
12025
notebooks/matching.ipynb
Normal file
12025
notebooks/matching.ipynb
Normal file
File diff suppressed because one or more lines are too long
78
scripts/field_derived.py
Normal file
78
scripts/field_derived.py
Normal file
|
|
@ -0,0 +1,78 @@
|
|||
# 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.CSiBORGPaths(**csiborgtools.paths_glamdring)
|
||||
|
||||
if args.ics is None or args.ics[0] == -1:
|
||||
ics = paths.get_ics()
|
||||
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.BoxUnits(nsnap, nsim, paths)
|
||||
density_gen = csiborgtools.field.DensityField(box, args.MAS)
|
||||
|
||||
rho = numpy.load(paths.field_path("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_path("potential", args.MAS, args.grid, nsim,
|
||||
args.in_rsp)
|
||||
print(f"{datetime.now()}: rank {rank} saving output to `{fout}`.")
|
||||
numpy.save(fout, field)
|
||||
|
|
@ -38,12 +38,16 @@ 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=["density", "velocity"],
|
||||
help="Calculate the density or velocity field?")
|
||||
parser.add_argument("--MAS", type=str, choices=["NGP", "CIC", "TSC", "PCS"],
|
||||
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 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)
|
||||
mpart = 1.1641532e-10 # Particle mass in CSiBORG simulations.
|
||||
|
||||
if args.ics is None or args.ics[0] == -1:
|
||||
ics = paths.get_ics()
|
||||
|
|
@ -59,11 +63,14 @@ for i in csiborgtools.fits.split_jobs(len(ics), nproc)[rank]:
|
|||
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)
|
||||
if args.kind == "density":
|
||||
gen = csiborgtools.field.DensityField(box, args.MAS)
|
||||
field = gen(parts, args.grid, in_rsp=args.in_rsp, verbose=verbose)
|
||||
else:
|
||||
gen = csiborgtools.field.VelocityField(box, args.MAS)
|
||||
field = gen(parts, args.grid, mpart, verbose=verbose)
|
||||
|
||||
fout = paths.density_field_path(args.MAS, nsim, args.in_rsp)
|
||||
fout = paths.field_path(args.kind, args.MAS, args.grid, nsim, args.in_rsp)
|
||||
print(f"{datetime.now()}: rank {rank} saving output to `{fout}`.")
|
||||
numpy.save(fout, rho)
|
||||
numpy.save(fout, field)
|
||||
Loading…
Reference in a new issue