From d3b4bfd29c8ca08b91180714f57f6c72f6f574f7 Mon Sep 17 00:00:00 2001 From: Richard Stiskalek Date: Fri, 21 Jun 2024 13:35:45 +0100 Subject: [PATCH] Remove overly verbose code (#129) * 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 --- csiborgtools/clustering/utils.py | 11 - csiborgtools/field/utils.py | 19 +- csiborgtools/flow/flow_model.py | 420 ++++++------------------------- csiborgtools/match/obs_to_box.py | 77 +----- csiborgtools/match/overlap.py | 14 +- csiborgtools/params.py | 63 +---- csiborgtools/read/catalogue.py | 226 +++-------------- csiborgtools/read/obs.py | 104 ++------ csiborgtools/read/paths.py | 155 +----------- csiborgtools/read/snapshot.py | 175 ++----------- csiborgtools/utils.py | 254 +++---------------- 11 files changed, 212 insertions(+), 1306 deletions(-) diff --git a/csiborgtools/clustering/utils.py b/csiborgtools/clustering/utils.py index db44b5e..a94c6fc 100644 --- a/csiborgtools/clustering/utils.py +++ b/csiborgtools/clustering/utils.py @@ -138,17 +138,6 @@ def wrapRA(ra, indeg): """ Wrap RA from :math:`[-180, 180)` to :math`[0, 360)` degrees if `indeg` or 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 if numpy.sum(mask) == 0: diff --git a/csiborgtools/field/utils.py b/csiborgtools/field/utils.py index 6033912..0f83c4d 100644 --- a/csiborgtools/field/utils.py +++ b/csiborgtools/field/utils.py @@ -22,9 +22,7 @@ import healpy def force_single_precision(x): - """ - Attempt to convert an array `x` to float 32. - """ + """Attempt to convert an array `x` to float32.""" if x.dtype != numpy.float32: x = x.astype(numpy.float32) return x @@ -32,9 +30,7 @@ def force_single_precision(x): @jit(nopython=True) 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." imax, jmax, kmax = field0.shape @@ -47,17 +43,8 @@ def divide_nonzero(field0, field1): 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. - - Parameters - ---------- - nside : int - HEALPix nside. - - Returns - ------- - angpos : 2-dimensional array of shape (npix, 2) """ pixs = numpy.arange(healpy.nside2npix(nside)) theta, phi = healpy.pix2ang(nside, pixs) diff --git a/csiborgtools/flow/flow_model.py b/csiborgtools/flow/flow_model.py index a1592cf..9517fe8 100644 --- a/csiborgtools/flow/flow_model.py +++ b/csiborgtools/flow/flow_model.py @@ -13,14 +13,15 @@ # 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]. +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 ---------- [1] https://arxiv.org/abs/1912.09383. """ from abc import ABC, abstractmethod -from datetime import datetime from warnings import catch_warnings, simplefilter, warn import numpy as np @@ -40,17 +41,12 @@ from scipy.optimize import fmin_powell from sklearn.model_selection import KFold from tqdm import trange -from ..params import simname2Omega_m -from ..utils import radec_to_galactic +from ..params import SPEED_OF_LIGHT, simname2Omega_m +from ..utils import fprint, radec_to_galactic -SPEED_OF_LIGHT = 299792.458 # km / s H0 = 100 # km / s / Mpc -def t(): - return datetime.now().strftime("%H:%M:%S") - - ############################################################################### # Data loader # ############################################################################### @@ -83,13 +79,11 @@ class DataLoader: """ def __init__(self, simname, ksim, catalogue, catalogue_fpath, paths, ksmooth=None, store_full_velocity=False, verbose=True): - if verbose: - print(f"{t()}: reading the catalogue.", flush=True) + fprint("reading the catalogue,", verbose) self._cat = self._read_catalogue(catalogue, catalogue_fpath) self._catname = catalogue - if verbose: - print(f"{t()}: reading the interpolated field.", flush=True) + fprint("reading the interpolated field,", verbose) self._field_rdist, self._los_density, self._los_velocity = self._read_field( # noqa simname, ksim, catalogue, ksmooth, paths) @@ -106,8 +100,7 @@ class DataLoader: raise ValueError("The number of objects in the catalogue does not " "match the number of objects in the field.") - if verbose: - print(f"{t()}: calculating the radial velocity.", flush=True) + fprint("calculating the radial velocity.", verbose) nobject = len(self._los_density) dtype = self._los_density.dtype @@ -144,74 +137,37 @@ class DataLoader: @property def cat(self): - """ - The distance indicators catalogue. - - Returns - ------- - structured array - """ + """The distance indicators catalogue (structured array).""" return self._cat[self._mask] @property def catname(self): - """ - Name of the catalogue. - - Returns - ------- - str - """ + """Catalogue name.""" return self._catname @property def rdist(self): - """ - Radial distances where the field was interpolated for each object. - - Returns - ------- - 1-dimensional array - """ + """Radial distances at which the field was interpolated.""" return self._field_rdist @property def los_density(self): - """ - Density field along the line of sight. - - Returns - ---------- - 2-dimensional array of shape (n_objects, n_steps) - """ + """Density field along the line of sight `(n_objects, n_steps)`.""" return self._los_density[self._mask] @property def los_velocity(self): - """ - Velocity field along the line of sight. - - Returns - ------- - 3-dimensional array of shape (3, n_objects, n_steps) - """ + """Velocity field along the line of sight `(3, n_objects, n_steps)`.""" if self._los_velocity is None: raise ValueError("The 3D velocities were not stored.") return self._los_velocity[self._mask] @property def los_radial_velocity(self): - """ - Radial velocity along the line of sight. - - Returns - ------- - 2-dimensional array of shape (n_objects, n_steps) - """ + """Radial velocity along the line of sight `(n_objects, n_steps)`.""" return self._los_radial_velocity[self._mask] def _read_field(self, simname, ksim, catalogue, ksmooth, paths): - """Read in the interpolated field.""" nsims = paths.get_ics(simname) if not (0 <= ksim < len(nsims)): raise ValueError("Invalid simulation index.") @@ -236,7 +192,6 @@ class DataLoader: return rdist, los_density, los_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()] @@ -279,20 +234,8 @@ class DataLoader: def make_jackknife_mask(self, i, n_splits, seed=42): """ - Set the jackknife mask to exclude the `i`-th split. - - 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. + Set the internal jackknife mask to exclude the `i`-th split out of + `n_splits`. """ cv = KFold(n_splits=n_splits, shuffle=True, random_state=seed) n = len(self._cat) @@ -320,20 +263,8 @@ class DataLoader: 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) + Calculate the radial velocity along the LOS from the 3D velocity + along the LOS `(3, n_steps)`. """ types = (float, np.float32, np.float64) 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 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)) @@ -378,17 +300,6 @@ 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.") @@ -400,17 +311,6 @@ 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 """ eta = 3 * Omega_m / 2 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 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 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 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 return SPEED_OF_LIGHT / H0 * (1 - z * (1 + q0)) 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 - """ + """Convert comoving distance to distance modulus, assuming z << 1.""" zcosmo = dist2redshift(dist, Omega_m) luminosity_distance = dist * (1 + zcosmo) 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, return_interpolator=False): """ - Convert distance modulus to comoving distance. Note that this is a costly - implementation, as it builts up the interpolator every time it is called - unless it is provided. + Convert distance modulus to comoving distance. This is costly as it builds + up the interpolator every time it is called, unless it is provided. 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, return_interpolator=False): """ - Convert distance modulus to cosmological redshift. Note that this is a - costly implementation, as it builts up the interpolator every time it is - called unless it is provided. + Convert distance modulus to cosmological redshift. This is costly as it + builts up the interpolator every time it is called, unless it is provided. Parameters ---------- @@ -556,49 +419,18 @@ def distmodulus2redsfhit(mu, Omega_m, ninterp=10000, zmax=0.1, mu2z=None, return mu2z(mu) -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 project_Vext(Vext_x, Vext_y, Vext_z, RA_radians, dec_radians): + """Project the external velocity vector onto the line of sight.""" + cos_dec = jnp.cos(dec_radians) + return (Vext_x * jnp.cos(RA_radians) * cos_dec + + Vext_y * jnp.sin(RA_radians) * cos_dec + + Vext_z * jnp.sin(dec_radians)) 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) @@ -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, is_err_squared=False): """ - Calculate `ptilde(r)` without any bias. + Calculate `ptilde(r)` without (im)homogeneous Malmquist bias. Parameters ---------- @@ -648,19 +480,6 @@ def calculate_likelihood_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 @@ -668,7 +487,7 @@ def calculate_likelihood_zobs(zobs, zobs_pred, sigma_v): 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. Parameters @@ -690,9 +509,6 @@ def stack_normal(mus, stds): class BaseFlowValidationModel(ABC): - """ - Base class for the flow validation models. - """ @property def ndata(self): @@ -813,17 +629,16 @@ class SD_PV_validation_model(BaseFlowValidationModel): def __init__(self, los_density, los_velocity, RA, dec, z_obs, r_hMpc, e_r_hMpc, r_xrange, Omega_m): - dt = jnp.float32 # Convert everything to JAX arrays. - self._los_density = jnp.asarray(los_density, dtype=dt) - self._los_velocity = jnp.asarray(los_velocity, dtype=dt) + self._los_density = jnp.asarray(los_density) + self._los_velocity = jnp.asarray(los_velocity) - 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._RA = jnp.asarray(np.deg2rad(RA)) + self._dec = jnp.asarray(np.deg2rad(dec)) + self._z_obs = jnp.asarray(z_obs) - self._r_hMpc = jnp.asarray(r_hMpc, dtype=dt) - self._e2_rhMpc = jnp.asarray(e_r_hMpc**2, dtype=dt) + self._r_hMpc = jnp.asarray(r_hMpc) + self._e2_rhMpc = jnp.asarray(e_r_hMpc**2) # Get radius squared 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, e_zobs, mB, x1, c, e_mB, e_x1, e_c, r_xrange, Omega_m): - dt = jnp.float32 # Convert everything to JAX arrays. - self._los_density = jnp.asarray(los_density, dtype=dt) - self._los_velocity = jnp.asarray(los_velocity, dtype=dt) + self._los_density = jnp.asarray(los_density) + self._los_velocity = jnp.asarray(los_velocity) - 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._RA = jnp.asarray(np.deg2rad(RA)) + self._dec = jnp.asarray(np.deg2rad(dec)) + self._z_obs = jnp.asarray(z_obs) if e_zobs is not None: - self._e2_cz_obs = jnp.asarray( - (SPEED_OF_LIGHT * e_zobs)**2, dtype=dt) + self._e2_cz_obs = jnp.asarray((SPEED_OF_LIGHT * e_zobs)**2) else: self._e2_cz_obs = jnp.zeros_like(self._z_obs) - self._mB = jnp.asarray(mB, dtype=dt) - self._x1 = jnp.asarray(x1, dtype=dt) - self._c = jnp.asarray(c, dtype=dt) - self._e2_mB = jnp.asarray(e_mB**2, dtype=dt) - self._e2_x1 = jnp.asarray(e_x1**2, dtype=dt) - self._e2_c = jnp.asarray(e_c**2, dtype=dt) + self._mB = jnp.asarray(mB) + self._x1 = jnp.asarray(x1) + self._c = jnp.asarray(c) + self._e2_mB = jnp.asarray(e_mB**2) + self._e2_x1 = jnp.asarray(e_x1**2) + self._e2_c = jnp.asarray(e_c**2) # Get radius squared r2_xrange = r_xrange**2 @@ -1001,32 +814,12 @@ class SN_PV_validation_model(BaseFlowValidationModel): def mu(self, mag_cal, alpha_cal, beta_cal): """ - Distance modulus of each object the given SALT2 calibration parameters. - - Parameters - ---------- - mag_cal, alpha_cal, beta_cal : floats - SALT2 calibration parameters. - - Returns - ------- - 1-dimensional array + Distance modulus of each object given SALT2 calibration parameters. """ 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): - """ - 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 - """ + """Linearly-propagated squared error on the SALT2 distance modulus.""" return (self._e2_mB + alpha_cal**2 * self._e2_x1 + beta_cal**2 * self._e2_c + e_mu_intrinsic**2) @@ -1122,10 +915,6 @@ class SN_PV_validation_model(BaseFlowValidationModel): sample_beta : bool, optional Whether to sample the velocity bias parameter `beta`, otherwise it is fixed to 1. - - Returns - ------- - None """ Vx = numpyro.sample("Vext_x", 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, mag, eta, e_mag, e_eta, r_xrange, Omega_m): - dt = jnp.float32 # Convert everything to JAX arrays. - self._los_density = jnp.asarray(los_density, dtype=dt) - self._los_velocity = jnp.asarray(los_velocity, dtype=dt) + self._los_density = jnp.asarray(los_density) + self._los_velocity = jnp.asarray(los_velocity) - 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._RA = jnp.asarray(np.deg2rad(RA)) + self._dec = jnp.asarray(np.deg2rad(dec)) + self._z_obs = jnp.asarray(z_obs) - self._mag = jnp.asarray(mag, dtype=dt) - self._eta = jnp.asarray(eta, dtype=dt) - self._e2_mag = jnp.asarray(e_mag**2, dtype=dt) - self._e2_eta = jnp.asarray(e_eta**2, dtype=dt) + self._mag = jnp.asarray(mag) + self._eta = jnp.asarray(eta) + self._e2_mag = jnp.asarray(e_mag**2) + self._e2_eta = jnp.asarray(e_eta**2) # Get radius squared r2_xrange = r_xrange**2 @@ -1243,34 +1031,11 @@ class TF_PV_validation_model(BaseFlowValidationModel): self._r_xrange = r_xrange def mu(self, a, b): - """ - Distance modulus of each object the given Tully-Fisher calibration. - - Parameters - ---------- - a, b : floats - Tully-Fisher calibration parameters. - - Returns - ------- - 1-dimensional array - """ - + """Distance modulus of each object given the TFR calibration.""" return self._mag - (a + b * self._eta) def squared_e_mu(self, b, e_mu_intrinsic): - """ - Squared error on the Tully-Fisher distance modulus. - - Parameters - ---------- - b, e_mu_intrinsic : floats - Tully-Fisher calibration parameters. - - Returns - ------- - 1-dimensional array - """ + """Linearly propagated squared error on the TFR distance modulus.""" return (self._e2_mag + b**2 * self._e2_eta + e_mu_intrinsic**2) def predict_zcosmo_from_calibration(self, **kwargs): @@ -1331,10 +1096,6 @@ class TF_PV_validation_model(BaseFlowValidationModel): sample_beta : bool, optional Whether to sample the velocity bias parameter `beta`, otherwise it is fixed to 1. - - Returns - ------- - None """ Vx = numpyro.sample("Vext_x", 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): """Base class for `Observed2CosmologicalRedshift`.""" def __init__(self, calibration_samples, r_xrange): - dt = jnp.float32 # Check calibration samples input. for i, key in enumerate(calibration_samples.keys()): x = calibration_samples[key] @@ -1687,14 +1447,13 @@ class BaseObserved2CosmologicalRedshift(ABC): if len(x) != ncalibratrion: raise ValueError("Calibration samples do not have the same length.") # noqa - # Enforce the same data type. - calibration_samples[key] = jnp.asarray(x, dtype=dt) + calibration_samples[key] = jnp.asarray(x) 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: - 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. dr = np.diff(r_xrange) @@ -1714,18 +1473,7 @@ class BaseObserved2CosmologicalRedshift(ABC): self._simps = jit(lambda y: simps(y, dr)) def get_calibration_samples(self, key): - """ - Get calibration samples for a given key. - - Parameters - ---------- - key : str - Key of the calibration samples. - - Returns - ------- - 1-dimensional array - """ + """Get calibration samples for a given key.""" if key not in self._calibration_samples: 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 def ncalibration_samples(self): - """ - Number of calibration samples. - - Returns - ------- - int - """ + """Number of calibration samples.""" return self._ncalibration_samples @property def calibration_keys(self): - """ - Calibration sample keys. - - Returns - ------- - list of str - """ + """Calibration sample keys.""" return list(self._calibration_samples.keys()) @@ -1785,22 +1521,8 @@ class Observed2CosmologicalRedshift(BaseObserved2CosmologicalRedshift): def posterior_mean_std(self, x, px): """ Calculate the mean and standard deviation of a 1-dimensional PDF. - Assumes that the PDF is already normalized. Assumes that the PDF - spacing is that of `r_xrange` which is inferred when initializing this - 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 + Assumes that the PDF is already normalized and that the spacing is that + of `r_xrange` which is inferred when initializing this class. """ mu = self._simps(x * px) std = (self._simps(x**2 * px) - mu**2)**0.5 diff --git a/csiborgtools/match/obs_to_box.py b/csiborgtools/match/obs_to_box.py index c39845b..dfe4bc3 100644 --- a/csiborgtools/match/obs_to_box.py +++ b/csiborgtools/match/obs_to_box.py @@ -13,7 +13,9 @@ # with this program; if not, write to the Free Software Foundation, Inc., # 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 @@ -30,17 +32,10 @@ from tqdm import trange class BaseMatchingProbability(ABC): - """Base class for `MatchingProbability`.""" @property def halo_pos(self): - """ - Halo positions in the constrained simulation. - - Returns - ------- - 2-dimensional array of shape `(n, 3)` - """ + """Halo positions in the constrained simulation.""" return self._halo_pos @halo_pos.setter @@ -51,13 +46,7 @@ class BaseMatchingProbability(ABC): @property def halo_log_mass(self): - """ - Halo log mass in the constrained simulation. - - Returns - ------- - 1-dimensional array of shape `(n,)` - """ + """Halo logarithmic mass in the constrained simulation.""" return self._halo_log_mass @halo_log_mass.setter @@ -68,30 +57,11 @@ class BaseMatchingProbability(ABC): @property def nhalo(self): - """" - Number of haloes in the constrained simulation that are used for - matching. - - Returns - ------- - int - """ + """"Number of haloes in the constrained simulation.""" return self.halo_log_mass.size def HMF(self, 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`. - """ + """Evaluate the halo mass function at a given 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 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) 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 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) 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 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) return (np.log(1 - cdf) / (-4 / 3 * np.pi * nd))**(1 / 3) diff --git a/csiborgtools/match/overlap.py b/csiborgtools/match/overlap.py index 1961ddb..766c93b 100644 --- a/csiborgtools/match/overlap.py +++ b/csiborgtools/match/overlap.py @@ -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. """ -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. """ from abc import ABC @@ -112,10 +112,6 @@ class RealisationsMatcher(BaseMatcher): """ Multiplier of the sum of the initial Lagrangian patch sizes of a halo pair. Determines the range within which neighbors are returned. - - Returns - ------- - float """ return self._nmult @@ -130,10 +126,6 @@ class RealisationsMatcher(BaseMatcher): """ Tolerance on the absolute logarithmic mass difference of potential matches. - - Returns - ------- - float """ 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 halo catalogue. - - Returns - ------- - str """ return self._mass_key diff --git a/csiborgtools/params.py b/csiborgtools/params.py index 4de3c5e..4358b3d 100644 --- a/csiborgtools/params.py +++ b/csiborgtools/params.py @@ -16,6 +16,7 @@ Various user parameters for CSiBORGTools. """ +SPEED_OF_LIGHT = 299792.458 # km / s CB2_REDSHIFT = [69.0000210000063, 40.250007218751264, 28.24050991940438, 21.6470609550175, 17.480001404480106, 14.608109099433955, 12.508772664512199, 10.90721705951751, 9.64516173673259, @@ -53,39 +54,18 @@ CB2_REDSHIFT = [69.0000210000063, 40.250007218751264, 28.24050991940438, def snap2redshift(snapnum, simname): - """ - Convert a snapshot number to redshift. - - Parameters - ---------- - snapnum : int - Snapshot number. - simname : str - Simulation name. - - Returns - ------- - float - """ + """Convert a snapshot number to redshift.""" if "csiborg2_" in simname: - return CB2_REDSHIFT[snapnum] + try: + return CB2_REDSHIFT[snapnum] + except KeyError: + raise ValueError(f"Unknown snapshot: `{snapnum}`.") else: - raise ValueError(f"Unknown simname: {simname}") + raise ValueError(f"Unsupported simulation: `{simname}`.") def simname2boxsize(simname): - """ - Return boxsize in `Mpc/h` for a given simname. - - Parameters - ---------- - simname : str - Simulation name. - - Returns - ------- - boxsize : float - """ + """Return boxsize in `Mpc/h` for a given simulation.""" d = {"csiborg1": 677.7, "csiborg2_main": 676.6, "csiborg2_varysmall": 676.6, @@ -97,28 +77,16 @@ def simname2boxsize(simname): "TNG300-1": 205., "Carrick2015": 400., } - boxsize = d.get(simname, None) if boxsize is None: - raise ValueError("Unknown simname: {}".format(simname)) + raise ValueError(f"Unknown simulation: `{simname}`.") return boxsize def simname2Omega_m(simname): - """ - Return Omega_m for a given simname. - - Parameters - ---------- - simname : str - Simulation name. - - Returns - ------- - Omega_m: float - """ + """Return `Omega_m` for a given simulation""" d = {"csiborg1": 0.307, "csiborg2_main": 0.3111, "csiborg2_random": 0.3111, @@ -132,7 +100,7 @@ def simname2Omega_m(simname): omega_m = d.get(simname, None) if omega_m is None: - raise ValueError("Unknown simname: {}".format(simname)) + raise ValueError(f"Unknown simulation: `{simname}`.") return omega_m @@ -141,7 +109,7 @@ paths_glamdring = { "csiborg1_srcdir": "/mnt/extraspace/rstiskalek/csiborg1", "csiborg2_main_srcdir": "/mnt/extraspace/rstiskalek/csiborg2_main", "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/", "quijote_dir": "/mnt/extraspace/rstiskalek/quijote", "borg1_dir": "/mnt/users/hdesmond/BORG_final", @@ -149,10 +117,3 @@ paths_glamdring = { "tng300_1_dir": "/mnt/extraspace/rstiskalek/TNG300-1/", "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} diff --git a/csiborgtools/read/catalogue.py b/csiborgtools/read/catalogue.py index b73a2bf..e463c36 100644 --- a/csiborgtools/read/catalogue.py +++ b/csiborgtools/read/catalogue.py @@ -17,7 +17,7 @@ Unified interface for simulation catalogues. Currently supports CSiBORG1, CSiBORG2 and Quijote. For specific implementation always check the relevant classes in this module. """ -from abc import ABC, abstractproperty +from abc import ABC, abstractmethod from collections import OrderedDict from functools import lru_cache from gc import collect @@ -104,13 +104,7 @@ class BaseCatalogue(ABC): @property def simname(self): - """ - Simulation name. - - Returns - ------- - str - """ + """Simulation name.""" if self._simname is None: raise RuntimeError("`simname` is not set!") return self._simname @@ -123,13 +117,7 @@ class BaseCatalogue(ABC): @property def nsim(self): - """ - Simulation IC realisation index. - - Returns - ------- - int - """ + """Simulation IC realisation index.""" if self._nsim is None: raise RuntimeError("`nsim` is not set!") return self._nsim @@ -142,13 +130,7 @@ class BaseCatalogue(ABC): @property def nsnap(self): - """ - Catalogue snapshot index. - - Returns - ------- - int - """ + """Catalogue snapshot index.""" if self._nsnap is None: raise RuntimeError("`nsnap` is not set!") return self._nsnap @@ -186,26 +168,14 @@ class BaseCatalogue(ABC): @property def paths(self): - """ - Paths manager. - - Returns - ------- - py:class:`csiborgtools.read.Paths` - """ + """Paths manager.""" if self._paths is None: return Paths(**paths_glamdring) return self._paths @property def boxsize(self): - """ - Box size in `cMpc / h`. - - Returns - ------- - float - """ + """Box size in `cMpc / h`.""" if self._boxsize is None: raise RuntimeError("`boxsize` is not set!") 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 observations. - - Returns - ------- - bool """ return self._flip_xz @@ -236,13 +202,7 @@ class BaseCatalogue(ABC): @property def cache_maxsize(self): - """ - Maximum length of the cache dictionary. - - Returns - ------- - int - """ + """Maximum length of the cache dictionary.""" if self._cache_maxsize is None: raise RuntimeError("`cache_maxsize` is not set!") return self._cache_maxsize @@ -253,111 +213,57 @@ class BaseCatalogue(ABC): self._cache_maxsize = cache_maxsize def cache_keys(self): - """ - Current keys of the cache dictionary. - - Parameters - ---------- - list of str - """ + """Current keys of the cache dictionary.""" return list(self._cache.keys()) def cache_length(self): - """ - Current length of the cache dictionary. - - Returns - ------- - int - """ + """Current length of the cache dictionary.""" return len(self._cache) - @abstractproperty + @property + @abstractmethod def coordinates(self): - """ - Halo coordinates. - - Returns - ------- - 2-dimensional array - """ + """Halo coordinates.""" pass - @abstractproperty + @property + @abstractmethod def velocities(self): - """ - Halo peculiar velocities. - - Returns - ------- - 2-dimensional array - """ + """Halo peculiar velocities.""" pass - @abstractproperty + @property + @abstractmethod def npart(self): - """ - Number of particles in a halo. - - Returns - ------- - 1-dimensional array - """ + """Number of particles in a halo.""" pass - @abstractproperty + @property + @abstractmethod def totmass(self): - """ - Total particle mass of a halo. - - Returns - ------- - 1-dimensional array - """ + """Total particle mass of a halo.""" pass - @abstractproperty + @property + @abstractmethod def index(self): - """ - Halo index. - - Returns - ------- - 1-dimensional array - """ + """Halo index.""" pass - @abstractproperty + @property + @abstractmethod def lagpatch_coordinates(self): - """ - Lagrangian patch coordinates. - - Returns - ------- - 2-dimensional array - """ + """Lagrangian patch coordinates.""" pass - @abstractproperty + @property + @abstractmethod def lagpatch_radius(self): - """ - Lagrangian patch radius. - - Returns - ------- - 1-dimensional array - """ + """Lagrangian patch radius.""" pass @property def observer_location(self): - """ - Observer location. - - Returns - ------- - 1-dimensional array - """ if self._observer_location is None: raise RuntimeError("`observer_location` is not set!") return self._observer_location @@ -371,13 +277,6 @@ class BaseCatalogue(ABC): @property def observer_velocity(self): - """ - Observer velocity. - - Returns - ------- - 1-dimensional array - """ if self._observer_velocity is None: raise RuntimeError("`observer_velocity` is not set!") return self._observer_velocity @@ -562,24 +461,12 @@ class BaseCatalogue(ABC): self._load_filtered = True def clear_cache(self): - """ - Clear the cache dictionary. - - Returns - ------- - None - """ + """Clear the cache dictionary.""" self._cache.clear() collect() def keys(self): - """ - Catalogue keys. - - Returns - ------- - list - """ + """Catalogue keys.""" return self._properties + self._custom_keys def pick_fiducial_observer(self, n, rmax): @@ -945,13 +832,7 @@ class CSiBORG2MergerTreeReader: @property def simname(self): - """ - Simulation name. - - Returns - ------- - str - """ + """Simulation name.""" if self._simname is None: raise RuntimeError("`simname` is not set!") return self._simname @@ -964,13 +845,7 @@ class CSiBORG2MergerTreeReader: @property def nsim(self): - """ - Simulation IC realisation index. - - Returns - ------- - int - """ + """Simulation IC realisation index.""" if self._nsim is None: raise RuntimeError("`nsim` is not set!") return self._nsim @@ -983,24 +858,12 @@ class CSiBORG2MergerTreeReader: @property def kind(self): - """ - Simulation kind. - - Returns - ------- - str - """ + """Simulation kind.""" return self._simname.split("_")[-1] @property def paths(self): - """ - Paths manager. - - Returns - ------- - py:class:`csiborgtools.read.Paths` - """ + """Paths manager.""" if self._paths is None: return Paths(**paths_glamdring) return self._paths @@ -1207,13 +1070,7 @@ class CSiBORG2SUBFINDCatalogue(BaseCatalogue): @property def kind(self): - """ - Simulation kind. - - Returns - ------- - str - """ + """Simulation kind.""" return self._simname.split("_")[-1] def _read_subfind_catalogue(self, kind): @@ -1396,24 +1253,19 @@ class QuijoteCatalogue(BaseCatalogue): class MDPL2Catalogue(BaseCatalogue): r""" - XXX + MDPL2 (FoF) halo catalogue at `z = 0`. Parameters ---------- - nsim : int - IC realisation index. paths : py:class`csiborgtools.read.Paths`, optional Paths object. - snapshot : subclass of py:class:`BaseSnapshot`, optional - Snapshot object corresponding to the catalogue. bounds : dict Parameter bounds; keys as parameter names, values as (min, max) 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 Maximum number of cached arrays. """ + def __init__(self, paths=None, bounds=None, cache_maxsize=64): boxsize = 1000. super().__init__() diff --git a/csiborgtools/read/obs.py b/csiborgtools/read/obs.py index aee5bf5..a32e9ac 100644 --- a/csiborgtools/read/obs.py +++ b/csiborgtools/read/obs.py @@ -15,7 +15,7 @@ """ Scripts to read in observation. """ -from abc import ABC, abstractproperty +from abc import ABC, abstractmethod from os.path import join from warnings import warn @@ -191,26 +191,14 @@ class FitsSurvey(ABC): @property def file(self): - """ - The survey FITS file. - - Returns - ------- - file : py:class:`astropy.io.fits.hdu.hdulist.HDUList` - """ + """The survey FITS file.""" if self._file is None: raise ValueError("`file` is not set!") return self._file @property def h(self): - """ - Little h. - - Returns - ------- - h : float - """ + """Little h.""" return self._h @h.setter @@ -240,15 +228,10 @@ class FitsSurvey(ABC): """ return self._routines - @abstractproperty + @property + @abstractmethod def size(self): - """ - Return the number of samples in the catalogue. - - Returns - ------- - size : int - """ + """Number of samples in the catalogue.""" pass @property @@ -259,18 +242,11 @@ class FitsSurvey(ABC): @property def selection_mask(self): - """ - Selection mask, generated with `fmask` when initialised. - - Returns - ------- - mask : 1-dimensional boolean array - """ + """Selection mask, generated with `fmask` when initialised.""" return self._selection_mask @selection_mask.setter def selection_mask(self, mask): - """Set the selection mask.""" if not (isinstance(mask, numpy.ndarray) and mask.ndim == 1 and mask.dtype == bool): @@ -280,50 +256,21 @@ class FitsSurvey(ABC): @property def fits_keys(self): - """ - Keys of the FITS file `self.file`. - - Parameters - ---------- - keys : list of str - """ + """Keys of the FITS file `self.file`.""" return self.file[1].data.columns.names @property def routine_keys(self): - """ - Routine keys. - - Parameters - ---------- - keys : list of str - """ + """Routine keys.""" return list(self.routines.keys()) def get_fitsitem(self, key): - """ - Get a column `key` from the FITS file `self.file`. - - Parameters - ---------- - key : str - FITS key. - - Returns - ------- - col : 1-dimensional array - """ + """Get a column `key` from the FITS file `self.file`.""" return self.file[1].data[key] @property def keys(self): - """ - Routine and FITS keys. - - Returns - ------- - keys : list of str - """ + """Routine and FITS keys.""" return self.routine_keys + self.fits_keys + ["INDEX"] def make_mask(self, steps): @@ -760,13 +707,7 @@ class BaseSingleObservation(ABC): @property def mass(self): - """ - Total mass estimate in Msun / h. - - Returns - ------- - float - """ + """Total mass estimate in Msun / h.""" if self._mass is None: raise ValueError("`mass` is not set!") return self._mass @@ -779,30 +720,15 @@ class BaseSingleObservation(ABC): def cartesian_pos(self, boxsize): """ - Cartesian position of the observation in Mpc / h, assuming the observer - is in the centre of the box. - - Parameters - ---------- - boxsize : float - Box size in Mpc / h. - - Returns - ------- - 1-dimensional array of shape (3,) + Cartesian position in Mpc / h, assuming the observer is in the centre + of the box. """ return radec_to_cartesian( self.spherical_pos.reshape(1, 3)).reshape(-1,) + boxsize / 2 @property def name(self): - """ - Observated object name. - - Returns - ------- - str - """ + """Object name.""" if self._name is None: raise ValueError("`name` is not set!") return self._name diff --git a/csiborgtools/read/paths.py b/csiborgtools/read/paths.py index 281a999..b4d9023 100644 --- a/csiborgtools/read/paths.py +++ b/csiborgtools/read/paths.py @@ -86,18 +86,7 @@ class Paths: self.aux_cat_dir = aux_cat_dir def get_ics(self, simname): - """ - Get available IC realisation IDs for a given simulation. - - Parameters - ---------- - simname : str - Simulation name. - - Returns - ------- - ids : 1-dimensional array - """ + """Get available IC realisation IDs for a given simulation.""" if simname == "csiborg1" or simname == "borg1": files = glob(join(self.csiborg1_srcdir, "chain_*")) files = [int(search(r'chain_(\d+)', f).group(1)) for f in files] @@ -126,20 +115,7 @@ class Paths: return numpy.sort(files) def get_snapshots(self, nsim, simname): - """ - List of available snapshots of simulation. - - Parameters - ---------- - nsim : int - IC realisation index. - simname : str - Simulation name. - - Returns - ------- - snapshots : 1-dimensional array - """ + """List available snapshots of simulation.""" if simname == "csiborg1": snaps = glob(join(self.csiborg1_srcdir, f"chain_{nsim}", "snapshot_*")) @@ -178,22 +154,7 @@ class Paths: return snaps def snapshot(self, nsnap, nsim, simname): - """ - 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 - """ + """Path to a simulation snapshot.""" if simname == "csiborg1": fpath = join(self.csiborg1_srcdir, f"chain_{nsim}", f"snapshot_{str(nsnap).zfill(5)}.hdf5") @@ -217,22 +178,7 @@ class Paths: return fpath def snapshot_catalogue(self, nsnap, nsim, simname): - """ - 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 - """ + """Path to the halo catalogue of a simulation snapshot.""" if simname == "csiborg1": return join(self.csiborg1_srcdir, f"chain_{nsim}", f"fof_{str(nsnap).zfill(5)}.hdf5") @@ -254,18 +200,7 @@ class Paths: raise ValueError(f"Unknown simulation name `{simname}`.") def external_halo_catalogue(self, name): - """ - Path to an external halo catalogue. - - Parameters - ---------- - name : str - Catalogue name. - - Returns - ------- - str - """ + """Path to an external halo catalogue.""" if name == "MDPL2": return join(self.aux_cat_dir, "MDPL2_FOF_125.hdf5") else: @@ -275,17 +210,6 @@ class Paths: """ Path to the Lagrangain patch information of a simulation for halos defined at z = 0. - - Parameters - ---------- - nsim : int - IC realisation index. - simname : str - Simulation name. - - Returns - ------- - str """ if simname == "csiborg1": return join(self.csiborg1_srcdir, f"chain_{nsim}", @@ -306,20 +230,7 @@ class Paths: raise ValueError(f"Unknown simulation name `{simname}`.") def trees(self, nsim, simname): - """ - Path to the halo trees of a simulation snapshot. - - Parameters - ---------- - nsim : int - IC realisation index. - simname : str - Simulation name. - - Returns - ------- - str - """ + """Path to the halo trees of a simulation snapshot.""" if simname == "csiborg1": raise ValueError("Trees not available for CSiBORG1.") elif simname == "csiborg2_main": @@ -377,20 +288,7 @@ class Paths: return join(fdir, fname) def random_mah(self, simname, nsim): - """ - Path to the files containing MAHs from random simulations. - - Parameters - ---------- - simname : str - Simulation name. - nsim0 : int - IC realisation index of the simulation. - - Returns - ------- - str - """ + """Path to the files containing MAHs from random simulations.""" fdir = join(self.postdir, "random_mah") try_create_directory(fdir) @@ -441,26 +339,7 @@ class Paths: return join(fdir, fname) def field(self, kind, MAS, grid, nsim, simname): - r""" - 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 - """ + """Path to the files containing the calculated fields.""" if simname == "borg2": return join(self.borg2_dir, f"mcmc_{nsim}.h5") @@ -493,22 +372,8 @@ class Paths: def observer_peculiar_velocity(self, MAS, grid, nsim, simname): """ - Path to the files containing the observer peculiar velocity. - - Parameters - ---------- - MAS : str - Mass-assignment scheme. - grid : int - Grid size. - nsim : int - IC realisation index. - simname : str - Simulation name. - - Returns - ------- - str + Path to the files containing the observer peculiar velocity defined as + the velocity of the centre of the box. """ fdir = join(self.postdir, "environment") try_create_directory(fdir) diff --git a/csiborgtools/read/snapshot.py b/csiborgtools/read/snapshot.py index a2a4ffa..64da99c 100644 --- a/csiborgtools/read/snapshot.py +++ b/csiborgtools/read/snapshot.py @@ -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 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 import numpy @@ -62,59 +62,29 @@ class BaseSnapshot(ABC): @property def nsim(self): - """ - Simulation index. - - Returns - ------- - int - """ + """Simulation index.""" return self._nsim @property def nsnap(self): - """ - Snapshot index. - - Returns - ------- - int - """ + """Snapshot index.""" return self._nsnap @property def simname(self): - """ - Simulation name. - - Returns - ------- - str - """ + """Simulation name.""" if self._simname is None: raise ValueError("Simulation name not set.") return self._simname @property def boxsize(self): - """ - Simulation boxsize in `cMpc/h`. - - Returns - ------- - float - """ + """Simulation boxsize in `cMpc/h`.""" return simname2boxsize(self.simname) @property def paths(self): - """ - Paths manager. - - Returns - ------- - Paths - """ + """Paths manager.""" if self._paths is None: self._paths = Paths(**paths_glamdring) return self._paths @@ -124,10 +94,6 @@ class BaseSnapshot(ABC): """ Whether to keep the snapshot file open when reading halo particles. This is useful for repeated access to the snapshot. - - Returns - ------- - bool """ 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 coordinates are consistent with observations. - - Returns - ------- - bool """ return self._flip_xz - @abstractproperty + @property + @abstractmethod def coordinates(self): - """ - Return the particle coordinates. - - Returns - ------- - coords : 2-dimensional array - """ + """Particle coordinates.""" pass - @abstractproperty + @property + @abstractmethod def velocities(self): - """ - Return the particle velocities. - - Returns - ------- - vel : 2-dimensional array - """ + """Particle velocities.""" pass - @abstractproperty + @property + @abstractmethod def masses(self): - """ - Return the particle masses. - - Returns - ------- - mass : 1-dimensional array - """ + """Particle masses.""" pass - @abstractproperty + @property + @abstractmethod def particle_ids(self): - """ - Return the particle IDs. - - Returns - ------- - ids : 1-dimensional array - """ + """Particle IDs.""" pass @abstractmethod def halo_coordinates(self, halo_id, is_group): """ - Return the halo particle coordinates. + Halo particle coordinates. Parameters ---------- @@ -253,19 +195,12 @@ class BaseSnapshot(ABC): @abstractmethod def _make_hid2offset(self): - """ - Private class function to make the halo ID to offset dictionary. - """ pass def open_snapshot(self): """ Open the snapshot file, particularly used in the context of loading in particles of individual haloes. - - Returns - ------- - h5py.File """ if not self.keep_snapshot_open: # Check if the snapshot path is set @@ -283,10 +218,6 @@ class BaseSnapshot(ABC): def close_snapshot(self): """ Close the snapshot file opened with `open_snapshot`. - - Returns - ------- - None """ if not self.keep_snapshot_open: return @@ -648,24 +579,12 @@ class BaseField(ABC): @property def nsim(self): - """ - Simulation index. - - Returns - ------- - int - """ + """Simulation index.""" return self._nsim @property def paths(self): - """ - Paths manager. - - Returns - ------- - Paths - """ + """Paths manager.""" if self._paths is None: self._paths = Paths(**paths_glamdring) 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 coordinates are consistent with observations. - - Returns - ------- - bool """ return self._flip_xz @abstractmethod def density_field(self, MAS, grid): - """ - Return the pre-computed density field. - - Parameters - ---------- - MAS : str - Mass assignment scheme. - grid : int - Grid size. - - Returns - ------- - field : 3-dimensional array - """ + """Precomputed density field.""" pass @abstractmethod def velocity_field(self, MAS, grid): - """ - Return the pre-computed velocity field. - - Parameters - ---------- - MAS : str - Mass assignment scheme. - grid : int - Grid size. - - Returns - ------- - field : 4-dimensional array - """ + """Precomputed velocity field.""" pass @abstractmethod def radial_velocity_field(self, MAS, grid): - """ - Return the pre-computed radial velocity field. - - Parameters - ---------- - MAS : str - Mass assignment scheme. - grid : int - Grid size. - - Returns - ------- - field : 3-dimensional array - """ + """Precomputed radial velocity field.""" pass @@ -1046,5 +922,4 @@ def is_instance_of_base_snapshot_subclass(obj): Check if `obj` is an instance of a subclass of `BaseSnapshot`. """ return isinstance(obj, BaseSnapshot) and any( - issubclass(cls, BaseSnapshot) for cls in obj.__class__.__bases__ - ) + issubclass(cls, BaseSnapshot) for cls in obj.__class__.__bases__) diff --git a/csiborgtools/utils.py b/csiborgtools/utils.py index 685c96d..57c471d 100644 --- a/csiborgtools/utils.py +++ b/csiborgtools/utils.py @@ -14,6 +14,10 @@ # 51 Franklin Street, Fifth Floor, Boston, MA 02110-1301, USA. """ 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 datetime import datetime @@ -32,23 +36,7 @@ from scipy.stats import multivariate_normal @jit(nopython=True, fastmath=True, boundscheck=False) 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. - - 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,)` - """ + """Calculate the CM, assuming periodic boundary conditions in a cube.""" cm = np.zeros(3, dtype=particle_positions.dtype) totmass = sum(particles_mass) @@ -71,21 +59,8 @@ def center_of_mass(particle_positions, particles_mass, boxsize): @jit(nopython=True, fastmath=True, boundscheck=False) def periodic_distance(points, reference_point, boxsize): """ - Compute the 3D distance between multiple points and a reference point using - periodic boundary conditions. - - 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,)` + Compute the 3D distance between multiple points and a reference point + using periodic boundary conditions. """ npoints = len(points) @@ -99,20 +74,7 @@ def periodic_distance(points, reference_point, boxsize): @jit(nopython=True, fastmath=True, boundscheck=False) def periodic_distance_two_points(p1, p2, boxsize): - """ - 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 - """ + """Compute the 3D distance between two points in a periodic box.""" half_box = boxsize / 2 dist = 0 @@ -129,20 +91,7 @@ def periodic_distance_two_points(p1, p2, boxsize): @jit(nopython=True, boundscheck=False) def periodic_wrap_grid(pos, boxsize=1): - """ - 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)` - """ + """Wrap positions in a periodic box. Overwrites the input array.""" for n in range(pos.shape[0]): for i in range(3): if pos[n, i] > boxsize: @@ -156,16 +105,8 @@ def periodic_wrap_grid(pos, boxsize=1): @jit(nopython=True, fastmath=True, boundscheck=False) def delta2ncells(field): """ - Calculate the number of cells in `field` 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 + Calculate the number of cells in `field` of shape `(nx, ny, nz)` that are + non-zero. """ tot = 0 imax, jmax, kmax = field.shape @@ -178,23 +119,7 @@ def delta2ncells(field): def cartesian_to_radec(X, return_degrees=True, origin=[0., 0., 0.]): - """ - 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. - """ + """Calculate the radial distance, RA and deg.""" x, y, z = X[:, 0], X[:, 1], X[:, 2] 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) dec = np.arcsin(z / dist) ra = np.arctan2(y, x) + # Wrapping to ensure RA is in [0, 2pi) (later converted to degrees). ra[ra < 0] += 2 * np.pi if return_degrees: @@ -220,17 +146,8 @@ def cartesian_to_radec(X, return_degrees=True, origin=[0., 0., 0.]): def radec_to_cartesian(X): """ - Calculate Cartesian coordinates from radial distance, RA [0, 360) deg and - dec [-90, 90] deg. - - Parameters - ---------- - X : 2-dimensional array of shape `(npoints, 3)` - Spherical coordinates: distance, RA and dec. - - Returns - ------- - 2-dimensional array of shape `(npoints, 3)` + Calculate Cartesian coordinates from radial distance, RA and dec + `(npoints, 3)`. """ dist, ra, dec = X[:, 0], X[:, 1], X[:, 2] @@ -243,19 +160,7 @@ def radec_to_cartesian(X): 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 - """ + """Convert right ascension and declination to galactic coordinates.""" c = SkyCoord(ra=ra*u.degree, dec=dec*u.degree, frame='icrs') 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) def great_circle_distance(x1, x2, in_degrees=True): """ - Great circle distance between two points on a sphere, defined by RA and - dec, both in degrees. - - 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 + Great circle distance between two points, each of shape `(2,)`, specified + by RA an dec. """ ra1, dec1 = x1 ra2, dec2 = x2 @@ -303,19 +197,8 @@ def great_circle_distance(x1, x2, in_degrees=True): def cosine_similarity(x, y): r""" Calculate the cosine similarity between two Cartesian vectors. Defined - as :math:`\Sum_{i} x_i y_{i} / (|x| * |y|)`. - - 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 + as :math:`\Sum_{i} x_i y_{i} / (|x| * |y|)`. Optionally, `y` can be a + 2-dimensional array of shape `(n_samples, 3)`. """ if x.ndim != 1: 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): - """ - Convert hours, minutes and seconds to degrees. - - Parameters - ---------- - hours, minutes, seconds : float - - Returns - ------- - float - """ + """Convert hours, minutes and seconds to degrees.""" return hours * 15 + (minutes or 0) / 60 * 15 + (seconds or 0) / 3600 * 15 def dms_to_degrees(degrees, arcminutes=None, arcseconds=None): - """ - Convert degrees, arcminutes and arcseconds to decimal degrees. - - Parameters - ---------- - degrees, arcminutes, arcseconds : float - - Returns - ------- - float - """ + """Convert degrees, arcminutes and arcseconds to decimal degrees.""" return degrees + (arcminutes or 0) / 60 + (arcseconds or 0) / 3600 def real2redshift(pos, vel, observer_location, observer_velocity, boxsize, periodic_wrap=True, make_copy=True): r""" - Convert real-space position to redshift space position. + Convert real space position to redshift space position. Parameters ---------- @@ -461,20 +324,7 @@ def heliocentric_to_cmb(z_helio, RA, dec, e_z_helio=None): @jit(nopython=True, fastmath=True, boundscheck=False) def number_counts(x, bin_edges): - """ - Calculate counts of samples in bins. - - Parameters - ---------- - x : 1-dimensional array - Samples to bin. - bin_edges : 1-dimensional array - Bin edges. - - Returns - ------- - 1-dimensional array - """ + """Calculate counts of samples in bins.""" out = np.full(bin_edges.size - 1, np.nan, dtype=np.float32) for i in range(bin_edges.size - 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): """ - Calculate a binned statistic. - - 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 + Calculate a binned statistic, `statistic` must be a callable `f(x)`. """ 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 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() @@ -553,19 +379,9 @@ def calculate_acf(data): 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 + package written by Collin Capano. Algorithm used is from: N. Madras and + A.D. Sokal, J. Stat. Phys. 50, 109 (1988). """ - # 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 @@ -585,20 +401,8 @@ def calculate_acl(data): 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. + Thin MCMC samples (dictionary with arrays of shape `(nchains, nsamples)`) + by the autocorrelation length of each chain and concatenate the chains. """ keys = list(samples.keys()) nchains = 1 if samples[keys[0]].ndim == 1 else samples[keys[0]].shape[0]