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 2023-05-12 13:07:58 +01:00 committed by GitHub
parent 26ef1661a4
commit b3fd14b81f
No known key found for this signature in database
GPG Key ID: 4AEE18F83AFDEB23
10 changed files with 12509 additions and 162 deletions

View File

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

View File

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

View 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

View File

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

View File

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

View File

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

File diff suppressed because one or more lines are too long

78
scripts/field_derived.py Normal file
View 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)

View File

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