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

Add simple distance flow model (#114)

Browse source 
* Add imports

* Add field LOS paths

* Add basic flow model

* Edit script

* Add nb

* Add nb

* Update nb

* Add some docs

* Add RA reading

* Add imoprts

* Updates to the flow model

* Update script

* Bring back A2

* Update imports

* Update imports

* Add Carrick to ICs

* Add Carrick boxsize

* Add Carrick and fix minor bugs

* Add Carrick box

* Update script

* Edit imports

* Add fixed flow!

* Update omega_m and add it

* Update nb

* Update nb

* Update nb

* Remove old print statements

* Update params

* Add thinning of chains

* Add import

* Add flow validation script

* Add submit script

* Add ksmooth

* Update nb

* Update params

* Update script

* Update string

* Move where distributions are defined

* Add density bias parameter

* Add lognorm mean

* Update scripts

* Update script
This commit is contained in:
Richard Stiskalek 2024-03-08 10:44:19 +00:00 • committed by GitHub
parent fb93f85543
commit a65e3cb15b
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
14 changed files with 1762 additions and 60 deletions
Download patch file Download diff file Expand all files Collapse all files

7
csiborgtools/__init__.py
View file

@ -12,14 +12,15 @@
# 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.
from csiborgtools import clustering, field, halo, match, read, summary # noqa
from csiborgtools import clustering, field, flow, halo, match, read, summary # noqa
from .utils import (center_of_mass, delta2ncells, number_counts, # noqa
periodic_distance, periodic_distance_two_points, # noqa
binned_statistic, cosine_similarity, fprint, # noqa
hms_to_degrees, dms_to_degrees, great_circle_distance, # noqa
radec_to_cartesian, cartesian_to_radec) # noqa
from .params import paths_glamdring, simname2boxsize # noqa
radec_to_cartesian, cartesian_to_radec, # noqa
thin_samples_by_acl) # noqa
from .params import paths_glamdring, simname2boxsize, simname2Omega_m # noqa
###############################################################################

17
csiborgtools/flow/__init__.py Normal file
View file

@ -0,0 +1,17 @@
# Copyright (C) 2024 Richard Stiskalek
# This program is free software; you can redistribute it and/or modify it
# under the terms of the GNU General Public License as published by the
# Free Software Foundation; either version 3 of the License, or (at your
# option) any later version.
#
# This program is distributed in the hope that it will be useful, but
# WITHOUT ANY WARRANTY; without even the implied warranty of
# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the GNU General
# Public License for more details.
#
# You should have received a copy of the GNU General Public License along
# with this program; if not, write to the Free Software Foundation, Inc.,
# 51 Franklin Street, Fifth Floor, Boston, MA 02110-1301, USA.
from .flow_model import (DataLoader, radial_velocity_los, dist2redshift, # noqa
dist2distmodulus, predict_zobs, project_Vext, # noqa
SD_PV_validation_model) # noqa

657
csiborgtools/flow/flow_model.py Normal file
View file

