clusters.py

import numpy as np, pyccl
from pixell import utils, bunch

# TODO
# 1. pyccl dependency is annoying. The main thing I need it for is chi(z) and d_A(z)
# 2. Maybe mark ProfileBattagliaFast as the default somehow? Or make a factory function?

class ProfileBase:
	"""Cluster profile evaluator base class. Defines the general
	interface and implements the simplest functions, but can't
	be used by itself. Instead use a subclass that defines the
	cluster pressure profile, like ProfileBattaglia"""
	def __init__(self, cosmology=None):
		self.cosmology = cosmology or pyccl.Cosmology(Omega_c=0.2589, Omega_b=0.0486, h=0.6774, sigma8=0.8159, n_s=0.9667, transfer_function="boltzmann_camb")
	def y    (self, m200, z, r=0, dist="angular"):
		"""Evaluate the (line-of-sight integrated) Compton y parameter
		for clusters with mass m200 and redshift z at a distance r
		from the center.
		* m200: M_200c in kg. M_200c is the mass inside R_200c, the
		  distance from the cluster center inside which the mean density
		  is 200 times the critical density.
		* z: The redshift
		* r: The distance from the cluster center
		* dist: The units of r
		  * "angular":  Angular distance in radians
		  * "physical": Physical distance in m
		  * "relative": Distance in units of R_200c
		m200, z and r can all be arrays, in which case their shapes
		must broadcast. For example, if you have nobj objects with
		masses m200[nobj] and redshifts z[nobj] and you want to evaluate
		all of them at the same set of distances r[nr], then use
		y(m200[:,None], z[:,None], r), which will return [nobj,nr]
		"""
		return self.Pe_los(m200, z, r=r, dist=dist) * (utils.sigma_T/(utils.m_e*utils.c**2))
	def Pe_los(self, m200, z, r=0, dist="angular"):
		"""Evaluate the line-of-sight-integrated electron pressure in Pa m.
		See y for the meaning of the arguments."""
		return 0.5176 * self.Pth_los(m200, z, r=r, dist=dist)
	def Pth_los(self, m200, z, r=0, dist="angular"):
		"""Evaluate the line-of-sight-integrated thermal pressure in Pa m.
		See y for the meaning of the arguments."""
		m200, z = [np.asanyarray(arr) for arr in [m200,z]]
		rho_c = calc_rho_c(z, self.cosmology)
		f_b   = self.cosmology["Omega_b"]/self.cosmology["Omega_m"]
		return utils.G * m200 * 200 * rho_c * f_b / 2 * self._raw_los(m200, z, r=r, dist=dist)
	def _raw_los(self, m200, z, r=0, dist="angular"): raise NotImplementedError
	def Pe(self, m200, z, r, dist="physical"):
		"""Evaluate the 3d electron pressure in Pa.
		See y for the meaning of the arguments."""
		return 0.5176 * self.Pth(m200, z, r=r, dist=dist)
	def Pth(self, m200, z, r=0, dist="physical"):
		"""Evaluate the 3d thermal pressure in Pa.
		See y for the meaning of the arguments."""
		m200, z = [np.asanyarray(arr) for arr in [m200,z]]
		r200  = calc_rdelta(m200, z, self.cosmology)
		rho_c = calc_rho_c(z, self.cosmology)
		f_b   = self.cosmology["Omega_b"]/self.cosmology["Omega_m"]
		return utils.G * m200 * 200 * rho_c * f_b / (2*r200) * self._raw(m200, z, r=r, dist=dist, r200=r200)
	def _raw(self, m200, z, r, dist="physical", r200=None): raise NotImplementedError
	def r2x(self, r200, z, r, dist="angular"):
		"""Transform distances to the relative units used internally"""
		r200, r, z = [np.asanyarray(arr) for arr in [r200, r, z]]
		if dist in ["rel",  "relative"]: return r
		if dist in ["ang",  "angular" ]:  return r/calc_angsize(r200, z, self.cosmology)
		if dist in ["phys", "physical"]:  return r/r200
		raise ValueError("Unrecognized distance type '%s'" % (str(dist)))

