cosmo_tracer.py

"""
coding: utf-8

Author : Edmond Chaussidon (CEA)

Code for Tracer modelisation: class tracer
"""

import os
import numpy as np
from scipy.interpolate import interp1d

from cosmoprimo.fiducial import DESI


# Reference cosmology --> LOAD DESI cosmology (fiducial cosmology for DESI measurment)
c_fid = DESI(engine='class')
# Compute background with engine as engine for computation (already set)
c_fid.get_background()
# Initialize the Fourier class to compute the power spectrum from a background
c_fid.get_fourier()


class DN_DZ(object):
    """
    Class to compute (interpolating) the dn_dz from different input.
    """

    def __init__(self, filename):
        """
        Initialize :class:`DN_DZ`.

        Ce n'est qu'un debut pour l'instant, on ne lit qu'a partir d'un fichier txt que j'ai construit

        Parameters
        ----------

        """

        self.filename = filename
        if self.filename.split('.')[-1] == 'txt':
            self.input_from_txt(self.filename)
        elif self.filename.split('.')[-1] == 'dat':
            print("ATTENTION POUR L'INSTANT C'est n'importe quoi !! -> on ne peut aps vraiment interpoler un nz normalise, en tout cas il faut le renormaliser ensuite !!")
            self.input_from_dat(self.filename)
        else:
            print("Filename has unexpected format")

    def input_from_txt(self, filename):
        """
        LOAD n(z) normalized from a .txt file ! see: f_NL/Compute_n_z.ipynb

        Parameters
        ----------
        filename : path of the .txt file

        """
        d = np.loadtxt(filename)
        self.interp = interp1d(d[0], d[1], kind='quadratic', bounds_error=False, fill_value=(0, 0))

    def input_from_dat(self, filename):
        """
        LOAD n(z) normalized from .dat file generated by Christophe as official KP1 product.

        Parameters
        ----------
        filename : path of the .dat file

        """
        d = np.loadtxt(filename, skiprows=1).T
        self.interp = interp1d((d[0] + d[1]) / 2, d[2], kind='quadratic', bounds_error=False, fill_value=(0, 0))

    def __call__(self, z):
        """Return (interpolated) dn_dz at redshift ``z`` (scalar or array)."""
        return self.interp(z)


class Tracer(object):
    """
    Implementation of the Tracer which probes the matter in the Universe
    """

    def __init__(self, name, cosmo, bias, pop, z0, dn_dz, area, density_deg2, z_width, shot_noise_limited):
        """
        Initialize :class:`Tracer`

        Parameters
        ----------
        name : name of the tracer, useful for legend
        cosmo : cosmo class from cosmoprimo
        bias : float
               bias of the consider tracer
        pop : float
              parameter to describe if the tracor is due to recent merger or not. Should be 1 (old merger) < pop < 1.6 (recent merger)
        z0 : float
            mean redshift of the sample of the consider tracer
        dn_dz : callable function
                function which describes the n(z) of the tracer
        area : float
               Surface in deg2 of the observation of the tracer
        density_deg2 : float
                       density of the tracer in deg2
        z_width : float
                  effective Delta_z of the dn_dz
        shot_noise_limited : bool
                             Are you in the shot noise regime ? Only usefull to compute sigma_P in cosmo_ps
        """

        self.name = name

        # fiducial cosmology
        self.cosmo = cosmo

        # tracor parameters
        self.z0 = z0
        self.dn_dz = dn_dz

        self.bias = bias
        # Only true for fnl=0 !!
        self.beta = self.cosmo.get_background().growth_rate(self.z0) / self.bias
        self.pop = pop

        # Compute b_phi (paramtrization from Slozar et al. 2008)
        # delta_c = 1.686  # the spherical collapse linear over-density
        self.bias_phi = 2 * 1.686 * (self.bias - self.pop)

        # Survey information
        self.area = area
        self.density_deg2 = density_deg2
        self.z_width = z_width

        self.V_survey = self.Volume()
        self.n_survey = self.Density()

        self.shot_noise_limited = shot_noise_limited

    def __str__(self):
        string = "\nBuild Tracer with the following parameters:"
        string += f"\n    * name: {self.name}"
        string += f"\n    * bias: {self.bias} -->  beta (constant value for fnl=0 only !): {self.beta:2.2f}"
        string += f"\n    * pop: {self.pop} -->  bias_phi: {self.bias_phi:2.3f}"
        string += f"\n    * z0: {self.z0:2.3f}"
        string += f"\n    * z_width: {self.z_width}"
        string += f"\n    * area: {self.area}"
        string += f"\n    * density_deg2: {self.density_deg2}"
        string += f"\n    * Survey Volume (Gpc/h)^3 = {self.V_survey/1.0e9:2.2f}"
        string += f"\n    * Survey density =  {self.n_survey:.2e}"
        string += f"\n    * Is in shoot noise limited region ? {self.shot_noise_limited} (only useful for sigma_P)\n"
        return string

    def Volume(self):
        f_sky = self.area / (4 * 180. * 180. / np.pi)
        V = 4.0 / 3.0 * np.pi * f_sky
        V *= self.cosmo.get_background().comoving_radial_distance(self.z0 + self.z_width / 2.0)**3 - self.cosmo.get_background().comoving_radial_distance(self.z0 - self.z_width / 2.0)**3
        return V

    def Density(self):
        n = self.area * self.density_deg2 / self.V_survey
        return n

    def __copy__(self):
        """
        Proper way to copy the class
        """
        new = self.__class__.__new__(self.__class__)
        new.__dict__.update(self.__dict__)
        return new

    def copy(self, **kwargs):  # super malin ca !
        new = self.__copy__()
        new.__dict__.update(kwargs)

        # On oublie pas de mettre à jours les quantites que l'on calcul pendant l'initialisation
        new.beta = new.cosmo.get_background().growth_rate(new.z0) / new.bias
        new.bias_phi = 2 * 1.686 * (new.bias - new.pop)
        new.V_survey = new.Volume()
        new.n_survey = new.Density()

        return new


