mirror of
https://github.com/Richard-Sti/csiborgtools.git
synced 2024-10-18 16:35:06 +02:00
Update field calculations (#51)
* Update density field routines * Add utils * Add possibility to return 2d array * Styling * Fixed density field * Fix bugs * add mpc2box * Fix evaluating sky * Fix sky map making * Rename file * Add paths if only positions * Add option to dump particles only * Add comments
This commit is contained in:
Richard Stiskalek
GitHub
parent
b5fefe4196
commit
553eec8228
@ -15,6 +15,7 @@
|
||||
"""
|
||||
Density field and cross-correlation calculations.
|
||||
"""
|
||||
from abc import ABC
|
||||
from warnings import warn
|
||||
|
||||
import MAS_library as MASL
|
||||
@ -23,73 +24,26 @@ 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
|
||||
|
||||
class DensityField:
|
||||
r"""
|
||||
Density field calculations. Based primarily on routines of Pylians [1].
|
||||
|
||||
Parameters
|
||||
----------
|
||||
particles : structured array
|
||||
Particle array. Must contain keys `['x', 'y', 'z', 'M']`. Particle
|
||||
coordinates are assumed to be :math:`\in [0, 1]` or in box units
|
||||
otherwise.
|
||||
boxsize : float
|
||||
Box length. Multiplies `particles` positions to fix the power spectum
|
||||
units.
|
||||
box : :py:class:`csiborgtools.units.BoxUnits`
|
||||
The simulation box information and transformations.
|
||||
MAS : str, optional
|
||||
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/
|
||||
"""
|
||||
_particles = None
|
||||
_boxsize = None
|
||||
class BaseField(ABC):
|
||||
"""Base class for density field calculations."""
|
||||
_box = None
|
||||
_MAS = None
|
||||
|
||||
def __init__(self, particles, boxsize, box, MAS="CIC"):
|
||||
self.particles = particles
|
||||
assert boxsize > 0
|
||||
self.boxsize = boxsize
|
||||
self.box = box
|
||||
assert MAS in ["NGP", "CIC", "TSC", "PCS"]
|
||||
self._MAS = MAS
|
||||
|
||||
@property
|
||||
def particles(self):
|
||||
"""
|
||||
Particles structured array.
|
||||
|
||||
Returns
|
||||
-------
|
||||
particles : structured array
|
||||
"""
|
||||
return self._particles
|
||||
|
||||
@particles.setter
|
||||
def particles(self, particles):
|
||||
"""Set `particles`, checking it has the right columns."""
|
||||
if any(p not in particles.dtype.names for p in ('x', 'y', 'z', 'M')):
|
||||
raise ValueError("`particles` must be a structured array "
|
||||
"containing `['x', 'y', 'z', 'M']`.")
|
||||
self._particles = particles
|
||||
|
||||
@property
|
||||
def boxsize(self):
|
||||
"""
|
||||
Box length. Determines the power spectrum units.
|
||||
Box size. Particle positions are always assumed to be in box units,
|
||||
therefore this is 1.
|
||||
|
||||
Returns
|
||||
-------
|
||||
boxsize : float
|
||||
"""
|
||||
return self._boxsize
|
||||
return 1.
|
||||
|
||||
@property
|
||||
def box(self):
|
||||
@ -119,323 +73,91 @@ class DensityField:
|
||||
-------
|
||||
MAS : str
|
||||
"""
|
||||
if self._MAS is None:
|
||||
raise ValueError("`mas` is not set.")
|
||||
return self._MAS
|
||||
|
||||
@staticmethod
|
||||
def _force_f32(x, name):
|
||||
if x.dtype != numpy.float32:
|
||||
warn("Converting `{}` to float32.".format(name), stacklevel=1)
|
||||
x = x.astype(numpy.float32)
|
||||
return x
|
||||
@MAS.setter
|
||||
def MAS(self, MAS):
|
||||
assert MAS in ["NGP", "CIC", "TSC", "PCS"]
|
||||
self._MAS = MAS
|
||||
|
||||
def density_field(self, grid, smooth_scale=None, verbose=True):
|
||||
def evaluate_cartesian(self, *fields, pos):
|
||||
"""
|
||||
Calculate the density field using a Pylians routine [1, 2]. Enforces
|
||||
float32 precision.
|
||||
|
||||
Parameters
|
||||
----------
|
||||
grid : int
|
||||
Grid size.
|
||||
smooth_scale : float, optional
|
||||
Scale to smoothen the density field, in units matching
|
||||
`self.boxsize`. By default no smoothing is applied.
|
||||
verbose : bool
|
||||
Verbosity flag.
|
||||
|
||||
Returns
|
||||
-------
|
||||
rho : 3-dimensional array of shape `(grid, grid, grid)`.
|
||||
|
||||
References
|
||||
----------
|
||||
[1] https://pylians3.readthedocs.io/
|
||||
[2] https://github.com/franciscovillaescusa/Pylians3/blob/master
|
||||
/library/MAS_library/MAS_library.pyx
|
||||
"""
|
||||
pos = numpy.vstack([self.particles[p] for p in ('x', 'y', 'z')]).T
|
||||
pos *= self.boxsize
|
||||
pos = self._force_f32(pos, "pos")
|
||||
weights = self._force_f32(self.particles['M'], 'M')
|
||||
|
||||
# Pre-allocate and do calculations
|
||||
rho = numpy.zeros((grid, grid, grid), dtype=numpy.float32)
|
||||
MASL.MA(pos, rho, self.boxsize, self.MAS, W=weights, verbose=verbose)
|
||||
if smooth_scale is not None:
|
||||
rho = self.smooth_field(rho, smooth_scale)
|
||||
return rho
|
||||
|
||||
def overdensity_field(self, grid, smooth_scale=None, verbose=True):
|
||||
r"""
|
||||
Calculate the overdensity field using Pylians routines.
|
||||
Defined as :math:`\rho/ <\rho> - 1`.
|
||||
|
||||
Parameters
|
||||
----------
|
||||
grid : int
|
||||
Grid size.
|
||||
smooth_scale : float, optional
|
||||
Scale to smoothen the density field, in units matching
|
||||
`self.boxsize`. By default no smoothing is applied.
|
||||
verbose : bool
|
||||
Verbosity flag.
|
||||
|
||||
Returns
|
||||
-------
|
||||
overdensity : 3-dimensional array of shape `(grid, grid, grid)`.
|
||||
"""
|
||||
# Get the overdensity
|
||||
delta = self.density_field(grid, smooth_scale, verbose)
|
||||
delta /= delta.mean()
|
||||
delta -= 1
|
||||
return delta
|
||||
|
||||
def potential_field(self, grid, smooth_scale=None, verbose=True):
|
||||
"""
|
||||
Calculate the potential field using Pylians routines.
|
||||
|
||||
Parameters
|
||||
----------
|
||||
grid : int
|
||||
Grid size.
|
||||
smooth_scale : float, optional
|
||||
Scale to smoothen the density field, in units matching
|
||||
`self.boxsize`. By default no smoothing is applied.
|
||||
verbose : bool
|
||||
Verbosity flag.
|
||||
|
||||
Returns
|
||||
-------
|
||||
potential : 3-dimensional array of shape `(grid, grid, grid)`.
|
||||
"""
|
||||
delta = self.overdensity_field(grid, smooth_scale, verbose)
|
||||
if verbose:
|
||||
print("Calculating potential from the overdensity..")
|
||||
return MASL.potential(
|
||||
delta, self.box._omega_m, self.box._aexp, self.MAS)
|
||||
|
||||
def gravitational_field(self, grid, smooth_scale=None, verbose=True):
|
||||
"""
|
||||
Calculate the gravitational vector field. Note that this method is
|
||||
only defined in a fork of `Pylians`.
|
||||
|
||||
Parameters
|
||||
----------
|
||||
grid : int
|
||||
Grid size.
|
||||
smooth_scale : float, optional
|
||||
Scale to smoothen the density field, in units matching
|
||||
`self.boxsize`. By default no smoothing is applied.
|
||||
verbose : bool
|
||||
Verbosity flag.
|
||||
|
||||
Returns
|
||||
-------
|
||||
grav_field_tensor : :py:class:`MAS_library.grav_field_tensor`
|
||||
Tidal tensor object, whose attributes `grav_field_tensor.gi`
|
||||
contain the relevant tensor components.
|
||||
"""
|
||||
delta = self.overdensity_field(grid, smooth_scale, verbose)
|
||||
return MASL.grav_field_tensor(
|
||||
delta, self.box._omega_m, self.box._aexp, self.MAS)
|
||||
|
||||
def tensor_field(self, grid, smooth_scale=None, verbose=True):
|
||||
"""
|
||||
Calculate the tidal tensor field.
|
||||
|
||||
Parameters
|
||||
----------
|
||||
grid : int
|
||||
Grid size.
|
||||
smooth_scale : float, optional
|
||||
Scale to smoothen the density field, in units matching
|
||||
`self.boxsize`. By default no smoothing is applied.
|
||||
verbose : bool, optional
|
||||
A verbosity flag.
|
||||
|
||||
Returns
|
||||
-------
|
||||
tidal_tensor : :py:class:`MAS_library.tidal_tensor`
|
||||
Tidal tensor object, whose attributes `tidal_tensor.Tij` contain
|
||||
the relevant tensor components.
|
||||
"""
|
||||
delta = self.overdensity_field(grid, smooth_scale, verbose)
|
||||
return MASL.tidal_tensor(
|
||||
delta, self.box._omega_m, self.box._aexp, self.MAS)
|
||||
|
||||
def auto_powerspectrum(self, grid, smooth_scale, verbose=True):
|
||||
"""
|
||||
Calculate the auto 1-dimensional power spectrum.
|
||||
|
||||
Parameters
|
||||
----------
|
||||
grid : int
|
||||
Grid size.
|
||||
smooth_scale : float, optional
|
||||
Scale to smoothen the density field, in units matching
|
||||
`self.boxsize`. By default no smoothing is applied.
|
||||
verbose : bool, optional
|
||||
Verbosity flag.
|
||||
|
||||
Returns
|
||||
-------
|
||||
pk : py:class`Pk_library.Pk`
|
||||
"""
|
||||
delta = self.overdensity_field(grid, smooth_scale, verbose)
|
||||
return PKL.Pk(
|
||||
delta, self.boxsize, axis=1, MAS=self.MAS, threads=1,
|
||||
verbose=verbose)
|
||||
|
||||
def smooth_field(self, field, smooth_scale, threads=1):
|
||||
"""
|
||||
Smooth a field with a Gaussian filter.
|
||||
|
||||
Parameters
|
||||
----------
|
||||
field : 3-dimensional array of shape `(grid, grid, grid)`
|
||||
The field to be smoothed.
|
||||
smooth_scale : float, optional
|
||||
Scale to smoothen the density field, in units matching
|
||||
`self.boxsize`. By default no smoothing is applied.
|
||||
threads : int, optional
|
||||
Number of threads. By default 1.
|
||||
|
||||
Returns
|
||||
-------
|
||||
smoothed_field : 3-dimensional array of shape `(grid, grid, grid)`
|
||||
"""
|
||||
Filter = "Gaussian"
|
||||
grid = field.shape[0]
|
||||
# FFT of the filter
|
||||
W_k = SL.FT_filter(self.boxsize, smooth_scale, grid, Filter, threads)
|
||||
return SL.field_smoothing(field, W_k, threads)
|
||||
|
||||
def evaluate_field(self, *field, pos):
|
||||
"""
|
||||
Evaluate the field at Cartesian coordinates using CIC interpolation.
|
||||
|
||||
Parameters
|
||||
----------
|
||||
field : (list of) 3-dimensional array of shape `(grid, grid, grid)`
|
||||
Density field that is to be interpolated.
|
||||
pos : 2-dimensional array of shape `(n_samples, 3)`
|
||||
Positions to evaluate the density field. The coordinates span range
|
||||
of [0, boxsize].
|
||||
|
||||
Returns
|
||||
-------
|
||||
interp_field : (list of) 1-dimensional array of shape `(n_samples,).
|
||||
"""
|
||||
self._force_f32(pos, "pos")
|
||||
|
||||
interp_field = [numpy.zeros(pos.shape[0], dtype=numpy.float32)
|
||||
for __ in range(len(field))]
|
||||
for i, f in enumerate(field):
|
||||
MASL.CIC_interp(f, self.boxsize, pos, interp_field[i])
|
||||
return interp_field
|
||||
|
||||
def evaluate_sky(self, *field, pos, isdeg=True):
|
||||
"""
|
||||
Evaluate the field at given distance, right ascension and declination.
|
||||
Assumes that the observed is in the centre of the box and uses CIC
|
||||
Evaluate a scalar field at Cartesian coordinates using CIC
|
||||
interpolation.
|
||||
|
||||
Parameters
|
||||
----------
|
||||
field : (list of) 3-dimensional array of shape `(grid, grid, grid)`
|
||||
Density field that is to be interpolated. Assumed to be defined
|
||||
on a Cartesian grid.
|
||||
Fields to be interpolated.
|
||||
pos : 2-dimensional array of shape `(n_samples, 3)`
|
||||
Spherical coordinates to evaluate the field. Should be distance,
|
||||
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_field : (list of) 1-dimensional array of shape `(n_samples,).
|
||||
interp_fields : (list of) 1-dimensional array of shape `(n_samples,).
|
||||
"""
|
||||
# TODO: implement this
|
||||
raise NotImplementedError("This method is not yet implemented.")
|
||||
# self._force_f32(pos, "pos")
|
||||
# X = numpy.vstack(
|
||||
# radec_to_cartesian(*(pos[:, i] for i in range(3)), isdeg)).T
|
||||
# X = X.astype(numpy.float32)
|
||||
# # Place the observer at the center of the box
|
||||
# X += 0.5 * self.boxsize
|
||||
# return self.evaluate_field(*field, pos=X)
|
||||
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)
|
||||
|
||||
@staticmethod
|
||||
def gravitational_field_norm(gx, gy, gz):
|
||||
"""
|
||||
Calculate the norm (magnitude) of a gravitational field.
|
||||
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
|
||||
----------
|
||||
gx, gy, gz : 1-dimensional arrays of shape `(n_samples,)`
|
||||
Gravitational field Cartesian components.
|
||||
|
||||
Returns
|
||||
-------
|
||||
g : 1-dimensional array of shape `(n_samples,)`
|
||||
"""
|
||||
return numpy.sqrt(gx * gx + gy * gy + gz * gz)
|
||||
|
||||
@staticmethod
|
||||
def tensor_field_eigvals(T00, T01, T02, T11, T12, T22):
|
||||
"""
|
||||
Calculate the eigenvalues of a symmetric tensor field. Eigenvalues are
|
||||
sorted in increasing order.
|
||||
|
||||
Parameters
|
||||
----------
|
||||
T00, T01, T02, T11, T12, T22 : 1-dim arrays of shape `(n_samples,)`
|
||||
Tensor field upper components evaluated for each sample.
|
||||
|
||||
Returns
|
||||
-------
|
||||
eigvals : 2-dimensional array of shape `(n_samples, 3)`
|
||||
"""
|
||||
n_samples = T00.size
|
||||
# Fill array of shape `(n_samples, 3, 3)` to calculate eigvals
|
||||
Teval = numpy.full((n_samples, 3, 3), numpy.nan, dtype=numpy.float32)
|
||||
Teval[:, 0, 0] = T00
|
||||
Teval[:, 0, 1] = T01
|
||||
Teval[:, 0, 2] = T02
|
||||
Teval[:, 1, 1] = T11
|
||||
Teval[:, 1, 2] = T12
|
||||
Teval[:, 2, 2] = T22
|
||||
|
||||
# Calculate the eigenvalues
|
||||
eigvals = numpy.full((n_samples, 3), numpy.nan, dtype=numpy.float32)
|
||||
for i in range(n_samples):
|
||||
eigvals[i, :] = numpy.linalg.eigvalsh(Teval[i, ...], 'U')
|
||||
eigvals[i, :] = numpy.sort(eigvals[i, :])
|
||||
|
||||
return eigvals
|
||||
|
||||
def make_sky_map(self, ra, dec, field, dist_marg, isdeg=True,
|
||||
verbose=True):
|
||||
"""
|
||||
Make a sky map of a density field. Places the observed in the center of
|
||||
the box and evaluates the field in directions `ra`, `dec`. At each such
|
||||
position evaluates the field at distances `dist_marg` and sums these
|
||||
interpolated values of the field.
|
||||
|
||||
NOTE: Supports only scalar fields.
|
||||
|
||||
Parameters
|
||||
----------
|
||||
ra, dec : 1-dimensional arrays of shape `(n_pos, )`
|
||||
Directions to evaluate the field. Assumes `dec` is in [-90, 90]
|
||||
degrees (or equivalently in radians).
|
||||
field : 3-dimensional array of shape `(grid, grid, grid)`
|
||||
Density field that is to be interpolated. Assumed to be defined
|
||||
on a Cartesian grid `[0, self.boxsize]^3`.
|
||||
dist_marg : 1-dimensional array
|
||||
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.
|
||||
isdeg : bool, optional
|
||||
Whether `ra` and `dec` are in degres. By default `True`.
|
||||
verbose : bool, optional
|
||||
Verbosity flag.
|
||||
|
||||
@ -443,23 +165,276 @@ class DensityField:
|
||||
-------
|
||||
interp_field : 1-dimensional array of shape `(n_pos, )`.
|
||||
"""
|
||||
# Angular positions at which to evaluate the field
|
||||
Nang = ra.size
|
||||
pos = numpy.vstack([ra, dec]).T
|
||||
|
||||
# Now loop over the angular positions, each time evaluating a vector
|
||||
# of distances. Pre-allocate arrays for speed
|
||||
ra_loop = numpy.ones_like(dist_marg)
|
||||
dec_loop = numpy.ones_like(dist_marg)
|
||||
pos_loop = numpy.ones((dist_marg.size, 3), dtype=numpy.float32)
|
||||
|
||||
out = numpy.zeros(Nang, dtype=numpy.float32)
|
||||
for i in trange(Nang) if verbose else range(Nang):
|
||||
# Get the position vector for this choice of theta, phi
|
||||
ra_loop[:] = pos[i, 0]
|
||||
dec_loop[:] = pos[i, 1]
|
||||
pos_loop[:] = numpy.vstack([dist_marg, ra_loop, dec_loop]).T
|
||||
# Evaluate and sum it up
|
||||
out[i] = numpy.sum(self.evaluate_sky(field, pos_loop, isdeg)[0, :])
|
||||
|
||||
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 #
|
||||
###############################################################################
|
||||
|
||||
|
||||
class DensityField(BaseField):
|
||||
r"""
|
||||
Density field calculations. Based primarily on routines of Pylians [1].
|
||||
|
||||
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
|
||||
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/
|
||||
"""
|
||||
_pos = None
|
||||
_mass = None
|
||||
|
||||
def __init__(self, pos, mass, box, MAS):
|
||||
self.pos = pos
|
||||
self.mass = mass
|
||||
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.
|
||||
|
||||
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.
|
||||
Defined as :math:`\rho/ <\rho> - 1`. Overwrites the input array.
|
||||
|
||||
|
||||
Parameters
|
||||
----------
|
||||
delta : 3-dimensional array of shape `(grid, grid, grid)`
|
||||
The density field.
|
||||
|
||||
Returns
|
||||
-------
|
||||
overdensity : 3-dimensional array of shape `(grid, grid, grid)`.
|
||||
"""
|
||||
delta /= delta.mean()
|
||||
delta -= 1
|
||||
return delta
|
||||
|
||||
def __call__(self, grid, smooth_scale=None, verbose=True):
|
||||
"""
|
||||
Calculate the density field using a Pylians routine [1, 2].
|
||||
|
||||
Parameters
|
||||
----------
|
||||
grid : int
|
||||
Grid size.
|
||||
smooth_scale : float, optional
|
||||
Gaussian kernal scale to smoothen the density field, in box units.
|
||||
verbose : bool
|
||||
Verbosity flag.
|
||||
|
||||
Returns
|
||||
-------
|
||||
rho : 3-dimensional array of shape `(grid, grid, grid)`.
|
||||
Density field.
|
||||
|
||||
References
|
||||
----------
|
||||
[1] https://pylians3.readthedocs.io/
|
||||
[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)
|
||||
return rho
|
||||
|
||||
|
||||
###############################################################################
|
||||
# Potential field calculation #
|
||||
###############################################################################
|
||||
|
||||
|
||||
class PotentialField(BaseField):
|
||||
"""
|
||||
Potential field calculation.
|
||||
|
||||
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).
|
||||
"""
|
||||
def __init__(self, box, MAS):
|
||||
self.box = box
|
||||
self.MAS = MAS
|
||||
|
||||
def __call__(self, overdensity_field):
|
||||
"""
|
||||
Calculate the potential field.
|
||||
|
||||
Parameters
|
||||
----------
|
||||
overdensity_field : 3-dimensional array of shape `(grid, grid, grid)`
|
||||
The overdensity field.
|
||||
|
||||
Returns
|
||||
-------
|
||||
potential : 3-dimensional array of shape `(grid, grid, grid)`.
|
||||
"""
|
||||
return MASL.potential(overdensity_field, self.box._omega_m,
|
||||
self.box._aexp, self.MAS)
|
||||
|
||||
|
||||
###############################################################################
|
||||
# Tidal tensor field calculation #
|
||||
###############################################################################
|
||||
|
||||
|
||||
class TidalTensorField(BaseField):
|
||||
"""
|
||||
Tidal tensor field calculation.
|
||||
|
||||
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).
|
||||
"""
|
||||
def __init__(self, box, MAS):
|
||||
self.box = box
|
||||
self.MAS = MAS
|
||||
|
||||
@staticmethod
|
||||
def tensor_field_eigvals(tidal_tensor):
|
||||
"""
|
||||
Calculate eigenvalues of the tidal tensor field, sorted in increasing
|
||||
order.
|
||||
|
||||
Parameters
|
||||
----------
|
||||
tidal_tensor : :py:class:`MAS_library.tidal_tensor`
|
||||
Tidal tensor object, whose attributes `tidal_tensor.Tij` contain
|
||||
the relevant tensor components.
|
||||
|
||||
Returns
|
||||
-------
|
||||
eigvals : 3-dimensional array of shape `(grid, grid, grid)`
|
||||
"""
|
||||
n_samples = tidal_tensor.T00.size
|
||||
# We create a array and then calculate the eigenvalues.
|
||||
Teval = numpy.full((n_samples, 3, 3), numpy.nan, dtype=numpy.float32)
|
||||
Teval[:, 0, 0] = tidal_tensor.T00
|
||||
Teval[:, 0, 1] = tidal_tensor.T01
|
||||
Teval[:, 0, 2] = tidal_tensor.T02
|
||||
Teval[:, 1, 1] = tidal_tensor.T11
|
||||
Teval[:, 1, 2] = tidal_tensor.T12
|
||||
Teval[:, 2, 2] = tidal_tensor.T22
|
||||
|
||||
eigvals = numpy.full((n_samples, 3), numpy.nan, dtype=numpy.float32)
|
||||
for i in range(n_samples):
|
||||
eigvals[i, :] = numpy.linalg.eigvalsh(Teval[i, ...], 'U')
|
||||
eigvals[i, :] = numpy.sort(eigvals[i, :])
|
||||
|
||||
return eigvals
|
||||
|
||||
def __call__(self, overdensity_field):
|
||||
"""
|
||||
Calculate the tidal tensor field.
|
||||
|
||||
Parameters
|
||||
----------
|
||||
overdensity_field : 3-dimensional array of shape `(grid, grid, grid)`
|
||||
The overdensity field.
|
||||
|
||||
Returns
|
||||
-------
|
||||
tidal_tensor : :py:class:`MAS_library.tidal_tensor`
|
||||
Tidal tensor object, whose attributes `tidal_tensor.Tij` contain
|
||||
the relevant tensor components.
|
||||
"""
|
||||
return MASL.tidal_tensor(overdensity_field, self.box._omega_m,
|
||||
self.box._aexp, self.MAS)
|
||||
|
||||
42
csiborgtools/field/utils.py
Normal file
42
csiborgtools/field/utils.py
Normal file
@ -0,0 +1,42 @@
|
||||
# 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.
|
||||
"""
|
||||
Utility functions for the field module.
|
||||
"""
|
||||
from warnings import warn
|
||||
|
||||
import numpy
|
||||
|
||||
|
||||
def force_single_precision(x, name):
|
||||
"""
|
||||
Convert `x` to float32 if it is not already.
|
||||
|
||||
Parameters
|
||||
----------
|
||||
x : array
|
||||
Array to convert.
|
||||
name : str
|
||||
Name of the array.
|
||||
|
||||
Returns
|
||||
-------
|
||||
x : array
|
||||
Converted array.
|
||||
"""
|
||||
if x.dtype != numpy.float32:
|
||||
warn(f"Converting `{name}` to float32.", UserWarning, stacklevel=1)
|
||||
x = x.astype(numpy.float32)
|
||||
return x
|
||||
@ -193,6 +193,23 @@ class BoxUnits:
|
||||
"""
|
||||
return length / (self._unit_l / units.kpc.to(units.cm) / self._aexp)
|
||||
|
||||
def mpc2box(self, length):
|
||||
r"""
|
||||
Convert length from :math:`\mathrm{cMpc}` (with :math:`h=0.705`) to
|
||||
box units.
|
||||
|
||||
Parameters
|
||||
----------
|
||||
length : float
|
||||
Length in :math:`\mathrm{cMpc}`
|
||||
|
||||
Returns
|
||||
-------
|
||||
length : foat
|
||||
Length in box units.
|
||||
"""
|
||||
return self.kpc2box(length * 1e3)
|
||||
|
||||
def box2mpc(self, length):
|
||||
r"""
|
||||
Convert length from box units to :math:`\mathrm{cMpc}` (with
|
||||
|
||||
@ -326,7 +326,7 @@ class CSiBORGPaths:
|
||||
fname = f"radpos_{str(nsim).zfill(5)}_{str(nsnap).zfill(5)}.npz"
|
||||
return join(fdir, fname)
|
||||
|
||||
def particle_h5py_path(self, nsim):
|
||||
def particle_h5py_path(self, nsim, with_vel):
|
||||
"""
|
||||
Path to the files containing all particles in a `.hdf5` file. Used for
|
||||
the SPH calculation.
|
||||
@ -335,6 +335,8 @@ class CSiBORGPaths:
|
||||
----------
|
||||
nsim : int
|
||||
IC realisation index.
|
||||
with_vel : bool
|
||||
Whether velocities are included.
|
||||
|
||||
Returns
|
||||
-------
|
||||
@ -344,7 +346,10 @@ class CSiBORGPaths:
|
||||
if not isdir(fdir):
|
||||
makedirs(fdir)
|
||||
warn(f"Created directory `{fdir}`.", UserWarning, stacklevel=1)
|
||||
fname = f"particles_{str(nsim).zfill(5)}.h5"
|
||||
if with_vel:
|
||||
fname = f"parts_{str(nsim).zfill(5)}.h5"
|
||||
else:
|
||||
fname = f"parts_pos_{str(nsim).zfill(5)}.h5"
|
||||
return join(fdir, fname)
|
||||
|
||||
def density_field_path(self, mas, nsim):
|
||||
|
||||
@ -191,7 +191,8 @@ class ParticleReader:
|
||||
"""
|
||||
return numpy.hstack([[0], numpy.cumsum(nparts[:-1])])
|
||||
|
||||
def read_particle(self, nsnap, nsim, pars_extract, verbose=True):
|
||||
def read_particle(self, nsnap, nsim, pars_extract, return_structured=True,
|
||||
verbose=True):
|
||||
"""
|
||||
Read particle files of a simulation at a given snapshot and return
|
||||
values of `pars_extract`.
|
||||
@ -204,17 +205,22 @@ class ParticleReader:
|
||||
IC realisation index.
|
||||
pars_extract : list of str
|
||||
Parameters to be extacted.
|
||||
return_structured : bool, optional
|
||||
Whether to return a structured array or a 2-dimensional array. If
|
||||
the latter, then the order of the columns is the same as the order
|
||||
of `pars_extract`. However, enforces single-precision floating
|
||||
point format for all columns.
|
||||
verbose : bool, optional
|
||||
Verbosity flag while for reading the CPU outputs.
|
||||
|
||||
Returns
|
||||
-------
|
||||
out : structured array
|
||||
out : array
|
||||
"""
|
||||
# Open the particle files
|
||||
nparts, partfiles = self.open_particle(nsnap, nsim, verbose=verbose)
|
||||
if verbose:
|
||||
print("Opened {} particle files.".format(nparts.size))
|
||||
print(f"Opened {nparts.size} particle files.")
|
||||
ncpu = nparts.size
|
||||
# Order in which the particles are written in the FortranFile
|
||||
forder = [("x", numpy.float32), ("y", numpy.float32),
|
||||
@ -229,27 +235,34 @@ class ParticleReader:
|
||||
pars_extract = [pars_extract]
|
||||
for p in pars_extract:
|
||||
if p not in fnames:
|
||||
raise ValueError(
|
||||
"Undefined parameter `{}`. Must be one of `{}`."
|
||||
.format(p, fnames))
|
||||
raise ValueError(f"Undefined parameter `{p}`.")
|
||||
|
||||
npart_tot = numpy.sum(nparts)
|
||||
# A dummy array is necessary for reading the fortran files.
|
||||
dum = numpy.full(npart_tot, numpy.nan, dtype=numpy.float16)
|
||||
# These are the data we read along with types
|
||||
dtype = {"names": pars_extract,
|
||||
"formats": [forder[fnames.index(p)][1] for p in pars_extract]}
|
||||
# Allocate the output structured array
|
||||
out = numpy.full(npart_tot, numpy.nan, dtype)
|
||||
# We allocate the output structured/2D array
|
||||
if return_structured:
|
||||
# These are the data we read along with types
|
||||
formats = [forder[fnames.index(p)][1] for p in pars_extract]
|
||||
dtype = {"names": pars_extract, "formats": formats}
|
||||
out = numpy.full(npart_tot, numpy.nan, dtype)
|
||||
else:
|
||||
par2arrpos = {par: i for i, par in enumerate(pars_extract)}
|
||||
out = numpy.full((npart_tot, len(pars_extract)), numpy.nan,
|
||||
dtype=numpy.float32)
|
||||
|
||||
start_ind = self.nparts_to_start_ind(nparts)
|
||||
iters = tqdm(range(ncpu)) if verbose else range(ncpu)
|
||||
for cpu in iters:
|
||||
i = start_ind[cpu]
|
||||
j = nparts[cpu]
|
||||
# trunk-ignore(ruff/B905)
|
||||
for (fname, fdtype) in zip(fnames, fdtypes):
|
||||
if fname in pars_extract:
|
||||
out[fname][i:i + j] = self.read_sp(fdtype, partfiles[cpu])
|
||||
single_part = self.read_sp(fdtype, partfiles[cpu])
|
||||
if return_structured:
|
||||
out[fname][i:i + j] = single_part
|
||||
else:
|
||||
out[i:i + j, par2arrpos[fname]] = single_part
|
||||
else:
|
||||
dum[i:i + j] = self.read_sp(fdtype, partfiles[cpu])
|
||||
# Close the fortran files
|
||||
@ -279,9 +292,8 @@ class ParticleReader:
|
||||
"""
|
||||
nsnap = str(nsnap).zfill(5)
|
||||
cpu = str(cpu + 1).zfill(5)
|
||||
fpath = join(self.paths.ic_path(nsim, tonew=False),
|
||||
"output_{}".format(nsnap),
|
||||
"unbinding_{}.out{}".format(nsnap, cpu))
|
||||
fpath = join(self.paths.ic_path(nsim, tonew=False), f"output_{nsnap}",
|
||||
f"unbinding_{nsnap}.out{cpu}")
|
||||
return FortranFile(fpath)
|
||||
|
||||
def read_clumpid(self, nsnap, nsim, verbose=True):
|
||||
@ -313,7 +325,7 @@ class ParticleReader:
|
||||
j = nparts[cpu]
|
||||
ff = self.open_unbinding(nsnap, nsim, cpu)
|
||||
clumpid[i:i + j] = ff.read_ints()
|
||||
# Close
|
||||
|
||||
ff.close()
|
||||
|
||||
return clumpid
|
||||
|
||||
@ -18,9 +18,9 @@ SPH density field calculation.
|
||||
|
||||
from datetime import datetime
|
||||
from gc import collect
|
||||
from distutils.util import strtobool
|
||||
|
||||
import h5py
|
||||
import numpy
|
||||
from mpi4py import MPI
|
||||
|
||||
try:
|
||||
@ -42,16 +42,22 @@ nproc = comm.Get_size()
|
||||
parser = ArgumentParser()
|
||||
parser.add_argument("--ics", type=int, nargs="+", default=None,
|
||||
help="IC realisatiosn. If `-1` processes all simulations.")
|
||||
parser.add_argument("--with_vel", type=lambda x: bool(strtobool(x)),
|
||||
help="Whether to include velocities in the particle file.")
|
||||
args = parser.parse_args()
|
||||
paths = csiborgtools.read.CSiBORGPaths(**csiborgtools.paths_glamdring)
|
||||
partreader = csiborgtools.read.ParticleReader(paths)
|
||||
pars_extract = ['x', 'y', 'z', 'vx', 'vy', 'vz', 'M']
|
||||
if args.with_vel:
|
||||
pars_extract = ['x', 'y', 'z', 'vx', 'vy', 'vz', 'M']
|
||||
else:
|
||||
pars_extract = ['x', 'y', 'z', 'M']
|
||||
if args.ics is None or args.ics == -1:
|
||||
ics = paths.get_ics(tonew=False)
|
||||
else:
|
||||
ics = args.ics
|
||||
|
||||
# We MPI loop over individual simulations.
|
||||
# MPI loop over individual simulations. We read in the particles from RAMSES
|
||||
# files and dump them to a HDF5 file.
|
||||
jobs = csiborgtools.fits.split_jobs(len(ics), nproc)[rank]
|
||||
for i in jobs:
|
||||
nsim = ics[i]
|
||||
@ -59,17 +65,11 @@ for i in jobs:
|
||||
print(f"{datetime.now()}: Rank {rank} completing simulation {nsim}.",
|
||||
flush=True)
|
||||
|
||||
# We read in the particles from RASMSES files, switch from a
|
||||
# structured array to 2-dimensional array and dump it.
|
||||
parts = partreader.read_particle(nsnap, nsim, pars_extract,
|
||||
verbose=nproc == 1)
|
||||
out = numpy.full((parts.size, len(pars_extract)), numpy.nan,
|
||||
dtype=numpy.float32)
|
||||
for j, par in enumerate(pars_extract):
|
||||
out[:, j] = parts[par]
|
||||
out = partreader.read_particle(
|
||||
nsnap, nsim, pars_extract, return_structured=False, verbose=nproc == 1)
|
||||
|
||||
with h5py.File(paths.particle_h5py_path(nsim), "w") as f:
|
||||
dset = f.create_dataset("particles", data=out)
|
||||
|
||||
del parts, out
|
||||
del out
|
||||
collect()
|
||||
@ -16,13 +16,18 @@
|
||||
Script to calculate the particle centre of mass and Lagrangian patch size in
|
||||
the initial snapshot. Optinally dumps the particle files, however this requires
|
||||
a lot of memory.
|
||||
|
||||
TODO:
|
||||
- stop saving the particle IDs. Unnecessary.
|
||||
- Switch to h5py files. This way can save the positions in the particle
|
||||
array only.
|
||||
"""
|
||||
from argparse import ArgumentParser
|
||||
from datetime import datetime
|
||||
from distutils.util import strtobool
|
||||
from gc import collect
|
||||
from os import remove
|
||||
from os.path import join
|
||||
from os.path import isfile, join
|
||||
|
||||
import numpy
|
||||
from mpi4py import MPI
|
||||
@ -129,6 +134,10 @@ for i, nsim in enumerate(paths.get_ics(tonew=True)):
|
||||
out = numpy.full(parent_ids.size, numpy.nan, dtype=dtype)
|
||||
for i, clid in enumerate(parent_ids):
|
||||
fpath = ftemp.format(nsim, clid, "fit")
|
||||
# There is no file if the halo was skipped due to too few
|
||||
# particles.
|
||||
if not isfile(fpath):
|
||||
continue
|
||||
with open(fpath, "rb") as f:
|
||||
inp = numpy.load(f)
|
||||
out["index"][i] = clid
|
||||
@ -151,8 +160,11 @@ for i, nsim in enumerate(paths.get_ics(tonew=True)):
|
||||
out = {}
|
||||
for clid in parent_ids:
|
||||
fpath = ftemp.format(nsim, clid, "particles")
|
||||
if not isfile(fpath):
|
||||
continue
|
||||
with open(fpath, "rb") as f:
|
||||
out.update({str(clid): numpy.load(f)})
|
||||
remove(fpath)
|
||||
|
||||
fout = paths.initmatch_path(nsim, "particles")
|
||||
print(f"{datetime.now()}: dumping particles to .. `{fout}`.",
|
||||
|
||||
Loading…
Reference in New Issue
Block a user