class ProfileBattaglia(ProfileBase):
	"""Battaglia cluster profile evaluator."""
	def _raw(self, m200, z, r, dist="physical", r200=None):
		"""Evaluate the dimensionless 3d pressure profiles for clusters
		with the the given masses m200 and redshifts z at the distances
		r from the center. m200, z and r must broadcast to the same shape."""
		m200, z, r = [np.asanyarray(arr) for arr in [m200, z, r]]
		if r200 is None: r200 = calc_rdelta(m200, z, self.cosmology)
		x      = self.r2x(m200, z, r, dist=dist)
		params = get_params_battaglia(m200, z, self.cosmology)
		return params.P0 * utils.gnfw(x, xc=params.xc, alpha=params.alpha, beta=params.beta, gamma=params.gamma)
	def _raw_los(self, m200, z, r=0, dist="angular"):
		"""Evaluate the dimensionless line-of-sight-integrated pressure
		profiles for clusters with the the given masses m200 and redshifts z
		at the distances r from the center. m200, z and r must broadcast
		to the same shape."""
		m200, z, r = np.broadcast_arrays(m200, z, r)
		shape      = m200.shape
		m200, z, r = [arr.reshape(-1) for arr in [m200,z,r]]
		n          = len(m200)
		# Compute our shape parameters and cluster size
		params = get_params_battaglia(m200, z, self.cosmology)
		r200   = calc_rdelta(m200, z, self.cosmology)
		x      = self.r2x(r200, z, r, dist=dist)
		res    = np.zeros(n, m200.dtype)
		for i in range(m200.size):
			res[i] = params.P0[i]*utils.tsz_profile_los_exact(x[i], xc=params.xc[i], alpha=params.alpha, beta=params.beta[i], gamma=params.gamma)
		return res.reshape(shape)

class ProfileBattagliaFast(ProfileBattaglia):
	def __init__(self, cosmology=None, beta_range=[-14,-3], obeta=6, nbeta=None, pad=0.1,
			alpha=1, gamma=-0.3, x1=1e-10, x2=1e8, npoint=200, zmax=1e5, _a=8):
		"""Initialize a Battaglia cluster profile evaluator. This exploits the
		Battaglia profile's constant alpha and gamma parameters to get away with
		a relatively simple interpolation scheme. This is further helped by xc
		only entering as a radial scaling, so the only real shape parameter is
		beta. We can therefore use a simple 1d interpolation over beta. Parameters:

		* cosmology: The pyccl cosmology object to use
		* beta_range: [beta_min, beta_max] used when building the interpolation
		* obeta: The interpolation order in beta
		* nbeta: The number of points to sample when building the beta interpolator.
		  Defaults to obeta.
		* pad: Log-factor to pad the beta range by. Defaults to 0.1
		* alpha, gamma: The alpha and gamma gnfw parameters. Must be consistent
		  with get_params_battaglia.
		* x1, x2: The minimal and maximal x = r/r200 to use in the radial
		  interpolator. This is done in log-space, so these should be positive
		  and separated by many orders of magnitude.
		* npoint: The number of points to use in the raidal interpolator.
		* zmax: The highest z = distance/r200 to consider in the line-of-sight
		  integral. This is done using a power-law-scaling, so large values are
		  ok.
		* _a: Implementation detail: Power-law scaling used to speed up
		  line-of-sight integral."""
		super().__init__(cosmology)
		# If the number of beta samples are not specified, use one point per basis function
		if nbeta is None: nbeta = obeta
		# The radial values to evlauate things at
		xp = np.linspace(np.log(x1),np.log(x2),npoint)
		# typically beta is */ 1.5. After log that's ± 0.4
		lbeta1, lbeta2 = np.log([-beta_range[1], -beta_range[0]])
		lbeta0 = 0.5*(lbeta2+lbeta1)
		dlbeta = 0.5*(lbeta2-lbeta1)+pad
		lbetas = np.linspace(lbeta0-dlbeta, lbeta0+dlbeta, nbeta)
		# Evaluate for all combinations
		yps = np.array([utils.tsz_profile_los_exact(np.exp(xp), xc=1, alpha=alpha, beta=-np.exp(lbeta), gamma=gamma, zmax=zmax, _a=_a) for lbeta in lbetas]) # [nbeta,npoint]
		yps = np.log(yps)
		# Build a polynomial interpolator in beta
		B = ((lbetas-lbeta0)/dlbeta)**(np.arange(obeta)[:,None]) # [obeta,nbeta]
		# Fit our data points to this basis
		a = np.linalg.solve(B.dot(B.T), B.dot(yps)) # [obeta,npoint]
		# Prefilter a to make utils.interpol faster
		a = utils.interpol_prefilter(a, npre=1, inplace=True)
		# We have everything we need for interpolation now
		self.lbeta0, self.dlbeta = lbeta0, dlbeta
		self.xp0,    self.dxp    = xp[0], xp[1]-xp[0]
		self.B, self.a, self.npoint, self.obeta = B, a, npoint, obeta
	def _raw(self, m200, z, r, dist="physical", r200=None):
		"""Evaluate the dimensionless 3d pressure profiles for clusters
		with the the given masses m200 and redshifts z at the distances
		r from the center. m200, z and r must broadcast to the same shape."""
		m200, z, r = [np.asanyarray(arr) for arr in [m200, z, r]]
		if r200 is None: r200 = calc_rdelta(m200, z, self.cosmology)
		x      = self.r2x(m200, z, r, dist=dist)
		params = get_params_battaglia(m200, z, self.cosmology)
		return params.P0 * utils.gnfw(x, xc=params.xc, alpha=params.alpha, beta=params.beta, gamma=params.gamma)
	def _raw_los(self, m200, z, r=0, dist="angular"):
		"""Evaluate the dimensionless line-of-sight-integrated pressure
		profiles for clusters with the the given masses m200 and redshifts z
		at the distances r from the center. m200, z and r must broadcast
		to the same shape."""
		m200, z, r = [np.asanyarray(arr) for arr in [m200, z, r]]
		# Compute our shape parameters and cluster size
		params = get_params_battaglia(m200, z, self.cosmology)
		r200   = calc_rdelta(m200, z, self.cosmology)
		# Nomralize beta so it matches what our polynomial basis expects
		b    = (np.log(-np.array(params.beta))-self.lbeta0)/self.dlbeta
		front= (slice(None),)+(None,)*(b.ndim or r.ndim)
		B    = b**(np.arange(self.obeta)[front]) # [obeta,...]
		# Get our dimensionless x values. Absorb xc into these since
		# we've built our interpolation around xc=1, but remember to compensate
		# for the unit change this implies before returning
		x    = self.r2x(r200*params.xc, z, r, dist=dist)
		xpix = (np.log(np.maximum(x,1e-10))-self.xp0)/self.dxp # [...]
		# This is inefficient if we have more interp points than output points
		ainter = utils.interpol(self.a, xpix[None], prefilter=False) # [obeta,...]
		profs  = np.sum(B*ainter,0) # [...]
		profs  = np.exp(profs)*params.xc*params.P0
		return profs