# ---------------------------------------------------------------------------------------------------- #
def LRG_tracer():
    """
    Define Standard DESI LRG tracer
    """
    # Param for LRG as tracer:
    bias = 2.3  # https://arxiv.org/pdf/1607.05383.pdf
    pop = 1.0
    z0 = 0.7
    z_width = 1.0

    # Survey info:
    density_deg2 = 500.0
    Area = 14000  # DESI geometry
    # Shot noise limited regime ?
    shot_noise_limited = True

    # Load n(z):
    dn_dz = DN_DZ(os.path.join(os.path.dirname(__file__), 'Data/dn_dz_lrg.txt'))

    return Tracer("LRG", c_fid, bias, pop, z0, dn_dz, Area, density_deg2, z_width, shot_noise_limited)


def ELG_VLO_tracer():
    """
    Define Standard DESI LRG tracer
    """
    # Param for LRG as tracer:
    bias = 1.3
    pop = 1.0
    z0 = 0.98
    z_width = 1.1

    # Survey info:
    density_deg2 = 500.0
    Area = 14000  # DESI geometry
    # Shot noise limited regime ?
    shot_noise_limited = True

    # Load n(z):
    dn_dz = DN_DZ(os.path.join(os.path.dirname(__file__), 'Data/dn_dz_elg_lp.txt'))

    return Tracer("LRG", c_fid, bias, pop, z0, dn_dz, Area, density_deg2, z_width, shot_noise_limited)


def QSO_tracer():
    """
    Define Standard DESI QSO tracer
    """
    # Param for QSO as tracer:
    bias = 2.5
    pop = 1.0
    z0 = 1.7
    z_width = 1.4

    # Survey info:
    density_deg2 = 200.0
    Area = 14000  # DESI geometry
    # Shot noise limited regime ?
    shot_noise_limited = True

    # Load n(z):
    dn_dz = DN_DZ(os.path.join(os.path.dirname(__file__), 'Data/dn_dz_qso.txt'))

    return Tracer("QSO", c_fid, bias, pop, z0, dn_dz, Area, density_deg2, z_width, shot_noise_limited)


def LBG_tracer():
    """
    Define LBG tracer
    """
    # Param for LBG tracer:
    bias = 3.2
    pop = 1.0
    z0 = 2.76
    z_width = 1.4

    # Survey info:
    density_deg2 = 500.0
    Area = 10000
    # Shot noise limited regime ?
    shot_noise_limited = True

    # Load n(z):
    dn_dz = DN_DZ(os.path.join(os.path.dirname(__file__), 'Data/dn_dz_lbg.txt'))

    return Tracer("LBG", c_fid, bias, pop, z0, dn_dz, Area, density_deg2, z_width, shot_noise_limited)


def BGS_tracer():
    """
    Define BGS tracer
    """
    # Param for BGS tracer:
    bias = 1.3
    pop = 1.0
    z0 = 0.25
    z_width = 0.5

    # Survey info:
    density_deg2 = 1400.0
    Area = 14000
    # Shot noise limited regime ?
    shot_noise_limited = True

    # Load n(z):
    dn_dz = DN_DZ(os.path.join(os.path.dirname(__file__), 'Data/nz_bgs_final.dat'))

    return Tracer("BGS", c_fid, bias, pop, z0, dn_dz, Area, density_deg2, z_width, shot_noise_limited)