mirror of
https://github.com/Richard-Sti/csiborgtools.git
synced 2024-12-23 01:58:03 +00:00
ce55a2b47e
* Parallelize over simulations * Update docs * Update dependency * Update imports * Add adtitional dependencies * Update .gitignore * Update ERADME * Simplify numpyro GOF * Speed up GOF * Deepcopy samples * Update scripts * Add GPU acceleration * Select boxes * Update script * Optionally sample beta * Fix old code * Simplify code * Start saving log posterior * Start popping log_likeliood * Add imports * Add converting samples * Fix sctipt name * Add evidence with harmonic * Remove comment * Update imports * Update imports so that pylians not required * Stop requiring Pylians to be installed * Update submission scripts for loops * Update nb * Update nb * Add Manticore boxes * Add verbosity flag * Add bulk flow * Update script * Update nb * Update normalization * Update submit * Update nb
252 lines
8.3 KiB
Python
252 lines
8.3 KiB
Python
# Copyright (C) 2023 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.
|
|
"""
|
|
Code to calculate the enclosed mass, momentum or bulk flow from various
|
|
radial-velocity based estimators (that may be flawed, be careful).
|
|
"""
|
|
import numpy as np
|
|
from numba import jit
|
|
from tqdm import tqdm
|
|
|
|
|
|
###############################################################################
|
|
# Enclosed mass at each distance from particles #
|
|
###############################################################################
|
|
|
|
|
|
@jit(nopython=True, boundscheck=False)
|
|
def _enclosed_mass(rdist, mass, rmax, start_index):
|
|
enclosed_mass = 0.
|
|
|
|
for i in range(start_index, len(rdist)):
|
|
if rdist[i] <= rmax:
|
|
enclosed_mass += mass[i]
|
|
else:
|
|
break
|
|
|
|
return enclosed_mass, i
|
|
|
|
|
|
def particles_enclosed_mass(rdist, mass, distances):
|
|
"""
|
|
Calculate the enclosed mass at each distance from a set of particles. Note
|
|
that the particles must be sorted by distance from the center of the box.
|
|
|
|
Parameters
|
|
----------
|
|
rdist : 1-dimensional array
|
|
Sorted distance of particles from the center of the box.
|
|
mass : 1-dimensional array
|
|
Sorted mass of particles.
|
|
distances : 1-dimensional array
|
|
Distances at which to calculate the enclosed mass.
|
|
|
|
Returns
|
|
-------
|
|
enclosed_mass : 1-dimensional array
|
|
Enclosed mass at each distance.
|
|
"""
|
|
enclosed_mass = np.full_like(distances, 0.)
|
|
start_index = 0
|
|
for i, dist in enumerate(distances):
|
|
if i > 0:
|
|
enclosed_mass[i] += enclosed_mass[i - 1]
|
|
|
|
m, start_index = _enclosed_mass(rdist, mass, dist, start_index)
|
|
enclosed_mass[i] += m
|
|
|
|
return enclosed_mass
|
|
|
|
|
|
###############################################################################
|
|
# Enclosed mass from a density field #
|
|
###############################################################################
|
|
|
|
|
|
@jit(nopython=True)
|
|
def _cell_rdist(i, j, k, Ncells, boxsize):
|
|
"""Radial distance of the center of a cell from the center of the box."""
|
|
xi = boxsize / Ncells * (i + 0.5) - boxsize / 2
|
|
yi = boxsize / Ncells * (j + 0.5) - boxsize / 2
|
|
zi = boxsize / Ncells * (k + 0.5) - boxsize / 2
|
|
|
|
return (xi**2 + yi**2 + zi**2)**0.5
|
|
|
|
|
|
@jit(nopython=True, boundscheck=False)
|
|
def _field_enclosed_mass(field, rmax, boxsize):
|
|
Ncells = field.shape[0]
|
|
cell_volume = (1000 * boxsize / Ncells)**3
|
|
|
|
mass = 0.
|
|
volume = 0.
|
|
for i in range(Ncells):
|
|
for j in range(Ncells):
|
|
for k in range(Ncells):
|
|
if _cell_rdist(i, j, k, Ncells, boxsize) < rmax:
|
|
mass += field[i, j, k]
|
|
volume += 1.
|
|
|
|
return mass * cell_volume, volume * cell_volume
|
|
|
|
|
|
def field_enclosed_mass(field, distances, boxsize, verbose=True):
|
|
"""
|
|
Calculate the approximate enclosed mass within a given radius from a
|
|
density field, counts the mass in cells and volume of cells whose
|
|
centers are within the radius.
|
|
|
|
Parameters
|
|
----------
|
|
field : 3-dimensional array
|
|
Density field in units of `h^2 Msun / kpc^3`.
|
|
rmax : 1-dimensional array
|
|
Radii to calculate the enclosed mass at in `Mpc / h`.
|
|
boxsize : float
|
|
Box size in `Mpc / h`.
|
|
verbose : bool
|
|
Verbosity flag.
|
|
|
|
Returns
|
|
-------
|
|
enclosed_mass : 1-dimensional array
|
|
Enclosed mass at each distance.
|
|
enclosed_volume : 1-dimensional array
|
|
Enclosed grid-like volume at each distance.
|
|
"""
|
|
enclosed_mass = np.zeros_like(distances)
|
|
enclosed_volume = np.zeros_like(distances)
|
|
|
|
for i, dist in enumerate(tqdm(distances, disable=not verbose)):
|
|
enclosed_mass[i], enclosed_volume[i] = _field_enclosed_mass(
|
|
field, dist, boxsize)
|
|
|
|
return enclosed_mass, enclosed_volume
|
|
|
|
|
|
###############################################################################
|
|
# Enclosed momentum at each distance from particles #
|
|
###############################################################################
|
|
|
|
|
|
@jit(nopython=True, boundscheck=False)
|
|
def _enclosed_momentum(rdist, mass, vel, rmax, start_index):
|
|
bulk_momentum = np.zeros(3, dtype=rdist.dtype)
|
|
|
|
for i in range(start_index, len(rdist)):
|
|
if rdist[i] <= rmax:
|
|
bulk_momentum += mass[i] * vel[i]
|
|
else:
|
|
break
|
|
|
|
return bulk_momentum, i
|
|
|
|
|
|
def particles_enclosed_momentum(rdist, mass, vel, distances):
|
|
"""
|
|
Calculate the enclosed momentum at each distance. Note that the particles
|
|
must be sorted by distance from the center of the box.
|
|
|
|
Parameters
|
|
----------
|
|
rdist : 1-dimensional array
|
|
Sorted distance of particles from the center of the box.
|
|
mass : 1-dimensional array
|
|
Sorted mass of particles.
|
|
vel : 2-dimensional array
|
|
Sorted velocity of particles.
|
|
distances : 1-dimensional array
|
|
Distances at which to calculate the enclosed momentum.
|
|
|
|
Returns
|
|
-------
|
|
bulk_momentum : 2-dimensional array
|
|
Enclosed momentum at each distance.
|
|
"""
|
|
bulk_momentum = np.zeros((len(distances), 3))
|
|
start_index = 0
|
|
for i, dist in enumerate(distances):
|
|
if i > 0:
|
|
bulk_momentum[i] += bulk_momentum[i - 1]
|
|
|
|
v, start_index = _enclosed_momentum(rdist, mass, vel, dist,
|
|
start_index)
|
|
bulk_momentum[i] += v
|
|
|
|
return bulk_momentum
|
|
|
|
|
|
###############################################################################
|
|
# Bulk flow estimators #
|
|
###############################################################################
|
|
|
|
|
|
def bulkflow_peery2018(rdist, mass, pos, vel, distances, weights,
|
|
verbose=True):
|
|
"""
|
|
Calculate the bulk flow from a set of particles using the estimator from
|
|
Peery+2018. Supports either `1/r^2` or `constant` weights. Particles
|
|
are assumed to be sorted by distance from the center of the box.
|
|
|
|
Parameters
|
|
----------
|
|
rdist : 1-dimensional array
|
|
Sorted distance of particles from the center of the box.
|
|
mass : 1-dimensional array
|
|
Sorted mass of particles.
|
|
pos : 2-dimensional array
|
|
Sorted position of particles.
|
|
vel : 2-dimensional array
|
|
Sorted velocity of particles.
|
|
distances : 1-dimensional array
|
|
Distances at which to calculate the bulk flow.
|
|
weights : str
|
|
Weights to use in the estimator, either `1/r^2` or `constant`.
|
|
verbose : bool
|
|
Verbosity flag.
|
|
|
|
Returns
|
|
-------
|
|
bulk_flow : 2-dimensional array
|
|
"""
|
|
# Select only the particles within the maximum distance to speed up the
|
|
# calculation.
|
|
if verbose:
|
|
print("Selecting particles within the maximum distance...")
|
|
kmax = np.searchsorted(rdist, np.max(distances))
|
|
rdist = rdist[:kmax]
|
|
mass = mass[:kmax]
|
|
pos = pos[:kmax]
|
|
vel = vel[:kmax]
|
|
|
|
if verbose:
|
|
print("Computing the cumulative quantities...")
|
|
if weights == "1/r^2":
|
|
cumulative_x = np.cumsum(mass[:, np.newaxis] * np.sum(vel * pos, axis=1)[:, np.newaxis] * pos / rdist[:, np.newaxis]**4, axis=0) # noqa
|
|
norm = lambda R: R**2 # noqa
|
|
elif weights == "constant":
|
|
cumulative_x = np.cumsum(mass[:, np.newaxis] * np.sum(vel * pos, axis=1)[:, np.newaxis] * pos / rdist[:, np.newaxis]**2, axis=0) # noqa
|
|
norm = lambda R: 3. # noqa
|
|
else:
|
|
raise ValueError("Invalid weights.")
|
|
cumulative_x /= np.cumsum(mass)[:, np.newaxis]
|
|
|
|
B = np.zeros((len(distances), 3))
|
|
for i in range(3):
|
|
for j, R in enumerate(distances):
|
|
k = np.searchsorted(rdist, R)
|
|
B[j, i] = norm(R) * cumulative_x[k - 1, i]
|
|
|
|
return B
|