# Battaglia 12 says:
#
# Δ = 200
#
# Pe  = (2X_H+2)/(5H_H+3) * Pth = 0.5176 Pth
# Pth =
#  P0 * (x/xc)**γ * [(1+(x/xc)**α]**-β' *
#  G M_Δ Δ ρ_cr(z) f_b / (2 R_Δ)
#
# P0(z) = 18.1  * m** 0.154   * (1+z)**-0.758
# xc(z) = 0.497 * m**-0.00865 * (1+z)** 0.731
# β'(z) = 4.35  * m** 0.0393  * (1+z)** 0.415
# α     = 1.00
# γ     = -0.3
#
# m = M_200/(1e14 M_sun)
#
# Relation to standard gnfw parameters:
# β = γ - αβ'
#
# y = k_B sigma_T/(m_e c²) * int dl n_e T_e
#   = sigma_T/(m_e c²) * int dl P_e
#
# Relation of dl and dx, which is what I have in my LOS integral:
# x = l/R_Δ

def websky_pkcs_nhalo(fname):
	"""Read rows offset:offset+num of raw data from the given pkcs file.
	if num==0, all values are read"""
	with open(fname, "r") as ifile:
		return np.fromfile(ifile, count=3, dtype=np.uint32)[0]

def websky_pkcs_read(fname, num=0, offset=0):
	"""Read rows offset:offset+num of raw data from the given pkcs file.
	if num==0, all values are read"""
	with open(fname, "r") as ifile:
		n   = np.fromfile(ifile, count=3, dtype=np.uint32)[0]-offset
		if num: n = num
		cat = np.fromfile(ifile, count=n*10, offset=offset*10*4, dtype=np.float32).reshape(n, 10)
		return cat

# This function seems to be just as good as mass_interp, but much simpler
def websky_m200m_to_m200c(m200m, z, cosmology):
	omegam  = cosmology["Omega_m"]
	omegamz = omegam*(1+z)**3/(omegam*(1+z)**3+1-omegam)
	m200c   = omegamz**0.35 * m200m # m200m to m200c conversion used for websky
	return m200c

def websky_decode(data, cosmology, mass_interp):
	"""Go from a raw websky catalog to pos, z and m200"""
	chi     = np.sum(data.T[:3]**2,0)**0.5 # comoving Mpc
	a       = pyccl.scale_factor_of_chi(cosmology, chi)
	z       = 1/a-1
	R       = data.T[6].astype(float) * 1e6*utils.pc # m. This is *not* r200!
	rho_m   = calc_rho_c(0, cosmology)*cosmology["Omega_m"]
	m200m   = 4/3*np.pi*rho_m*R**3
	m200    = mass_interp(m200m, z)
	ra, dec = utils.rect2ang(data.T[:3])
	return bunch.Bunch(z=z, ra=ra, dec=dec, m200=m200)

def get_H0(cosmology): return cosmology["h"]*100*1e3/(1e6*utils.pc)