@ -0,0 +1,657 @@
# Copyright (C) 2024 Richard Stiskalek
# This program is free software; you can redistribute it and/or modify it
# under the terms of the GNU General Public License as published by the
# Free Software Foundation; either version 3 of the License, or (at your
# option) any later version.
#
# This program is distributed in the hope that it will be useful, but
# WITHOUT ANY WARRANTY; without even the implied warranty of
# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the GNU General
# Public License for more details.
#
# You should have received a copy of the GNU General Public License along
# with this program; if not, write to the Free Software Foundation, Inc.,
# 51 Franklin Street, Fifth Floor, Boston, MA 02110-1301, USA.
"""
Validation of the CSiBORG velocity field against PV measurements. Based on [1].
References
----------
[1] https://arxiv.org/abs/1912.09383.
"""
from datetime import datetime
from warnings import warn
import numpy as np
import numpyro
import numpyro.distributions as dist
from astropy import units as u
from astropy.coordinates import SkyCoord
from astropy.cosmology import FlatLambdaCDM
from h5py import File
from jax import numpy as jnp
from jax import vmap
from tqdm import tqdm, trange
from ..params import simname2Omega_m
from ..read import CSiBORG1Catalogue
SPEED_OF_LIGHT = 299792.458 # km / s
def t():
"""Shortcut to get the current time."""
return datetime.now().strftime("%H:%M:%S")
def radec_to_galactic(ra, dec):
"""
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')
return c.galactic.l.degree, c.galactic.b.degree
###############################################################################
# Data loader #
###############################################################################
class DataLoader:
"""
Data loader for the line of sight (LOS) interpolated fields and the
corresponding catalogues.
Parameters
----------
simname : str
Simulation name.
catalogue : str
Name of the catalogue with LOS objects.
catalogue_fpath : str
Path to the LOS catalogue file.
paths : csiborgtools.read.Paths
Paths object.
ksmooth : int, optional
Smoothing index.
store_full_velocity : bool, optional
Whether to store the full 3D velocity field. Otherwise stores only
the radial velocity.
"""
def __init__(self, simname, catalogue, catalogue_fpath, paths,
ksmooth=None, store_full_velocity=False):
print(f"{t()}: reading the catalogue.")
self._cat = self._read_catalogue(catalogue, catalogue_fpath)
self._catname = catalogue
print(f"{t()}: reading the interpolated field.")
self._field_rdist, self._los_density, self._los_velocity = self._read_field( # noqa
simname, catalogue, ksmooth, paths)
if len(self._field_rdist) % 2 == 0:
warn(f"The number of radial steps is even. Skipping the first "
f"step at {self._field_rdist[0]} because Simpson's rule "
"requires an odd number of steps.")
self._field_rdist = self._field_rdist[1:]
self._los_density = self._los_density[..., 1:]
self._los_velocity = self._los_velocity[..., 1:]
if len(self._cat) != len(self._los_density):
raise ValueError("The number of objects in the catalogue does not "
"match the number of objects in the field.")
print(f"{t()}: calculating the radial velocity.")
nobject, nsim = self._los_density.shape[:2]
# In case of Carrick 2015 the box is in galactic coordinates..
if simname == "Carrick2015":
d1, d2 = radec_to_galactic(self._cat["RA"], self._cat["DEC"])
else:
d1, d2 = self._cat["RA"], self._cat["DEC"]
radvel = np.empty((nobject, nsim, len(self._field_rdist)),
self._los_velocity.dtype)
for i in trange(nobject):
for j in range(nsim):
radvel[i, j, :] = radial_velocity_los(
self._los_velocity[:, i, j, ...], d1[i], d2[i])
self._los_radial_velocity = radvel
if not store_full_velocity:
self._los_velocity = None
Omega_m = simname2Omega_m(simname)
# Normalize the CSiBORG density by the mean matter density
if "csiborg" in simname:
cosmo = FlatLambdaCDM(H0=100, Om0=Omega_m)
mean_rho_matter = cosmo.critical_density0.to("Msun/kpc^3").value
mean_rho_matter *= Omega_m
self._los_density /= mean_rho_matter
# Since Carrick+2015 provide `rho / <rho> - 1`
if simname == "Carrick2015":
self._los_density += 1
@property
def cat(self):
"""
The distance indicators catalogue.
Returns
-------
structured array
"""
return self._cat
@property
def catname(self):
"""
Name of the catalogue.
Returns
-------
str
"""
return self._catname
@property
def rdist(self):
"""
Radial distances where the field was interpolated for each object.
Returns
-------
1-dimensional array
"""
return self._field_rdist
@property
def los_density(self):
"""
Density field along the line of sight.
Returns
----------
3-dimensional array of shape (n_objects, n_simulations, n_steps)
"""
return self._los_density
@property
def los_velocity(self):
"""
Velocity field along the line of sight.
Returns
-------
4-dimensional array of shape (n_objects, n_simulations, 3, n_steps)
"""
if self._los_velocity is None:
raise ValueError("The 3D velocities were not stored.")
return self._los_velocity
@property
def los_radial_velocity(self):
"""
Radial velocity along the line of sight.
Returns
-------
3-dimensional array of shape (n_objects, n_simulations, n_steps)
"""
return self._los_radial_velocity
def _read_field(self, simname, catalogue, k, paths):
"""Read in the interpolated field."""
out_density = None
out_velocity = None
has_smoothed = False
nsims = paths.get_ics(simname)
with File(paths.field_los(simname, catalogue), 'r') as f:
has_smoothed = True if f[f"density_{nsims[0]}"].ndim > 2 else False
if has_smoothed and (k is None or not isinstance(k, int)):
raise ValueError("The output contains smoothed field but "
"`ksmooth` is None. It must be provided.")
for i, nsim in enumerate(tqdm(nsims)):
if out_density is None:
nobject, nstep = f[f"density_{nsim}"].shape[:2]
out_density = np.empty(
(nobject, len(nsims), nstep), dtype=np.float32)
out_velocity = np.empty(
(3, nobject, len(nsims), nstep), dtype=np.float32)
indx = (..., k) if has_smoothed else (...)
out_density[:, i, :] = f[f"density_{nsim}"][indx]
out_velocity[:, :, i, :] = f[f"velocity_{nsim}"][indx]
rdist = f[f"rdist_{nsims[0]}"][:]
return rdist, out_density, out_velocity
def _read_catalogue(self, catalogue, catalogue_fpath):
"""
Read in the distance indicator catalogue.
"""
if catalogue == "A2":
with File(catalogue_fpath, 'r') as f:
dtype = [(key, np.float32) for key in f.keys()]
arr = np.empty(len(f["RA"]), dtype=dtype)
for key in f.keys():
arr[key] = f[key][:]
elif catalogue == "LOSS" or catalogue == "Foundation":
with File(catalogue_fpath, 'r') as f:
grp = f[catalogue]
dtype = [(key, np.float32) for key in grp.keys()]
arr = np.empty(len(grp["RA"]), dtype=dtype)
for key in grp.keys():
arr[key] = grp[key][:]
elif "csiborg1" in catalogue:
nsim = int(catalogue.split("_")[-1])
cat = CSiBORG1Catalogue(nsim, bounds={"totmass": (1e13, None)})
seed = 42
gen = np.random.default_rng(seed)
mask = gen.choice(len(cat), size=100, replace=False)
keys = ["r_hMpc", "RA", "DEC"]
dtype = [(key, np.float32) for key in keys]
arr = np.empty(len(mask), dtype=dtype)
sph_pos = cat["spherical_pos"]
arr["r_hMpc"] = sph_pos[mask, 0]
arr["RA"] = sph_pos[mask, 1]
arr["DEC"] = sph_pos[mask, 2]
# TODO: add peculiar velocit
else:
raise ValueError(f"Unknown catalogue: `{catalogue}`.")
return arr
###############################################################################
# Supplementary flow functions #
###############################################################################
def radial_velocity_los(los_velocity, ra, dec):
"""
Calculate the radial velocity along the line of sight.
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)
if not isinstance(ra, types) and not isinstance(dec, types):
raise ValueError("RA and dec must be floats.")
if los_velocity.ndim != 2 and los_velocity.shape[0] != 3:
raise ValueError("The shape of `los_velocity` must be (3, n_steps).")
ra_rad, dec_rad = np.deg2rad(ra), np.deg2rad(dec)
vx, vy, vz = los_velocity
return (vx * np.cos(ra_rad) * np.cos(dec_rad)
+ vy * np.sin(ra_rad) * np.cos(dec_rad)
+ vz * np.sin(dec_rad))
###############################################################################
# JAX Flow model #
###############################################################################
def lognorm_mean_std_to_loc_scale(mu, std):
"""
Calculate the location and scale parameters for the log-normal distribution
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)
scale = np.sqrt(np.log(1 + (std / mu) ** 2))
return loc, scale
def simps(y, dx):
"""
Simpson's rule 1D integration, assuming that the number of steps is even
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:
raise ValueError("The number of steps must be odd.")
return dx / 3 * jnp.sum(y[0:-1:2] + 4 * y[1::2] + y[2::2])
def dist2redshift(dist, Omega_m):
"""
Convert comoving distance to cosmological redshift if the Universe is
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
"""
H0 = 100
eta = 3 * Omega_m / 2
return 1 / eta * (1 - (1 - 2 * H0 * dist / SPEED_OF_LIGHT * eta)**0.5)
def dist2distmodulus(dist, Omega_m):
"""
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)
luminosity_distance = dist * (1 + zcosmo)
return 5 * jnp.log10(luminosity_distance) + 25
def project_Vext(Vext_x, Vext_y, Vext_z, RA, dec):
"""
Project the external velocity onto the line of sight along direction
specified by RA/dec. Note that the angles must be in radians.
Parameters
----------
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):
"""
Predict the observed redshift at a given comoving distance given some
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)
vrad = beta * vpec_radial + Vext_radial
return (1 + zcosmo) * (1 + vrad / SPEED_OF_LIGHT) - 1
###############################################################################
# Flow validation models #
###############################################################################
def calculate_ptilde_wo_bias(xrange, mu, err, r_squared_xrange=None):
"""
Calculate `ptilde(r)` without any bias.
Parameters
----------
xrange : 1-dimensional array
Radial distances where the field was interpolated for each object.
mu : float
Comoving distance in `Mpc / h`.
err : float
Error on the comoving distance in `Mpc / h`.
r_squared_xrange : 1-dimensional array, optional
Radial distances squared where the field was interpolated for each
object. If not provided, the `r^2` correction is not applied.
Returns
-------
1-dimensional array
"""
ptilde = jnp.exp(-0.5 * ((xrange - mu) / err)**2)
if r_squared_xrange is not None:
ptilde *= r_squared_xrange
return ptilde
def calculate_ll_zobs(zobs, zobs_pred, sigma_v):
"""
Calculate the likelihood of the observed redshift given the predicted
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)
return jnp.exp(-0.5 * (dcz / sigma_v)**2) / jnp.sqrt(2 * np.pi) / sigma_v
class SD_PV_validation_model:
"""
Simple distance peculiar velocity (PV) validation model, assuming that
we already have a calibrated estimate of the comoving distance to the
objects.
Parameters
----------
los_density : 2-dimensional array of shape (n_objects, n_steps)
LOS density field.
los_velocity : 3-dimensional array of shape (n_objects, n_steps)
LOS radial velocity field.
RA, dec : 1-dimensional arrays of shape (n_objects)
Right ascension and declination in degrees.
z_obs : 1-dimensional array of shape (n_objects)
Observed redshifts.
r_hMpc : 1-dimensional array of shape (n_objects)
Estimated comoving distances in `h^-1 Mpc`.
e_r_hMpc : 1-dimensional array of shape (n_objects)
Errors on the estimated comoving distances in `h^-1 Mpc`.
r_xrange : 1-dimensional array
Radial distances where the field was interpolated for each object.
Omega_m : float
Matter density parameter.
"""
def __init__(self, los_density, los_velocity, RA, dec, z_obs,
r_hMpc, e_r_hMpc, r_xrange, Omega_m):
# Convert everything to JAX arrays.
dt = jnp.float32
self._los_density = jnp.asarray(los_density, dtype=dt)
self._los_velocity = jnp.asarray(los_velocity, dtype=dt)
self._RA = jnp.asarray(np.deg2rad(RA), dtype=dt)
self._dec = jnp.asarray(np.deg2rad(dec), dtype=dt)
self._z_obs = jnp.asarray(z_obs, dtype=dt)
self._r_hMpc = jnp.asarray(r_hMpc, dtype=dt)
self._e_rhMpc = jnp.asarray(e_r_hMpc, dtype=dt)
# Get radius squared
r2_xrange = r_xrange**2
r2_xrange /= r2_xrange.mean()
# Get the stepsize, we need it to be constant for Simpson's rule.
dr = np.diff(r_xrange)
if not np.all(np.isclose(dr, dr[0], atol=1e-5)):
raise ValueError("The radial step size must be constant.")
dr = dr[0]
# Get the various vmapped functions
self._vmap_ptilde_wo_bias = vmap(lambda mu, err: calculate_ptilde_wo_bias(r_xrange, mu, err, r2_xrange)) # noqa
self._vmap_simps = vmap(lambda y: simps(y, dr))
self._vmap_zobs = vmap(lambda beta, Vr, vpec_rad: predict_zobs(r_xrange, beta, Vr, vpec_rad, Omega_m), in_axes=(None, 0, 0)) # noqa
self._vmap_ll_zobs = vmap(lambda zobs, zobs_pred, sigma_v: calculate_ll_zobs(zobs, zobs_pred, sigma_v), in_axes=(0, 0, None)) # noqa
# Vext_x, Vext_y, Vext_z: external velocity components
self._dist_Vext = dist.Uniform(-1000, 1000)
# We want sigma_v to be 150 +- 100 km / s (lognormal)
self._dist_sigma_v = dist.LogNormal(
*lognorm_mean_std_to_loc_scale(150, 100))
# Density power-law bias
self._dist_alpha = dist.LogNormal(
*lognorm_mean_std_to_loc_scale(1.0, 0.5))
# Velocity bias
self._dist_beta = dist.Normal(1., 0.5)
def __call__(self):
"""
The simple distance NumPyro PV validation model. Samples the following
parameters:
- `Vext_x`, `Vext_y`, `Vext_z`: external velocity components
- `alpha`: density bias parameter
- `beta`: velocity bias parameter
- `sigma_v`: velocity uncertainty
"""
Vx = numpyro.sample("Vext_x", self._dist_Vext)
Vy = numpyro.sample("Vext_y", self._dist_Vext)
Vz = numpyro.sample("Vext_z", self._dist_Vext)
alpha = numpyro.sample("alpha", self._dist_alpha)
beta = numpyro.sample("beta", self._dist_beta)
sigma_v = numpyro.sample("sigma_v", self._dist_sigma_v)
Vext_rad = project_Vext(Vx, Vy, Vz, self._RA, self._dec)
# Calculate p(r) and multiply it by the galaxy bias
ptilde = self._vmap_ptilde_wo_bias(self._r_hMpc, self._e_rhMpc)
ptilde *= self._los_density**alpha
# Normalization of p(r)
pnorm = self._vmap_simps(ptilde)
# Calculate p(z_obs) and multiply it by p(r)
zobs_pred = self._vmap_zobs(beta, Vext_rad, self._los_velocity)
ptilde *= self._vmap_ll_zobs(self._z_obs, zobs_pred, sigma_v)
ll = jnp.sum(jnp.log(self._vmap_simps(ptilde) / pnorm))
numpyro.factor("ll", ll)
# def SN_PV_wcal_validation_model(los_overdensity=None, los_velocity=None,
# RA=None, dec=None, z_CMB=None,
# mB=None, x1=None, c=None,
# e_mB=None, e_x1=None, e_c=None,
# mu_xrange=None, r_xrange=None,
# norm_r2_xrange=None, Omega_m=None, dr=None):
# """
# Pass
# """
# Vx = numpyro.sample("Vext_x", dist.Uniform(-1000, 1000))
# Vy = numpyro.sample("Vext_y", dist.Uniform(-1000, 1000))
# Vz = numpyro.sample("Vext_z", dist.Uniform(-1000, 1000))
# beta = numpyro.sample("beta", dist.Uniform(-10, 10))
#
# # TODO: Later sample these as well.
# e_mu_intrinsic = 0.064
# alpha_cal = 0.135
# beta_cal = 2.9
# mag_cal = -18.555
# sigma_v = 112
#
# # TODO: Check these for fiducial values.
# mu = mB - mag_cal + alpha_cal * x1 - beta_cal * c
# squared_e_mu = e_mB**2 + alpha_cal**2 * e_x1**2 + beta_cal**2 * e_c**2
#
# squared_e_mu += e_mu_intrinsic**2
# ll = 0.
# for i in range(len(los_overdensity)):
# # Project the external velocity for this galaxy.
# Vext_rad = project_Vext(Vx, Vy, Vz, RA[i], dec[i])
#
# dmu = mu_xrange - mu[i]
# ptilde = norm_r2_xrange * jnp.exp(-0.5 * dmu**2 / squared_e_mu[i])
# # TODO: Add some bias
# ptilde *= (1 + los_overdensity[i])
#
# zobs_pred = predict_zobs(r_xrange, beta, Vext_rad, los_velocity[i],
# Omega_m)
#
# dczobs = SPEED_OF_LIGHT * (z_CMB[i] - zobs_pred)
#
# ll_zobs = jnp.exp(-0.5 * (dczobs / sigma_v)**2) / sigma_v
#
# ll += jnp.log(simps(ptilde * ll_zobs, dr))
# ll -= jnp.log(simps(ptilde, dr))
#
# numpyro.factor("ll", ll)

