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

Remove overly verbose code (#129)

Browse source 
* Remove not necessary comments

* Simplify

* Remove unnecessary comments

* Remove verbose comments

* Remove unnecessary verbosity

* Update params

* Remove verbose

* Simplify verbosity

* Simploify comments

* Remove more silly verbosity
This commit is contained in:
Richard Stiskalek 2024-06-21 13:35:45 +01:00 committed by GitHub
parent 779f2e76ac
commit d3b4bfd29c
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
11 changed files with 212 additions and 1306 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

11
csiborgtools/clustering/utils.py
View file

@ -138,17 +138,6 @@ def wrapRA(ra, indeg):
""" """
Wrap RA from :math:`[-180, 180)` to :math`[0, 360)` degrees if `indeg` or Wrap RA from :math:`[-180, 180)` to :math`[0, 360)` degrees if `indeg` or
equivalently in radians otherwise. equivalently in radians otherwise.
Paramaters
----------
ra : 1-dimensional array
Right ascension.
indeg : bool
Whether the right ascension is in degrees.
Returns
-------
wrapped_ra : 1-dimensional array
""" """
mask = ra < 0 mask = ra < 0
if numpy.sum(mask) == 0: if numpy.sum(mask) == 0:

19
csiborgtools/field/utils.py
View file

@ -22,9 +22,7 @@ import healpy
def force_single_precision(x): def force_single_precision(x):
""" """Attempt to convert an array `x` to float32."""
Attempt to convert an array `x` to float 32.
"""
if x.dtype != numpy.float32: if x.dtype != numpy.float32:
x = x.astype(numpy.float32) x = x.astype(numpy.float32)
return x return x
@ -32,9 +30,7 @@ def force_single_precision(x):
@jit(nopython=True) @jit(nopython=True)
def divide_nonzero(field0, field1): def divide_nonzero(field0, field1):
""" """Perform in-place `field0 /= field1` but only where `field1 != 0`."""
Perform in-place `field0 /= field1` but only where `field1 != 0`.
"""
assert field0.shape == field1.shape, "Field shapes must match." assert field0.shape == field1.shape, "Field shapes must match."
imax, jmax, kmax = field0.shape imax, jmax, kmax = field0.shape
@ -47,17 +43,8 @@ def divide_nonzero(field0, field1):
def nside2radec(nside): def nside2radec(nside):
""" """
Generate RA [0, 360] deg. and declination [-90, 90] deg for HEALPix pixel Generate RA [0, 360] deg and declination [-90, 90] deg for HEALPix pixel
centres at a given nside. centres at a given nside.
Parameters
----------
nside : int
HEALPix nside.
Returns
-------
angpos : 2-dimensional array of shape (npix, 2)
""" """
pixs = numpy.arange(healpy.nside2npix(nside)) pixs = numpy.arange(healpy.nside2npix(nside))
theta, phi = healpy.pix2ang(nside, pixs) theta, phi = healpy.pix2ang(nside, pixs)

420
csiborgtools/flow/flow_model.py
View file

@ -13,14 +13,15 @@
# with this program; if not, write to the Free Software Foundation, Inc., # with this program; if not, write to the Free Software Foundation, Inc.,
# 51 Franklin Street, Fifth Floor, Boston, MA 02110-1301, USA. # 51 Franklin Street, Fifth Floor, Boston, MA 02110-1301, USA.
""" """
Validation of the CSiBORG velocity field against PV measurements. Based on [1]. Validation of the CSiBORG velocity field against PV measurements. A lot here
is based on [1], though with many modifications. Throughout, comoving distances
are in `Mpc / h` and velocities in `km / s`.
References References
---------- ----------
[1] https://arxiv.org/abs/1912.09383. [1] https://arxiv.org/abs/1912.09383.
""" """
from abc import ABC, abstractmethod from abc import ABC, abstractmethod
from datetime import datetime
from warnings import catch_warnings, simplefilter, warn from warnings import catch_warnings, simplefilter, warn
import numpy as np import numpy as np
@ -40,17 +41,12 @@ from scipy.optimize import fmin_powell
from sklearn.model_selection import KFold from sklearn.model_selection import KFold
from tqdm import trange from tqdm import trange
from ..params import simname2Omega_m from ..params import SPEED_OF_LIGHT, simname2Omega_m
from ..utils import radec_to_galactic from ..utils import fprint, radec_to_galactic
SPEED_OF_LIGHT = 299792.458 # km / s
H0 = 100 # km / s / Mpc H0 = 100 # km / s / Mpc
def t():
return datetime.now().strftime("%H:%M:%S")
############################################################################### ###############################################################################
# Data loader # # Data loader #
############################################################################### ###############################################################################
@ -83,13 +79,11 @@ class DataLoader:
""" """
def __init__(self, simname, ksim, catalogue, catalogue_fpath, paths, def __init__(self, simname, ksim, catalogue, catalogue_fpath, paths,
ksmooth=None, store_full_velocity=False, verbose=True): ksmooth=None, store_full_velocity=False, verbose=True):
if verbose: fprint("reading the catalogue,", verbose)
print(f"{t()}: reading the catalogue.", flush=True)
self._cat = self._read_catalogue(catalogue, catalogue_fpath) self._cat = self._read_catalogue(catalogue, catalogue_fpath)
self._catname = catalogue self._catname = catalogue
if verbose: fprint("reading the interpolated field,", verbose)
print(f"{t()}: reading the interpolated field.", flush=True)
self._field_rdist, self._los_density, self._los_velocity = self._read_field( # noqa self._field_rdist, self._los_density, self._los_velocity = self._read_field( # noqa
simname, ksim, catalogue, ksmooth, paths) simname, ksim, catalogue, ksmooth, paths)
@ -106,8 +100,7 @@ class DataLoader:
raise ValueError("The number of objects in the catalogue does not " raise ValueError("The number of objects in the catalogue does not "
"match the number of objects in the field.") "match the number of objects in the field.")
if verbose: fprint("calculating the radial velocity.", verbose)
print(f"{t()}: calculating the radial velocity.", flush=True)
nobject = len(self._los_density) nobject = len(self._los_density)
dtype = self._los_density.dtype dtype = self._los_density.dtype
@ -144,74 +137,37 @@ class DataLoader:
@property @property
def cat(self): def cat(self):
""" """The distance indicators catalogue (structured array)."""
The distance indicators catalogue.
Returns
-------
structured array
"""
return self._cat[self._mask] return self._cat[self._mask]
@property @property
def catname(self): def catname(self):
""" """Catalogue name."""
Name of the catalogue.
Returns
-------
str
"""
return self._catname return self._catname
@property @property
def rdist(self): def rdist(self):
""" """Radial distances at which the field was interpolated."""
Radial distances where the field was interpolated for each object.
Returns
-------
1-dimensional array
"""
return self._field_rdist return self._field_rdist
@property @property
def los_density(self): def los_density(self):
""" """Density field along the line of sight `(n_objects, n_steps)`."""
Density field along the line of sight.
Returns
----------
2-dimensional array of shape (n_objects, n_steps)
"""
return self._los_density[self._mask] return self._los_density[self._mask]
@property @property
def los_velocity(self): def los_velocity(self):
""" """Velocity field along the line of sight `(3, n_objects, n_steps)`."""
Velocity field along the line of sight.
Returns
-------
3-dimensional array of shape (3, n_objects, n_steps)
"""
if self._los_velocity is None: if self._los_velocity is None:
raise ValueError("The 3D velocities were not stored.") raise ValueError("The 3D velocities were not stored.")
return self._los_velocity[self._mask] return self._los_velocity[self._mask]
@property @property
def los_radial_velocity(self): def los_radial_velocity(self):
""" """Radial velocity along the line of sight `(n_objects, n_steps)`."""
Radial velocity along the line of sight.
Returns
-------
2-dimensional array of shape (n_objects, n_steps)
"""
return self._los_radial_velocity[self._mask] return self._los_radial_velocity[self._mask]
def _read_field(self, simname, ksim, catalogue, ksmooth, paths): def _read_field(self, simname, ksim, catalogue, ksmooth, paths):
"""Read in the interpolated field."""
nsims = paths.get_ics(simname) nsims = paths.get_ics(simname)
if not (0 <= ksim < len(nsims)): if not (0 <= ksim < len(nsims)):
raise ValueError("Invalid simulation index.") raise ValueError("Invalid simulation index.")
@ -236,7 +192,6 @@ class DataLoader:
return rdist, los_density, los_velocity return rdist, los_density, los_velocity
def _read_catalogue(self, catalogue, catalogue_fpath): def _read_catalogue(self, catalogue, catalogue_fpath):
"""Read in the distance indicator catalogue."""
if catalogue == "A2": if catalogue == "A2":
with File(catalogue_fpath, 'r') as f: with File(catalogue_fpath, 'r') as f:
dtype = [(key, np.float32) for key in f.keys()] dtype = [(key, np.float32) for key in f.keys()]
@ -279,20 +234,8 @@ class DataLoader:
def make_jackknife_mask(self, i, n_splits, seed=42): def make_jackknife_mask(self, i, n_splits, seed=42):
""" """
Set the jackknife mask to exclude the `i`-th split. Set the internal jackknife mask to exclude the `i`-th split out of
`n_splits`.
Parameters
----------
i : int
Index of the split to exclude.
n_splits : int
Number of splits.
seed : int, optional
Random seed.
Returns
-------
None, sets `mask` internally.
""" """
cv = KFold(n_splits=n_splits, shuffle=True, random_state=seed) cv = KFold(n_splits=n_splits, shuffle=True, random_state=seed)
n = len(self._cat) n = len(self._cat)
@ -320,20 +263,8 @@ class DataLoader:
def radial_velocity_los(los_velocity, ra, dec): def radial_velocity_los(los_velocity, ra, dec):
""" """
Calculate the radial velocity along the line of sight. Calculate the radial velocity along the LOS from the 3D velocity
along the LOS `(3, n_steps)`.
Parameters
----------
los_velocity : 2-dimensional array of shape (3, n_steps)
Line of sight velocity field.
ra, dec : floats
Right ascension and declination of the line of sight.
is_degrees : bool, optional
Whether the angles are in degrees.
Returns
-------
1-dimensional array of shape (n_steps)
""" """
types = (float, np.float32, np.float64) types = (float, np.float32, np.float64)
if not isinstance(ra, types) and not isinstance(dec, types): if not isinstance(ra, types) and not isinstance(dec, types):
@ -359,15 +290,6 @@ def lognorm_mean_std_to_loc_scale(mu, std):
""" """
Calculate the location and scale parameters for the log-normal distribution Calculate the location and scale parameters for the log-normal distribution
from the mean and standard deviation. from the mean and standard deviation.
Parameters
----------
mu, std : float
Mean and standard deviation.
Returns
-------
loc, scale : float
""" """
loc = np.log(mu) - 0.5 * np.log(1 + (std / mu) ** 2) loc = np.log(mu) - 0.5 * np.log(1 + (std / mu) ** 2)
scale = np.sqrt(np.log(1 + (std / mu) ** 2)) scale = np.sqrt(np.log(1 + (std / mu) ** 2))
@ -378,17 +300,6 @@ def simps(y, dx):
""" """
Simpson's rule 1D integration, assuming that the number of steps is even Simpson's rule 1D integration, assuming that the number of steps is even
and that the step size is constant. and that the step size is constant.
Parameters
----------
y : 1-dimensional array
Function values.
dx : float
Step size.
Returns
-------
float
""" """
if len(y) % 2 == 0: if len(y) % 2 == 0:
raise ValueError("The number of steps must be odd.") raise ValueError("The number of steps must be odd.")
@ -400,17 +311,6 @@ def dist2redshift(dist, Omega_m):
""" """
Convert comoving distance to cosmological redshift if the Universe is Convert comoving distance to cosmological redshift if the Universe is
flat and z << 1. flat and z << 1.
Parameters
----------
dist : float or 1-dimensional array
Comoving distance in `Mpc / h`.
Omega_m : float
Matter density parameter.
Returns
-------
float or 1-dimensional array
""" """
eta = 3 * Omega_m / 2 eta = 3 * Omega_m / 2
return 1 / eta * (1 - (1 - 2 * H0 * dist / SPEED_OF_LIGHT * eta)**0.5) return 1 / eta * (1 - (1 - 2 * H0 * dist / SPEED_OF_LIGHT * eta)**0.5)
@ -420,17 +320,6 @@ def redshift2dist(z, Omega_m):
""" """
Convert cosmological redshift to comoving distance if the Universe is Convert cosmological redshift to comoving distance if the Universe is
flat and z << 1. flat and z << 1.
Parameters
----------
z : float or 1-dimensional array
Cosmological redshift.
Omega_m : float
Matter density parameter.
Returns
-------
float or 1-dimensional array
""" """
q0 = 3 * Omega_m / 2 - 1 q0 = 3 * Omega_m / 2 - 1
return SPEED_OF_LIGHT * z / (2 * H0) * (2 - z * (1 + q0)) return SPEED_OF_LIGHT * z / (2 * H0) * (2 - z * (1 + q0))
@ -440,37 +329,13 @@ def gradient_redshift2dist(z, Omega_m):
""" """
Gradient of the redshift to comoving distance conversion if the Universe is Gradient of the redshift to comoving distance conversion if the Universe is
flat and z << 1. flat and z << 1.
Parameters
----------
z : float or 1-dimensional array
Cosmological redshift.
Omega_m : float
Matter density parameter.
Returns
-------
float or 1-dimensional array
""" """
q0 = 3 * Omega_m / 2 - 1 q0 = 3 * Omega_m / 2 - 1
return SPEED_OF_LIGHT / H0 * (1 - z * (1 + q0)) return SPEED_OF_LIGHT / H0 * (1 - z * (1 + q0))
def dist2distmodulus(dist, Omega_m): def dist2distmodulus(dist, Omega_m):
""" """Convert comoving distance to distance modulus, assuming z << 1."""
Convert comoving distance to distance modulus, assuming z << 1.
Parameters
----------
dist : float or 1-dimensional array
Comoving distance in `Mpc / h`.
Omega_m : float
Matter density parameter.
Returns
-------
float or 1-dimensional array
"""
zcosmo = dist2redshift(dist, Omega_m) zcosmo = dist2redshift(dist, Omega_m)
luminosity_distance = dist * (1 + zcosmo) luminosity_distance = dist * (1 + zcosmo)
return 5 * jnp.log10(luminosity_distance) + 25 return 5 * jnp.log10(luminosity_distance) + 25
@ -479,9 +344,8 @@ def dist2distmodulus(dist, Omega_m):
def distmodulus2dist(mu, Omega_m, ninterp=10000, zmax=0.1, mu2comoving=None, def distmodulus2dist(mu, Omega_m, ninterp=10000, zmax=0.1, mu2comoving=None,
return_interpolator=False): return_interpolator=False):
""" """
Convert distance modulus to comoving distance. Note that this is a costly Convert distance modulus to comoving distance. This is costly as it builds
implementation, as it builts up the interpolator every time it is called up the interpolator every time it is called, unless it is provided.
unless it is provided.
Parameters Parameters
---------- ----------
@ -520,9 +384,8 @@ def distmodulus2dist(mu, Omega_m, ninterp=10000, zmax=0.1, mu2comoving=None,
def distmodulus2redsfhit(mu, Omega_m, ninterp=10000, zmax=0.1, mu2z=None, def distmodulus2redsfhit(mu, Omega_m, ninterp=10000, zmax=0.1, mu2z=None,
return_interpolator=False): return_interpolator=False):
""" """
Convert distance modulus to cosmological redshift. Note that this is a Convert distance modulus to cosmological redshift. This is costly as it
costly implementation, as it builts up the interpolator every time it is builts up the interpolator every time it is called, unless it is provided.
called unless it is provided.
Parameters Parameters
---------- ----------
@ -556,49 +419,18 @@ def distmodulus2redsfhit(mu, Omega_m, ninterp=10000, zmax=0.1, mu2z=None,
return mu2z(mu) return mu2z(mu)
def project_Vext(Vext_x, Vext_y, Vext_z, RA, dec): def project_Vext(Vext_x, Vext_y, Vext_z, RA_radians, dec_radians):
""" """Project the external velocity vector onto the line of sight."""
Project the external velocity onto the line of sight along direction cos_dec = jnp.cos(dec_radians)
specified by RA/dec. Note that the angles must be in radians. return (Vext_x * jnp.cos(RA_radians) * cos_dec
+ Vext_y * jnp.sin(RA_radians) * cos_dec
Parameters + Vext_z * jnp.sin(dec_radians))
----------
Vext_x, Vext_y, Vext_z : floats
Components of the external velocity.
RA, dec : floats
Right ascension and declination in radians
Returns
-------
float
"""
cos_dec = jnp.cos(dec)
return (Vext_x * jnp.cos(RA) * cos_dec
+ Vext_y * jnp.sin(RA) * cos_dec
+ Vext_z * jnp.sin(dec))
def predict_zobs(dist, beta, Vext_radial, vpec_radial, Omega_m): def predict_zobs(dist, beta, Vext_radial, vpec_radial, Omega_m):
""" """
Predict the observed redshift at a given comoving distance given some Predict the observed redshift at a given comoving distance given some
velocity field. velocity field.
Parameters
----------
dist : float
Comoving distance in `Mpc / h`.
beta : float
Velocity bias parameter.
Vext_radial : float
Radial component of the external velocity along the LOS.
vpec_radial : float
Radial component of the peculiar velocity along the LOS.
Omega_m : float
Matter density parameter.
Returns
-------
float
""" """
zcosmo = dist2redshift(dist, Omega_m) zcosmo = dist2redshift(dist, Omega_m)
@ -613,7 +445,7 @@ def predict_zobs(dist, beta, Vext_radial, vpec_radial, Omega_m):
def calculate_ptilde_wo_bias(xrange, mu, err, r_squared_xrange=None, def calculate_ptilde_wo_bias(xrange, mu, err, r_squared_xrange=None,
is_err_squared=False): is_err_squared=False):
""" """
Calculate `ptilde(r)` without any bias. Calculate `ptilde(r)` without (im)homogeneous Malmquist bias.
Parameters Parameters
---------- ----------
@ -648,19 +480,6 @@ def calculate_likelihood_zobs(zobs, zobs_pred, sigma_v):
""" """
Calculate the likelihood of the observed redshift given the predicted Calculate the likelihood of the observed redshift given the predicted
redshift. redshift.
Parameters
----------
zobs : float
Observed redshift.
zobs_pred : float
Predicted redshift.
sigma_v : float
Velocity uncertainty.
Returns
-------
float
""" """
dcz = SPEED_OF_LIGHT * (zobs - zobs_pred) dcz = SPEED_OF_LIGHT * (zobs - zobs_pred)
return jnp.exp(-0.5 * (dcz / sigma_v)**2) / jnp.sqrt(2 * np.pi) / sigma_v return jnp.exp(-0.5 * (dcz / sigma_v)**2) / jnp.sqrt(2 * np.pi) / sigma_v
@ -668,7 +487,7 @@ def calculate_likelihood_zobs(zobs, zobs_pred, sigma_v):
def stack_normal(mus, stds): def stack_normal(mus, stds):
""" """
Stack the normal distributions and approximate the stacked distribution Stack normal distributions and approximate the stacked distribution
by a single Gaussian. by a single Gaussian.
Parameters Parameters
@ -690,9 +509,6 @@ def stack_normal(mus, stds):
class BaseFlowValidationModel(ABC): class BaseFlowValidationModel(ABC):
"""
Base class for the flow validation models.
"""
@property @property
def ndata(self): def ndata(self):
@ -813,17 +629,16 @@ class SD_PV_validation_model(BaseFlowValidationModel):
def __init__(self, los_density, los_velocity, RA, dec, z_obs, def __init__(self, los_density, los_velocity, RA, dec, z_obs,
r_hMpc, e_r_hMpc, r_xrange, Omega_m): r_hMpc, e_r_hMpc, r_xrange, Omega_m):
dt = jnp.float32
# Convert everything to JAX arrays. # Convert everything to JAX arrays.
self._los_density = jnp.asarray(los_density, dtype=dt) self._los_density = jnp.asarray(los_density)
self._los_velocity = jnp.asarray(los_velocity, dtype=dt) self._los_velocity = jnp.asarray(los_velocity)
self._RA = jnp.asarray(np.deg2rad(RA), dtype=dt) self._RA = jnp.asarray(np.deg2rad(RA))
self._dec = jnp.asarray(np.deg2rad(dec), dtype=dt) self._dec = jnp.asarray(np.deg2rad(dec))
self._z_obs = jnp.asarray(z_obs, dtype=dt) self._z_obs = jnp.asarray(z_obs)
self._r_hMpc = jnp.asarray(r_hMpc, dtype=dt) self._r_hMpc = jnp.asarray(r_hMpc)
self._e2_rhMpc = jnp.asarray(e_r_hMpc**2, dtype=dt) self._e2_rhMpc = jnp.asarray(e_r_hMpc**2)
# Get radius squared # Get radius squared
r2_xrange = r_xrange**2 r2_xrange = r_xrange**2
@ -941,26 +756,24 @@ class SN_PV_validation_model(BaseFlowValidationModel):
def __init__(self, los_density, los_velocity, RA, dec, z_obs, def __init__(self, los_density, los_velocity, RA, dec, z_obs,
e_zobs, mB, x1, c, e_mB, e_x1, e_c, r_xrange, Omega_m): e_zobs, mB, x1, c, e_mB, e_x1, e_c, r_xrange, Omega_m):
dt = jnp.float32
# Convert everything to JAX arrays. # Convert everything to JAX arrays.
self._los_density = jnp.asarray(los_density, dtype=dt) self._los_density = jnp.asarray(los_density)
self._los_velocity = jnp.asarray(los_velocity, dtype=dt) self._los_velocity = jnp.asarray(los_velocity)
self._RA = jnp.asarray(np.deg2rad(RA), dtype=dt) self._RA = jnp.asarray(np.deg2rad(RA))
self._dec = jnp.asarray(np.deg2rad(dec), dtype=dt) self._dec = jnp.asarray(np.deg2rad(dec))
self._z_obs = jnp.asarray(z_obs, dtype=dt) self._z_obs = jnp.asarray(z_obs)
if e_zobs is not None: if e_zobs is not None:
self._e2_cz_obs = jnp.asarray( self._e2_cz_obs = jnp.asarray((SPEED_OF_LIGHT * e_zobs)**2)
(SPEED_OF_LIGHT * e_zobs)**2, dtype=dt)
else: else:
self._e2_cz_obs = jnp.zeros_like(self._z_obs) self._e2_cz_obs = jnp.zeros_like(self._z_obs)
self._mB = jnp.asarray(mB, dtype=dt) self._mB = jnp.asarray(mB)
self._x1 = jnp.asarray(x1, dtype=dt) self._x1 = jnp.asarray(x1)
self._c = jnp.asarray(c, dtype=dt) self._c = jnp.asarray(c)
self._e2_mB = jnp.asarray(e_mB**2, dtype=dt) self._e2_mB = jnp.asarray(e_mB**2)
self._e2_x1 = jnp.asarray(e_x1**2, dtype=dt) self._e2_x1 = jnp.asarray(e_x1**2)
self._e2_c = jnp.asarray(e_c**2, dtype=dt) self._e2_c = jnp.asarray(e_c**2)
# Get radius squared # Get radius squared
r2_xrange = r_xrange**2 r2_xrange = r_xrange**2
@ -1001,32 +814,12 @@ class SN_PV_validation_model(BaseFlowValidationModel):
def mu(self, mag_cal, alpha_cal, beta_cal): def mu(self, mag_cal, alpha_cal, beta_cal):
""" """
Distance modulus of each object the given SALT2 calibration parameters. Distance modulus of each object given SALT2 calibration parameters.
Parameters
----------
mag_cal, alpha_cal, beta_cal : floats
SALT2 calibration parameters.
Returns
-------
1-dimensional array
""" """
return self._mB - mag_cal + alpha_cal * self._x1 - beta_cal * self._c return self._mB - mag_cal + alpha_cal * self._x1 - beta_cal * self._c
def squared_e_mu(self, alpha_cal, beta_cal, e_mu_intrinsic): def squared_e_mu(self, alpha_cal, beta_cal, e_mu_intrinsic):
""" """Linearly-propagated squared error on the SALT2 distance modulus."""
Linearly-propagated squared error on the SALT2 distance modulus.
Parameters
----------
alpha_cal, beta_cal, e_mu_intrinsic : floats
SALT2 calibration parameters.
Returns
-------
1-dimensional array
"""
return (self._e2_mB + alpha_cal**2 * self._e2_x1 return (self._e2_mB + alpha_cal**2 * self._e2_x1
+ beta_cal**2 * self._e2_c + e_mu_intrinsic**2) + beta_cal**2 * self._e2_c + e_mu_intrinsic**2)
@ -1122,10 +915,6 @@ class SN_PV_validation_model(BaseFlowValidationModel):
sample_beta : bool, optional sample_beta : bool, optional
Whether to sample the velocity bias parameter `beta`, otherwise Whether to sample the velocity bias parameter `beta`, otherwise
it is fixed to 1. it is fixed to 1.
Returns
-------
None
""" """
Vx = numpyro.sample("Vext_x", self._Vext) Vx = numpyro.sample("Vext_x", self._Vext)
Vy = numpyro.sample("Vext_y", self._Vext) Vy = numpyro.sample("Vext_y", self._Vext)
@ -1192,19 +981,18 @@ class TF_PV_validation_model(BaseFlowValidationModel):
def __init__(self, los_density, los_velocity, RA, dec, z_obs, def __init__(self, los_density, los_velocity, RA, dec, z_obs,
mag, eta, e_mag, e_eta, r_xrange, Omega_m): mag, eta, e_mag, e_eta, r_xrange, Omega_m):
dt = jnp.float32
# Convert everything to JAX arrays. # Convert everything to JAX arrays.
self._los_density = jnp.asarray(los_density, dtype=dt) self._los_density = jnp.asarray(los_density)
self._los_velocity = jnp.asarray(los_velocity, dtype=dt) self._los_velocity = jnp.asarray(los_velocity)
self._RA = jnp.asarray(np.deg2rad(RA), dtype=dt) self._RA = jnp.asarray(np.deg2rad(RA))
self._dec = jnp.asarray(np.deg2rad(dec), dtype=dt) self._dec = jnp.asarray(np.deg2rad(dec))
self._z_obs = jnp.asarray(z_obs, dtype=dt) self._z_obs = jnp.asarray(z_obs)
self._mag = jnp.asarray(mag, dtype=dt) self._mag = jnp.asarray(mag)
self._eta = jnp.asarray(eta, dtype=dt) self._eta = jnp.asarray(eta)
self._e2_mag = jnp.asarray(e_mag**2, dtype=dt) self._e2_mag = jnp.asarray(e_mag**2)
self._e2_eta = jnp.asarray(e_eta**2, dtype=dt) self._e2_eta = jnp.asarray(e_eta**2)
# Get radius squared # Get radius squared
r2_xrange = r_xrange**2 r2_xrange = r_xrange**2
@ -1243,34 +1031,11 @@ class TF_PV_validation_model(BaseFlowValidationModel):
self._r_xrange = r_xrange self._r_xrange = r_xrange
def mu(self, a, b): def mu(self, a, b):
""" """Distance modulus of each object given the TFR calibration."""
Distance modulus of each object the given Tully-Fisher calibration.
Parameters
----------
a, b : floats
Tully-Fisher calibration parameters.
Returns
-------
1-dimensional array
"""
return self._mag - (a + b * self._eta) return self._mag - (a + b * self._eta)
def squared_e_mu(self, b, e_mu_intrinsic): def squared_e_mu(self, b, e_mu_intrinsic):
""" """Linearly propagated squared error on the TFR distance modulus."""
Squared error on the Tully-Fisher distance modulus.
Parameters
----------
b, e_mu_intrinsic : floats
Tully-Fisher calibration parameters.
Returns
-------
1-dimensional array
"""
return (self._e2_mag + b**2 * self._e2_eta + e_mu_intrinsic**2) return (self._e2_mag + b**2 * self._e2_eta + e_mu_intrinsic**2)
def predict_zcosmo_from_calibration(self, **kwargs): def predict_zcosmo_from_calibration(self, **kwargs):
@ -1331,10 +1096,6 @@ class TF_PV_validation_model(BaseFlowValidationModel):
sample_beta : bool, optional sample_beta : bool, optional
Whether to sample the velocity bias parameter `beta`, otherwise Whether to sample the velocity bias parameter `beta`, otherwise
it is fixed to 1. it is fixed to 1.
Returns
-------
None
""" """
Vx = numpyro.sample("Vext_x", self._Vext) Vx = numpyro.sample("Vext_x", self._Vext)
Vy = numpyro.sample("Vext_y", self._Vext) Vy = numpyro.sample("Vext_y", self._Vext)
@ -1671,7 +1432,6 @@ def _posterior_element(r, beta, Vext_radial, los_velocity, Omega_m, zobs,
class BaseObserved2CosmologicalRedshift(ABC): class BaseObserved2CosmologicalRedshift(ABC):
"""Base class for `Observed2CosmologicalRedshift`.""" """Base class for `Observed2CosmologicalRedshift`."""
def __init__(self, calibration_samples, r_xrange): def __init__(self, calibration_samples, r_xrange):
dt = jnp.float32
# Check calibration samples input. # Check calibration samples input.
for i, key in enumerate(calibration_samples.keys()): for i, key in enumerate(calibration_samples.keys()):
x = calibration_samples[key] x = calibration_samples[key]
@ -1687,14 +1447,13 @@ class BaseObserved2CosmologicalRedshift(ABC):
if len(x) != ncalibratrion: if len(x) != ncalibratrion:
raise ValueError("Calibration samples do not have the same length.") # noqa raise ValueError("Calibration samples do not have the same length.") # noqa
# Enforce the same data type. calibration_samples[key] = jnp.asarray(x)
calibration_samples[key] = jnp.asarray(x, dtype=dt)
if "alpha" not in calibration_samples: if "alpha" not in calibration_samples:
calibration_samples["alpha"] = jnp.ones(ncalibratrion, dtype=dt) calibration_samples["alpha"] = jnp.ones(ncalibratrion)
if "beta" not in calibration_samples: if "beta" not in calibration_samples:
calibration_samples["beta"] = jnp.ones(ncalibratrion, dtype=dt) calibration_samples["beta"] = jnp.ones(ncalibratrion)
# Get the stepsize, we need it to be constant for Simpson's rule. # Get the stepsize, we need it to be constant for Simpson's rule.
dr = np.diff(r_xrange) dr = np.diff(r_xrange)
@ -1714,18 +1473,7 @@ class BaseObserved2CosmologicalRedshift(ABC):
self._simps = jit(lambda y: simps(y, dr)) self._simps = jit(lambda y: simps(y, dr))
def get_calibration_samples(self, key): def get_calibration_samples(self, key):
""" """Get calibration samples for a given key."""
Get calibration samples for a given key.
Parameters
----------
key : str
Key of the calibration samples.
Returns
-------
1-dimensional array
"""
if key not in self._calibration_samples: if key not in self._calibration_samples:
raise ValueError(f"Key `{key}` not found in calibration samples. Available keys are: `{self.calibration_keys}`.") # noqa raise ValueError(f"Key `{key}` not found in calibration samples. Available keys are: `{self.calibration_keys}`.") # noqa
@ -1733,24 +1481,12 @@ class BaseObserved2CosmologicalRedshift(ABC):
@property @property
def ncalibration_samples(self): def ncalibration_samples(self):
""" """Number of calibration samples."""
Number of calibration samples.
Returns
-------
int
"""
return self._ncalibration_samples return self._ncalibration_samples
@property @property
def calibration_keys(self): def calibration_keys(self):
""" """Calibration sample keys."""
Calibration sample keys.
Returns
-------
list of str
"""
return list(self._calibration_samples.keys()) return list(self._calibration_samples.keys())
@ -1785,22 +1521,8 @@ class Observed2CosmologicalRedshift(BaseObserved2CosmologicalRedshift):
def posterior_mean_std(self, x, px): def posterior_mean_std(self, x, px):
""" """
Calculate the mean and standard deviation of a 1-dimensional PDF. Calculate the mean and standard deviation of a 1-dimensional PDF.
Assumes that the PDF is already normalized. Assumes that the PDF Assumes that the PDF is already normalized and that the spacing is that
spacing is that of `r_xrange` which is inferred when initializing this of `r_xrange` which is inferred when initializing this class.
class.
Parameters
----------
x : 1-dimensional array
Values at which the PDF is evaluated. Note that the PDF must be
normalized.
px : 1-dimensional array
PDF values.
dx
Returns
-------
mu, std : floats
""" """
mu = self._simps(x * px) mu = self._simps(x * px)
std = (self._simps(x**2 * px) - mu**2)**0.5 std = (self._simps(x**2 * px) - mu**2)**0.5

77
csiborgtools/match/obs_to_box.py
View file

@ -13,7 +13,9 @@
# with this program; if not, write to the Free Software Foundation, Inc., # with this program; if not, write to the Free Software Foundation, Inc.,
# 51 Franklin Street, Fifth Floor, Boston, MA 02110-1301, USA. # 51 Franklin Street, Fifth Floor, Boston, MA 02110-1301, USA.
""" """
Code to match observations to a constrained simulation. Code to match observations to a constrained simulation. Throughout, masses are
assumed to be in `Msun / h`, distances in `Mpc / h` and the HMF in
`h^3 Mpc^-3 dex^-1`.
""" """
from abc import ABC from abc import ABC
@ -30,17 +32,10 @@ from tqdm import trange
class BaseMatchingProbability(ABC): class BaseMatchingProbability(ABC):
"""Base class for `MatchingProbability`."""
@property @property
def halo_pos(self): def halo_pos(self):
""" """Halo positions in the constrained simulation."""
Halo positions in the constrained simulation.
Returns
-------
2-dimensional array of shape `(n, 3)`
"""
return self._halo_pos return self._halo_pos
@halo_pos.setter @halo_pos.setter
@ -51,13 +46,7 @@ class BaseMatchingProbability(ABC):
@property @property
def halo_log_mass(self): def halo_log_mass(self):
""" """Halo logarithmic mass in the constrained simulation."""
Halo log mass in the constrained simulation.
Returns
-------
1-dimensional array of shape `(n,)`
"""
return self._halo_log_mass return self._halo_log_mass
@halo_log_mass.setter @halo_log_mass.setter
@ -68,30 +57,11 @@ class BaseMatchingProbability(ABC):
@property @property
def nhalo(self): def nhalo(self):
"""" """"Number of haloes in the constrained simulation."""
Number of haloes in the constrained simulation that are used for
matching.
Returns
-------
int
"""
return self.halo_log_mass.size return self.halo_log_mass.size
def HMF(self, log_mass): def HMF(self, log_mass):
""" """Evaluate the halo mass function at a given log mass."""
Evaluate the halo mass function at a given mass.
Parameters
----------
log_mass : float
Logarithmic mass of the halo in `Msun / h`.
Returns
-------
HMF : float
The HMF in `h^3 Mpc^-3 dex^-1`.
"""
return self._hmf(log_mass) return self._hmf(log_mass)
@ -140,17 +110,6 @@ class MatchingProbability(BaseMatchingProbability):
""" """
Calculate the PDF of finding a halo of a given mass at a given distance Calculate the PDF of finding a halo of a given mass at a given distance
from a random point. from a random point.
Parameters
----------
r : float
Distance from the random point in `Mpc / h`.
log_mass : float
Logarithmic mass of the halo in `Msun / h`.
Returns
-------
float
""" """
nd = self.HMF(log_mass) nd = self.HMF(log_mass)
return 4 * np.pi * r**2 * nd * np.exp(-4 / 3 * np.pi * r**3 * nd) return 4 * np.pi * r**2 * nd * np.exp(-4 / 3 * np.pi * r**3 * nd)
@ -159,17 +118,6 @@ class MatchingProbability(BaseMatchingProbability):
""" """
Calculate the CDF of finding a halo of a given mass at a given distance Calculate the CDF of finding a halo of a given mass at a given distance
from a random point. from a random point.
Parameters
----------
r : float
Distance from the random point in `Mpc / h`.
log_mass : float
Logarithmic mass of the halo in `Msun / h`.
Returns
-------
float
""" """
nd = self.HMF(log_mass) nd = self.HMF(log_mass)
return 1 - np.exp(-4 / 3 * np.pi * r**3 * nd) return 1 - np.exp(-4 / 3 * np.pi * r**3 * nd)
@ -178,17 +126,6 @@ class MatchingProbability(BaseMatchingProbability):
""" """
Calculate the inverse CDF of finding a halo of a given mass at a given Calculate the inverse CDF of finding a halo of a given mass at a given
distance from a random point. distance from a random point.
Parameters
----------
cdf : float
CDF of finding a halo of a given mass at a given distance.
log_mass : float
Logarithmic mass of the halo in `Msun / h`.
Returns
-------
float
""" """
nd = self.HMF(log_mass) nd = self.HMF(log_mass)
return (np.log(1 - cdf) / (-4 / 3 * np.pi * nd))**(1 / 3) return (np.log(1 - cdf) / (-4 / 3 * np.pi * nd))**(1 / 3)

14
csiborgtools/match/overlap.py
View file

@ -13,7 +13,7 @@
# with this program; if not, write to the Free Software Foundation, Inc., # with this program; if not, write to the Free Software Foundation, Inc.,
# 51 Franklin Street, Fifth Floor, Boston, MA 02110-1301, USA. # 51 Franklin Street, Fifth Floor, Boston, MA 02110-1301, USA.
""" """
Support for matching halos between CSiBORG IC realisations based on their Code for matching halos between CSiBORG IC realisations based on their
Lagrangian patch overlap. Lagrangian patch overlap.
""" """
from abc import ABC from abc import ABC
@ -112,10 +112,6 @@ class RealisationsMatcher(BaseMatcher):
""" """
Multiplier of the sum of the initial Lagrangian patch sizes of a halo Multiplier of the sum of the initial Lagrangian patch sizes of a halo
pair. Determines the range within which neighbors are returned. pair. Determines the range within which neighbors are returned.
Returns
-------
float
""" """
return self._nmult return self._nmult
@ -130,10 +126,6 @@ class RealisationsMatcher(BaseMatcher):
""" """
Tolerance on the absolute logarithmic mass difference of potential Tolerance on the absolute logarithmic mass difference of potential
matches. matches.
Returns
-------
float
""" """
return self._dlogmass return self._dlogmass
@ -148,10 +140,6 @@ class RealisationsMatcher(BaseMatcher):
""" """
Mass key whose similarity is to be checked. Must be a valid key in the Mass key whose similarity is to be checked. Must be a valid key in the
halo catalogue. halo catalogue.
Returns
-------
str
""" """
return self._mass_key return self._mass_key

63
csiborgtools/params.py
View file

@ -16,6 +16,7 @@
Various user parameters for CSiBORGTools. Various user parameters for CSiBORGTools.
""" """
SPEED_OF_LIGHT = 299792.458 # km / s
CB2_REDSHIFT = [69.0000210000063, 40.250007218751264, 28.24050991940438, CB2_REDSHIFT = [69.0000210000063, 40.250007218751264, 28.24050991940438,
21.6470609550175, 17.480001404480106, 14.608109099433955, 21.6470609550175, 17.480001404480106, 14.608109099433955,
12.508772664512199, 10.90721705951751, 9.64516173673259, 12.508772664512199, 10.90721705951751, 9.64516173673259,
@ -53,39 +54,18 @@ CB2_REDSHIFT = [69.0000210000063, 40.250007218751264, 28.24050991940438,
def snap2redshift(snapnum, simname): def snap2redshift(snapnum, simname):
""" """Convert a snapshot number to redshift."""
Convert a snapshot number to redshift.
Parameters
----------
snapnum : int
Snapshot number.
simname : str
Simulation name.
Returns
-------
float
"""
if "csiborg2_" in simname: if "csiborg2_" in simname:
return CB2_REDSHIFT[snapnum] try:
return CB2_REDSHIFT[snapnum]
except KeyError:
raise ValueError(f"Unknown snapshot: `{snapnum}`.")
else: else:
raise ValueError(f"Unknown simname: {simname}") raise ValueError(f"Unsupported simulation: `{simname}`.")
def simname2boxsize(simname): def simname2boxsize(simname):
""" """Return boxsize in `Mpc/h` for a given simulation."""
Return boxsize in `Mpc/h` for a given simname.
Parameters
----------
simname : str
Simulation name.
Returns
-------
boxsize : float
"""
d = {"csiborg1": 677.7, d = {"csiborg1": 677.7,
"csiborg2_main": 676.6, "csiborg2_main": 676.6,
"csiborg2_varysmall": 676.6, "csiborg2_varysmall": 676.6,
@ -97,28 +77,16 @@ def simname2boxsize(simname):
"TNG300-1": 205., "TNG300-1": 205.,
"Carrick2015": 400., "Carrick2015": 400.,
} }
boxsize = d.get(simname, None) boxsize = d.get(simname, None)
if boxsize is None: if boxsize is None:
raise ValueError("Unknown simname: {}".format(simname)) raise ValueError(f"Unknown simulation: `{simname}`.")
return boxsize return boxsize
def simname2Omega_m(simname): def simname2Omega_m(simname):
""" """Return `Omega_m` for a given simulation"""
Return Omega_m for a given simname.
Parameters
----------
simname : str
Simulation name.
Returns
-------
Omega_m: float
"""
d = {"csiborg1": 0.307, d = {"csiborg1": 0.307,
"csiborg2_main": 0.3111, "csiborg2_main": 0.3111,
"csiborg2_random": 0.3111, "csiborg2_random": 0.3111,
@ -132,7 +100,7 @@ def simname2Omega_m(simname):
omega_m = d.get(simname, None) omega_m = d.get(simname, None)
if omega_m is None: if omega_m is None:
raise ValueError("Unknown simname: {}".format(simname)) raise ValueError(f"Unknown simulation: `{simname}`.")
return omega_m return omega_m
@ -141,7 +109,7 @@ paths_glamdring = {
"csiborg1_srcdir": "/mnt/extraspace/rstiskalek/csiborg1", "csiborg1_srcdir": "/mnt/extraspace/rstiskalek/csiborg1",
"csiborg2_main_srcdir": "/mnt/extraspace/rstiskalek/csiborg2_main", "csiborg2_main_srcdir": "/mnt/extraspace/rstiskalek/csiborg2_main",
"csiborg2_varysmall_srcdir": "/mnt/extraspace/rstiskalek/csiborg2_varysmall", # noqa "csiborg2_varysmall_srcdir": "/mnt/extraspace/rstiskalek/csiborg2_varysmall", # noqa
"csiborg2_random_srcdir": "/mnt/extraspace/rstiskalek/csiborg2_random", # noqa "csiborg2_random_srcdir": "/mnt/extraspace/rstiskalek/csiborg2_random",
"postdir": "/mnt/extraspace/rstiskalek/csiborg_postprocessing/", "postdir": "/mnt/extraspace/rstiskalek/csiborg_postprocessing/",
"quijote_dir": "/mnt/extraspace/rstiskalek/quijote", "quijote_dir": "/mnt/extraspace/rstiskalek/quijote",
"borg1_dir": "/mnt/users/hdesmond/BORG_final", "borg1_dir": "/mnt/users/hdesmond/BORG_final",
@ -149,10 +117,3 @@ paths_glamdring = {
"tng300_1_dir": "/mnt/extraspace/rstiskalek/TNG300-1/", "tng300_1_dir": "/mnt/extraspace/rstiskalek/TNG300-1/",
"aux_cat_dir": "/mnt/extraspace/rstiskalek/catalogs", "aux_cat_dir": "/mnt/extraspace/rstiskalek/catalogs",
} }
# neighbour_kwargs = {"rmax_radial": 155 / 0.705,
# "nbins_radial": 50,
# "rmax_neighbour": 100.,
# "nbins_neighbour": 150,
# "paths_kind": paths_glamdring}

226
csiborgtools/read/catalogue.py
View file

@ -17,7 +17,7 @@ Unified interface for simulation catalogues. Currently supports CSiBORG1,
CSiBORG2 and Quijote. For specific implementation always check the relevant CSiBORG2 and Quijote. For specific implementation always check the relevant
classes in this module. classes in this module.
""" """
from abc import ABC, abstractproperty from abc import ABC, abstractmethod
from collections import OrderedDict from collections import OrderedDict
from functools import lru_cache from functools import lru_cache
from gc import collect from gc import collect
@ -104,13 +104,7 @@ class BaseCatalogue(ABC):
@property @property
def simname(self): def simname(self):
""" """Simulation name."""
Simulation name.
Returns
-------
str
"""
if self._simname is None: if self._simname is None:
raise RuntimeError("`simname` is not set!") raise RuntimeError("`simname` is not set!")
return self._simname return self._simname
@ -123,13 +117,7 @@ class BaseCatalogue(ABC):
@property @property
def nsim(self): def nsim(self):
""" """Simulation IC realisation index."""
Simulation IC realisation index.
Returns
-------
int
"""
if self._nsim is None: if self._nsim is None:
raise RuntimeError("`nsim` is not set!") raise RuntimeError("`nsim` is not set!")
return self._nsim return self._nsim
@ -142,13 +130,7 @@ class BaseCatalogue(ABC):
@property @property
def nsnap(self): def nsnap(self):
""" """Catalogue snapshot index."""
Catalogue snapshot index.
Returns
-------
int
"""
if self._nsnap is None: if self._nsnap is None:
raise RuntimeError("`nsnap` is not set!") raise RuntimeError("`nsnap` is not set!")
return self._nsnap return self._nsnap
@ -186,26 +168,14 @@ class BaseCatalogue(ABC):
@property @property
def paths(self): def paths(self):
""" """Paths manager."""
Paths manager.
Returns
-------
py:class:`csiborgtools.read.Paths`
"""
if self._paths is None: if self._paths is None:
return Paths(**paths_glamdring) return Paths(**paths_glamdring)
return self._paths return self._paths
@property @property
def boxsize(self): def boxsize(self):
""" """Box size in `cMpc / h`."""
Box size in `cMpc / h`.
Returns
-------
float
"""
if self._boxsize is None: if self._boxsize is None:
raise RuntimeError("`boxsize` is not set!") raise RuntimeError("`boxsize` is not set!")
return self._boxsize return self._boxsize
@ -221,10 +191,6 @@ class BaseCatalogue(ABC):
""" """
Whether to flip the x- and z-coordinates to undo the MUSIC bug to match Whether to flip the x- and z-coordinates to undo the MUSIC bug to match
observations. observations.
Returns
-------
bool
""" """
return self._flip_xz return self._flip_xz
@ -236,13 +202,7 @@ class BaseCatalogue(ABC):
@property @property
def cache_maxsize(self): def cache_maxsize(self):
""" """Maximum length of the cache dictionary."""
Maximum length of the cache dictionary.
Returns
-------
int
"""
if self._cache_maxsize is None: if self._cache_maxsize is None:
raise RuntimeError("`cache_maxsize` is not set!") raise RuntimeError("`cache_maxsize` is not set!")
return self._cache_maxsize return self._cache_maxsize
@ -253,111 +213,57 @@ class BaseCatalogue(ABC):
self._cache_maxsize = cache_maxsize self._cache_maxsize = cache_maxsize
def cache_keys(self): def cache_keys(self):
""" """Current keys of the cache dictionary."""
Current keys of the cache dictionary.
Parameters
----------
list of str
"""
return list(self._cache.keys()) return list(self._cache.keys())
def cache_length(self): def cache_length(self):
""" """Current length of the cache dictionary."""
Current length of the cache dictionary.
Returns
-------
int
"""
return len(self._cache) return len(self._cache)
@abstractproperty @property
@abstractmethod
def coordinates(self): def coordinates(self):
""" """Halo coordinates."""
Halo coordinates.
Returns
-------
2-dimensional array
"""
pass pass
@abstractproperty @property
@abstractmethod
def velocities(self): def velocities(self):
""" """Halo peculiar velocities."""
Halo peculiar velocities.
Returns
-------
2-dimensional array
"""
pass pass
@abstractproperty @property
@abstractmethod
def npart(self): def npart(self):
""" """Number of particles in a halo."""
Number of particles in a halo.
Returns
-------
1-dimensional array
"""
pass pass
@abstractproperty @property
@abstractmethod
def totmass(self): def totmass(self):
""" """Total particle mass of a halo."""
Total particle mass of a halo.
Returns
-------
1-dimensional array
"""
pass pass
@abstractproperty @property
@abstractmethod
def index(self): def index(self):
""" """Halo index."""
Halo index.
Returns
-------
1-dimensional array
"""
pass pass
@abstractproperty @property
@abstractmethod
def lagpatch_coordinates(self): def lagpatch_coordinates(self):
""" """Lagrangian patch coordinates."""
Lagrangian patch coordinates.
Returns
-------
2-dimensional array
"""
pass pass
@abstractproperty @property
@abstractmethod
def lagpatch_radius(self): def lagpatch_radius(self):
""" """Lagrangian patch radius."""
Lagrangian patch radius.
Returns
-------
1-dimensional array
"""
pass pass
@property @property
def observer_location(self): def observer_location(self):
"""
Observer location.
Returns
-------
1-dimensional array
"""
if self._observer_location is None: if self._observer_location is None:
raise RuntimeError("`observer_location` is not set!") raise RuntimeError("`observer_location` is not set!")
return self._observer_location return self._observer_location
@ -371,13 +277,6 @@ class BaseCatalogue(ABC):
@property @property
def observer_velocity(self): def observer_velocity(self):
"""
Observer velocity.
Returns
-------
1-dimensional array
"""
if self._observer_velocity is None: if self._observer_velocity is None:
raise RuntimeError("`observer_velocity` is not set!") raise RuntimeError("`observer_velocity` is not set!")
return self._observer_velocity return self._observer_velocity
@ -562,24 +461,12 @@ class BaseCatalogue(ABC):
self._load_filtered = True self._load_filtered = True
def clear_cache(self): def clear_cache(self):
""" """Clear the cache dictionary."""
Clear the cache dictionary.
Returns
-------
None
"""
self._cache.clear() self._cache.clear()
collect() collect()
def keys(self): def keys(self):
""" """Catalogue keys."""
Catalogue keys.
Returns
-------
list
"""
return self._properties + self._custom_keys return self._properties + self._custom_keys
def pick_fiducial_observer(self, n, rmax): def pick_fiducial_observer(self, n, rmax):
@ -945,13 +832,7 @@ class CSiBORG2MergerTreeReader:
@property @property
def simname(self): def simname(self):
""" """Simulation name."""
Simulation name.
Returns
-------
str
"""
if self._simname is None: if self._simname is None:
raise RuntimeError("`simname` is not set!") raise RuntimeError("`simname` is not set!")
return self._simname return self._simname
@ -964,13 +845,7 @@ class CSiBORG2MergerTreeReader:
@property @property
def nsim(self): def nsim(self):
""" """Simulation IC realisation index."""
Simulation IC realisation index.
Returns
-------
int
"""
if self._nsim is None: if self._nsim is None:
raise RuntimeError("`nsim` is not set!") raise RuntimeError("`nsim` is not set!")
return self._nsim return self._nsim
@ -983,24 +858,12 @@ class CSiBORG2MergerTreeReader:
@property @property
def kind(self): def kind(self):
""" """Simulation kind."""
Simulation kind.
Returns
-------
str
"""
return self._simname.split("_")[-1] return self._simname.split("_")[-1]
@property @property
def paths(self): def paths(self):
""" """Paths manager."""
Paths manager.
Returns
-------
py:class:`csiborgtools.read.Paths`
"""
if self._paths is None: if self._paths is None:
return Paths(**paths_glamdring) return Paths(**paths_glamdring)
return self._paths return self._paths
@ -1207,13 +1070,7 @@ class CSiBORG2SUBFINDCatalogue(BaseCatalogue):
@property @property
def kind(self): def kind(self):
""" """Simulation kind."""
Simulation kind.
Returns
-------
str
"""
return self._simname.split("_")[-1] return self._simname.split("_")[-1]
def _read_subfind_catalogue(self, kind): def _read_subfind_catalogue(self, kind):
@ -1396,24 +1253,19 @@ class QuijoteCatalogue(BaseCatalogue):
class MDPL2Catalogue(BaseCatalogue): class MDPL2Catalogue(BaseCatalogue):
r""" r"""
XXX MDPL2 (FoF) halo catalogue at `z = 0`.
Parameters Parameters
---------- ----------
nsim : int
IC realisation index.
paths : py:class`csiborgtools.read.Paths`, optional paths : py:class`csiborgtools.read.Paths`, optional
Paths object. Paths object.
snapshot : subclass of py:class:`BaseSnapshot`, optional
Snapshot object corresponding to the catalogue.
bounds : dict bounds : dict
Parameter bounds; keys as parameter names, values as (min, max) Parameter bounds; keys as parameter names, values as (min, max)
tuples. Use `dist` for radial distance, `None` for no bound. tuples. Use `dist` for radial distance, `None` for no bound.
observer_velocity : array, optional
Observer's velocity in :math:`\mathrm{km} / \mathrm{s}`.
cache_maxsize : int, optional cache_maxsize : int, optional
Maximum number of cached arrays. Maximum number of cached arrays.
""" """
def __init__(self, paths=None, bounds=None, cache_maxsize=64): def __init__(self, paths=None, bounds=None, cache_maxsize=64):
boxsize = 1000. boxsize = 1000.
super().__init__() super().__init__()

104
csiborgtools/read/obs.py
View file

@ -15,7 +15,7 @@
""" """
Scripts to read in observation. Scripts to read in observation.
""" """
from abc import ABC, abstractproperty from abc import ABC, abstractmethod
from os.path import join from os.path import join
from warnings import warn from warnings import warn
@ -191,26 +191,14 @@ class FitsSurvey(ABC):
@property @property
def file(self): def file(self):
""" """The survey FITS file."""
The survey FITS file.
Returns
-------
file : py:class:`astropy.io.fits.hdu.hdulist.HDUList`
"""
if self._file is None: if self._file is None:
raise ValueError("`file` is not set!") raise ValueError("`file` is not set!")
return self._file return self._file
@property @property
def h(self): def h(self):
""" """Little h."""
Little h.
Returns
-------
h : float
"""
return self._h return self._h
@h.setter @h.setter
@ -240,15 +228,10 @@ class FitsSurvey(ABC):
""" """
return self._routines return self._routines
@abstractproperty @property
@abstractmethod
def size(self): def size(self):
""" """Number of samples in the catalogue."""
Return the number of samples in the catalogue.
Returns
-------
size : int
"""
pass pass
@property @property
@ -259,18 +242,11 @@ class FitsSurvey(ABC):
@property @property
def selection_mask(self): def selection_mask(self):
""" """Selection mask, generated with `fmask` when initialised."""
Selection mask, generated with `fmask` when initialised.
Returns
-------
mask : 1-dimensional boolean array
"""
return self._selection_mask return self._selection_mask
@selection_mask.setter @selection_mask.setter
def selection_mask(self, mask): def selection_mask(self, mask):
"""Set the selection mask."""
if not (isinstance(mask, numpy.ndarray) if not (isinstance(mask, numpy.ndarray)
and mask.ndim == 1 and mask.ndim == 1
and mask.dtype == bool): and mask.dtype == bool):
@ -280,50 +256,21 @@ class FitsSurvey(ABC):
@property @property
def fits_keys(self): def fits_keys(self):
""" """Keys of the FITS file `self.file`."""
Keys of the FITS file `self.file`.
Parameters
----------
keys : list of str
"""
return self.file[1].data.columns.names return self.file[1].data.columns.names
@property @property
def routine_keys(self): def routine_keys(self):
""" """Routine keys."""
Routine keys.
Parameters
----------
keys : list of str
"""
return list(self.routines.keys()) return list(self.routines.keys())
def get_fitsitem(self, key): def get_fitsitem(self, key):
""" """Get a column `key` from the FITS file `self.file`."""
Get a column `key` from the FITS file `self.file`.
Parameters
----------
key : str
FITS key.
Returns
-------
col : 1-dimensional array
"""
return self.file[1].data[key] return self.file[1].data[key]
@property @property
def keys(self): def keys(self):
""" """Routine and FITS keys."""
Routine and FITS keys.
Returns
-------
keys : list of str
"""
return self.routine_keys + self.fits_keys + ["INDEX"] return self.routine_keys + self.fits_keys + ["INDEX"]
def make_mask(self, steps): def make_mask(self, steps):
@ -760,13 +707,7 @@ class BaseSingleObservation(ABC):
@property @property
def mass(self): def mass(self):
""" """Total mass estimate in Msun / h."""
Total mass estimate in Msun / h.
Returns
-------
float
"""
if self._mass is None: if self._mass is None:
raise ValueError("`mass` is not set!") raise ValueError("`mass` is not set!")
return self._mass return self._mass
@ -779,30 +720,15 @@ class BaseSingleObservation(ABC):
def cartesian_pos(self, boxsize): def cartesian_pos(self, boxsize):
""" """
Cartesian position of the observation in Mpc / h, assuming the observer Cartesian position in Mpc / h, assuming the observer is in the centre
is in the centre of the box. of the box.
Parameters
----------
boxsize : float
Box size in Mpc / h.
Returns
-------
1-dimensional array of shape (3,)
""" """
return radec_to_cartesian( return radec_to_cartesian(
self.spherical_pos.reshape(1, 3)).reshape(-1,) + boxsize / 2 self.spherical_pos.reshape(1, 3)).reshape(-1,) + boxsize / 2
@property @property
def name(self): def name(self):
""" """Object name."""
Observated object name.
Returns
-------
str
"""
if self._name is None: if self._name is None:
raise ValueError("`name` is not set!") raise ValueError("`name` is not set!")
return self._name return self._name

155
csiborgtools/read/paths.py
View file

@ -86,18 +86,7 @@ class Paths:
self.aux_cat_dir = aux_cat_dir self.aux_cat_dir = aux_cat_dir
def get_ics(self, simname): def get_ics(self, simname):
""" """Get available IC realisation IDs for a given simulation."""
Get available IC realisation IDs for a given simulation.
Parameters
----------
simname : str
Simulation name.
Returns
-------
ids : 1-dimensional array
"""
if simname == "csiborg1" or simname == "borg1": if simname == "csiborg1" or simname == "borg1":
files = glob(join(self.csiborg1_srcdir, "chain_*")) files = glob(join(self.csiborg1_srcdir, "chain_*"))
files = [int(search(r'chain_(\d+)', f).group(1)) for f in files] files = [int(search(r'chain_(\d+)', f).group(1)) for f in files]
@ -126,20 +115,7 @@ class Paths:
return numpy.sort(files) return numpy.sort(files)
def get_snapshots(self, nsim, simname): def get_snapshots(self, nsim, simname):
""" """List available snapshots of simulation."""
List of available snapshots of simulation.
Parameters
----------
nsim : int
IC realisation index.
simname : str
Simulation name.
Returns
-------
snapshots : 1-dimensional array
"""
if simname == "csiborg1": if simname == "csiborg1":
snaps = glob(join(self.csiborg1_srcdir, f"chain_{nsim}", snaps = glob(join(self.csiborg1_srcdir, f"chain_{nsim}",
"snapshot_*")) "snapshot_*"))
@ -178,22 +154,7 @@ class Paths:
return snaps return snaps
def snapshot(self, nsnap, nsim, simname): def snapshot(self, nsnap, nsim, simname):
""" """Path to a simulation snapshot."""
Path to a simulation snapshot.
Parameters
----------
nsnap : int
Snapshot index. For Quijote, `ICs` indicates the IC snapshot.
nsim : int
IC realisation index.
simname : str
Simulation name.
Returns
-------
str
"""
if simname == "csiborg1": if simname == "csiborg1":
fpath = join(self.csiborg1_srcdir, f"chain_{nsim}", fpath = join(self.csiborg1_srcdir, f"chain_{nsim}",
f"snapshot_{str(nsnap).zfill(5)}.hdf5") f"snapshot_{str(nsnap).zfill(5)}.hdf5")
@ -217,22 +178,7 @@ class Paths:
return fpath return fpath
def snapshot_catalogue(self, nsnap, nsim, simname): def snapshot_catalogue(self, nsnap, nsim, simname):
""" """Path to the halo catalogue of a simulation snapshot."""
Path to the halo catalogue of a simulation snapshot.
Parameters
----------
nsnap : int
Snapshot index.
nsim : int
IC realisation index.
simname : str
Simulation name.
Returns
-------
str
"""
if simname == "csiborg1": if simname == "csiborg1":
return join(self.csiborg1_srcdir, f"chain_{nsim}", return join(self.csiborg1_srcdir, f"chain_{nsim}",
f"fof_{str(nsnap).zfill(5)}.hdf5") f"fof_{str(nsnap).zfill(5)}.hdf5")
@ -254,18 +200,7 @@ class Paths:
raise ValueError(f"Unknown simulation name `{simname}`.") raise ValueError(f"Unknown simulation name `{simname}`.")
def external_halo_catalogue(self, name): def external_halo_catalogue(self, name):
""" """Path to an external halo catalogue."""
Path to an external halo catalogue.
Parameters
----------
name : str
Catalogue name.
Returns
-------
str
"""
if name == "MDPL2": if name == "MDPL2":
return join(self.aux_cat_dir, "MDPL2_FOF_125.hdf5") return join(self.aux_cat_dir, "MDPL2_FOF_125.hdf5")
else: else:
@ -275,17 +210,6 @@ class Paths:
""" """
Path to the Lagrangain patch information of a simulation for halos Path to the Lagrangain patch information of a simulation for halos
defined at z = 0. defined at z = 0.
Parameters
----------
nsim : int
IC realisation index.
simname : str
Simulation name.
Returns
-------
str
""" """
if simname == "csiborg1": if simname == "csiborg1":
return join(self.csiborg1_srcdir, f"chain_{nsim}", return join(self.csiborg1_srcdir, f"chain_{nsim}",
@ -306,20 +230,7 @@ class Paths:
raise ValueError(f"Unknown simulation name `{simname}`.") raise ValueError(f"Unknown simulation name `{simname}`.")
def trees(self, nsim, simname): def trees(self, nsim, simname):
""" """Path to the halo trees of a simulation snapshot."""
Path to the halo trees of a simulation snapshot.
Parameters
----------
nsim : int
IC realisation index.
simname : str
Simulation name.
Returns
-------
str
"""
if simname == "csiborg1": if simname == "csiborg1":
raise ValueError("Trees not available for CSiBORG1.") raise ValueError("Trees not available for CSiBORG1.")
elif simname == "csiborg2_main": elif simname == "csiborg2_main":
@ -377,20 +288,7 @@ class Paths:
return join(fdir, fname) return join(fdir, fname)
def random_mah(self, simname, nsim): def random_mah(self, simname, nsim):
""" """Path to the files containing MAHs from random simulations."""
Path to the files containing MAHs from random simulations.
Parameters
----------
simname : str
Simulation name.
nsim0 : int
IC realisation index of the simulation.
Returns
-------
str
"""
fdir = join(self.postdir, "random_mah") fdir = join(self.postdir, "random_mah")
try_create_directory(fdir) try_create_directory(fdir)
@ -441,26 +339,7 @@ class Paths:
return join(fdir, fname) return join(fdir, fname)
def field(self, kind, MAS, grid, nsim, simname): def field(self, kind, MAS, grid, nsim, simname):
r""" """Path to the files containing the calculated fields."""
Path to the files containing the calculated fields in CSiBORG.
Parameters
----------
kind : str
Field type.
MAS : str
Mass-assignment scheme.
grid : int
Grid size.
nsim : int
IC realisation index.
simname : str
Simulation name.
Returns
-------
str
"""
if simname == "borg2": if simname == "borg2":
return join(self.borg2_dir, f"mcmc_{nsim}.h5") return join(self.borg2_dir, f"mcmc_{nsim}.h5")
@ -493,22 +372,8 @@ class Paths:
def observer_peculiar_velocity(self, MAS, grid, nsim, simname): def observer_peculiar_velocity(self, MAS, grid, nsim, simname):
""" """
Path to the files containing the observer peculiar velocity. Path to the files containing the observer peculiar velocity defined as
the velocity of the centre of the box.
Parameters
----------
MAS : str
Mass-assignment scheme.
grid : int
Grid size.
nsim : int
IC realisation index.
simname : str
Simulation name.
Returns
-------
str
""" """
fdir = join(self.postdir, "environment") fdir = join(self.postdir, "environment")
try_create_directory(fdir) try_create_directory(fdir)

175
csiborgtools/read/snapshot.py
View file

@ -17,7 +17,7 @@ Classes for reading in snapshots and unifying the snapshot interface. Here
should be implemented things such as flipping x- and z-axes, to make sure that should be implemented things such as flipping x- and z-axes, to make sure that
observed RA-dec can be mapped into the simulation box. observed RA-dec can be mapped into the simulation box.
""" """
from abc import ABC, abstractmethod, abstractproperty from abc import ABC, abstractmethod
from os.path import join from os.path import join
import numpy import numpy
@ -62,59 +62,29 @@ class BaseSnapshot(ABC):
@property @property
def nsim(self): def nsim(self):
""" """Simulation index."""
Simulation index.
Returns
-------
int
"""
return self._nsim return self._nsim
@property @property
def nsnap(self): def nsnap(self):
""" """Snapshot index."""
Snapshot index.
Returns
-------
int
"""
return self._nsnap return self._nsnap
@property @property
def simname(self): def simname(self):
""" """Simulation name."""
Simulation name.
Returns
-------
str
"""
if self._simname is None: if self._simname is None:
raise ValueError("Simulation name not set.") raise ValueError("Simulation name not set.")
return self._simname return self._simname
@property @property
def boxsize(self): def boxsize(self):
""" """Simulation boxsize in `cMpc/h`."""
Simulation boxsize in `cMpc/h`.
Returns
-------
float
"""
return simname2boxsize(self.simname) return simname2boxsize(self.simname)
@property @property
def paths(self): def paths(self):
""" """Paths manager."""
Paths manager.
Returns
-------
Paths
"""
if self._paths is None: if self._paths is None:
self._paths = Paths(**paths_glamdring) self._paths = Paths(**paths_glamdring)
return self._paths return self._paths
@ -124,10 +94,6 @@ class BaseSnapshot(ABC):
""" """
Whether to keep the snapshot file open when reading halo particles. Whether to keep the snapshot file open when reading halo particles.
This is useful for repeated access to the snapshot. This is useful for repeated access to the snapshot.
Returns
-------
bool
""" """
return self._keep_snapshot_open return self._keep_snapshot_open
@ -136,61 +102,37 @@ class BaseSnapshot(ABC):
""" """
Whether to flip the x- and z-axes to undo the MUSIC bug so that the Whether to flip the x- and z-axes to undo the MUSIC bug so that the
coordinates are consistent with observations. coordinates are consistent with observations.
Returns
-------
bool
""" """
return self._flip_xz return self._flip_xz
@abstractproperty @property
@abstractmethod
def coordinates(self): def coordinates(self):
""" """Particle coordinates."""
Return the particle coordinates.
Returns
-------
coords : 2-dimensional array
"""
pass pass
@abstractproperty @property
@abstractmethod
def velocities(self): def velocities(self):
""" """Particle velocities."""
Return the particle velocities.
Returns
-------
vel : 2-dimensional array
"""
pass pass
@abstractproperty @property
@abstractmethod
def masses(self): def masses(self):
""" """Particle masses."""
Return the particle masses.
Returns
-------
mass : 1-dimensional array
"""
pass pass
@abstractproperty @property
@abstractmethod
def particle_ids(self): def particle_ids(self):
""" """Particle IDs."""
Return the particle IDs.
Returns
-------
ids : 1-dimensional array
"""
pass pass
@abstractmethod @abstractmethod
def halo_coordinates(self, halo_id, is_group): def halo_coordinates(self, halo_id, is_group):
""" """
Return the halo particle coordinates. Halo particle coordinates.
Parameters Parameters
---------- ----------
@ -253,19 +195,12 @@ class BaseSnapshot(ABC):
@abstractmethod @abstractmethod
def _make_hid2offset(self): def _make_hid2offset(self):
"""
Private class function to make the halo ID to offset dictionary.
"""
pass pass
def open_snapshot(self): def open_snapshot(self):
""" """
Open the snapshot file, particularly used in the context of loading in Open the snapshot file, particularly used in the context of loading in
particles of individual haloes. particles of individual haloes.
Returns
-------
h5py.File
""" """
if not self.keep_snapshot_open: if not self.keep_snapshot_open:
# Check if the snapshot path is set # Check if the snapshot path is set
@ -283,10 +218,6 @@ class BaseSnapshot(ABC):
def close_snapshot(self): def close_snapshot(self):
""" """
Close the snapshot file opened with `open_snapshot`. Close the snapshot file opened with `open_snapshot`.
Returns
-------
None
""" """
if not self.keep_snapshot_open: if not self.keep_snapshot_open:
return return
@ -648,24 +579,12 @@ class BaseField(ABC):
@property @property
def nsim(self): def nsim(self):
""" """Simulation index."""
Simulation index.
Returns
-------
int
"""
return self._nsim return self._nsim
@property @property
def paths(self): def paths(self):
""" """Paths manager."""
Paths manager.
Returns
-------
Paths
"""
if self._paths is None: if self._paths is None:
self._paths = Paths(**paths_glamdring) self._paths = Paths(**paths_glamdring)
return self._paths return self._paths
@ -675,65 +594,22 @@ class BaseField(ABC):
""" """
Whether to flip the x- and z-axes to undo the MUSIC bug so that the Whether to flip the x- and z-axes to undo the MUSIC bug so that the
coordinates are consistent with observations. coordinates are consistent with observations.
Returns
-------
bool
""" """
return self._flip_xz return self._flip_xz
@abstractmethod @abstractmethod
def density_field(self, MAS, grid): def density_field(self, MAS, grid):
""" """Precomputed density field."""
Return the pre-computed density field.
Parameters
----------
MAS : str
Mass assignment scheme.
grid : int
Grid size.
Returns
-------
field : 3-dimensional array
"""
pass pass
@abstractmethod @abstractmethod
def velocity_field(self, MAS, grid): def velocity_field(self, MAS, grid):
""" """Precomputed velocity field."""
Return the pre-computed velocity field.
Parameters
----------
MAS : str
Mass assignment scheme.
grid : int
Grid size.
Returns
-------
field : 4-dimensional array
"""
pass pass
@abstractmethod @abstractmethod
def radial_velocity_field(self, MAS, grid): def radial_velocity_field(self, MAS, grid):
""" """Precomputed radial velocity field."""
Return the pre-computed radial velocity field.
Parameters
----------
MAS : str
Mass assignment scheme.
grid : int
Grid size.
Returns
-------
field : 3-dimensional array
"""
pass pass
@ -1046,5 +922,4 @@ def is_instance_of_base_snapshot_subclass(obj):
Check if `obj` is an instance of a subclass of `BaseSnapshot`. Check if `obj` is an instance of a subclass of `BaseSnapshot`.
""" """
return isinstance(obj, BaseSnapshot) and any( return isinstance(obj, BaseSnapshot) and any(
issubclass(cls, BaseSnapshot) for cls in obj.__class__.__bases__ issubclass(cls, BaseSnapshot) for cls in obj.__class__.__bases__)
)

254
csiborgtools/utils.py
View file

@ -14,6 +14,10 @@
# 51 Franklin Street, Fifth Floor, Boston, MA 02110-1301, USA. # 51 Franklin Street, Fifth Floor, Boston, MA 02110-1301, USA.
""" """
Collection of stand-off utility functions used in the scripts. Collection of stand-off utility functions used in the scripts.
Right ascension and declination is always assumed to be in degrees such that
RA is in [0, 360) and dec is in [-90, 90]. Galactic coordinates are also in
degrees.
""" """
from copy import deepcopy from copy import deepcopy
from datetime import datetime from datetime import datetime
@ -32,23 +36,7 @@ from scipy.stats import multivariate_normal
@jit(nopython=True, fastmath=True, boundscheck=False) @jit(nopython=True, fastmath=True, boundscheck=False)
def center_of_mass(particle_positions, particles_mass, boxsize): def center_of_mass(particle_positions, particles_mass, boxsize):
""" """Calculate the CM, assuming periodic boundary conditions in a cube."""
Calculate the center of mass of a halo while assuming periodic boundary
conditions of a cubical box.
Parameters
----------
particle_positions : 2-dimensional array of shape `(nparticles, 3)`
Particle positions in the box.
particles_mass : 1-dimensional array of shape `(nparticles,)`
Particle masses.
boxsize : float
Box size.
Returns
-------
1-dimensional array of shape `(3,)`
"""
cm = np.zeros(3, dtype=particle_positions.dtype) cm = np.zeros(3, dtype=particle_positions.dtype)
totmass = sum(particles_mass) totmass = sum(particles_mass)
@ -71,21 +59,8 @@ def center_of_mass(particle_positions, particles_mass, boxsize):
@jit(nopython=True, fastmath=True, boundscheck=False) @jit(nopython=True, fastmath=True, boundscheck=False)
def periodic_distance(points, reference_point, boxsize): def periodic_distance(points, reference_point, boxsize):
""" """
Compute the 3D distance between multiple points and a reference point using Compute the 3D distance between multiple points and a reference point
periodic boundary conditions. using periodic boundary conditions.
Parameters
----------
points : 2-dimensional array of shape `(npoints, 3)`
Points to calculate the distance from.
reference_point : 1-dimensional array of shape `(3,)`
Reference point.
boxsize : float
Box size.
Returns
-------
1-dimensional array of shape `(npoints,)`
""" """
npoints = len(points) npoints = len(points)
@ -99,20 +74,7 @@ def periodic_distance(points, reference_point, boxsize):
@jit(nopython=True, fastmath=True, boundscheck=False) @jit(nopython=True, fastmath=True, boundscheck=False)
def periodic_distance_two_points(p1, p2, boxsize): def periodic_distance_two_points(p1, p2, boxsize):
""" """Compute the 3D distance between two points in a periodic box."""
Compute the 3D distance between two points in a periodic box.
Parameters
----------
p1, p2 : 1-dimensional array of shape `(3,)`
Points to calculate the distance between.
boxsize : float
Box size.
Returns
-------
float
"""
half_box = boxsize / 2 half_box = boxsize / 2
dist = 0 dist = 0
@ -129,20 +91,7 @@ def periodic_distance_two_points(p1, p2, boxsize):
@jit(nopython=True, boundscheck=False) @jit(nopython=True, boundscheck=False)
def periodic_wrap_grid(pos, boxsize=1): def periodic_wrap_grid(pos, boxsize=1):
""" """Wrap positions in a periodic box. Overwrites the input array."""
Wrap positions in a periodic box. Overwrites the input array.
Parameters
----------
pos : 2-dimensional array of shape `(npoints, 3)`
Positions to wrap.
boxsize : float, optional
Box size.
Returns
-------
2-dimensional array of shape `(npoints, 3)`
"""
for n in range(pos.shape[0]): for n in range(pos.shape[0]):
for i in range(3): for i in range(3):
if pos[n, i] > boxsize: if pos[n, i] > boxsize:
@ -156,16 +105,8 @@ def periodic_wrap_grid(pos, boxsize=1):
@jit(nopython=True, fastmath=True, boundscheck=False) @jit(nopython=True, fastmath=True, boundscheck=False)
def delta2ncells(field): def delta2ncells(field):
""" """
Calculate the number of cells in `field` that are non-zero. Calculate the number of cells in `field` of shape `(nx, ny, nz)` that are
non-zero.
Parameters
----------
field : 3-dimensional array of shape `(nx, ny, nz)`
Field to calculate the number of non-zero cells.
Returns
-------
int
""" """
tot = 0 tot = 0
imax, jmax, kmax = field.shape imax, jmax, kmax = field.shape
@ -178,23 +119,7 @@ def delta2ncells(field):
def cartesian_to_radec(X, return_degrees=True, origin=[0., 0., 0.]): def cartesian_to_radec(X, return_degrees=True, origin=[0., 0., 0.]):
""" """Calculate the radial distance, RA and deg."""
Calculate the radial distance, RA [0, 360) deg and dec [-90, 90] deg.
Parameters
----------
X : 2-dimensional array of shape `(npoints, 3)`
Cartesian coordinates.
return_degrees : bool, optional
Whether to return the angles in degrees.
origin : 1-dimensional array of shape `(3,)`, optional
Origin of the coordinate system.
Returns
-------
out : 2-dimensional array of shape `(npoints, 3)`
Spherical coordinates: distance, RA and dec.
"""
x, y, z = X[:, 0], X[:, 1], X[:, 2] x, y, z = X[:, 0], X[:, 1], X[:, 2]
x -= origin[0] x -= origin[0]
@ -204,6 +129,7 @@ def cartesian_to_radec(X, return_degrees=True, origin=[0., 0., 0.]):
dist = np.linalg.norm(X, axis=1) dist = np.linalg.norm(X, axis=1)
dec = np.arcsin(z / dist) dec = np.arcsin(z / dist)
ra = np.arctan2(y, x) ra = np.arctan2(y, x)
# Wrapping to ensure RA is in [0, 2pi) (later converted to degrees).
ra[ra < 0] += 2 * np.pi ra[ra < 0] += 2 * np.pi
if return_degrees: if return_degrees:
@ -220,17 +146,8 @@ def cartesian_to_radec(X, return_degrees=True, origin=[0., 0., 0.]):
def radec_to_cartesian(X): def radec_to_cartesian(X):
""" """
Calculate Cartesian coordinates from radial distance, RA [0, 360) deg and Calculate Cartesian coordinates from radial distance, RA and dec
dec [-90, 90] deg. `(npoints, 3)`.
Parameters
----------
X : 2-dimensional array of shape `(npoints, 3)`
Spherical coordinates: distance, RA and dec.
Returns
-------
2-dimensional array of shape `(npoints, 3)`
""" """
dist, ra, dec = X[:, 0], X[:, 1], X[:, 2] dist, ra, dec = X[:, 0], X[:, 1], X[:, 2]
@ -243,19 +160,7 @@ def radec_to_cartesian(X):
def radec_to_galactic(ra, dec): def radec_to_galactic(ra, dec):
""" """Convert right ascension and declination to galactic coordinates."""
Convert right ascension and declination to galactic coordinates (all in
degrees.)
Parameters
----------
ra, dec : float or 1-dimensional array
Right ascension and declination in degrees.
Returns
-------
l, b : float or 1-dimensional array
"""
c = SkyCoord(ra=ra*u.degree, dec=dec*u.degree, frame='icrs') c = SkyCoord(ra=ra*u.degree, dec=dec*u.degree, frame='icrs')
return c.galactic.l.degree, c.galactic.b.degree return c.galactic.l.degree, c.galactic.b.degree
@ -263,19 +168,8 @@ def radec_to_galactic(ra, dec):
@jit(nopython=True, fastmath=True, boundscheck=False) @jit(nopython=True, fastmath=True, boundscheck=False)
def great_circle_distance(x1, x2, in_degrees=True): def great_circle_distance(x1, x2, in_degrees=True):
""" """
Great circle distance between two points on a sphere, defined by RA and Great circle distance between two points, each of shape `(2,)`, specified
dec, both in degrees. by RA an dec.
Parameters
----------
x1, x2 : 1-dimensional arrays of shape `(2,)`
RA and dec in degrees.
in_degrees : bool, optional
Whether the input is in degrees.
Returns
-------
float
""" """
ra1, dec1 = x1 ra1, dec1 = x1
ra2, dec2 = x2 ra2, dec2 = x2
@ -303,19 +197,8 @@ def great_circle_distance(x1, x2, in_degrees=True):
def cosine_similarity(x, y): def cosine_similarity(x, y):
r""" r"""
Calculate the cosine similarity between two Cartesian vectors. Defined Calculate the cosine similarity between two Cartesian vectors. Defined
as :math:`\Sum_{i} x_i y_{i} / (|x| * |y|)`. as :math:`\Sum_{i} x_i y_{i} / (|x| * |y|)`. Optionally, `y` can be a
2-dimensional array of shape `(n_samples, 3)`.
Parameters
----------
x : 1-dimensional array
The first vector.
y : 1- or 2-dimensional array
The second vector. Can be 2-dimensional of shape `(n_samples, 3)`,
in which case the calculation is broadcasted.
Returns
-------
out : float or 1-dimensional array
""" """
if x.ndim != 1: if x.ndim != 1:
raise ValueError("`x` must be a 1-dimensional array.") raise ValueError("`x` must be a 1-dimensional array.")
@ -330,39 +213,19 @@ def cosine_similarity(x, y):
def hms_to_degrees(hours, minutes=None, seconds=None): def hms_to_degrees(hours, minutes=None, seconds=None):
""" """Convert hours, minutes and seconds to degrees."""
Convert hours, minutes and seconds to degrees.
Parameters
----------
hours, minutes, seconds : float
Returns
-------
float
"""
return hours * 15 + (minutes or 0) / 60 * 15 + (seconds or 0) / 3600 * 15 return hours * 15 + (minutes or 0) / 60 * 15 + (seconds or 0) / 3600 * 15
def dms_to_degrees(degrees, arcminutes=None, arcseconds=None): def dms_to_degrees(degrees, arcminutes=None, arcseconds=None):
""" """Convert degrees, arcminutes and arcseconds to decimal degrees."""
Convert degrees, arcminutes and arcseconds to decimal degrees.
Parameters
----------
degrees, arcminutes, arcseconds : float
Returns
-------
float
"""
return degrees + (arcminutes or 0) / 60 + (arcseconds or 0) / 3600 return degrees + (arcminutes or 0) / 60 + (arcseconds or 0) / 3600
def real2redshift(pos, vel, observer_location, observer_velocity, boxsize, def real2redshift(pos, vel, observer_location, observer_velocity, boxsize,
periodic_wrap=True, make_copy=True): periodic_wrap=True, make_copy=True):
r""" r"""
Convert real-space position to redshift space position. Convert real space position to redshift space position.
Parameters Parameters
---------- ----------
@ -461,20 +324,7 @@ def heliocentric_to_cmb(z_helio, RA, dec, e_z_helio=None):
@jit(nopython=True, fastmath=True, boundscheck=False) @jit(nopython=True, fastmath=True, boundscheck=False)
def number_counts(x, bin_edges): def number_counts(x, bin_edges):
""" """Calculate counts of samples in bins."""
Calculate counts of samples in bins.
Parameters
----------
x : 1-dimensional array
Samples to bin.
bin_edges : 1-dimensional array
Bin edges.
Returns
-------
1-dimensional array
"""
out = np.full(bin_edges.size - 1, np.nan, dtype=np.float32) out = np.full(bin_edges.size - 1, np.nan, dtype=np.float32)
for i in range(bin_edges.size - 1): for i in range(bin_edges.size - 1):
out[i] = np.sum((x >= bin_edges[i]) & (x < bin_edges[i + 1])) out[i] = np.sum((x >= bin_edges[i]) & (x < bin_edges[i + 1]))
@ -483,22 +333,7 @@ def number_counts(x, bin_edges):
def binned_statistic(x, y, left_edges, bin_width, statistic): def binned_statistic(x, y, left_edges, bin_width, statistic):
""" """
Calculate a binned statistic. Calculate a binned statistic, `statistic` must be a callable `f(x)`.
Parameters
----------
x, y : 1-dimensional arrays
Values by which to bin and calculate the statistic on, respectively.
left_edges : 1-dimensional array
Left edges of the bins.
bin_width : float
Width of the bins.
statistic : callable
Function to calculate the statistic, must be `f(x)`.
Returns
-------
1-dimensional array
""" """
out = np.full(left_edges.size, np.nan, dtype=x.dtype) out = np.full(left_edges.size, np.nan, dtype=x.dtype)
@ -525,15 +360,6 @@ def calculate_acf(data):
""" """
Calculates the autocorrelation of some data. Taken from `epsie` package Calculates the autocorrelation of some data. Taken from `epsie` package
written by Collin Capano. written by Collin Capano.
Parameters
----------
data : 1-dimensional array
The data to calculate the autocorrelation of.
Returns
-------
acf : 1-dimensional array
""" """
# zero the mean # zero the mean
data = data - data.mean() data = data - data.mean()
@ -553,19 +379,9 @@ def calculate_acf(data):
def calculate_acl(data): def calculate_acl(data):
""" """
Calculate the autocorrelation length of some data. Taken from `epsie` Calculate the autocorrelation length of some data. Taken from `epsie`
package written by Collin Capano. Algorithm used is from: package written by Collin Capano. Algorithm used is from: N. Madras and
N. Madras and A.D. Sokal, J. Stat. Phys. 50, 109 (1988). A.D. Sokal, J. Stat. Phys. 50, 109 (1988).
Parameters
----------
data : 1-dimensional array
The data to calculate the autocorrelation length of.
Returns
-------
acl : int
""" """
# calculate the acf
acf = calculate_acf(data) acf = calculate_acf(data)
# now the ACL: Following from Sokal, this is estimated # now the ACL: Following from Sokal, this is estimated
# as the first point where M*tau[k] <= k, where # as the first point where M*tau[k] <= k, where
@ -585,20 +401,8 @@ def calculate_acl(data):
def thin_samples_by_acl(samples): def thin_samples_by_acl(samples):
""" """
Thin MCMC samples by the autocorrelation length of each chain. Thin MCMC samples (dictionary with arrays of shape `(nchains, nsamples)`)
by the autocorrelation length of each chain and concatenate the chains.
Parameters
----------
samples : dict
Dictionary of samples. Each value is a 2-dimensional array of shape
`(nchains, nsamples)`.
Returns
-------
thinned_samples : dict
Dictionary of thinned samples. Each value is a 1-dimensional array of
shape `(n_thinned_samples)`, where the samples are concatenated across
the chain.
""" """
keys = list(samples.keys()) keys = list(samples.keys())
nchains = 1 if samples[keys[0]].ndim == 1 else samples[keys[0]].shape[0] nchains = 1 if samples[keys[0]].ndim == 1 else samples[keys[0]].shape[0]