def get_H(z, cosmology):
	z = np.asanyarray(z)
	return get_H0(cosmology)*pyccl.h_over_h0(cosmology, 1/(z.reshape(-1)+1)).reshape(z.shape)

def calc_rho_c(z, cosmology):
	H     = get_H(z, cosmology)
	rho_c = 3*H**2/(8*np.pi*utils.G)
	return rho_c

def calc_rho_m(z, cosmology):
	rho_c = calc_rho_c(z, cosmology)
	rho_m = rho_c * cosmology["Omega_m"]
	return rho_m

def calc_rdelta(mdelta, z, cosmology, delta=200, type="critical"):
	"""Given M_delta in kg, returns R_delta in m"""
	if   type == "critical": rho = calc_rho_c(z, cosmology)
	elif type == "mean":     rho = calc_rho_m(z, cosmology)
	else: raise ValueError("Unknown density type '%s'" % str(type))
	rdelta = (mdelta/(4/3*np.pi*delta*rho))**(1/3)
	return rdelta

def calc_angsize(physsize, z, cosmology):
	"""Given a physical size in m, returns the angular size in radians"""
	z   = np.asanyarray(z)
	d_A = pyccl.angular_diameter_distance(cosmology, 1/(z.reshape(-1)+1)).reshape(z.shape)*1e6*utils.pc
	return np.arctan2(physsize,d_A)

def get_params_battaglia(m200, z, cosmology):
	"""Return a bunch of xc, alpha, beta, gamma for a cluster with
	the given m200 in SI units."""
	# First get the gnfw parameters. utils.gnfw has the opposite sign for
	# beta and gamma as nemo, but otherwise the same convention
	z1    = z+1
	m     = m200/(1e14*utils.M_sun)
	P0    =  18.1     * m** 0.154   * z1**-0.758
	xc    =  0.497    * m**-0.00865 * z1** 0.731
	beta  =  4.35     * m** 0.0393  * z1** 0.415
	alpha =  1; gamma = -0.3
	# Go from battaglia convention to standard gnfw
	beta  = gamma - alpha*beta
	return bunch.Bunch(xc=xc, alpha=alpha, beta=beta, gamma=gamma, P0=P0)

class MdeltaTranslator:
	def __init__(self, cosmology,
			type1="matter", delta1=200, type2="critical", delta2=200,
			zlim=[0,20], mlim=[1e11*utils.M_sun,5e16*utils.M_sun], step=0.1):
		"""Construct a functor that translates from one M_delta defintion to
		another.
		* type1, type2: Type of M_delta, e.g. m200c vs m200m.
		  * "matter": The mass inside the region where the average density is
		    delta times higher than the current matter density
		  * "critical": The same, but for the critical density instead. This
		    differs due to the presence of dark energy.
		* delta1, delta2: The delta value used in type1, type2.
		* zlim: The z-range to build the interpolator for.
		* mlim: The Mass range to build the interpolator for, in kg
		* step: The log-spacing of the interpolators.
		Some combinations of delta and type may not be supported, limited by
		support in pyccl. The main thing this object does beyond pyccl is to
		allow one to vectorize over both z and mass."""
		idef = pyccl.halos.MassDef(delta1, type1, c_m_relation="Bhattacharya13")
		odef = pyccl.halos.MassDef(delta2, type2, c_m_relation="Bhattacharya13")
		# Set up our sample grid, which will be log-spaced in both z and mass direction
		lz1, lz2 = np.log(1+np.array(zlim)) # lz = log(1+z) = -log(a)
		lm1, lm2 = np.log(np.array(mlim))   # lm = log(m)
		nz  = utils.ceil((lz2-lz1)/step)
		nm  = utils.ceil((lm2-lm1)/step)
		lzs = np.linspace(lz1, lz2, nz)
		lms = np.linspace(lm1, lm2, nm)
		olms = np.zeros((len(lzs),len(lms)))
		for ai, lz in enumerate(lzs):
			moo = np.exp(lms[-1])/utils.M_sun
			olms[ai] = idef.translate_mass(cosmology, np.exp(lms)/utils.M_sun, np.exp(-lz), odef)
		olms = np.log(olms*utils.M_sun)
		olms = utils.interpol_prefilter(olms, order=3)
		# Save parameters
		self.lz1, self.lz2, self.dlz = lz1, lz2, (lz2-lz1)/(nz-1)
		self.lm1, self.lm2, self.dlm = lm1, lm2, (lm2-lm1)/(nm-1)
		self.olms = olms
	def __call__(self, m, z):
		zpix = (np.log(1+np.array(z))-self.lz1)/self.dlz
		mpix = (np.log(m)-self.lm1)/self.dlm
		pix  = np.array([zpix,mpix])
		return np.exp(utils.interpol(self.olms, pix, order=3, prefilter=False))