31
csiborgtools/params.py
View file

@ -13,7 +13,7 @@
# with this program; if not, write to the Free Software Foundation, Inc.,
# 51 Franklin Street, Fifth Floor, Boston, MA 02110-1301, USA.
"""
Various user parameters for csiborgtools.
Various user parameters for CSiBORGTools.
"""
@ -37,7 +37,8 @@ def simname2boxsize(simname):
"borg1": 677.7,
"borg2": 676.6,
"quijote": 1000.,
"TNG300-1": 205.
"TNG300-1": 205.,
"Carrick2015": 400.,
}
boxsize = d.get(simname, None)
@ -48,6 +49,32 @@ def simname2boxsize(simname):
return boxsize
def simname2Omega_m(simname):
"""
Return Omega_m for a given simname.
Parameters
----------
simname : str
Simulation name.
Returns
-------
Omega_m: float
"""
d = {"csiborg1": 0.307,
"borg1": 0.307,
"Carrick2015": 0.3,
}
omega_m = d.get(simname, None)
if omega_m is None:
raise ValueError("Unknown simname: {}".format(simname))
return omega_m
paths_glamdring = {
"csiborg1_srcdir": "/mnt/extraspace/rstiskalek/csiborg1",
"csiborg2_main_srcdir": "/mnt/extraspace/rstiskalek/csiborg2_main",

21
csiborgtools/read/paths.py
View file

@ -111,6 +111,8 @@ class Paths:
files = glob(join(self.quijote_dir, "fiducial_processed",
"chain_*"))
files = [int(search(r'chain_(\d+)', f).group(1)) for f in files]
elif simname == "Carrick2015":
return [0]
else:
raise ValueError(f"Unknown simulation name `{simname}`.")
@ -653,3 +655,22 @@ class Paths:
str
"""
return self.tng300_1_dir
def field_los(self, simnname, catalogue):
"""
Path to the files containing the line-of-sight fields.
Parameters
----------
simname : str
Simulation name.
catalogue : str
Catalogue name.
Returns
-------
str
"""
fdir = join(self.postdir, "field_los")
try_create_directory(fdir)
return join(fdir, f"los_{catalogue}_{simnname}.hdf5")

181
csiborgtools/utils.py
View file

@ -15,7 +15,9 @@
"""
Collection of stand-off utility functions used in the scripts.
"""
import numpy
from copy import deepcopy
import numpy as np
from numba import jit
from datetime import datetime
@ -30,17 +32,17 @@ 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.
"""
cm = numpy.zeros(3, dtype=particle_positions.dtype)
cm = np.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(particles_mass * numpy.exp(
2j * numpy.pi * particle_positions[:, i] / boxsize))
cm_i = sum(particles_mass * np.exp(
2j * np.pi * particle_positions[:, i] / boxsize))
cm_i /= totmass
cm_i = numpy.arctan2(cm_i.imag, cm_i.real) * boxsize / (2 * numpy.pi)
cm_i = np.arctan2(cm_i.imag, cm_i.real) * boxsize / (2 * np.pi)
if cm_i < 0:
cm_i += boxsize
@ -58,7 +60,7 @@ def periodic_distance(points, reference_point, boxsize):
npoints = len(points)
half_box = boxsize / 2
dist = numpy.zeros(npoints, dtype=points.dtype)
dist = np.zeros(npoints, dtype=points.dtype)
for i in range(npoints):
for j in range(3):
dist_1d = abs(points[i, j] - reference_point[j])
@ -124,15 +126,15 @@ def cartesian_to_radec(X):
"""
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
dist = np.linalg.norm(X, axis=1)
dec = np.arcsin(z / dist)
ra = np.arctan2(y, x)
ra[ra < 0] += 2 * np.pi
ra *= 180 / numpy.pi
dec *= 180 / numpy.pi
ra *= 180 / np.pi
dec *= 180 / np.pi
return numpy.vstack([dist, ra, dec]).T
return np.vstack([dist, ra, dec]).T
def radec_to_cartesian(X):
@ -142,11 +144,11 @@ def radec_to_cartesian(X):
"""
dist, ra, dec = X[:, 0], X[:, 1], X[:, 2]
cdec = numpy.cos(dec * numpy.pi / 180)
return numpy.vstack([
dist * cdec * numpy.cos(ra * numpy.pi / 180),
dist * cdec * numpy.sin(ra * numpy.pi / 180),
dist * numpy.sin(dec * numpy.pi / 180)
cdec = np.cos(dec * np.pi / 180)
return np.vstack([
dist * cdec * np.cos(ra * np.pi / 180),
dist * cdec * np.sin(ra * np.pi / 180),
dist * np.sin(dec * np.pi / 180)
]).T
@ -159,14 +161,14 @@ def great_circle_distance(x1, x2):
ra1, dec1 = x1
ra2, dec2 = x2
ra1 *= numpy.pi / 180
dec1 *= numpy.pi / 180
ra2 *= numpy.pi / 180
dec2 *= numpy.pi / 180
ra1 *= np.pi / 180
dec1 *= np.pi / 180
ra2 *= np.pi / 180
dec2 *= np.pi / 180
return 180 / numpy.pi * numpy.arccos(
numpy.sin(dec1) * numpy.sin(dec2)
+ numpy.cos(dec1) * numpy.cos(dec2) * numpy.cos(ra1 - ra2)
return 180 / np.pi * np.arccos(
np.sin(dec1) * np.sin(dec2)
+ np.cos(dec1) * np.cos(dec2) * np.cos(ra1 - ra2)
)
@ -193,8 +195,8 @@ def cosine_similarity(x, y):
if y.ndim == 1:
y = y.reshape(1, -1)
out = numpy.sum(x * y, axis=1)
out /= numpy.linalg.norm(x) * numpy.linalg.norm(y, axis=1)
out = np.sum(x * y, axis=1)
out /= np.linalg.norm(x) * np.linalg.norm(y, axis=1)
return out[0] if out.size == 1 else out
@ -258,8 +260,8 @@ def real2redshift(pos, vel, observer_location, observer_velocity, boxsize,
Redshift-space Cartesian position in `Mpc / h`.
"""
if make_copy:
pos = numpy.copy(pos)
vel = numpy.copy(vel)
pos = np.copy(pos)
vel = np.copy(vel)
H0_inv = 1. / 100
@ -267,8 +269,8 @@ def real2redshift(pos, vel, observer_location, observer_velocity, boxsize,
pos -= observer_location
vel -= observer_velocity
vr_dot = numpy.einsum('ij,ij->i', pos, vel)
norm2 = numpy.einsum('ij,ij->i', pos, pos)
vr_dot = np.einsum('ij,ij->i', pos, vel)
norm2 = np.einsum('ij,ij->i', pos, pos)
pos *= (1 + H0_inv * vr_dot / norm2).reshape(-1, 1)
@ -293,9 +295,9 @@ def number_counts(x, bin_edges):
"""
Calculate counts of samples in bins.
"""
out = numpy.full(bin_edges.size - 1, numpy.nan, dtype=numpy.float32)
out = np.full(bin_edges.size - 1, np.nan, dtype=np.float32)
for i in range(bin_edges.size - 1):
out[i] = numpy.sum((x >= bin_edges[i]) & (x < bin_edges[i + 1]))
out[i] = np.sum((x >= bin_edges[i]) & (x < bin_edges[i + 1]))
return out
@ -303,12 +305,12 @@ def binned_statistic(x, y, left_edges, bin_width, statistic):
"""
Calculate a binned statistic.
"""
out = numpy.full(left_edges.size, numpy.nan, dtype=x.dtype)
out = np.full(left_edges.size, np.nan, dtype=x.dtype)
for i in range(left_edges.size):
mask = (x >= left_edges[i]) & (x < left_edges[i] + bin_width)
if numpy.any(mask):
if np.any(mask):
out[i] = statistic(y[mask])
return out
@ -317,3 +319,112 @@ def fprint(msg, verbose=True):
"""Print and flush a message with a timestamp."""
if verbose:
print(f"{datetime.now()}: {msg}", flush=True)
###############################################################################
# ACL of MCMC chains #
###############################################################################
def calculate_acf(data):
"""
Calculates the autocorrelation of some data. Taken from `epsie` package
written by Collin Capano.
Parameters
----------
data : 1-dimensional array
The data to calculate the autocorrelation of.
Returns
-------
acf : 1-dimensional array
"""
# zero the mean
data = data - data.mean()
# zero-pad to 2 * nearest power of 2
newlen = int(2**(1 + np.ceil(np.log2(len(data)))))
x = np.zeros(newlen)
x[:len(data)] = data[:]
# correlate
acf = np.correlate(x, x, mode='full')
# drop corrupted region
acf = acf[len(acf)//2:]
# normalize
acf /= acf[0]
return acf
def calculate_acl(data):
"""
Calculate the autocorrelation length of some data. Taken from `epsie`
package written by Collin Capano. Algorithm used is from:
N. Madras and 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)
# now the ACL: Following from Sokal, this is estimated
# as the first point where M*tau[k] <= k, where
# tau = 2*cumsum(acf) - 1, and M is a tuneable parameter,
# generally chosen to be = 5 (which we use here)
m = 5
cacf = 2. * np.cumsum(acf) - 1.
win = m * cacf <= np.arange(len(cacf))
if win.any():
acl = int(np.ceil(cacf[np.where(win)[0][0]]))
else:
# data is too short to estimate the ACL, just choose
# the length of the data
acl = len(data)
return acl
def thin_samples_by_acl(samples):
"""
Thin MCMC samples by the autocorrelation length of each chain.
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())
nchains = 1 if samples[keys[0]].ndim == 1 else samples[keys[0]].shape[0]
samples = deepcopy(samples)
if nchains == 1:
for key in keys:
samples[key] = samples[key].reshape(1, -1)
# Calculate the ACL of each chain.
acl = np.zeros(nchains, dtype=int)
for i in range(nchains):
acl[i] = max(calculate_acl(samples[key][i]) for key in keys)
thinned_samples = {}
for key in keys:
key_samples = []
for i in range(nchains):
key_samples.append(samples[key][i, ::acl[i]])
thinned_samples[key] = np.hstack(key_samples)
return thinned_samples