Aquila-Consortium/csiborgtools
0
0
Fork 0
mirror of https://github.com/Richard-Sti/csiborgtools_public.git synced 2025-05-13 06:01:13 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions

Velocity observer (#86)

Browse source 
* Continue if r200c not defined

* Remove smooth scale

* Remove smooth scale

* Edit Max Matching plot

* Add peculiar velocity

* Add Vobs calculation

* Edit docs

* Add Vobs plot

* Improve plotting

* Edit naming convention

* Make a note

* Add new cat options

* Update density field RSP calculation

* Update field 2 rsp

* Move functions and shorten documentation

* Improve transforms and comments

* Update docs

* Update imports

* Edit calculation

* Add docs

* Remove imports

* Add Quijote flags

* Edit documentation

* Shorten documentation

* Edit func calls

* Shorten

* Docs edits

* Edit docs

* Shorten docs

* Short docs edits

* Simplify docs a little bit

* Save plotting

* Update env
This commit is contained in:
Richard Stiskalek 2023-08-30 23:27:20 +01:00 • committed by GitHub
parent 8e3127f4d9
commit ae92fd9b72
No known key found for this signature in database
GPG key ID: 4AEE18F83AFDEB23
18 changed files with 761 additions and 788 deletions
Download patch file Download diff file Expand all files Collapse all files

65
csiborgtools/field/density.py
View file

@ -22,7 +22,6 @@ 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
@ -32,18 +31,6 @@ class BaseField(ABC):
_box = None
_MAS = None
@property
def boxsize(self):
"""
Box size. Particle positions are always assumed to be in box units,
therefore this is 1.
Returns
-------
boxsize : float
"""
return 1.
@property
def box(self):
"""
@ -51,7 +38,7 @@ class BaseField(ABC):
Returns
-------
box : :py:class:`csiborgtools.units.CSiBORGBox`
:py:class:`csiborgtools.units.CSiBORGBox`
"""
return self._box
@ -70,10 +57,10 @@ class BaseField(ABC):
Returns
-------
MAS : str
str
"""
if self._MAS is None:
raise ValueError("`mas` is not set.")
raise ValueError("`MAS` is not set.")
return self._MAS
@MAS.setter
@ -99,6 +86,8 @@ class DensityField(BaseField):
Mass assignment scheme. Options are Options are: 'NGP' (nearest grid
point), 'CIC' (cloud-in-cell), 'TSC' (triangular-shape cloud), 'PCS'
(piecewise cubic spline).
paths : :py:class:`csiborgtools.read.Paths`
The simulation paths.
References
----------
@ -128,13 +117,12 @@ class DensityField(BaseField):
delta -= 1
return delta
def __call__(self, parts, grid, in_rsp, flip_xz=True, nbatch=30,
verbose=True):
def __call__(self, parts, grid, flip_xz=True, nbatch=30, verbose=True):
"""
Calculate the density field using a Pylians routine [1, 2].
Iteratively loads the particles into memory, flips their `x` and `z`
coordinates. Particles are assumed to be in box units, with positions
in [0, 1] and observer in the centre of the box if RSP is applied.
in [0, 1]
Parameters
----------
@ -143,8 +131,6 @@ class DensityField(BaseField):
Columns are: `x`, `y`, `z`, `vx`, `vy`, `vz`, `M`.
grid : int
Grid size.
in_rsp : bool
Whether to calculate the density field in redshift space.
flip_xz : bool, optional
Whether to flip the `x` and `z` coordinates.
nbatch : int, optional
@ -173,21 +159,14 @@ class DensityField(BaseField):
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")
pos = force_single_precision(pos)
vel = force_single_precision(vel)
mass = force_single_precision(mass)
if flip_xz:
pos[:, [0, 2]] = pos[:, [2, 0]]
vel[:, [0, 2]] = vel[:, [2, 0]]
if in_rsp:
raise NotImplementedError("Redshift space needs to be fixed.")
# TODO change how called + units.
pos = real2redshift(pos, vel, [0.5, 0.5, 0.5], self.box,
in_box_units=True, periodic_wrap=True,
make_copy=False)
MASL.MA(pos, rho, self.boxsize, self.MAS, W=mass, verbose=False)
MASL.MA(pos, rho, 1., self.MAS, W=mass, verbose=False)
if end == nparts:
break
start = end
@ -223,7 +202,7 @@ class VelocityField(BaseField):
@staticmethod
@jit(nopython=True)
def radial_velocity(rho_vel):
def radial_velocity(rho_vel, observer_velocity):
"""
Calculate the radial velocity field around the observer in the centre
of the box.
@ -232,6 +211,8 @@ class VelocityField(BaseField):
----------
rho_vel : 4-dimensional array of shape `(3, grid, grid, grid)`.
Velocity field along each axis.
observer_velocity : 3-dimensional array of shape `(3,)`
Observer velocity.
Returns
-------
@ -240,13 +221,19 @@ class VelocityField(BaseField):
"""
grid = rho_vel.shape[1]
radvel = numpy.zeros((grid, grid, grid), dtype=numpy.float32)
vx0, vy0, vz0 = observer_velocity
for i in range(grid):
px = i - 0.5 * (grid - 1)
for j in range(grid):
py = j - 0.5 * (grid - 1)
for k in range(grid):
pz = k - 0.5 * (grid - 1)
vx, vy, vz = rho_vel[:, i, j, k]
vx = rho_vel[0, i, j, k] - vx0
vy = rho_vel[1, i, j, k] - vy0
vz = rho_vel[2, i, j, k] - vz0
radvel[i, j, k] = ((px * vx + py * vy + pz * vz)
/ numpy.sqrt(px**2 + py**2 + pz**2))
return radvel
@ -297,19 +284,19 @@ class VelocityField(BaseField):
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")
pos = force_single_precision(pos)
vel = force_single_precision(vel)
mass = force_single_precision(mass)
if flip_xz:
pos[:, [0, 2]] = pos[:, [2, 0]]
vel[:, [0, 2]] = vel[:, [2, 0]]
vel *= mass.reshape(-1, 1)
for i in range(3):
MASL.MA(pos, rho_vel[i], self.boxsize, self.MAS, W=vel[:, i],
MASL.MA(pos, rho_vel[i], 1., self.MAS, W=vel[:, i],
verbose=False)
MASL.MA(pos, cellcounts, self.boxsize, self.MAS, W=mass,
MASL.MA(pos, cellcounts, 1., self.MAS, W=mass,
verbose=False)
if end == nparts:
break

244
csiborgtools/field/interp.py
View file

@ -20,23 +20,22 @@ import numpy
from numba import jit
from tqdm import trange
from ..read.utils import radec_to_cartesian, real2redshift
from .utils import force_single_precision
from ..utils import periodic_wrap_grid, radec_to_cartesian
def evaluate_cartesian(*fields, pos, interp="CIC"):
"""
Evaluate a scalar field at Cartesian coordinates.
Evaluate a scalar field(s) at Cartesian coordinates `pos`.
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.
Query positions in box units.
interp : str, optional
Interpolation method. Can be either `CIC` or `NGP`.
Interpolation method, `NGP` or `CIC`.
Returns
-------
@ -44,7 +43,7 @@ def evaluate_cartesian(*fields, pos, interp="CIC"):
"""
assert interp in ["CIC", "NGP"]
boxsize = 1.
pos = force_single_precision(pos, "pos")
pos = force_single_precision(pos)
nsamples = pos.shape[0]
interp_fields = [numpy.full(nsamples, numpy.nan, dtype=numpy.float32)
@ -64,60 +63,71 @@ def evaluate_cartesian(*fields, pos, interp="CIC"):
return interp_fields
def evaluate_sky(*fields, pos, box, isdeg=True):
def evaluate_sky(*fields, pos, box):
"""
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.
Evaluate a scalar field(s) at radial distance `Mpc / h`, right ascensions
[0, 360) deg and declinations [-90, 90] deg.
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.
Query spherical coordinates.
box : :py:class:`csiborgtools.read.CSiBORGBox`
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.
pos = force_single_precision(pos)
pos[:, 0] = box.mpc2box(pos[:, 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)
cart_pos = radec_to_cartesian(pos) + 0.5
return evaluate_cartesian(*fields, pos=cart_pos)
def observer_vobs(velocity_field):
"""
Calculate the observer velocity from a velocity field. Assumes an observer
in the centre of the box.
Parameters
----------
velocity_field : 4-dimensional array of shape `(3, grid, grid, grid)`
Returns
-------
1-dimensional array of shape `(3,)`
"""
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 make_sky(field, angpos, dist, box, volume_weight=True, 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.
box the field is evaluated along directions `angpos` (RA [0, 360) deg,
dec [-90, 90] deg). Along each direction, the field is evaluated distances
`dist` (`Mpc / h`) 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.
Directions to evaluate the field.
dist : 1-dimensional array
Uniformly spaced radial distances to evaluate the field.
box : :py:class:`csiborgtools.read.CSiBORGBox`
The simulation box information and transformations.
volume_weight : bool, optional
Whether to weight the field by the volume of the pixel, i.e. a
:math:`r^2` correction.
Whether to weight the field by the volume of the pixel.
verbose : bool, optional
Verbosity flag.
@ -152,22 +162,9 @@ def make_sky(field, angpos, dist, box, volume_weight=True, verbose=True):
@jit(nopython=True)
def divide_nonzero(field0, field1):
"""
Divide two fields where the second one is not zero. If the second field
is zero, the first one is left unchanged. Operates in-place.
Parameters
----------
field0 : 3-dimensional array of shape `(grid, grid, grid)`
Field to be divided.
field1 : 3-dimensional array of shape `(grid, grid, grid)`
Field to divide by.
Returns
-------
field0 : 3-dimensional array of shape `(grid, grid, grid)`
Field divided by the second one.
Perform in-place `field0 /= field1` but only where `field1 != 0`.
"""
assert field0.shape == field1.shape
assert field0.shape == field1.shape, "Field shapes must match."
imax, jmax, kmax = field0.shape
for i in range(imax):
@ -177,134 +174,89 @@ def divide_nonzero(field0, field1):
field0[i, j, k] /= field1[i, j, k]
def observer_vobs(velocity_field):
@jit(nopython=True)
def make_gridpos(grid_size):
"""Make a regular grid of positions and distances from the center."""
grid_pos = numpy.full((grid_size**3, 3), numpy.nan, dtype=numpy.float32)
grid_dist = numpy.full(grid_size**3, numpy.nan, dtype=numpy.float32)
n = 0
for i in range(grid_size):
px = (i - 0.5 * (grid_size - 1)) / grid_size
px2 = px**2
for j in range(grid_size):
py = (j - 0.5 * (grid_size - 1)) / grid_size
py2 = py**2
for k in range(grid_size):
pz = (k - 0.5 * (grid_size - 1)) / grid_size
pz2 = pz**2
grid_pos[n, 0] = px
grid_pos[n, 1] = py
grid_pos[n, 2] = pz
grid_dist[n] = (px2 + py2 + pz2)**0.5
n += 1
return grid_pos, grid_dist
def field2rsp(field, radvel_field, box, MAS, init_value=0.):
"""
Calculate the observer velocity from a velocity field. Assumes the observer
is in the centre of the box.
Forward model a real space scalar field to redshift space.
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
positions reconstructs back the field on a grid by NGP interpolation.
Parameters
----------
fields : (list of) 3-dimensional array of shape `(grid, grid, grid)`
Real space fields to be evolved to redshift space.
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`.
field : 3-dimensional array of shape `(grid, grid, grid)`
Real space field to be evolved to redshift space.
radvel_field : 3-dimensional array of shape `(grid, grid, grid)`
Radial velocity field in `km / s`. Expected to account for the observer
velocity.
box : :py:class:`csiborgtools.read.CSiBORGBox`
The simulation box information and transformations.
nbatch : int, optional
Number of batches to use when moving particles to redshift space.
Particles are assumed to be lazily loaded to memory.
flip_partsxz : bool, optional
Whether to flip the x and z coordinates of the particles. This is
because of a RAMSES bug.
MAS : str
Mass assignment. Must be one of `NGP`, `CIC`, `TSC` or `PCS`.
init_value : float, optional
Initial value of the RSP field on the grid.
verbose : bool, optional
Verbosity flag.
Returns
-------
rsp_fields : (list of) 3-dimensional array of shape `(grid, grid, grid)`
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:
assert all(fields[0].shape == fields[i].shape
for i in range(1, nfields))
grid = field.shape[0]
H0_inv = 1. / 100 / box.box2mpc(1.)
rsp_fields = [numpy.full(field.shape, init_value, dtype=numpy.float32)
for field in fields]
cellcounts = numpy.zeros(rsp_fields[0].shape, dtype=numpy.float32)
# Calculate the regular grid positions and distances from the center.
grid_pos, grid_dist = make_gridpos(grid)
grid_dist = grid_dist.reshape(-1, 1)
nparts = parts.shape[0]
batch_size = nparts // nbatch
start = 0
for __ in trange(nbatch + 1) if verbose else range(nbatch + 1):
# We first load the batch of particles into memory and flip x and z.
end = min(start + batch_size, nparts)
pos = parts[start:end]
pos, vel = pos[:, :3], pos[:, 3:6]
if flip_partsxz:
pos[:, [0, 2]] = pos[:, [2, 0]]
vel[:, [0, 2]] = vel[:, [2, 0]]
# Then move the particles to redshift space.
# TODO here the function is now called differently and pos assumes
# different units.
rsp_pos = real2redshift(pos, vel, [0.5, 0.5, 0.5], box,
in_box_units=True, periodic_wrap=True,
make_copy=True)
# ... and count the number of particles in each grid cell.
MASL.MA(rsp_pos, cellcounts, 1., "NGP")
# Move the grid positions to redshift space.
grid_pos *= (1 + H0_inv * radvel_field.reshape(-1, 1) / grid_dist)
grid_pos += 0.5
grid_pos = periodic_wrap_grid(grid_pos)
# Now finally we evaluate the field at the particle positions in real
# space and then assign the values to the grid in redshift space.
for i in range(nfields):
values = evaluate_cartesian(fields[i], pos=pos)
MASL.MA(rsp_pos, rsp_fields[i], 1., "NGP", W=values)
if end == nparts:
break
start = end
rsp_field = numpy.full(field.shape, init_value, dtype=numpy.float32)
cell_counts = numpy.zeros(rsp_field.shape, dtype=numpy.float32)
# We divide by the number of particles in each cell.
for i in range(len(fields)):
divide_nonzero(rsp_fields[i], cellcounts)
# Interpolate the field to the grid positions.
MASL.MA(grid_pos, rsp_field, 1., MAS, W=field.reshape(-1,))
MASL.MA(grid_pos, cell_counts, 1., MAS)
divide_nonzero(rsp_field, cell_counts)
if len(fields) == 1:
return rsp_fields[0]
return rsp_fields
return rsp_field
@jit(nopython=True)
def fill_outside(field, fill_value, rmax, boxsize):
"""
Fill cells outside of a sphere of radius `rmax` with `fill_value`. Centered
in the middle of the box.
Parameters
----------
field : 3-dimensional array of shape `(grid, grid, grid)`
Field to be filled.
fill_value : float
Value to fill the field with.
rmax : float
Radius outside of which to fill the field..
boxsize : float
Size of the box.
Returns
-------
field : 3-dimensional array of shape `(grid, grid, grid)`
Fill cells outside of a sphere of radius `rmax` around the box centre with
`fill_value`.
"""
imax, jmax, kmax = field.shape
assert imax == jmax == kmax
N = imax
# Squared radial distance from the center of the box in box units.
rmax_box2 = (N * rmax / boxsize)**2

49
csiborgtools/field/utils.py
View file

@ -15,31 +15,16 @@
"""
Utility functions for the field module.
"""
from warnings import warn
import healpy
import numpy
import smoothing_library as SL
def force_single_precision(x, name):
def force_single_precision(x):
"""
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.
Attempt to convert an array `x` to float 32.
"""
if x.dtype != numpy.float32:
warn(f"Converting `{name}` to float32.", UserWarning, stacklevel=1)
x = x.astype(numpy.float32)
return x
@ -47,21 +32,6 @@ def force_single_precision(x, name):
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)
@ -70,19 +40,10 @@ def smoothen_field(field, smooth_scale, boxsize, threads=1):
def nside2radec(nside):
"""
Generate RA and declination for HEALPix pixel centres.
Parameters
----------
nside : int
HEALPix nside parameter.
Returns
-------
radec : array of shape `(npix, 2)`
RA and declination in degrees.
Generate RA [0, 360] deg. and declination [-90, 90] deg. for HEALPix pixel
centres at a given nside.
"""
pixs = numpy.arange(healpy.nside2npix(nside))
theta, phi = healpy.pix2ang(nside, pixs)
theta -= numpy.pi / 2
return numpy.rad2deg(numpy.vstack([phi, theta]).T)
return 180 / numpy.pi * numpy.vstack([phi, theta]).T

43
csiborgtools/match/match.py
View file

@ -151,7 +151,7 @@ class RealisationsMatcher(BaseMatcher):
Returns
-------
dlogmass : float
float
"""
return self._dlogmass
@ -169,7 +169,7 @@ class RealisationsMatcher(BaseMatcher):
Returns
-------
mass_kind : str
str
"""
return self._mass_kind
@ -186,7 +186,7 @@ class RealisationsMatcher(BaseMatcher):
Returns
-------
overlapper : :py:class:`csiborgtools.match.ParticleOverlap`
:py:class:`csiborgtools.match.ParticleOverlap`
"""
return self._overlapper
@ -661,20 +661,7 @@ class ParticleOverlap(BaseMatcher):
def pos2cell(pos, ncells):
"""
Convert position to cell number. If `pos` is in
`numpy.typecodes["AllInteger"]` assumes it to already be the cell
number.
Parameters
----------
pos : 1-dimensional array
Array of positions along an axis in the box.
ncells : int
Number of cells along the axis.
Returns
-------
cells : 1-dimensional array
Convert position to cell number if there are `ncells` cells along the axis.
"""
if pos.dtype.char in numpy.typecodes["AllInteger"]:
return pos
@ -684,17 +671,7 @@ def pos2cell(pos, ncells):
def read_nshift(smooth_kwargs):
"""
Determine the number of cells to pad the density field if smoothing is
applied. It defaults to the ceiling of three times the smoothing scale.
Parameters
----------
smooth_kwargs : dict or None
Arguments to be passed to :py:func:`scipy.ndimage.gaussian_filter`.
If `None`, no smoothing is applied.
Returns
-------
nshift : int
applied. Defaults to the ceiling of three times the smoothing scale.
"""
return 0 if smooth_kwargs is None else ceil(3 * smooth_kwargs["sigma"])
@ -702,8 +679,7 @@ def read_nshift(smooth_kwargs):
@jit(nopython=True)
def fill_delta(delta, xcell, ycell, zcell, xmin, ymin, zmin, weights):
"""
Fill array `delta` by adding `weights` to the specified cells. This is a
JIT implementation.
Fill array `delta` by adding `weights` to the specified cells.
Parameters
----------
@ -730,8 +706,8 @@ def fill_delta(delta, xcell, ycell, zcell, xmin, ymin, zmin, weights):
@jit(nopython=True)
def fill_delta_indxs(delta, xcell, ycell, zcell, xmin, ymin, zmin, weights):
"""
Fill array `delta` by adding `weights` to the specified cells and return
indices where `delta` was assigned a value. This is a JIT implementation.
Fill array `delta` by adding `weights` to the specified cells. Returns
the indices where `delta` was assigned a value.
Parameters
----------
@ -1161,6 +1137,9 @@ def matching_max(cat0, catx, mass_kind, mult, periodic, overlap=None,
out[i]["hid0"] = hid0
out[i]["mass0"] = 10**mass0[i]
if not numpy.isfinite(rad0[i]):
continue
neigh_dists, neigh_inds = knnx.radius_neighbors(pos0[i].reshape(1, -1),
mult * rad0[i])
neigh_dists, neigh_inds = neigh_dists[0], neigh_inds[0]

3
csiborgtools/read/__init__.py
View file

@ -26,5 +26,4 @@ from .pk_summary import PKReader # noqa
from .readsim import (MmainReader, CSiBORGReader, QuijoteReader, halfwidth_mask, # noqa
load_halo_particles) # noqa
from .tpcf_summary import TPCFReader # noqa
from .utils import (cartesian_to_radec, cols_to_structured, radec_to_cartesian, # noqa
read_h5, real2redshift) # noqa
from .utils import (cols_to_structured, read_h5) # noqa

20
csiborgtools/read/halo_cat.py
View file

@ -31,9 +31,9 @@ from sklearn.neighbors import NearestNeighbors
from .box_units import CSiBORGBox, QuijoteBox
from .paths import Paths
from .readsim import CSiBORGReader
from .utils import (add_columns, cartesian_to_radec, cols_to_structured,
flip_cols, radec_to_cartesian, real2redshift)
from ..utils import periodic_distance_two_points
from .utils import add_columns, cols_to_structured, flip_cols
from ..utils import (periodic_distance_two_points, real2redshift,
cartesian_to_radec, radec_to_cartesian)
class BaseCatalogue(ABC):
@ -631,6 +631,9 @@ class QuijoteHaloCatalogue(BaseCatalogue):
Load initial positions from `fit_init.py`.
with_lagpatch : bool, optional
Load halos with a resolved Lagrangian patch.
load_backup : bool, optional
Load halos from the backup catalogue that do not have corresponding
snapshots.
"""
_nsnap = None
_origin = None
@ -638,14 +641,15 @@ class QuijoteHaloCatalogue(BaseCatalogue):
def __init__(self, nsim, paths, nsnap,
observer_location=[500., 500., 500.],
bounds=None, load_fitted=True, load_initial=True,
with_lagpatch=False):
with_lagpatch=False, load_backup=False):
self.nsim = nsim
self.paths = paths
self.nsnap = nsnap
self.observer_location = observer_location
self._box = QuijoteBox(nsnap, nsim, paths)
# NOTE watch out about here setting nsim = 0
self._box = QuijoteBox(nsnap, 0, paths)
fpath = self.paths.fof_cat(nsim, "quijote")
fpath = self.paths.fof_cat(nsim, "quijote", load_backup)
fof = FoF_catalog(fpath, self.nsnap, long_ids=False, swap=False,
SFR=False, read_IDs=False)
@ -667,6 +671,10 @@ class QuijoteHaloCatalogue(BaseCatalogue):
# particles unassigned to any FoF group.
data["index"] = 1 + numpy.arange(data.size, dtype=numpy.int32)
if load_backup and (load_initial or load_fitted):
raise ValueError("Cannot load initial/fitted data for the backup "
"catalogues.")
if load_initial:
data = self.load_initial(data, paths, "quijote")
if load_fitted:

65
csiborgtools/read/paths.py
View file

@ -214,7 +214,7 @@ class Paths:
fout = fout.replace(".npy", "_sorted.npy")
return fout
def fof_cat(self, nsim, simname):
def fof_cat(self, nsim, simname, from_quijote_backup=False):
r"""
Path to the :math:`z = 0` FoF halo catalogue.
@ -224,13 +224,24 @@ class Paths:
IC realisation index.
simname : str
Simulation name. Must be one of `csiborg` or `quijote`.
from_quijote_backup : bool, optional
Whether to return the path to the Quijote FoF catalogue from the
backup.
Returns
-------
path : str
"""
if simname == "csiborg":
fdir = join(self.postdir, "FoF_membership", )
try_create_directory(fdir)
return join(fdir, f"halo_catalog_{nsim}_FOF.txt")
elif simname == "quijote":
return join(self.quijote_dir, "Halos_fiducial", str(nsim))
if from_quijote_backup:
return join(self.quijote_dir, "halos_backup", str(nsim))
else:
return join(self.quijote_dir, "Halos_fiducial", str(nsim))
else:
raise ValueError(f"Unknown simulation name `{simname}`.")
@ -285,7 +296,7 @@ class Paths:
try_create_directory(fdir)
return join(fdir, f"{kind}_{str(nsim).zfill(5)}.{ftype}")
def get_ics(self, simname):
def get_ics(self, simname, from_quijote_backup=False):
"""
Get available IC realisation IDs for either the CSiBORG or Quijote
simulation suite.
@ -294,6 +305,8 @@ class Paths:
----------
simname : str
Simulation name. Must be `csiborg` or `quijote`.
from_quijote_backup : bool, optional
Whether to return the ICs from the Quijote backup.
Returns
-------
@ -311,10 +324,13 @@ class Paths:
except ValueError:
pass
elif simname == "quijote" or simname == "quijote_full":
files = glob(
"/mnt/extraspace/rstiskalek/Quijote/Snapshots_fiducial/*")
if from_quijote_backup:
files = glob(join(self.quijote_dir, "halos_backup", "*"))
else:
warn(("Taking only the snapshots that also have "
"a FoF catalogue!"))
files = glob(join(self.quijote_dir, "Snapshots_fiducial", "*"))
files = [int(f.split("/")[-1]) for f in files]
warn("Taking only the snapshots that also have a FoF catalogue!")
files = [f for f in files if f < 100]
else:
raise ValueError(f"Unknown simulation name `{simname}`.")
@ -545,7 +561,7 @@ class Paths:
return join(fdir, fname)
def field(self, kind, MAS, grid, nsim, in_rsp, smooth_scale=None):
def field(self, kind, MAS, grid, nsim, in_rsp):
r"""
Path to the files containing the calculated density fields in CSiBORG.
@ -562,8 +578,6 @@ class Paths:
IC realisation index.
in_rsp : bool
Whether the calculation is performed in redshift space.
smooth_scale : float, optional
Smoothing scale in :math:`\mathrm{Mpc}/h`
Returns
-------
@ -579,12 +593,32 @@ class Paths:
kind = kind + "_rsp"
fname = f"{kind}_{MAS}_{str(nsim).zfill(5)}_grid{grid}.npy"
if smooth_scale is not None and smooth_scale > 0:
smooth_scale = float(smooth_scale)
fname = fname.replace(".npy", f"smooth{smooth_scale}.npy")
return join(fdir, fname)
def halo_counts(self, simname, nsim):
def observer_peculiar_velocity(self, MAS, grid, nsim):
"""
Path to the files containing the observer peculiar velocity.
Parameters
----------
MAS : str
Mass-assignment scheme.
grid : int
Grid size.
nsim : int
IC realisation index.
Returns
-------
path : str
"""
fdir = join(self.postdir, "environment")
try_create_directory(fdir)
fname = f"obs_vp_{MAS}_{str(nsim).zfill(5)}_{grid}.npz"
return join(fdir, fname)
def halo_counts(self, simname, nsim, from_quijote_backup=False):
"""
Path to the files containing the binned halo counts.
@ -594,6 +628,9 @@ class Paths:
Simulation name. Must be `csiborg`, `quijote` or `quijote_full`.
nsim : int
IC realisation index.
from_quijote_backup : bool, optional
Whether to return the path to the Quijote halo counts from the
backup catalogues.
Returns
-------
@ -602,6 +639,8 @@ class Paths:
fdir = join(self.postdir, "HMF")
try_create_directory(fdir)
fname = f"halo_counts_{simname}_{str(nsim).zfill(5)}.npz"
if from_quijote_backup:
fname = fname.replace("halo_counts", "halo_counts_backup")
return join(fdir, fname)
def cross_nearest(self, simname, run, kind, nsim=None, nobs=None):

164
csiborgtools/read/utils.py
View file

@ -12,126 +12,11 @@
# 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.
"""
Various coordinate transformations.
"""
from os.path import isfile
import numpy
from h5py import File
###############################################################################
# Coordinate transforms #
###############################################################################
def cartesian_to_radec(X, indeg=True):
"""
Calculate the radial distance, RA, dec from Cartesian coordinates. Note,
RA is in range [0, 360) degrees and dec in range [-90, 90] degrees.
Parameters
----------
X : 2-dimensional array `(nsamples, 3)`
Cartesian coordinates.
indeg : bool, optional
Whether to return RA and DEC in degrees.
Returns
-------
out : 2-dimensional array `(nsamples, 3)`
Radial distance, RA and dec.
"""
x, y, z = X[:, 0], X[:, 1], X[:, 2]
dist = numpy.linalg.norm(X, axis=1)
dec = numpy.arcsin(z / dist)
ra = numpy.arctan2(y, x)
ra[ra < 0] += 2 * numpy.pi # Wrap RA to [0, 2pi)
if indeg:
ra = numpy.rad2deg(ra)
dec = numpy.rad2deg(dec)
return numpy.vstack([dist, ra, dec]).T
def radec_to_cartesian(X, isdeg=True):
"""
Calculate Cartesian coordinates from radial distance, RA, dec. Note, RA is
expected in range [0, 360) degrees and dec in range [-90, 90] degrees.
Parameters
----------
X : 2-dimensional array `(nsamples, 3)`
Radial distance, RA and dec.
isdeg : bool, optional
Whether to return RA and DEC in degrees.
Returns
-------
out : 2-dimensional array `(nsamples, 3)`
Cartesian coordinates.
"""
dist, ra, dec = X[:, 0], X[:, 1], X[:, 2]
if isdeg:
ra = numpy.deg2rad(ra)
dec = numpy.deg2rad(dec)
x = dist * numpy.cos(dec) * numpy.cos(ra)
y = dist * numpy.cos(dec) * numpy.sin(ra)
z = dist * numpy.sin(dec)
return numpy.vstack([x, y, z]).T
def real2redshift(pos, vel, observer_location, box, periodic_wrap=True,
subtract_observer=False, make_copy=True):
r"""
Convert real-space position to redshift space position.
Parameters
----------
pos : 2-dimensional array `(nsamples, 3)`
Real-space Cartesian components in :math:`\mathrm{cMpc} / h`.
vel : 2-dimensional array `(nsamples, 3)`
Cartesian velocity in :math:`\mathrm{km} \mathrm{s}^{-1}`.
observer_location: 1-dimensional array `(3,)`
Observer location in :math:`\mathrm{cMpc} / h`.
box : py:class:`csiborg.read.CSiBORGBox`
Box units.
periodic_wrap : bool, optional
Whether to wrap around the box, particles may be outside the default
bounds once RSD is applied.
subtract_observer : bool, optional
If True, subtract the observer's location from the returned
positions.
make_copy : bool, optional
Whether to make a copy of `pos` before modifying it.
Returns
-------
pos : 2-dimensional array `(nsamples, 3)`
Redshift-space Cartesian position in :math:`\mathrm{cMpc} / h`.
"""
if make_copy:
pos = numpy.copy(pos)
# Place the observer at the origin
pos -= observer_location
# Dot product of position vector and velocity
vr_dot = numpy.sum(pos * vel, axis=1)
# Compute the norm squared of the displacement
norm2 = numpy.sum(pos**2, axis=1)
pos *= (1 + box._aexp / box.H0 * vr_dot / norm2).reshape(-1, 1)
# Place the observer back at the original location
if not subtract_observer:
pos += observer_location
if periodic_wrap:
boxsize = box.box2mpc(1.)
# Wrap around the box.
pos = numpy.where(pos > boxsize, pos - boxsize, pos)
pos = numpy.where(pos < 0, pos + boxsize, pos)
return pos
###############################################################################
# Array manipulation #
@ -140,19 +25,7 @@ def real2redshift(pos, vel, observer_location, box, periodic_wrap=True,
def cols_to_structured(N, cols):
"""
Allocate a structured array from `cols`.
Parameters
----------
N : int
Structured array size.
cols: list of tuples
Column names and dtypes. Each tuple must be written as `(name, dtype)`.
Returns
-------
out : structured array
Initialized structured array.
Allocate a structured array from `cols`, a list of (name, dtype) tuples.
"""
if not (isinstance(cols, list)
and all(isinstance(c, tuple) and len(c) == 2 for c in cols)):
@ -167,19 +40,6 @@ def cols_to_structured(N, cols):
def add_columns(arr, X, cols):
"""
Add new columns `X` to a record array `arr`. Creates a new array.
Parameters
----------
arr : structured array
Structured array to add columns to.
X : (list of) 1-dimensional array(s)
Columns to be added.
cols : str or list of str
Column names to be added.
Returns
-------
out : structured array
"""
cols = [cols] if isinstance(cols, str) else cols
@ -214,17 +74,6 @@ def add_columns(arr, X, cols):
def rm_columns(arr, cols):
"""
Remove columns `cols` from a structured array `arr`. Allocates a new array.
Parameters
----------
arr : structured array
Structured array to remove columns from.
cols : str or list of str
Column names to be removed.
Returns
-------
out : structured array
"""
# Ensure cols is a list
cols = [cols] if isinstance(cols, str) else cols
@ -247,16 +96,7 @@ def rm_columns(arr, cols):
def flip_cols(arr, col1, col2):
"""
Flip values in columns `col1` and `col2`. `arr` is modified in place.
Parameters
----------
arr : structured array
Array whose columns are to be flipped.
col1 : str
First column name.
col2 : str
Second column name.
Flip values in columns `col1` and `col2` of a structured array `arr`.
"""
if col1 not in arr.dtype.names or col2 not in arr.dtype.names:
raise ValueError(f"Both `{col1}` and `{col2}` must exist in `arr`.")

206
csiborgtools/utils.py
View file

@ -16,34 +16,25 @@
import numpy
from numba import jit
###############################################################################
# Positions #
###############################################################################
@jit(nopython=True, fastmath=True, boundscheck=False)
def center_of_mass(points, mass, boxsize):
def center_of_mass(particle_positions, particles_mass, boxsize):
"""
Calculate the center of mass of a halo while assuming periodic boundary
conditions of a cubical box. Assuming that particle positions are in
`[0, boxsize)` range. This is a JIT implementation.
Parameters
----------
points : 2-dimensional array of shape (n_particles, 3)
Particle position array.
mass : 1-dimensional array of shape `(n_particles, )`
Particle mass array.
boxsize : float
Box size in the same units as `parts` coordinates.
Returns
-------
cm : 1-dimensional array of shape `(3, )`
conditions of a cubical box.
"""
cm = numpy.zeros(3, dtype=points.dtype)
totmass = sum(mass)
cm = numpy.zeros(3, dtype=particle_positions.dtype)
totmass = sum(particles_mass)
# Convert positions to unit circle coordinates in the complex plane,
# calculate the weighted average and convert it back to box coordinates.
for i in range(3):
cm_i = sum(mass * numpy.exp(2j * numpy.pi * points[:, i] / boxsize))
cm_i = sum(particles_mass * numpy.exp(
2j * numpy.pi * particle_positions[:, i] / boxsize))
cm_i /= totmass
cm_i = numpy.arctan2(cm_i.imag, cm_i.real) * boxsize / (2 * numpy.pi)
@ -55,24 +46,11 @@ def center_of_mass(points, mass, boxsize):
return cm
@jit(nopython=True)
def periodic_distance(points, reference, boxsize):
@jit(nopython=True, fastmath=True, boundscheck=False)
def periodic_distance(points, reference_point, boxsize):
"""
Compute the 3D distance between multiple points and a reference point using
periodic boundary conditions. This is an optimized JIT implementation.
Parameters
----------
points : 2-dimensional array of shape `(n_points, 3)`
Points to calculate the distance from the reference point.
reference : 1-dimensional array of shape `(3, )`
Reference point.
boxsize : float
Box size.
Returns
-------
dist : 1-dimensional array of shape `(n_points, )`
periodic boundary conditions.
"""
npoints = len(points)
half_box = boxsize / 2
@ -80,7 +58,7 @@ def periodic_distance(points, reference, boxsize):
dist = numpy.zeros(npoints, dtype=points.dtype)
for i in range(npoints):
for j in range(3):
dist_1d = abs(points[i, j] - reference[j])
dist_1d = abs(points[i, j] - reference_point[j])
if dist_1d > (half_box):
dist_1d = boxsize - dist_1d
@ -94,23 +72,7 @@ def periodic_distance(points, reference, boxsize):
@jit(nopython=True, fastmath=True, boundscheck=False)
def periodic_distance_two_points(p1, p2, boxsize):
"""
Compute the 3D distance between two points using periodic boundary
conditions. This is an optimized JIT implementation.
Parameters
----------
p1 : 1-dimensional array of shape `(3, )`
First point.
p2 : 1-dimensional array of shape `(3, )`
Second point.
boxsize : float
Box size.
Returns
-------
dist : 1-dimensional array of shape `(n_points, )`
"""
"""Compute the 3D distance between two points in a periodic box."""
half_box = boxsize / 2
dist = 0
@ -125,47 +87,133 @@ def periodic_distance_two_points(p1, p2, boxsize):
return dist**0.5
@jit(nopython=True)
def periodic_wrap_grid(pos, boxsize=1):
"""Wrap positions in a periodic box."""
for n in range(pos.shape[0]):
for i in range(3):
if pos[n, i] > boxsize:
pos[n, i] -= boxsize
elif pos[n, i] < 0:
pos[n, i] += boxsize
return pos
@jit(nopython=True, fastmath=True, boundscheck=False)
def delta2ncells(delta):
def delta2ncells(field):
"""
Calculate the number of cells in `delta` that are non-zero.
Parameters
----------
delta : 3-dimensional array
Halo density field.
Returns
-------
ncells : int
Number of non-zero cells.
Calculate the number of cells in `field` that are non-zero.
"""
tot = 0
imax, jmax, kmax = delta.shape
imax, jmax, kmax = field.shape
for i in range(imax):
for j in range(jmax):
for k in range(kmax):
if delta[i, j, k] > 0:
if field[i, j, k] > 0:
tot += 1
return tot
def cartesian_to_radec(X):
"""
Calculate the radial distance, RA [0, 360) deg and dec [-90, 90] deg.
"""
x, y, z = X[:, 0], X[:, 1], X[:, 2]
dist = numpy.linalg.norm(X, axis=1)
dec = numpy.arcsin(z / dist)
ra = numpy.arctan2(y, x)
ra[ra < 0] += 2 * numpy.pi
ra *= 180 / numpy.pi
dec *= 180 / numpy.pi
return numpy.vstack([dist, ra, dec]).T
def radec_to_cartesian(X):
"""
Calculate Cartesian coordinates from radial distance, RA [0, 360) deg and
dec [-90, 90] deg.
"""
dist, ra, dec = X[:, 0], X[:, 1], X[:, 2]
ra *= numpy.pi / 180
dec *= numpy.pi / 180
cdec = numpy.cos(dec)
return numpy.vstack([
dist * cdec * numpy.cos(ra),
dist * cdec * numpy.sin(ra),
dist * numpy.sin(dec)
]).T
def real2redshift(pos, vel, observer_location, observer_velocity, box,
periodic_wrap=True, make_copy=True):
r"""
Convert real-space position to redshift space position.
Parameters
----------
pos : 2-dimensional array `(nsamples, 3)`
Real-space Cartesian components in `Mpc / h`.
vel : 2-dimensional array `(nsamples, 3)`
Cartesian velocity in `km / s`.
observer_location: 1-dimensional array `(3,)`
Observer location in `Mpc / h`.
observer_velocity: 1-dimensional array `(3,)`
Observer velocity in `km / s`.
box : py:class:`csiborg.read.CSiBORGBox`
Box units.
periodic_wrap : bool, optional
Whether to wrap around the box, particles may be outside the default
bounds once RSD is applied.
make_copy : bool, optional
Whether to make a copy of `pos` before modifying it.
Returns
-------
pos : 2-dimensional array `(nsamples, 3)`
Redshift-space Cartesian position in `Mpc / h`.
"""
if make_copy:
pos = numpy.copy(pos)
vel = numpy.copy(vel)
H0_inv = 1. / 100
# Place the observer at the origin
pos -= observer_location
vel -= observer_velocity
vr_dot = numpy.einsum('ij,ij->i', pos, vel)
norm2 = numpy.einsum('ij,ij->i', pos, pos)
pos *= (1 + H0_inv * vr_dot / norm2).reshape(-1, 1)
# Place the observer back
pos += observer_location
if not make_copy:
vel += observer_velocity
if periodic_wrap:
boxsize = box.box2mpc(1.)
pos = periodic_wrap_grid(pos, boxsize)
return pos
###############################################################################
# Statistics #
###############################################################################
@jit(nopython=True, fastmath=True, boundscheck=False)
def number_counts(x, bin_edges):
"""
Calculate counts of samples in bins.
Parameters
----------
x : 1-dimensional array
Samples' values.
bin_edges : 1-dimensional array
Bin edges.
Returns
-------
counts : 1-dimensional array
Bin counts.
"""
out = numpy.full(bin_edges.size - 1, numpy.nan, dtype=numpy.float32)
for i in range(bin_edges.size - 1):